Compile optimize frobenius_norm op and trace op (#57327)

* add sqrtkerenl * update * add trace_kernel * update --------- Co-authored-by: niuliling123 <sdust_niull@163.com>
PaddlePaddle · Sep 19, 2023 · 92f5c10 · 92f5c10
1 parent 86308a3
commit 92f5c10
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 6 deletions.
diff --git a/paddle/phi/kernels/gpu/frobenius_norm_kernel.cu b/paddle/phi/kernels/gpu/frobenius_norm_kernel.cu
@@ -15,8 +15,10 @@
 #include "paddle/phi/kernels/frobenius_norm_kernel.h"
 
 #include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/activation_kernel.h"
 #include "paddle/phi/kernels/funcs/activation_functor.h"
 #include "paddle/phi/kernels/gpu/reduce.h"
+
 namespace phi {
 
 template <typename T, typename Context>
@@ -30,10 +32,8 @@ void FrobeniusNormKernel(const Context& dev_ctx,
   auto out_dtype = x.dtype();
   phi::Reduce<T, kps::AddFunctor, kps::SquareFunctor>(
       dev_ctx, x, reduce_all, dims, keep_dim, out_dtype, out);
-  std::vector<const DenseTensor*> ins = {out};
-  std::vector<DenseTensor*> outs = {out};
-  auto functor = funcs::CudaSqrtFunctor<T>();
-  funcs::ElementwiseKernel<T>(dev_ctx, ins, &outs, functor);
+
+  SqrtKernel<T, Context>(dev_ctx, *out, out);
 }
 
 }  // namespace phi

diff --git a/paddle/phi/kernels/gpu/trace_kernel.cu b/paddle/phi/kernels/gpu/trace_kernel.cu
@@ -18,6 +18,7 @@
 #include "paddle/phi/core/kernel_registry.h"
 #include "paddle/phi/kernels/funcs/diagonal.h"
 #include "paddle/phi/kernels/funcs/reduce_function.h"
+#include "paddle/phi/kernels/reduce_sum_kernel.h"
 
 namespace phi {
 
@@ -36,8 +37,8 @@ void TraceKernel(const Context& ctx,
     auto out_dim_size = out->dims().size();
     if (out_dim_size == 0) out_dim_size = 1;
     reduce_dims.push_back(out_dim_size);
-    funcs::ReduceKernel<T, T, kps::AddFunctor, kps::IdentityFunctor<T>>(
-        ctx, diag, out, kps::IdentityFunctor<T>(), reduce_dims);
+    phi::SumKernel<T, Context>(
+        ctx, diag, reduce_dims, diag.dtype(), false, out);
   } else {
     phi::funcs::SetConstant<Context, T> functor;
     functor(ctx, out, static_cast<T>(0));