Add API to get ep graph partitioning info

[adrianlizarraga]

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):

    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):

  // 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

[Copilot]

Pull request overview

This PR adds a new API to retrieve information about how the graph is partitioned across execution providers (EPs). The feature allows applications to query which nodes are assigned to which EPs after session creation.

Key changes:

• Adds C API functions (Session_GetEpGraphPartitioningInfo and related accessors) to query EP graph partitioning information
• Implements C++ and Python bindings for the new API
• Introduces OrtEpAssignedSubgraph and OrtEpAssignedNode structures to represent partitioning data

Reviewed changes

Copilot reviewed 16 out of 16 changed files in this pull request and generated 5 comments.

Show a summary per file

File	Description
include/onnxruntime/core/session/onnxruntime_c_api.h	Declares new C API functions for retrieving EP partitioning info and accessing subgraph/node details
include/onnxruntime/core/session/onnxruntime_cxx_api.h	Adds C++ wrapper classes ConstEpAssignedSubgraph and ConstEpAssignedNode with GetEpGraphPartitioningInfo() method
include/onnxruntime/core/session/onnxruntime_cxx_inline.h	Implements inline methods for the C++ wrapper classes
include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h	Adds config key to enable recording of EP graph partitioning information
onnxruntime/core/session/ort_apis.h	Declares implementation functions for the new C API
onnxruntime/core/session/onnxruntime_c_api.cc	Implements the C API functions with proper null checks and ORT_MINIMAL_BUILD guards
onnxruntime/core/session/ep_graph_partition_info.h	Defines data structures (OrtEpAssignedNode, OrtEpAssignedSubgraph) for storing partitioning information
onnxruntime/core/session/inference_session.h	Adds member variables to store graph partitioning info and getter method
onnxruntime/core/session/inference_session.cc	Implements RecordPartitionInfo() function and integrates it into the graph transformation pipeline
onnxruntime/core/framework/graph_partitioner.h	Adds OnPartitionAssignmentFunction callback type and constructor parameter
onnxruntime/core/framework/graph_partitioner.cc	Invokes the callback when a subgraph is assigned to an EP
onnxruntime/python/onnxruntime_pybind_state.cc	Exposes Python bindings for OrtEpAssignedNode and OrtEpAssignedSubgraph classes
onnxruntime/python/onnxruntime_inference_collection.py	Adds get_provider_graph_partitioning_info() method to InferenceSession Python API
onnxruntime/__init__.py	Exports the new Python classes
onnxruntime/test/python/onnxruntime_test_python.py	Adds Python unit test validating the new API with CPU EP
onnxruntime/test/autoep/test_execution.cc	Adds C++ unit test validating the API with both plugin EP and CPU EP

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 16 out of 16 changed files in this pull request and generated 35 comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.