[ET-VK][ops] Add batch support to q8ta convolutions

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

@@ -77,14 +77,18 @@ void main() {
   // Thread mapping
   int oc4 = int(gl_GlobalInvocationID.z);
   int w4 = int(gl_GlobalInvocationID.x);
+  const int H = int(outp.sizes[0][1]);
+  const int hn = int(gl_GlobalInvocationID.y);
+  const int n = hn / H;
+  const int h = hn % H;
 
   // Initialize output tensor index (WHCN order)
   // Each thread handles 4 adjacent widths starting at base_out_w
   TensorIndex4D outp_tidx;
   outp_tidx.data[0] = w4 * 4;
-  outp_tidx.data[1] = int(gl_GlobalInvocationID.y);
+  outp_tidx.data[1] = h;
   outp_tidx.data[2] = oc4 * 4;
-  outp_tidx.data[3] = 0;
+  outp_tidx.data[3] = n;
 
   const int W = int(outp.sizes[0][0]);
   const int OC = int(outp.sizes[0][2]);
@@ -113,6 +117,7 @@ void main() {
   const int inp_w_stride = int(inp.strides[0][0]);
   const int inp_h_stride = int(inp.strides[0][1]);
   const int inp_c_stride = int(inp.strides[0][2]);
+  const int inp_n_stride = int(inp.strides[0][3]);
   const int w_texel_step = conv2d_params.dilation.x * inp_w_stride;
   const int h_texel_step = conv2d_params.dilation.y * inp_h_stride;
   const int subtile_w_step = conv2d_params.stride.x * inp_w_stride;
@@ -122,7 +127,7 @@ void main() {
   inp_tidx.data[0] = outp_tidx.data[0] * conv2d_params.stride.x - conv2d_params.padding.x;
   inp_tidx.data[1] = outp_tidx.data[1] * conv2d_params.stride.y - conv2d_params.padding.y;
   inp_tidx.data[2] = ic_group_start;
-  inp_tidx.data[3] = 0;
+  inp_tidx.data[3] = n;
 
   int base_inp_texel_idx;
   if (get_outer_packed_dim_block_size(inp_layout) == 1) {
@@ -172,7 +177,11 @@ void main() {
               // inp_texel_idx = tensor4d_idx_to_texel_idx(inp, inp_tidx, inp_layout);
               const int w4 = div_4(inp_tidx.data[0]);
               const int inp_c4 = div_4(inp_tidx.data[2]);
-              inp_texel_idx = (inp_tidx.data[1] * inp_h_stride + w4 * inp_w_stride + inp_c4) * 4 + mod_4(inp_tidx.data[0]);
+              inp_texel_idx =
+                  (n * inp_n_stride + inp_tidx.data[1] * inp_h_stride +
+                   w4 * inp_w_stride + inp_c4) *
+                      4 +
+                  mod_4(inp_tidx.data[0]);
             }
             packed_input = t_packed_int8_input[inp_texel_idx];
           }

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

@@ -71,14 +71,18 @@ ivec4 quantize(const vec4 texel, const float inv_scale, const int zp) {
 
 void main() {
   const int c4 = int(gl_GlobalInvocationID.z);
+  const int H = int(outp.sizes[0][1]);
+  const int hn = int(gl_GlobalInvocationID.y);
+  const int n = hn / H;
+  const int h = hn % H;
 
   // Initialize output tensor index (WHCN order)
   // Each thread handles 4 adjacent widths starting at base_out_w
   TensorIndex4D outp_tidx;
   outp_tidx.data[0] = int(gl_GlobalInvocationID.x) * 4;
-  outp_tidx.data[1] = int(gl_GlobalInvocationID.y);
+  outp_tidx.data[1] = h;
   outp_tidx.data[2] = c4 * 4;
-  outp_tidx.data[3] = 0;
+  outp_tidx.data[3] = n;
 
   const int W = int(outp.sizes[0][0]);
   const int C4 = int(div_up_4(outp.sizes[0][2]));
@@ -94,6 +98,7 @@ void main() {
   // Get strides for width and height dimensions (in texel space)
   const int w_stride = int(inp.strides[0][0]);
   const int h_stride = int(inp.strides[0][1]);
+  const int n_stride = int(inp.strides[0][3]);
 
   // Pre-compute step sizes for efficient indexing
   const int w_texel_step = conv2d_params.dilation.x * w_stride;
@@ -106,7 +111,7 @@ void main() {
   inp_tidx.data[0] = outp_tidx.data[0] * conv2d_params.stride.x - conv2d_params.padding.x;
   inp_tidx.data[1] = outp_tidx.data[1] * conv2d_params.stride.y - conv2d_params.padding.y;
   inp_tidx.data[2] = outp_tidx.data[2];
-  inp_tidx.data[3] = 0;  // batch = 0 since N == 1
+  inp_tidx.data[3] = n;
 
   int base_inp_texel_idx;
   if (get_outer_packed_dim_block_size(inp_layout) == 1) {
@@ -152,7 +157,11 @@ void main() {
             // inp_texel_idx = base_inp_texel_idx + div_4(w_offset) * w_stride + mod_4(w_offset);
             // inp_texel_idx = tensor4d_idx_to_texel_idx(inp, inp_tidx, inp_layout);
             const int w4 = div_4(inp_tidx.data[0]);
-            inp_texel_idx = (inp_tidx.data[1] * h_stride + w4 * w_stride + c4) * 4 + mod_4(inp_tidx.data[0]);
+            inp_texel_idx =
+                (n * n_stride + inp_tidx.data[1] * h_stride +
+                 w4 * w_stride + c4) *
+                    4 +
+                mod_4(inp_tidx.data[0]);
           }
           const int packed_input = t_packed_int8_input[inp_texel_idx];
           input_4c = unpack_int8x4(packed_input);

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

@@ -73,14 +73,17 @@ ${layout_declare_spec_const(C, "int", "inp_layout", "CONTIG_LAYOUT_INT")}
 int compute_outp_buffer_idx(
     const int w_block_idx,
     const int h_idx,
-    const int c_block_idx) {
+    const int c_block_idx,
+    const int n_idx) {
   if (get_outer_packed_dim_block_size(outp_layout) == 1) {
-    return h_idx * int(outp.strides[0][1])
+    return n_idx * int(outp.strides[0][3])
+           + h_idx * int(outp.strides[0][1])
            + mul_4(w_block_idx) * int(outp.strides[0][0])
            + c_block_idx * int(outp.strides[0][2]);
   } else {
     return mul_4(
-      h_idx * int(outp.strides[0][1])
+      n_idx * int(outp.strides[0][3])
+      + h_idx * int(outp.strides[0][1])
       + w_block_idx * int(outp.strides[0][0])
       + c_block_idx * int(outp.strides[0][2]));
   }
@@ -91,20 +94,22 @@ void main() {
   // Thread mapping: each thread handles TILE_M widths x TILE_N output channels.
   // gl_GlobalInvocationID.x -> output channel blocks.
   // gl_GlobalInvocationID.y -> width blocks.
-  // gl_GlobalInvocationID.z → batch (or height * batch combined)
+  // gl_GlobalInvocationID.z -> height * batch.
   const int oc_block_idx = int(gl_GlobalInvocationID.x) * TILE_N4;
   const int ow_block_idx = int(gl_GlobalInvocationID.y) * TILE_M4;
-  const int oh = int(gl_GlobalInvocationID.z);
 
   // Get output extents in block space (div_up_4 for packed dimensions)
   const int W = int(outp.sizes[0][0]);
   const int W4 = div_up_4(int(outp.sizes[0][0]));
   const int H = int(outp.sizes[0][1]);
   const int OC4 = div_up_4(int(outp.sizes[0][2]));
+  const int hn = int(gl_GlobalInvocationID.z);
+  const int n = hn / H;
+  const int oh = hn % H;
 
   // Bounds check in block space
   if (ow_block_idx >= W4 ||
-      oh >= H ||
+      n >= int(outp.sizes[0][3]) ||
       oc_block_idx >= OC4) {
     return;
   }
@@ -118,6 +123,7 @@ void main() {
   const int inp_w_stride = int(inp.strides[0][0]);
   const int inp_h_stride = int(inp.strides[0][1]);
   const int inp_c_stride = int(inp.strides[0][2]);
+  const int inp_n_stride = int(inp.strides[0][3]);
 
   // Initialize int32 accumulator
   ivec4 out_accum[TILE_M][TILE_N4];
@@ -133,7 +139,8 @@ void main() {
   // Compute initial input tile index with group offset
   // For grouped im2col, each group's K range starts at group_idx * K4_per_group
   // For non-grouped (groups=1), group_idx is always 0 so offset is 0
-  int input_idx = oh * inp_h_stride
+  int input_idx = n * inp_n_stride
+                + oh * inp_h_stride
                 + ow_block_idx * inp_w_stride
                 + group_idx * K4_per_group;
 
@@ -256,7 +263,8 @@ void main() {
       const int base_outp_buffer_idx = compute_outp_buffer_idx(
           ow_block_idx + m4,
           oh,
-          oc_block_idx + n4);
+          oc_block_idx + n4,
+          n);
       if (oc_block_idx + n4 < OC4) {
         // Store individual ints from the ivec4
         const int subtile_w_limit = min(4, W - mul_4(ow_block_idx + m4));

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

@@ -50,15 +50,19 @@ void main() {
   const int im2col_W4 = div_up_4(im2col_sizes.x);
   const int im2col_H = im2col_sizes.y;
   const int im2col_Z4 = div_up_4(im2col_sizes.z);
+  const int im2col_N = im2col_sizes.w;
 
   // im2col block index from linear output buffer index
   const int c4_idx = out_buf_idx % im2col_Z4;
   const int row = out_buf_idx / im2col_Z4;
   const int w4_idx = row % im2col_W4;
-  const int h_idx = row / im2col_W4;
+  const int hn_idx = row / im2col_W4;
+  const int h_idx = hn_idx % im2col_H;
+  const int n_idx = hn_idx / im2col_H;
 
   // out of bounds check
-  if (w4_idx >= im2col_W4 || h_idx >= im2col_H || c4_idx >= im2col_Z4) {
+  if (w4_idx >= im2col_W4 || h_idx >= im2col_H ||
+      c4_idx >= im2col_Z4 || n_idx >= im2col_N) {
     return;
   }
 
@@ -108,12 +112,14 @@ void main() {
       const int x_mod = mod_4(x);
       int scalar_idx;
       if (get_outer_packed_dim_block_size(inp_layout) == 1) {
-        scalar_idx = input_y * int(inp.strides[0][1])
+        scalar_idx = n_idx * int(inp.strides[0][3])
+                     + input_y * int(inp.strides[0][1])
                      + x * int(inp.strides[0][0])
                      + z4 * int(inp.strides[0][2]);
       } else {
         scalar_idx = mul_4(
-            input_y * int(inp.strides[0][1])
+            n_idx * int(inp.strides[0][3])
+            + input_y * int(inp.strides[0][1])
             + x4 * int(inp.strides[0][0])
             + z4) + x_mod;
       }
@@ -122,7 +128,8 @@ void main() {
   }
 
   // store_packed_int8_output_tile (with TILE_M4=1, TILE_N4=1)
-  const int buffer_idx = h_idx * int(im2col_outp.strides[0][1])
+  const int buffer_idx = n_idx * int(im2col_outp.strides[0][3])
+                         + h_idx * int(im2col_outp.strides[0][1])
                          + w4_idx * int(im2col_outp.strides[0][0])
                          + c4_idx;
 

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

@@ -42,6 +42,52 @@ bool q8ta_conv2d_check_4w4c_packed_dim_info(const api::PackedDimInfo& info) {
       info.outer_packed_dim_block_size == 4;
 }
 
+namespace {
+
+uint64_t q8ta_conv2d_im2col_scratch_limit(ComputeGraph& graph) {
+  constexpr uint64_t kMaxBatchedIm2ColScratchBytes = 32ULL * 1024ULL * 1024ULL;
+  const uint64_t device_scratch_limit =
+      graph.context()->adapter_ptr()->max_buffer_numel();
+  return device_scratch_limit < kMaxBatchedIm2ColScratchBytes
+      ? device_scratch_limit
+      : kMaxBatchedIm2ColScratchBytes;
+}
+
+bool should_use_q8ta_conv2d_im2col(
+    ComputeGraph& graph,
+    const int64_t batch,
+    const int64_t groups,
+    const int64_t in_channels_per_group,
+    const int64_t flattened_kernel_size,
+    const int64_t out_height,
+    const int64_t out_width) {
+  const bool im2col_eligible = in_channels_per_group % 4 == 0;
+  if (!im2col_eligible) {
+    return false;
+  }
+
+  const int64_t spatial_out = out_height * out_width;
+  if (batch > 1) {
+    constexpr int64_t kMinFlattenedKernelSize = 1024;
+    constexpr int64_t kMaxSpatialOutput = 64;
+    const uint64_t scratch_bytes = static_cast<uint64_t>(batch) *
+        static_cast<uint64_t>(flattened_kernel_size) *
+        static_cast<uint64_t>(out_height) *
+        static_cast<uint64_t>(utils::align_up_4(out_width));
+    return groups == 1 && flattened_kernel_size >= kMinFlattenedKernelSize &&
+        spatial_out <= kMaxSpatialOutput &&
+        scratch_bytes <= q8ta_conv2d_im2col_scratch_limit(graph);
+  }
+
+  if (graph.device_is_mali()) {
+    return true;
+  }
+
+  return groups == 1 && (in_channels_per_group >= 32 || spatial_out <= 4096);
+}
+
+} // namespace
+
 //
 // Workgroup size selection functions
 //
@@ -70,13 +116,16 @@ GlobalWorkGrid pick_q8ta_conv2d_gwg(
   const uint32_t W = graph->size_at<uint32_t>(-1, output);
   const uint32_t H = graph->size_at<uint32_t>(-2, output);
   const uint32_t C = graph->size_at<uint32_t>(-3, output);
+  const uint32_t N = graph->size_at<uint32_t>(-4, output);
 
   // Each thread processes 4 adjacent width positions and 4 channels (4Wx4C
   // tile)
   const uint32_t W4 = utils::div_up_4(W);
   const uint32_t C4 = utils::div_up_4(C);
 
-  return GlobalWorkGrid({W4, H, C4}, kTiledWorkGrid);
+  return GlobalWorkGrid(
+      {W4, utils::safe_downcast<uint32_t>(static_cast<uint64_t>(H) * N), C4},
+      kTiledWorkGrid);
 }
 
 /**
@@ -466,33 +515,33 @@ void q8ta_conv2d_general(
 
 void q8ta_conv2d(ComputeGraph& graph, const std::vector<ValueRef>& args) {
   const ValueRef input = args.at(0);
+  const ValueRef kernel_size_ref = args.at(9);
   const ValueRef groups_ref = args.at(13);
   const ValueRef output = args.at(15);
 
   const int64_t groups = graph.extract_scalar<int64_t>(groups_ref);
   const int64_t in_channels = graph.size_at<int64_t>(-3, input);
   const int64_t in_channels_per_group = in_channels / groups;
+  const int64_t batch = graph.size_at<int64_t>(-4, input);
 
   const int64_t H_out = graph.size_at<int64_t>(-2, output);
   const int64_t W_out = graph.size_at<int64_t>(-1, output);
-  const int64_t spatial_out = H_out * W_out;
-
-  // Im2col requires input channels per group to be a multiple of 4
-  const bool im2col_eligible = in_channels_per_group % 4 == 0;
-
-  bool use_im2col = false;
-  if (graph.device_is_mali()) {
-    // On Mali, im2col is faster than the general shader across the board.
-    use_im2col = im2col_eligible;
-  } else {
-    // Default: on Adreno and unknown GPU architectures, im2col is only
-    // beneficial for ungrouped convolutions with sufficient channel depth or
-    // small spatial output. For grouped convolutions, the general shader is
-    // more efficient (0.7-0.95x regression measured on Adreno).
-    use_im2col = im2col_eligible && groups == 1 &&
-        (in_channels_per_group >= 32 || spatial_out <= 4096);
+  int64_t flattened_kernel_size;
+  {
+    const auto kernel_size = graph.get_int_list(kernel_size_ref);
+    flattened_kernel_size = utils::align_up_4(
+        in_channels_per_group * kernel_size->at(0) * kernel_size->at(1));
   }
 
+  const bool use_im2col = should_use_q8ta_conv2d_im2col(
+      graph,
+      batch,
+      groups,
+      in_channels_per_group,
+      flattened_kernel_size,
+      H_out,
+      W_out);
+
   if (use_im2col) {
     q8ta_conv2d_im2col(graph, args);
   } else {

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

@@ -34,13 +34,16 @@ GlobalWorkGrid pick_q8ta_conv2d_dw_gwg(
   const uint32_t W = graph->size_at<uint32_t>(-1, output);
   const uint32_t H = graph->size_at<uint32_t>(-2, output);
   const uint32_t C = graph->size_at<uint32_t>(-3, output);
+  const uint32_t N = graph->size_at<uint32_t>(-4, output);
 
   // Each thread processes 4 adjacent width positions and 4 channels (4Wx4C
   // tile)
   const uint32_t W4 = utils::div_up_4(W);
   const uint32_t C4 = utils::div_up_4(C);
 
-  return GlobalWorkGrid({W4, H, C4}, kTiledWorkGrid);
+  return GlobalWorkGrid(
+      {W4, utils::safe_downcast<uint32_t>(static_cast<uint64_t>(H) * N), C4},
+      kTiledWorkGrid);
 }
 
 LocalWorkGroup pick_q8ta_conv2d_dw_lwg(