[NPU]speed up bn and momentum (#1487)

Author: LielinJiang

backends/npu/kernels/batch_norm_kernel.cc

@@ -32,6 +32,13 @@ void TransposeKernel(const Context& dev_ctx,
                      const std::vector<int>& axis,
                      phi::DenseTensor* out);
 
+template <typename T, typename Context>
+void FullLikeKernel(const Context& dev_ctx,
+                    const phi::DenseTensor& x,
+                    const phi::Scalar& val,
+                    phi::DataType dtype,
+                    phi::DenseTensor* out);
+
 template <typename T, typename Context>
 void AclopBatchNormKernel(const Context& dev_ctx,
                           const phi::DenseTensor& x,
@@ -536,18 +543,25 @@ void BatchNormKernel(const Context& dev_ctx,
         aclnnInplaceAdd, dev_ctx, *variance_out, *saved_variance, momentum_p);
     auto stream = dev_ctx.stream();
 
-    const auto& adds_runner =
-        NpuOpRunner("Adds",
-                    {*saved_variance},
-                    {*saved_variance},
-                    {{"value", static_cast<float>(epsilon)}});
-    adds_runner.Run(stream);
-    const auto& inv_runner =
-        NpuOpRunner("Inv", {*saved_variance}, {*saved_variance}, {});
-    inv_runner.Run(stream);
-    const auto& sqrt_ruuner =
-        NpuOpRunner("Sqrt", {*saved_variance}, {*saved_variance}, {});
-    sqrt_ruuner.Run(stream);
+    phi::Scalar one_scalar = static_cast<float>(1.0);
+
+    phi::DenseTensor epsilon_tensor;
+    epsilon_tensor.set_meta(saved_variance->meta());
+    custom_kernel::FullLikeKernel<T, Context>(
+        dev_ctx,
+        *saved_variance,
+        phi::Scalar(static_cast<float>(epsilon)),
+        saved_variance->dtype(),
+        &epsilon_tensor);
+
+    EXEC_NPU_CMD(aclnnAdd,
+                 dev_ctx,
+                 *saved_variance,
+                 epsilon_tensor,
+                 one_scalar,
+                 *saved_variance);
+
+    EXEC_NPU_CMD(aclnnInplaceRsqrt, dev_ctx, *saved_variance);
   }
 }
 

backends/npu/kernels/momentum_kernel.cc

@@ -48,12 +48,13 @@ void MomentumKernel(const Context& dev_ctx,
     regularized_grad.Resize(grad.dims());
     dev_ctx.template Alloc<T>(&regularized_grad);
 
-    const auto& runner1 = NpuOpRunner(
-        "Muls", {param}, {regularized_grad}, {{"value", regularization_coeff}});
-    runner1.Run(dev_ctx.stream());
-    const auto& runner2 =
-        NpuOpRunner("Add", {regularized_grad, grad}, {regularized_grad}, {});
-    runner2.Run(dev_ctx.stream());
+    phi::Scalar regularization_coeff_scalar = regularization_coeff;
+    EXEC_NPU_CMD(aclnnAdd,
+                 dev_ctx,
+                 grad,
+                 param,
+                 regularization_coeff_scalar,
+                 regularized_grad);
   } else {
     regularized_grad = grad;
   }