[pytorch-vulkan] conv1d, only handle special case (#112880)

yipjustin · facebook-github-bot · commit 6551fe6b485b · 2023-11-05T18:38:35.000-08:00
Summary: Just enough to cover the requirement for our target use-case. Will add complete implementation later. bypass-github-pytorch-ci-checks periodic / linux-focal-cuda11.8-py3.9-gcc9 / test is unrelated to this diff. Test Plan: ``` (base) yipjustin@yipjustin-mac fbsource % buck2 run --target-platforms ovr_config//platform/macos:arm64-fbsource //xplat/caffe2:pt_vulkan_api_test_binAppleMac\#macosx-arm64 -- --gtest_filter="*conv1d*" File changed: fbsource//xplat/caffe2/aten/src/ATen/test/vulkan_api_test.cpp Buck UI: https://www.internalfb.com/buck2/27291bfe-940a-4bed-9616-8f3b4f2a3fc7 Network: Up: 20MiB Down: 142B (reSessionID-5632e058-9f48-40eb-8157-30e2db104272) Jobs completed: 6. Time elapsed: 13.5s. Cache hits: 0%. Commands: 2 (cached: 0, remote: 0, local: 2) BUILD SUCCEEDED Running main() from third-party/googletest/1.11.0/googletest/googletest/src/gtest_main.cc Note: Google Test filter = *conv1d* [==========] Running 2 tests from 1 test suite. [----------] Global test environment set-up. [----------] 2 tests from VulkanAPITest [ RUN ] VulkanAPITest.conv1d_simple [ OK ] VulkanAPITest.conv1d_simple (37 ms) [ RUN ] VulkanAPITest.conv1d [ OK ] VulkanAPITest.conv1d (2 ms) [----------] 2 tests from VulkanAPITest (39 ms total) [----------] Global test environment tear-down [==========] 2 tests from 1 test suite ran. (39 ms total) [ PASSED ] 2 tests. ``` Reviewed By: manuelcandales Differential Revision: D50914117
diff --git a/aten/src/ATen/native/vulkan/glsl/conv1d.glsl b/aten/src/ATen/native/vulkan/glsl/conv1d.glsl
@@ -0,0 +1,85 @@
+#version 450 core
+#define PRECISION $precision
+#define FORMAT $format
+
+layout(std430) buffer;
+
+/*
+ * Output Image
+ */
+layout(set = 0, binding = 0, FORMAT) uniform PRECISION restrict writeonly image3D uOutput;
+
+/*
+ * Input Textures
+ */
+layout(set = 0, binding = 1) uniform PRECISION sampler3D uInput;
+layout(set = 0, binding = 2) uniform PRECISION sampler3D uKernel;
+layout(set = 0, binding = 3) uniform PRECISION sampler3D uBias;
+
+layout(set = 0, binding = 4) uniform PRECISION restrict Block {
+  int out_channels;
+  int in_lengths;
+  int kernel_size;
+}
+uBlock;
+
+// This implementation optimize for simplicity (and partially performance) for a
+// (1, C, L) where C == groups. Hence we only focus on calculating the rolling
+// kernel of the L dimension.
+void main() {
+  const ivec3 pos = ivec3(gl_GlobalInvocationID);
+
+  const int out_channels = uBlock.out_channels;
+  const int in_lengths = uBlock.in_lengths;
+  const int kernel_size = uBlock.kernel_size;
+
+  // The global workgroup should have taken care of it. We only perform one
+  // work item for each 1d tensor on lengths
+  if (pos.x >= 1) {
+    return;
+  }
+
+  int c = pos.y;
+  if (c >= out_channels) {
+    return;
+  }
+
+  // Assume n = 1, do not handle n > 1 case for now.
+  int n = pos.z;
+  if (n >= 1) {
+    return;
+  }
+
+  vec4 bias = texelFetch(uBias, ivec3(c, 0, 0), 0);
+
+  for (int i = 0; i < in_lengths - kernel_size + 1; i++) {
+    vec4 v = vec4(0,0,0,0);
+    for (int k = 0; k < kernel_size; k++) {
+      const ivec3 in_pos = ivec3(i+k, c, 0);
+      const vec4 input_value = texelFetch(uInput, in_pos, 0);
+
+      // Note that we are reading weight in the inner loop, this could be
+      // improved by moving it before the outer loop. Since the weight vector is
+      // contant for the entire call.
+
+      // weight in input-space: (c, 0, k);
+      // notice that c is 4-packed. We need to mod 4 to get the actual weight.
+      const ivec3 w_pos = ivec3(k, 0, c / 4);
+      const vec4 weight = texelFetch(uKernel, w_pos, 0);
+
+      float w = weight.x;
+      if (c % 4 == 1) {
+        w = weight.y;
+      } else if (c % 4 == 2) {
+        w = weight.z;
+      } else if (c % 4 == 3) {
+        w = weight.w;
+      }
+
+      v += w * input_value.x;
+    }
+
+    ivec3 out_pos = ivec3(i, c, 0);
+    imageStore(uOutput, out_pos, vec4(v.x + bias.x, 0, 0, 0));
+  }
+}
diff --git a/aten/src/ATen/native/vulkan/ops/Convolution.cpp b/aten/src/ATen/native/vulkan/ops/Convolution.cpp
@@ -835,6 +835,120 @@ Tensor quantized_convolution(
 
 } // namespace
 
+namespace conv1d {
+
+// This implementation only cover a special case. It only supports
+// input = (n=1, input_channel, lengths)
+// ouput = (n=1, output_channel, lengths - kernel_size + 1)
+// stride=1, padding=0, dilation=1, groups=input_channels=output_channels
+//
+// Hence:
+// weight's shape should be (output_channel, 1, kernel_size)
+// bias's shape (if applicable) should be (output_channel,)
+//
+// In this implementation, it reduces to running a 1d convolution for reach
+// channel.
+// There are multiple perf improvement opportunities: e.g. width-packing
+// input and weight tensors, batch reading when groups is low.
+
+Tensor conv1d(
+    const Tensor& input_arg,
+    const Tensor& weight_arg,
+    const c10::optional<Tensor>& bias_arg_opt,
+    IntArrayRef stride,
+    IntArrayRef padding,
+    IntArrayRef dilation,
+    int64_t groups) {
+  api::Context* const context = api::context();
+  const Tensor input = input_arg.is_vulkan() ? input_arg : input_arg.vulkan();
+  const Tensor weight =
+      weight_arg.is_vulkan() ? weight_arg : weight_arg.vulkan();
+
+  const IntArrayRef& input_sizes = input.sizes();
+  const IntArrayRef& weight_sizes = weight.sizes();
+
+  int32_t weight_out_channels = static_cast<int32_t>(weight_sizes[0]);
+  int32_t kernel_size = static_cast<int32_t>(weight_sizes[2]);
+
+  Tensor bias;
+  if (bias_arg_opt) {
+    if (bias_arg_opt->is_vulkan()) {
+      bias = bias_arg_opt.value();
+    } else {
+      bias = bias_arg_opt.value().vulkan();
+    }
+  } else {
+    bias = at::zeros({weight_out_channels}).vulkan();
+  }
+
+  TORCH_CHECK(input.dim() == 3, "input must be a 3-dim tensor");
+  TORCH_CHECK(weight.dim() == 3, "weight must be a 3-dim tensor");
+  TORCH_CHECK(stride == IntArrayRef(1), "stride must be 1");
+  TORCH_CHECK(padding == IntArrayRef(0), "padding must be 1");
+  TORCH_CHECK(dilation == IntArrayRef(1), "dilation must be 1");
+
+  TORCH_CHECK(input_sizes[0] == 1, "Only support single batch");
+  TORCH_CHECK(input_sizes[1] == groups, "input_channel must equals to groups");
+  TORCH_CHECK(
+      weight_sizes[0] == groups, "weight.sizes()[0] must equals to groups");
+  TORCH_CHECK(weight_sizes[1] == 1, "weight.sizes()[1] must equals to 1");
+
+  const vTensor& v_input = convert(input);
+  const IntArrayRef v_input_sizes = v_input.sizes();
+
+  const vTensor& v_weight = convert(weight);
+  const vTensor& v_bias = convert(bias);
+
+  vTensor v_output{
+      context,
+      {
+          v_input_sizes[0],
+          weight_out_channels,
+          v_input_sizes[2] - kernel_size + 1,
+      },
+      input_arg.scalar_type(),
+  };
+
+  const struct Block final {
+    int32_t out_channels;
+    int32_t in_lengths;
+    int32_t kernel_size;
+  } block{
+      weight_out_channels,
+      static_cast<int32_t>(input_sizes[2]),
+      kernel_size,
+  };
+
+  api::UniformParamsBuffer params(context, block);
+  api::PipelineBarrier pipeline_barrier{};
+
+  context->submit_compute_job(
+      // shader descriptor
+      VK_KERNEL(conv1d),
+      // pipeline barrier
+      pipeline_barrier,
+      // global work group size
+      {1, static_cast<uint32_t>(weight_out_channels), 1},
+      // local work group size
+      {1, 1, 1},
+      // 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),
+      v_weight.image(pipeline_barrier, api::PipelineStage::COMPUTE),
+      v_bias.image(pipeline_barrier, api::PipelineStage::COMPUTE),
+      // params buffer
+      params.buffer());
+
+  return convert(v_output);
+}
+
+} // namespace conv1d
+
 Conv2dPackedContext::Conv2dPackedContext(
     const Tensor& weight,
     const c10::optional<Tensor>& bias,
@@ -1252,6 +1366,7 @@ Tensor conv2d_clamp_run(
 
 TORCH_LIBRARY_IMPL(aten, Vulkan, m) {
   m.impl("convolution_overrideable", convolution);
+  m.impl(TORCH_SELECTIVE_NAME("aten::conv1d"), TORCH_FN(conv1d::conv1d));
 }
 
 } // namespace ops
diff --git a/aten/src/ATen/test/vulkan_api_test.cpp b/aten/src/ATen/test/vulkan_api_test.cpp
@@ -7,6 +7,7 @@
 #include <ATen/native/vulkan/ops/Copy.h>
 #include <ATen/native/vulkan/ops/Convolution.h>
 #include <c10/util/irange.h>
+#include <c10/util/ArrayRef.h>
 
 // TODO: These functions should move to a common place.
 
@@ -1179,6 +1180,93 @@ TEST_F(VulkanAPITest, clamp_) {
   ASSERT_TRUE(check);
 }
 
+TEST_F(VulkanAPITest, conv1d_simple) {
+  // This is a simple case using arange for input, ones for weights, and arange
+  // for bias. This makes debugging easiser.
+  int64_t kernel_size = 3;
+  int64_t channels = 5;
+  int64_t lengths = 9;
+
+  c10::InferenceMode mode;
+
+  const auto input_cpu = at::arange(lengths * channels, at::kFloat).reshape({1, channels, lengths});
+  const auto weights_cpu = at::ones({channels, 1, kernel_size}, at::device(at::kCPU).dtype(at::kFloat));
+  const auto bias_cpu = at::arange(channels, at::kFloat);
+
+  const auto input_vk = input_cpu.vulkan();
+  const auto weights_vk = weights_cpu.vulkan();
+  const auto bias_vk = bias_cpu.vulkan();
+
+  std::array<int64_t, 1> stride{1};
+  std::array<int64_t, 1> padding{0};
+  std::array<int64_t, 1> dilation{1};
+
+  const auto output_cpu = at::conv1d(
+      input_cpu, weights_cpu, bias_cpu,
+      stride, padding, dilation, channels);
+
+  const auto output_vk = at::conv1d(
+      input_vk, weights_vk, bias_vk,
+      stride,
+      padding,
+      dilation,
+      channels);
+  const auto output_vk_cpu = output_vk.cpu();
+
+  const bool check = almostEqual(output_cpu, output_vk_cpu);
+  if (!check) {
+    showRtol(output_cpu, output_vk_cpu);
+  }
+
+  ASSERT_TRUE(check);
+}
+
+void test_conv1d(int64_t kernel_size, int64_t channels, int64_t lengths) {
+  c10::InferenceMode mode;
+
+  const auto input_cpu = at::rand({1, channels, lengths}, at::kFloat);
+  const auto weights_cpu = at::rand({channels, 1, kernel_size}, at::kFloat);
+  const auto bias_cpu = at::rand({channels,}, at::kFloat);
+
+  const auto input_vk = input_cpu.vulkan();
+  const auto weights_vk = weights_cpu.vulkan();
+  const auto bias_vk = bias_cpu.vulkan();
+
+  std::array<int64_t, 1> stride{1};
+  std::array<int64_t, 1> padding{0};
+  std::array<int64_t, 1> dilation{1};
+  int64_t groups = channels;
+
+  const auto output_cpu = at::conv1d(
+      input_cpu, weights_cpu, bias_cpu,
+      stride, padding, dilation, groups);
+
+  const auto output_vk = at::conv1d(
+      input_vk, weights_vk, bias_vk,
+      stride,
+      padding,
+      dilation,
+      channels);
+  const auto output_vk_cpu = output_vk.cpu();
+
+  const bool check = almostEqual(output_cpu, output_vk_cpu);
+  if (!check) {
+    showRtol(output_cpu, output_vk_cpu);
+  }
+
+  ASSERT_TRUE(check);
+}
+
+TEST_F(VulkanAPITest, conv1d) {
+  test_conv1d(3, 5, 8);
+  test_conv1d(9, 5, 9);
+  test_conv1d(1, 12, 3);
+  test_conv1d(1, 12, 1);
+  test_conv1d(10, 12, 20);
+}
+
+
+
 void test_conv2d_context(
     const at::IntArrayRef input_shape,
     const at::IntArrayRef weight_shape,
