support channels last inputs in xnnpack

Author: leafs1

.Package.swift/.watchman-cookie-madragna-mac-84335-443

@@ -40,3 +40,19 @@ xcuserdata/
 .swiftpm/
 *.xcworkspace/
 *.xcframework/
+
+# misc
+/.vscode/
+*.so
+*.dylib
+/cmake_wrapper.sh
+/data/
+/devtools/bundled_program/serialize/
+/exir/_serialize/program.fbs
+/exir/_serialize/scalar_type.fbs
+/include/
+/share/
+/version.py
+
+# Android
+*.aar

.ci/.watchman-cookie-madragna-mac-84335-443

@@ -12,6 +12,15 @@
             "args": [
                 "--model_path=./add.pte",
             ]
-        }
+        },
+        {
+            "name": "Debug python proj",
+            "type": "debugpy",
+            "request": "launch",
+            "module": "unittest",
+            "args": [
+                "./backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py"
+            ]
+        },
     ]
 }

.github/.watchman-cookie-madragna-mac-84335-443

@@ -62,8 +62,22 @@
         "algorithm": "cpp",
         "iterator": "cpp",
         "tuple": "cpp",
-        "span": "cpp"
+        "span": "cpp",
+        "*.inc": "cpp",
+        "alignedvector3": "cpp"
     },
     "C_Cpp.default.compilerPath": "/library/developer/commandlinetools/usr/bin/c++",
-    "python.analysis.typeCheckingMode": "off"
+    "python.analysis.typeCheckingMode": "off",
+    "python.testing.unittestArgs": [
+        "-v",
+        "-s",
+        "./backends",
+        "-p",
+        "test_*.py"
+    ],
+    "python.testing.pytestEnabled": true,
+    "python.testing.unittestEnabled": false,
+    "python.testing.pytestArgs": [
+        "."
+    ]
 }

.gitignore

@@ -0,0 +1,49 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <executorch/backends/vulkan/runtime/graph/ops/DynamicDispatchNode.h>
+
+#include <executorch/backends/vulkan/runtime/graph/ComputeGraph.h>
+
+namespace vkcompute {
+
+DynamicDispatchNode::DynamicDispatchNode(
+    ComputeGraph& graph,
+    const PickShaderFn& pick_shader_fn,
+    const PickGlobalFn& pick_global_wg_fn,
+    const PickLocalFn& pick_local_wg_fn,
+    const std::vector<ArgGroup>& args,
+    const vkapi::ParamsBindList& params,
+    const std::vector<PushConstantDataInfo>& push_constants,
+    const vkapi::SpecVarList& spec_vars,
+    const std::vector<ValueRef>& resize_args,
+    const ResizeFunction& resize_fn)
+    : DispatchNode(
+          graph,
+          pick_shader_fn(&graph, args, resize_args),
+          pick_global_wg_fn(&graph, args, resize_args),
+          pick_local_wg_fn(&graph, args, resize_args),
+          args,
+          params,
+          push_constants,
+          spec_vars,
+          resize_args,
+          resize_fn),
+      pick_shader_fn_(pick_shader_fn),
+      pick_global_wg_fn_(pick_global_wg_fn),
+      pick_local_wg_fn_(pick_local_wg_fn) {}
+
+void DynamicDispatchNode::encode(ComputeGraph* graph) {
+  shader_ = pick_shader_fn_(graph, args_, resize_args_);
+  global_workgroup_size_ = pick_global_wg_fn_(graph, args_, resize_args_);
+  local_workgroup_size_ =
+      utils::WorkgroupSize(pick_local_wg_fn_(graph, args_, resize_args_));
+  DispatchNode::encode(graph);
+}
+
+} // namespace vkcompute

.vscode/launch.json

@@ -0,0 +1,69 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <executorch/backends/vulkan/runtime/api/api.h>
+
+#include <executorch/backends/vulkan/runtime/graph/containers/PushConstantData.h>
+#include <executorch/backends/vulkan/runtime/graph/containers/Value.h>
+
+#include <executorch/backends/vulkan/runtime/graph/ops/DispatchNode.h>
+
+namespace vkcompute {
+
+class ComputeGraph;
+
+/*
+ * Represents a single shader execution op in a ML model.
+ */
+class DynamicDispatchNode final : public DispatchNode {
+  friend class ComputeGraph;
+
+ public:
+  using PickShaderFn = const std::function<vkapi::ShaderInfo(
+      ComputeGraph*,
+      const std::vector<ArgGroup>&,
+      const std::vector<ValueRef>&)>;
+  using PickGlobalFn = const std::function<utils::uvec3(
+      ComputeGraph*,
+      const std::vector<ArgGroup>&,
+      const std::vector<ValueRef>&)>;
+  using PickLocalFn = const std::function<utils::uvec3(
+      ComputeGraph*,
+      const std::vector<ArgGroup>&,
+      const std::vector<ValueRef>&)>;
+
+  explicit DynamicDispatchNode(
+      ComputeGraph& graph,
+      const PickShaderFn& pick_shader_fn,
+      const PickGlobalFn& pick_global_wg_fn,
+      const PickLocalFn& pick_local_wg_fn,
+      const std::vector<ArgGroup>& args,
+      const vkapi::ParamsBindList& params,
+      const std::vector<PushConstantDataInfo>& push_constants,
+      const vkapi::SpecVarList& spec_vars,
+      const std::vector<ValueRef>& resize_args,
+      const ResizeFunction& resize_fn = nullptr);
+
+  ~DynamicDispatchNode() override = default;
+
+  void encode(ComputeGraph* graph) override;
+
+ protected:
+  const PickShaderFn pick_shader_fn_;
+  const PickGlobalFn pick_global_wg_fn_;
+  const PickLocalFn pick_local_wg_fn_;
+
+ public:
+  operator bool() const {
+    return shader_;
+  }
+};
+
+} // namespace vkcompute

.vscode/settings.json

@@ -0,0 +1,45 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#version 450 core
+
+#define PRECISION ${PRECISION}
+
+layout(std430) buffer;
+
+${layout_declare_tensor(0, "w", "t_out", "float", "texture3d")}
+${layout_declare_tensor(1, "r", "t_in1", "float", "texture3d")}
+${layout_declare_tensor(2, "r", "t_in2", "float", "texture3d")}
+
+layout(push_constant) uniform restrict Block {
+  ivec4 out_sizes;
+  ivec4 in1_sizes;
+  ivec4 in2_sizes;
+};
+
+layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
+
+void main() {
+  const ivec3 pos = ivec3(gl_GlobalInvocationID);
+
+  if (any(greaterThanEqual(pos, out_sizes.xyz))) {
+    return;
+  }
+
+
+  vec4 out_texel = vec4(0.0);
+  for (int row = 0; row < in1_sizes.y; ++row) {
+    ivec3 in_pos = ivec3(pos.x, row, pos.z);
+    vec4 in1_texel = texelFetch(t_in1, in_pos, 0);
+    vec4 in2_texel = texelFetch(t_in2, in_pos, 0);
+
+    out_texel += in1_texel * in2_texel;
+  }
+
+  imageStore(t_out, pos, out_texel + ${OFFSET});
+}