make nhwc transformer only apply to CPU ep. (#11882)

QLinearConv does not work with DML EP because this optimizer intended for CPU EP is wrongfully applied to it.

Limit NHWC optimizer to nodes assigned to the CPU EP

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 <provider>.
+        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();

onnxruntime/core/optimizer/graph_transformer_utils.cc

@@ -293,11 +293,12 @@ InlinedVector<std::unique_ptr<GraphTransformer>> GenerateTransformers(
       }
       auto cpu_allocator = cpu_execution_provider.GetAllocator(0, OrtMemTypeDefault);
       transformers.emplace_back(std::make_unique<NhwcTransformer>(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<ConvAddActivationFusion>(cpu_ep));
 #endif
     } break;
@@ -361,6 +362,7 @@ InlinedVector<std::unique_ptr<GraphTransformer>> GenerateTransformersForMinimalB
       // currently the only level 3 optimizer is the NhwcTransformer which is fully supported at runtime
       if (!saving) {
 #ifndef DISABLE_CONTRIB_OPS
+        const InlinedHashSet<std::string_view> cpu_ep = {onnxruntime::kCpuExecutionProvider};
         auto cpu_allocator = cpu_execution_provider.GetAllocator(0, OrtMemTypeDefault);
         transformers.emplace_back(std::make_unique<NhwcTransformer>(std::move(cpu_allocator)));
 #else

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<api::NodeRef>& 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") {

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:

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_(