diff --git a/onnxruntime/core/framework/graph_partitioner.cc b/onnxruntime/core/framework/graph_partitioner.cc index 53e8a87c48..c96d932497 100644 --- a/onnxruntime/core/framework/graph_partitioner.cc +++ b/onnxruntime/core/framework/graph_partitioner.cc @@ -546,7 +546,12 @@ static Status PartitionOrtFormatModelImpl(Graph& graph, FuncManager& func_mgr, const IndexedSubGraph& indexed_sub_graph = *capability->sub_graph; const IndexedSubGraph::MetaDef* metadef = indexed_sub_graph.GetMetaDef(); if (!metadef) { - // Static kernel - use the kernel hash that was saved in the ORT format model + // Static kernel - use the kernel hash that was saved in the ORT format model. + auto* node = graph.GetNode(indexed_sub_graph.nodes[0]); + if (nullptr != node && node->GetExecutionProviderType().empty()) { + // The node was not fused or assigned. Assign it to this . + node->SetExecutionProviderType(type); + } continue; } @@ -602,6 +607,21 @@ static Status PartitionOrtFormatModelImpl(Graph& graph, FuncManager& func_mgr, return Status::OK(); } +// If this is an ORT format model the hashes will be for CPU EP kernels, so set the EP of any unassigned nodes +// to kCpuExecutionProvider. +static void AssignRemainingNodesToCpuEp(Graph& graph) { + for (auto& node : graph.Nodes()) { + for (auto& entry : node.GetAttributeNameToMutableSubgraphMap()) { + Graph* subgraph = entry.second; + AssignRemainingNodesToCpuEp(*subgraph); + } + + if (node.GetExecutionProviderType().empty()) { + node.SetExecutionProviderType(kCpuExecutionProvider); + } + } +} + // Simplified partitioning where custom EPs may produce compiled nodes. // EPs with static kernels do not need to be processed as their kernels are matched via hash information serialized // as part of the ORT format model. @@ -615,13 +635,13 @@ Status GraphPartitioner::PartitionOrtFormatModel( for (const auto& ep : providers_) { if (ep->Type() == kCpuExecutionProvider) { // hash for kernel is stored in session state for EPs that have pre-registered kernels - // (vs. runtime fused kernels) so nothing to do here. - continue; + // (vs. runtime fused kernels) so we can simply assign any remaining nodes to the CPU EP + AssignRemainingNodesToCpuEp(graph); + } else { + ORT_RETURN_IF_ERROR(PartitionOrtFormatModelImpl(graph, func_mgr, kernel_registry_mgr_, fused_kernel_registry, + *ep, compiled_kernel_hashes, fused_node_unique_id, + transform_layout_function)); } - - ORT_RETURN_IF_ERROR(PartitionOrtFormatModelImpl(graph, func_mgr, kernel_registry_mgr_, fused_kernel_registry, - *ep, compiled_kernel_hashes, fused_node_unique_id, - transform_layout_function)); } return Status::OK(); diff --git a/onnxruntime/core/optimizer/graph_transformer_utils.cc b/onnxruntime/core/optimizer/graph_transformer_utils.cc index 4b029f9176..b6cdacd994 100644 --- a/onnxruntime/core/optimizer/graph_transformer_utils.cc +++ b/onnxruntime/core/optimizer/graph_transformer_utils.cc @@ -293,11 +293,12 @@ InlinedVector> GenerateTransformers( } auto cpu_allocator = cpu_execution_provider.GetAllocator(0, OrtMemTypeDefault); transformers.emplace_back(std::make_unique(std::move(cpu_allocator))); - // NCHWCtransformer should have a higher priority versus this. Because NCHWCtransformer also do the similiar things + // NCHWCtransformer should have a higher priority versus this. Because NCHWCtransformer also do the similar things // of fusion patterns and target on CPU. However, NCHWCtransformer will reorder the layout to nchwc which is only available for // x86-64 cpu, not edge cpu like arm. But This tranformer could be used by opencl-ep/cpu-ep. So // we will prefer NhwcTransformer once ort runs on x86-64 CPU, otherwise ConvAddActivationFusion is enabled. - // this PR #6351 implemented similiar fusion-pattern but only for CUDA, and can only fuse conv-add-relu, while we can fuse more activation. + // PR #6351 implemented similar fusion-pattern for CUDA only, and can only fuse conv-add-relu, + // while we can fuse more activation. transformers.emplace_back(std::make_unique(cpu_ep)); #endif } break; @@ -361,6 +362,7 @@ InlinedVector> GenerateTransformersForMinimalB // currently the only level 3 optimizer is the NhwcTransformer which is fully supported at runtime if (!saving) { #ifndef DISABLE_CONTRIB_OPS + const InlinedHashSet cpu_ep = {onnxruntime::kCpuExecutionProvider}; auto cpu_allocator = cpu_execution_provider.GetAllocator(0, OrtMemTypeDefault); transformers.emplace_back(std::make_unique(std::move(cpu_allocator))); #else diff --git a/onnxruntime/core/optimizer/nhwc_transformer.cc b/onnxruntime/core/optimizer/nhwc_transformer.cc index be5e7aa39f..48e0a5413c 100644 --- a/onnxruntime/core/optimizer/nhwc_transformer.cc +++ b/onnxruntime/core/optimizer/nhwc_transformer.cc @@ -31,6 +31,11 @@ Status NhwcTransformer::ApplyImpl(Graph& graph, bool& modified, int graph_level, modified = false; for (std::unique_ptr& node : api_graph->Nodes()) { + // If the node is not supported in the CPU EP, skip it + if (node->GetExecutionProviderType() != kCpuExecutionProvider) { + continue; + } + // Only QLinearConv needs to be handled explicitly. The rest will be transformed if needed during transpose // optimization. if (node->OpType() == "QLinearConv") { diff --git a/onnxruntime/core/optimizer/nhwc_transformer.h b/onnxruntime/core/optimizer/nhwc_transformer.h index 0c981a59ce..f1871a3781 100644 --- a/onnxruntime/core/optimizer/nhwc_transformer.h +++ b/onnxruntime/core/optimizer/nhwc_transformer.h @@ -20,7 +20,7 @@ class NhwcTransformer : public GraphTransformer { AllocatorPtr cpu_allocator_; public: - explicit NhwcTransformer(AllocatorPtr cpu_allocator) noexcept + explicit NhwcTransformer(AllocatorPtr cpu_allocator) noexcept : GraphTransformer("NhwcTransformer"), cpu_allocator_(std::move(cpu_allocator)){}; private: diff --git a/onnxruntime/core/session/inference_session.cc b/onnxruntime/core/session/inference_session.cc index cbc96b0ad8..c7cfa14ff4 100644 --- a/onnxruntime/core/session/inference_session.cc +++ b/onnxruntime/core/session/inference_session.cc @@ -1399,12 +1399,8 @@ common::Status InferenceSession::Initialize() { #if !defined(ORT_MINIMAL_BUILD) || defined(ORT_EXTENDED_MINIMAL_BUILD) // nodes are already partitioned, but a custom EP may compile some at runtime. // run the partitioning to allow that to happen. - // - // We always have the CPU EP, so only need to run this if some other EP is enabled - if (execution_providers_.NumProviders() > 1) { - ORT_RETURN_IF_ERROR(PartitionOrtFormatModel(graph, execution_providers_, kernel_registry_manager_, - *session_state_)); - } + ORT_RETURN_IF_ERROR(PartitionOrtFormatModel(graph, execution_providers_, kernel_registry_manager_, + *session_state_)); const auto& cpu_ep = *execution_providers_.Get(onnxruntime::kCpuExecutionProvider); ORT_RETURN_IF_ERROR_SESSIONID_(