diff --git a/onnxruntime/core/session/plugin_ep/ep_plugin_provider_interfaces.cc b/onnxruntime/core/session/plugin_ep/ep_plugin_provider_interfaces.cc index c8829423fbe26..55245420db37a 100644 --- a/onnxruntime/core/session/plugin_ep/ep_plugin_provider_interfaces.cc +++ b/onnxruntime/core/session/plugin_ep/ep_plugin_provider_interfaces.cc @@ -3,6 +3,7 @@ #include "core/session/plugin_ep/ep_plugin_provider_interfaces.h" +#include #include #include #include @@ -117,6 +118,17 @@ static OrtDevice GetOrtDeviceForPluginEp(gsl::span ep_ return device_memory_info != nullptr ? device_memory_info->device : OrtDevice(); } +static const Node* FindFirstNodeAssignedToOtherEP(const std::string& ep_type, + gsl::span ep_nodes) { + auto node_iter = std::find_if(ep_nodes.begin(), ep_nodes.end(), + [&ep_type](const EpNode* node) -> bool { + const auto& node_ep_type = node->GetInternalNode().GetExecutionProviderType(); + return !node_ep_type.empty() && node_ep_type != ep_type; + }); + + return node_iter != ep_nodes.end() ? &(*node_iter)->GetInternalNode() : nullptr; +} + PluginExecutionProvider::PluginExecutionProvider(UniqueOrtEp ep, const OrtSessionOptions& session_options, OrtEpFactory& ep_factory, gsl::span ep_devices, @@ -158,9 +170,11 @@ PluginExecutionProvider::GetCapability(const onnxruntime::GraphViewer& graph_vie ORT_UNUSED_PARAMETER(resource_accountant); // TODO: Add support? Not used by prioritized EPs ORT_UNUSED_PARAMETER(kernel_lookup); // TODO: Add support? Not used by prioritized EPs, so probably not needed? + const logging::Logger& logger = GetLogger() != nullptr ? *GetLogger() : logging::LoggingManager::DefaultLogger(); + std::unique_ptr ep_graph = nullptr; if (Status status = EpGraph::Create(graph_viewer, ep_graph); !status.IsOK()) { - LOGS_DEFAULT(ERROR) << "Failed to create OrtGraph: " << status.ToString(); + LOGS(logger, ERROR) << "Failed to create OrtGraph for " << Type() << ": " << status.ToString(); return {}; } @@ -168,7 +182,7 @@ PluginExecutionProvider::GetCapability(const onnxruntime::GraphViewer& graph_vie Status status = ToStatusAndRelease(ort_ep_->GetCapability(ort_ep_.get(), ep_graph->ToExternal(), &api_graph_support_info)); if (!status.IsOK()) { - LOGS_DEFAULT(ERROR) << "OrtEp::GetCapability() failed with error: " << status.ToString(); + LOGS(logger, ERROR) << "OrtEp::GetCapability() for " << Type() << " failed with error: " << status.ToString(); return {}; } @@ -182,12 +196,39 @@ PluginExecutionProvider::GetCapability(const onnxruntime::GraphViewer& graph_vie // Create ComputeCapability instances from OrtEpGraphSupportInfo::NodeGrouping instances. for (const OrtEpGraphSupportInfo::NodeGrouping& node_grouping : api_graph_support_info.node_groupings) { + // Skip this node grouping if any node has already been assigned to another EP. + if (const Node* node_for_other_ep = FindFirstNodeAssignedToOtherEP(Type(), node_grouping.nodes); + node_for_other_ep != nullptr) { + LOGS(logger, WARNING) << "OrtEp::GetCapability() specified nodes that cannot be assigned to " << Type() << ". " + << "Found one or more nodes that were already assigned to a different EP named '" + << node_for_other_ep->GetExecutionProviderType() << "'. Ex: " + << node_for_other_ep->OpType() << " node with name '" + << node_for_other_ep->Name() << "'."; + continue; + } + if (node_grouping.kind == OrtEpGraphSupportInfo::NodeGroupingKind::kSingleAssignedNode) { + if (node_grouping.nodes.size() != 1) { + // The EpGraphSupportInfo_AddSingleNode() C API should already return an error if the EP tries to provide + // an invalid node. However, we check here too just in case this changes. + LOGS(logger, ERROR) << "OrtEp::GetCapability() for " << Type() << " did not specify exactly one valid node " + << "when calling EpGraphSupportInfo_AddSingleNode()."; + return {}; + } + auto indexed_sub_graph = std::make_unique(); indexed_sub_graph->nodes.push_back(node_grouping.nodes[0]->GetInternalNode().Index()); result.push_back(std::make_unique(std::move(indexed_sub_graph))); } else if (node_grouping.kind == OrtEpGraphSupportInfo::NodeGroupingKind::kFusedNode) { + if (node_grouping.nodes.empty()) { + // The EpGraphSupportInfo_AddNodesToFuse() C API should already return an error if the EP tries to provide + // an empty array of nodes from OrtEp::GetCapability(). However, we check here too just in case this changes. + LOGS(logger, ERROR) << "OrtEp::GetCapability() for " << Type() << " set an empty array of nodes " + << "when specifying supported nodes."; + return {}; + } + std::unordered_set node_set; node_set.reserve(node_grouping.nodes.size()); @@ -207,27 +248,29 @@ PluginExecutionProvider::GetCapability(const onnxruntime::GraphViewer& graph_vie this->Type(), this->Type(), /*node_unit_map*/ nullptr, node_grouping.fusion_options.drop_constant_initializers); - if (capabilities.size() > 1) { - LOGS_DEFAULT(ERROR) << "OrtEp::GetCapability() set nodes that cannot be fused together. " - << "Please ensure that the nodes provided to EpGraphSupportInfo_AddFusedNodes() do not " + if (capabilities.size() != 1) { + LOGS(logger, ERROR) << "OrtEp::GetCapability() for " << Type() << " set nodes that cannot be fused together. " + << "Please ensure that the nodes provided to EpGraphSupportInfo_AddNodesToFuse() do not " << "have an unsupported node in any path between two of the supported nodes."; return {}; } - // Enforce that the nodes in node_set match the nodes in capabilities[0] + // Log an error if the nodes in node_set do not match the nodes in capabilities[0]. We expect this to always + // be true because we've already checked that the EP did not try to claim nodes already assigned to another EP. // TODO(adrianlizarraga): This check can be removed when we stop using utils::CreateSupportedPartitions() above. std::vector& capability_node_indices = capabilities[0]->sub_graph->nodes; std::unordered_set capability_node_indices_set(capability_node_indices.begin(), capability_node_indices.end()); - ORT_ENFORCE(node_set.size() == capability_node_indices_set.size()); - ORT_ENFORCE(std::all_of(node_set.begin(), node_set.end(), [&capability_node_indices_set](const Node* node) { - return capability_node_indices_set.count(node->Index()) != 0; - })); + if (node_set.size() != capability_node_indices_set.size()) { + LOGS(logger, ERROR) << "OrtEp::GetCapability() for " << Type() + << " set nodes that cannot all be fused together."; + return {}; + } result.push_back(std::move(capabilities[0])); } else { - LOGS_DEFAULT(ERROR) << "PluginExecutionProvider::GetCapability() has invalid NodeGroupingKind: " + LOGS(logger, ERROR) << "PluginExecutionProvider::GetCapability() has invalid NodeGroupingKind: " << static_cast(node_grouping.kind); return {}; } diff --git a/onnxruntime/test/framework/ep_plugin_provider_test.cc b/onnxruntime/test/framework/ep_plugin_provider_test.cc index 35f7d06fb0912..30595d5ce97b2 100644 --- a/onnxruntime/test/framework/ep_plugin_provider_test.cc +++ b/onnxruntime/test/framework/ep_plugin_provider_test.cc @@ -3,9 +3,14 @@ #include "core/session/plugin_ep/ep_plugin_provider_interfaces.h" +#include #include "gsl/gsl" #include "gtest/gtest.h" +#include "core/common/logging/sinks/file_sink.h" +#include "core/graph/graph_viewer.h" +#include "core/graph/model.h" +#include "core/optimizer/graph_optimizer_registry.h" #include "core/session/abi_devices.h" #include "core/session/onnxruntime_cxx_api.h" #include "test/util/include/asserts.h" @@ -23,6 +28,14 @@ struct ApiPtrs { const gsl::not_null ep_api; }; +static void CheckStringInFile(const PathString& filename, const std::string& look_for) { + std::ifstream ifs{filename}; + std::string content(std::istreambuf_iterator{ifs}, + std::istreambuf_iterator{}); + + EXPECT_NE(content.find(look_for), std::string::npos); +} + // Normally, a plugin EP would be implemented in a separate library. // The `test_plugin_ep` namespace contains a local implementation intended for unit testing. namespace test_plugin_ep { @@ -114,6 +127,10 @@ MakeTestOrtEpResult MakeTestOrtEp(std::vector ep_devices = { return result; } +class MockKernelLookup : public IExecutionProvider::IKernelLookup { + const KernelCreateInfo* LookUpKernel(const Node& /*node*/) const override { return nullptr; } +}; + } // namespace test_plugin_ep TEST(PluginExecutionProviderTest, GetPreferredLayout) { @@ -317,4 +334,218 @@ TEST(PluginExecutionProviderTest, InferOrtDeviceFromDeviceMemoryInfo) { #endif // !defined(ORT_NO_EXCEPTIONS) } +static void LoadModelAndAssignNodesToEp(const ORTCHAR_T* model_path, + const char* ep_name, + const std::unordered_set& ep_node_names, + /*out*/ std::shared_ptr& model) { + ASSERT_STATUS_OK(Model::Load(model_path, model, nullptr, + DefaultLoggingManager().DefaultLogger())); + + Graph& graph = model->MainGraph(); + + for (Node& node : graph.Nodes()) { + if (ep_node_names.count(node.Name()) > 0) { + node.SetExecutionProviderType(ep_name); + } + } +} + +static OrtStatus* ORT_API_CALL GetCapabilityTakeAllNodesOneGroup(OrtEp* this_ptr, const OrtGraph* graph, + OrtEpGraphSupportInfo* graph_support_info) noexcept { + auto* this_ep = static_cast(this_ptr); + + size_t num_nodes = 0; + if (OrtStatus* st = this_ep->ort_api->Graph_GetNumNodes(graph, &num_nodes); st != nullptr) { + return st; + } + + std::vector nodes(num_nodes); + if (OrtStatus* st = this_ep->ort_api->Graph_GetNodes(graph, nodes.data(), nodes.size()); st != nullptr) { + return st; + } + + if (OrtStatus* st = this_ep->ep_api->EpGraphSupportInfo_AddNodesToFuse(graph_support_info, + nodes.data(), nodes.size(), nullptr); + st != nullptr) { + return st; + } + + return nullptr; +} + +static OrtStatus* ORT_API_CALL GetCapabilityTakeAllNodesTwoGroups(OrtEp* this_ptr, const OrtGraph* graph, + OrtEpGraphSupportInfo* graph_support_info) noexcept { + auto* this_ep = static_cast(this_ptr); + + size_t num_nodes = 0; + if (OrtStatus* st = this_ep->ort_api->Graph_GetNumNodes(graph, &num_nodes); st != nullptr) { + return st; + } + + std::vector nodes(num_nodes); + if (OrtStatus* st = this_ep->ort_api->Graph_GetNodes(graph, nodes.data(), nodes.size()); st != nullptr) { + return st; + } + + // Expect at least 2 nodes. If not, this is really a testing/setup error. + if (num_nodes < 2) { + return this_ep->ort_api->CreateStatus(OrtErrorCode::ORT_FAIL, + "Expected at least two nodes in call to GetCapability"); + } + + std::vector node_group1; + std::vector node_group2; + + for (size_t i = 0; i < num_nodes; i++) { + if (i < num_nodes / 2) { + node_group1.push_back(nodes[i]); + } else { + node_group2.push_back(nodes[i]); + } + } + + if (OrtStatus* st = this_ep->ep_api->EpGraphSupportInfo_AddNodesToFuse(graph_support_info, + node_group1.data(), node_group1.size(), + nullptr); + st != nullptr) { + return st; + } + + if (OrtStatus* st = this_ep->ep_api->EpGraphSupportInfo_AddNodesToFuse(graph_support_info, + node_group2.data(), node_group2.size(), + nullptr); + st != nullptr) { + return st; + } + + return nullptr; +} + +static OrtStatus* ORT_API_CALL GetCapabilityTakeSingleNode(OrtEp* this_ptr, const OrtGraph* graph, + OrtEpGraphSupportInfo* graph_support_info) noexcept { + auto* this_ep = static_cast(this_ptr); + + size_t num_nodes = 0; + if (OrtStatus* st = this_ep->ort_api->Graph_GetNumNodes(graph, &num_nodes); st != nullptr) { + return st; + } + + std::vector nodes(num_nodes); + if (OrtStatus* st = this_ep->ort_api->Graph_GetNodes(graph, nodes.data(), nodes.size()); st != nullptr) { + return st; + } + + // Take only the first node using EpGraphSupportInfo_AddSingleNode(). + if (OrtStatus* st = this_ep->ep_api->EpGraphSupportInfo_AddSingleNode(graph_support_info, nodes[0]); + st != nullptr) { + return st; + } + + return nullptr; +} + +// Tests that GetCapability() doesn't crash if a plugin EP tries to claim a mix of unassigned nodes and +// nodes that are already assigned to another EP. +TEST(PluginExecutionProviderTest, GetCapability_ClaimNodesAssignedToOtherEP) { + std::filesystem::path log_file = ORT_TSTR("log_get_capability.txt"); + + // Helper function that loads a model (Add -> Mul -> Add) and assigns some or all of the nodes to another EP. + // Then, IExecutionProvider::GetCapability() is called to test the expected behavior. + auto run_test = [&log_file](IExecutionProvider& ep, + const std::unordered_set& nodes_for_other_ep, + const std::unordered_set& nodes_for_this_ep, + const char* expected_log_string) { + std::shared_ptr model; + ASSERT_NO_FATAL_FAILURE(LoadModelAndAssignNodesToEp(ORT_TSTR("testdata/add_mul_add.onnx"), + "OtherEp", nodes_for_other_ep, model)); + + std::filesystem::remove(log_file); + + // Call IExecutionProvider::GetCapability and check results + logs. + { + logging::LoggingManager log_manager{std::make_unique(log_file, false, false), + logging::Severity::kWARNING, false, + logging::LoggingManager::InstanceType::Temporal}; + auto file_logger = log_manager.CreateLogger("FileLogger"); + ep.SetLogger(file_logger.get()); // Make EP log to a file. + + GraphViewer graph_viewer(model->MainGraph()); + auto compute_capabilities = ep.GetCapability(graph_viewer, + test_plugin_ep::MockKernelLookup{}, + GraphOptimizerRegistry(nullptr, nullptr, file_logger.get()), + nullptr); + + ASSERT_EQ(compute_capabilities.size(), nodes_for_this_ep.empty() ? 0 : 1); + + if (compute_capabilities.size() == 1) { + ASSERT_EQ(compute_capabilities[0]->sub_graph->nodes.size(), nodes_for_this_ep.size()); + + for (NodeIndex node_index : compute_capabilities[0]->sub_graph->nodes) { + const Node* node = graph_viewer.GetNode(node_index); + ASSERT_NE(node, nullptr); + EXPECT_EQ(nodes_for_this_ep.count(node->Name()), 1); + } + } + } + + ASSERT_TRUE(std::filesystem::exists(log_file)); + EXPECT_NO_FATAL_FAILURE(CheckStringInFile(log_file, expected_log_string)); + }; + + constexpr std::array node_names = {"add_0", "mul_0", "add_1"}; + + auto [ep, ort_ep] = test_plugin_ep::MakeTestOrtEp(); + + // Load a model and assign all of its nodes to another EP named 'OtherEp'. + // The plugin EP tries to claim all nodes in a single group via EpGraphSupportInfo_AddNodesToFuse. + // IExecutionProvider::GetCapability() should return an empty result and log a warning. + ort_ep->GetCapability = GetCapabilityTakeAllNodesOneGroup; + std::unordered_set nodes_for_other_ep = {"add_0", "mul_0", "add_1"}; + std::unordered_set nodes_for_this_ep; + run_test(*ep, nodes_for_other_ep, nodes_for_this_ep, + "Found one or more nodes that were already assigned to a different EP named 'OtherEp'"); + + // Load a model and assign only one node to another EP named 'OtherEp'. + // The plugin EP tries to claim all nodes in a single group. + // IExecutionProvider::GetCapability() should return an empty result and log a warning. + ort_ep->GetCapability = GetCapabilityTakeAllNodesOneGroup; + for (const char* node_name : node_names) { + nodes_for_other_ep = std::unordered_set{node_name}; + nodes_for_this_ep = std::unordered_set{}; + run_test(*ep, nodes_for_other_ep, nodes_for_this_ep, + "Found one or more nodes that were already assigned to a different EP named 'OtherEp'"); + } + + // Load a model and assign only the last Add node to another EP named 'OtherEp'. + // The plugin EP tries to claim all nodes in the following 2 groups: (add_0), (mul_0, add_1). + // IExecutionProvider::GetCapability() will only return (add_0) because the second group has a node + // that was assigned to 'OtherEp'. + ort_ep->GetCapability = GetCapabilityTakeAllNodesTwoGroups; + nodes_for_other_ep = std::unordered_set{"add_1"}; + nodes_for_this_ep = std::unordered_set{"add_0"}; + run_test(*ep, nodes_for_other_ep, nodes_for_this_ep, + "Found one or more nodes that were already assigned to a different EP named 'OtherEp'"); + + // Load a model and assign only the first Add node to another EP named 'OtherEp'. + // The plugin EP tries to claim all nodes in the following 2 groups: (add_0), (mul_0, add_1). + // IExecutionProvider::GetCapability() will only return (mul_0, add_1) because the first group has a node + // that was assigned to 'OtherEp'. + ort_ep->GetCapability = GetCapabilityTakeAllNodesTwoGroups; + nodes_for_other_ep = std::unordered_set{"add_0"}; + nodes_for_this_ep = std::unordered_set{"mul_0", "add_1"}; + run_test(*ep, nodes_for_other_ep, nodes_for_this_ep, + "Found one or more nodes that were already assigned to a different EP named 'OtherEp'"); + + // Load a model and assign the first Add node to another EP named 'OtherEp'. + // The plugin EP will try to take only the first Add node with a single call to EpGraphSupportInfo_AddSingleNode. + // IExecutionProvider::GetCapability() will return an empty result and log a warning. + ort_ep->GetCapability = GetCapabilityTakeSingleNode; + nodes_for_other_ep = std::unordered_set{"add_0"}; + nodes_for_this_ep = std::unordered_set{}; + run_test(*ep, nodes_for_other_ep, nodes_for_this_ep, + "Found one or more nodes that were already assigned to a different EP named 'OtherEp'"); + + std::filesystem::remove(log_file); +} + } // namespace onnxruntime::test