chore: import upstream snapshot with attribution

paddle/phi/kernels/sparse/gpu/sum_kernel.cu

@@ -0,0 +1,461 @@
+// 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/sparse/unary_kernel.h"
+
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/backends/gpu/gpu_launch_config.h"
+#include "paddle/phi/backends/gpu/gpu_primitives.h"
+#include "paddle/phi/common/scalar.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/core/visit_type.h"
+#include "paddle/phi/kernels/cast_kernel.h"
+#include "paddle/phi/kernels/cum_kernel.h"
+#include "paddle/phi/kernels/empty_kernel.h"
+#include "paddle/phi/kernels/funcs/elementwise_base.h"
+#include "paddle/phi/kernels/index_select_kernel.h"
+#include "paddle/phi/kernels/reduce_sum_kernel.h"
+#include "paddle/phi/kernels/reshape_kernel.h"
+#include "paddle/phi/kernels/sparse/empty_kernel.h"
+#include "paddle/phi/kernels/sparse/sparse_utils_kernel.h"
+
+namespace phi {
+namespace sparse {
+
+template <typename T, typename IntT>
+__global__ void SumCooCudaKernel(const IntT* x_indices_data,
+                                 const T* x_values_data,
+                                 const int64_t x_nnz,
+                                 const int64_t dense_dim,
+                                 const int64_t sparse_dim,
+                                 const int64_t axis,
+                                 const bool keep_dim,
+                                 IntT* out_indices_data,
+                                 T* out_values_data) {
+  CUDA_KERNEL_LOOP_TYPE(index_i, x_nnz, int64_t) {
+    int64_t i = 0;
+    for (int j = 0; j < dense_dim; ++j) {
+      out_values_data[j + index_i * dense_dim] = 0;
+    }
+
+    int64_t _index_j_ =
+        static_cast<int64_t>(blockIdx.y) * blockDim.y + threadIdx.y;
+    for (auto index_j = _index_j_; index_j < x_nnz;
+         index_j += static_cast<int64_t>(blockDim.y) * gridDim.y) {
+      // Determine whether the index_i and index_j elements have the same
+      // indices in all dimensions except for the specified axis dimension.
+      bool same = true;
+      for (int j = 0; j < sparse_dim + !keep_dim; ++j) {
+        if (j != axis && x_indices_data[index_i + j * x_nnz] !=
+                             x_indices_data[index_j + j * x_nnz]) {
+          same = false;
+          break;
+        }
+      }
+      if (same) {
+        for (int j = 0; j < dense_dim; ++j) {
+          CudaAtomicAdd(&out_values_data[j + index_i * dense_dim],
+                        x_values_data[j + index_j * dense_dim]);
+        }
+      }
+    }
+    if (_index_j_ != 0) {
+      return;
+    }
+    if (keep_dim) {
+      for (int j = 0; j < sparse_dim; ++j) {
+        if (j == axis) {
+          out_indices_data[index_i + j * x_nnz] = 0;
+        } else {
+          out_indices_data[index_i + j * x_nnz] =
+              x_indices_data[index_i + j * x_nnz];
+        }
+      }
+      return;
+    }
+    for (int j = 0; j < sparse_dim; ++j) {
+      // out_indices_data [sparse_dim, x.nnz()]
+      int64_t x_indices_data_offset;
+      if (j < axis) {
+        x_indices_data_offset = index_i + j * x_nnz;
+      } else {
+        x_indices_data_offset = index_i + (j + 1) * x_nnz;
+      }
+      out_indices_data[index_i + j * x_nnz] =
+          x_indices_data[x_indices_data_offset];
+    }
+  }
+}
+
+__global__ void SumAllCsrCudaKernel(int64_t* out_crows_data,
+                                    int64_t* out_cols_data) {
+  CUDA_KERNEL_LOOP_TYPE(index, 2, int64_t) {
+    out_crows_data[index] = index;
+    if (index == 0) {
+      out_cols_data[0] = 0;
+    }
+  }
+}
+
+template <typename T>
+__global__ void SumCsr2DCudaKernel(const int64_t* x_crows_data,
+                                   const T* x_values_data,
+                                   const int64_t x_dim0,
+                                   int64_t* out_crows_data,
+                                   int64_t* out_cols_data,
+                                   T* out_values_data) {
+  CUDA_KERNEL_LOOP_TYPE(index, x_dim0 + 1, int64_t) {
+    out_crows_data[index] = index;
+    if (index != x_dim0) {
+      out_cols_data[index] = 0;
+      T sum_value = 0;
+      for (auto j = x_crows_data[index]; j < x_crows_data[index + 1]; ++j) {
+        sum_value += x_values_data[j];
+      }
+      out_values_data[index] = sum_value;
+    }
+  }
+}
+
+template <typename T>
+__global__ void SumCsr3DCudaKernel(const int64_t* x_crows_data,
+                                   const T* x_values_data,
+                                   const int64_t x_dim0,
+                                   const int64_t x_dim1,
+                                   const int64_t* batch_nnz_data,
+                                   int64_t* out_crows_data,
+                                   int64_t* out_cols_data,
+                                   T* out_values_data) {
+  {
+    CUDA_KERNEL_LOOP_TYPE(index, x_dim0 * x_dim1, int64_t) {
+      out_cols_data[index] = 0;
+    }
+  }
+
+  CUDA_KERNEL_LOOP_TYPE(index, x_dim0 * (x_dim1 + 1), int64_t) {
+    int64_t batch = index / (x_dim1 + 1);
+    int64_t number = index % (x_dim1 + 1);
+    out_crows_data[index] = number;
+
+    if (number != x_dim1) {
+      T sum_value = 0;
+      int64_t x_values_data_offset;
+      if (batch == 0) {
+        x_values_data_offset = 0;
+      } else {
+        x_values_data_offset = batch_nnz_data[batch - 1];
+      }
+      for (int64_t j = x_crows_data[index]; j < x_crows_data[index + 1]; ++j) {
+        sum_value += x_values_data[j + x_values_data_offset];
+      }
+
+      // `index - batch` would never exceed x_dim0 * x_dim1.
+      out_values_data[index - batch] = sum_value;
+    }
+  }
+}
+
+template <typename T, typename IntT, typename Context>
+void SumCooGPU0Kernel(const Context& dev_ctx,
+                      const SparseCooTensor& x,
+                      const IntArray& axis,
+                      DataType dtype,
+                      bool keep_dim,
+                      SparseCooTensor* out) {
+  auto sparse_dim = x.sparse_dim();
+  // create out sparse tensor
+  const auto& x_dims = x.dims();
+  const auto& x_indices = x.indices();
+  const auto& x_values = x.values();
+  DDim out_dims;
+  DenseTensor out_indices;
+  DenseTensor out_values;
+  if (keep_dim) {
+    out_dims = make_ddim(std::vector<int64_t>(x_dims.size(), 1));
+    out_indices = Empty<IntT, Context>(dev_ctx, {sparse_dim, 1});
+  } else {
+    out_dims = make_ddim({1});
+    out_indices = Empty<IntT, Context>(dev_ctx, {1, 1});
+  }
+  funcs::SetConstant<Context, IntT> set_out_indices;
+  set_out_indices(dev_ctx, &out_indices, static_cast<IntT>(0));
+  out_values = phi::Sum<T>(dev_ctx, x.values(), {}, dtype, keep_dim);
+  out->SetMember(out_indices, out_values, out_dims, x.coalesced());
+}
+
+template <typename T, typename IntT, typename Context>
+void SumCooGPU1Kernel(const Context& dev_ctx,
+                      const SparseCooTensor& x,
+                      const IntArray& axis,
+                      DataType dtype,
+                      bool keep_dim,
+                      SparseCooTensor* out) {
+  auto sparse_dim = x.sparse_dim();
+  // create out sparse tensor
+  const auto& x_dims = x.dims();
+  const auto& x_indices = x.indices();
+  const auto& x_values = x.values();
+  DDim out_dims;
+  DenseTensor out_indices;
+  DenseTensor out_values;
+  auto n_dim = x.dims().size();
+  auto dim = axis[0] < 0 ? x_dims.size() + axis[0] : axis[0];
+
+  std::vector<int64_t> dims;
+  for (int i = 0; i < n_dim; ++i) {
+    if (i != dim) {
+      dims.emplace_back(x.dims()[i]);
+    } else if (keep_dim || (dim < sparse_dim && sparse_dim == 1)) {
+      dims.emplace_back(1);
+    }
+  }
+  out_dims = make_ddim(dims);
+
+  if (dim >= sparse_dim) {
+    out_indices = x_indices;
+    dim = dim - sparse_dim + 1;
+    out_values = phi::Sum<T>(dev_ctx, x.values(), {dim}, dtype, keep_dim);
+    out->SetMember(out_indices, out_values, out_dims, x.coalesced());
+    return;
+  }
+
+  // Ensure the sparse_dim is not less than 1.
+  if (sparse_dim == 1) {
+    keep_dim = true;
+  }
+  // if axis in sparse_dim and keep_dim, sparse_dim will be reduced.
+  if (!keep_dim) {
+    sparse_dim -= 1;
+  }
+
+  std::vector<int> out_values_dims;
+  out_values_dims.push_back(x.nnz());
+  for (auto i = 1; i < x.values().dims().size(); ++i) {
+    out_values_dims.push_back(static_cast<int>(x.values().dims()[i]));
+  }
+  int64_t dense_dim = std::accumulate(out_values_dims.begin() + 1,
+                                      out_values_dims.end(),
+                                      1,
+                                      std::multiplies<int64_t>());
+
+  out_indices = Empty<IntT, Context>(dev_ctx, {sparse_dim, x.nnz()});
+  out_values = Empty<T, Context>(dev_ctx, out_values_dims);
+
+  const auto* x_indices_data = x_indices.data<IntT>();
+  const auto* x_values_data = x_values.data<T>();
+  auto* out_indices_data = out_indices.data<IntT>();
+  auto* out_values_data = out_values.data<T>();
+
+  auto config = backends::gpu::GetGpuLaunchConfig2D(dev_ctx, x.nnz(), x.nnz());
+  SumCooCudaKernel<T, IntT><<<config.block_per_grid.x,
+                              config.thread_per_block.x,
+                              0,
+                              dev_ctx.stream()>>>(x_indices_data,
+                                                  x_values_data,
+                                                  x.nnz(),
+                                                  dense_dim,
+                                                  sparse_dim,
+                                                  dim,
+                                                  keep_dim,
+                                                  out_indices_data,
+                                                  out_values_data);
+  if (dtype != phi::DataType::UNDEFINED && dtype != x.dtype()) {
+    out_values = Cast<T, Context>(dev_ctx, out_values, dtype);
+  }
+  out->SetMember(out_indices, out_values, out_dims, x.coalesced());
+}
+
+template <typename T, typename Context>
+void SumCooKernel(const Context& dev_ctx,
+                  const SparseCooTensor& x,
+                  const IntArray& axis,
+                  DataType dtype,
+                  bool keep_dim,
+                  SparseCooTensor* out) {
+  const size_t n_dim = axis.size();
+  if (n_dim == 0) {
+    PD_VISIT_BASE_INTEGRAL_TYPES(x.indices().dtype(), "SumCooGPUKernel", ([&] {
+                                   SumCooGPU0Kernel<T, data_t, Context>(
+                                       dev_ctx, x, axis, dtype, keep_dim, out);
+                                 }));
+  } else {
+    PD_VISIT_BASE_INTEGRAL_TYPES(x.indices().dtype(), "SumCooGPUKernel", ([&] {
+                                   SumCooGPU1Kernel<T, data_t, Context>(
+                                       dev_ctx, x, axis, dtype, keep_dim, out);
+                                 }));
+  }
+}
+
+template <typename T, typename Context>
+void SumCsr0Kernel(const Context& dev_ctx,
+                   const SparseCsrTensor& x,
+                   const IntArray& axis,
+                   DataType dtype,
+                   bool keep_dim,
+                   SparseCsrTensor* out) {
+  auto x_dim0 = x.dims()[0];
+  auto x_dim1 = x.dims()[1];
+  const auto& x_crows = x.crows();
+  const auto& x_values = x.values();
+  const auto* x_crows_data = x_crows.data<int64_t>();
+  const auto* x_values_data = x_values.data<T>();
+
+  DenseTensor out_crows, out_cols, out_values;
+  DDim out_dims;
+  if (keep_dim && x.dims().size() == 3) {
+    out_dims = make_ddim({1, 1, 1});
+  } else {
+    out_dims = make_ddim({1, 1});
+  }
+  out_crows = Empty<int64_t, Context>(dev_ctx, {2});  // crows = [0, 1]
+  out_cols = Empty<int64_t, Context>(dev_ctx, {1});   // crows = [0]
+  auto* out_crows_data = out_crows.data<int64_t>();
+  auto* out_cols_data = out_cols.data<int64_t>();
+
+  auto config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, 2, 1);
+  SumAllCsrCudaKernel<<<config.block_per_grid.x,
+                        config.thread_per_block.x,
+                        0,
+                        dev_ctx.stream()>>>(out_crows_data, out_cols_data);
+
+  out_values = phi::Sum<T>(dev_ctx, x.values(), {}, dtype, true);
+  out->SetMember(out_crows, out_cols, out_values, out_dims);
+}
+
+template <typename T, typename Context>
+void SumCsr1Kernel(const Context& dev_ctx,
+                   const SparseCsrTensor& x,
+                   const IntArray& axis,
+                   DataType dtype,
+                   bool keep_dim,
+                   SparseCsrTensor* out) {
+  auto x_dim0 = x.dims()[0];
+  auto x_dim1 = x.dims()[1];
+  const auto& x_crows = x.crows();
+  const auto& x_values = x.values();
+  const auto* x_crows_data = x_crows.data<int64_t>();
+  const auto* x_values_data = x_values.data<T>();
+
+  DenseTensor out_crows, out_cols, out_values;
+  DDim out_dims;
+  out_crows = EmptyLike<int64_t, Context>(dev_ctx, x.crows());
+  auto* out_crows_data = out_crows.data<int64_t>();
+
+  if (x.dims().size() == 2) {
+    out_cols = Empty<int64_t, Context>(dev_ctx, {x_dim0});
+    out_values = Empty<T, Context>(dev_ctx, {x_dim0});
+    auto* out_cols_data = out_cols.data<int64_t>();
+    auto* out_values_data = out_values.data<T>();
+    out_dims = make_ddim({x_dim0, 1});
+    auto config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, x_dim0 + 1, 1);
+    SumCsr2DCudaKernel<T><<<config.block_per_grid.x,
+                            config.thread_per_block.x,
+                            0,
+                            dev_ctx.stream()>>>(x_crows_data,
+                                                x_values_data,
+                                                x_dim0,
+                                                out_crows_data,
+                                                out_cols_data,
+                                                out_values_data);
+
+  } else {
+    out_cols = Empty<int64_t, Context>(dev_ctx, {x_dim0 * x_dim1});
+    out_values = Empty<T, Context>(dev_ctx, {x_dim0 * x_dim1});
+    auto* out_cols_data = out_cols.data<int64_t>();
+    auto* out_values_data = out_values.data<T>();
+    if (keep_dim) {
+      out_dims = make_ddim({x_dim0, x_dim1, 1});
+    } else {
+      out_dims = make_ddim({x_dim0, x_dim1});
+    }
+
+    DenseTensor x_crows_reshape =
+        Reshape<int64_t, Context>(dev_ctx, x_crows, {x_dim0, x_dim1 + 1});
+    DenseTensor last_indices = Empty<int64_t, Context>(dev_ctx, {1});
+    funcs::SetConstant<Context, int64_t> set_constant;
+    set_constant(dev_ctx, &last_indices, static_cast<int64_t>(x_dim1));
+
+    DenseTensor x_crows_last = Empty<int64_t, Context>(dev_ctx, {x_dim0, 1});
+    IndexSelectKernel<int64_t, Context>(
+        dev_ctx, x_crows_reshape, last_indices, 1, &x_crows_last);
+
+    DenseTensor batch_nnz = Empty<int64_t, Context>(dev_ctx, {x_dim0, 1});
+    CumsumKernel<int64_t, Context>(
+        dev_ctx, x_crows_last, Scalar(0), false, false, false, &batch_nnz);
+    auto* batch_nnz_data = batch_nnz.data<int64_t>();
+
+    auto config = backends::gpu::GetGpuLaunchConfig1D(
+        dev_ctx, x.dims()[0] * (x.dims()[1] + 1), 1);
+    SumCsr3DCudaKernel<T><<<config.block_per_grid.x,
+                            config.thread_per_block.x,
+                            0,
+                            dev_ctx.stream()>>>(x_crows_data,
+                                                x_values_data,
+                                                x_dim0,
+                                                x_dim1,
+                                                batch_nnz_data,
+                                                out_crows_data,
+                                                out_cols_data,
+                                                out_values_data);
+  }
+  if (dtype != phi::DataType::UNDEFINED && dtype != x.dtype()) {
+    out_values = Cast<T, Context>(dev_ctx, out_values, dtype);
+  }
+  out->SetMember(out_crows, out_cols, out_values, out_dims);
+}
+
+template <typename T, typename Context>
+void SumCsrKernel(const Context& dev_ctx,
+                  const SparseCsrTensor& x,
+                  const IntArray& axis,
+                  DataType dtype,
+                  bool keep_dim,
+                  SparseCsrTensor* out) {
+  size_t n_dim = axis.size();
+  if (n_dim == 0) {
+    SumCsr0Kernel<T, Context>(dev_ctx, x, axis, dtype, keep_dim, out);
+  } else {
+    PADDLE_ENFORCE_EQ(axis[0],
+                      -1,
+                      common::errors::Unimplemented(
+                          "`axis` of SumCsrKernel only support None or -1 now."
+                          "More number will be supported in the future."));
+    SumCsr1Kernel<T, Context>(dev_ctx, x, axis, dtype, keep_dim, out);
+  }
+}
+
+}  // namespace sparse
+}  // namespace phi
+
+PD_REGISTER_KERNEL(sum_coo,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::SumCooKernel,
+                   float,
+                   double,
+                   int,
+                   int64_t) {
+  kernel->OutputAt(0).SetDataType(phi::DataType::UNDEFINED);
+}
+
+PD_REGISTER_KERNEL(sum_csr,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sparse::SumCsrKernel,
+                   float,
+                   double,
+                   int,
+                   int64_t) {
+  kernel->OutputAt(0).SetDataType(phi::DataType::UNDEFINED);
+}