[vulkan] jit passes for vulkan conv2 prepack and fuse with clamp (pytorch#39282)

Summary: Pull Request resolved: pytorch#39282

Test Plan: Imported from OSS

Differential Revision: D21962424

Pulled By: IvanKobzarev

fbshipit-source-id: 2d20e827d2c3836b7e6b443293377c68dc1ffa5a

Author: IvanKobzarev

aten/src/ATen/native/vulkan/VulkanAten.cpp

@@ -179,7 +179,7 @@ at::Tensor vulkan_convolution(
       voutput,
       vinput,
       weight.data_ptr<float>(),
-      bias.defined() ? c10::make_optional<float*>(bias.data_ptr<float>())
+      bias.defined() ? c10::make_optional<const float*>(bias.data_ptr<float>())
                      : c10::nullopt,
       params);
   return new_with_vtensor_vulkan(std::move(voutput), input.options());
@@ -242,7 +242,8 @@ at::Tensor vulkan_convolution_prepacked(
         voutput,
         vinput,
         vweight,
-        hasBias ? c10::make_optional((*bias).data_ptr<float>()) : c10::nullopt,
+        hasBias ? c10::make_optional<const float*>((*bias).data_ptr<float>())
+                : c10::nullopt,
         params,
         output_min,
         output_max);

aten/src/ATen/native/vulkan/VulkanConvolution.cpp

@@ -66,14 +66,13 @@ ContextConv2D create(
   const auto stride_expanded = expand_param_if_needed(stride, "stride", 2);
   const auto dilation_expanded =
       expand_param_if_needed(dilation, "dilation", 2);
-  const Tensor weight_nchw = weight.contiguous();
+  Tensor weight_nchw = weight.contiguous();
+  auto ws = weight_nchw.sizes();
   return ContextConv2D{
-      at::native::vulkan_convolution_prepack_weights(weight),
+      groups == 1 ? at::native::vulkan_convolution_prepack_weights(weight_nchw)
+                  : weight_nchw.vulkan(),
       bias.has_value() ? c10::make_optional((*bias).vulkan()) : c10::nullopt,
-      {weight_nchw.sizes()[0],
-       weight_nchw.sizes()[1],
-       weight_nchw.sizes()[2],
-       weight_nchw.sizes()[3]},
+      {{ws[0], ws[1], ws[2], ws[3]}},
       {padding_expanded[0], padding_expanded[1]},
       {stride_expanded[0], stride_expanded[1]},
       {dilation_expanded[0], dilation_expanded[1]},

aten/src/ATen/native/vulkan/VulkanOps.cpp

@@ -176,7 +176,7 @@ VBuffer kernelNCHW_OCHW_repack_O4C4HWi4o4(
 }
 
 VBuffer bufferFromOptionalHostData(
-    c10::optional<float*> data,
+    c10::optional<const float*> data,
     const uint32_t size) {
   const auto sizeAligned =
       ROUND_UP(size, context().limits().minStorageBufferOffsetAlignment);
@@ -202,17 +202,15 @@ uint32_t conv2d_biasBufferSize(uint32_t oc) {
 void conv2d_depthwise(
     VulkanTensor& output,
     const VulkanTensor& input,
-    const float* weight,
-    const c10::optional<float*> bias,
-    const Conv2DParams params,
+    const VulkanTensor& weight,
+    const VBuffer& biasBuffer,
+    const Conv2DParams& params,
     c10::optional<float> output_min,
     c10::optional<float> output_max) {
   TORCH_INTERNAL_ASSERT(params.G == params.C);
   auto osizes = output.sizes();
   TORCH_INTERNAL_ASSERT(osizes[2] == params.OH);
   TORCH_INTERNAL_ASSERT(osizes[3] == params.OW);
-  auto biasBuffer =
-      bufferFromOptionalHostData(bias, conv2d_biasBufferSize(params.OC));
   struct ConstBlock {
     int32_t padding[2];
     int32_t kernelSize[2];
@@ -234,9 +232,6 @@ void conv2d_depthwise(
       output_max ? *output_max : std::numeric_limits<float>::infinity()};
   VBuffer constBuffer = makeUniformConstBuffer((void*)&cb, sizeof(cb));
 
-  VulkanTensor kernel{{params.OC, params.KH, params.KW}};
-  kernel.set_data_from_host(weight);
-
   VkDescriptorSetLayout descriptorSetLayout{};
   VkDescriptorPool descriptorPool{};
   VkDescriptorSet descriptorSet{};
@@ -256,7 +251,7 @@ void conv2d_depthwise(
 
   output.image()->bindStorageImage(descriptorSet, 0);
   input.image()->bindShaderRead(descriptorSet, 1);
-  kernel.image()->bindShaderRead(descriptorSet, 2);
+  weight.image()->bindShaderRead(descriptorSet, 2);
   biasBuffer.bind(descriptorSet, 3);
   constBuffer.bind(descriptorSet, 4);
 
@@ -269,7 +264,7 @@ void conv2d_depthwise(
   auto commandBuffer = computeUnit.commandBuffer();
   output.image()->addImageMemoryBarrierToGeneral(commandBuffer);
   input.image()->addImageMemoryBarrierToShaderRead(commandBuffer);
-  kernel.image()->addImageMemoryBarrierToShaderRead(commandBuffer);
+  weight.image()->addImageMemoryBarrierToShaderRead(commandBuffer);
   computeUnit.dispatchCommandBuffer(
       params.OW, params.OH, params.OC_4, workGroupSize);
   computeUnit.endCommandBuffer();
@@ -279,6 +274,44 @@ void conv2d_depthwise(
   vkDestroyDescriptorSetLayout(device, descriptorSetLayout, nullptr);
 }
 
+void conv2d_depthwise(
+    VulkanTensor& output,
+    const VulkanTensor& input,
+    const VulkanTensor& weight,
+    const c10::optional<const float*> bias,
+    const Conv2DParams params,
+    c10::optional<float> output_min,
+    c10::optional<float> output_max) {
+  conv2d_depthwise(
+      output,
+      input,
+      weight,
+      bufferFromOptionalHostData(bias, conv2d_biasBufferSize(params.OC)),
+      params,
+      output_min,
+      output_max);
+}
+
+void conv2d_depthwise(
+    VulkanTensor& output,
+    const VulkanTensor& input,
+    const float* weight,
+    const c10::optional<const float*> bias,
+    const Conv2DParams params,
+    c10::optional<float> output_min,
+    c10::optional<float> output_max) {
+  VulkanTensor weightTensor{{params.OC, params.KH, params.KW}};
+  weightTensor.set_data_from_host(weight);
+  conv2d_depthwise(
+      output,
+      input,
+      weightTensor,
+      bufferFromOptionalHostData(bias, conv2d_biasBufferSize(params.OC)),
+      params,
+      output_min,
+      output_max);
+}
+
 ImageSizes conv2d_prepack_weights_image_sizes(
     int64_t OC,
     int64_t C,
@@ -463,7 +496,7 @@ void conv2d(
     VulkanTensor& output,
     const VulkanTensor& input,
     const VImage& kernelImage,
-    const c10::optional<float*> bias,
+    const c10::optional<const float*> bias,
     const Conv2DParams& params,
     c10::optional<float> output_min,
     c10::optional<float> output_max) {
@@ -483,10 +516,22 @@ void conv2d(
     VulkanTensor& output,
     const VulkanTensor& input,
     const VulkanTensor& weight_prepacked,
-    c10::optional<float*> bias,
+    c10::optional<const float*> bias,
     const Conv2DParams params,
     c10::optional<float> output_min,
     c10::optional<float> output_max) {
+  if (params.G > 1) {
+    conv2d_depthwise(
+        output,
+        input,
+        weight_prepacked,
+        bufferFromOptionalHostData(bias, conv2d_biasBufferSize(params.OC)),
+        params,
+        output_min,
+        output_max);
+    return;
+  }
+
   conv2d(
       output,
       input,
@@ -505,6 +550,18 @@ void conv2d(
     const Conv2DParams params,
     c10::optional<float> output_min,
     c10::optional<float> output_max) {
+  if (params.G > 1) {
+    conv2d_depthwise(
+        output,
+        input,
+        weight_prepacked,
+        *(bias.buffer()),
+        params,
+        output_min,
+        output_max);
+    return;
+  }
+
   conv2d(
       output,
       input,
@@ -519,7 +576,7 @@ void conv2d(
     VulkanTensor& output,
     const VulkanTensor& input,
     const float* weight,
-    const c10::optional<float*> bias,
+    const c10::optional<const float*> bias,
     const Conv2DParams params,
     c10::optional<float> output_min,
     c10::optional<float> output_max) {

aten/src/ATen/native/vulkan/VulkanOps.h

@@ -37,7 +37,7 @@ void conv2d(
     VulkanTensor& output,
     const VulkanTensor& input,
     const float* weight,
-    const c10::optional<float*> bias,
+    const c10::optional<const float*> bias,
     const Conv2DParams params,
     c10::optional<float> output_min = c10::nullopt,
     c10::optional<float> output_max = c10::nullopt);
@@ -46,7 +46,7 @@ void conv2d(
     VulkanTensor& output,
     const VulkanTensor& input,
     const VulkanTensor& weight_prepacked,
-    const c10::optional<float*> bias,
+    const c10::optional<const float*> bias,
     const Conv2DParams params,
     c10::optional<float> output_min = c10::nullopt,
     c10::optional<float> output_max = c10::nullopt);

aten/src/ATen/test/vulkan_test.cpp

@@ -496,7 +496,7 @@ TEST(VulkanTest, conv2dPrepack) {
   ASSERT_TRUE(no_prepack_check);
 
   auto prepack = callOpByName(
-      "vulkan::conv2d_clamp_prepack",
+      "vulkan_prepack::conv2d_clamp_prepack",
       "",
       t_w,
       t_b,
@@ -507,7 +507,7 @@ TEST(VulkanTest, conv2dPrepack) {
       output_min,
       output_max);
   auto tv_out_prepack_ivalues =
-      callOpByName("vulkan::conv2d_clamp_run", "", tv_in, prepack[0]);
+      callOpByName("vulkan_prepack::conv2d_clamp_run", "", tv_in, prepack[0]);
   auto tv_out_prepack = tv_out_prepack_ivalues[0].toTensor();
   auto t_out_prepack = tv_out_prepack.cpu();
   const auto prepack_check = almostEqual(t_out_prepack, t_out_expected);

binaries/CMakeLists.txt

@@ -103,3 +103,4 @@ endif()
 caffe2_binary_target("tutorial_blob.cc")
 
 caffe2_binary_target("dump_operator_names.cc")
+caffe2_binary_target("optimize_for_mobile.cc")

binaries/optimize_for_mobile.cc

@@ -17,6 +17,7 @@
 #include <string>
 
 #include "torch/csrc/jit/api/module.h"
+#include "torch/csrc/jit/passes/vulkan_rewrite.h"
 #include "torch/csrc/jit/passes/xnnpack_rewrite.h"
 #include "torch/csrc/jit/serialization/import.h"
 
@@ -29,6 +30,7 @@ C10_DEFINE_bool(
     save_for_mobile,
     false,
     "Save the model with bytecode format compatible with lite inteprter.");
+C10_DEFINE_bool(vulkan, false, "Vulkan optimize_for_mobile");
 
 int main(int argc, char** argv) {
   c10::SetUsageMessage(
@@ -52,7 +54,10 @@ int main(int argc, char** argv) {
   }
 
   auto module = torch::jit::load(FLAGS_model);
-  auto optimized_module = torch::jit::optimizeForMobile(module);
+
+  auto optimized_module = FLAGS_vulkan
+      ? torch::jit::vulkanOptimizeForMobile(module)
+      : torch::jit::optimizeForMobile(module);
 
   if (FLAGS_save_for_mobile) {
     optimized_module._save_for_mobile(output_model_name);

test/run_test.py

@@ -48,6 +48,7 @@
     'test_optim',
     'test_mobile_optimizer',
     'test_xnnpack_integration',
+    'test_vulkan',
     'test_quantization',
     'test_sparse',
     'test_serialization',