aten.select.int

backends/vulkan/runtime/api/Tensor.h

@@ -255,6 +255,14 @@ class vTensor final {
     return sizes_;
   }
 
+  inline const int64_t size(size_t dim) const {
+    return sizes().at(dim);
+  }
+
+  inline const int64_t dim() const {
+    return sizes_.size();
+  }
+
   inline const std::vector<int64_t>& strides() const {
     return strides_;
   }

backends/vulkan/runtime/graph/Logging.h

@@ -8,6 +8,8 @@
 
 #pragma once
 
+#include <executorch/backends/vulkan/runtime/api/Utils.h>
+
 #include <ostream>
 #include <vector>
 
@@ -23,4 +25,8 @@ inline std::ostream& operator<<(std::ostream& os, const std::vector<T>& vec) {
   return os; // Return the ostream to allow chaining
 }
 
+inline std::ostream& operator<<(std::ostream& os, const api::utils::uvec3& v) {
+  return api::utils::operator<<(os, v);
+}
+
 } // namespace vkcompute

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

@@ -0,0 +1,53 @@
+/*
+ * 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}
+
+layout(std430) buffer;
+
+#include "indexing_utils.h"
+
+layout(set = 0, binding = 0, ${IMAGE_FORMAT[DTYPE]}) uniform PRECISION restrict writeonly ${IMAGE_T[NDIM][DTYPE]} image_out;
+layout(set = 0, binding = 1) uniform PRECISION sampler3D image_in;
+
+layout(set = 0, binding = 2) uniform PRECISION restrict OutSizes {
+  uvec4 data;
+}
+out_sizes;
+
+layout(set = 0, binding = 3) uniform PRECISION restrict SelectVal {
+  // data.x: index along batch dim to select
+  // data.y: number of batches
+  // data.z: number of texels per batch
+  // data.w: unused
+  ivec4 data;
+}
+select_info;
+
+layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
+
+void main() {
+  const int num_batches = select_info.data.y;
+  const int num_texel_per_batch = select_info.data.z;
+  const int index = select_info.data.x;
+
+  const ivec3 pos = ivec3(gl_GlobalInvocationID);
+
+  const ivec4 idx = to_tensor_idx_C_packed(pos, out_sizes.data);
+
+  if (any(greaterThanEqual(idx, out_sizes.data))) {
+    return;
+  }
+
+  const uint src_pos_z = (num_texel_per_batch * index) + pos.z;
+  imageStore(
+      image_out, pos, texelFetch(image_in, ivec3(pos.x, pos.y, src_pos_z), 0));
+}
+

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

@@ -0,0 +1,10 @@
+select_batch_4d:
+  parameter_names_with_default_values:
+    DTYPE: float
+    NDIM: 3
+  generate_variant_forall:
+    DTYPE:
+      - VALUE: half
+      - VALUE: float
+  shader_variants:
+    - NAME: select_batch_4d

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

@@ -0,0 +1,50 @@
+/*
+ * 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 ${texel_component_type(DTYPE)}
+
+layout(std430) buffer;
+
+#include "indexing_utils.h"
+
+layout(set = 0, binding = 0, ${IMAGE_FORMAT[DTYPE]}) uniform PRECISION restrict writeonly ${IMAGE_T[NDIM][DTYPE]} image_out;
+layout(set = 0, binding = 1) uniform PRECISION sampler3D image_in;
+
+layout(set = 0, binding = 2) uniform PRECISION restrict OutSizes {
+  uvec4 data;
+}
+out_sizes;
+
+// index to select
+layout(set = 0, binding = 3) uniform PRECISION restrict IndexVal {
+  int data;
+}
+index;
+
+layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
+
+void main() {
+  const ivec3 pos = ivec3(gl_GlobalInvocationID);
+
+  const ivec4 idx = to_tensor_idx_C_packed(pos, out_sizes.data);
+
+  if (any(greaterThanEqual(idx, out_sizes.data))) {
+    return;
+  }
+
+  const int tex = index.data / 4;
+  const int ind = index.data % 4;
+  const T v = VEC4_T(texelFetch(image_in, ivec3(pos.x, pos.y, tex), 0))[ind];
+
+  imageStore(image_out, ivec3(pos.x, pos.y, 0), VEC4_T(v, 0, 0, 0));
+}

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

@@ -0,0 +1,10 @@
+select_channel_3d:
+  parameter_names_with_default_values:
+    DTYPE: float
+    NDIM: 3
+  generate_variant_forall:
+    DTYPE:
+      - VALUE: half
+      - VALUE: float
+  shader_variants:
+    - NAME: select_channel_3d

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

@@ -0,0 +1,64 @@
+/*
+ * 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)}
+
+layout(std430) buffer;
+
+#include "indexing_utils.h"
+
+layout(set = 0, binding = 0, ${IMAGE_FORMAT[DTYPE]}) uniform PRECISION restrict writeonly ${IMAGE_T[NDIM][DTYPE]} image_out;
+layout(set = 0, binding = 1) uniform PRECISION sampler3D image_in;
+
+layout(set = 0, binding = 2) uniform PRECISION restrict OutSizes {
+  uvec4 data;
+}
+out_sizes;
+
+layout(set = 0, binding = 3) uniform PRECISION restrict SelectVal {
+  // data.x: index along channel dim to select
+  // data.y: number of batches
+  // data.z: number of texels per batch
+  // data.w: unused
+  ivec4 data;
+}
+select_info;
+
+layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
+
+void main() {
+  const ivec3 pos = ivec3(gl_GlobalInvocationID);
+  const ivec4 idx = to_tensor_idx_C_packed(pos, out_sizes.data);
+
+  if (any(greaterThanEqual(idx, out_sizes.data))) {
+    return;
+  }
+
+  const int num_batches = select_info.data.y;
+  const int num_texel_per_batch = select_info.data.z;
+  const int index = select_info.data.x;
+
+  // read in the same channel from 4 separate batches
+  VEC4_T out_texel = VEC4_T(0, 0, 0, 0);
+  for (int k = 0; k < 4; k++) {
+    if ((k + pos.z * 4) >=
+        num_batches) { 
+      break;
+    }
+    const uint src_pos_z = (4 * num_texel_per_batch * pos.z) +
+        (k * num_texel_per_batch) + (index / 4);
+    const uint src_pos_t = index % 4;
+    out_texel[k] =
+        VEC4_T(texelFetch(image_in, ivec3(pos.x, pos.y, src_pos_z), 0))[src_pos_t];
+  }
+
+  imageStore(image_out, pos, out_texel);
+}

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

@@ -0,0 +1,10 @@
+select_channel_4d:
+  parameter_names_with_default_values:
+    DTYPE: float
+    NDIM: 3
+  generate_variant_forall:
+    DTYPE:
+      - VALUE: half
+      - VALUE: float
+  shader_variants:
+    - NAME: select_channel_4d

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

@@ -0,0 +1,62 @@
+/*
+ * 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)}
+
+layout(std430) buffer;
+
+#include "indexing_utils.h"
+
+layout(set = 0, binding = 0, ${IMAGE_FORMAT[DTYPE]}) uniform PRECISION restrict writeonly ${IMAGE_T[NDIM][DTYPE]} image_out;
+layout(set = 0, binding = 1) uniform PRECISION sampler3D image_in;
+
+layout(set = 0, binding = 2) uniform PRECISION restrict OutSizes {
+  uvec4 data;
+}
+out_sizes;
+
+// index to select
+layout(set = 0, binding = 3) uniform PRECISION restrict IndexVal {
+  int data;
+}
+index;
+
+layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
+
+void main() {
+  const ivec3 pos = ivec3(gl_GlobalInvocationID);
+
+  const ivec4 idx = to_tensor_idx_C_packed(pos, out_sizes.data);
+
+  if (any(greaterThanEqual(idx, out_sizes.data))) {
+    return;
+  }
+
+  // w
+  const int src_x = pos.x;
+  // h
+  const int src_y = index.data;
+  // c
+  const int src_z = pos.y;
+
+  const VEC4_T v = VEC4_T(texelFetch(image_in, ivec3(src_x, src_y, src_z), 0));
+
+  for (int i = 0; i < 4; i++) {
+    ivec3 new_pos = ivec3(pos.x, pos.y * 4 + i, 0);
+
+    // When the C-channel exceeds original block size, exit early
+    if (new_pos.y >= out_sizes.data.y) {
+      return;
+    }
+
+    imageStore(image_out, new_pos, VEC4_T(v[i], 0, 0, 0));
+  }
+}

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

@@ -0,0 +1,10 @@
+select_height_3d:
+  parameter_names_with_default_values:
+    DTYPE: float
+    NDIM: 3
+  generate_variant_forall:
+    DTYPE:
+      - VALUE: half
+      - VALUE: float
+  shader_variants:
+    - NAME: select_height_3d

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

@@ -0,0 +1,62 @@
+/*
+ * 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)}
+
+layout(std430) buffer;
+
+#include "indexing_utils.h"
+
+layout(set = 0, binding = 0, ${IMAGE_FORMAT[DTYPE]}) uniform PRECISION restrict writeonly ${IMAGE_T[NDIM][DTYPE]} image_out;
+layout(set = 0, binding = 1) uniform PRECISION sampler3D image_in;
+
+layout(set = 0, binding = 2) uniform PRECISION restrict OutSizes {
+  uvec4 data;
+}
+out_sizes;
+
+// index to select
+layout(set = 0, binding = 3) uniform PRECISION restrict IndexVal {
+  // data.x: index along height dim to select
+  // data.y: number of batches
+  // data.z: number of texels per batch
+  // data.w: unused
+  ivec4 data;
+}
+select_info;
+
+layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
+
+void main() {
+  const ivec3 pos = ivec3(gl_GlobalInvocationID);
+  const ivec4 idx = to_tensor_idx_C_packed(pos, out_sizes.data);
+  if (any(greaterThanEqual(idx, out_sizes.data))) {
+    return;
+  }
+
+  const int num_batches = select_info.data.y;
+  const int num_texel_per_batch = select_info.data.z;
+  const int index = select_info.data.x;
+
+  VEC4_T out_texel = VEC4_T(0, 0, 0, 0);
+  // read in the same channel from 4 separate batches
+  for (int k = 0; k < 4; k++) {
+    if ((k + pos.z * 4) >= num_batches
+        ) { // < 4 batches for this texel, exit early
+      break;
+    }
+    const uint src_pos_z = (pos.z * num_texel_per_batch * 4) +
+        k * num_texel_per_batch + (pos.y / 4);
+    out_texel[k] = VEC4_T(texelFetch(
+        image_in, ivec3(pos.x, index, src_pos_z), 0))[pos.y % 4];
+  }
+  imageStore(image_out, pos, out_texel);
+}

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

@@ -0,0 +1,10 @@
+select_height_4d:
+  parameter_names_with_default_values:
+    DTYPE: float
+    NDIM: 3
+  generate_variant_forall:
+    DTYPE:
+      - VALUE: half
+      - VALUE: float
+  shader_variants:
+    - NAME: select_height_4d