Add API to get ep graph partitioning info (#26781)

### Description
- Adds API functions to get information about the subgraphs/nodes
assigned to the EPs in the session.
- `Session_GetEpGraphAssignmentInfo`: Returns a list of "subgraphs",
each with information about the assigned EP and nodes.
- Note: App must enable session configuration
`"session.record_ep_graph_assignment_info"` to signal ORT to collect
this information. If not enabled, API returns empty results.
- `EpAssignedSubgraph_GetEpName`: Returns the name of the EP to which
the subgraph is assigned
  - `EpAssignedSubgraph_GetNodes`: Returns a list of assigned nodes
  - `EpAssignedNode_GetName`: Returns the assigned node's name
  - `EpAssignedNode_GetDomain`: Returns the assigned node's domain
- `EpAssignedNode_GetOperatorType`: Returns the assigned node's operator
type
- Also adds C++ and Python bindings

#### Structure of returned information
The API returns a list of "subgraphs". Each subgraph has the following
information:
- Subgraph info:
- EP name: The name of the execution provider to which this subgraph is
assigned.
- nodes: Name and operator type of each node. Ex: `[{"multiply", "Mul"},
...]`

Python example program (taken from unit tests):
```python
    def test_get_graph_provider_assignment_info(self):
        """
        Tests querying for information about the nodes assigned to the CPU EP.
        """

        # Create session options that enables recording EP graph partitioning info.
        session_options = onnxrt.SessionOptions()
        session_options.add_session_config_entry("session.record_ep_graph_assignment_info", "1")

        session = onnxrt.InferenceSession(get_name("add_mul_add.onnx"), sess_options=session_options)

        # Query session for information on each subgraph assigned to an EP.
        ep_subgraphs = session.get_provider_graph_assignment_info()

        # Check that all 3 nodes are assigned to CPU EP (each in its own subgraph)
        self.assertEqual(len(ep_subgraphs), 3)
        for ep_subgraph in ep_subgraphs:
            self.assertEqual(ep_subgraph.ep_name, "CPUExecutionProvider")
            self.assertEqual(len(ep_subgraph.get_nodes()), 1)

        # Serialize each node to an identifier (concatenates operator type and node name)
        node_ids: list[str] = [f"{n.op_type}/{n.name}" for s in ep_subgraphs for n in s.get_nodes()]

        # Should have 1 Mul and 2 Adds.
        self.assertEqual(len(node_ids), 3)
        self.assertIn("Add/add_0", node_ids)
        self.assertIn("Add/add_1", node_ids)
        self.assertIn("Mul/mul_0", node_ids)
```

C++ program (taken from unit test):
```c++
  // Check the ep graph partitioning (Mul on plugin EP, others on CPU EP).
  // Model has 3 subgraphs (in no particular order):
  // - Subgraph 1: Add assigned to CPU EP.
  // - Subgraph 2: Mul assigned to plugin EP.
  // - Subgraph 3: Add assigned to CPU EP.
  std::vector<Ort::ConstEpAssignedSubgraph> ep_subgraphs = session.GetEpGraphAssignmentInfo();
  ASSERT_EQ(ep_subgraphs.size(), 3);

  for (Ort::ConstEpAssignedSubgraph ep_subgraph : ep_subgraphs) {
    std::string ep_name = ep_subgraph.EpName();
    ASSERT_TRUE(ep_name == Utils::example_ep_info.ep_name || ep_name == kCpuExecutionProvider);

    const std::vector<Ort::ConstEpAssignedNode> ep_nodes = ep_subgraph.GetNodes();
    ASSERT_GE(ep_nodes.size(), 1);  // All of these subgraphs just have one node.

    if (ep_name == kCpuExecutionProvider) {
      std::string op_type = ep_nodes[0].OpType();
      std::string node_name = ep_nodes[0].Name();

      ASSERT_EQ(op_type, "Add");
      ASSERT_TRUE(node_name == "add_0" || node_name == "add_1");
    } else {
      ASSERT_TRUE(ep_name == Utils::example_ep_info.ep_name);

      std::string op_type = ep_nodes[0].OpType();
      std::string node_name = ep_nodes[0].Name();
      ASSERT_EQ(op_type, "Mul");
      ASSERT_EQ(node_name, "mul_0");
    }
  }
```

### Motivation and Context
<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: adrianlizarraga

include/onnxruntime/core/session/onnxruntime_c_api.h

@@ -337,6 +337,8 @@ ORT_RUNTIME_CLASS(ExternalResourceImporter);  // Capability object for external
 ORT_RUNTIME_CLASS(ExternalMemoryHandle);      // EP-imported view of shared external allocation
 ORT_RUNTIME_CLASS(ExternalSemaphoreHandle);   // EP-imported view of shared external semaphore
 ORT_RUNTIME_CLASS(DeviceEpIncompatibilityDetails);
+ORT_RUNTIME_CLASS(EpAssignedSubgraph);
+ORT_RUNTIME_CLASS(EpAssignedNode);
 
 #ifdef _MSC_VER
 typedef _Return_type_success_(return == 0) OrtStatus* OrtStatusPtr;
@@ -7016,6 +7018,101 @@ struct OrtApi {
    * \since Version 1.24
    */
   ORT_API2_STATUS(CreateEnvWithOptions, _In_ const OrtEnvCreationOptions* options, _Outptr_ OrtEnv** out);
+
+  /** \brief Get information about the subgraphs assigned to each execution provider (EP) and the nodes within.
+   *
+   * Each returned OrtEpAssignedSubgraph instance contains details of the subgraph/nodes assigned to an execution
+   * provider, including the execution provider's name, and the name, domain, and operator type for every node.
+   *
+   * For compiling execution providers, a single OrtEpAssignedSubgraph instance contains information about the
+   * nodes that are fused and compiled within a single subgraph assigned to the execution provider.
+   *
+   * For execution providers that use kernel registration (e.g., CPU EP), each node with a registered kernel is
+   * contained in its own OrtEpAssignedSubgraph instance.
+   *
+   * \note The caller must enable the collection of this information by enabling the session
+   *       configuration entry "session.record_ep_graph_assignment_info" during session creation.
+   *       Refer to onnxruntime_session_options_config_keys.h. Otherwise, if not enabled, this function returns a
+   *       status with error code ORT_FAIL.
+   *
+   * \note The information reported by this function is obtained immediately after running basic optimizations on the
+   *       original graph if the session optimization level is set to ORT_ENABLE_BASIC or higher. If the session
+   *       optimization level is set to ORT_DISABLE_ALL, only minimal/required optimizations are run before
+   *       the information is collected.
+   *
+   * \param[in] session The OrtSession instance.
+   * \param[out] ep_subgraphs Output parameter set to the array of OrtEpAssignedSubgraph instances.
+   * \param[out] num_ep_subgraphs Output parameter set to the number of elements in the `ep_subgraphs` array.
+   *
+   * \snippet{doc} snippets.dox OrtStatus Return Value
+   *
+   * \since Version 1.24.
+   */
+  ORT_API2_STATUS(Session_GetEpGraphAssignmentInfo, _In_ const OrtSession* session,
+                  _Outptr_ const OrtEpAssignedSubgraph* const** ep_subgraphs,
+                  _Out_ size_t* num_ep_subgraphs);
+
+  /** \brief Get the name of the execution provider to which the subgraph was assigned.
+   *
+   * \param[in] ep_subgraph The OrtEpAssignedSubgraph instance.
+   * \param[out] out Output parameter set to the execution provider's name as a UTF-8 null-terminated string.
+   *                 Owned by the OrtEpAssignedSubgraph instance (do not free).
+   *
+   * \snippet{doc} snippets.dox OrtStatus Return Value
+   *
+   * \since Version 1.24.
+   */
+  ORT_API2_STATUS(EpAssignedSubgraph_GetEpName, _In_ const OrtEpAssignedSubgraph* ep_subgraph,
+                  _Outptr_ const char** out);
+
+  /** \brief Get the nodes in a subgraph assigned to a specific execution provider.
+   *
+   * \param[in] ep_subgraph The OrtEpAssignedSubgraph instance.
+   * \param[out] ep_nodes Output parameter set to the array of OrtEpAssignedNode instances.
+   * \param[out] num_ep_nodes Output parameter set to the number of OrtEpAssignedNode instance returned.
+   *
+   * \snippet{doc} snippets.dox OrtStatus Return Value
+   *
+   * \since Version 1.24.
+   */
+  ORT_API2_STATUS(EpAssignedSubgraph_GetNodes, _In_ const OrtEpAssignedSubgraph* ep_subgraph,
+                  _Outptr_ const OrtEpAssignedNode* const** ep_nodes, _Out_ size_t* num_ep_nodes);
+
+  /** \brief Get the name of the node assigned to an execution provider.
+   *
+   * \param[in] ep_node The OrtEpAssignedNode instance.
+   * \param[out] out Output parameter set to the node's name as a UTF-8 null-terminated string.
+   *                 Owned by the OrtEpAssignedNode instance (do not free).
+   *
+   * \snippet{doc} snippets.dox OrtStatus Return Value
+   *
+   * \since Version 1.24.
+   */
+  ORT_API2_STATUS(EpAssignedNode_GetName, _In_ const OrtEpAssignedNode* ep_node, _Outptr_ const char** out);
+
+  /** \brief Get the domain of the node assigned to an execution provider.
+   *
+   * \param[in] ep_node The OrtEpAssignedNode instance.
+   * \param[out] out Output parameter set to the node's domain as a UTF-8 null-terminated string.
+   *                 Owned by the OrtEpAssignedNode instance (do not free).
+   *
+   * \snippet{doc} snippets.dox OrtStatus Return Value
+   *
+   * \since Version 1.24.
+   */
+  ORT_API2_STATUS(EpAssignedNode_GetDomain, _In_ const OrtEpAssignedNode* ep_node, _Outptr_ const char** out);
+
+  /** \brief Get the operator type of the node assigned to an execution provider.
+   *
+   * \param[in] ep_node The OrtEpAssignedNode instance.
+   * \param[out] out Output parameter set to the node's operator type as a UTF-8 null-terminated string.
+   *                 Owned by the OrtEpAssignedNode instance (do not free).
+   *
+   * \snippet{doc} snippets.dox OrtStatus Return Value
+   *
+   * \since Version 1.24.
+   */
+  ORT_API2_STATUS(EpAssignedNode_GetOperatorType, _In_ const OrtEpAssignedNode* ep_node, _Outptr_ const char** out);
 };
 
 /*

include/onnxruntime/core/session/onnxruntime_cxx_api.h

@@ -1164,6 +1164,39 @@ OrtCompiledModelCompatibility GetModelCompatibilityForEpDevices(
     const std::vector<ConstEpDevice>& ep_devices,
     const char* compatibility_info);
 
+namespace detail {
+template <typename T>
+struct EpAssignedNodeImpl : Ort::detail::Base<T> {
+  using B = Ort::detail::Base<T>;
+  using B::B;
+
+  std::string GetName() const;
+  std::string GetDomain() const;
+  std::string GetOperatorType() const;
+};
+}  // namespace detail
+
+/** \brief Constant wrapper around ::OrtEpAssignedNode
+ * \remarks EpAssignedNode is always read-only for ORT API users.
+ */
+using ConstEpAssignedNode = detail::EpAssignedNodeImpl<Ort::detail::Unowned<const OrtEpAssignedNode>>;
+
+namespace detail {
+template <typename T>
+struct EpAssignedSubgraphImpl : Ort::detail::Base<T> {
+  using B = Ort::detail::Base<T>;
+  using B::B;
+
+  std::string GetEpName() const;
+  std::vector<ConstEpAssignedNode> GetNodes() const;
+};
+}  // namespace detail
+
+/** \brief Constant wrapper around ::OrtEpAssignedSubgraph
+ * \remarks EpAssignedSubgraph is always read-only for ORT API users.
+ */
+using ConstEpAssignedSubgraph = detail::EpAssignedSubgraphImpl<Ort::detail::Unowned<const OrtEpAssignedSubgraph>>;
+
 /** \brief The Env (Environment)
  *
  * The Env holds the logging state used by all other objects.
@@ -1665,9 +1698,14 @@ struct ConstSessionImpl : Base<T> {
 
   int GetOpset(const std::string& domain) const;  ///< Wraps OrtApi::SessionGetOpsetForDomain
 
-  // Will move before checkin if that's the case.
   std::vector<ValueInfo> GetInputs() const;
   std::vector<ValueInfo> GetOutputs() const;
+
+  /** \brief Returns information on the subgraph/nodes assigned to execution providers in the session.
+   *
+   * \return A list of ConstEpAssignedSubgraph instances.
+   */
+  std::vector<ConstEpAssignedSubgraph> GetEpGraphAssignmentInfo() const;
 };
 
 template <typename T>

include/onnxruntime/core/session/onnxruntime_cxx_inline.h

@@ -747,6 +747,61 @@ inline EpDevice::EpDevice(OrtEpFactory& ep_factory, ConstHardwareDevice& hardwar
   ThrowOnError(GetEpApi().CreateEpDevice(&ep_factory, hardware_device, ep_metadata, ep_options, &p_));
 }
 
+namespace detail {
+template <typename T>
+inline std::string EpAssignedSubgraphImpl<T>::GetEpName() const {
+  const char* ep_name = nullptr;
+
+  // Returned null-terminated string will not be null if API function returns successfully.
+  ThrowOnError(GetApi().EpAssignedSubgraph_GetEpName(this->p_, &ep_name));
+  return std::string(ep_name);
+}
+
+template <typename T>
+inline std::vector<ConstEpAssignedNode> EpAssignedSubgraphImpl<T>::GetNodes() const {
+  size_t num_ep_nodes = 0;
+  const OrtEpAssignedNode* const* ep_node_ptrs = nullptr;
+  ThrowOnError(GetApi().EpAssignedSubgraph_GetNodes(this->p_, &ep_node_ptrs, &num_ep_nodes));
+
+  std::vector<ConstEpAssignedNode> ep_nodes;
+  if (num_ep_nodes > 0) {
+    ep_nodes.reserve(num_ep_nodes);
+    for (size_t i = 0; i < num_ep_nodes; ++i) {
+      ep_nodes.emplace_back(ep_node_ptrs[i]);
+    }
+  }
+
+  return ep_nodes;
+}
+
+template <typename T>
+inline std::string EpAssignedNodeImpl<T>::GetName() const {
+  const char* node_name = nullptr;
+
+  // Returned null-terminated string will not be null if API function returns successfully.
+  ThrowOnError(GetApi().EpAssignedNode_GetName(this->p_, &node_name));
+  return std::string(node_name);
+}
+
+template <typename T>
+inline std::string EpAssignedNodeImpl<T>::GetDomain() const {
+  const char* domain = nullptr;
+
+  // Returned null-terminated string will not be null if API function returns successfully.
+  ThrowOnError(GetApi().EpAssignedNode_GetDomain(this->p_, &domain));
+  return std::string(domain);
+}
+
+template <typename T>
+inline std::string EpAssignedNodeImpl<T>::GetOperatorType() const {
+  const char* op_type = nullptr;
+
+  // Returned null-terminated string will not be null if API function returns successfully.
+  ThrowOnError(GetApi().EpAssignedNode_GetOperatorType(this->p_, &op_type));
+  return std::string(op_type);
+}
+}  // namespace detail
+
 inline Env::Env(OrtLoggingLevel logging_level, _In_ const char* logid) {
   ThrowOnError(GetApi().CreateEnv(logging_level, logid, &p_));
   if (strcmp(logid, "onnxruntime-node") == 0) {
@@ -1756,6 +1811,23 @@ std::vector<ValueInfo> ConstSessionImpl<T>::GetOutputs() const {
   return outputs;
 }
 
+template <typename T>
+inline std::vector<ConstEpAssignedSubgraph> ConstSessionImpl<T>::GetEpGraphAssignmentInfo() const {
+  size_t num_ep_subgraphs = 0;
+  const OrtEpAssignedSubgraph* const* ep_subgraph_ptrs = nullptr;
+  ThrowOnError(GetApi().Session_GetEpGraphAssignmentInfo(this->p_, &ep_subgraph_ptrs, &num_ep_subgraphs));
+
+  std::vector<ConstEpAssignedSubgraph> ep_subgraphs;
+  if (num_ep_subgraphs > 0) {
+    ep_subgraphs.reserve(num_ep_subgraphs);
+    for (size_t i = 0; i < num_ep_subgraphs; ++i) {
+      ep_subgraphs.emplace_back(ep_subgraph_ptrs[i]);
+    }
+  }
+
+  return ep_subgraphs;
+}
+
 template <typename T>
 inline std::vector<Value> SessionImpl<T>::Run(const RunOptions& run_options, const char* const* input_names, const Value* input_values, size_t input_count,
                                               const char* const* output_names, size_t output_count) {

include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h

@@ -421,3 +421,11 @@ static const char* const kOrtSessionOptionsFailOnSuboptimalCompiledModel =
 // "high_power_saver", "low_balanced", "extreme_power_saver", "low_power_saver", "power_saver",
 // "sustained_high_performance". Default to "default".
 static const char* const kOrtEpDynamicOptionsQnnHtpPerformanceMode = "ep.dynamic.qnn_htp_performance_mode";
+
+// Enables the session to record information about the subgraphs/nodes assigned to execution providers.
+// When enabled, an application may call Session_GetEpGraphAssignmentInfo() to retrieve the information.
+//
+// Option values:
+// - "0": Recording of EP graph assignment information is disabled. [DEFAULT]
+// - "1": Recording of EP graph assignment information is enabled.
+static const char* const kOrtSessionOptionsRecordEpGraphAssignmentInfo = "session.record_ep_graph_assignment_info";

onnxruntime/__init__.py

@@ -34,6 +34,8 @@
         OrtArenaCfg,  # noqa: F401
         OrtCompileApiFlags,  # noqa: F401
         OrtDeviceMemoryType,  # noqa: F401
+        OrtEpAssignedNode,  # noqa: F401
+        OrtEpAssignedSubgraph,  # noqa: F401
         OrtEpDevice,  # noqa: F401
         OrtExecutionProviderDevicePolicy,  # noqa: F401
         OrtExternalInitializerInfo,  # noqa: F401

onnxruntime/core/framework/graph_partitioner.cc

@@ -68,6 +68,7 @@ struct PartitionParams {
   std::reference_wrapper<const layout_transformation::TransformLayoutFunction> transform_layout_function;
   std::reference_wrapper<const layout_transformation::DebugGraphFn> debug_graph_fn;
 #endif  // !defined(ORT_MINIMAL_BUILD) || defined(ORT_EXTENDED_MINIMAL_BUILD)
+  std::reference_wrapper<const OnPartitionAssignmentFunction> on_partition_assignment_fn;
 };
 }  // namespace
 
@@ -426,6 +427,7 @@ static Status PartitionOnnxFormatModelImpl(Graph& graph, FuncManager& func_mgr,
                                            const layout_transformation::TransformLayoutFunction& transform_layout_fn,
                                            const layout_transformation::DebugGraphFn& debug_graph_fn,
                                            const CheckLoadCancellationFn& check_load_cancellation_fn,
+                                           const OnPartitionAssignmentFunction& on_partition_assignment_fn,
                                            const logging::Logger& logger, IResourceAccountant* resource_accountant,
                                            const GraphOptimizerRegistry& graph_optimizer_registry,
                                            bool disable_model_compile) {
@@ -444,6 +446,7 @@ static Status PartitionOnnxFormatModelImpl(Graph& graph, FuncManager& func_mgr,
                                                        fused_kernel_registry, current_ep, mode, fused_node_unique_id,
                                                        transform_layout_fn, debug_graph_fn,
                                                        check_load_cancellation_fn,
+                                                       on_partition_assignment_fn,
                                                        logger, resource_accountant,
                                                        graph_optimizer_registry, disable_model_compile));
     }
@@ -518,6 +521,12 @@ static Status PartitionOnnxFormatModelImpl(Graph& graph, FuncManager& func_mgr,
 
     Node* n = nullptr;
     if (sub_graph_available_for_assignment) {
+      if (on_partition_assignment_fn) {
+        // Call custom function provided by owner of GraphPartitioner whenever a subgraph is assigned to an EP.
+        // This can be used, for example, to collect partitioning information.
+        on_partition_assignment_fn(graph, *capability, type);
+      }
+
       n = PlaceNode(graph, *capability->sub_graph, fusion_style, type, mode, fused_node_unique_id);
     }
 
@@ -1018,6 +1027,7 @@ static Status PartitionOnnxFormatModel(const PartitionParams& partition_params,
   auto& fused_node_unique_id = partition_params.fused_node_unique_id.get();
   const auto& transform_layout_function = partition_params.transform_layout_function;
   const CheckLoadCancellationFn& check_load_cancellation_fn = partition_params.check_load_cancellation_fn;
+  const OnPartitionAssignmentFunction& on_partition_assignment_fn = partition_params.on_partition_assignment_fn;
 
   do {
     // process full graph with each EP
@@ -1034,6 +1044,7 @@ static Status PartitionOnnxFormatModel(const PartitionParams& partition_params,
                                                        transform_layout_function,
                                                        partition_params.debug_graph_fn,
                                                        check_load_cancellation_fn,
+                                                       on_partition_assignment_fn,
                                                        logger, resource_accountant, graph_optimizer_registry,
                                                        disable_model_compile));
     }
@@ -1280,7 +1291,8 @@ Status GraphPartitioner::Partition(Graph& graph, FuncManager& func_mgr,
       std::ref(*fused_kernel_registry),
       std::ref(fused_node_unique_id),
       std::cref(transform_layout_function),
-      std::cref(debug_graph_fn)};
+      std::cref(debug_graph_fn),
+      std::cref(on_partition_assignment_fn_)};
 
 #else  // !defined(ORT_MINIMAL_BUILD) || defined(ORT_EXTENDED_MINIMAL_BUILD)
 
@@ -1290,6 +1302,7 @@ Status GraphPartitioner::Partition(Graph& graph, FuncManager& func_mgr,
   PartitionParams partition_params{
       std::ref(graph),
       std::cref(check_load_cancellation_fn),
+      std::cref(on_partition_assignment_fn_),
   };
 
 #endif  // !defined(ORT_MINIMAL_BUILD) || defined(ORT_EXTENDED_MINIMAL_BUILD)

onnxruntime/core/framework/graph_partitioner.h

@@ -20,6 +20,12 @@ namespace epctx {
 struct ModelGenOptions;
 }
 
+// OnPartitionAssignmentFunction is called by GraphPartitioner when a subgraph is assigned to
+// an execution provider. Can be used to collect partitioning information.
+using OnPartitionAssignmentFunction = std::function<void(const Graph& graph,
+                                                         const ComputeCapability& assigned_subgraph,
+                                                         const std::string& assigned_ep_type)>;
+
 class GraphPartitioner {
  public:
   enum class Mode {
@@ -40,11 +46,13 @@ class GraphPartitioner {
   GraphPartitioner(KernelRegistryManager& kernel_registry_mgr,
                    const ExecutionProviders& providers,
                    std::unique_ptr<GraphOptimizerRegistry> graph_optimizer_registry,
-                   CheckLoadCancellationFn check_load_cancellation_fn)
+                   CheckLoadCancellationFn check_load_cancellation_fn,
+                   OnPartitionAssignmentFunction on_partition_assignment_fn = {})
       : kernel_registry_mgr_(kernel_registry_mgr),
         providers_(providers),
         graph_optimizer_registry_(std::move(graph_optimizer_registry)),
-        check_load_cancellation_fn_(std::move(check_load_cancellation_fn)) {
+        check_load_cancellation_fn_(std::move(check_load_cancellation_fn)),
+        on_partition_assignment_fn_(std::move(on_partition_assignment_fn)) {
   }
 
   // Run partitioning.
@@ -89,6 +97,7 @@ class GraphPartitioner {
   const ExecutionProviders& providers_;
   std::unique_ptr<GraphOptimizerRegistry> graph_optimizer_registry_;
   CheckLoadCancellationFn check_load_cancellation_fn_;
+  OnPartitionAssignmentFunction on_partition_assignment_fn_;
 };
 
 }  // namespace onnxruntime