[ET-VK][runtime][1/3] Prepare linear shaders for 2D dispatch

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

@@ -36,8 +36,10 @@ layout(push_constant) uniform restrict Block {
 
 layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 
+#include "dispatch.glslh"
+
 void main() {
-  const int i = int(gl_GlobalInvocationID.x);
+  const int i = int(linear_idx_from_gid());
   if (i >= numel) {
     return;
   }

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

@@ -28,8 +28,10 @@ ${layout_declare_ubo(B, "float", "step")}
 
 layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 
+#include "dispatch.glslh"
+
 void main() {
-  const uint out_bufi = gl_GlobalInvocationID.x;
+  const uint out_bufi = linear_idx_from_gid();
   if (out_of_bounds(out_bufi, outp)) {
     return;
   }

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

@@ -53,6 +53,8 @@ layout(push_constant) uniform restrict Block {
 
 layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 
+#include "dispatch.glslh"
+
 ${layout_declare_spec_const(C, "int", "out_layout", "CONTIG_LAYOUT_INT")}
 ${layout_declare_spec_const(C, "int", "in_layout", "CONTIG_LAYOUT_INT")}
 ${layout_declare_spec_const(C, "int", "other_layout", "CONTIG_LAYOUT_INT")}
@@ -61,7 +63,7 @@ ${layout_declare_spec_const(C, "int", "in_broadcast_packed_dim", "0")}
 ${layout_declare_spec_const(C, "int", "other_broadcast_packed_dim", "0")}
 
 void main() {
-  const uint out_bufi = gl_GlobalInvocationID.x;
+  const uint out_bufi = linear_idx_from_gid();
   if (out_of_bounds(out_bufi, outp)) {
     return;
   }

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

@@ -57,11 +57,13 @@ layout(push_constant) uniform restrict Block {
 
 layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 
+#include "dispatch.glslh"
+
 $if not IS_COMPARISON_OP:
   #include "binary_op_defs.glslh"
 
 void main() {
-  const uint out_bufi = gl_GlobalInvocationID.x;
+  const uint out_bufi = linear_idx_from_gid();
   if (out_of_bounds(out_bufi, outp)) {
     return;
   }

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

@@ -32,12 +32,14 @@ $else:
 
 layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 
+#include "dispatch.glslh"
+
 ${layout_declare_spec_const(C, "int", "t_layout", "DEFAULT_LAYOUT")}
 const lowp ivec4 axis_map = unhash_axis_map(t_layout);
 const lowp int packed_dim = unhash_packed_dim(t_layout);
 
 void main() {
-  const int out_buf_idx = int(gl_GlobalInvocationID.x);
+  const int out_buf_idx = int(linear_idx_from_gid());
   // On the CPU, the number of elements is determined based on a buffer of int8
   // elements. However, on the GPU, since the int8 data type is not supported
   // each group of 4 elements is interepreted as 1 int32 element. Thus each

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

@@ -14,8 +14,10 @@ ${layout_declare_ubo(2, "int", "numel")}
 
 layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 
+#include "dispatch.glslh"
+
 void main() {
-  int tid = int(gl_GlobalInvocationID.x);
+  int tid = int(linear_idx_from_gid());
   if (tid >= numel) {
     return;
   }

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

@@ -19,12 +19,14 @@ ${layout_declare_ubo(B, "BufferMetadata", "inp")}
 
 layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 
+#include "dispatch.glslh"
+
 // This constant is unused in this shader but is kept so that the signature is
 // consistent with image_to_nchw.
 ${layout_declare_spec_const(C, "int", "unused", "0")}
 
 void main() {
-  uint inp_bufi = gl_GlobalInvocationID.x;
+  uint inp_bufi = linear_idx_from_gid();
   if (inp_bufi>= numel(inp)) {
     return;
   }

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

@@ -37,6 +37,8 @@ ${layout_declare_spec_const(C, "int", "out_layout", "CONTIG_LAYOUT_INT")}
 
 layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 
+#include "dispatch.glslh"
+
 #define NUM_INPUTS ${NUM_INPUTS}
 
 /*
@@ -45,7 +47,7 @@ layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
  * be inserted along the output's concat_dim starting at concat_offset.
  */
 void main() {
-  const int tid = int(gl_GlobalInvocationID.x);
+  const int tid = int(linear_idx_from_gid());
 
   // The 1-3 input tensors are interpreted as one concatenated tensor ("volume")
   // along the concat_dim for the purposes of tensor indexing. Each thread is

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

@@ -43,6 +43,8 @@ const int out_packed_dim = get_packed_dim(out_layout);
 
 layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 
+#include "dispatch.glslh"
+
 #define NUM_INPUTS ${NUM_INPUTS}
 
 /*
@@ -55,7 +57,7 @@ layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
  * input tensor.
  */
 void main() {
-  const int tid = int(gl_GlobalInvocationID.x);
+  const int tid = int(linear_idx_from_gid());
 
   // Compute inp_volume_sizes from output sizes, replacing concat_dim with the
   // sum of all input sizes along that dimension.

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

@@ -30,16 +30,19 @@ ${layout_declare_ubo(8, "float", "out_min", "float", "out_max")}
 
 layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 
+#include "dispatch.glslh"
+
 ${layout_declare_spec_const(C, "int", "ngroups", "1")}
 
 /*
  * Computes a depthwise convolution. Each shader invocation calculates the
  * output at a single output location.
  */
 void main() {
-  const uint div_by_x = gl_GlobalInvocationID.x / out_limits.x;
+  const uint linear_idx = linear_idx_from_gid();
+  const uint div_by_x = linear_idx / out_limits.x;
   const ivec3 pos = ivec3(
-    gl_GlobalInvocationID.x % out_limits.x,
+    linear_idx % out_limits.x,
     div_by_x % out_limits.y,
     div_by_x / out_limits.y);
 

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

@@ -46,6 +46,8 @@ layout(push_constant) uniform restrict Block {
 
 layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 
+#include "dispatch.glslh"
+
 ${layout_declare_spec_const(C, "int", "apply_bias", "1")}
 ${layout_declare_spec_const(C, "int", "conv2d_params_stride_x", "1")}
 ${layout_declare_spec_const(C, "int", "conv2d_params_stride_y", "1")}
@@ -69,7 +71,7 @@ ${layout_declare_spec_const(C, "int", "conv2d_params_groups", "1")}
 #include "conv2d_dw_q8_utils.glslh"
 
 void main() {
-  const int tid = int(gl_GlobalInvocationID.x);
+  const int tid = int(linear_idx_from_gid());
   Conv2dBlockExtents out_block_extents = make_block_extents(output_sizes);
 
   Conv2dBlockIndex out_block_idx = linear_idx_to_block_idx(

backends/vulkan/runtime/graph/ops/glsl/dispatch.glslh

@@ -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 DISPATCH_GLSLH
+#define DISPATCH_GLSLH
+
+uint linear_idx_from_gid() {
+  // Extend this to include Z if two-dimensional linear dispatches become
+  // insufficient.
+  return gl_GlobalInvocationID.x +
+      gl_GlobalInvocationID.y * gl_NumWorkGroups.x * gl_WorkGroupSize.x;
+}
+
+#endif // DISPATCH_GLSLH

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

@@ -32,6 +32,8 @@ ${layout_declare_ubo(B, "BufferMetadata", "weight")}
 
 layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 
+#include "dispatch.glslh"
+
 TensorIndex out_tidx_to_indices_tidx(const TensorIndex out_tidx) {
   TensorIndex indices_tidx;
   int d = 0;
@@ -56,7 +58,7 @@ T load_weight_elem(const int embedding_idx, const uint dim_idx) {
 }
 
 void main() {
-  const uint out_bufi = gl_GlobalInvocationID.x;
+  const uint out_bufi = linear_idx_from_gid();
   if (out_of_bounds(out_bufi, outp)) {
     return;
   }

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

@@ -27,8 +27,10 @@ ${layout_declare_ubo(B, "BufferMetadata", "inp")}
 
 layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 
+#include "dispatch.glslh"
+
 void main() {
-  const uint outp_bufi = gl_GlobalInvocationID.x;
+  const uint outp_bufi = linear_idx_from_gid();
   if (outp_bufi >= numel(outp)) {
     return;
   }

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

@@ -26,8 +26,10 @@ ${layout_declare_ubo(B, "float", "fill_value")}
 
 layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 
+#include "dispatch.glslh"
+
 void main() {
-  const uint out_bufi = gl_GlobalInvocationID.x;
+  const uint out_bufi = linear_idx_from_gid();
   if (out_of_bounds(out_bufi, outp)) {
     return;
   }

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

@@ -32,10 +32,12 @@ ${layout_declare_ubo(B, "BufferMetadata", "index")}
 
 layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 
+#include "dispatch.glslh"
+
 ${layout_declare_spec_const(C, "int", "gather_dim", "0")}
 
 void main() {
-  const uint out_bufi = gl_GlobalInvocationID.x;
+  const uint out_bufi = linear_idx_from_gid();
   if (out_of_bounds(out_bufi, outp)) {
     return;
   }

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

@@ -60,11 +60,13 @@ layout(push_constant) uniform restrict Block {
 
 layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 
+#include "dispatch.glslh"
+
 #include "conv2d_int8_output_tile_store.glslh"
 #include "im2col_packed_int8_utils.glslh"
 
 void main() {
-  const int out_buf_idx = int(gl_GlobalInvocationID.x);
+  const int out_buf_idx = int(linear_idx_from_gid());
 
   Conv2dBlockExtents im2col_block_extents = make_block_extents(im2col_sizes);
 

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

@@ -46,6 +46,8 @@ $if not TO_STAGING:
 
 layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 
+#include "dispatch.glslh"
+
 ${layout_declare_spec_const(C, "int", "in_layout", "CONTIG_LAYOUT_INT")}
 
 $if not TO_STAGING:
@@ -60,7 +62,7 @@ void main() {
   // the (VRAM-cached) texture reads -- even though each texel is now fetched up
   // to 4 times, once per component. On unified-memory (mobile) GPUs the extra
   // fetches are a net loss, so this variant is gated to discrete GPUs.
-  const int oi = int(gl_GlobalInvocationID.x);
+  const int oi = int(linear_idx_from_gid());
   const int W = inp.sizes.x;
   const int H = inp.sizes.y;
   const int C = inp.sizes.z;

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

@@ -30,12 +30,14 @@ ${layout_declare_ubo(B, "BufferMetadata", "index")}
 
 layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 
+#include "dispatch.glslh"
+
 // Implements aten.index.Tensor for the case where self is 1D and there is
 // exactly one index tensor. Each output element is:
 //   output[...] = self[index[...]]
 
 void main() {
-  const uint out_bufi = gl_GlobalInvocationID.x;
+  const uint out_bufi = linear_idx_from_gid();
   if (out_of_bounds(out_bufi, outp)) {
     return;
   }

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

@@ -26,12 +26,14 @@ ${layout_declare_ubo(B, "BufferMetadata", "inp")}
 
 layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 
+#include "dispatch.glslh"
+
 ${layout_declare_spec_const(C, "int", "inp_layout", "CONTIG_LAYOUT_INT")}
 
 void main() {
   // One thread per output int32 in the NCHW staging buffer.
   // Each output int32 holds 4 consecutive NCHW bytes.
-  const uint out_int32_idx = gl_GlobalInvocationID.x;
+  const uint out_int32_idx = linear_idx_from_gid();
 
   const uint W = inp.sizes[0][0];
   const uint H = inp.sizes[0][1];

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

@@ -28,6 +28,8 @@ ${layout_declare_ubo(B, "ivec4", "x_sizes")}
 
 layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 
+#include "dispatch.glslh"
+
 void main() {
   // dW[N, K] = sum_m d_out[m, N] * x[m, K]; contraction over the flattened M.
   const int N = dout_sizes.x;
@@ -38,7 +40,7 @@ void main() {
   const int nkt = (K + TILE_K - 1) / TILE_K;
   const int tiles = nnt * nkt;
 
-  const int tile_idx = int(gl_GlobalInvocationID.x);
+  const int tile_idx = int(linear_idx_from_gid());
   if (tile_idx >= tiles) {
     return;
   }

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

@@ -57,6 +57,8 @@ ${layout_declare_ubo(B, "ivec4", "input_sizes")}
 
 layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 
+#include "dispatch.glslh"
+
 ${layout_declare_spec_const(C, "int", "apply_bias", "0")}
 ${layout_declare_spec_const(C, "int", "K4_per_group", "0")}
 
@@ -73,7 +75,7 @@ shared FPOutTile partial_sums[WGS];
 
 void main() {
   const int lid = int(gl_LocalInvocationID.z);
-  const int n8 = int(gl_GlobalInvocationID.x);
+  const int n8 = int(linear_idx_from_gid());
 
   // The output tensor will have a shape of [n, 1, 1, 1]. Each thread computes
   // 8 output elements, so each thread will write to 8 elements starting at the

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

@@ -46,6 +46,8 @@ $else:
 
 layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 
+#include "dispatch.glslh"
+
 // This header file must be defined after the layout descriptors have been
 // declared because the functions in the header assume some variables have been
 // declared as layout descriptors.
@@ -57,7 +59,7 @@ layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 #endif
 
 void main() {
-  const int out_bufi = int(gl_GlobalInvocationID.x);
+  const int out_bufi = int(linear_idx_from_gid());
   if (out_bufi >= out_numel) {
     return;
   }
@@ -88,9 +90,10 @@ void main() {
 #else // USING_TEXTURE
 
 void main() {
+  const uint linear_idx = linear_idx_from_gid();
   const ivec2 out_pos = ivec2(
-    gl_GlobalInvocationID.x % out_limits.x,
-    gl_GlobalInvocationID.x / out_limits.x);
+    linear_idx % out_limits.x,
+    linear_idx / out_limits.x);
 
   if (out_pos.y >= out_limits.y) {
     return;