[NPU] add npu kernel for truncated_gaussian_random op (#31654)

* init

* add todo

* add npu kernel for truncated_gaussian_random

* add sync

* fix concat_grad

* fix typo

paddle/fluid/operators/concat_op_npu.cc

@@ -80,7 +80,6 @@ class ConcatGradNPUKernel : public framework::OpKernel<T> {
     axis = ComputeAxis(static_cast<int64_t>(axis),
                        static_cast<int64_t>(ins[0]->dims().size()));
 
-    std::vector<int> sizes;
     int offset = 0;
     auto stream =
         ctx.template device_context<paddle::platform::NPUDeviceContext>()
@@ -91,7 +90,6 @@ class ConcatGradNPUKernel : public framework::OpKernel<T> {
       if (out_var_names[j] != framework::kEmptyVarName &&
           outs[j]->numel() != 0UL) {
         outs[j]->mutable_data<T>(ctx.GetPlace());
-        sizes.push_back(outs[j]->dims()[axis]);
         std::vector<int> offsets;
         std::vector<int> sizes;
         for (int dim = 0; dim < ins[j]->dims().size(); ++dim) {
@@ -103,9 +101,8 @@ class ConcatGradNPUKernel : public framework::OpKernel<T> {
             sizes.push_back(ins[j]->dims()[dim]);
           }
         }
-        auto runner =
-            NpuOpRunner("SliceD", {*out_grad}, {*outs[j]},
-                        {{"offsets", offset}, {"size", ins[j]->dims()[axis]}});
+        auto runner = NpuOpRunner("SliceD", {*out_grad}, {*outs[j]},
+                                  {{"offsets", offsets}, {"size", sizes}});
         runner.Run(stream);
       }
       if (ins[j]->numel() != 0UL) {

paddle/fluid/operators/truncated_gaussian_random_op_npu.cc

+/* Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#include "paddle/fluid/operators/truncated_gaussian_random_op.h"
+#include <memory>
+#include <string>
+#include "paddle/fluid/operators/npu_op_runner.h"
+
+namespace paddle {
+namespace operators {
+
+using Tensor = framework::Tensor;
+
+template <typename DeviceContext, typename T>
+class TruncatedGaussianRandomNPUKernel : public framework::OpKernel<T> {
+ public:
+  void Compute(const framework::ExecutionContext& ctx) const override {
+    // TODO(zhiqiu): support dynamic shape and call ParameterizedTruncatedNormal
+    std::vector<int> shape = ctx.Attr<std::vector<int>>("shape");
+    Tensor shape_tensor(framework::proto::VarType::INT32);
+    shape_tensor.mutable_data<int32_t>({static_cast<int>(shape.size())},
+                                       ctx.GetPlace());
+    TensorFromVector(shape, ctx.device_context(), &shape_tensor);
+    float mean = ctx.Attr<float>("mean");
+    Tensor mean_tensor(framework::proto::VarType::FP32);
+    mean_tensor.mutable_data<float>({1}, ctx.GetPlace());
+    TensorFromVector(std::vector<float>{mean}, ctx.device_context(),
+                     &mean_tensor);
+
+    float std = ctx.Attr<float>("std");
+    Tensor std_tensor(framework::proto::VarType::FP32);
+    std_tensor.mutable_data<float>({1}, ctx.GetPlace());
+    TensorFromVector(std::vector<float>{std}, ctx.device_context(),
+                     &std_tensor);
+
+    int32_t seed_var = ctx.Attr<int32_t>("seed");
+
+    Tensor min_tensor(framework::proto::VarType::FP32);
+    min_tensor.mutable_data<float>({1}, ctx.GetPlace());
+    float min_value = mean - std * 2.0;
+    TensorFromVector(std::vector<float>{min_value}, ctx.device_context(),
+                     &min_tensor);
+
+    Tensor max_tensor(framework::proto::VarType::FP32);
+    max_tensor.mutable_data<float>({1}, ctx.GetPlace());
+    float max_value = mean + std * 2.0;
+    TensorFromVector(std::vector<float>{max_value}, ctx.device_context(),
+                     &max_tensor);
+
+    auto* out = ctx.Output<framework::Tensor>("Out");
+    out->mutable_data<T>(ctx.GetPlace());
+    auto stream =
+        ctx.template device_context<paddle::platform::NPUDeviceContext>()
+            .stream();
+    auto runner = NpuOpRunner(
+        "ParameterizedTruncatedNormal",
+        {shape_tensor, mean_tensor, std_tensor, min_tensor, max_tensor}, {*out},
+        {{"seed", seed_var}});
+    runner.Run(stream);
+  }
+};
+
+// NOTE(zhiqiu): actually, this is cpu version kernel, and we need to make the
+// above
+// npu version work in the future.
+template <typename T>
+class NPUTruncatedGaussianRandomKernel : public framework::OpKernel<T> {
+ public:
+  void Compute(const framework::ExecutionContext& context) const override {
+    float mean = context.Attr<float>("mean");
+    float std = context.Attr<float>("std");
+    auto* tensor = context.Output<framework::Tensor>("Out");
+    tensor->mutable_data<T>(context.GetPlace());
+
+    Tensor cpu_tensor(tensor->type());
+    cpu_tensor.Resize(tensor->dims());
+    T* cpu_data = cpu_tensor.mutable_data<T>(platform::CPUPlace());
+    std::uniform_real_distribution<T> dist(std::numeric_limits<float>::min(),
+                                           1.0);
+    TruncatedNormal<T> truncated_normal(mean, std);
+    int64_t size = tensor->numel();
+
+    unsigned int seed = static_cast<unsigned int>(context.Attr<int>("seed"));
+    auto engine = framework::GetCPURandomEngine(seed);
+    for (int64_t i = 0; i < size; ++i) {
+      cpu_data[i] = truncated_normal(dist(*engine));
+    }
+    framework::TensorCopy(
+        cpu_tensor, context.GetPlace(),
+        context.template device_context<platform::DeviceContext>(), tensor);
+    context.template device_context<paddle::platform::NPUDeviceContext>()
+        .Wait();
+  }
+};
+
+}  // namespace operators
+}  // namespace paddle
+
+namespace ops = paddle::operators;
+
+REGISTER_OP_NPU_KERNEL(truncated_gaussian_random,
+                       ops::NPUTruncatedGaussianRandomKernel<float>);

python/paddle/fluid/tests/unittests/npu/test_truncated_gaussian_random_op_npu.py

+#  Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import print_function
+
+import numpy as np
+import unittest
+import sys
+sys.path.append("..")
+from op_test import OpTest
+import paddle
+import paddle.fluid as fluid
+import paddle.fluid.core as core
+from paddle.fluid.op import Operator
+from paddle.fluid.executor import Executor
+
+paddle.enable_static()
+SEED = 2021
+
+
+@unittest.skipIf(not paddle.is_compiled_with_npu(),
+                 "core is not compiled with NPU")
+class TestTruncatedNormal(unittest.TestCase):
+    def _test(self, run_npu=True):
+        main_prog = paddle.static.Program()
+        startup_prog = paddle.static.Program()
+        scope = paddle.fluid.core.Scope()
+
+        main_prog.random_seed = SEED
+        startup_prog.random_seed = SEED
+        np.random.seed(SEED)
+        paddle.seed(SEED)
+
+        with fluid.scope_guard(scope):
+            with paddle.static.program_guard(main_prog, startup_prog):
+                weight_attr = paddle.framework.ParamAttr(
+                    name="linear_weight",
+                    initializer=paddle.nn.initializer.TruncatedNormal(
+                        mean=0.0, std=2.0))
+                linear = paddle.nn.Linear(
+                    2, 2, weight_attr=weight_attr, bias_attr=False)
+
+            if run_npu:
+                place = paddle.NPUPlace(0)
+            else:
+                place = paddle.CPUPlace()
+
+            exe = paddle.static.Executor(place)
+            w = exe.run(startup_prog, fetch_list=['linear_weight'])
+            return w
+
+    def test_npu(self):
+        cpu_w = self._test(False)
+        npu_w = self._test(True)
+
+        self.assertTrue(np.allclose(npu_w, cpu_w))
+
+
+if __name__ == '__main__':
+    unittest.main()