[Pytorch][Vulkan] Call binary_op_scalar when 'other' is a 0-dim tensor (#109035)

lucylq · facebook-github-bot · commit e03bc22fc9ad · 2023-09-11T13:04:22.000-07:00
Summary: 0-dim tensor are not supported in Vulkan. If a binary_op_tensor is called with 'other_arg' being a 0-dim tensor, then we extract the scalar out and call binary_op_scalar. Used to run the [FD model]( https://www.internalfb.com/manifold/explorer/wrist-camera-ml/tree/models/fd-ted-pi/fd-hybrid/fd_hybrid_vulkan.ptl) on [CLI](https://www.internalfb.com/intern/wiki/Malibu/Software/Machine_Learning/PyTorch_On_Device_Catalog/#build-and-run-native-pyt) Test Plan: Run vulkan_api_tests ``` lfq@lfq-mbp fbsource % 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 ... [ RUN ] VulkanAPITest.querypool_flushed_shader_log xplat/caffe2/aten/src/ATen/test/vulkan_api_test.cpp:6891: Skipped QueryPool is not available [ SKIPPED ] VulkanAPITest.querypool_flushed_shader_log (0 ms) [----------] 339 tests from VulkanAPITest (5308 ms total) [----------] Global test environment tear-down [==========] 339 tests from 1 test suite ran. (5308 ms total) [ PASSED ] 338 tests. [ SKIPPED ] 1 test, listed below: [ SKIPPED ] VulkanAPITest.querypool_flushed_shader_log YOU HAVE 5 DISABLED TESTS ``` Reviewed By: manuelcandales Differential Revision: D48672535
diff --git a/aten/src/ATen/native/vulkan/api/Tensor.cpp b/aten/src/ATen/native/vulkan/api/Tensor.cpp
@@ -169,7 +169,8 @@ c10::SmallVector<int64_t, 6u> calc_gpu_sizes(
   else {
     TORCH_CHECK(
         ndim >= 1 && ndim <= 4,
-        "Texture storage only valid for 1 <= ndim <= 4!");
+        "Texture storage only valid for 1 <= ndim <= 4, received: ",
+        ndim);
 
     c10::SmallVector<int64_t, 6u> gpu_sizes(ndim == 4 ? 4 : 3);
 
diff --git a/aten/src/ATen/native/vulkan/ops/BinaryOp.cpp b/aten/src/ATen/native/vulkan/ops/BinaryOp.cpp
@@ -402,6 +402,13 @@ Tensor add_tensor(
     const Tensor& self_arg,
     const Tensor& other_arg,
     const Scalar& alpha) {
+  if (other_arg.dim() == 0) {
+    return binary_op_scalar(
+        self_arg,
+        other_arg.item(),
+        c10::optional<Scalar>(),
+        VK_KERNEL(add_scalar));
+  }
   return binary_op_tensor(
       self_arg, other_arg, c10::optional<Scalar>(alpha), VK_KERNEL(add));
 }
@@ -437,6 +444,13 @@ Tensor sub_tensor(
     const Tensor& self_arg,
     const Tensor& other_arg,
     const Scalar& alpha) {
+  if (other_arg.dim() == 0) {
+    return binary_op_scalar(
+        self_arg,
+        other_arg.item(),
+        c10::optional<Scalar>(-1 * alpha.to<float>()),
+        VK_KERNEL(add_scalar));
+  }
   return binary_op_tensor(
       self_arg, other_arg, c10::optional<Scalar>(alpha), VK_KERNEL(sub));
 }
@@ -460,6 +474,13 @@ Tensor& mul_scalar_(Tensor& self, const Scalar& other) {
 }
 
 Tensor mul_tensor(const Tensor& self_arg, const Tensor& other_arg) {
+  if (other_arg.dim() == 0) {
+    return binary_op_scalar(
+        self_arg,
+        other_arg.item(),
+        c10::optional<Scalar>(),
+        VK_KERNEL(mul_scalar));
+  }
   return binary_op_tensor(
       self_arg, other_arg, c10::optional<Scalar>(), VK_KERNEL(mul));
 }
@@ -486,6 +507,13 @@ Tensor& div_scalar_(Tensor& self, const Scalar& other) {
 }
 
 Tensor div_tensor(const Tensor& self_arg, const Tensor& other_arg) {
+  if (other_arg.dim() == 0) {
+    return binary_op_scalar(
+        self_arg,
+        1.0 / other_arg.item().to<float>(),
+        c10::optional<Scalar>(),
+        VK_KERNEL(mul_scalar));
+  }
   return binary_op_tensor(
       self_arg, other_arg, c10::optional<Scalar>(), VK_KERNEL(div));
 }
