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

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

@@ -14,7 +14,7 @@
 
 #define op(X, A, B) ${OPERATOR}
 
-#include "indexing_utils.h"
+#include "indexing_utils_u16.h"
 
 layout(std430) buffer;
 
@@ -35,10 +35,7 @@ layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
  * output at a single output location.
  */
 void main() {
-  const ivec3 pos = ivec3(
-    gl_GlobalInvocationID.x % out_limits.x,
-    (gl_GlobalInvocationID.x / out_limits.x) % out_limits.y,
-    gl_GlobalInvocationID.x / (out_limits.x * out_limits.y));
+  const ivec3 pos = idx_to_u16pos_x_wise(gl_GlobalInvocationID.x, out_limits.x, out_limits.y);
 
   if (any(greaterThanEqual(pos, out_limits))) {
     return;

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

@@ -14,6 +14,8 @@
 
 #define TILE_SIZE ${TILE_SIZE}
 
+#define BATCH_SIZE_X ${BATCH_SIZE_X}
+
 #define BATCH_SIZE_Y ${BATCH_SIZE_Y}
 
 #define op(X, A, B) ${OPERATOR}
@@ -41,73 +43,79 @@ layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
  * output at a single output location.
  */
 void main() {
-  // y divided up by batch size is used to determine 3d position
+  // x and y are divided by batch size to determine 3d position
   // since work size is calculated by x * ((y + B_Y - 1) / B_Y) * z
-  const uint out_limits_y_scaled = (out_limits.y + BATCH_SIZE_Y - 1) / BATCH_SIZE_Y;
+  const ivec2 out_limits_xy_scaled = (out_limits.xy + ivec2(BATCH_SIZE_X, BATCH_SIZE_Y) - 1) / ivec2(BATCH_SIZE_X, BATCH_SIZE_Y);
 
-  u16vec3 pos = u16vec3(
-    gl_GlobalInvocationID.x % out_limits.x,
-    ((gl_GlobalInvocationID.x / out_limits.x) % out_limits_y_scaled),
-    gl_GlobalInvocationID.x / (out_limits.x * out_limits_y_scaled));
+  ivec3 pos = idx_to_ipos_x_wise(gl_GlobalInvocationID.x, out_limits_xy_scaled.x, out_limits_xy_scaled.y);
 
-  // scale pos.y by batch size, because that's the top pixel to be processed
-  pos.y *= uint16_t(BATCH_SIZE_Y);
+  // scale pos.xy by batch sizes, because that's the top pixel to be processed
+  pos.x *= BATCH_SIZE_X;
+  pos.y *= BATCH_SIZE_Y;
 
   // do not process if top pixel does not fit within the output range
-  if (any(greaterThanEqual(u16vec3(pos.x, pos.y, pos.z), out_limits))) {
+  if (any(greaterThanEqual(pos, out_limits))) {
     return;
   }
 
   // Compute the index of the top-left element of the overlay region. Negative
   // indices indicate that the top-left element is in a region added by padding.
-  const u16vec2 ipos = pos.xy * u16vec2(stride) - u16vec2(padding);
+  const ivec2 ipos = pos.xy * stride - padding;
 
   // Compute the start and end of the input indices to load. Padding is assumed
   // to be constant 0 padding, so any reads from the padding region is skipped.
-  const u16vec2 start = ipos;
-  const u16vec2 end = ipos + u16vec2(overlay_region.xy);
+  const ivec2 start = ipos;
+  const ivec2 end = ipos + overlay_region.xy;
 
   // sum outputs
-  VEC4_T sum[BATCH_SIZE_Y];
+  VEC4_T sum[BATCH_SIZE_Y][BATCH_SIZE_X];
 
-  sum[0] = texelFetch(t_bias, u16vec2(pos.z, 0), 0);
-  for (int i = 1; i < BATCH_SIZE_Y; i++) {
-    sum[i] = sum[0];
+  sum[0][0] = texelFetch(t_bias, ivec2(pos.z, 0), 0);
+  for (int y = 0; y < BATCH_SIZE_Y; y++) {
+    for (int x = 0; x < BATCH_SIZE_X; x++) {
+      sum[y][x] = sum[0][0];
+    }
   }
 
   // array to store input texels
-  VEC4_T in_texels[TILE_SIZE];
+  VEC4_T in_texels[TILE_SIZE + BATCH_SIZE_X - 1];
 
   // array to store kernel data of previous y
   VEC4_T prev_kernel_line[TILE_SIZE];
 
-  uint16_t kx = uint16_t(0);
-  for (uint16_t y = start.y, i = uint16_t(0); i < uint16_t(TILE_SIZE + BATCH_SIZE_Y - 1); y += uint16_t(dilation.y), i++) {
-    for (uint16_t x = start.x, j = uint16_t(0); j < uint16_t(TILE_SIZE); x += uint16_t(dilation.x), j++) {
-      in_texels[int(j)] = texelFetch(t_in, u16vec3(x, y, pos.z), 0);
+  int kx = 0;
+  for (int y = start.y, i = 0; i < TILE_SIZE + BATCH_SIZE_Y - 1; y += dilation.y, i++) {
+    for (int x = start.x, j = 0; j < TILE_SIZE + BATCH_SIZE_X - 1; x += dilation.x, j++) {
+      in_texels[j] = texelFetch(t_in, ivec3(x, y, pos.z), 0);
     }
 
     // from 2nd iteration onwards accumulate dot product in 2nd sum
     // based on kernel line data fetched in previous iteration and input texel from this iteration
-    if (i > uint16_t(0)) {
-      for (uint16_t j = uint16_t(0); j < uint16_t(TILE_SIZE); j++) {
-        sum[1] = fma(in_texels[int(j)], prev_kernel_line[int(j)], sum[1]);
+    if (i > 0) {
+      for (int j = 0; j < TILE_SIZE; j++) {
+        for (int s = 0; s < BATCH_SIZE_X; s++) {
+          sum[1][s] = fma(in_texels[j + s], prev_kernel_line[j], sum[1][s]);
+        }
       }
     }
 
     // accumulate dot product in 1st sum only until tile size
-    if (i < uint16_t(TILE_SIZE)) {
-      for (uint16_t j = uint16_t(0); j < uint16_t(TILE_SIZE); j++, kx++) {
-        prev_kernel_line[int(j)] = texelFetch(t_kernel, u16vec2(kx, pos.z), 0);
-        sum[0] = fma(in_texels[int(j)], prev_kernel_line[int(j)], sum[0]);
+    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++) {
+            sum[0][s] = fma(in_texels[j + s], prev_kernel_line[j], sum[0][s]);
+        }
       }
     }
   }
 
-  for (int i = 0; i < BATCH_SIZE_Y; i++) {
-    if (any(greaterThanEqual(u16vec3(pos.x, pos.y + i, pos.z), out_limits))) {
-      continue;
+  for (int y = 0; y < BATCH_SIZE_Y; y++) {
+    for (int x = 0; x < BATCH_SIZE_X; x++) {
+      if (any(greaterThanEqual(ivec3(pos.x + x, pos.y + y, pos.z), out_limits))) {
+        continue;
+      }
+      imageStore(t_out, ivec3(pos.x + x, pos.y + y, pos.z), op(sum[y][x], out_min, out_max));
     }
-    imageStore(t_out, u16vec3(pos.x, pos.y + i, pos.z), op(sum[i], out_min, out_max));
   }
 }

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

@@ -10,6 +10,7 @@ conv2d_dw_output_tile:
     NDIM: 3
     DTYPE: float
     TILE_SIZE: 3
+    BATCH_SIZE_X: 4
     BATCH_SIZE_Y: 2
   generate_variant_forall:
     DTYPE:

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

@@ -16,7 +16,7 @@
 
 #define op(X, A, B) ${OPERATOR}
 
-#include "indexing_utils.h"
+#include "indexing_utils_u16.h"
 
 layout(std430) buffer;
 
@@ -43,13 +43,10 @@ shared u16vec2 pos_shared[gl_WorkGroupSize.x * gl_WorkGroupSize.y * gl_WorkGroup
  * size is only 1x1, making it easier to re-use loaded texels from t_kernel.
  */
 void main() {
-  const uvec2 out_limits_scaled = (out_limits.xy + TILE_SIZE - 1) / TILE_SIZE;
+  const ivec2 out_limits_scaled = (out_limits.xy + TILE_SIZE - 1) / TILE_SIZE;
   const uint shared_mem_stride = gl_WorkGroupSize.x * gl_WorkGroupSize.y * gl_WorkGroupSize.z;
 
-  const u16vec3 gpos = u16vec3(
-    gl_GlobalInvocationID.x % out_limits_scaled.x,
-    (gl_GlobalInvocationID.x / out_limits_scaled.x) % out_limits_scaled.y,
-    gl_GlobalInvocationID.x / (out_limits_scaled.x * out_limits_scaled.y));
+  const u16vec3 gpos = idx_to_u16pos_x_wise(gl_GlobalInvocationID.x, out_limits_scaled.x, out_limits_scaled.y);
 
   // Output position for TILE_SIZE = 2
   // +--------+--------+
@@ -98,7 +95,6 @@ void main() {
     const vec4 ktex_2 = texelFetchOffset(t_kernel, u16vec2(z, gpos.z), 0, u16vec2(2, 0));
     const vec4 ktex_3 = texelFetchOffset(t_kernel, u16vec2(z, gpos.z), 0, u16vec2(3, 0));
 
-
 #pragma unroll
     for (int i = 0; i < TILE_SIZE * TILE_SIZE; ++i) {
       const vec4 in_tex = texelFetch(t_in, u16vec3(ipos[i], z4), 0);

backends/vulkan/runtime/graph/ops/glsl/indexing_utils.h

@@ -223,6 +223,11 @@ ivec3 lpos_to_pos(const ivec3 lpos, const ivec4 axis_map) {
   return pos;
 }
 
+ivec3 idx_to_ipos_x_wise(uint idx, int size_x, int size_y) {
+  const uint div_by_x = idx / size_x;
+  return ivec3(idx % size_x, div_by_x % size_y, div_by_x / size_y);
+}
+
 #ifdef USING_BUFFER
 #define load_texel(buf, idx) buf[idx]
 #elif defined(USING_TEXTURE2D)

backends/vulkan/runtime/graph/ops/glsl/indexing_utils_u16.h

@@ -0,0 +1,19 @@
+/*
+ * 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.
+ */
+
+#ifndef INDEXING_UTILS_U16_H
+#define INDEXING_UTILS_U16_H
+
+#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require
+
+u16vec3 idx_to_u16pos_x_wise(uint idx, int size_x, int size_y) {
+  const uint div_by_x = idx / size_x;
+  return u16vec3(idx % size_x, div_by_x % size_y, div_by_x / size_y);
+}
+
+#endif // INDEXING_UTILS_U16_H

backends/vulkan/runtime/graph/ops/impl/Convolution.cpp

@@ -299,7 +299,7 @@ utils::uvec3 create_conv2d_global_wg_size(
   } else if (method == Conv2dMethod::Depthwise) {
     const utils::uvec3 image_extents = graph.logical_limits_of(out);
     return {
-        utils::div_up(image_extents[0u], 1u),
+        utils::div_up(image_extents[0u], 4u),
         utils::div_up(image_extents[1u], 2u),
         image_extents[2u]};
   } else {