add asinh on diopi torch impl

DeepLink-org · Aug 14, 2024 · 766a976 · 766a976
1 parent dca4fcc
commit 766a976
Show file tree

Hide file tree

Showing 4 changed files with 74 additions and 0 deletions.
diff --git a/diopi_test/python/configs/diopi_configs.py b/diopi_test/python/configs/diopi_configs.py
@@ -1307,6 +1307,44 @@
         ),
     ),
 
+    'asinh': dict(
+        name=['asinh'],
+        interface=['torch'],
+        is_inplace=True,
+        saved_args=dict(output=0),
+        dtype=[np.float16, np.float32, np.float64],
+        tensor_para=dict(
+            gen_fn='Genfunc.randn',
+            args=[
+                {
+                    "ins": ['input'],
+                    "shape": ((), (1, ), (1024,), (364800, 4), (2, 128, 3072),
+                              (256, 128, 3, 3),
+                              (2, 31, 512, 6, 40),
+                              (0,), (16, 0), (1, 0, 6)),
+                },
+            ],
+        ),
+    ),
+
+    'asinh_not_float': dict(
+        name=['asinh'],
+        interface=['torch'],
+        dtype=[np.int16, np.int32, np.int64, np.uint8, np.int8, np.bool_],
+        tensor_para=dict(
+            gen_fn='Genfunc.randn',
+            args=[
+                {
+                    "ins": ['input'],
+                    "shape": ((), (1, ), (1024,), (364800, 4), (2, 128, 3072),
+                              (256, 128, 3, 3),
+                              (2, 31, 512, 6, 40),
+                              (0,), (16, 0), (1, 0, 6)),
+                },
+            ],
+        ),
+    ),
+
     'sign': dict(
         name=['sign'],
         interface=['torch'],

diff --git a/diopi_test/python/conformance/diopi_functions.py b/diopi_test/python/conformance/diopi_functions.py
@@ -429,6 +429,10 @@ def atan(input, inplace=False) -> Tensor:
     return unary_op(input, inplace, "diopiAtan", promote_type(input, Dtype.float32))
 
 
+def asinh(input, inplace=False) -> Tensor:
+    return unary_op(input, inplace, "diopiAsinh", promote_type(input, Dtype.float32))
+
+
 def exp(input, inplace=False) -> Tensor:
     return unary_op(input, inplace, "diopiExp", promote_type(input, Dtype.float32))
 

diff --git a/impl/torch/functions/functions.cpp b/impl/torch/functions/functions.cpp
@@ -1026,6 +1026,23 @@ diopiError_t diopiAtanInp(diopiContextHandle_t ctx, diopiTensorHandle_t input) {
     return diopiSuccess;
 }
 
+diopiError_t diopiAsinh(diopiContextHandle_t ctx, diopiTensorHandle_t out, diopiConstTensorHandle_t input) {
+    impl::aten::setCurStream(ctx);
+    auto atInput = impl::aten::buildATen(input);
+    auto atOut = impl::aten::buildATen(out);
+    CALL_ATEN_CUDA_FUNC(asinh_out, atOut, atInput);
+
+    return diopiSuccess;
+}
+
+diopiError_t diopiAsinhInp(diopiContextHandle_t ctx, diopiTensorHandle_t input) {
+    impl::aten::setCurStream(ctx);
+    auto atInput = impl::aten::buildATen(input);
+    CALL_ATEN_CUDA_FUNC(asinh_, atInput);
+
+    return diopiSuccess;
+}
+
 diopiError_t diopiSigmoid(diopiContextHandle_t ctx, diopiTensorHandle_t out, diopiConstTensorHandle_t input) {
     impl::aten::setCurStream(ctx);
     auto atInput = impl::aten::buildATen(input);

diff --git a/proto/include/diopi/functions.h b/proto/include/diopi/functions.h
@@ -977,6 +977,21 @@ DIOPI_API diopiError_t diopiAtan(diopiContextHandle_t ctx, diopiTensorHandle_t o
  */
 DIOPI_API diopiError_t diopiAtanInp(diopiContextHandle_t ctx, diopiTensorHandle_t input);
 
+/**
+ * @brief The in-place version of diopiAsinh().
+ * @param[in] ctx Context environment.
+ * @param[inout] input the input tensor and will be stroed reuslt tensor. type = [float16, float32, float64].
+ */
+DIOPI_API diopiError_t diopiAsinhInp(diopiContextHandle_t ctx, diopiTensorHandle_t input);
+
+/**
+ * @brief Returns a new tensor with the arc hyperbolic sine of the elements of input.
+ * @param[in] ctx Context environment.
+ * @param[in] input the input tensor. type = [float16, float32, float64].
+ * @param[out] out the output tensor. type = [float16, float32, float64].
+ */
+DIOPI_API diopiError_t diopiAsinh(diopiContextHandle_t ctx, diopiTensorHandle_t out, diopiConstTensorHandle_t input);
+
 /**
  * @brief The in-place version of diopiSigmoid().
  * @param[in] ctx Context environment.