Address review comments: Add API for domain, clean up test, clean up docs, etc

adrianlizarraga · adrianlizarraga · commit 57aafc5d0844 · 2026-01-15T02:11:42.000-08:00
diff --git a/include/onnxruntime/core/session/onnxruntime_c_api.h b/include/onnxruntime/core/session/onnxruntime_c_api.h
@@ -6915,20 +6915,30 @@ struct OrtApi {
 
   /// @}
 
-  /** \brief Get information about the subgraphs assigned to each EP and the nodes within.
+  /** \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 and operator type for every node. For compiling EPs,
-   * a subgraph contains one or more nodes. Alternatively, for EPs that use kernel registration (e.g., CPU EP), each
-   * registered kernel for a node is contained in its own subgraph (i.e., a subgraph contains one node).
+   * 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.
    *
-   * \note Application must enable the recording of graph partitioning information by enabling the session configuration
-   *       for the key "session.record_ep_graph_assignment_info". Refer to onnxruntime_session_options_config_keys.h.
-   *       If the session configuration is not enabled, this function returns an empty result.
+   * 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.
    *
-   * \param[in] session The OrtSession instance to query.
-   * \param[out] ep_subgraphs The OrtEpAssignedSubgraph instances denoting the EP graph partitioning.
-   * \param[out] num_ep_subgraphs The number of OrtEpAssignedSubgraph instances returned.
+   * 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
    *
@@ -6940,18 +6950,18 @@ struct OrtApi {
 
   /** \brief Get the name of the execution provider to which the subgraph was assigned.
    *
-   * \param[in] ep_subgraph The OrtEpAssignedSubgraph instance to query.
-   * \return The execution provider name.
+   * \param[in] ep_subgraph The OrtEpAssignedSubgraph instance.
+   * \return The execution provider's name as UTF-8 null-terminated string.
    *
    * \since Version 1.24.
    */
-  const char*(ORT_API_CALL* EpAssignedSubgraph_GetEpName)(_In_ const OrtEpAssignedSubgraph* ep_subgraph);
+  ORT_API_T(const char*, EpAssignedSubgraph_GetEpName, _In_ const OrtEpAssignedSubgraph* ep_subgraph);
 
-  /** \brief Get the list of nodes assigned to an execution provider.
+  /** \brief Get the nodes in a subgraph assigned to a specific execution provider.
    *
-   * \param[in] ep_subgraph The OrtEpAssignedSubgraph instance to query.
-   * \param[out] ep_nodes Output parameter set to the list of OrtEpAssignedNode instances.
-   * \param[out] num_ep_nodes Output parameter set to the number of OrtEpAssignedNode instances returned.
+   * \param[in] ep_subgraph The OrtEpAssignedSubgraph instance.
+   * \param[out] ep_nodes Output parameter set to the array of OrtEpAssignedNode instances.
+   * \param[in] num_ep_nodes Output parameter set to the number of OrtEpAssignedNode instance returned.
    *
    * \snippet{doc} snippets.dox OrtStatus Return Value
    *
@@ -6962,21 +6972,30 @@ struct OrtApi {
 
   /** \brief Get the name of the node assigned to an execution provider.
    *
-   * \param[in] ep_node The OrtEpAssignedNode instance to query.
-   * \return The node's name.
+   * \param[in] ep_node The OrtEpAssignedNode instance.
+   * \return The node's name as a UTF-8 null-terminated string.
+   *
+   * \since Version 1.24.
+   */
+  ORT_API_T(const char*, EpAssignedNode_GetName, _In_ const OrtEpAssignedNode* ep_node);
+
+  /** \brief Get the domain of the node assigned to an execution provider.
+   *
+   * \param[in] ep_node The OrtEpAssignedNode instance.
+   * \return The node's domain as a UTF-8 null-terminated string.
    *
    * \since Version 1.24.
    */
-  const char*(ORT_API_CALL* EpAssignedNode_GetName)(_In_ const OrtEpAssignedNode* ep_node);
+  ORT_API_T(const char*, EpAssignedNode_GetDomain, _In_ const OrtEpAssignedNode* ep_node);
 
   /** \brief Get the operator type of the node assigned to an execution provider.
    *
-   * \param[in] ep_node The OrtEpAssignedNode instance to query.
-   * \return The node's operator type.
+   * \param[in] ep_node The OrtEpAssignedNode instance.
+   * \return The node's operator type as a UTF-8 null-terminated string.
    *
    * \since Version 1.24.
    */
-  const char*(ORT_API_CALL* EpAssignedNode_GetOperatorType)(_In_ const OrtEpAssignedNode* ep_node);
+  ORT_API_T(const char*, EpAssignedNode_GetOperatorType, _In_ const OrtEpAssignedNode* ep_node);
 };
 
 /*
diff --git a/include/onnxruntime/core/session/onnxruntime_cxx_api.h b/include/onnxruntime/core/session/onnxruntime_cxx_api.h
@@ -1170,8 +1170,9 @@ struct EpAssignedNodeImpl : Ort::detail::Base<T> {
   using B = Ort::detail::Base<T>;
   using B::B;
 
-  const char* GetName() const;
-  const char* GetOperatorType() const;
+  std::string GetName() const;
+  std::string GetDomain() const;
+  std::string GetOperatorType() const;
 };
 }  // namespace detail
 
@@ -1186,7 +1187,7 @@ struct EpAssignedSubgraphImpl : Ort::detail::Base<T> {
   using B = Ort::detail::Base<T>;
   using B::B;
 
-  const char* GetEpName() const;
+  std::string GetEpName() const;
   std::vector<ConstEpAssignedNode> GetNodes() const;
 };
 }  // namespace detail
diff --git a/include/onnxruntime/core/session/onnxruntime_cxx_inline.h b/include/onnxruntime/core/session/onnxruntime_cxx_inline.h
@@ -749,8 +749,8 @@ inline EpDevice::EpDevice(OrtEpFactory& ep_factory, ConstHardwareDevice& hardwar
 
 namespace detail {
 template <typename T>
-inline const char* EpAssignedSubgraphImpl<T>::GetEpName() const {
-  return GetApi().EpAssignedSubgraph_GetEpName(this->p_);
+inline std::string EpAssignedSubgraphImpl<T>::GetEpName() const {
+  return std::string(GetApi().EpAssignedSubgraph_GetEpName(this->p_));
 }
 
 template <typename T>
@@ -771,13 +771,18 @@ inline std::vector<ConstEpAssignedNode> EpAssignedSubgraphImpl<T>::GetNodes() co
 }
 
 template <typename T>
-inline const char* EpAssignedNodeImpl<T>::GetName() const {
-  return GetApi().EpAssignedNode_GetName(this->p_);
+inline std::string EpAssignedNodeImpl<T>::GetName() const {
+  return std::string(GetApi().EpAssignedNode_GetName(this->p_));
 }
 
 template <typename T>
-inline const char* EpAssignedNodeImpl<T>::GetOperatorType() const {
-  return GetApi().EpAssignedNode_GetOperatorType(this->p_);
+inline std::string EpAssignedNodeImpl<T>::GetDomain() const {
+  return std::string(GetApi().EpAssignedNode_GetDomain(this->p_));
+}
+
+template <typename T>
+inline std::string EpAssignedNodeImpl<T>::GetOperatorType() const {
+  return std::string(GetApi().EpAssignedNode_GetOperatorType(this->p_));
 }
 }  // namespace detail
 
diff --git a/onnxruntime/core/session/ep_graph_assignment_info.h b/onnxruntime/core/session/ep_graph_assignment_info.h
@@ -13,6 +13,7 @@
 /// </summary>
 struct OrtEpAssignedNode {
   std::string name;
+  std::string domain;
   std::string op_type;
 };
 
diff --git a/onnxruntime/core/session/inference_session.cc b/onnxruntime/core/session/inference_session.cc
@@ -1288,10 +1288,9 @@ common::Status InferenceSession::ApplyUpdates(const OrtModel& model_editor_api_m
 }
 
 #if !defined(ORT_MINIMAL_BUILD)
-static void RecordEpGraphPartitionAssignment(std::vector<std::unique_ptr<OrtEpAssignedSubgraph>>& ep_assigned_subgraphs,
-                                             const Graph& graph,
-                                             const ComputeCapability& capability,
-                                             const std::string& ep_name) {
+static std::unique_ptr<OrtEpAssignedSubgraph> CreateEpAssignedSubgraph(const Graph& graph,
+                                                                       const ComputeCapability& capability,
+                                                                       const std::string& ep_name) {
   auto assigned_subgraph = std::make_unique<OrtEpAssignedSubgraph>();
   assigned_subgraph->ep_name = ep_name;
 
@@ -1300,18 +1299,17 @@ static void RecordEpGraphPartitionAssignment(std::vector<std::unique_ptr<OrtEpAs
   for (NodeIndex node_index : node_indices) {
     const Node* node = graph.GetNode(node_index);
     if (node != nullptr) {
-      const std::string& op_type = node->OpType();
-
       auto assigned_node = std::make_unique<OrtEpAssignedNode>();
       assigned_node->name = node->Name();
-      assigned_node->op_type = op_type;
+      assigned_node->domain = node->Domain();
+      assigned_node->op_type = node->OpType();
 
       assigned_subgraph->nodes.push_back(assigned_node.get());
       assigned_subgraph->nodes_storage.push_back(std::move(assigned_node));
     }
   }
 
-  ep_assigned_subgraphs.push_back(std::move(assigned_subgraph));
+  return assigned_subgraph;
 }
 #endif  // !defined(ORT_MINIMAL_BUILD)
 
@@ -1329,15 +1327,19 @@ common::Status InferenceSession::TransformGraph(onnxruntime::Graph& graph, bool
   // 8. Repeat steps 5 to 7 depending on the graph optimizations loop level.
   // 9. insert copy nodes (required transformer).
 
-  OnPartitionAssignmentFunction on_partition_assign_fn;
+  OnPartitionAssignmentFunction on_partition_assignment_fn;
 #if !defined(ORT_MINIMAL_BUILD)
-  bool record_ep_graph_partitioning =
+  bool record_ep_graph_assignment =
       session_options_.config_options.GetConfigOrDefault(kOrtSessionOptionsRecordEpGraphAssignmentInfo, "0") == "1";
-  if (record_ep_graph_partitioning) {
-    on_partition_assign_fn = [this](const Graph& graph, const ComputeCapability& assigned_subgraph,
-                                    const std::string& assigned_ep_type) {
-      RecordEpGraphPartitionAssignment(this->ep_graph_assignment_info_storage_, graph, assigned_subgraph,
-                                       assigned_ep_type);
+  if (record_ep_graph_assignment) {
+    on_partition_assignment_fn = [this](const Graph& graph, const ComputeCapability& compute_capability,
+                                        const std::string& ep_name) {
+      std::unique_ptr<OrtEpAssignedSubgraph> assigned_subgraph = CreateEpAssignedSubgraph(graph,
+                                                                                          compute_capability,
+                                                                                          ep_name);
+
+      this->ep_graph_assignment_info_.push_back(assigned_subgraph.get());
+      this->ep_graph_assignment_info_storage_.push_back(std::move(assigned_subgraph));
     };
   }
 #endif  // !defined(ORT_MINIMAL_BUILD)
@@ -1347,7 +1349,7 @@ common::Status InferenceSession::TransformGraph(onnxruntime::Graph& graph, bool
                                                                            execution_providers_.Get(onnxruntime::kCpuExecutionProvider),
                                                                            session_logger_);
   GraphPartitioner partitioner(kernel_registry_manager_, execution_providers_, std::move(graph_optimizer_registry),
-                               check_load_cancellation_fn_, on_partition_assign_fn);
+                               check_load_cancellation_fn_, on_partition_assignment_fn);
 
   // Run Ahead Of time function inlining
   if (const bool disable_aot_function_inlining =
@@ -1489,14 +1491,6 @@ common::Status InferenceSession::TransformGraph(onnxruntime::Graph& graph, bool
                                                        session_options_.config_options, *session_logger_,
                                                        mode, session_options_.GetEpContextGenerationOptions(), debug_graph_fn));
 
-#if !defined(ORT_MINIMAL_BUILD)
-  if (record_ep_graph_partitioning) {
-    for (std::unique_ptr<OrtEpAssignedSubgraph>& ep_subgraph : ep_graph_assignment_info_storage_) {
-      ep_graph_assignment_info_.push_back(ep_subgraph.get());
-    }
-  }
-#endif  // !defined(ORT_MINIMAL_BUILD)
-
   // Get graph optimizations loop level from session config, if not present, set to default value of 1 as per
   // the definition of kOrtSessionOptionsGraphOptimizationsLoopLevel.
   unsigned int graph_optimizations_loop_level = static_cast<unsigned int>(std::stoi(
diff --git a/onnxruntime/core/session/onnxruntime_c_api.cc b/onnxruntime/core/session/onnxruntime_c_api.cc
@@ -48,6 +48,7 @@
 #include "core/session/lora_adapters.h"
 #include "core/session/model_editor_api.h"
 #include "core/session/onnxruntime_c_api.h"
+#include "core/session/onnxruntime_session_options_config_keys.h"
 #include "core/session/ort_apis.h"
 #include "core/session/ort_env.h"
 #include "core/session/utils.h"
@@ -878,6 +879,18 @@ ORT_API_STATUS_IMPL(OrtApis::Session_GetEpGraphAssignmentInfo, _In_ const OrtSes
                     _Out_ size_t* num_ep_subgraphs) {
   API_IMPL_BEGIN
 #if !defined(ORT_MINIMAL_BUILD)
+  const auto* inference_session = reinterpret_cast<const onnxruntime::InferenceSession*>(session);
+  const auto& session_options = inference_session->GetSessionOptions();
+  bool is_enabled =
+      session_options.config_options.GetConfigOrDefault(kOrtSessionOptionsRecordEpGraphAssignmentInfo, "0") == "1";
+
+  if (!is_enabled) {
+    std::ostringstream oss;
+    oss << "Session configuration entry '" << kOrtSessionOptionsRecordEpGraphAssignmentInfo
+        << "' must be set to \"1\" to retrieve EP graph assignment information.";
+    return OrtApis::CreateStatus(ORT_FAIL, oss.str().c_str());
+  }
+
   if (ep_subgraphs == nullptr) {
     return OrtApis::CreateStatus(ORT_INVALID_ARGUMENT, "'ep_subgraphs' argument is null");
   }
@@ -886,7 +899,6 @@ ORT_API_STATUS_IMPL(OrtApis::Session_GetEpGraphAssignmentInfo, _In_ const OrtSes
     return OrtApis::CreateStatus(ORT_INVALID_ARGUMENT, "'num_ep_subgraphs' argument is null");
   }
 
-  auto inference_session = reinterpret_cast<const onnxruntime::InferenceSession*>(session);
   const std::vector<const OrtEpAssignedSubgraph*>& ep_assignment_info = inference_session->GetEpGraphAssignmentInfo();
 
   *ep_subgraphs = ep_assignment_info.data();
@@ -896,7 +908,7 @@ ORT_API_STATUS_IMPL(OrtApis::Session_GetEpGraphAssignmentInfo, _In_ const OrtSes
   ORT_UNUSED_PARAMETER(session);
   ORT_UNUSED_PARAMETER(ep_subgraphs);
   ORT_UNUSED_PARAMETER(num_ep_subgraphs);
-  return OrtApis::CreateStatus(ORT_NOT_IMPLEMENTED, "EP graph partitioning information is not supported in this build");
+  return OrtApis::CreateStatus(ORT_NOT_IMPLEMENTED, "EP graph assignment information is not supported in this build");
 #endif  // !defined(ORT_MINIMAL_BUILD)
   API_IMPL_END
 }
@@ -906,7 +918,7 @@ ORT_API(const char*, OrtApis::EpAssignedSubgraph_GetEpName, _In_ const OrtEpAssi
   return ep_subgraph->ep_name.c_str();
 #else
   ORT_UNUSED_PARAMETER(ep_subgraph);
-  fprintf(stderr, "EP graph partitioning information is not supported in this build\n");
+  fprintf(stderr, "EP graph assignment information is not supported in this build\n");
   return nullptr;
 #endif  // !defined(ORT_MINIMAL_BUILD)
 }
@@ -916,11 +928,15 @@ ORT_API_STATUS_IMPL(OrtApis::EpAssignedSubgraph_GetNodes, _In_ const OrtEpAssign
   API_IMPL_BEGIN
 #if !defined(ORT_MINIMAL_BUILD)
   if (ep_nodes == nullptr) {
-    return OrtApis::CreateStatus(ORT_INVALID_ARGUMENT, "'ep_nodes' argument is null");
+    return OrtApis::CreateStatus(ORT_INVALID_ARGUMENT,
+                                 "EpAssignedSubgraph_GetNodes requires a valid (non-null) `ep_nodes` output parameter "
+                                 "into which to store the pointer to the node array.");
   }
 
   if (num_ep_nodes == nullptr) {
-    return OrtApis::CreateStatus(ORT_INVALID_ARGUMENT, "'num_ep_nodes' argument is null");
+    return OrtApis::CreateStatus(ORT_INVALID_ARGUMENT,
+                                 "EpAssignedSubgraph_GetNodes requires a valid (non-null) `num_ep_nodes` "
+                                 "output parameter into which to store the number of nodes.");
   }
 
   *ep_nodes = ep_subgraph->nodes.data();
@@ -930,7 +946,7 @@ ORT_API_STATUS_IMPL(OrtApis::EpAssignedSubgraph_GetNodes, _In_ const OrtEpAssign
   ORT_UNUSED_PARAMETER(ep_subgraph);
   ORT_UNUSED_PARAMETER(ep_nodes);
   ORT_UNUSED_PARAMETER(num_ep_nodes);
-  return OrtApis::CreateStatus(ORT_NOT_IMPLEMENTED, "EP graph partitioning information is not supported in this build");
+  return OrtApis::CreateStatus(ORT_NOT_IMPLEMENTED, "EP graph assignment information is not supported in this build");
 #endif  // !defined(ORT_MINIMAL_BUILD)
   API_IMPL_END
 }
@@ -940,7 +956,17 @@ ORT_API(const char*, OrtApis::EpAssignedNode_GetName, _In_ const OrtEpAssignedNo
   return ep_node->name.c_str();
 #else
   ORT_UNUSED_PARAMETER(ep_node);
-  fprintf(stderr, "EP graph partitioning information is not supported in this build\n");
+  fprintf(stderr, "EP graph assignment information is not supported in this build\n");
+  return nullptr;
+#endif  // !defined(ORT_MINIMAL_BUILD)
+}
+
+ORT_API(const char*, OrtApis::EpAssignedNode_GetDomain, _In_ const OrtEpAssignedNode* ep_node) {
+#if !defined(ORT_MINIMAL_BUILD)
+  return ep_node->domain.c_str();
+#else
+  ORT_UNUSED_PARAMETER(ep_node);
+  fprintf(stderr, "EP graph assignment information is not supported in this build\n");
   return nullptr;
 #endif  // !defined(ORT_MINIMAL_BUILD)
 }
@@ -950,7 +976,7 @@ ORT_API(const char*, OrtApis::EpAssignedNode_GetOperatorType, _In_ const OrtEpAs
   return ep_node->op_type.c_str();
 #else
   ORT_UNUSED_PARAMETER(ep_node);
-  fprintf(stderr, "EP graph partitioning information is not supported in this build\n");
+  fprintf(stderr, "EP graph assignment information is not supported in this build\n");
   return nullptr;
 #endif  // !defined(ORT_MINIMAL_BUILD)
 }
@@ -4386,6 +4412,7 @@ static constexpr OrtApi ort_api_1_to_24 = {
     &OrtApis::EpAssignedSubgraph_GetEpName,
     &OrtApis::EpAssignedSubgraph_GetNodes,
     &OrtApis::EpAssignedNode_GetName,
+    &OrtApis::EpAssignedNode_GetDomain,
     &OrtApis::EpAssignedNode_GetOperatorType,
 };
 
diff --git a/onnxruntime/core/session/ort_apis.h b/onnxruntime/core/session/ort_apis.h
@@ -790,5 +790,6 @@ ORT_API(const char*, EpAssignedSubgraph_GetEpName, _In_ const OrtEpAssignedSubgr
 ORT_API_STATUS_IMPL(EpAssignedSubgraph_GetNodes, _In_ const OrtEpAssignedSubgraph* ep_subgraph,
                     _Outptr_ const OrtEpAssignedNode* const** ep_nodes, _Out_ size_t* num_ep_nodes);
 ORT_API(const char*, EpAssignedNode_GetName, _In_ const OrtEpAssignedNode* ep_node);
+ORT_API(const char*, EpAssignedNode_GetDomain, _In_ const OrtEpAssignedNode* ep_node);
 ORT_API(const char*, EpAssignedNode_GetOperatorType, _In_ const OrtEpAssignedNode* ep_node);
 }  // namespace OrtApis
diff --git a/onnxruntime/test/autoep/test_execution.cc b/onnxruntime/test/autoep/test_execution.cc
