[ET-VK] Modernize permute op with safe indexing and unified dispatch

Modernize the permute operator to follow current best practices, fixing an
Adreno 740 driver crash caused by dynamic UBO indexing in the texture shader.

Texture shader changes:
- Replace old indexing_utils.h with indexing.glslh
- Use TextureMetadata UBOs instead of push constant sizes
- Use texture_pos_to_tensor4d_idx_simple() and related helpers
- Replace permute_dims[out_packed_dim] with safe_idx() to avoid
  dynamic indexing of push constant with spec-const-derived index
- Use TextureElementIndex pattern for the slow path

C++ dispatch changes:
- Merge add_permute_node() and add_permute_buffer_node() into a single
  unified function using graph.meta_ubo() and conditional logic
- Remove unused channel_info computation
- Move WHCNPermuteDims struct into anonymous namespace
- Guard texture path with VK_CHECK_COND(permute_ndim <= 4)

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

ghstack-source-id: 357844381
Pull Request resolved: #18511

Author: ssjia

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

@@ -9,97 +9,97 @@
 #version 450 core
 
 ${define_required_extensions("texture3d", DTYPE)}
-${define_explicit_type_extensions(DTYPE)}
 
 #define PRECISION ${PRECISION}
 
-#define VEC4_T ${texel_type(DTYPE)}
-#define T ${buffer_scalar_type(DTYPE)}
+#define VEC4_T ${texel_load_type(DTYPE, "texture3d")}
+#define T ${texel_load_component_type(DTYPE, "texture3d")}
 
 ${define_active_storage_type("texture3d")}
 
+#extension GL_EXT_control_flow_attributes : require
+
 layout(std430) buffer;
 
-#include "indexing_utils.h"
+#include "common.glslh"
+#include "indexing.glslh"
 
 ${layout_declare_tensor(B, "w", "t_out", DTYPE, "texture3d")}
 ${layout_declare_tensor(B, "r", "t_in", DTYPE, "texture3d")}
 
+${layout_declare_ubo(B, "TextureMetadata", "outp")}
+${layout_declare_ubo(B, "TextureMetadata", "inp")}
+
 layout(push_constant) uniform restrict Block {
-  ivec4 out_sizes;
-  ivec4 in_sizes;
-  ivec4 permute_dims; // Permutation mapping: permute_dims[i] = j means output dim i comes from input dim j
+  ivec4 permute_dims;
 };
 
-${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);
-
-${layout_declare_spec_const(C, "int", "in_layout", "DEFAULT_LAYOUT")}
-const lowp ivec4 in_axis_map = unhash_axis_map(in_layout);
-const lowp int in_packed_dim = unhash_packed_dim(in_layout);
+${layout_declare_spec_const(C, "int", "out_layout", "CONTIG_LAYOUT_INT")}
+${layout_declare_spec_const(C, "int", "in_layout", "CONTIG_LAYOUT_INT")}
+const int out_packed_dim = get_packed_dim(out_layout);
+const int in_packed_dim = get_packed_dim(in_layout);
 
 layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 
-// Convert output tensor index to input tensor index based on permutation
+// Convert output tensor index to input tensor index based on permutation.
+// permute_dims[i] = j means output dim i comes from input dim j.
+// We write: in_tidx[permute_dims.{x,y,z,w}] = out_tidx.{x,y,z,w}
+// This uses literal component access on the push constant (safe) and dynamic
+// indexing into the local in_tidx variable (also safe).
 ivec4 out_tidx_to_in_tidx(const ivec4 out_tidx) {
   ivec4 in_tidx;
-
-  // Apply the permutation mapping: in_tidx[permute_dims[i]] = out_tidx[i]
   in_tidx[permute_dims.x] = out_tidx.x;
   in_tidx[permute_dims.y] = out_tidx.y;
   in_tidx[permute_dims.z] = out_tidx.z;
   in_tidx[permute_dims.w] = out_tidx.w;
-
   return in_tidx;
 }
 
-// Check if we can use the fast path where texels from the input tensor can be
-// copied directly into the output tensor. This occurs when the packed dimension
-// is preserved in the permutation, i.e. reading a texel from the output tensor
-// produces 4 texels along the same dimension as reading a texel from the input
-// tensor.
-bool can_use_fast_path() {
-  // Fast path is possible when the packed dimension is preserved in the permutation
-  // This means permute_dims[out_packed_dim] == in_packed_dim
-  return permute_dims[out_packed_dim] == in_packed_dim;
-}
-
 void main() {
-  const ivec3 lpos = ivec3(gl_GlobalInvocationID);
-  ivec4 out_tidx = lpos_to_tidx(lpos, out_sizes, out_axis_map.w, out_packed_dim);
+  const ivec3 out_pos = ivec3(gl_GlobalInvocationID);
 
-  if (any(greaterThanEqual(out_tidx, out_sizes))) {
+  if (out_of_bounds(out_pos, outp)) {
     return;
   }
 
-  if (can_use_fast_path()) {
+  TensorIndex4D out_tidx =
+      texture_pos_to_tensor4d_idx_simple(outp, out_pos, out_layout);
+
+  // Check if packed dimension is preserved in the permutation. Use safe_idx
+  // to avoid dynamic indexing of push constant with spec-const-derived index.
+  const bool fast_path =
+      safe_idx(permute_dims, out_packed_dim) == in_packed_dim;
+
+  if (fast_path) {
     // Fast path: packed dimension is preserved, so we can copy texels directly
-    ivec4 in_tidx = out_tidx_to_in_tidx(out_tidx);
-    ivec3 in_pos = tidx_to_pos(in_tidx, in_sizes, in_axis_map, in_packed_dim);
-    VEC4_T in_texel = VEC4_T(load_texel(t_in, in_pos));
+    ivec4 in_tidx_data = out_tidx_to_in_tidx(out_tidx.data);
+    TensorIndex4D in_tidx;
+    in_tidx.data = in_tidx_data;
 
-    write_texel_lpos(t_out, lpos, in_texel, out_axis_map);
-  }
-  else {
+    ivec3 in_pos =
+        tensor4d_idx_to_texel_pos_simple(inp, in_tidx, in_layout);
+    VEC4_T in_texel = texelFetch(t_in, in_pos, 0);
+
+    imageStore(t_out, out_pos, in_texel);
+  } else {
     // Slow path: packed dimension is not preserved, so each element of the
-    // output texel may be "sourced" from a different texel in the input tensor.
-    // Therefore each output texel element is processed individually.
+    // output texel may come from a different texel in the input tensor.
     VEC4_T out_texel = VEC4_T(0);
 
-    for (int texel_i = 0; texel_i < 4; ++texel_i) {
-      ivec4 in_tidx = out_tidx_to_in_tidx(out_tidx);
-      ivec3 in_pos = tidx_to_pos(in_tidx, in_sizes, in_axis_map, in_packed_dim);
-      int element_idx = in_tidx[in_packed_dim] % 4;
+    for (int comp = 0; comp < 4; comp++) {
+      ivec4 in_tidx_data = out_tidx_to_in_tidx(out_tidx.data);
+      TensorIndex4D in_tidx;
+      in_tidx.data = in_tidx_data;
 
-      VEC4_T in_texel = VEC4_T(load_texel(t_in, in_pos));
-      T selected_value = T(in_texel[element_idx]);
+      TextureElementIndex in_elem =
+          tensor4d_idx_to_texture_element_idx_simple(inp, in_tidx, in_layout);
 
-      out_texel[texel_i] = selected_value;
+      VEC4_T in_texel = texelFetch(t_in, in_elem.pos, 0);
+      out_texel[comp] = in_texel[in_elem.comp];
 
-      out_tidx[out_packed_dim]++;
+      out_tidx.data[out_packed_dim]++;
     }
 
-    write_texel_lpos(t_out, lpos, out_texel, out_axis_map);
+    imageStore(t_out, out_pos, out_texel);
   }
 }