[ascend]zq/update LayerNorm by aclnn (#1204)

Author: NeosZhang

impl/ascend/aclnn/adaptor.hpp

@@ -12,6 +12,7 @@
 #include <dlfcn.h>
 
 #include <array>
+#include <cassert>
 #include <cstddef>
 #include <cstdint>
 #include <functional>
@@ -108,6 +109,17 @@ inline aclScalar* createAclScalarFromDiopiScalar(const diopiScalar_t* scalar) {
 
 inline aclIntArray* createAclIntArrayFromDiopiSize(const diopiSize_t size) { return ::aclCreateIntArray(size.data, size.len); }
 
+template <size_t N>
+inline aclBoolArray* createAclBoolArrayFromVector(const std::array<bool, N>& vec) {
+    return ::aclCreateBoolArray(vec.data(), vec.size());
+}
+
+template <typename T>
+struct IsBoolStdArray : std::false_type {};
+
+template <std::size_t N>
+struct IsBoolStdArray<std::array<bool, N>> : std::true_type {};
+
 inline aclIntArray* createAclIntArrayFromIntVector(const std::vector<int64_t>& vec) { return ::aclCreateIntArray(vec.data(), vec.size()); }
 
 inline aclTensorList* createAclTensorListFromDiopiTensorVector(const std::vector<diopiTensorHandle_t>& tensorsVec) {
@@ -144,6 +156,8 @@ decltype(auto) convertType(T&& param) {
         return createAclIntArrayFromIntVector(std::forward<T>(param));
     } else if constexpr (std::is_same_v<U, diopiDtype_t>) {
         return diopiDtypeToAclDataType(std::forward<T>(param));
+    } else if constexpr (IsBoolStdArray<U>::value) {
+        return createAclBoolArrayFromVector<std::tuple_size_v<U>>(std::forward<T>(param));
     } else {
         static_assert(!std::is_class_v<U> && !std::is_pointer_v<U>);
         return std::forward<T>(param);

impl/ascend/functions/layer_norm.cpp

@@ -4,7 +4,8 @@
  * @copyright  (c) 2023, DeepLink.
  */
 
-#include "../common/acloprunner.hpp"
+#include "../aclnn/acl_scalar.hpp"
+#include "../aclnn/adaptor.hpp"
 
 namespace impl {
 namespace ascend {
@@ -13,8 +14,10 @@ diopiError_t diopiLayerNorm(diopiContextHandle_t ctx, diopiTensorHandle_t out, d
                             diopiConstTensorHandle_t input, diopiConstTensorHandle_t weight, diopiConstTensorHandle_t bias, diopiSize_t normalizedShape,
                             double eps) {
     AscendTensor inputAt(input);
+    AscendTensor outAt(out);
     if (0 == inputAt.numel()) {
-        AclOpRunner<1, 1>("Fills", ctx).addInput(out).setAttr<float>("value", 0).addOutput(out).run();
+        diopiScalar_t zeroScalar = constructDiopiScalarT(outAt.dtype(), 0.0);
+        DIOPI_ASCEND_CALL_ACLNN(aclnnInplaceFillScalar, ctx, out, &zeroScalar);
         return diopiSuccess;
     }
 
@@ -27,18 +30,8 @@ diopiError_t diopiLayerNorm(diopiContextHandle_t ctx, diopiTensorHandle_t out, d
     const int axis = inShape.len - normalizedShape.len;
     int64_t beginDim = axis;
 
-    // call acl op
-    AclOpRunner<3, 3>("LayerNorm", ctx)
-        .addInput(input)
-        .addInput(weightTemp)
-        .addInput(biasTemp)
-        .addOutput(out)
-        .addOutput(saveMean)
-        .addOutput(saveInvstd)
-        .setAttr("begin_norm_axis", beginDim)
-        .setAttr("begin_params_axis", beginDim)
-        .setAttr<float>("epsilon", eps)
-        .run();
+    // call aclnnLayerNorm
+    DIOPI_ASCEND_CALL_ACLNN(aclnnLayerNorm, ctx, input, normalizedShape, weightTemp, biasTemp, eps, out, saveMean, saveInvstd);
     return diopiSuccess;
 }
 
@@ -47,6 +40,7 @@ diopiError_t diopiLayerNormBackward(diopiContextHandle_t ctx, diopiTensorHandle_
                                     diopiConstTensorHandle_t bias, diopiConstTensorHandle_t mean, diopiConstTensorHandle_t rstd, diopiSize_t normalizedShape) {
     AscendTensor inputAt(input);
     diopiTensorHandle_t weightTemp = createTensorIfNullptrOrConstCast(ctx, weight, normalizedShape, inputAt.dtype(), true, 1);
+    diopiTensorHandle_t biasTemp = createTensorIfNullptrOrConstCast(ctx, bias, normalizedShape, inputAt.dtype(), true, 0);
     diopiTensorHandle_t gradWeightTemp = createTensorIfNullptrOrConstCast(ctx, gradWeight, normalizedShape, inputAt.dtype(), false, 0);
     diopiTensorHandle_t gradBiasTemp = createTensorIfNullptrOrConstCast(ctx, gradBias, normalizedShape, inputAt.dtype(), false, 0);
 
@@ -57,16 +51,30 @@ diopiError_t diopiLayerNormBackward(diopiContextHandle_t ctx, diopiTensorHandle_
         rstdAt.unsqueeze(rstdAt.dim());
     }
 
-    AclOpRunner<5, 3>("LayerNormGrad", ctx)
-        .addInput(gradOutput)
-        .addInput(input)
-        .addInput(rstdAt)
-        .addInput(meanAt)
-        .addInput(weightTemp)
-        .addOutput(gradInput)
-        .addOutput(gradWeightTemp)
-        .addOutput(gradBiasTemp)
-        .run();
+    std::array<bool, 3> gradMask = {true, true, true};
+    if (nullptr == gradInput) {
+        gradMask[0] = false;
+    }
+    if (nullptr == gradWeight) {
+        gradMask[1] = false;
+    }
+    if (nullptr == gradBias) {
+        gradMask[2] = false;
+    }
+
+    DIOPI_ASCEND_CALL_ACLNN(aclnnLayerNormBackward,
+                            ctx,
+                            gradOutput,
+                            inputAt,
+                            normalizedShape,
+                            meanAt,
+                            rstdAt,
+                            weightTemp,
+                            biasTemp,
+                            gradMask,
+                            gradInput,
+                            gradWeightTemp,
+                            gradBiasTemp);
     return diopiSuccess;
 }