Prototype getting EP graph partitioning info from OrtSession

[adrianlizarraga]

Description

• Adds API functions to get information about the subgraphs/nodes assigned to the EPs in the session.
  • Also adds C++ and Python bindings
• TODO: Make returning information about the nodes optional. User may only want to know which EPs/devices are used.

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.
  • OrtHardwareDevice: The device the EP uses to execute the subgraph. Only available if session created with auto EP feature.
  • OpTypeCounts: The number of operators in the subgraph by type. Ex: {"Conv": 10, "Transpose": 3, ...}
  • [optional] nodes: Name and operator type of each node. Ex: [{"multiply", "Mul"}, ...]

Python example program (taken from unit tests):

    def test_get_graph_provider_partitioning_info_qnn(self):
        """
        Tests querying for information about the nodes assigned to the CPU and QNN EPs.
        """

        if "QNNExecutionProvider" not in onnxrt.get_available_providers():
            self.skipTest("Skipping test because it requires QNN 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_partitioning_info", "1")
        session_options.add_provider("QNNExecutionProvider", {"backend_type": "htp"})

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

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

        # QNN EP offloads QuantizeLinear and DequantizeLinear nodes that are attached to model inputs/outputs to the cpu
        # for performance. So, with a graph with 2 inputs and 2 outputs, we should have 5 subgraphs/partitions:
        #  - Subgraph 1: QuantizeLinear for input 1 is assigned to CPU EP
        #  - Subgraph 2: QuantizeLinear for input 2 is assigned to CPU EP
        #  - Subgraph 3: Bulk of the graph is assigned to QNN EP
        #  - Subgraph 4: DequantizeLinear for output 1 is assigned to CPU EP
        #  - Subgraph 5: DequantizeLinear for output 2 is assigned to CPU EP
        self.assertEqual(len(ep_subgraphs), 5)

        for ep_subgraph in ep_subgraphs:
            self.assertIn(ep_subgraph.ep_name, ("QNNExecutionProvider", "CPUExecutionProvider"))

            if ep_subgraph.ep_name == "CPUExecutionProvider":
                op_type_counts: dict[str, int] = ep_subgraph.get_op_type_counts()
                self.assertEqual(len(op_type_counts), 1)
                if "QuantizeLinear" in op_type_counts:
                    self.assertEqual(op_type_counts["QuantizeLinear"], 1)
                else:
                    self.assertEqual(op_type_counts["DequantizeLinear"], 1)

                ep_nodes = ep_subgraph.get_nodes()
                self.assertEqual(len(ep_nodes), 1)
                self.assertIn(ep_nodes[0].op_type, ("QuantizeLinear", "DequantizeLinear"))
            else:
                # QNNExecutionProvider
                op_type_counts: dict[str, int] = ep_subgraph.get_op_type_counts()
                self.assertEqual(op_type_counts["Transpose"], 1)
                self.assertEqual(op_type_counts["DequantizeLinear"], 7)
                self.assertEqual(op_type_counts["Mul"], 2)
                self.assertEqual(op_type_counts["Conv"], 1)
                self.assertEqual(op_type_counts["QuantizeLinear"], 4)

                ep_nodes = ep_subgraph.get_nodes()
                self.assertEqual(len(ep_nodes), 15)  # Won't match the original model due to layout transformation
                conv_node = next((n for n in ep_nodes if n.op_type == "Conv"), None)
                self.assertIsNotNone(conv_node)
                self.assertTrue(conv_node.name.startswith("conv_node"))  # Layout transformation adds name suffixes

C++ program (taken from unit test):

// Tests querying session for information about which nodes where assigned to the EPs in the session.
TEST_F(QnnHTPBackendTests, Session_GetEpGraphPartitioningInfo) {
  Ort::SessionOptions session_options;
  session_options.AddConfigEntry(kOrtSessionOptionsRecordEpGraphPartitioningInfo, "1");
  session_options.AppendExecutionProvider(kQnnExecutionProvider, ProviderOptions{{"backend_type", "htp"}});

  const ORTCHAR_T* ort_model_path = ORT_MODEL_FOLDER "layout_transform_const_folding.qdq.onnx";
  Ort::Session session(*ort_env, ort_model_path, session_options);
  ASSERT_TRUE(SessionHasEp(session, kQnnExecutionProvider));

  std::vector<Ort::ConstEpAssignedSubgraph> ep_subgraphs = session.GetEpGraphPartitioningInfo();

  // QNN EP offloads QuantizeLinear and DequantizeLinear nodes that are attached to model inputs/outputs to the cpu
  // for performance. So, with a graph with 2 inputs and 2 outputs, we should have 5 subgraphs/partitions:
  //  - Subgraph 1: QuantizeLinear for input 1 is assigned to CPU EP
  //  - Subgraph 2: QuantizeLinear for input 2 is assigned to CPU EP
  //  - Subgraph 3: Bulk of the graph is assigned to QNN EP
  //  - Subgraph 4: DequantizeLinear for output 1 is assigned to CPU EP
  //  - Subgraph 5: DequantizeLinear for output 2 is assigned to CPU EP
  ASSERT_EQ(ep_subgraphs.size(), 5);

  for (auto subgraph : ep_subgraphs) {
    std::string ep_name = subgraph.EpName();
    ASSERT_TRUE(ep_name == kQnnExecutionProvider || ep_name == kCpuExecutionProvider);

    std::unordered_map<std::string, size_t> op_type_counts = subgraph.GetOpTypeCounts();
    const std::vector<Ort::ConstEpAssignedNode> ep_nodes = subgraph.GetNodes();

    if (ep_name == kCpuExecutionProvider) {
      ASSERT_EQ(op_type_counts.size(), 1);
      if (auto it = op_type_counts.find("QuantizeLinear"); it != op_type_counts.end()) {
        ASSERT_EQ(it->second, 1);
      } else {
        it = op_type_counts.find("DequantizeLinear");
        ASSERT_TRUE(it != op_type_counts.end());
        ASSERT_EQ(it->second, 1);
      }

      ASSERT_EQ(ep_nodes.size(), 1);

      std::string op_type = ep_nodes[0].OpType();
      ASSERT_TRUE(op_type == "QuantizeLinear" || op_type == "DequantizeLinear");
    } else /*if (ep_name == kQnnExecutionProvider)*/ {
      ASSERT_EQ(op_type_counts.size(), 5);
      ASSERT_EQ(op_type_counts["Transpose"], 1);
      ASSERT_EQ(op_type_counts["DequantizeLinear"], 7);
      ASSERT_EQ(op_type_counts["Mul"], 2);
      ASSERT_EQ(op_type_counts["Conv"], 1);
      ASSERT_EQ(op_type_counts["QuantizeLinear"], 4);

      ASSERT_EQ(ep_nodes.size(), 15);
      auto it = std::find_if(ep_nodes.begin(), ep_nodes.end(),
                             [](const Ort::ConstEpAssignedNode& n) -> bool { return std::strncmp(n.OpType(), "Conv", 4) == 0; });
      ASSERT_TRUE(it != ep_nodes.end());
      std::string conv_name = it->Name();
      ASSERT_EQ(conv_name.rfind("conv_node", 0), 0);  // Check node name starts with "conv_node".
                                                      // Note that layout transformation may add a suffix to node names
    }
  }
}

Motivation and Context

[Craigacp]

How stable is this C API? I'd like to build a Java frontend on top of it to help me run down #23154.

[adrianlizarraga]

How stable is this C API? I'd like to build a Java frontend on top of it to help me run down #23154.

Hi @Craigacp, These C APIs are not stable. This draft PR has not been reviewed so APIs are likely to change.

[yuslepukhin]

const OnPartitionAssignmentFunction& on_partition_assign_fn =

nit: it seems like the function signature is growing. Would you consider making it a member of the class?
Functions can be be null or non-null.

Another option is to make it a part of partition params

[yuslepukhin]

ComputeCapability__SetHardwareDevic

Naming is off, it is usually not exposed like this.

[yuslepukhin]

py::class_ py_ep_node(m

It may warrant a separate compilation unit. There is already a sizable chunk of code dedicated to the topic. hard to maintain.

[adrianlizarraga]

Closing this in favor of #26781