[pytorch-vulkan] log, log_softmax

yipjustin · facebook-github-bot · commit 07820ec2503d · 2023-11-02T23:20:26.000-07:00
Summary: tsia Test Plan: ``` [yipjustin@189650.od ~/fbsource (631468db3)]$ LD_LIBRARY_PATH=third-party/swiftshader/lib/linux-x64/ buck2 run fbcode/mode/dev-nosan //xplat/caffe2:pt_vulkan_api_test_bin -- --gtest_filter="*softmax*" File changed: fbsource//xplat/caffe2/aten/src/ATen/test/vulkan_api_test.cpp File changed: fbcode//caffe2/aten/src/ATen/native/vulkan/ops/Softmax.cpp File changed: fbcode//caffe2/aten/src/ATen/test/vulkan_api_test.cpp 1 additional file change events Buck UI: https://www.internalfb.com/buck2/d4f62e52-aba9-448a-a181-cf8881affb14 Network: Up: 0B Down: 0B Jobs completed: 4. Time elapsed: 0.5s. BUILD SUCCEEDED Running main() from third-party/googletest/1.11.0/googletest/googletest/src/gtest_main.cc Note: Google Test filter = *softmax* [==========] Running 2 tests from 1 test suite. [----------] Global test environment set-up. [----------] 2 tests from VulkanAPITest [ RUN ] VulkanAPITest.softmax [ OK ] VulkanAPITest.softmax (467 ms) [ RUN ] VulkanAPITest.log_softmax [ OK ] VulkanAPITest.log_softmax (95 ms) [----------] 2 tests from VulkanAPITest (563 ms total) [----------] Global test environment tear-down [==========] 2 tests from 1 test suite ran. (563 ms total) [ PASSED ] 2 tests. YOU HAVE 1 DISABLED TEST [yipjustin@189650.od ~/fbsource (631468db3)]$ LD_LIBRARY_PATH=third-party/swiftshader/lib/linux-x64/ buck2 run fbcode/mode/dev-nosan //xplat/caffe2:pt_vulkan_api_test_bin -- --gtest_filter="*log*" Buck UI: https://www.internalfb.com/buck2/e8210eb5-fd56-45f7-bf6c-5024931e778e Network: Up: 0B Down: 0B Jobs completed: 4. Time elapsed: 0.2s. BUILD SUCCEEDED Running main() from third-party/googletest/1.11.0/googletest/googletest/src/gtest_main.cc Note: Google Test filter = *log* [==========] Running 4 tests from 1 test suite. [----------] Global test environment set-up. [----------] 4 tests from VulkanAPITest [ RUN ] VulkanAPITest.log_softmax [ OK ] VulkanAPITest.log_softmax (572 ms) [ RUN ] VulkanAPITest.unary_op_log [ OK ] VulkanAPITest.unary_op_log (0 ms) [ RUN ] VulkanAPITest.unary_op_log_ [ OK ] VulkanAPITest.unary_op_log_ (59 ms) [ RUN ] VulkanAPITest.querypool_flushed_shader_log xplat/caffe2/aten/src/ATen/test/vulkan_api_test.cpp:7677: Skipped QueryPool is not available [ SKIPPED ] VulkanAPITest.querypool_flushed_shader_log (0 ms) [----------] 4 tests from VulkanAPITest (633 ms total) [----------] Global test environment tear-down [==========] 4 tests from 1 test suite ran. (633 ms total) [ PASSED ] 3 tests. [ SKIPPED ] 1 test, listed below: [ SKIPPED ] VulkanAPITest.querypool_flushed_shader_log YOU HAVE 1 DISABLED TEST ``` Differential Revision: D50961359
diff --git a/aten/src/ATen/native/vulkan/glsl/templates/unary_op_params.yaml b/aten/src/ATen/native/vulkan/glsl/templates/unary_op_params.yaml
@@ -5,6 +5,8 @@ unary_op:
   parameter_values:
     - NAME: sqrt
       OPERATOR: sqrt(X)
+    - NAME: log
+      OPERATOR: log(X)
 
 unary_op_inplace:
   parameter_names_with_default_values:
@@ -13,3 +15,5 @@ unary_op_inplace:
   parameter_values:
     - NAME: sqrt_
       OPERATOR: sqrt(X)
+    - NAME: log_
+      OPERATOR: log(X)
diff --git a/aten/src/ATen/native/vulkan/ops/Softmax.cpp b/aten/src/ATen/native/vulkan/ops/Softmax.cpp
@@ -87,12 +87,12 @@ void set_softmax_kernel_params(
 
 Tensor softmax_internal(
     const at::Tensor& input_arg,
-    const int64_t dim,
-    const bool half_to_float,
-    const bool log_softmax) {
+    const int64_t dim_arg,
+    const bool half_to_float) {
   TORCH_CHECK(
       input_arg.dim() >= 1 && input_arg.dim() <= 4,
       "Vulkan softmax expects 1,2,3 or 4-dimensional input!");
+  int64_t dim = utils::normalize(dim_arg, input_arg.dim());
   TORCH_CHECK(
       dim >= 0 && dim < input_arg.dim(),
       "Softmax dim input was ",
@@ -140,25 +140,15 @@ Tensor softmax_internal(
       0,
   };
   api::ShaderInfo shader_descriptor;
-  if (log_softmax) {
-    if (dim == 1) {
-      shader_descriptor = VK_KERNEL(log_softmax);
-    } else {
-      TORCH_CHECK(
-          dim == 1,
-          "Vulkan log_softmax expects 4-dimensional input with dim=1!");
-    }
-  } else {
-    set_softmax_kernel_params(
-        input_arg.dim(),
-        dim,
-        v_input_sizes,
-        shader_descriptor,
-        input_shader_extents,
-        early_exit,
-        input_dim_stride,
-        input_tensor_dims);
-  }
+  set_softmax_kernel_params(
+      input_arg.dim(),
+      dim,
+      v_input_sizes,
+      shader_descriptor,
+      input_shader_extents,
+      early_exit,
+      input_dim_stride,
+      input_tensor_dims);
 
   const struct Block final {
     ivec4 input_shader_extents;
@@ -197,14 +187,14 @@ Tensor softmax(
     const at::Tensor& input_arg,
     const int64_t dim,
     const bool half_to_float) {
-  return softmax_internal(input_arg, dim, half_to_float, false);
+  return softmax_internal(input_arg, dim, half_to_float);
 }
 
 Tensor log_softmax(
     const at::Tensor& input_arg,
     const int64_t dim,
     const bool half_to_float) {
-  return softmax_internal(input_arg, dim, half_to_float, true);
+  return softmax_internal(input_arg, dim, half_to_float).log();
 }
 
 #ifdef USE_VULKAN_API
diff --git a/aten/src/ATen/native/vulkan/ops/UnaryOp.cpp b/aten/src/ATen/native/vulkan/ops/UnaryOp.cpp
@@ -117,13 +117,23 @@ Tensor& sqrt_(Tensor& self_arg) {
   return unary_op_(self_arg, VK_KERNEL(sqrt_));
 }
 
+Tensor log(const Tensor& self_arg) {
+  return unary_op(self_arg, VK_KERNEL(log));
+}
+
+Tensor& log_(Tensor& self_arg) {
+  return unary_op_(self_arg, VK_KERNEL(log_));
+}
+
 #ifdef USE_VULKAN_API
 
 TORCH_LIBRARY_IMPL(aten, Vulkan, m) {
   m.impl(TORCH_SELECTIVE_NAME("aten::exp"), TORCH_FN(exp));
   m.impl(TORCH_SELECTIVE_NAME("aten::exp_"), TORCH_FN(exp_));
   m.impl(TORCH_SELECTIVE_NAME("aten::sqrt"), TORCH_FN(sqrt));
   m.impl(TORCH_SELECTIVE_NAME("aten::sqrt_"), TORCH_FN(sqrt_));
+  m.impl(TORCH_SELECTIVE_NAME("aten::log"), TORCH_FN(log));
+  m.impl(TORCH_SELECTIVE_NAME("aten::log_"), TORCH_FN(log_));
 }
 
 #endif /* USE_VULKAN_API */
diff --git a/aten/src/ATen/test/vulkan_api_test.cpp b/aten/src/ATen/test/vulkan_api_test.cpp
@@ -3952,33 +3952,65 @@ TEST_F(VulkanAPITest, sigmoid_) {
   ASSERT_TRUE(check);
 }
 
+void test_softmax(const at::IntArrayRef shape, bool log_softmax = false) {
+  at::Tensor in_cpu =
+      at::rand(shape, at::TensorOptions(at::kCPU).dtype(at::kFloat));
+  const at::Tensor in_vulkan = in_cpu.vulkan();
+
+  // Cast to signed to test negative index for dim
+  int64_t size = static_cast<int64_t>(shape.size());
+
+  // Test on all dim
+  for (auto dim = -size; dim < size; dim++) {
+    const at::Tensor out_cpu =
+        log_softmax ? at::log_softmax(in_cpu, dim) : at::softmax(in_cpu, dim);
+
+    const at::Tensor out_vulkan = log_softmax ? at::log_softmax(in_vulkan, dim)
+                                              : at::softmax(in_vulkan, dim);
+    const bool check = almostEqual(out_cpu, out_vulkan.cpu());
+
+    if (!check) {
+      std::cout << "Softmax test failed on axis " << dim << "for tensor dims {";
+      for (uint32_t place = 0; place < shape.size() - 1; place++) {
+        std::cout << shape[place] << " ";
+      }
+      std::cout << shape.back() << "}" << std::endl;
+      showRtol(out_cpu, out_vulkan.cpu());
+    }
+    ASSERT_TRUE(check);
+  }
+}
 
 TEST_F(VulkanAPITest, softmax) {
   c10::InferenceMode mode;
   std::vector<std::vector<int64_t>> test_in_dims = {
-    {1, 3, 4, 2},
-    {4, 8, 5, 7},
-    {9, 11, 12, 12},
+      {1, 3, 4, 2},
+      {4, 8, 5, 7},
+      {9, 11, 12, 12},
   };
-  for (const std::vector<int64_t >& dim_vec : test_in_dims) {
+  bool log_softmax = false;
+  for (const std::vector<int64_t>& dim_vec : test_in_dims) {
     for (uint32_t trunc = 0; trunc < dim_vec.size(); trunc++) {
-      const std::vector<int64_t> trunc_dim_vec = std::vector<int64_t>(dim_vec.begin(), dim_vec.end() - trunc);
-      at::Tensor in_cpu = at::rand(trunc_dim_vec, at::TensorOptions(at::kCPU).dtype(at::kFloat));
-      for (uint32_t dim = 0; dim < trunc_dim_vec.size(); dim++) {
-        const at::Tensor out_cpu = at::softmax(in_cpu, dim);
-        const at::Tensor in_vulkan = in_cpu.vulkan();
-        const at::Tensor out_vulkan = at::softmax(in_vulkan, dim);
-        const bool check = almostEqual(out_cpu, out_vulkan.cpu());
-        if (!check) {
-          std::cout << "Softmax test failed on axis " << dim << "for tensor dims {";
-          for (uint32_t place = 0; place < trunc_dim_vec.size() - 1; place++) {
-            std::cout << trunc_dim_vec[place] << " ";
-          }
-          std::cout << trunc_dim_vec.back() << "}" << std::endl;
-          showRtol(out_cpu, out_vulkan.cpu());
-        }
-        ASSERT_TRUE(check);
-      }
+      const std::vector<int64_t> trunc_dim_vec =
+          std::vector<int64_t>(dim_vec.begin(), dim_vec.end() - trunc);
+      test_softmax(trunc_dim_vec, log_softmax);
+    }
+  }
+}
+
+TEST_F(VulkanAPITest, log_softmax) {
+  c10::InferenceMode mode;
+  std::vector<std::vector<int64_t>> test_in_dims = {
+      {1, 3, 4, 2},
+      {4, 8, 5, 7},
+      {9, 11, 12, 12},
+  };
+  bool log_softmax = true;
+  for (const std::vector<int64_t>& dim_vec : test_in_dims) {
+    for (uint32_t trunc = 0; trunc < dim_vec.size(); trunc++) {
+      const std::vector<int64_t> trunc_dim_vec =
+          std::vector<int64_t>(dim_vec.begin(), dim_vec.end() - trunc);
+      test_softmax(trunc_dim_vec, log_softmax);
     }
   }
 }
@@ -4819,6 +4851,58 @@ TEST_F(VulkanAPITest, unary_op_sqrt_) {
   test_sqrt_({11, 1, 4, 2});
 }
 
+void test_log(const at::IntArrayRef input_shape) {
+  c10::InferenceMode mode;
+  // Need to add a very small constant to avoid 0.
+  const auto in_cpu =
+      at::rand(input_shape, at::device(at::kCPU).dtype(at::kFloat)) + 0.0001;
+  const auto out_cpu = at::log(in_cpu);
+
+  const auto in_vulkan = in_cpu.vulkan();
+  const auto out_vulkan = at::log(in_vulkan);
+
+  const auto check = almostEqual(out_cpu, out_vulkan.cpu());
+  if (!check) {
+    showRtol(out_cpu, out_vulkan.cpu());
+    std::cout << "log test failed with input shape: " << input_shape
+              << std::endl;
+  }
+  ASSERT_TRUE(check);
+}
+
+TEST_F(VulkanAPITest, unary_op_log) {
+  test_log({5});
+  test_log({5, 6});
+  test_log({7, 3, 5});
+  test_log({11, 1, 4, 2});
+}
+
+void test_log_(const at::IntArrayRef input_shape) {
+  c10::InferenceMode mode;
+  // Need to add a very small constant to avoid 0.
+  const auto cpu =
+      at::rand(input_shape, at::device(at::kCPU).dtype(at::kFloat)) + 0.0001;
+  const auto vulkan = cpu.vulkan();
+
+  cpu.log_();
+  vulkan.log_();
+
+  const auto check = almostEqual(cpu, vulkan.cpu());
+  if (!check) {
+    showRtol(cpu, vulkan.cpu());
+    std::cout << "log_ test failed with input shape: " << input_shape
+              << std::endl;
+  }
+  ASSERT_TRUE(check);
+}
+
+TEST_F(VulkanAPITest, unary_op_log_) {
+  test_log_({5});
+  test_log_({5, 6});
+  test_log_({7, 3, 5});
+  test_log_({11, 1, 4, 2});
+}
+
 void test_unsqueeze(const at::IntArrayRef input_shape, int64_t dim) {
   c10::InferenceMode mode;
   const auto in_cpu = at::rand(input_shape, at::device(at::kCPU).dtype(at::kFloat));
