wehub-resource-sync
376 lines
16 KiB
Plaintext
376 lines
16 KiB
Plaintext
// 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) {}
|