optimized expand_as_kernel (PaddlePaddle#57509)

danleifeng · Sep 22, 2023 · 0717491 · 0717491
1 parent b4bcf93
commit 0717491
Showing 1 changed file with 2 additions and 5 deletions.
diff --git a/paddle/phi/kernels/gpu/expand_as_kernel.cu b/paddle/phi/kernels/gpu/expand_as_kernel.cu
@@ -18,6 +18,7 @@
 #include "paddle/phi/common/scalar.h"
 #include "paddle/phi/core/dense_tensor.h"
 #include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/expand_kernel.h"
 #include "paddle/phi/kernels/funcs/broadcast_function.h"
 
 namespace phi {
@@ -70,11 +71,7 @@ void ExpandAsKernel(const Context& ctx,
     }
   }
 
-  out->Resize(phi::make_ddim(target_shape));
-  ctx.template Alloc<T>(out);
-  std::vector<const DenseTensor*> ins = {&x};
-  std::vector<DenseTensor*> outs = {out};
-  phi::funcs::BroadcastKernel<T>(ctx, ins, &outs, kps::IdentityFunctor<T>());
+  ExpandKernel<T, Context>(ctx, x, target_shape, out);
 }
 
 }  // namespace phi