Revert "[pytorch] Add support for "height" and "width" dimension for …

…the "select" operator on pytorch vulkan backend (pytorch#94612)" This reverts commit f2c2642.
ROCm · May 2, 2023 · 1a686de · 1a686de
1 parent 6e3439d
commit 1a686de
Show file tree

Hide file tree

Showing 4 changed files with 2 additions and 240 deletions.
diff --git a/aten/src/ATen/native/vulkan/glsl/select_height.glsl b/aten/src/ATen/native/vulkan/glsl/select_height.glsl
diff --git a/aten/src/ATen/native/vulkan/glsl/select_width.glsl b/aten/src/ATen/native/vulkan/glsl/select_width.glsl
diff --git a/aten/src/ATen/native/vulkan/ops/Select.cpp b/aten/src/ATen/native/vulkan/ops/Select.cpp
@@ -53,125 +53,9 @@ Tensor select_depth(const Tensor& input_arg, uint32_t index) {
   return convert(v_output);
 }
 
-Tensor select_height(const Tensor& input_arg, uint32_t index) {
-  api::Context* const context = api::context();
-
-  const Tensor input = input_arg.is_vulkan() ? input_arg : input_arg.vulkan();
-  const vTensor& v_input = convert(input);
-  const IntArrayRef v_input_sizes = v_input.sizes();
-
-  vTensor v_output{
-      context,
-      {v_input_sizes[0], v_input_sizes[2]},
-      input_arg.scalar_type(),
-  };
-
-  const struct Block final {
-    uvec3 size; // output texture size
-    uint32_t index;
-  } block{v_output.extents(), index};
-
-  // Input tensor is a (c, h, w)
-  // Output tensor is a (c, w)
-  // In shader, the input texture's coordinate is (w, h, c)
-  // In shader, the output texture's coordinate is (w, c, 1)
-  uint32_t w = v_output.extents().data[0u];
-  uint32_t c = v_output.extents().data[1u];
-
-  // Encoding of c-channel is packed into texel, hence we only call ceil(c/4)
-  // times to minimize invocation and read.
-  // For the last dimension, it is the selected height. Shader will do a direct
-  // lookup based on block.index.
-  uvec3 global_workgroup_size{w, api::utils::div_up(c, 4u), 1};
-
-  api::UniformParamsBuffer params(context, block);
-  api::PipelineBarrier pipeline_barrier{};
-
-  context->submit_compute_job(
-      // shader descriptor
-      VK_KERNEL(select_height),
-      // pipeline barrier
-      pipeline_barrier,
-      // global work group size
-      global_workgroup_size,
-      // local work group size
-      adaptive_work_group_size(global_workgroup_size),
-      // fence handle
-      VK_NULL_HANDLE,
-      // shader arguments
-      v_output.image(
-          pipeline_barrier,
-          api::PipelineStage::COMPUTE,
-          api::MemoryAccessType::WRITE),
-      v_input.image(pipeline_barrier, api::PipelineStage::COMPUTE),
-      // params buffer
-      params.buffer());
-
-  return convert(v_output);
-}
-
-Tensor select_width(const Tensor& input_arg, uint32_t index) {
-  api::Context* const context = api::context();
-
-  const Tensor input = input_arg.is_vulkan() ? input_arg : input_arg.vulkan();
-  const vTensor& v_input = convert(input);
-  const IntArrayRef v_input_sizes = v_input.sizes();
-
-  vTensor v_output{
-      context,
-      {v_input_sizes[0], v_input_sizes[1]},
-      input_arg.scalar_type(),
-  };
-
-  const struct Block final {
-    uvec3 size; // output texture size
-    uint32_t index;
-  } block{v_output.extents(), index};
-
-  // Input tensor is a (c, h, w)
-  // Output tensor is a (c, h)
-  // In shader, the input texture's coordinate is (w, h, c)
-  // In shader, the output texture's coordinate is (h, c, 1)
-  uint32_t h = v_output.extents().data[0u];
-  uint32_t c = v_output.extents().data[1u];
-
-  // Encoding of c-channel is packed into texel, hence we only call ceil(c/4)
-  // times to minimize invocation and read.
-  // For the last dimension, it is the selected width. Shader will do a direct
-  // lookup based on block.index.
-  uvec3 global_workgroup_size{h, api::utils::div_up(c, 4u), 1};
-
-  api::UniformParamsBuffer params(context, block);
-  api::PipelineBarrier pipeline_barrier{};
-
-  context->submit_compute_job(
-      // shader descriptor
-      VK_KERNEL(select_width),
-      // pipeline barrier
-      pipeline_barrier,
-      // global work group size
-      global_workgroup_size,
-      // local work group size
-      adaptive_work_group_size(global_workgroup_size),
-      // fence handle
-      VK_NULL_HANDLE,
-      // shader arguments
-      v_output.image(
-          pipeline_barrier,
-          api::PipelineStage::COMPUTE,
-          api::MemoryAccessType::WRITE),
-      v_input.image(pipeline_barrier, api::PipelineStage::COMPUTE),
-      // params buffer
-      params.buffer());
-
-  return convert(v_output);
-}
-
 Tensor select(const Tensor& self, int64_t dim, int64_t index) {
   TORCH_CHECK(self.dim() == 3, "Vulkan select only supports 3d tensors!");
-  TORCH_CHECK(
-      0 <= dim && dim <= 2,
-      "Vulkan select only supports one of the dim (0, 1, 2)");
+  TORCH_CHECK(dim == 0, "Vulkan select only supports dim = 0!");
 
   const int64_t size = self.size(dim);
 
@@ -189,13 +73,7 @@ Tensor select(const Tensor& self, int64_t dim, int64_t index) {
     index += size;
   }
 
-  if (dim == 0) {
-    return select_depth(self, index);
-  } else if (dim == 1) {
-    return select_height(self, index);
-  } else {
-    return select_width(self, index);
-  }
+  return select_depth(self, index);
 }
 
 #ifdef USE_VULKAN_API

diff --git a/aten/src/ATen/test/vulkan_api_test.cpp b/aten/src/ATen/test/vulkan_api_test.cpp
@@ -2803,42 +2803,6 @@ TEST_F(VulkanAPITest, select_3d_depth_large) {
   test_select({100, 1, 144}, 0, 50);
 }
 
-TEST_F(VulkanAPITest, select_3d_height_small) {
-  test_select({1, 1, 1}, 1, 0);
-}
-
-TEST_F(VulkanAPITest, select_3d_height_medium) {
-  test_select({3, 5, 2}, 1, 2);
-}
-
-TEST_F(VulkanAPITest, select_3d_height_medium1) {
-  test_select({16, 16, 5}, 1, 6);
-}
-
-TEST_F(VulkanAPITest, select_3d_height_medium2) {
-  test_select({17, 17, 5}, 1, 6);
-}
-
-TEST_F(VulkanAPITest, select_3d_height_large) {
-  test_select({100, 144, 5}, 1, 50);
-}
-
-TEST_F(VulkanAPITest, select_3d_width_small) {
-  test_select({1, 1, 1}, 2, 0);
-}
-
-TEST_F(VulkanAPITest, select_3d_width_medium) {
-  test_select({3, 5, 3}, 2, 2);
-}
-
-TEST_F(VulkanAPITest, select_3d_width_medium2) {
-  test_select({17, 17, 8}, 2, 6);
-}
-
-TEST_F(VulkanAPITest, select_3d_width_large) {
-  test_select({100, 3, 144}, 2, 50);
-}
-
 TEST_F(VulkanAPITest, sigmoid) {
   const auto in_cpu = at::rand({17, 197, 302, 5}, at::device(at::kCPU).dtype(at::kFloat));
   const auto in_vulkan = in_cpu.vulkan();