wehub-resource-sync
466 lines
18 KiB
C++
466 lines
18 KiB
C++
// 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.
|
|
|
|
#pragma once
|
|
#include <thrust/functional.h>
|
|
|
|
#include <algorithm>
|
|
#include <functional>
|
|
#include <string>
|
|
|
|
#include "paddle/common/hostdevice.h"
|
|
#include "paddle/phi/backends/gpu/gpu_primitives.h"
|
|
#include "paddle/phi/core/dense_tensor.h"
|
|
#include "paddle/phi/kernels/funcs/math.h"
|
|
|
|
namespace phi {
|
|
static constexpr int kNumCUDAThreads = 512;
|
|
static constexpr int kNumMaximumNumBlocks = 4096;
|
|
static const int NTHREADS = 32;
|
|
static inline int64_t NumBlocks(const int64_t N) {
|
|
return std::min((N + kNumCUDAThreads - 1) / kNumCUDAThreads,
|
|
static_cast<int64_t>(kNumMaximumNumBlocks));
|
|
}
|
|
|
|
template <typename T>
|
|
__global__ void GPUNLLLossForward1D_no_reduce(T* out_data,
|
|
const T* x_data,
|
|
const int64_t* label_data,
|
|
const T* weight_data,
|
|
const int64_t batch_size,
|
|
const int64_t n_classes,
|
|
const int64_t ignore_index) {
|
|
CUDA_KERNEL_LOOP(i, batch_size) {
|
|
const int64_t cur_label = label_data[i];
|
|
if (cur_label == ignore_index) {
|
|
out_data[i] = 0;
|
|
continue;
|
|
}
|
|
PADDLE_ENFORCE(cur_label >= 0 && cur_label < n_classes,
|
|
"label should not be out of bounds.");
|
|
const T cur_weight = weight_data ? weight_data[cur_label] : (T)1;
|
|
out_data[i] = -x_data[i * n_classes + cur_label] * cur_weight;
|
|
}
|
|
}
|
|
|
|
template <typename T, typename AccT>
|
|
__global__ void GPUNLLLossForward1D_with_reduce(T* out_data,
|
|
T* total_weight_data,
|
|
const T* x_data,
|
|
const int64_t* label_data,
|
|
const T* weight_data,
|
|
const int64_t batch_size,
|
|
const int64_t n_classes,
|
|
const int64_t size_average,
|
|
const int64_t ignore_index) {
|
|
__shared__ T sharedInputs[NTHREADS], sharedWeights[NTHREADS];
|
|
sharedInputs[threadIdx.x] = 0;
|
|
sharedWeights[threadIdx.x] = 0;
|
|
int64_t i;
|
|
for (i = threadIdx.x; i < batch_size; i += NTHREADS) {
|
|
const auto cur_label = label_data[i];
|
|
if (cur_label != ignore_index) {
|
|
PADDLE_ENFORCE(cur_label >= 0 && cur_label < n_classes,
|
|
"label should not be out of bounds.");
|
|
const auto cur_weight = weight_data ? weight_data[cur_label] : (T)1;
|
|
sharedInputs[threadIdx.x] -=
|
|
x_data[i * n_classes + cur_label] * cur_weight;
|
|
sharedWeights[threadIdx.x] += cur_weight;
|
|
}
|
|
}
|
|
__syncthreads();
|
|
|
|
if (threadIdx.x == 0) {
|
|
*out_data = *total_weight_data = 0;
|
|
AccT output_val = 0;
|
|
AccT total_weight_val = 0;
|
|
for (i = 0; i < NTHREADS; ++i) {
|
|
output_val += sharedInputs[i];
|
|
total_weight_val += sharedWeights[i];
|
|
}
|
|
*total_weight_data = total_weight_val;
|
|
*out_data = output_val;
|
|
|
|
if (size_average && *total_weight_data != 0) {
|
|
*out_data = output_val / total_weight_val;
|
|
}
|
|
}
|
|
}
|
|
|
|
// Compute thread count: clamp(round_pow2(N/16), 32, 1024).
|
|
inline int nll_loss_threads(int64_t batch_size) {
|
|
int x = static_cast<int>((batch_size + 15) / 16);
|
|
// Round to nearest power of 2
|
|
x = std::max(x, 1);
|
|
int log2_val = 0;
|
|
int tmp = x;
|
|
while (tmp > 1) {
|
|
tmp >>= 1;
|
|
log2_val++;
|
|
}
|
|
// Round: check if x is closer to (1 << log2_val) or (1 << (log2_val + 1))
|
|
int lower = 1 << log2_val;
|
|
int upper = 1 << (log2_val + 1);
|
|
int rounded = (x - lower <= upper - x) ? lower : upper;
|
|
// Clamp to [32, 1024]
|
|
return std::min(std::max(rounded, 32), 1024);
|
|
}
|
|
|
|
// Accuracy-compatible NLL loss with tree reduction in shared memory.
|
|
template <typename T, typename AccT>
|
|
__global__ void GPUNLLLossForward1D_with_reduce_compatible(
|
|
T* out_data,
|
|
T* total_weight_data,
|
|
const T* x_data,
|
|
const int64_t* label_data,
|
|
const T* weight_data,
|
|
const int64_t batch_size,
|
|
const int64_t n_classes,
|
|
const int64_t size_average,
|
|
const int64_t ignore_index) {
|
|
// Dynamic shared memory: first nthreads AccT for loss, next nthreads AccT
|
|
// for weight
|
|
extern __shared__ char smem[];
|
|
AccT* sh_loss = reinterpret_cast<AccT*>(smem);
|
|
AccT* sh_weight = sh_loss + blockDim.x;
|
|
|
|
// Thread-strided sequential accumulation
|
|
AccT thread_loss = AccT(0);
|
|
AccT thread_weight = AccT(0);
|
|
for (int64_t i = threadIdx.x; i < batch_size; i += blockDim.x) {
|
|
const int64_t cur_label = label_data[i];
|
|
if (cur_label != ignore_index) {
|
|
PADDLE_ENFORCE(cur_label >= 0 && cur_label < n_classes,
|
|
"label should not be out of bounds.");
|
|
const AccT cur_weight =
|
|
weight_data ? static_cast<AccT>(weight_data[cur_label]) : AccT(1);
|
|
thread_loss -=
|
|
static_cast<AccT>(x_data[i * n_classes + cur_label]) * cur_weight;
|
|
thread_weight += cur_weight;
|
|
}
|
|
}
|
|
|
|
sh_loss[threadIdx.x] = thread_loss;
|
|
sh_weight[threadIdx.x] = thread_weight;
|
|
__syncthreads();
|
|
|
|
// Power-of-2 tree reduction
|
|
for (unsigned int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
|
|
if (threadIdx.x < stride) {
|
|
sh_loss[threadIdx.x] += sh_loss[threadIdx.x + stride];
|
|
sh_weight[threadIdx.x] += sh_weight[threadIdx.x + stride];
|
|
}
|
|
__syncthreads();
|
|
}
|
|
|
|
if (threadIdx.x == 0) {
|
|
if (size_average && sh_weight[0] != AccT(0)) {
|
|
*out_data = static_cast<T>(sh_loss[0] / sh_weight[0]);
|
|
} else {
|
|
*out_data = static_cast<T>(sh_loss[0]);
|
|
}
|
|
*total_weight_data = static_cast<T>(sh_weight[0]);
|
|
}
|
|
}
|
|
|
|
// Reduce N values concurrently, i.e. suppose N = 2, and there are 4 threads:
|
|
// (1, 2), (3, 4), (5, 6), (7, 8), then the return in threadVals for thread 0
|
|
// is (1 + 3 + 5 + 7, 2 + 4 + 6 + 8) = (16, 20)
|
|
//
|
|
// If smem is not used again, there is no need to __syncthreads before this
|
|
// call. However, if smem will be used, e.g., this function is called in a loop,
|
|
// then __syncthreads is needed either before or afterwards to prevent non-0
|
|
// threads overriding smem in the next loop before num-0 thread reads from it.
|
|
template <typename T, typename ReduceOp, int64_t N>
|
|
__device__ void reduceNValuesInBlock(T* smem,
|
|
T threadVals[N],
|
|
const unsigned int numVals,
|
|
ReduceOp reduceOp,
|
|
T init) {
|
|
if (numVals == 0) {
|
|
#pragma unroll
|
|
for (int64_t i = 0; i < N; ++i) {
|
|
threadVals[i] = init;
|
|
}
|
|
return;
|
|
}
|
|
|
|
// We store each of the N values contiguously, so if N = 2, all values for
|
|
// the first threadVal for each thread in the block are stored followed by
|
|
// all of the values for the second threadVal for each thread in the block
|
|
if (threadIdx.x < numVals) {
|
|
#pragma unroll
|
|
for (int64_t i = 0; i < N; ++i) {
|
|
smem[i * numVals + threadIdx.x] = threadVals[i];
|
|
}
|
|
}
|
|
__syncthreads();
|
|
|
|
// Number of lanes in the final reduction --> this is used to determine
|
|
// where to put the outputs of each of the n things we are reducing. If
|
|
// nLP = 32, then we have the 32 outputs for the first threadVal,
|
|
// followed by the 32 outputs for the second threadVal, etc.
|
|
const unsigned int numLanesParticipating = min(numVals, warpSize);
|
|
|
|
if (numVals > warpSize && ((threadIdx.x / warpSize) == 0)) {
|
|
#pragma unroll
|
|
for (int64_t i = 0; i < N; ++i) {
|
|
threadVals[i] = threadIdx.x < numVals ? threadVals[i] : init;
|
|
}
|
|
|
|
for (int64_t i = warpSize + static_cast<int64_t>(threadIdx.x); i < numVals;
|
|
i += warpSize) {
|
|
#pragma unroll
|
|
for (int64_t j = 0; j < N; ++j) {
|
|
threadVals[j] = reduceOp(threadVals[j], smem[j * numVals + i]);
|
|
}
|
|
}
|
|
|
|
#pragma unroll
|
|
for (int64_t i = 0; i < N; ++i) {
|
|
smem[i * numLanesParticipating + threadIdx.x] = threadVals[i];
|
|
}
|
|
}
|
|
__syncthreads();
|
|
|
|
if (threadIdx.x == 0) {
|
|
if (numLanesParticipating == 32) {
|
|
#pragma unroll
|
|
for (int64_t i = 0; i < N; ++i) {
|
|
#pragma unroll
|
|
for (int64_t j = 1; j < 32; ++j) {
|
|
threadVals[i] = reduceOp(threadVals[i], smem[i * 32 + j]);
|
|
}
|
|
}
|
|
} else {
|
|
#pragma unroll
|
|
for (int64_t i = 0; i < N; ++i) {
|
|
for (int64_t j = 1; j < numLanesParticipating; ++j) {
|
|
threadVals[i] = reduceOp(threadVals[i], smem[i * numVals + j]);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Block-wide reduction in shared memory helper; only threadIdx.x == 0 will
|
|
// return the reduced value
|
|
//
|
|
// If smem is not used again, there is no need to __syncthreads before this
|
|
// call. However, if smem will be used, e.g., this function is called in a loop,
|
|
// then __syncthreads is needed either before or afterwards to prevent non-0
|
|
// threads overriding smem in the next loop before num-0 thread reads from it.
|
|
template <typename T, typename ReduceOp>
|
|
__device__ T reduceBlock(T* smem,
|
|
const unsigned int numVals,
|
|
T threadVal,
|
|
ReduceOp reduceOp,
|
|
T init) {
|
|
reduceNValuesInBlock<T, ReduceOp, 1>(
|
|
smem, &threadVal, numVals, reduceOp, init);
|
|
return threadVal;
|
|
}
|
|
|
|
template <typename T>
|
|
__global__ void GPUNLLLossForward2D_no_reduce(T* out_data,
|
|
const T* x_data,
|
|
const int64_t* label_data,
|
|
const T* weight_data,
|
|
const int64_t batch_size,
|
|
const int64_t n_classes,
|
|
const int64_t in_dim2,
|
|
const int64_t in_dim3,
|
|
const int64_t ignore_index) {
|
|
const int64_t map_size = in_dim2 * in_dim3;
|
|
const int64_t sample_size = n_classes * map_size;
|
|
const int64_t out_numel = batch_size * map_size;
|
|
CUDA_KERNEL_LOOP(i, out_numel) {
|
|
const int64_t b = i % batch_size;
|
|
const int64_t h = (i / batch_size) % in_dim2;
|
|
const int64_t w = (i / (batch_size * in_dim2)) % in_dim3;
|
|
|
|
const int64_t index = b * map_size + h * in_dim3 + w;
|
|
const int64_t cur_label = label_data[index];
|
|
if (cur_label == ignore_index) {
|
|
out_data[index] = 0;
|
|
continue;
|
|
}
|
|
PADDLE_ENFORCE(cur_label >= 0 && cur_label < n_classes,
|
|
"label should not be out of bounds.");
|
|
const T cur_weight = weight_data ? weight_data[cur_label] : (T)1;
|
|
out_data[index] =
|
|
-x_data[b * sample_size + cur_label * map_size + h * in_dim3 + w] *
|
|
cur_weight;
|
|
}
|
|
}
|
|
|
|
template <typename T, typename AccT>
|
|
__global__ void GPUNLLLossForward2D_with_reduce(T* out_data,
|
|
T* total_weight_data,
|
|
const T* x_data,
|
|
const int64_t* label_data,
|
|
const T* weight_data,
|
|
const int64_t batch_size,
|
|
const int64_t n_classes,
|
|
const int64_t map_nelem,
|
|
const int64_t blocks_per_sample,
|
|
const int64_t ignore_index) {
|
|
__shared__ AccT partial_sums[kNumCUDAThreads];
|
|
int64_t i;
|
|
AccT input_sum = 0;
|
|
AccT acc_weight = 0;
|
|
*out_data = 0;
|
|
*total_weight_data = 0;
|
|
|
|
int64_t sample = static_cast<int64_t>(blockIdx.x) / blocks_per_sample;
|
|
int64_t toffset = sample * map_nelem;
|
|
int64_t ioffset = sample * map_nelem * n_classes;
|
|
int64_t step = static_cast<int64_t>(blockDim.x) * blocks_per_sample;
|
|
for (i = (blockIdx.x % blocks_per_sample) * blockDim.x + threadIdx.x;
|
|
i < map_nelem;
|
|
i += step) {
|
|
const int64_t cur_label = label_data[toffset + i];
|
|
if (cur_label != ignore_index) {
|
|
PADDLE_ENFORCE(cur_label >= 0 && cur_label < n_classes,
|
|
"label should not be out of bounds.");
|
|
const AccT cur_weight = weight_data ? weight_data[cur_label] : (T)1;
|
|
input_sum -= x_data[ioffset + i + map_nelem * cur_label] * cur_weight;
|
|
acc_weight += cur_weight;
|
|
}
|
|
}
|
|
|
|
input_sum = reduceBlock(
|
|
partial_sums, blockDim.x, input_sum, thrust::plus<AccT>(), (AccT)0);
|
|
__syncthreads();
|
|
acc_weight = reduceBlock(
|
|
partial_sums, blockDim.x, acc_weight, thrust::plus<AccT>(), (AccT)0);
|
|
|
|
if (threadIdx.x == 0) {
|
|
CudaAtomicAdd(total_weight_data, acc_weight);
|
|
CudaAtomicAdd(out_data, input_sum);
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
__global__ void GPUNLLLossForward2D_size_average(T* out_data,
|
|
T* total_weight_data) {
|
|
if (*total_weight_data != 0) {
|
|
*out_data /= *total_weight_data;
|
|
}
|
|
}
|
|
template <typename T>
|
|
__global__ void GPUNLLLossBackward1D_no_reduce(T* dx_data,
|
|
const int64_t* label_data,
|
|
const T* weight_data,
|
|
const T* dout_data,
|
|
const int64_t batch_size,
|
|
const int64_t n_classes,
|
|
const int64_t ignore_index) {
|
|
CUDA_KERNEL_LOOP(i, batch_size) {
|
|
const int64_t cur_label = label_data[i];
|
|
if (cur_label == ignore_index) {
|
|
continue;
|
|
}
|
|
const T cur_weight = weight_data ? weight_data[cur_label] : (T)1;
|
|
dx_data[i * n_classes + cur_label] = -dout_data[i] * cur_weight;
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
__global__ void GPUNLLLossBackward1D_with_reduce(T* dx_data,
|
|
const T* total_weight_data,
|
|
const int64_t* label_data,
|
|
const T* weight_data,
|
|
const T* dout_data,
|
|
const int64_t batch_size,
|
|
const int64_t n_classes,
|
|
const int64_t size_average,
|
|
const int64_t ignore_index) {
|
|
if (*total_weight_data <= 0) {
|
|
return;
|
|
}
|
|
int64_t i;
|
|
const T norm = size_average ? (T)(1 / *total_weight_data) : (T)1;
|
|
for (i = threadIdx.x; i < batch_size; i += NTHREADS) {
|
|
const int64_t cur_label = label_data[i];
|
|
if (cur_label != ignore_index) {
|
|
const T cur_weight = weight_data ? weight_data[cur_label] : (T)1;
|
|
dx_data[i * n_classes + cur_label] = -cur_weight * dout_data[0] * norm;
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
__global__ void GPUNLLLossBackward2D_no_reduce(T* dx_data,
|
|
const int64_t* label_data,
|
|
const T* weight_data,
|
|
const T* dout_data,
|
|
const int64_t batch_size,
|
|
const int64_t n_classes,
|
|
const int64_t in_dim2,
|
|
const int64_t in_dim3,
|
|
const int64_t ignore_index) {
|
|
const int64_t map_size = in_dim2 * in_dim3;
|
|
const int64_t sample_size = n_classes * map_size;
|
|
const int64_t out_numel = batch_size * map_size;
|
|
CUDA_KERNEL_LOOP(i, out_numel) {
|
|
const int64_t b = i % batch_size;
|
|
const int64_t h = (i / batch_size) % in_dim2;
|
|
const int64_t w = (i / (batch_size * in_dim2)) % in_dim3;
|
|
const int64_t index = b * map_size + h * in_dim3 + w;
|
|
const int64_t cur_label = label_data[index];
|
|
if (cur_label == ignore_index) {
|
|
continue;
|
|
}
|
|
const T cur_weight = weight_data ? weight_data[cur_label] : (T)1;
|
|
dx_data[b * sample_size + cur_label * map_size + h * in_dim3 + w] =
|
|
-dout_data[index] * cur_weight;
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
__global__ void GPUNLLLossBackward2D_with_reduce(
|
|
T* dx_data,
|
|
const T* total_weight_data,
|
|
const int64_t* label_data,
|
|
const T* weight_data,
|
|
const T* dout_data,
|
|
const int64_t batch_size,
|
|
const int64_t n_classes,
|
|
const int64_t map_nelem,
|
|
const int64_t blocks_per_sample,
|
|
const int64_t size_average,
|
|
const int64_t ignore_index) {
|
|
if (*total_weight_data <= 0) {
|
|
return;
|
|
}
|
|
int64_t i;
|
|
const T norm = size_average ? (T)(1 / *total_weight_data) : (T)1;
|
|
int64_t sample = static_cast<int64_t>(blockIdx.x) / blocks_per_sample;
|
|
int64_t step = static_cast<int64_t>(blockDim.x) * blocks_per_sample;
|
|
int64_t toffset = sample * map_nelem;
|
|
int64_t ioffset = sample * map_nelem * n_classes;
|
|
for (i = (blockIdx.x % blocks_per_sample) * blockDim.x + threadIdx.x;
|
|
i < map_nelem;
|
|
i += step) {
|
|
const int64_t cur_label = label_data[toffset + i];
|
|
if (cur_label != ignore_index) {
|
|
dx_data[ioffset + i + map_nelem * cur_label] =
|
|
-(weight_data ? weight_data[cur_label] : (T)1) * norm * dout_data[0];
|
|
}
|
|
}
|
|
}
|
|
|
|
} // namespace phi
|