|
#include <ATen/ATen.h> |
|
#include <ATen/AccumulateType.h> |
|
#include <ATen/cuda/CUDAContext.h> |
|
#include <ATen/cuda/Exceptions.h> |
|
|
|
|
|
|
|
#include <assert.h> |
|
|
|
#include <sstream> |
|
|
|
#include "type_shim.h" |
|
#include "multi_tensor_apply.cuh" |
|
|
|
#define BLOCK_SIZE 512 |
|
#define ILP 4 |
|
|
|
template<typename in_t, typename out_t> |
|
struct ScaleFunctor |
|
{ |
|
__device__ __forceinline__ void operator()( |
|
int chunk_size, |
|
volatile int* noop_gmem, |
|
TensorListMetadata<2>& tl, |
|
float scale) |
|
{ |
|
|
|
|
|
|
|
|
|
int tensor_loc = tl.block_to_tensor[blockIdx.x]; |
|
int chunk_idx = tl.block_to_chunk[blockIdx.x]; |
|
int n = tl.sizes[tensor_loc]; |
|
|
|
in_t* in = (in_t*)tl.addresses[0][tensor_loc]; |
|
in += chunk_idx*chunk_size; |
|
|
|
out_t* out = (out_t*)tl.addresses[1][tensor_loc]; |
|
out += chunk_idx*chunk_size; |
|
|
|
n -= chunk_idx*chunk_size; |
|
|
|
|
|
float incoming_vals[ILP]; |
|
for(int i_start = 0; |
|
i_start < n && i_start < chunk_size; |
|
i_start += blockDim.x*ILP) |
|
{ |
|
#pragma unroll |
|
for(int ii = 0; ii < ILP; ii++) |
|
{ |
|
incoming_vals[ii] = 0; |
|
int i = i_start + threadIdx.x + ii*blockDim.x; |
|
if(i < n && i < chunk_size) |
|
incoming_vals[ii] = static_cast<float>(in[i]); |
|
} |
|
|
|
|
|
|
|
|
|
|
|
|
|
#pragma unroll |
|
for(int ii = 0; ii < ILP; ii++) |
|
{ |
|
int i = i_start + threadIdx.x + ii*blockDim.x; |
|
if(i < n && i < chunk_size) |
|
{ |
|
out[i] = static_cast<out_t>(incoming_vals[ii]*scale); |
|
if(!isfinite(incoming_vals[ii])) |
|
*noop_gmem = 1; |
|
} |
|
} |
|
} |
|
} |
|
}; |
|
|
|
void multi_tensor_scale_cuda( |
|
int chunk_size, |
|
at::Tensor noop_flag, |
|
std::vector<std::vector<at::Tensor>> tensor_lists, |
|
float scale) |
|
{ |
|
using namespace at; |
|
|
|
|
|
|
|
|
|
DISPATCH_FLOAT_AND_HALF(tensor_lists[0][0].scalar_type(), 0, "multi_tensor_scale_cuda", |
|
DISPATCH_FLOAT_AND_HALF(tensor_lists[1][0].scalar_type(), 1, "multi_tensor_scale_cuda", |
|
multi_tensor_apply<2>( |
|
BLOCK_SIZE, |
|
chunk_size, |
|
noop_flag, |
|
tensor_lists, |
|
ScaleFunctor<scalar_t_0, scalar_t_1>(), |
|
scale); )) |
|
AT_CUDA_CHECK(cudaGetLastError()); |
|
|
|
|
|
} |
|
|