chore: import upstream snapshot with attribution

paddle/phi/kernels/gpu/c_split_kernel.cu

@@ -0,0 +1,113 @@
+// Copyright (c) 2023 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/phi/kernels/c_split_kernel.h"
+
+#include "paddle/phi/core/kernel_registry.h"
+
+namespace phi {
+
+static constexpr int64_t kNumCUDAThreads = 512;
+static constexpr int64_t kNumMaximumNumBlocks = 4096;
+
+static inline int64_t NumBlocks(const int64_t N) {
+  return std::min((N + kNumCUDAThreads - 1) / kNumCUDAThreads,
+                  kNumMaximumNumBlocks);
+}
+
+template <typename T>
+__global__ void SplitFromRank(const T* input,
+                              T* output,
+                              const int64_t rows,
+                              const int64_t columns,
+                              const int rank,
+                              const int nranks,
+                              const int64_t limit) {
+  CUDA_KERNEL_LOOP_TYPE(i, limit, int64_t) {
+    int64_t row = i / columns;
+    int64_t col = i % columns;
+
+    int64_t block = columns / nranks;
+    int64_t start = block * rank;
+    int64_t end = start + block;
+
+    if (col >= start && col < end) {
+      int64_t idx = block * row + col % block;
+      output[idx] = input[i];
+    }
+  }
+}
+
+template <typename T, typename Context>
+void CSplitKernel(const Context& dev_ctx,
+                  const DenseTensor& x,
+                  int rank,
+                  int nranks,
+                  bool use_model_parallel,
+                  DenseTensor* out) {
+  auto place = dev_ctx.GetPlace();
+
+  PADDLE_ENFORCE_GE(rank,
+                    0,
+                    common::errors::PreconditionNotMet(
+                        "The value of rank (%d) for c_split must be "
+                        "greater than or equal to 0.",
+                        rank));
+  PADDLE_ENFORCE_GE(nranks,
+                    2,
+                    common::errors::PreconditionNotMet(
+                        "The value of nranks (%d) for c_split must be "
+                        "greater than or equal to 2.",
+                        nranks));
+  PADDLE_ENFORCE_LT(rank,
+                    nranks,
+                    common::errors::PreconditionNotMet(
+                        "The value of rank (%d) for c_split must be "
+                        "less than that of nranks (%d).",
+                        rank,
+                        nranks));
+
+  auto dims = x.dims();
+  auto dims_size = dims.size();
+  // final dim
+  int64_t end_size = dims[dims_size - 1];
+
+  // remain dim
+  auto remain_ddim = slice_ddim(dims, 0, dims_size - 1);
+  int64_t remain_numel = common::product(remain_ddim);
+
+  int64_t limit = x.numel();
+  int64_t blocks = NumBlocks(limit);
+  int64_t threads = kNumCUDAThreads;
+
+  dims[dims_size - 1] /= nranks;
+  out->Resize(dims);
+  dev_ctx.template Alloc<T>(out);
+
+  SplitFromRank<T><<<blocks, threads, 0, dev_ctx.stream()>>>(
+      x.data<T>(), out->data<T>(), remain_numel, end_size, rank, nranks, limit);
+}
+
+}  // namespace phi
+
+PD_REGISTER_KERNEL(c_split,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::CSplitKernel,
+                   float,
+                   double,
+                   int,
+                   int64_t,
+                   phi::bfloat16,
+                   phi::float16) {}