From 31cb3cb254615a1414e7323a4da9ff0b1a0f727d Mon Sep 17 00:00:00 2001
From: Patrice Vignola <vignola.patrice@gmail.com>
Date: Thu, 10 Nov 2022 00:56:23 -0800
Subject: [PATCH] [DML EP] Revert DML's cpu fallback logic (#13605)

### Description
Revert DML's CPU fallback logic from
https://github.com/microsoft/onnxruntime/pull/13442.

### Motivation and Context
Although the logic works great in many models that have good DML
coverage, it makes perf worse in some models where many operators are
missing DML coverage (e.g. int64). Overall, the right fix seems to
instead implement the operator on DML even though it almost always falls
back to the CPU, just for the sake of having a registration.
---
 .../core/framework/fallback_cpu_capability.cc    | 16 ++++++----------
 .../core/framework/fallback_cpu_capability.h     |  4 +---
 .../src/ExecutionProvider.cpp                    | 15 ++++++++++++++-
 3 files changed, 21 insertions(+), 14 deletions(-)

diff --git a/onnxruntime/core/framework/fallback_cpu_capability.cc b/onnxruntime/core/framework/fallback_cpu_capability.cc
index 8ee958fdbc..ab8266cf93 100644
--- a/onnxruntime/core/framework/fallback_cpu_capability.cc
+++ b/onnxruntime/core/framework/fallback_cpu_capability.cc
@@ -40,8 +40,7 @@ static bool IsSmallInitializer(const onnxruntime::GraphViewer& graph, const Node
 
 std::unordered_set<NodeIndex> GetCpuPreferredNodes(const onnxruntime::GraphViewer& graph,
                                                    const IExecutionProvider::IKernelLookup& kernel_lookup,
-                                                   gsl::span<const NodeIndex> tentative_nodes,
-                                                   bool requires_matching_kernel) {
+                                                   gsl::span<const NodeIndex> tentative_nodes) {
   // automatic conversion from const std::vector&
   const auto& ordered_nodes = graph.GetNodesInTopologicalOrder();
   InlinedVector<size_t> node_id_to_order_map(graph.MaxNodeIndex());
@@ -70,18 +69,15 @@ std::unordered_set<NodeIndex> GetCpuPreferredNodes(const onnxruntime::GraphViewe
     const Node* node = graph.GetNode(node_id);
 
     const KernelCreateInfo* kernel_info = kernel_lookup.LookUpKernel(*node);
-
-    if (requires_matching_kernel) {
-      // at least one registry has a target provider's kernel for this node
-      ORT_ENFORCE(kernel_info != nullptr);
-      node_to_kernel.insert({node_id, kernel_info});
-    }
+    // at least one registry has a target provider's kernel for this node
+    ORT_ENFORCE(kernel_info != nullptr);
+    node_to_kernel.insert({node_id, kernel_info});
 
     // first, find all the direct consumer of cpu tensors.
     ORT_THROW_IF_ERROR(node->ForEachWithIndex(
         node->OutputDefs(),
         [&](const NodeArg& node_arg, size_t out_index) {
-          if ((!requires_matching_kernel && kernel_info == nullptr) || kernel_info->kernel_def->IsOutputOnCpu(out_index)) {
+          if (kernel_info->kernel_def->IsOutputOnCpu(out_index)) {
             cpu_output_args.insert(&node_arg);
             auto consumer_nodes = graph.GetConsumerNodes(node_arg.Name());
             for (auto& consumer_node : consumer_nodes) {
@@ -139,7 +135,7 @@ std::unordered_set<NodeIndex> GetCpuPreferredNodes(const onnxruntime::GraphViewe
       }
 
       // input is a CPU tensor, but it's intended to be consumed as CPU input by the target EP
-      if (requires_matching_kernel && node_to_kernel[cur]->kernel_def->IsInputOnCpu(i)) {
+      if (node_to_kernel[cur]->kernel_def->IsInputOnCpu(i)) {
         place_in_cpu = false;
         break;
       }
diff --git a/onnxruntime/core/framework/fallback_cpu_capability.h b/onnxruntime/core/framework/fallback_cpu_capability.h
index f8e8fdc359..afc518d28a 100644
--- a/onnxruntime/core/framework/fallback_cpu_capability.h
+++ b/onnxruntime/core/framework/fallback_cpu_capability.h
@@ -16,11 +16,9 @@ namespace onnxruntime {
   @param graph Graph viewer
   @param kernel_lookup The kernel lookup for the target execution provider
   @param tentative_nodes Nodes that are tentative to be placed on on target EP
-  @param requires_matching_kernel If no matching kernels are found, assumes that the kernel should be on the CPU
   */
   std::unordered_set<NodeIndex> GetCpuPreferredNodes(const GraphViewer& graph,
                                                     const IExecutionProvider::IKernelLookup& kernel_lookup,
-                                                    gsl::span<const NodeIndex> tentative_nodes,
-                                                    bool requires_matching_kernel = true);
+                                                    gsl::span<const NodeIndex> tentative_nodes);
 
 }  // namespace onnxruntime
diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.cpp
index b2681cd82f..09314c4db3 100644
--- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.cpp
+++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.cpp
@@ -668,8 +668,21 @@ namespace Dml
         // downstream nodes consuming them.
         const std::vector<onnxruntime::NodeIndex>& toplogicalOrder = graph.GetNodesInTopologicalOrder();
 
+        std::vector<onnxruntime::NodeIndex> tentativeNodes;
+        tentativeNodes.reserve(toplogicalOrder.size());
+
+        for (onnxruntime::NodeIndex nodeIndex : toplogicalOrder)
+        {
+            const onnxruntime::Node& node = *graph.GetNode(nodeIndex);
+            const auto* kernelInfo = kernel_lookup.LookUpKernel(node);
+            if (kernelInfo != nullptr)
+            {
+                tentativeNodes.push_back(nodeIndex);
+            }
+        }
+
         // Get the list of nodes that should stay on the CPU
-        auto cpuPreferredNodes = GetCpuPreferredNodes(graph, kernel_lookup, toplogicalOrder, false);
+        auto cpuPreferredNodes = GetCpuPreferredNodes(graph, kernel_lookup, tentativeNodes);
 
         for (size_t nodeIndex : toplogicalOrder)
         {