chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,318 @@
|
||||
/* Copyright (c) 2022 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/softmax_grad_kernel.h"
|
||||
|
||||
#include <thrust/binary_search.h>
|
||||
#include <thrust/device_ptr.h>
|
||||
#include <thrust/equal.h>
|
||||
#include <thrust/iterator/constant_iterator.h>
|
||||
#include <thrust/iterator/discard_iterator.h>
|
||||
#include <thrust/sequence.h>
|
||||
#include <thrust/sort.h>
|
||||
#include <thrust/transform.h>
|
||||
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/core/visit_type.h"
|
||||
#include "paddle/phi/kernels/empty_kernel.h"
|
||||
#include "paddle/phi/kernels/full_kernel.h"
|
||||
#include "paddle/phi/kernels/funcs/math_cuda_utils.h"
|
||||
#include "paddle/phi/kernels/funcs/math_function.h"
|
||||
#include "paddle/phi/kernels/funcs/sparse/softmax.cu.h"
|
||||
#include "paddle/phi/kernels/softmax_grad_kernel.h"
|
||||
#include "paddle/phi/kernels/sparse/empty_kernel.h"
|
||||
|
||||
namespace phi {
|
||||
namespace sparse {
|
||||
|
||||
template <typename T, typename IntT = int>
|
||||
__global__ void SoftmaxGradGpuKernel(const IntT* out_crows,
|
||||
const T* out_values,
|
||||
const T* dout_values,
|
||||
T* dx_values,
|
||||
int row_number,
|
||||
int total_row_number) {
|
||||
// dx = (dout - sum(dout * out)) * out
|
||||
int row = blockIdx.x * blockDim.y + threadIdx.y;
|
||||
int non_zero_idx = threadIdx.x;
|
||||
if (row >= total_row_number) return;
|
||||
int cur_batch = row / row_number;
|
||||
int crow_idx = cur_batch * (row_number + 1) + (row % row_number);
|
||||
int cur_batch_offset = 0;
|
||||
for (int i = 1; i < cur_batch + 1; ++i) {
|
||||
cur_batch_offset += out_crows[i * (row_number + 1) - 1];
|
||||
}
|
||||
int row_first = cur_batch_offset + static_cast<int>(out_crows[crow_idx]);
|
||||
int row_nnz = static_cast<int>(out_crows[crow_idx + 1] - out_crows[crow_idx]);
|
||||
if (row_nnz == 0) return;
|
||||
|
||||
int kIteration = (row_nnz + warpSize - 1) / warpSize;
|
||||
|
||||
T mul_result = 0;
|
||||
for (int i = 0; i < kIteration; ++i) {
|
||||
int idx = non_zero_idx + i * warpSize;
|
||||
if (idx >= row_nnz) break;
|
||||
|
||||
mul_result += out_values[row_first + idx] * dout_values[row_first + idx];
|
||||
}
|
||||
T sum = funcs::WarpReduceSum<T>(mul_result, 0xFFFFFFFF);
|
||||
|
||||
for (int i = 0; i < kIteration; ++i) {
|
||||
int idx = non_zero_idx + i * warpSize;
|
||||
if (idx >= row_nnz) break;
|
||||
|
||||
dx_values[row_first + idx] =
|
||||
(dout_values[row_first + idx] - sum) * out_values[row_first + idx];
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void SoftmaxCsrGradKernel(const Context& dev_ctx,
|
||||
const SparseCsrTensor& out,
|
||||
const SparseCsrTensor& dout,
|
||||
int axis,
|
||||
SparseCsrTensor* dx) {
|
||||
PADDLE_ENFORCE_EQ(axis,
|
||||
-1,
|
||||
common::errors::Unimplemented(
|
||||
"SparseCsrTensor only support axis=-1 for softmax, "
|
||||
"which is faster when reading data by row (axis=-1)"));
|
||||
EmptyLikeCsrKernel<T, Context>(dev_ctx, dout, dx);
|
||||
|
||||
auto out_dim = out.dims();
|
||||
auto out_rank = out_dim.size();
|
||||
|
||||
int total_row_number = 1;
|
||||
int row_number = 1;
|
||||
for (int i = 0; i < out_rank - 1; ++i) {
|
||||
total_row_number *= out_dim[i];
|
||||
if (i == out_rank - 2) {
|
||||
row_number = out_dim[i];
|
||||
}
|
||||
}
|
||||
|
||||
dim3 grid((total_row_number + 3) / 4);
|
||||
dim3 block(32, 4);
|
||||
|
||||
PD_VISIT_BASE_INTEGRAL_TYPES(
|
||||
out.crows().dtype(), "SoftmaxCsrGradKernel", ([&] {
|
||||
SoftmaxGradGpuKernel<T, data_t><<<grid, block, 0, dev_ctx.stream()>>>(
|
||||
out.crows().data<data_t>(),
|
||||
out.values().data<T>(),
|
||||
dout.values().data<T>(),
|
||||
dx->mutable_values()->data<T>(),
|
||||
row_number,
|
||||
total_row_number);
|
||||
}));
|
||||
}
|
||||
|
||||
template <typename T, typename IntT>
|
||||
__global__ void SoftmaxCooGradGPURawKernel(IntT* sorted_pool_indices,
|
||||
IntT size,
|
||||
IntT* pool_sizes,
|
||||
IntT* pool_offsets,
|
||||
IntT nvalues,
|
||||
IntT grad_nnz,
|
||||
IntT* grad_offsets,
|
||||
IntT* out_offsets,
|
||||
IntT* lower_bound_values,
|
||||
T* values,
|
||||
T* out_values,
|
||||
T* grad_values,
|
||||
int total_rows) {
|
||||
int row = blockIdx.x * blockDim.y + threadIdx.y;
|
||||
if (row >= total_rows) return;
|
||||
|
||||
int tid = threadIdx.x;
|
||||
int index = row / nvalues;
|
||||
int nval = row % nvalues;
|
||||
IntT offset = pool_offsets[index];
|
||||
IntT* pool_indices = sorted_pool_indices + offset;
|
||||
IntT pool_indices_size = pool_sizes[index];
|
||||
|
||||
int kIteration = (pool_indices_size + warpSize - 1) / warpSize;
|
||||
T mul_result = 0;
|
||||
for (int k = 0; k < kIteration; ++k) {
|
||||
int idx = tid + k * warpSize;
|
||||
if (idx >= pool_indices_size) break;
|
||||
|
||||
auto i = pool_indices[idx];
|
||||
auto cur_out_value = out_values + i * nvalues;
|
||||
auto j = lower_bound_values[i];
|
||||
if (j < grad_nnz && (out_offsets[i] == grad_offsets[j])) {
|
||||
auto cur_grad_value = grad_values + j * nvalues;
|
||||
mul_result += (*(cur_out_value + nval)) * (*(cur_grad_value + nval));
|
||||
}
|
||||
}
|
||||
T sum = funcs::WarpReduceSum<T>(mul_result, 0xFFFFFFFF);
|
||||
|
||||
for (int k = 0; k < kIteration; ++k) {
|
||||
int idx = tid + k * warpSize;
|
||||
if (idx >= pool_indices_size) break;
|
||||
|
||||
auto i = pool_indices[idx];
|
||||
auto j = lower_bound_values[i];
|
||||
auto cur_out_value = out_values + i * nvalues;
|
||||
auto cur_value = values + i * nvalues;
|
||||
auto cur_grad_value = grad_values + j * nvalues;
|
||||
if (j < grad_nnz && (out_offsets[i] == grad_offsets[j])) {
|
||||
cur_value[nval] =
|
||||
(*(cur_out_value + nval)) * (*(cur_grad_value + nval) - sum);
|
||||
} else {
|
||||
cur_value[nval] = -(*(cur_out_value + nval)) * sum;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename IntT, typename Context>
|
||||
void SoftmaxCooGradGPUKernel(const Context& dev_ctx,
|
||||
const SparseCooTensor& out,
|
||||
const SparseCooTensor& dout,
|
||||
int axis,
|
||||
SparseCooTensor* dx) {
|
||||
using thrust_ptr = thrust::device_ptr<IntT>;
|
||||
auto out_indices = out.indices();
|
||||
auto out_values = out.values();
|
||||
auto out_values_ptr = out_values.data<T>();
|
||||
const auto output_indices_dims = out.indices().dims();
|
||||
const auto out_dims = out.dims();
|
||||
auto sparse_dim = out.sparse_dim();
|
||||
auto sizes = vectorize<IntT>(out_dims);
|
||||
auto grad_indices = dout.indices();
|
||||
auto grad_values = dout.values();
|
||||
auto grad_values_ptr = grad_values.data<T>();
|
||||
auto out_nnz = out.nnz();
|
||||
auto grad_nnz = dout.nnz();
|
||||
auto place = dev_ctx.GetPlace();
|
||||
auto stream = dev_ctx.stream();
|
||||
|
||||
*(dx->mutable_indices()) = out_indices;
|
||||
DenseTensor* values = dx->mutable_values();
|
||||
values->Resize(out_dims);
|
||||
values->set_meta(out_values.meta());
|
||||
dev_ctx.template Alloc<T>(values);
|
||||
funcs::SetConstant<GPUContext, T> set_zero;
|
||||
set_zero(dev_ctx, values, static_cast<T>(0.0f));
|
||||
|
||||
DenseTensor out_offsets = funcs::sparse::GetOffsets<IntT, Context>(
|
||||
dev_ctx, out_indices, sizes, static_cast<IntT>(-1));
|
||||
auto out_offsets_ptr = out_offsets.data<IntT>();
|
||||
DenseTensor grad_offsets = funcs::sparse::GetOffsets<IntT, Context>(
|
||||
dev_ctx, grad_indices, sizes, static_cast<IntT>(-1));
|
||||
auto grad_offsets_ptr = grad_offsets.data<IntT>();
|
||||
|
||||
#ifdef PADDLE_WITH_HIP
|
||||
const auto& policy = thrust::hip::par.on(dev_ctx.stream());
|
||||
bool is_same_offset = thrust::equal(thrust::hip::par.on(dev_ctx.stream()),
|
||||
#else
|
||||
const auto& policy = thrust::cuda::par.on(dev_ctx.stream());
|
||||
bool is_same_offset = thrust::equal(thrust::cuda::par.on(dev_ctx.stream()),
|
||||
#endif
|
||||
out_offsets_ptr,
|
||||
out_offsets_ptr + out_offsets.numel(),
|
||||
grad_offsets_ptr);
|
||||
|
||||
int dim = axis < 0 ? out_dims.size() + axis : axis;
|
||||
if (dim >= sparse_dim) {
|
||||
PADDLE_ENFORCE_EQ(
|
||||
is_same_offset,
|
||||
true,
|
||||
common::errors::Unimplemented(
|
||||
"SparseCooTensor only support same offsets for softmax."));
|
||||
SoftmaxGradKernel<T, Context>(
|
||||
dev_ctx, out_values, grad_values, dim - sparse_dim + 1, values);
|
||||
return;
|
||||
}
|
||||
|
||||
auto nnz = out.nnz();
|
||||
IntT nvalues = std::accumulate(sizes.begin() + sparse_dim,
|
||||
sizes.end(),
|
||||
static_cast<IntT>(1),
|
||||
std::multiplies<>());
|
||||
|
||||
DenseTensor values_2(*values);
|
||||
values_2.Resize({nnz, nvalues});
|
||||
|
||||
DenseTensor sorted_indices;
|
||||
DenseTensor pool_offsets;
|
||||
DenseTensor pool_sizes;
|
||||
|
||||
std::tie(sorted_indices, pool_offsets, pool_sizes, std::ignore) =
|
||||
funcs::sparse::ComputePoolMax<T, IntT, Context, false>(
|
||||
dev_ctx, out_indices, values_2, sizes, nvalues, dim);
|
||||
|
||||
DenseTensor bound =
|
||||
Empty<IntT>(dev_ctx, {static_cast<IntT>(out_offsets.dims()[0])});
|
||||
IntT* bound_ptr = bound.data<IntT>();
|
||||
thrust::lower_bound(policy,
|
||||
thrust_ptr(grad_offsets_ptr),
|
||||
thrust_ptr(grad_offsets_ptr + grad_offsets.dims()[0]),
|
||||
thrust_ptr(out_offsets_ptr),
|
||||
thrust_ptr(out_offsets_ptr) + out_offsets.dims()[0],
|
||||
thrust_ptr(bound.data<IntT>()));
|
||||
|
||||
auto pool_size = pool_offsets.dims()[0];
|
||||
int total_rows = pool_size * nvalues;
|
||||
dim3 grid((total_rows + 15) / 16);
|
||||
dim3 block(32, 16);
|
||||
SoftmaxCooGradGPURawKernel<T, IntT>
|
||||
<<<grid, block, 0, stream>>>(sorted_indices.data<IntT>(),
|
||||
pool_size,
|
||||
pool_sizes.data<IntT>(),
|
||||
pool_offsets.data<IntT>(),
|
||||
nvalues,
|
||||
grad_nnz,
|
||||
grad_offsets.data<IntT>(),
|
||||
out_offsets.data<IntT>(),
|
||||
bound_ptr,
|
||||
values_2.data<T>(),
|
||||
out_values.data<T>(),
|
||||
grad_values.data<T>(),
|
||||
total_rows);
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void SoftmaxCooGradKernel(const Context& dev_ctx,
|
||||
const SparseCooTensor& out,
|
||||
const SparseCooTensor& dout,
|
||||
int axis,
|
||||
SparseCooTensor* dx) {
|
||||
PD_VISIT_BASE_INTEGRAL_TYPES(
|
||||
out.indices().dtype(), "SoftmaxCooGradGPUKernel", ([&] {
|
||||
SoftmaxCooGradGPUKernel<T, data_t, Context>(
|
||||
dev_ctx, out, dout, axis, dx);
|
||||
}));
|
||||
}
|
||||
|
||||
} // namespace sparse
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(softmax_csr_grad,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::sparse::SoftmaxCsrGradKernel,
|
||||
float,
|
||||
double) {
|
||||
kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR);
|
||||
}
|
||||
|
||||
PD_REGISTER_KERNEL(softmax_coo_grad,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::sparse::SoftmaxCooGradKernel,
|
||||
float,
|
||||
double) {
|
||||
kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
|
||||
}
|
||||
Reference in New Issue
Block a user