Update on "[ET-VK] Adding batch processing in x axis to conv2d dw shader by caching input texel for reuse."

This diff adds batch processing in the x axis to the conv2d dw shader by reusing input texel overlapping between consecutive tiles. The changes include modifying the glsl code for the conv2d dw output tile, adding a new parameter to the yaml file, and modifying the Convolution.cpp file to use the new parameter.

Differential Revision: [D67868671](https://our.internmc.facebook.com/intern/diff/D67868671/)

[ghstack-poisoned]

Author: trivedivivek

.lintrunner.toml

@@ -302,8 +302,8 @@ include_patterns = [
     'profiler/**/*.py',
     'runtime/**/*.py',
     'scripts/**/*.py',
-    # 'test/**/*.py',
-    # 'util/**/*.py',
+    'test/**/*.py',
+    'util/**/*.py',
     '*.py',
 ]
 exclude_patterns = [

.mypy.ini

@@ -21,10 +21,14 @@ files =
     profiler,
     runtime,
     scripts,
+    test,
     util
 
 mypy_path = executorch
 
+[mypy-executorch.backends.*]
+follow_untyped_imports = True
+
 [mypy-executorch.codegen.*]
 follow_untyped_imports = True
 
@@ -46,6 +50,12 @@ follow_untyped_imports = True
 [mypy-executorch.runtime.*]
 follow_untyped_imports = True
 
+[mypy-executorch.test.*]
+follow_untyped_imports = True
+
+[mypy-functorch.*]
+follow_untyped_imports = True
+
 [mypy-requests.*]
 follow_untyped_imports = True
 
@@ -80,4 +90,4 @@ ignore_missing_imports = True
 ignore_missing_imports = True
 
 [mypy-zstd]
-ignore_missing_imports = True
+ignore_missing_imports = True

backends/cadence/fusion_g3/operators/op_add.cpp

@@ -10,6 +10,7 @@
 
 #include <xa_nnlib_kernels_api.h>
 
+#include <executorch/backends/cadence/fusion_g3/operators/xt_macros.h>
 #include <executorch/kernels/portable/cpu/scalar_utils.h>
 #include <executorch/kernels/portable/cpu/util/elementwise_util.h>
 #include <executorch/kernels/portable/cpu/util/kernel_ops_util.h>
@@ -28,15 +29,6 @@ namespace impl {
 namespace G3 {
 namespace native {
 
-#define XT_KERNEL_CHECK(ctx, out, kernel, ...) \
-  const auto ret = kernel(__VA_ARGS__);        \
-  ET_KERNEL_CHECK_MSG(                         \
-      ctx,                                     \
-      ret == 0,                                \
-      InvalidArgument,                         \
-      out,                                     \
-      "Failed to run kernel: " #kernel "(" #__VA_ARGS__ ")");
-
 Tensor& add_out(
     KernelRuntimeContext& ctx,
     const Tensor& a,

backends/cadence/fusion_g3/operators/targets.bzl

@@ -27,6 +27,7 @@ def define_operator(name: str, deps: list[str] | None = None) -> None:
         deps = deps + common_deps,
         exported_deps = [
             ":operators_header",
+            ":xt_macros",
         ],
     )
 
@@ -61,5 +62,17 @@ def define_common_targets():
         ],
     )
 
+    runtime.cxx_library(
+        name = "xt_macros",
+        exported_headers = ["xt_macros.h"],
+        visibility = [
+            "//executorch/backends/cadence/...",
+        ],
+        exported_deps = [
+            "//executorch/runtime/core/exec_aten:lib",
+            "//executorch/runtime/kernel:kernel_runtime_context",
+        ],
+    )
+
     for op in OPERATORS:
         define_operator(op)

backends/cadence/fusion_g3/operators/xt_macros.h

@@ -0,0 +1,20 @@
+/*
+ * 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/runtime/core/exec_aten/util/tensor_util.h>
+
+#define XT_KERNEL_CHECK(ctx, out, kernel, ...) \
+  const auto ret = kernel(__VA_ARGS__);        \
+  ET_KERNEL_CHECK_MSG(                         \
+      ctx,                                     \
+      ret == 0,                                \
+      InvalidArgument,                         \
+      out,                                     \
+      "Failed to run kernel: " #kernel "(" #__VA_ARGS__ ")");

backends/vulkan/runtime/graph/ops/glsl/conv2d_dw_output_tile.glsl

@@ -100,7 +100,7 @@ void main() {
     }
 
     // accumulate dot product in 1st sum only until tile size
-    if (i < int(TILE_SIZE)) {
+    if (i < TILE_SIZE) {
       for (int j = 0; j < TILE_SIZE; j++, kx++) {
         prev_kernel_line[j] = texelFetch(t_kernel, ivec2(kx, pos.z), 0);
         for (int s = 0; s < BATCH_SIZE_X; s++) {

docs/TARGETS

@@ -1,7 +1,7 @@
 load("@fbcode_macros//build_defs:native_rules.bzl", "buck_filegroup", "buck_sh_test")
 load("@fbcode_macros//build_defs:python_binary.bzl", "python_binary")
 
-oncall("pytorch_r2p")
+oncall("executorch")
 
 python_binary(
     name = "sphinx",

extension/flat_tensor/serialize/TARGETS

@@ -0,0 +1,16 @@
+load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime")
+load(":targets.bzl", "define_common_targets")
+
+oncall("executorch")
+
+define_common_targets()
+
+runtime.python_library(
+    name = "schema",
+    srcs = [
+        "flat_tensor_schema.py",
+    ],
+    visibility = [
+        "//executorch/...",
+    ],
+)

extension/flat_tensor/flat_tensor.fbs renamed to extension/flat_tensor/serialize/flat_tensor.fbs

@@ -13,7 +13,7 @@ table TensorMetadata {
   scalar_type: executorch_flatbuffer.ScalarType;
 
   // Size of each dimension.
-  dim_sizes: [int32];
+  sizes: [int32];
 
   // Specifies in what order the dimensions are laid out in memory (from outer
   // to inner).

extension/flat_tensor/flat_tensor_schema.py renamed to extension/flat_tensor/serialize/flat_tensor_schema.py

@@ -17,7 +17,7 @@
 class TensorMetadata:
     fully_qualified_name: str
     scalar_type: ScalarType
-    dim_sizes: List[int]
+    sizes: List[int]
     dim_order: List[bytes]
 
     segment_index: int