wehub-resource-sync
355 lines
11 KiB
Plaintext
355 lines
11 KiB
Plaintext
/**
|
||
* Copyright 2017-present, Facebook, Inc.
|
||
* All rights reserved.
|
||
*
|
||
* This source code is licensed under the license found in the
|
||
* LICENSE file in the root directory of this source tree.
|
||
*/
|
||
|
||
#include <ATen/ATen.h>
|
||
#include <ATen/cuda/CUDAContext.h> // @manual=//caffe2/aten:ATen-cu
|
||
#include <cuda_runtime.h>
|
||
#include <algorithm> // std::min/max
|
||
#include <cub/cub.cuh>
|
||
|
||
#include "alignment_train_cuda.h"
|
||
#include "utils.h"
|
||
|
||
namespace {
|
||
|
||
// The thread block length in threads along the X dimension
|
||
constexpr int BLOCK_DIM_X = 128;
|
||
// The thread block length in threads along the Y dimension
|
||
constexpr int BLOCK_DIM_Y = 8;
|
||
// The thread block length in threads for scan operation
|
||
constexpr int SCAN_BLOCK = 512;
|
||
|
||
#define gpuErrchk(ans) \
|
||
{ gpuAssert((ans), __FILE__, __LINE__); }
|
||
|
||
inline void
|
||
gpuAssert(cudaError_t code, const char* file, int line, bool abort = true) {
|
||
if (code != cudaSuccess) {
|
||
fprintf(
|
||
stderr,
|
||
"\nGPUassert: %s %s %d\n",
|
||
cudaGetErrorString(code),
|
||
file,
|
||
line);
|
||
if (abort)
|
||
exit(code);
|
||
}
|
||
}
|
||
|
||
template <typename T>
|
||
struct Prod {
|
||
/// prod operator, returns <tt>a * b</tt>
|
||
__host__ __device__ __forceinline__ T
|
||
operator()(const T& a, const T& b) const {
|
||
return a * b;
|
||
}
|
||
};
|
||
|
||
template <typename T>
|
||
struct BlockPrefixProdCallbackOp {
|
||
// Running prefix
|
||
T running_total;
|
||
|
||
// Constructor
|
||
__device__ BlockPrefixProdCallbackOp(T running_total)
|
||
: running_total(running_total) {}
|
||
|
||
// Callback operator to be entered by the first warp of threads in the block.
|
||
// Thread-0 is responsible for returning a value for seeding the block-wide
|
||
// scan.
|
||
__device__ T operator()(const T block_aggregate) {
|
||
T old_prefix = running_total;
|
||
running_total *= block_aggregate;
|
||
return old_prefix;
|
||
}
|
||
};
|
||
|
||
template <typename T>
|
||
struct BlockPrefixSumCallbackOp {
|
||
// Running prefix
|
||
T running_total;
|
||
|
||
// Constructor
|
||
__device__ BlockPrefixSumCallbackOp(T running_total)
|
||
: running_total(running_total) {}
|
||
|
||
// Callback operator to be entered by the first warp of threads in the block.
|
||
// Thread-0 is responsible for returning a value for seeding the block-wide
|
||
// scan.
|
||
__device__ T operator()(const T block_aggregate) {
|
||
T old_prefix = running_total;
|
||
running_total += block_aggregate;
|
||
return old_prefix;
|
||
}
|
||
};
|
||
|
||
template <typename T>
|
||
__global__ void oneMinusPKernel(
|
||
const T* __restrict__ p_choose,
|
||
T* __restrict__ cumprod_1mp,
|
||
uint32_t bsz,
|
||
uint32_t tgt_len,
|
||
uint32_t src_len) {
|
||
for (uint32_t b = blockIdx.x; b < bsz; b += gridDim.x) {
|
||
for (uint32_t tgt = threadIdx.y; tgt < tgt_len; tgt += blockDim.y) {
|
||
for (uint32_t src = threadIdx.x; src < src_len; src += blockDim.x) {
|
||
uint32_t idx = b * tgt_len * src_len + tgt * src_len + src;
|
||
cumprod_1mp[idx] = 1 - p_choose[idx];
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
template <typename T, int TPB>
|
||
__global__ void innermostScanKernel(
|
||
T* __restrict__ cumprod_1mp,
|
||
uint32_t bsz,
|
||
uint32_t tgt_len,
|
||
uint32_t src_len) {
|
||
for (uint32_t b = blockIdx.y; b < bsz; b += gridDim.y) {
|
||
for (uint32_t tgt = blockIdx.x; tgt < tgt_len; tgt += gridDim.x) {
|
||
// Specialize BlockScan for a 1D block of TPB threads on type T
|
||
typedef cub::BlockScan<T, TPB> BlockScan;
|
||
// Allocate shared memory for BlockScan
|
||
__shared__ typename BlockScan::TempStorage temp_storage;
|
||
// Initialize running total
|
||
BlockPrefixProdCallbackOp<T> prefix_op(1);
|
||
|
||
const uint32_t tid = threadIdx.x;
|
||
for (uint32_t block_src = 0; block_src < src_len;
|
||
block_src += blockDim.x) {
|
||
uint32_t src = block_src + tid;
|
||
uint32_t idx = b * tgt_len * src_len + tgt * src_len + src;
|
||
T thread_data = (src < src_len) ? cumprod_1mp[idx] : (T)0;
|
||
|
||
// Collectively compute the block-wide inclusive prefix sum
|
||
BlockScan(temp_storage)
|
||
.ExclusiveScan(thread_data, thread_data, Prod<T>(), prefix_op);
|
||
__syncthreads();
|
||
|
||
// write the scanned value to output
|
||
if (src < src_len) {
|
||
cumprod_1mp[idx] = thread_data;
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
template <typename T>
|
||
__global__ void clampKernel(
|
||
const T* __restrict__ cumprod_1mp,
|
||
T* __restrict__ cumprod_1mp_clamp,
|
||
uint32_t bsz,
|
||
uint32_t tgt_len,
|
||
uint32_t src_len,
|
||
T min_val,
|
||
T max_val) {
|
||
for (uint32_t b = blockIdx.x; b < bsz; b += gridDim.x) {
|
||
for (uint32_t tgt = threadIdx.y; tgt < tgt_len; tgt += blockDim.y) {
|
||
for (uint32_t src = threadIdx.x; src < src_len; src += blockDim.x) {
|
||
uint32_t idx = b * tgt_len * src_len + tgt * src_len + src;
|
||
if (cumprod_1mp[idx] < min_val) {
|
||
cumprod_1mp_clamp[idx] = min_val;
|
||
} else if (cumprod_1mp[idx] > max_val) {
|
||
cumprod_1mp_clamp[idx] = max_val;
|
||
} else {
|
||
cumprod_1mp_clamp[idx] = cumprod_1mp[idx];
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
template <typename T>
|
||
__global__ void initAlphaCUDAKernel(
|
||
T* alpha,
|
||
uint32_t bsz,
|
||
uint32_t tgt_len,
|
||
uint32_t src_len) {
|
||
// alpha[:, 0, 0] = 1.0
|
||
for (uint32_t b = blockIdx.x; b < bsz; b += gridDim.x) {
|
||
alpha[b * tgt_len * src_len] = (T)1.0;
|
||
}
|
||
}
|
||
|
||
template <typename T, int TPB>
|
||
__global__ void alignmentTrainCUDAKernel(
|
||
const T* __restrict__ p_choose,
|
||
const T* __restrict__ cumprod_1mp,
|
||
const T* __restrict__ cumprod_1mp_clamp,
|
||
T* __restrict__ alpha,
|
||
uint32_t bsz,
|
||
uint32_t tgt_len,
|
||
uint32_t src_len,
|
||
uint32_t tgt) {
|
||
for (uint32_t b = blockIdx.x; b < bsz; b += gridDim.x) {
|
||
// Specialize BlockScan for a 1D block of TPB threads on type T
|
||
typedef cub::BlockScan<T, TPB> BlockScan;
|
||
|
||
// Allocate shared memory for BlockScan
|
||
__shared__ typename BlockScan::TempStorage temp_storage;
|
||
// Initialize running total
|
||
BlockPrefixSumCallbackOp<T> prefix_op(0);
|
||
|
||
uint32_t b_offset = b * tgt_len * src_len;
|
||
const uint32_t tid = threadIdx.x;
|
||
for (uint32_t block_src = 0; block_src < src_len; block_src += blockDim.x) {
|
||
uint32_t src = block_src + tid;
|
||
// Obtain a segment of consecutive items that are blocked across threads
|
||
uint32_t inout_idx, alpha_idx;
|
||
if (tgt == 0) {
|
||
// both alpha and other input index is [b][0][src]
|
||
alpha_idx = b_offset + src;
|
||
} else {
|
||
// alpha index is [b][tgt-1][src]
|
||
alpha_idx = b_offset + (tgt - 1) * src_len + src;
|
||
}
|
||
inout_idx = b_offset + tgt * src_len + src;
|
||
T thread_data = (T)0;
|
||
if (src < src_len) {
|
||
thread_data = alpha[alpha_idx] / cumprod_1mp_clamp[inout_idx];
|
||
}
|
||
|
||
// Collectively compute the block-wide inclusive prefix sum
|
||
BlockScan(temp_storage).InclusiveSum(thread_data, thread_data, prefix_op);
|
||
__syncthreads();
|
||
|
||
if (src < src_len) {
|
||
T out = thread_data * p_choose[inout_idx] * cumprod_1mp[inout_idx];
|
||
// Clamps all elements into the range [ 0, 1.0 ]
|
||
alpha[inout_idx] = std::min<T>(std::max<T>(out, 0), (T)1.0);
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
template <typename T>
|
||
void exclusiveCumprod(
|
||
const T* p_choose,
|
||
T* cumprod_1mp,
|
||
uint32_t bsz,
|
||
uint32_t tgt_len,
|
||
uint32_t src_len,
|
||
uint32_t max_grid_x,
|
||
uint32_t max_grid_y,
|
||
cudaStream_t& stream) {
|
||
// cumprod_1mp = 1 - p_choose
|
||
dim3 grid(std::min<T>(max_grid_x, bsz), 1, 1);
|
||
dim3 block(BLOCK_DIM_X, BLOCK_DIM_Y, 1);
|
||
oneMinusPKernel<T><<<grid, block, 0, stream>>>(
|
||
p_choose, cumprod_1mp, bsz, tgt_len, src_len);
|
||
gpuErrchk(cudaGetLastError());
|
||
|
||
// scan on the innermost dimension of cumprod_1mp
|
||
// cumprod_1mp = cumprod(cumprod_1mp)
|
||
dim3 grid_scan(
|
||
std::min<T>(max_grid_x, tgt_len), std::min<T>(max_grid_y, bsz), 1);
|
||
innermostScanKernel<T, SCAN_BLOCK><<<grid_scan, SCAN_BLOCK, 0, stream>>>(
|
||
cumprod_1mp, bsz, tgt_len, src_len);
|
||
gpuErrchk(cudaGetLastError());
|
||
}
|
||
|
||
template <typename T>
|
||
void alignmentTrainCUDAImpl(
|
||
const T* p_choose,
|
||
T* alpha,
|
||
uint32_t bsz,
|
||
uint32_t tgt_len,
|
||
uint32_t src_len,
|
||
float eps) {
|
||
// p_choose: bsz , tgt_len, src_len
|
||
// cumprod_1mp: bsz , tgt_len, src_len
|
||
// cumprod_1mp_clamp : bsz, tgt_len, src_len
|
||
// alpha: bsz, tgt_len, src_len
|
||
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
|
||
uint32_t max_grid_x = at::cuda::getCurrentDeviceProperties()->maxGridSize[0];
|
||
uint32_t max_grid_y = at::cuda::getCurrentDeviceProperties()->maxGridSize[1];
|
||
|
||
// Implementing exclusive cumprod.
|
||
// cumprod_1mp = cumprod(1 - p_choose)
|
||
// There is cumprod in pytorch, however there is no exclusive mode.
|
||
// cumprod(x) = [x1, x1x2, x2x3x4, ..., prod_{i=1}^n x_i]
|
||
// exclusive means
|
||
// cumprod(x) = [1, x1, x1x2, x1x2x3, ..., prod_{i=1}^{n-1} x_i]
|
||
uint32_t elements = bsz * tgt_len * src_len;
|
||
T* cumprod_1mp;
|
||
gpuErrchk(cudaMalloc(&cumprod_1mp, elements * sizeof(T)));
|
||
exclusiveCumprod<T>(
|
||
p_choose,
|
||
cumprod_1mp,
|
||
bsz,
|
||
tgt_len,
|
||
src_len,
|
||
max_grid_x,
|
||
max_grid_y,
|
||
stream);
|
||
|
||
// clamp cumprod_1mp to the range [eps, 1.0]
|
||
T* cumprod_1mp_clamp;
|
||
gpuErrchk(cudaMalloc(&cumprod_1mp_clamp, elements * sizeof(T)));
|
||
dim3 grid_clamp(std::min<T>(max_grid_x, bsz), 1, 1);
|
||
dim3 block_clamp(BLOCK_DIM_X, BLOCK_DIM_Y, 1);
|
||
clampKernel<T><<<grid_clamp, block_clamp, 0, stream>>>(
|
||
cumprod_1mp, cumprod_1mp_clamp, bsz, tgt_len, src_len, (T)eps, (T)1.0);
|
||
gpuErrchk(cudaGetLastError());
|
||
|
||
// ai = p_i * cumprod(1 − pi) * cumsum(a_i / cumprod(1 − pi))
|
||
dim3 grid_init(std::min<int>(max_grid_x, bsz), 1, 1);
|
||
initAlphaCUDAKernel<T>
|
||
<<<grid_init, 1, 0, stream>>>(alpha, bsz, tgt_len, src_len);
|
||
gpuErrchk(cudaGetLastError());
|
||
|
||
const int grid = std::min(bsz, max_grid_x);
|
||
|
||
for (uint32_t i = 0; i < tgt_len; i++) {
|
||
alignmentTrainCUDAKernel<T, SCAN_BLOCK><<<grid, SCAN_BLOCK, 0, stream>>>(
|
||
p_choose,
|
||
cumprod_1mp,
|
||
cumprod_1mp_clamp,
|
||
alpha,
|
||
bsz,
|
||
tgt_len,
|
||
src_len,
|
||
i);
|
||
gpuErrchk(cudaGetLastError());
|
||
}
|
||
|
||
gpuErrchk(cudaFree(cumprod_1mp));
|
||
gpuErrchk(cudaFree(cumprod_1mp_clamp));
|
||
}
|
||
|
||
} // namespace
|
||
|
||
void alignmentTrainCUDAWrapper(
|
||
const torch::Tensor& p_choose,
|
||
torch::Tensor& alpha,
|
||
float eps) {
|
||
// p_choose dimension: bsz, tgt_len, src_len
|
||
uint32_t bsz = p_choose.size(0);
|
||
uint32_t tgt_len = p_choose.size(1);
|
||
uint32_t src_len = p_choose.size(2);
|
||
|
||
cudaSetDevice(p_choose.get_device());
|
||
|
||
AT_DISPATCH_FLOATING_TYPES_AND2(
|
||
torch::ScalarType::Half,
|
||
torch::ScalarType::BFloat16,
|
||
p_choose.scalar_type(),
|
||
"alignmentTrainCUDAImpl",
|
||
[&]() {
|
||
alignmentTrainCUDAImpl<scalar_t>(
|
||
p_choose.data_ptr<scalar_t>(),
|
||
alpha.data_ptr<scalar_t>(),
|
||
bsz,
|
||
tgt_len,
|
||
src_len,
|
||
eps);
|
||
});
|
||
}
|