impl std with aclnn api for ascend

DeepLink-org · Jun 27, 2024 · f962070 · f962070
1 parent 1399f5f
commit f962070
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 0 deletions.
diff --git a/impl/ascend/convert_config.yaml b/impl/ascend/convert_config.yaml
@@ -458,3 +458,7 @@
 - diopiMaxPool2dBackward:
     tensor_dtype:
         indices: (int64)->int32
+
+- diopiStd:
+    dtype: (float64)->float32
+    layout: ND
diff --git a/impl/ascend/functions/reduce.cpp b/impl/ascend/functions/reduce.cpp
@@ -3,6 +3,7 @@
  * @author DeepLink
  * @copyright  (c) 2023, DeepLink.
  */
+#include <numeric>
 
 #include "../aclnn/acl_scalar.hpp"
 #include "../aclnn/adaptor.hpp"
@@ -52,6 +53,31 @@ diopiError_t diopiMean(diopiContextHandle_t ctx, diopiTensorHandle_t out, diopiC
     return diopiSuccess;
 }
 
+diopiError_t diopiStd(diopiContextHandle_t ctx, diopiTensorHandle_t out, diopiConstTensorHandle_t input, diopiSize_t dim, bool unbiased) {
+    AscendTensor inputAt(input);
+    AscendTensor outAt(out);
+
+    bool keepdim = false;
+    if (inputAt.dim() == outAt.dim()) {
+        keepdim = true;
+    }
+
+    int64_t correction = 0;
+    if (unbiased) {
+        correction = 1;
+    }
+
+    if (dim.data == nullptr || dim.len == 0) {
+        std::vector<int64_t> allDim(inputAt.dim());
+        std::iota(allDim.begin(), allDim.end(), 0);
+        diopiSize_t rDim = vectorToDiopiSize(allDim);
+        DIOPI_ASCEND_CALL_ACLNN(aclnnStd, ctx, input, rDim, correction, keepdim, out);
+    } else {
+        DIOPI_ASCEND_CALL_ACLNN(aclnnStd, ctx, input, dim, correction, keepdim, out);
+    }
+    return diopiSuccess;
+}
+
 diopiError_t diopiAll(diopiContextHandle_t ctx, diopiTensorHandle_t out, diopiConstTensorHandle_t input, const int64_t* dim) {
     diopiSize_t inputSize, outSize;
     diopiGetTensorShape(input, &inputSize);

diff --git a/impl/ascend_npu/ascend_config.yaml b/impl/ascend_npu/ascend_config.yaml
@@ -223,6 +223,7 @@ ascend:
 - diopiSqrt
 - diopiSqrtInp
 - diopiStack
+- diopiStd
 - diopiTanh
 - diopiTanhBackward
 - diopiTanhInp