[PyTorch] Add Vulkan support and tests for at::select.int operator, 4 dim/rank tensor case (#96228)

liuk22 · facebook-github-bot · commit 22fd2a753db6 · 2023-03-08T11:30:42.000-08:00
Summary: Pull Request resolved: #96228 Currently, selection along a dimension/rank is only supported for 3D/rank tensors in PyTorch Vulkan. This adds support for 4D/rank tensors at selection along batch, channel (depth), height, and width. Additionally: - The existing implementations have been name-refactored to reflect whether they operate on 3d or 4d tensors. - The params buffer for all select operations now use `ivec2` or `ivec4` only, for memory alignment safety. Test Plan: **Internal:** 1. `buck run --target-platforms ovr_config//platform/macos:arm64-fbsource //xplat/caffe2:pt_vulkan_api_test_binAppleMac\#macosx-arm64 -c pt.vulkan_full_precision=1` on Apple M1 MacBook 2. Confirm all tests pass with no regression, and the directly affected tests `select_4d_*`, refactored `select_3d_`, pass 3. Test output P636928908, in particular: ``` [...bunch of other tests...] [ RUN ] VulkanAPITest.select_3d_depth_small [ OK ] VulkanAPITest.select_3d_depth_small (1 ms) [ RUN ] VulkanAPITest.select_3d_depth_medium [ OK ] VulkanAPITest.select_3d_depth_medium (0 ms) [ RUN ] VulkanAPITest.select_3d_depth_large [ OK ] VulkanAPITest.select_3d_depth_large (1 ms) [ RUN ] VulkanAPITest.select_3d_height_small [ OK ] VulkanAPITest.select_3d_height_small (0 ms) [ RUN ] VulkanAPITest.select_3d_height_medium [ OK ] VulkanAPITest.select_3d_height_medium (0 ms) [ RUN ] VulkanAPITest.select_3d_height_medium1 [ OK ] VulkanAPITest.select_3d_height_medium1 (0 ms) [ RUN ] VulkanAPITest.select_3d_height_medium2 [ OK ] VulkanAPITest.select_3d_height_medium2 (0 ms) [ RUN ] VulkanAPITest.select_3d_height_large [ OK ] VulkanAPITest.select_3d_height_large (1 ms) [ RUN ] VulkanAPITest.select_3d_width_small [ OK ] VulkanAPITest.select_3d_width_small (0 ms) [ RUN ] VulkanAPITest.select_3d_width_medium [ OK ] VulkanAPITest.select_3d_width_medium (0 ms) [ RUN ] VulkanAPITest.select_3d_width_medium2 [ OK ] VulkanAPITest.select_3d_width_medium2 (0 ms) [ RUN ] VulkanAPITest.select_3d_width_large [ OK ] VulkanAPITest.select_3d_width_large (1 ms) [ RUN ] VulkanAPITest.select_4d_batch_small [ OK ] VulkanAPITest.select_4d_batch_small (0 ms) [ RUN ] VulkanAPITest.select_4d_batch_medium [ OK ] VulkanAPITest.select_4d_batch_medium (0 ms) [ RUN ] VulkanAPITest.select_4d_batch_large [ OK ] VulkanAPITest.select_4d_batch_large (1 ms) [ RUN ] VulkanAPITest.select_4d_depth_small [ OK ] VulkanAPITest.select_4d_depth_small (1 ms) [ RUN ] VulkanAPITest.select_4d_depth_medium [ OK ] VulkanAPITest.select_4d_depth_medium (0 ms) [ RUN ] VulkanAPITest.select_4d_depth_large [ OK ] VulkanAPITest.select_4d_depth_large (1 ms) [ RUN ] VulkanAPITest.select_4d_height_small [ OK ] VulkanAPITest.select_4d_height_small (0 ms) [ RUN ] VulkanAPITest.select_4d_height_medium [ OK ] VulkanAPITest.select_4d_height_medium (0 ms) [ RUN ] VulkanAPITest.select_4d_height_large [ OK ] VulkanAPITest.select_4d_height_large (1 ms) [ RUN ] VulkanAPITest.select_4d_width_small [ OK ] VulkanAPITest.select_4d_width_small (0 ms) [ RUN ] VulkanAPITest.select_4d_width_medium [ OK ] VulkanAPITest.select_4d_width_medium (0 ms) [ RUN ] VulkanAPITest.select_4d_width_large [ OK ] VulkanAPITest.select_4d_width_large (1 ms) [...bunch of other tests...] [ FAILED ] 7 tests, listed below: [ FAILED ] VulkanAPITest.cat_dim1_singledepth_success [ FAILED ] VulkanAPITest.gru_success [ FAILED ] VulkanAPITest.gru_mclareninputs_success [ FAILED ] VulkanAPITest.gru_prepack_success [ FAILED ] VulkanAPITest.lstm_success [ FAILED ] VulkanAPITest.lstm_mclareninputs_success [ FAILED ] VulkanAPITest.lstm_prepack_success ``` Reviewed By: SS-JIA Differential Revision: D42623181 fbshipit-source-id: 5b42fe7f2ceb3d4d3dddd7a7389ccc343320da7d
diff --git a/aten/src/ATen/native/vulkan/glsl/select_batch_4d.glsl b/aten/src/ATen/native/vulkan/glsl/select_batch_4d.glsl
@@ -0,0 +1,39 @@
+#version 450 core
+#define PRECISION $precision
+#define FORMAT $format
+
+layout(std430) buffer;
+
+/* Qualifiers: layout - storage - precision - memory */
+
+/*
+ * Output Image
+ */
+layout(set = 0, binding = 0, FORMAT) uniform PRECISION restrict writeonly image3D uOutput;
+
+/*
+ * Input Buffer
+ */
+layout(set = 0, binding = 1) uniform PRECISION sampler3D uInput;
+
+/*
+ * Params Buffer
+ */
+layout(set = 0, binding = 2) uniform PRECISION restrict Block {
+  // batch_info.x: number of texels per batch
+  // batch_info.y: index along batch dim to select
+  ivec2 batch_info;
+}
+uBlock;
+
+/*
+ * Local Work Group Size
+ */
+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 uint src_pos_z = (uBlock.batch_info.y * uBlock.batch_info.x) + pos.z;
+  imageStore(
+      uOutput, pos, texelFetch(uInput, ivec3(pos.x, pos.y, src_pos_z), 0));
+}
diff --git a/aten/src/ATen/native/vulkan/glsl/select_depth.glsl b/aten/src/ATen/native/vulkan/glsl/select_depth.glsl
diff --git a/aten/src/ATen/native/vulkan/glsl/select_depth_3d.glsl b/aten/src/ATen/native/vulkan/glsl/select_depth_3d.glsl
@@ -0,0 +1,32 @@
+#version 450 core
+#define PRECISION $precision
+#define FORMAT $format
+
+layout(std430) buffer;
+
+/* Qualifiers: layout - storage - precision - memory */
+
+layout(set = 0, binding = 0, FORMAT) uniform PRECISION restrict writeonly image3D uOutput;
+layout(set = 0, binding = 1) uniform PRECISION sampler3D uInput;
+layout(set = 0, binding = 2) uniform PRECISION restrict Block {
+  // depth_info.x: output texture x extent
+  // depth_info.y: output texture y extent
+  // depth_info.z: output texture z extent
+  // depth_info.w: output texture w extent
+  ivec4 depth_info;
+}
+uBlock;
+
+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);
+
+  if (all(lessThan(pos, uBlock.depth_info.xyz))) {
+    const int tex = uBlock.depth_info.w / 4;
+    const int ind = uBlock.depth_info.w % 4;
+    const float v = texelFetch(uInput, ivec3(pos.x, pos.y, tex), 0)[ind];
+
+    imageStore(uOutput, ivec3(pos.x, pos.y, 0), vec4(v, 0, 0, 0));
+  }
+}
diff --git a/aten/src/ATen/native/vulkan/glsl/select_depth_4d.glsl b/aten/src/ATen/native/vulkan/glsl/select_depth_4d.glsl
@@ -0,0 +1,53 @@
+#version 450 core
+#define PRECISION $precision
+#define FORMAT $format
+
+layout(std430) buffer;
+
+/* Qualifiers: layout - storage - precision - memory */
+
+/*
+ * Output Image
+ */
+layout(set = 0, binding = 0, FORMAT) uniform PRECISION restrict writeonly image3D uOutput;
+
+/*
+ * Input Buffer
+ */
+layout(set = 0, binding = 1) uniform PRECISION sampler3D uInput;
+
+/*
+ * Params Buffer
+ */
+layout(set = 0, binding = 2) uniform PRECISION restrict Block {
+  // depth_info.x: number of batches
+  // depth_info.y: number of texels per batch
+  // depth_info.z: index along channel dim to select
+  // depth_info.w: zero pad for alignment
+  ivec4 depth_info;
+}
+uBlock;
+
+/*
+ * Local Work Group Size
+ */
+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);
+  // read in the same channel from 4 separate batches
+  vec4 out_texel = vec4(0, 0, 0, 0);
+  for (int k = 0; k < 4; k++) {
+    if ((k + pos.z * 4) >=
+        uBlock.depth_info.x) { // < 4 batches for this texel, exit early
+      break;
+    }
+    const uint src_pos_z = (4 * uBlock.depth_info.y * pos.z) +
+        (k * uBlock.depth_info.y) + (uBlock.depth_info.z / 4);
+    const uint src_pos_t = uBlock.depth_info.z % 4;
+    out_texel[k] =
+        texelFetch(uInput, ivec3(pos.x, pos.y, src_pos_z), 0)[src_pos_t];
+  }
+
+  imageStore(uOutput, pos, out_texel);
+}
diff --git a/aten/src/ATen/native/vulkan/glsl/select_height_3d.glsl b/aten/src/ATen/native/vulkan/glsl/select_height_3d.glsl
@@ -9,9 +9,13 @@ layout(std430) buffer;
 layout(set = 0, binding = 0, FORMAT) uniform PRECISION restrict writeonly image3D uOutput;
 layout(set = 0, binding = 1) uniform PRECISION sampler3D uInput;
 layout(set = 0, binding = 2) uniform PRECISION restrict Block {
-  ivec3 size;
-  int index;
-} uBlock;
+  // height_info.x: output texture x extent
+  // height_info.y: output texture y extent
+  // height_info.z: output texture z extent
+  // height_info.w: output texture w extent
+  ivec4 height_info;
+}
+uBlock;
 
 layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 
@@ -21,7 +25,7 @@ void main() {
   // w
   const int src_x = pos.x;
   // h
-  const int src_y = uBlock.index;
+  const int src_y = uBlock.height_info.w;
   // c
   const int src_z = pos.y;
 
@@ -31,7 +35,7 @@ void main() {
     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 >= uBlock.size.y) {
+    if (new_pos.y >= uBlock.height_info.y) {
       return;
     }
 
diff --git a/aten/src/ATen/native/vulkan/glsl/select_height_4d.glsl b/aten/src/ATen/native/vulkan/glsl/select_height_4d.glsl
@@ -0,0 +1,51 @@
+#version 450 core
+#define PRECISION $precision
+#define FORMAT $format
+
+layout(std430) buffer;
+
+/* Qualifiers: layout - storage - precision - memory */
+
+/*
+ * Output Image
+ */
+layout(set = 0, binding = 0, FORMAT) uniform PRECISION restrict writeonly image3D uOutput;
+
+/*
+ * Input Buffer
+ */
+layout(set = 0, binding = 1) uniform PRECISION sampler3D uInput;
+
+/*
+ * Params Buffer
+ */
+layout(set = 0, binding = 2) uniform PRECISION restrict Block {
+  // height_info.x: number of batches
+  // height_info.y: number of texels per batch
+  // height_info.z: index along height dim to select
+  // height_info.w: zero pad for alignment
+  ivec4 height_info;
+}
+uBlock;
+
+/*
+ * Local Work Group Size
+ */
+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);
+  vec4 out_texel = vec4(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) >=
+        uBlock.height_info.x) { // < 4 batches for this texel, exit early
+      break;
+    }
+    const uint src_pos_z = (pos.z * uBlock.height_info.y * 4) +
+        k * uBlock.height_info.y + (pos.y / 4);
+    out_texel[k] = texelFetch(
+        uInput, ivec3(pos.x, uBlock.height_info.z, src_pos_z), 0)[pos.y % 4];
+  }
+  imageStore(uOutput, pos, out_texel);
+}
diff --git a/aten/src/ATen/native/vulkan/glsl/select_width_3d.glsl b/aten/src/ATen/native/vulkan/glsl/select_width_3d.glsl
@@ -9,17 +9,21 @@ layout(std430) buffer;
 layout(set = 0, binding = 0, FORMAT) uniform PRECISION restrict writeonly image3D uOutput;
 layout(set = 0, binding = 1) uniform PRECISION sampler3D uInput;
 layout(set = 0, binding = 2) uniform PRECISION restrict Block {
-  ivec3 size;
-  int index;
-} uBlock;
+  // width_info.x: output texture x extent
+  // width_info.y: output texture y extent
+  // width_info.z: output texture z extent
+  // width_info.w: output texture w extent
+  ivec4 width_info;
+}
+uBlock;
 
 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);
 
   // w
-  const int src_x = uBlock.index;
+  const int src_x = uBlock.width_info.w;
   // h
   const int src_y = pos.x;
   // c
@@ -31,7 +35,7 @@ void main() {
     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 >= uBlock.size.y) {
+    if (new_pos.y >= uBlock.width_info.y) {
       return;
     }
 
diff --git a/aten/src/ATen/native/vulkan/glsl/select_width_4d.glsl b/aten/src/ATen/native/vulkan/glsl/select_width_4d.glsl
@@ -0,0 +1,51 @@
+#version 450 core
+#define PRECISION $precision
+#define FORMAT $format
+
+layout(std430) buffer;
+
+/* Qualifiers: layout - storage - precision - memory */
+
+/*
+ * Output Image
+ */
+layout(set = 0, binding = 0, FORMAT) uniform PRECISION restrict writeonly image3D uOutput;
+
+/*
+ * Input Buffer
+ */
+layout(set = 0, binding = 1) uniform PRECISION sampler3D uInput;
+
+/*
+ * Params Buffer
+ */
+layout(set = 0, binding = 2) uniform PRECISION restrict Block {
+  // width_info.x: number of batches
+  // width_info.y: number of texels per batch
+  // width_info.z: index along width dim to select
+  // width_info.w: zero pad for alignment
+  ivec4 width_info;
+}
+uBlock;
+
+/*
+ * Local Work Group Size
+ */
+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);
+  vec4 out_texel = vec4(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) >=
+        uBlock.width_info.x) { // < 4 batches for this texel, exit early
+      break;
+    }
+    const uint src_pos_z = (pos.z * uBlock.width_info.y * 4) +
+        k * uBlock.width_info.y + (pos.y / 4);
+    out_texel[k] = texelFetch(
+        uInput, ivec3(uBlock.width_info.z, pos.x, src_pos_z), 0)[pos.y % 4];
+  }
+  imageStore(uOutput, pos, out_texel);
+}
diff --git a/aten/src/ATen/native/vulkan/ops/Select.cpp b/aten/src/ATen/native/vulkan/ops/Select.cpp
diff --git a/aten/src/ATen/test/vulkan_api_test.cpp b/aten/src/ATen/test/vulkan_api_test.cpp
