chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,375 @@
|
||||
// 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_grad_kernel.h"
|
||||
#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/memory_utils.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/kernels/empty_kernel.h"
|
||||
#include "paddle/phi/kernels/funcs/slice_utils.h"
|
||||
|
||||
namespace phi {
|
||||
namespace sparse {
|
||||
|
||||
template <typename T>
|
||||
__global__ void GetCooInputGradCudaKernel(const int64_t* out_grad_indices_data,
|
||||
const T* out_grad_values_data,
|
||||
const int64_t* axes,
|
||||
const int64_t* starts,
|
||||
const int64_t axes_size,
|
||||
const int64_t sparse_dim,
|
||||
const int64_t out_grad_nnz,
|
||||
int64_t* dx_indices_data,
|
||||
T* dx_values_data) {
|
||||
CUDA_KERNEL_LOOP_TYPE(j, out_grad_nnz, int64_t) {
|
||||
// set indices
|
||||
for (int64_t i = 0; i < sparse_dim; ++i) {
|
||||
dx_indices_data[i * out_grad_nnz + j] =
|
||||
out_grad_indices_data[i * out_grad_nnz + j];
|
||||
}
|
||||
for (size_t ii = 0; ii < axes_size; ++ii) {
|
||||
int64_t i = axes[ii];
|
||||
dx_indices_data[i * out_grad_nnz + j] += starts[ii];
|
||||
}
|
||||
// set value
|
||||
dx_values_data[j] = out_grad_values_data[j];
|
||||
}
|
||||
}
|
||||
template <typename T, typename Context>
|
||||
void SliceCooGradCompute(const Context& dev_ctx,
|
||||
const SparseCooTensor& x,
|
||||
const SparseCooTensor& out_grad,
|
||||
const std::vector<int64_t>& axes,
|
||||
const std::vector<int64_t>& starts,
|
||||
const std::vector<int64_t>& ends,
|
||||
SparseCooTensor* x_grad) {
|
||||
const DDim& x_dims = x.dims();
|
||||
|
||||
// copy axes to device
|
||||
auto d_axes_tensor = memory_utils::Alloc(
|
||||
dev_ctx.GetPlace(),
|
||||
sizeof(int64_t) * axes.size(),
|
||||
phi::Stream(reinterpret_cast<phi::StreamId>(dev_ctx.stream())));
|
||||
int64_t* d_axes = reinterpret_cast<int64_t*>(d_axes_tensor->ptr());
|
||||
memory_utils::Copy(dev_ctx.GetPlace(),
|
||||
d_axes,
|
||||
phi::CPUPlace(),
|
||||
axes.data(),
|
||||
sizeof(int64_t) * axes.size(),
|
||||
dev_ctx.stream());
|
||||
|
||||
// copy starts to device
|
||||
auto d_starts_tensor = memory_utils::Alloc(
|
||||
dev_ctx.GetPlace(),
|
||||
sizeof(int64_t) * starts.size(),
|
||||
phi::Stream(reinterpret_cast<phi::StreamId>(dev_ctx.stream())));
|
||||
int64_t* d_starts = reinterpret_cast<int64_t*>(d_starts_tensor->ptr());
|
||||
memory_utils::Copy(dev_ctx.GetPlace(),
|
||||
d_starts,
|
||||
phi::CPUPlace(),
|
||||
starts.data(),
|
||||
sizeof(int64_t) * starts.size(),
|
||||
dev_ctx.stream());
|
||||
|
||||
// Step2: Set indices and values of x_grad
|
||||
const int64_t out_grad_nnz = out_grad.nnz();
|
||||
auto sparse_dim = static_cast<int64_t>(out_grad.sparse_dim());
|
||||
DenseTensor dx_indices =
|
||||
Empty<int64_t, Context>(dev_ctx, {sparse_dim, out_grad_nnz});
|
||||
DenseTensor dx_values = Empty<T, Context>(dev_ctx, {out_grad_nnz});
|
||||
auto* dx_indices_data = dx_indices.data<int64_t>();
|
||||
auto* dx_values_data = dx_values.data<T>();
|
||||
|
||||
const auto* out_grad_indices_data = out_grad.indices().data<int64_t>();
|
||||
const auto* out_grad_values_data = out_grad.values().data<T>();
|
||||
|
||||
x_grad->SetMember(dx_indices, dx_values, x.dims(), x.coalesced());
|
||||
|
||||
auto config =
|
||||
backends::gpu::GetGpuLaunchConfig1D(dev_ctx, out_grad_nnz + 1, 1);
|
||||
GetCooInputGradCudaKernel<T><<<config.block_per_grid.x,
|
||||
config.thread_per_block.x,
|
||||
0,
|
||||
dev_ctx.stream()>>>(out_grad_indices_data,
|
||||
out_grad_values_data,
|
||||
d_axes,
|
||||
d_starts,
|
||||
axes.size(),
|
||||
sparse_dim,
|
||||
out_grad_nnz,
|
||||
dx_indices_data,
|
||||
dx_values_data);
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void SliceCooGradKernel(const Context& dev_ctx,
|
||||
const SparseCooTensor& x,
|
||||
const SparseCooTensor& out_grad,
|
||||
const phi::IntArray& axes,
|
||||
const phi::IntArray& starts,
|
||||
const phi::IntArray& ends,
|
||||
SparseCooTensor* x_grad) {
|
||||
const DDim& x_dims = x.dims();
|
||||
std::vector<int64_t> axes_vec = axes.GetData();
|
||||
std::vector<int64_t> starts_vec = starts.GetData();
|
||||
std::vector<int64_t> ends_vec = ends.GetData();
|
||||
// Check and update sparse slice attrs
|
||||
funcs::CheckAndUpdateSparseSliceAttrs<int64_t>(
|
||||
x_dims, &axes_vec, &starts_vec, &ends_vec);
|
||||
|
||||
SliceCooGradCompute<T, Context>(
|
||||
dev_ctx, x, out_grad, axes_vec, starts_vec, ends_vec, x_grad);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__global__ void GetCsrInputColsValuesCudaKernel(
|
||||
const int64_t* out_grad_cols_data,
|
||||
const T* out_grad_values_data,
|
||||
const int64_t out_grad_nnz,
|
||||
const int64_t cols_start,
|
||||
int64_t* dx_cols_data,
|
||||
T* dx_values_data) {
|
||||
CUDA_KERNEL_LOOP_TYPE(i, out_grad_nnz, int64_t) {
|
||||
dx_cols_data[i] = out_grad_cols_data[i] + cols_start;
|
||||
dx_values_data[i] = out_grad_values_data[i];
|
||||
}
|
||||
}
|
||||
|
||||
__global__ void GetCsrInputCrowsCudaKernel(
|
||||
const int64_t* out_grad_crows_data,
|
||||
const int64_t out_grad_n_rows,
|
||||
const int64_t out_grad_nnz,
|
||||
const int64_t x_n_rows,
|
||||
const int64_t rows_start,
|
||||
const int64_t rows_end,
|
||||
int64_t* dx_crows_data,
|
||||
const int64_t dx_crows_offset = 0,
|
||||
const int64_t out_grad_crows_offset = 0) {
|
||||
CUDA_KERNEL_LOOP_TYPE(i, x_n_rows + 1, int64_t) {
|
||||
int64_t idx = i + dx_crows_offset;
|
||||
if (i < rows_start) {
|
||||
dx_crows_data[idx] = 0;
|
||||
} else if (i < rows_start + out_grad_n_rows + 1) {
|
||||
int64_t out_grad_idx = out_grad_crows_offset + (i - rows_start);
|
||||
dx_crows_data[idx] = out_grad_crows_data[out_grad_idx];
|
||||
} else {
|
||||
int64_t out_grad_idx = out_grad_crows_offset + out_grad_n_rows;
|
||||
dx_crows_data[idx] = out_grad_crows_data[out_grad_idx];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void SliceCsrGrad2D(const Context& dev_ctx,
|
||||
const SparseCsrTensor& x,
|
||||
const SparseCsrTensor& out_grad,
|
||||
const std::vector<int64_t>& axes,
|
||||
const std::vector<int64_t>& starts,
|
||||
const std::vector<int64_t>& ends,
|
||||
SparseCsrTensor* x_grad) {
|
||||
const int64_t out_grad_nnz = out_grad.nnz();
|
||||
const int64_t n_rows = x.dims()[0];
|
||||
const auto* out_grad_crows_data = out_grad.crows().data<int64_t>();
|
||||
const auto* out_grad_cols_data = out_grad.cols().data<int64_t>();
|
||||
const auto* out_grad_values_data = out_grad.values().data<T>();
|
||||
|
||||
DenseTensor dx_crows = Empty<int64_t>(dev_ctx, {n_rows + 1});
|
||||
DenseTensor dx_cols = Empty<int64_t>(dev_ctx, {out_grad_nnz});
|
||||
DenseTensor dx_values = Empty<T, Context>(dev_ctx, {out_grad_nnz});
|
||||
auto* dx_crows_data = dx_crows.data<int64_t>();
|
||||
auto* dx_cols_data = dx_cols.data<int64_t>();
|
||||
auto* dx_values_data = dx_values.data<T>();
|
||||
x_grad->SetMember(dx_crows, dx_cols, dx_values, x.dims());
|
||||
|
||||
// set cols and values
|
||||
auto config =
|
||||
backends::gpu::GetGpuLaunchConfig1D(dev_ctx, out_grad_nnz + 1, 1);
|
||||
GetCsrInputColsValuesCudaKernel<T><<<config.block_per_grid.x,
|
||||
config.thread_per_block.x,
|
||||
0,
|
||||
dev_ctx.stream()>>>(out_grad_cols_data,
|
||||
out_grad_values_data,
|
||||
out_grad_nnz,
|
||||
starts[1],
|
||||
dx_cols_data,
|
||||
dx_values_data);
|
||||
config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, n_rows + 1, 1);
|
||||
GetCsrInputCrowsCudaKernel<<<config.block_per_grid.x,
|
||||
config.thread_per_block.x,
|
||||
0,
|
||||
dev_ctx.stream()>>>(out_grad_crows_data,
|
||||
out_grad.dims()[0],
|
||||
out_grad_nnz,
|
||||
x.dims()[0],
|
||||
starts[0],
|
||||
ends[0],
|
||||
dx_crows_data,
|
||||
0,
|
||||
0);
|
||||
}
|
||||
|
||||
__global__ void GetCsrInputCrowsPart1CudaKernel(const int64_t n_rows,
|
||||
const int64_t dim0_idx,
|
||||
int64_t* dx_crows_data) {
|
||||
CUDA_KERNEL_LOOP_TYPE(j, n_rows + 1, int64_t) {
|
||||
dx_crows_data[dim0_idx * (n_rows + 1) + j] = 0;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void SliceCsrGrad3D(const Context& dev_ctx,
|
||||
const SparseCsrTensor& x,
|
||||
const SparseCsrTensor& out_grad,
|
||||
const std::vector<int64_t>& axes,
|
||||
const std::vector<int64_t>& starts,
|
||||
const std::vector<int64_t>& ends,
|
||||
SparseCsrTensor* x_grad) {
|
||||
const int64_t dim0 = x.dims()[0], n_rows = x.dims()[1];
|
||||
const int64_t out_grad_nnz = out_grad.nnz();
|
||||
const auto* out_grad_crows_data = out_grad.crows().data<int64_t>();
|
||||
const auto* out_grad_cols_data = out_grad.cols().data<int64_t>();
|
||||
const auto* out_grad_values_data = out_grad.values().data<T>();
|
||||
|
||||
DenseTensor dx_crows = Empty<int64_t>(dev_ctx, {dim0 * (n_rows + 1)});
|
||||
DenseTensor dx_cols = Empty<int64_t>(dev_ctx, {out_grad_nnz});
|
||||
DenseTensor dx_values = Empty<T, Context>(dev_ctx, {out_grad_nnz});
|
||||
auto* dx_crows_data = dx_crows.data<int64_t>();
|
||||
auto* dx_cols_data = dx_cols.data<int64_t>();
|
||||
auto* dx_values_data = dx_values.data<T>();
|
||||
x_grad->SetMember(dx_crows, dx_cols, dx_values, x.dims());
|
||||
|
||||
// set cols and values
|
||||
auto config =
|
||||
backends::gpu::GetGpuLaunchConfig1D(dev_ctx, out_grad_nnz + 1, 1);
|
||||
GetCsrInputColsValuesCudaKernel<T><<<config.block_per_grid.x,
|
||||
config.thread_per_block.x,
|
||||
0,
|
||||
dev_ctx.stream()>>>(out_grad_cols_data,
|
||||
out_grad_values_data,
|
||||
out_grad_nnz,
|
||||
starts[2],
|
||||
dx_cols_data,
|
||||
dx_values_data);
|
||||
// set crows
|
||||
int64_t out_grad_n_rows = out_grad.dims()[1];
|
||||
for (int64_t i = 0; i < dim0; ++i) {
|
||||
if (i < starts[0] || i >= ends[0]) {
|
||||
config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, n_rows + 1, 1);
|
||||
GetCsrInputCrowsPart1CudaKernel<<<config.block_per_grid.x,
|
||||
config.thread_per_block.x,
|
||||
0,
|
||||
dev_ctx.stream()>>>(
|
||||
n_rows, i, dx_crows_data);
|
||||
} else {
|
||||
int64_t dx_crows_offset = i * (n_rows + 1);
|
||||
int64_t out_grad_crows_offset = (i - starts[0]) * (out_grad_n_rows + 1);
|
||||
config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, n_rows + 1, 1);
|
||||
GetCsrInputCrowsCudaKernel<<<config.block_per_grid.x,
|
||||
config.thread_per_block.x,
|
||||
0,
|
||||
dev_ctx.stream()>>>(out_grad_crows_data,
|
||||
out_grad_n_rows,
|
||||
out_grad_nnz,
|
||||
n_rows,
|
||||
starts[1],
|
||||
ends[1],
|
||||
dx_crows_data,
|
||||
dx_crows_offset,
|
||||
out_grad_crows_offset);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void SliceCsrGradCompute(const Context& dev_ctx,
|
||||
const SparseCsrTensor& x,
|
||||
const SparseCsrTensor& out_grad,
|
||||
const std::vector<int64_t>& axes,
|
||||
const std::vector<int64_t>& starts,
|
||||
const std::vector<int64_t>& ends,
|
||||
SparseCsrTensor* x_grad) {
|
||||
const DDim& x_dims = x.dims();
|
||||
|
||||
// construct new axes, starts, and ends
|
||||
std::vector<int64_t> new_axes(3), new_starts(3), new_ends(3);
|
||||
funcs::ConstructNewSliceAttrs(
|
||||
x_dims, axes, starts, ends, &new_axes, &new_starts, &new_ends);
|
||||
|
||||
const int64_t sparse_dim = x_dims.size();
|
||||
if (sparse_dim == 2) {
|
||||
SliceCsrGrad2D<T, Context>(
|
||||
dev_ctx, x, out_grad, new_axes, new_starts, new_ends, x_grad);
|
||||
} else if (sparse_dim == 3) {
|
||||
SliceCsrGrad3D<T, Context>(
|
||||
dev_ctx, x, out_grad, new_axes, new_starts, new_ends, x_grad);
|
||||
} else {
|
||||
// throw exception
|
||||
common::errors::InvalidArgument(
|
||||
"Slice grad for Sparse CSR Tensor only support 2-D or 3-D, but got "
|
||||
"%d-D.",
|
||||
x_dims.size());
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void SliceCsrGradKernel(const Context& dev_ctx,
|
||||
const SparseCsrTensor& x,
|
||||
const SparseCsrTensor& out_grad,
|
||||
const phi::IntArray& axes,
|
||||
const phi::IntArray& starts,
|
||||
const phi::IntArray& ends,
|
||||
SparseCsrTensor* x_grad) {
|
||||
const DDim& x_dims = x.dims();
|
||||
std::vector<int64_t> axes_vec = axes.GetData();
|
||||
std::vector<int64_t> starts_vec = starts.GetData();
|
||||
std::vector<int64_t> ends_vec = ends.GetData();
|
||||
// update starts and ends
|
||||
funcs::CheckAndUpdateSparseSliceAttrs<int64_t>(
|
||||
x_dims, &axes_vec, &starts_vec, &ends_vec);
|
||||
|
||||
SliceCsrGradCompute<T, Context>(
|
||||
dev_ctx, x, out_grad, axes_vec, starts_vec, ends_vec, x_grad);
|
||||
}
|
||||
|
||||
} // namespace sparse
|
||||
} // namespace phi
|
||||
PD_REGISTER_KERNEL(slice_coo_grad,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::sparse::SliceCooGradKernel,
|
||||
float,
|
||||
double,
|
||||
int8_t,
|
||||
uint8_t,
|
||||
int16_t,
|
||||
int,
|
||||
int64_t,
|
||||
bool) {}
|
||||
|
||||
PD_REGISTER_KERNEL(slice_csr_grad,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::sparse::SliceCsrGradKernel,
|
||||
float,
|
||||
double,
|
||||
int8_t,
|
||||
uint8_t,
|
||||
int16_t,
|
||||
int,
|
||||
int64_t,
|
||||
bool) {}
|
||||
Reference in New Issue
Block a user