Enable xnnpack in aten mode

Summary: This diff enables XNNPack delegate in ATen mode resolving some compilation issues and also making sure that the deps that XNNPack depends on are portable + aten friendly.

Differential Revision: D70704202

Author: tarun292

backends/xnnpack/runtime/XNNExecutor.cpp

@@ -88,11 +88,16 @@ ET_NODISCARD Error XNNExecutor::prepare_args(EValue** args) {
     Tensor* tensor = &args[ext_id]->toTensor();
     externals_[i].data = tensor->mutable_data_ptr<float>();
 
+    executorch::aten::DimOrderType dim_order[kTensorDimensionLimit];
+
     // Reshape runtime inputs
     if (i < input_ids_.size()) {
       size_t num_dims = tensor->dim();
+      ET_CHECK_OR_RETURN_FALSE(
+        runtime::get_dim_order(*tensor, other_dim_order, num_dims),
+        "Failed to retrieve dim order from tensor!");
       ET_CHECK_OR_RETURN_ERROR(
-          is_contiguous_dim_order(tensor->dim_order().data(), tensor->dim()),
+          is_contiguous_dim_order(other_dim_order, tensor->dim()),
           Internal,
           "Expecting default dim_order but got a non default dim_order tensor for external input %u",
           i);
@@ -213,7 +218,7 @@ ET_NODISCARD Error XNNExecutor::resize_outputs(EValue** args) const {
         expected_output_size, static_cast<size_t>(num_dim)};
 
     ET_LOG(Debug, "Resizing output tensor to a new shape");
-    Error err = resize_tensor(*out_tensor, output_size);
+    Error err = runtime::resize_tensor(*out_tensor, output_size);
     if (err != Error::Ok) {
       ET_LOG(Error, "Failed to resize output tensor for XNNExecutor");
       return err;

backends/xnnpack/targets.bzl

@@ -1,5 +1,5 @@
 load("@fbsource//xplat/executorch/backends/xnnpack/third-party:third_party_libs.bzl", "third_party_dep")
-load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime")
+load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "get_aten_mode_options", "runtime")
 
 def _get_preprocessor_flags():
     """
@@ -29,39 +29,41 @@ def define_common_targets():
         ],
     )
 
-    runtime.cxx_library(
-        name = "xnnpack_backend",
-        srcs = native.glob([
-            "runtime/*.cpp",
-            "runtime/profiling/*.cpp",
-        ]),
-        headers = native.glob([
-            "runtime/*.h",
-            "runtime/profiling/*.h",
-        ]),
-        visibility = [
-            "//executorch/exir/backend:backend_lib",
-            "//executorch/exir/backend/test/...",
-            "//executorch/backends/xnnpack/test/...",
-            "//executorch/extension/pybindings/...",
-            "@EXECUTORCH_CLIENTS",
-        ],
-        preprocessor_flags = [
-            # Uncomment to enable per operator timings
-            # "-DENABLE_XNNPACK_PROFILING",
-            # Uncomment to enable using KleidiAI Kernels
-            # "-DENABLE_XNNPACK_KLEIDI"
-        ] + _get_preprocessor_flags(),
-        exported_deps = [
-            "//executorch/runtime/backend:interface",
-        ],
-        deps = [
-            third_party_dep("XNNPACK"),
-            "//executorch/backends/xnnpack/serialization:xnnpack_flatbuffer_header",
-            "//executorch/extension/threadpool:threadpool",
-            "//executorch/runtime/core/exec_aten/util:tensor_util",
-        ],
-        # XnnpackBackend.cpp needs to compile with executor as whole
-        # @lint-ignore BUCKLINT: Avoid `link_whole=True` (https://fburl.com/avoid-link-whole)
-        link_whole = True,
-    )
+    for aten_mode in get_aten_mode_options():
+        aten_suffix = "_aten" if aten_mode else ""
+        runtime.cxx_library(
+            name = "xnnpack_backend" + aten_suffix,
+            srcs = native.glob([
+                "runtime/*.cpp",
+                "runtime/profiling/*.cpp",
+            ]),
+            headers = native.glob([
+                "runtime/*.h",
+                "runtime/profiling/*.h",
+            ]),
+            visibility = [
+                "//executorch/exir/backend:backend_lib",
+                "//executorch/exir/backend/test/...",
+                "//executorch/backends/xnnpack/test/...",
+                "//executorch/extension/pybindings/...",
+                "@EXECUTORCH_CLIENTS",
+            ],
+            preprocessor_flags = [
+                # Uncomment to enable per operator timings
+                # "-DENABLE_XNNPACK_PROFILING",
+                # Uncomment to enable using KleidiAI Kernels
+                # "-DENABLE_XNNPACK_KLEIDI"
+            ] + _get_preprocessor_flags(),
+            exported_deps = [
+                "//executorch/runtime/backend:interface",
+            ],
+            deps = [
+                third_party_dep("XNNPACK"),
+                "//executorch/backends/xnnpack/serialization:xnnpack_flatbuffer_header",
+                "//executorch/extension/threadpool:threadpool",
+                "//executorch/runtime/core/exec_aten/util:tensor_util" + aten_suffix,
+            ],
+            # XnnpackBackend.cpp needs to compile with executor as whole
+            # @lint-ignore BUCKLINT: Avoid `link_whole=True` (https://fburl.com/avoid-link-whole)
+            link_whole = True,
+        )