【PaddlePaddle Hackathon 4 No.36】为 Paddle 优化 tile op 在 GPU 上的计算性能

paddle/phi/kernels/gpu/tile_kernel.cu

@@ -13,10 +13,90 @@
 // limitations under the License.
 
 #include "paddle/phi/kernels/tile_kernel.h"
-
 #include "paddle/phi/backends/gpu/gpu_context.h"
 #include "paddle/phi/core/kernel_registry.h"
-#include "paddle/phi/kernels/impl/tile_kernel_impl.h"
+#include "paddle/phi/kernels/funcs/broadcast_function.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void TileKernel(const Context& dev_ctx,
+                const DenseTensor& x,
+                const IntArray& repeat_times,
+                DenseTensor* out) {
+  auto x_dims = x.dims();
+  auto rank = x_dims.size();
+  auto repeat_times_data = repeat_times.GetData();
+  int repeat_times_size = repeat_times_data.size();
+  rank = std::max(rank, repeat_times_size);
+
+  if (rank == 0) {
+    phi::Copy<DeviceContext>(dev_ctx, x, dev_ctx.GetPlace(), false, out);
+    return;
+  }
+
+  for (size_t i = 0; i < repeat_times_data.size(); ++i) {
+    PADDLE_ENFORCE_GT(
+        repeat_times_data[i],
+        0,
+        errors::InvalidArgument(
+            "All elements of the input 'repeat_times' for tile op must "
+            "be positive integers, but the value received is %d.",
+            repeat_times_data[i]));
+  }
+
+  auto vec_x_dims = phi::vectorize<int>(x_dims);
+  if (repeat_times_data.size() < vec_x_dims.size()) {
+    int diff = vec_x_dims.size() - repeat_times_data.size();
+    repeat_times_data.insert(repeat_times_data.begin(), diff, 1);
+  } else {
+    int diff = repeat_times_data.size() - vec_x_dims.size();
+    vec_x_dims.insert(vec_x_dims.begin(), diff, 1);
+  }
+
+  PADDLE_ENFORCE_EQ(
+      repeat_times_data.size(),
+      vec_x_dims.size(),
+      errors::InvalidArgument(
+          "The rank (%d) of the input 'x' and the rank (%d) of the input "
+          "'repeat_times' for tile op must match after promotion.",
+          vec_x_dims.size(),
+          repeat_times_data.size()));
+
+  DDim new_x_dims = make_ddim(vec_x_dims);
+  DDim out_dims(new_x_dims);
+  DenseTensor new_x = x;
+  vec_x_dims.insert(vec_x_dims.begin(), 1, 1);
+  for (size_t i = 0; i < repeat_times_data.size(); ++i) {
+    out_dims[i] *= repeat_times_data[i];
+    new_x.Resize(make_ddim(vec_x_dims));
+    std::vector<const DenseTensor*> ins = {&new_x};
+    vec_x_dims[i] *= repeat_times_data[i];
+    if (i != repeat_times_data.size() - 1) {
+      if (repeat_times_data[i] != 1) {
+        DenseTensor tmp_out;
+        tmp_out.Resize(make_ddim(vec_x_dims));
+        dev_ctx.template Alloc<T>(&tmp_out);
+        std::vector<DenseTensor*> outs = {&tmp_out};
+        phi::funcs::BroadcastKernel<ElementwiseType::kUnary, T, T>(
+            dev_ctx, ins, &outs, i, kps::IdentityFunctor<T>());
+        tmp_out.Resize(out_dims);
+        new_x = tmp_out;
+      }
+      vec_x_dims[i] *= vec_x_dims[i + 1];
+      vec_x_dims[i + 1] = 1;
+    } else {
+      out->Resize(make_ddim(vec_x_dims));
+      dev_ctx.template Alloc<T>(out);
+      std::vector<DenseTensor*> outs = {out};
+      phi::funcs::BroadcastKernel<ElementwiseType::kUnary, T, T>(
+          dev_ctx, ins, &outs, i, kps::IdentityFunctor<T>());
+      out->Resize(out_dims);
+    }
+  }
+}
+
+}  // namespace phi
 
 PD_REGISTER_KERNEL(tile,
                    GPU,