[ET-VK] New implementation of cat operator

## Changes

* Introduce `concat_texture.glsl` and `concat_buffer.glsl` to implement the `torch.cat` operator
* Introduce `Concat.cpp` to replace `Cat.cpp`
* Fix a bug with channels-packed buffer tensors where input data would be copied incorrectly with multiple dims have a stride of 1

## Motivation

> * Introduce `concat_texture.glsl` and `concat_buffer.glsl` to implement the `torch.cat` operator
> * Introduce `Concat.cpp` to replace `Cat.cpp`

The existing implementation of `torch.cat` uses the copy_channel_offset` shaders. However, these shaders have a critical bug where the output tensor is passed in separately with difference access types, i.e.

```
graph.execute_nodes().emplace_back(new DispatchNode(
        graph,
        VK_KERNEL_FROM_STR(kernel_name),
        global_size,
        local_size,
        // Inputs and Outputs
        {
            {out, vkapi::kWrite},
            {out, vkapi::kRead},
            {in, vkapi::kRead},
        },
```

This creates many validation layer errors because the memory barriers for the resource cannot be formed properly. The shader essentially relies on undefined behaviour to work correctly. The result is that the `cat` operator produces incorrect result on many platforms.

Rather than fix the `copy_offset` shaders, I decided to just introduce new shaders to perform the concat operation. The new implementation handles both buffer and texture inputs and is agnostic to memory layout.

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

[ghstack-poisoned]

Author: SS-JIA

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

@@ -0,0 +1,67 @@
+/*
+ * 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}
+
+#define VEC4_T ${texel_type(DTYPE)}
+#define T ${buffer_scalar_type(DTYPE)}
+
+${define_active_storage_type("buffer")}
+${define_required_extensions(DTYPE)}
+
+layout(std430) buffer;
+
+#include "indexing_utils.h"
+
+${layout_declare_tensor(B, "w", "t_out", DTYPE, "buffer")}
+
+$for i in range(NUM_INPUTS):
+  ${layout_declare_tensor(B, "r", "t_in" + str(i + 1), DTYPE, "buffer")}
+
+${layout_declare_ubo(B, "int", "concat_dim")}
+
+${layout_declare_ubo(B, "ivec4", "out_sizes")}
+${layout_declare_ubo(B, "ivec4", "out_strides")}
+
+$for i in range(NUM_INPUTS):
+  ${layout_declare_ubo(B, "ivec4", "in" + str(i+1) + "_sizes")}
+  ${layout_declare_ubo(B, "ivec4", "in" + str(i+1) + "_strides")}
+
+${layout_declare_ubo(B, "int", "out_numel")}
+
+${layout_declare_spec_const(C, "int", "out_packed_dim", "DEFAULT_LAYOUT")}
+
+layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
+
+void main() {
+  const int out_bufi = ivec3(gl_GlobalInvocationID).x;
+  if (out_bufi >= out_numel) {
+    return;
+  }
+
+  // Convert buffer linear index to 4-D tensor index for output
+  const ivec4 out_tidx = bufi_to_tidx(out_bufi, out_strides, out_packed_dim);
+
+  // Determine which input tensor to read from
+  ivec4 in_tidx = out_tidx;
+
+  $for i in range(NUM_INPUTS):
+    // Check if the index at the concat dim is within bounds of the input tensor
+    // If so, read from that input tensor and write to output
+    if (in_tidx[concat_dim] < in${i+1}_sizes[concat_dim]) {
+      int in_bufi = tidx_to_bufi(in_tidx, in${i+1}_strides);
+      t_out[out_bufi] = t_in${i+1}[in_bufi];
+      return;
+    }
+    // otherwise, decrement the index at the concat dim
+    else {
+      in_tidx[concat_dim] -= in${i+1}_sizes[concat_dim];
+    }
+}

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

@@ -0,0 +1,14 @@
+concat_buffer:
+  parameter_names_with_default_values:
+    DTYPE: float
+    NUM_INPUTS: 2
+  generate_variant_forall:
+    DTYPE:
+      - VALUE: half
+      - VALUE: float
+  shader_variants:
+    - NAME: concat_1_buffer
+      NUM_INPUTS: 1
+    - NAME: concat_2_buffer
+    - NAME: concat_3_buffer
+      NUM_INPUTS: 3

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

@@ -0,0 +1,129 @@
+/*
+ * 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}
+
+#define VEC4_T ${texel_type(DTYPE)}
+#define T ${buffer_scalar_type(DTYPE)}
+
+#define USING_TEXTURE3D
+
+layout(std430) buffer;
+
+#include "indexing_utils.h"
+
+${layout_declare_tensor(B, "w", "t_out", DTYPE, "texture3d")}
+
+$for i in range(NUM_INPUTS):
+  ${layout_declare_tensor(B, "r", "t_in" + str(i + 1), DTYPE, "texture3d")}
+
+${layout_declare_ubo(B, "int", "concat_dim")}
+
+$in_metadata = ""
+$for i in range(NUM_INPUTS):
+  $in_metadata += "ivec4 in" + str(i + 1) + "_sizes;\n"
+
+layout(push_constant) uniform restrict Block {
+  ivec4 out_sizes;
+  ${in_metadata}
+};
+
+${layout_declare_spec_const(C, "int", "out_layout", "DEFAULT_LAYOUT")}
+const lowp ivec4 out_axis_map = unhash_axis_map(out_layout);
+const lowp int out_packed_dim = unhash_packed_dim(out_layout);
+
+$for i in range(NUM_INPUTS):
+  ${layout_declare_spec_const(C, "int", "in" + str(i+1) + "_layout", "DEFAULT_LAYOUT")}
+  const lowp ivec4 in${i+1}_axis_map = unhash_axis_map(in${i+1}_layout);
+  const lowp int in${i+1}_packed_dim = unhash_packed_dim(in${i+1}_layout);
+
+layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
+
+// Check if we can use the fast path (no texel merging required)
+bool can_use_fast_path() {
+  // Fast path is possible when:
+  // 1. The concat dimension is not the packed dimension, or
+  // 2. The concat dimension is the packed dimension but both input tensors have dimensions
+  //    that are multiples of 4 along the packed dimension
+  if (concat_dim != out_packed_dim) {
+    return true;
+  }
+
+  // Check if all input tensors have dimensions that are multiples of 4 along the packed dimension
+  bool all_concat_dim_size_multiple_of_4 = true;
+  $for i in range(NUM_INPUTS):
+    all_concat_dim_size_multiple_of_4 =
+        all_concat_dim_size_multiple_of_4 &&
+        (in${i+1}_sizes[concat_dim] % 4 == 0);
+
+  return all_concat_dim_size_multiple_of_4;
+}
+
+void main() {
+  const ivec3 lpos = ivec3(gl_GlobalInvocationID);
+  ivec4 out_tidx = lpos_to_tidx(lpos, out_sizes, out_axis_map.w, out_packed_dim);
+
+  if (any(greaterThanEqual(out_tidx, out_sizes))) {
+    return;
+  }
+
+  if (can_use_fast_path()) {
+    // Fast path: No texel merging required
+    ivec4 in_tidx = out_tidx;
+
+    $for i in range(NUM_INPUTS):
+      // For each input tensor, check if the tensor index is within bounds. If
+      // so, read the texel from the input tensor and write it to the output
+      if (in_tidx[concat_dim] < in${i+1}_sizes[concat_dim]) {
+        const ivec3 in_pos = tidx_to_pos(in_tidx, in${i+1}_sizes, in${i+1}_axis_map, in${i+1}_packed_dim);
+        const VEC4_T in_texel = load_texel(t_in${i+1}, in_pos);
+        write_texel_lpos(t_out, lpos, in_texel, out_axis_map);
+        return;
+      }
+      // Otherwise, adjust the index along the concat dimension and try the next
+      // input tensor.
+      else {
+        in_tidx[concat_dim] -= in${i+1}_sizes[concat_dim];
+      }
+  }
+  else {
+    // Slow path: Texel merging required
+    VEC4_T out_texel = VEC4_T(0);
+
+    // Process each element in the output texel individually
+    for (int texel_i = 0; texel_i < 4; ++texel_i) {
+      ivec4 curr_out_tidx = out_tidx;
+      curr_out_tidx[out_packed_dim] += texel_i;
+
+      // Skip if we're out of bounds
+      if (curr_out_tidx[out_packed_dim] >= out_sizes[out_packed_dim]) {
+        continue;
+      }
+
+      ivec4 in_tidx = curr_out_tidx;
+      $for i in range(NUM_INPUTS):
+        // For each input tensor, check if the tensor index is within bounds. If
+        // so, read the corresponding texel element from the input tensor and
+        // write it to the output texel.
+        if (in_tidx[concat_dim] < in${i+1}_sizes[concat_dim]) {
+          const ivec4 in_posi = tidx_to_posi(in_tidx, in${i+1}_sizes, in${i+1}_axis_map, in${i+1}_packed_dim);
+          out_texel[texel_i] = load_texel(t_in${i+1}, in_posi.xyz)[in_posi.w];
+          continue;
+        }
+        // Otherwise, adjust the index along the concat dimension and try the
+        // next input tensor.
+        else {
+          in_tidx[concat_dim] -= in${i+1}_sizes[concat_dim];
+        }
+    }
+
+    write_texel_lpos(t_out, lpos, out_texel, out_axis_map);
+  }
+}

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

@@ -0,0 +1,14 @@
+concat_texture:
+  parameter_names_with_default_values:
+    DTYPE: float
+    NUM_INPUTS: 2
+  generate_variant_forall:
+    DTYPE:
+      - VALUE: half
+      - VALUE: float
+  shader_variants:
+    - NAME: concat_1_texture3d
+      NUM_INPUTS: 1
+    - NAME: concat_2_texture3d
+    - NAME: concat_3_texture3d
+      NUM_INPUTS: 3

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

@@ -68,21 +68,6 @@
  */
 #define mod4(x) ((x) & 3)
 
-/*
- * Find the packed dimension of a tensor given its strides. The packed dimension
- * is the "fastest moving" dimension which will have a stride of 1.
- */
-int find_packed_dim(const ivec4 strides) {
-  int packed_dim = 0;
-  for (int i = 0; i <= 3; i++) {
-    if (strides[i] == 1) {
-      packed_dim = i;
-      break;
-    }
-  }
-  return packed_dim;
-}
-
 /*
  * Get the staging buffer indices that contain the data of the texel that
  * corresponds to the provided tensor index. Since the texel have 4 elements,
@@ -144,14 +129,6 @@ ivec4 bufi_to_tidx(int bufi, const ivec4 strides, const int packed_dim) {
   return idx;
 }
 
-// Convenience overload of the above function, which will determine the packed
-// dim from the strides automatically so it doesn't have to be passed in as a
-// function argument.
-ivec4 bufi_to_tidx(const int bufi, const ivec4 strides) {
-  int packed_dim = find_packed_dim(strides);
-  return bufi_to_tidx(bufi, strides, packed_dim);
-}
-
 int tidx_to_bufi(const ivec4 tidx, ivec4 strides) {
   return tidx.x * strides.x + tidx.y * strides.y + tidx.z * strides.z +
       tidx.w * strides.w;

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

@@ -28,7 +28,7 @@ layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 
 // This constant is unused in this shader but is kept so that the signature is
 // consistent with nchw_to_image.
-${layout_declare_spec_const(C, "int", "UNUSED_layout", "0")}
+${layout_declare_spec_const(C, "int", "packed_dim", "0")}
 ${layout_declare_spec_const(C, "int", "transpose_hw", "0")}
 
 void main() {
@@ -37,7 +37,7 @@ void main() {
     return;
   }
 
-  ivec4 out_tidx = bufi_to_tidx(out_bufi, out_strides);
+  ivec4 out_tidx = bufi_to_tidx(out_bufi, out_strides, packed_dim);
 
   ivec4 sizes = out_sizes;
   if (transpose_hw == 1) {