From c3f7c5422e4c979e944a85e6d0a1b327c33a67be Mon Sep 17 00:00:00 2001
From: hjhee <hjhee@users.noreply.github.com>
Date: Thu, 17 Oct 2024 19:02:01 +0800
Subject: [PATCH] Add aten::log_normal_ (#766)

- log_normal_

---------

Co-authored-by: Yutao Xu <yutao.xu@intel.com>
---
 src/ATen/native/xpu/Distributions.cpp         |  1 +
 src/ATen/native/xpu/XPUFallback.template      |  1 -
 .../native/xpu/sycl/DistributionKernels.h     |  6 +++
 .../xpu/sycl/DistributionLogNormalKernel.cpp  | 17 +++++++++
 .../native/xpu/sycl/DistributionTemplates.h   | 38 +++++++++++++++++++
 test/xpu/xpu_test_utils.py                    |  1 +
 yaml/native/native_functions.yaml             |  8 ++++
 7 files changed, 71 insertions(+), 1 deletion(-)
 create mode 100644 src/ATen/native/xpu/sycl/DistributionLogNormalKernel.cpp

diff --git a/src/ATen/native/xpu/Distributions.cpp b/src/ATen/native/xpu/Distributions.cpp
index bce51bbf8..264a368f8 100644
--- a/src/ATen/native/xpu/Distributions.cpp
+++ b/src/ATen/native/xpu/Distributions.cpp
@@ -28,5 +28,6 @@ REGISTER_XPU_DISPATCH(
 REGISTER_XPU_DISPATCH(
     multinomial_with_replacement_stub,
     &xpu::multinomial_kernel);
+REGISTER_XPU_DISPATCH(log_normal_stub, &xpu::log_normal_kernel);
 } // namespace native
 } // namespace at
diff --git a/src/ATen/native/xpu/XPUFallback.template b/src/ATen/native/xpu/XPUFallback.template
index 7226a6c4e..6819e9c51 100644
--- a/src/ATen/native/xpu/XPUFallback.template
+++ b/src/ATen/native/xpu/XPUFallback.template
@@ -213,7 +213,6 @@ TORCH_LIBRARY_IMPL(aten, XPU, m) {
     "_linalg_svd.U",
     "linspace.out",
     "_logcumsumexp",
-    "log_normal_",
     "logspace.out",
     "lu_unpack.out",
     "max_pool3d_with_indices",
diff --git a/src/ATen/native/xpu/sycl/DistributionKernels.h b/src/ATen/native/xpu/sycl/DistributionKernels.h
index 88d62933f..494033dbb 100644
--- a/src/ATen/native/xpu/sycl/DistributionKernels.h
+++ b/src/ATen/native/xpu/sycl/DistributionKernels.h
@@ -45,4 +45,10 @@ TORCH_XPU_API void exponential_kernel(
     double lambda,
     c10::optional<Generator> gen);
 
+TORCH_XPU_API void log_normal_kernel(
+    TensorIteratorBase& iter,
+    double mean,
+    double std,
+    std::optional<Generator> gen);
+
 } // namespace at::native::xpu
diff --git a/src/ATen/native/xpu/sycl/DistributionLogNormalKernel.cpp b/src/ATen/native/xpu/sycl/DistributionLogNormalKernel.cpp
new file mode 100644
index 000000000..4726209e0
--- /dev/null
+++ b/src/ATen/native/xpu/sycl/DistributionLogNormalKernel.cpp
@@ -0,0 +1,17 @@
+#include <ATen/native/TensorIterator.h>
+#include <ATen/native/xpu/sycl/DistributionTemplates.h>
+#include <ATen/xpu/XPUGeneratorImpl.h>
+
+namespace at::native::xpu {
+
+void log_normal_kernel(
+    TensorIteratorBase& iter,
+    double mean,
+    double std,
+    std::optional<Generator> gen) {
+  auto generator = get_generator_or_default<at::XPUGeneratorImpl>(
+      gen, at::xpu::detail::getDefaultXPUGenerator());
+  at::native::templates::xpu::log_normal_kernel(iter, mean, std, generator);
+}
+
+} // namespace at::native::xpu
diff --git a/src/ATen/native/xpu/sycl/DistributionTemplates.h b/src/ATen/native/xpu/sycl/DistributionTemplates.h
index f5a5efdb5..f696e50d1 100644
--- a/src/ATen/native/xpu/sycl/DistributionTemplates.h
+++ b/src/ATen/native/xpu/sycl/DistributionTemplates.h
@@ -698,6 +698,44 @@ void exponential_kernel(TensorIteratorBase& iter, double lambda, RNG gen) {
       });
 }
 
+// ====================== LogNormal ======================
+
+template <typename scalar_t, typename accscalar_t>
+struct LogNormalFunctor {
+  scalar_t operator()(accscalar_t rand) const {
+    return static_cast<scalar_t>(transformation::log_normal<accscalar_t>(
+        transformation::normal<accscalar_t>(rand, mean_, std_)));
+  }
+  LogNormalFunctor(accscalar_t mean, accscalar_t std)
+      : mean_(mean), std_(std) {}
+
+ private:
+  accscalar_t mean_;
+  accscalar_t std_;
+};
+
+template <typename RNG>
+void log_normal_kernel(
+    TensorIteratorBase& iter,
+    double mean,
+    double std,
+    RNG gen) {
+  AT_DISPATCH_FLOATING_TYPES_AND2(
+      at::ScalarType::Half,
+      at::ScalarType::BFloat16,
+      iter.dtype(),
+      "log_normal_xpu_",
+      [&] {
+        using accscalar_t = at::acc_type_device<scalar_t, kXPU>;
+        auto mean_ = static_cast<accscalar_t>(mean);
+        auto std_ = static_cast<accscalar_t>(std);
+        // define functor to multiply std and add mean
+        LogNormalFunctor<scalar_t, accscalar_t> log_normal_functor(mean_, std_);
+        normal_and_transform<scalar_t, accscalar_t, rand4_engine_calls>(
+            iter, gen, log_normal_functor);
+      });
+}
+
 } // namespace xpu
 } // namespace templates
 } // namespace native
diff --git a/test/xpu/xpu_test_utils.py b/test/xpu/xpu_test_utils.py
index 42f6ab715..80922171a 100644
--- a/test/xpu/xpu_test_utils.py
+++ b/test/xpu/xpu_test_utils.py
@@ -249,6 +249,7 @@
     "square",
     "heaviside",
     "argsort",
+    "log_normal",
 ]
 
 _ops_without_cuda_support = [
diff --git a/yaml/native/native_functions.yaml b/yaml/native/native_functions.yaml
index 71fd4f342..2e9b5aabb 100644
--- a/yaml/native/native_functions.yaml
+++ b/yaml/native/native_functions.yaml
@@ -4100,6 +4100,14 @@
     XPU: fmin_out
   tags: pointwise
 
+- func: log_normal_(Tensor(a!) self, float mean=1, float std=2, *, Generator? generator=None) -> Tensor(a!)
+  device_check: NoCheck   # TensorIterator
+  tags: nondeterministic_seeded
+  variants: method
+  dispatch:
+    XPU: log_normal_
+  autogen: log_normal, log_normal.out
+
 - func: exponential_(Tensor(a!) self, float lambd=1, *, Generator? generator=None) -> Tensor(a!)
   device_check: NoCheck   # TensorIterator
   tags: nondeterministic_seeded