mirror of https://github.com/hpcaitech/ColossalAI
[NFC] polish colossalai/kernel/cuda_native/csrc/multi_tensor_l2norm_kernel.cu code style (#635)
Wangbo Zhao
binmakeswell
parent
8a5d526e95
commit
6fcb381801
|
|
@ -1,4 +1,5 @@
|
||||||
// modified from https://github.com/NVIDIA/apex/blob/master/csrc/multi_tensor_l2norm_kernel.cu
|
// modified from
|
||||||
|
// https://github.com/NVIDIA/apex/blob/master/csrc/multi_tensor_l2norm_kernel.cu
|
||||||
#include <ATen/ATen.h>
|
#include <ATen/ATen.h>
|
||||||
#include <ATen/AccumulateType.h>
|
#include <ATen/AccumulateType.h>
|
||||||
#include <ATen/cuda/CUDAContext.h>
|
#include <ATen/cuda/CUDAContext.h>
|
||||||
|
|
@ -9,356 +10,311 @@
|
||||||
|
|
||||||
#include <assert.h>
|
#include <assert.h>
|
||||||
|
|
||||||
#include "type_shim.h"
|
|
||||||
#include "multi_tensor_apply.cuh"
|
#include "multi_tensor_apply.cuh"
|
||||||
|
#include "type_shim.h"
|
||||||
|
|
||||||
#define BLOCK_SIZE 512
|
#define BLOCK_SIZE 512
|
||||||
#define ILP 4
|
#define ILP 4
|
||||||
|
|
||||||
template <typename T>
|
template <typename T> __device__ __forceinline__ bool is_aligned(T *p) {
|
||||||
__device__ __forceinline__ bool is_aligned(T *p)
|
return ((uint64_t)p) % (ILP * sizeof(T)) == 0;
|
||||||
{
|
|
||||||
return ((uint64_t)p) % (ILP * sizeof(T)) == 0;
|
|
||||||
}
|
}
|
||||||
|
|
||||||
template <typename T>
|
template <typename T>
|
||||||
__device__ __forceinline__ void load_store(T *dst, T *src, int dst_offset, int src_offset)
|
__device__ __forceinline__ void load_store(T *dst, T *src, int dst_offset,
|
||||||
{
|
int src_offset) {
|
||||||
typedef typename std::aligned_storage<ILP * sizeof(T), ILP * alignof(T)>::type LT;
|
typedef
|
||||||
((LT *)dst)[dst_offset] = ((LT *)src)[src_offset];
|
typename std::aligned_storage<ILP * sizeof(T), ILP * alignof(T)>::type LT;
|
||||||
|
((LT *)dst)[dst_offset] = ((LT *)src)[src_offset];
|
||||||
}
|
}
|
||||||
|
|
||||||
template <typename x_t>
|
template <typename x_t> struct L2NormFunctor {
|
||||||
struct L2NormFunctor
|
__device__ __forceinline__ void
|
||||||
{
|
operator()(int chunk_size, volatile int *noop_gmem, TensorListMetadata<1> &tl,
|
||||||
__device__ __forceinline__ void operator()(
|
float *output, float *output_per_tensor, bool per_tensor,
|
||||||
int chunk_size,
|
int max_chunks_per_tensor) {
|
||||||
volatile int *noop_gmem,
|
// I'd like this kernel to propagate infs/nans.
|
||||||
TensorListMetadata<1> &tl,
|
// if(*noop_gmem == 1)
|
||||||
float *output,
|
// return;
|
||||||
float *output_per_tensor,
|
|
||||||
bool per_tensor,
|
|
||||||
int max_chunks_per_tensor)
|
|
||||||
{
|
|
||||||
// I'd like this kernel to propagate infs/nans.
|
|
||||||
// if(*noop_gmem == 1)
|
|
||||||
// return;
|
|
||||||
|
|
||||||
int tensor_loc = tl.block_to_tensor[blockIdx.x];
|
int tensor_loc = tl.block_to_tensor[blockIdx.x];
|
||||||
int chunk_idx = tl.block_to_chunk[blockIdx.x];
|
int chunk_idx = tl.block_to_chunk[blockIdx.x];
|
||||||
int n = tl.sizes[tensor_loc];
|
int n = tl.sizes[tensor_loc];
|
||||||
|
|
||||||
x_t *x = (x_t *)tl.addresses[0][tensor_loc];
|
x_t *x = (x_t *)tl.addresses[0][tensor_loc];
|
||||||
x += chunk_idx * chunk_size;
|
x += chunk_idx * chunk_size;
|
||||||
|
|
||||||
n -= chunk_idx * chunk_size;
|
n -= chunk_idx * chunk_size;
|
||||||
|
|
||||||
__shared__ float s_vals[512];
|
__shared__ float s_vals[512];
|
||||||
|
|
||||||
float vals[ILP]; // = {0}; // this probably works too but I want to be sure...
|
float
|
||||||
x_t r_x[ILP];
|
vals[ILP]; // = {0}; // this probably works too but I want to be sure...
|
||||||
for (int i = 0; i < ILP; i++)
|
x_t r_x[ILP];
|
||||||
{
|
for (int i = 0; i < ILP; i++) {
|
||||||
vals[i] = 0.f;
|
vals[i] = 0.f;
|
||||||
r_x[i] = 0;
|
r_x[i] = 0;
|
||||||
}
|
|
||||||
|
|
||||||
// to make things simple, we put aligned case in a different code path
|
|
||||||
if (n % ILP == 0 && chunk_size % ILP == 0 && is_aligned(x))
|
|
||||||
{
|
|
||||||
for (int i_start = threadIdx.x; i_start * ILP < n && i_start * ILP < chunk_size; i_start += blockDim.x)
|
|
||||||
{
|
|
||||||
// load
|
|
||||||
load_store(r_x, x, 0, i_start);
|
|
||||||
#pragma unroll
|
|
||||||
for (int ii = 0; ii < ILP; ii++)
|
|
||||||
{
|
|
||||||
float next = static_cast<float>(r_x[ii]);
|
|
||||||
vals[ii] += next * next;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
else
|
|
||||||
{
|
|
||||||
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++)
|
|
||||||
{
|
|
||||||
int i = i_start + threadIdx.x + ii * blockDim.x;
|
|
||||||
if (i < n && i < chunk_size)
|
|
||||||
{
|
|
||||||
float next = static_cast<float>(x[i]);
|
|
||||||
vals[ii] += next * next;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
float val = 0.f;
|
|
||||||
for (int i = 0; i < ILP; i++)
|
|
||||||
val += vals[i];
|
|
||||||
|
|
||||||
float final = reduce_block_into_lanes(s_vals, val);
|
|
||||||
|
|
||||||
if (threadIdx.x == 0)
|
|
||||||
{
|
|
||||||
if (!isfinite(final))
|
|
||||||
*noop_gmem = 1; // Blindly fire off a write. These will race but that's ok.
|
|
||||||
output[blockIdx.x] += final;
|
|
||||||
if (per_tensor)
|
|
||||||
output_per_tensor[(tl.start_tensor_this_launch + tensor_loc) * max_chunks_per_tensor + chunk_idx] = final;
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// to make things simple, we put aligned case in a different code path
|
||||||
|
if (n % ILP == 0 && chunk_size % ILP == 0 && is_aligned(x)) {
|
||||||
|
for (int i_start = threadIdx.x;
|
||||||
|
i_start * ILP < n && i_start * ILP < chunk_size;
|
||||||
|
i_start += blockDim.x) {
|
||||||
|
// load
|
||||||
|
load_store(r_x, x, 0, i_start);
|
||||||
|
#pragma unroll
|
||||||
|
for (int ii = 0; ii < ILP; ii++) {
|
||||||
|
float next = static_cast<float>(r_x[ii]);
|
||||||
|
vals[ii] += next * next;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
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++) {
|
||||||
|
int i = i_start + threadIdx.x + ii * blockDim.x;
|
||||||
|
if (i < n && i < chunk_size) {
|
||||||
|
float next = static_cast<float>(x[i]);
|
||||||
|
vals[ii] += next * next;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
float val = 0.f;
|
||||||
|
for (int i = 0; i < ILP; i++)
|
||||||
|
val += vals[i];
|
||||||
|
|
||||||
|
float final = reduce_block_into_lanes(s_vals, val);
|
||||||
|
|
||||||
|
if (threadIdx.x == 0) {
|
||||||
|
if (!isfinite(final))
|
||||||
|
*noop_gmem =
|
||||||
|
1; // Blindly fire off a write. These will race but that's ok.
|
||||||
|
output[blockIdx.x] += final;
|
||||||
|
if (per_tensor)
|
||||||
|
output_per_tensor[(tl.start_tensor_this_launch + tensor_loc) *
|
||||||
|
max_chunks_per_tensor +
|
||||||
|
chunk_idx] = final;
|
||||||
|
}
|
||||||
|
}
|
||||||
};
|
};
|
||||||
|
|
||||||
// Probably better to template, but since we are not likely to support other norm
|
// Probably better to template, but since we are not likely to support other
|
||||||
template <typename x_t>
|
// norm
|
||||||
struct MaxNormFunctor
|
template <typename x_t> struct MaxNormFunctor {
|
||||||
{
|
__device__ __forceinline__ void
|
||||||
__device__ __forceinline__ void operator()(
|
operator()(int chunk_size, volatile int *noop_gmem, TensorListMetadata<1> &tl,
|
||||||
int chunk_size,
|
float *output, float *output_per_tensor, bool per_tensor,
|
||||||
volatile int *noop_gmem,
|
int max_chunks_per_tensor) {
|
||||||
TensorListMetadata<1> &tl,
|
// I'd like this kernel to propagate infs/nans.
|
||||||
float *output,
|
// if(*noop_gmem == 1)
|
||||||
float *output_per_tensor,
|
// return;
|
||||||
bool per_tensor,
|
|
||||||
int max_chunks_per_tensor)
|
|
||||||
{
|
|
||||||
// I'd like this kernel to propagate infs/nans.
|
|
||||||
// if(*noop_gmem == 1)
|
|
||||||
// return;
|
|
||||||
|
|
||||||
int tensor_loc = tl.block_to_tensor[blockIdx.x];
|
int tensor_loc = tl.block_to_tensor[blockIdx.x];
|
||||||
int chunk_idx = tl.block_to_chunk[blockIdx.x];
|
int chunk_idx = tl.block_to_chunk[blockIdx.x];
|
||||||
int n = tl.sizes[tensor_loc];
|
int n = tl.sizes[tensor_loc];
|
||||||
|
|
||||||
x_t *x = (x_t *)tl.addresses[0][tensor_loc];
|
x_t *x = (x_t *)tl.addresses[0][tensor_loc];
|
||||||
x += chunk_idx * chunk_size;
|
x += chunk_idx * chunk_size;
|
||||||
|
|
||||||
n -= chunk_idx * chunk_size;
|
n -= chunk_idx * chunk_size;
|
||||||
|
|
||||||
__shared__ float s_vals[512];
|
__shared__ float s_vals[512];
|
||||||
|
|
||||||
float vals[ILP]; // = {0}; // this probably works too but I want to be sure...
|
float
|
||||||
x_t r_x[ILP];
|
vals[ILP]; // = {0}; // this probably works too but I want to be sure...
|
||||||
for (int i = 0; i < ILP; i++)
|
x_t r_x[ILP];
|
||||||
{
|
for (int i = 0; i < ILP; i++) {
|
||||||
vals[i] = 0.f;
|
vals[i] = 0.f;
|
||||||
r_x[i] = 0;
|
r_x[i] = 0;
|
||||||
}
|
|
||||||
|
|
||||||
// to make things simple, we put aligned case in a different code path
|
|
||||||
if (n % ILP == 0 && chunk_size % ILP == 0 && is_aligned(x))
|
|
||||||
{
|
|
||||||
for (int i_start = threadIdx.x; i_start * ILP < n && i_start * ILP < chunk_size; i_start += blockDim.x)
|
|
||||||
{
|
|
||||||
// load
|
|
||||||
load_store(r_x, x, 0, i_start);
|
|
||||||
#pragma unroll
|
|
||||||
for (int ii = 0; ii < ILP; ii++)
|
|
||||||
{
|
|
||||||
float next = static_cast<float>(r_x[ii]);
|
|
||||||
vals[ii] = fmaxf(fabsf(vals[ii]), fabsf(next));
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
else
|
|
||||||
{
|
|
||||||
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++)
|
|
||||||
{
|
|
||||||
int i = i_start + threadIdx.x + ii * blockDim.x;
|
|
||||||
if (i < n && i < chunk_size)
|
|
||||||
{
|
|
||||||
float next = static_cast<float>(x[i]);
|
|
||||||
vals[ii] = fmaxf(fabsf(vals[ii]), fabsf(next));
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
float val = 0.f;
|
|
||||||
for (int i = 0; i < ILP; i++)
|
|
||||||
val = fmaxf(fabsf(val), fabsf(vals[i]));
|
|
||||||
|
|
||||||
float final = reduce_block_into_lanes_max_op(s_vals, val);
|
|
||||||
|
|
||||||
if (threadIdx.x == 0)
|
|
||||||
{
|
|
||||||
if (!isfinite(final))
|
|
||||||
*noop_gmem = 1; // Blindly fire off a write. These will race but that's ok.
|
|
||||||
output[blockIdx.x] = fmaxf(fabsf(output[blockIdx.x]), fabsf(final));
|
|
||||||
if (per_tensor)
|
|
||||||
output_per_tensor[(tl.start_tensor_this_launch + tensor_loc) * max_chunks_per_tensor + chunk_idx] = final;
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// to make things simple, we put aligned case in a different code path
|
||||||
|
if (n % ILP == 0 && chunk_size % ILP == 0 && is_aligned(x)) {
|
||||||
|
for (int i_start = threadIdx.x;
|
||||||
|
i_start * ILP < n && i_start * ILP < chunk_size;
|
||||||
|
i_start += blockDim.x) {
|
||||||
|
// load
|
||||||
|
load_store(r_x, x, 0, i_start);
|
||||||
|
#pragma unroll
|
||||||
|
for (int ii = 0; ii < ILP; ii++) {
|
||||||
|
float next = static_cast<float>(r_x[ii]);
|
||||||
|
vals[ii] = fmaxf(fabsf(vals[ii]), fabsf(next));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
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++) {
|
||||||
|
int i = i_start + threadIdx.x + ii * blockDim.x;
|
||||||
|
if (i < n && i < chunk_size) {
|
||||||
|
float next = static_cast<float>(x[i]);
|
||||||
|
vals[ii] = fmaxf(fabsf(vals[ii]), fabsf(next));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
float val = 0.f;
|
||||||
|
for (int i = 0; i < ILP; i++)
|
||||||
|
val = fmaxf(fabsf(val), fabsf(vals[i]));
|
||||||
|
|
||||||
|
float final = reduce_block_into_lanes_max_op(s_vals, val);
|
||||||
|
|
||||||
|
if (threadIdx.x == 0) {
|
||||||
|
if (!isfinite(final))
|
||||||
|
*noop_gmem =
|
||||||
|
1; // Blindly fire off a write. These will race but that's ok.
|
||||||
|
output[blockIdx.x] = fmaxf(fabsf(output[blockIdx.x]), fabsf(final));
|
||||||
|
if (per_tensor)
|
||||||
|
output_per_tensor[(tl.start_tensor_this_launch + tensor_loc) *
|
||||||
|
max_chunks_per_tensor +
|
||||||
|
chunk_idx] = final;
|
||||||
|
}
|
||||||
|
}
|
||||||
};
|
};
|
||||||
|
|
||||||
__global__ void cleanup(
|
__global__ void cleanup(float *output, float *output_per_tensor, float *ret,
|
||||||
float *output,
|
float *ret_per_tensor, bool per_tensor,
|
||||||
float *output_per_tensor,
|
int max_chunks_per_tensor) {
|
||||||
float *ret,
|
__shared__ float vals[512];
|
||||||
float *ret_per_tensor,
|
|
||||||
bool per_tensor,
|
|
||||||
int max_chunks_per_tensor)
|
|
||||||
{
|
|
||||||
__shared__ float vals[512];
|
|
||||||
|
|
||||||
if (blockIdx.x == 0)
|
if (blockIdx.x == 0) {
|
||||||
{
|
float val = 0;
|
||||||
float val = 0;
|
if (threadIdx.x < 320)
|
||||||
if (threadIdx.x < 320)
|
val = output[threadIdx.x];
|
||||||
val = output[threadIdx.x];
|
|
||||||
|
|
||||||
float final = reduce_block_into_lanes(vals, val);
|
float final = reduce_block_into_lanes(vals, val);
|
||||||
|
|
||||||
if (threadIdx.x == 0)
|
if (threadIdx.x == 0)
|
||||||
*ret = sqrt(final);
|
*ret = sqrt(final);
|
||||||
}
|
}
|
||||||
|
|
||||||
if (per_tensor)
|
if (per_tensor) {
|
||||||
{
|
float *output_this_tensor =
|
||||||
float *output_this_tensor = output_per_tensor + blockIdx.x * max_chunks_per_tensor;
|
output_per_tensor + blockIdx.x * max_chunks_per_tensor;
|
||||||
|
|
||||||
float val = 0;
|
float val = 0;
|
||||||
for (int i = threadIdx.x; i < max_chunks_per_tensor; i += blockDim.x)
|
for (int i = threadIdx.x; i < max_chunks_per_tensor; i += blockDim.x)
|
||||||
val += output_this_tensor[i];
|
val += output_this_tensor[i];
|
||||||
|
|
||||||
float final = reduce_block_into_lanes(vals, val);
|
float final = reduce_block_into_lanes(vals, val);
|
||||||
|
|
||||||
if (threadIdx.x == 0)
|
if (threadIdx.x == 0)
|
||||||
ret_per_tensor[blockIdx.x] = sqrt(final);
|
ret_per_tensor[blockIdx.x] = sqrt(final);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
__global__ void cleanup_v2(
|
__global__ void cleanup_v2(float *output, float *output_per_tensor, float *ret,
|
||||||
float *output,
|
float *ret_per_tensor, bool per_tensor,
|
||||||
float *output_per_tensor,
|
int max_chunks_per_tensor, int norm_type,
|
||||||
float *ret,
|
float alpha, float beta) {
|
||||||
float *ret_per_tensor,
|
__shared__ float vals[512];
|
||||||
bool per_tensor,
|
|
||||||
int max_chunks_per_tensor,
|
|
||||||
int norm_type,
|
|
||||||
float alpha,
|
|
||||||
float beta)
|
|
||||||
{
|
|
||||||
__shared__ float vals[512];
|
|
||||||
|
|
||||||
if (blockIdx.x == 0)
|
if (blockIdx.x == 0) {
|
||||||
{
|
float val = 0;
|
||||||
float val = 0;
|
if (threadIdx.x < 320)
|
||||||
if (threadIdx.x < 320)
|
val = output[threadIdx.x];
|
||||||
val = output[threadIdx.x];
|
|
||||||
|
|
||||||
if (norm_type == 0)
|
if (norm_type == 0) {
|
||||||
{
|
float final = reduce_block_into_lanes_max_op(vals, val);
|
||||||
float final = reduce_block_into_lanes_max_op(vals, val);
|
if (threadIdx.x == 0)
|
||||||
if (threadIdx.x == 0)
|
*ret = alpha * (*ret) + beta * final;
|
||||||
*ret = alpha * (*ret) + beta * final;
|
} else {
|
||||||
}
|
float final = reduce_block_into_lanes(vals, val);
|
||||||
else
|
if (threadIdx.x == 0)
|
||||||
{
|
*ret = sqrt(alpha * (*ret) * (*ret) + beta * final);
|
||||||
float final = reduce_block_into_lanes(vals, val);
|
|
||||||
if (threadIdx.x == 0)
|
|
||||||
*ret = sqrt(alpha * (*ret) * (*ret) + beta * final);
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
}
|
||||||
|
|
||||||
if (per_tensor)
|
if (per_tensor) {
|
||||||
{
|
float *output_this_tensor =
|
||||||
float *output_this_tensor = output_per_tensor + blockIdx.x * max_chunks_per_tensor;
|
output_per_tensor + blockIdx.x * max_chunks_per_tensor;
|
||||||
|
|
||||||
if (norm_type == 0)
|
if (norm_type == 0) {
|
||||||
{
|
float val = 0;
|
||||||
float val = 0;
|
for (int i = threadIdx.x; i < max_chunks_per_tensor; i += blockDim.x)
|
||||||
for (int i = threadIdx.x; i < max_chunks_per_tensor; i += blockDim.x)
|
val = fmaxf(fabsf(val), fabsf(output_this_tensor[i]));
|
||||||
val = fmaxf(fabsf(val), fabsf(output_this_tensor[i]));
|
|
||||||
|
|
||||||
float final = reduce_block_into_lanes_max_op(vals, val);
|
float final = reduce_block_into_lanes_max_op(vals, val);
|
||||||
|
|
||||||
if (threadIdx.x == 0)
|
if (threadIdx.x == 0)
|
||||||
ret_per_tensor[blockIdx.x] = alpha * ret_per_tensor[blockIdx.x] + beta * final;
|
ret_per_tensor[blockIdx.x] =
|
||||||
}
|
alpha * ret_per_tensor[blockIdx.x] + beta * final;
|
||||||
else
|
} else {
|
||||||
{
|
float val = 0;
|
||||||
float val = 0;
|
for (int i = threadIdx.x; i < max_chunks_per_tensor; i += blockDim.x)
|
||||||
for (int i = threadIdx.x; i < max_chunks_per_tensor; i += blockDim.x)
|
val += output_this_tensor[i];
|
||||||
val += output_this_tensor[i];
|
|
||||||
|
|
||||||
float final = reduce_block_into_lanes(vals, val);
|
float final = reduce_block_into_lanes(vals, val);
|
||||||
|
|
||||||
if (threadIdx.x == 0)
|
if (threadIdx.x == 0)
|
||||||
ret_per_tensor[blockIdx.x] = sqrt(alpha * ret_per_tensor[blockIdx.x] * ret_per_tensor[blockIdx.x] + beta * final);
|
ret_per_tensor[blockIdx.x] = sqrt(alpha * ret_per_tensor[blockIdx.x] *
|
||||||
}
|
ret_per_tensor[blockIdx.x] +
|
||||||
|
beta * final);
|
||||||
}
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
std::tuple<at::Tensor, at::Tensor> multi_tensor_l2norm_cuda(
|
std::tuple<at::Tensor, at::Tensor>
|
||||||
int chunk_size,
|
multi_tensor_l2norm_cuda(int chunk_size, at::Tensor noop_flag,
|
||||||
at::Tensor noop_flag,
|
std::vector<std::vector<at::Tensor>> tensor_lists,
|
||||||
std::vector<std::vector<at::Tensor>> tensor_lists,
|
at::optional<bool> per_tensor_python) {
|
||||||
at::optional<bool> per_tensor_python)
|
bool per_tensor =
|
||||||
{
|
per_tensor_python.has_value() ? per_tensor_python.value() : false;
|
||||||
bool per_tensor = per_tensor_python.has_value() ? per_tensor_python.value() : false;
|
|
||||||
|
|
||||||
auto float_options = tensor_lists[0][0].options().dtype(at::kFloat);
|
auto float_options = tensor_lists[0][0].options().dtype(at::kFloat);
|
||||||
auto output = at::zeros({320}, float_options);
|
auto output = at::zeros({320}, float_options);
|
||||||
|
|
||||||
at::Tensor output_per_tensor;
|
at::Tensor output_per_tensor;
|
||||||
at::Tensor ret_per_tensor;
|
at::Tensor ret_per_tensor;
|
||||||
|
|
||||||
int ntensors = tensor_lists[0].size();
|
int ntensors = tensor_lists[0].size();
|
||||||
int max_chunks_per_tensor = -1;
|
int max_chunks_per_tensor = -1;
|
||||||
|
|
||||||
if (per_tensor)
|
if (per_tensor) {
|
||||||
{
|
for (int t = 0; t < ntensors; t++) {
|
||||||
for (int t = 0; t < ntensors; t++)
|
int max_chunks_this_tensor =
|
||||||
{
|
(tensor_lists[0][t].numel() + chunk_size - 1) / chunk_size;
|
||||||
int max_chunks_this_tensor = (tensor_lists[0][t].numel() + chunk_size - 1) / chunk_size;
|
if (max_chunks_this_tensor > max_chunks_per_tensor)
|
||||||
if (max_chunks_this_tensor > max_chunks_per_tensor)
|
max_chunks_per_tensor = max_chunks_this_tensor;
|
||||||
max_chunks_per_tensor = max_chunks_this_tensor;
|
|
||||||
}
|
|
||||||
output_per_tensor = at::zeros({ntensors * max_chunks_per_tensor}, float_options);
|
|
||||||
ret_per_tensor = at::empty({ntensors}, float_options);
|
|
||||||
}
|
|
||||||
else
|
|
||||||
{
|
|
||||||
ret_per_tensor = at::empty({0}, float_options);
|
|
||||||
}
|
}
|
||||||
|
output_per_tensor =
|
||||||
|
at::zeros({ntensors * max_chunks_per_tensor}, float_options);
|
||||||
|
ret_per_tensor = at::empty({ntensors}, float_options);
|
||||||
|
} else {
|
||||||
|
ret_per_tensor = at::empty({0}, float_options);
|
||||||
|
}
|
||||||
|
|
||||||
DISPATCH_FLOAT_AND_HALF(tensor_lists[0][0].scalar_type(), 0, "multi_tensor_l2norm_cuda",
|
DISPATCH_FLOAT_AND_HALF(
|
||||||
multi_tensor_apply<1>(
|
tensor_lists[0][0].scalar_type(), 0, "multi_tensor_l2norm_cuda",
|
||||||
BLOCK_SIZE,
|
multi_tensor_apply<1>(
|
||||||
chunk_size,
|
BLOCK_SIZE, chunk_size, noop_flag, tensor_lists,
|
||||||
noop_flag,
|
L2NormFunctor<scalar_t_0>(), output.DATA_PTR<float>(),
|
||||||
tensor_lists,
|
per_tensor ? output_per_tensor.DATA_PTR<float>() : nullptr,
|
||||||
L2NormFunctor<scalar_t_0>(),
|
per_tensor, max_chunks_per_tensor);)
|
||||||
output.DATA_PTR<float>(),
|
|
||||||
per_tensor ? output_per_tensor.DATA_PTR<float>() : nullptr,
|
|
||||||
per_tensor,
|
|
||||||
max_chunks_per_tensor);)
|
|
||||||
|
|
||||||
AT_CUDA_CHECK(cudaGetLastError());
|
AT_CUDA_CHECK(cudaGetLastError());
|
||||||
// AT_CUDA_CHECK(cudaDeviceSynchronize());
|
// AT_CUDA_CHECK(cudaDeviceSynchronize());
|
||||||
|
|
||||||
// This involves one more small kernel launches, but will be negligible end to end.
|
// This involves one more small kernel launches, but will be negligible end to
|
||||||
// I could get rid of these by hacking the functor + multi tensor harness with persistence
|
// end. I could get rid of these by hacking the functor + multi tensor harness
|
||||||
// logic, but keeping it simple for now
|
// with persistence logic, but keeping it simple for now
|
||||||
auto ret = at::empty({1}, output.options());
|
auto ret = at::empty({1}, output.options());
|
||||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(output));
|
const at::cuda::OptionalCUDAGuard device_guard(device_of(output));
|
||||||
auto stream = at::cuda::getCurrentCUDAStream();
|
auto stream = at::cuda::getCurrentCUDAStream();
|
||||||
cleanup<<<per_tensor ? ntensors : 1, 512, 0, stream>>>(
|
cleanup<<<per_tensor ? ntensors : 1, 512, 0, stream>>>(
|
||||||
output.DATA_PTR<float>(),
|
output.DATA_PTR<float>(),
|
||||||
per_tensor ? output_per_tensor.DATA_PTR<float>() : nullptr,
|
per_tensor ? output_per_tensor.DATA_PTR<float>() : nullptr,
|
||||||
ret.DATA_PTR<float>(),
|
ret.DATA_PTR<float>(),
|
||||||
per_tensor ? ret_per_tensor.DATA_PTR<float>() : nullptr,
|
per_tensor ? ret_per_tensor.DATA_PTR<float>() : nullptr, per_tensor,
|
||||||
per_tensor,
|
max_chunks_per_tensor);
|
||||||
max_chunks_per_tensor);
|
|
||||||
|
|
||||||
return std::tuple<at::Tensor, at::Tensor>(ret, ret_per_tensor);
|
return std::tuple<at::Tensor, at::Tensor>(ret, ret_per_tensor);
|
||||||
}
|
}
|
||||||
|
|
||||||
// Compute and update grad norm
|
// Compute and update grad norm
|
||||||
|
|
@ -366,90 +322,66 @@ std::tuple<at::Tensor, at::Tensor> multi_tensor_l2norm_cuda(
|
||||||
// L-2: gn = sqrt(a * gn^2 + b * n^2)
|
// L-2: gn = sqrt(a * gn^2 + b * n^2)
|
||||||
// L-inf: gn = a * gn + b * n
|
// L-inf: gn = a * gn + b * n
|
||||||
void multi_tensor_norm_out_cuda(
|
void multi_tensor_norm_out_cuda(
|
||||||
int chunk_size,
|
int chunk_size, at::Tensor noop_flag,
|
||||||
at::Tensor noop_flag,
|
std::vector<std::vector<at::Tensor>> tensor_lists, at::Tensor out,
|
||||||
std::vector<std::vector<at::Tensor>> tensor_lists,
|
const float alpha, const float beta, const int norm_type) {
|
||||||
at::Tensor out,
|
auto float_options = tensor_lists[0][0].options().dtype(at::kFloat);
|
||||||
const float alpha,
|
TORCH_CHECK(tensor_lists[0][0].device() == noop_flag.device(),
|
||||||
const float beta,
|
"noop flag should be on the same device as tensors");
|
||||||
const int norm_type)
|
// we don't need global thus uses empty here
|
||||||
{
|
auto output = at::empty({320}, float_options);
|
||||||
auto float_options = tensor_lists[0][0].options().dtype(at::kFloat);
|
|
||||||
TORCH_CHECK(tensor_lists[0][0].device() == noop_flag.device(), "noop flag should be on the same device as tensors");
|
|
||||||
// we don't need global thus uses empty here
|
|
||||||
auto output = at::empty({320}, float_options);
|
|
||||||
|
|
||||||
at::Tensor output_per_tensor;
|
at::Tensor output_per_tensor;
|
||||||
at::Tensor ret_per_tensor;
|
at::Tensor ret_per_tensor;
|
||||||
|
|
||||||
int ntensors = tensor_lists[0].size();
|
int ntensors = tensor_lists[0].size();
|
||||||
int max_chunks_per_tensor = -1;
|
int max_chunks_per_tensor = -1;
|
||||||
|
|
||||||
for (int t = 0; t < ntensors; t++)
|
for (int t = 0; t < ntensors; t++) {
|
||||||
{
|
int max_chunks_this_tensor =
|
||||||
int max_chunks_this_tensor = (tensor_lists[0][t].numel() + chunk_size - 1) / chunk_size;
|
(tensor_lists[0][t].numel() + chunk_size - 1) / chunk_size;
|
||||||
if (max_chunks_this_tensor > max_chunks_per_tensor)
|
if (max_chunks_this_tensor > max_chunks_per_tensor)
|
||||||
max_chunks_per_tensor = max_chunks_this_tensor;
|
max_chunks_per_tensor = max_chunks_this_tensor;
|
||||||
}
|
}
|
||||||
|
|
||||||
// Although it is single write then read, still need to be zero
|
// Although it is single write then read, still need to be zero
|
||||||
// Since tailing element also participate cleanup
|
// Since tailing element also participate cleanup
|
||||||
output_per_tensor = at::zeros({ntensors * max_chunks_per_tensor}, float_options);
|
output_per_tensor =
|
||||||
|
at::zeros({ntensors * max_chunks_per_tensor}, float_options);
|
||||||
|
|
||||||
if (norm_type == 0)
|
if (norm_type == 0) {
|
||||||
{
|
DISPATCH_FLOAT_AND_HALF(
|
||||||
DISPATCH_FLOAT_AND_HALF(
|
tensor_lists[0][0].scalar_type(), 0, "multi_tensor_maxnorm_cuda",
|
||||||
tensor_lists[0][0].scalar_type(), 0, "multi_tensor_maxnorm_cuda",
|
multi_tensor_apply<1>(
|
||||||
multi_tensor_apply<1>(
|
BLOCK_SIZE, chunk_size, noop_flag, tensor_lists,
|
||||||
BLOCK_SIZE,
|
MaxNormFunctor<scalar_t_0>(), output.DATA_PTR<float>(),
|
||||||
chunk_size,
|
output_per_tensor.DATA_PTR<float>(), true, max_chunks_per_tensor);)
|
||||||
noop_flag,
|
} else {
|
||||||
tensor_lists,
|
DISPATCH_FLOAT_AND_HALF(
|
||||||
MaxNormFunctor<scalar_t_0>(),
|
tensor_lists[0][0].scalar_type(), 0, "multi_tensor_l2norm_cuda",
|
||||||
output.DATA_PTR<float>(),
|
multi_tensor_apply<1>(
|
||||||
output_per_tensor.DATA_PTR<float>(),
|
BLOCK_SIZE, chunk_size, noop_flag, tensor_lists,
|
||||||
true,
|
L2NormFunctor<scalar_t_0>(), output.DATA_PTR<float>(),
|
||||||
max_chunks_per_tensor);)
|
output_per_tensor.DATA_PTR<float>(), true, max_chunks_per_tensor);)
|
||||||
}
|
}
|
||||||
else
|
AT_CUDA_CHECK(cudaGetLastError());
|
||||||
{
|
|
||||||
DISPATCH_FLOAT_AND_HALF(
|
|
||||||
tensor_lists[0][0].scalar_type(), 0, "multi_tensor_l2norm_cuda",
|
|
||||||
multi_tensor_apply<1>(
|
|
||||||
BLOCK_SIZE,
|
|
||||||
chunk_size,
|
|
||||||
noop_flag,
|
|
||||||
tensor_lists,
|
|
||||||
L2NormFunctor<scalar_t_0>(),
|
|
||||||
output.DATA_PTR<float>(),
|
|
||||||
output_per_tensor.DATA_PTR<float>(),
|
|
||||||
true,
|
|
||||||
max_chunks_per_tensor);)
|
|
||||||
}
|
|
||||||
AT_CUDA_CHECK(cudaGetLastError());
|
|
||||||
|
|
||||||
// AT_CUDA_CHECK(cudaDeviceSynchronize());
|
// AT_CUDA_CHECK(cudaDeviceSynchronize());
|
||||||
|
|
||||||
// This involves one more small kernel launches, but will be negligible end to end.
|
// This involves one more small kernel launches, but will be negligible end to
|
||||||
// I could get rid of these by hacking the functor + multi tensor harness with persistence
|
// end. I could get rid of these by hacking the functor + multi tensor harness
|
||||||
// logic, but keeping it simple for now
|
// with persistence logic, but keeping it simple for now
|
||||||
auto ret = at::empty({1}, output.options());
|
auto ret = at::empty({1}, output.options());
|
||||||
|
|
||||||
// Adding the following device guard since it happens sometimes that the
|
// Adding the following device guard since it happens sometimes that the
|
||||||
// tensors are on one device and the cuda stream is on another device which
|
// tensors are on one device and the cuda stream is on another device which
|
||||||
// results in ILLEGAL MEM ACCESS error.
|
// results in ILLEGAL MEM ACCESS error.
|
||||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(output));
|
const at::cuda::OptionalCUDAGuard device_guard(device_of(output));
|
||||||
auto stream = at::cuda::getCurrentCUDAStream();
|
auto stream = at::cuda::getCurrentCUDAStream();
|
||||||
cleanup_v2<<<ntensors, 512, 0, stream>>>(
|
cleanup_v2<<<ntensors, 512, 0, stream>>>(
|
||||||
output.DATA_PTR<float>(),
|
output.DATA_PTR<float>(), output_per_tensor.DATA_PTR<float>(),
|
||||||
output_per_tensor.DATA_PTR<float>(),
|
ret.DATA_PTR<float>(), out.DATA_PTR<float>(), true, max_chunks_per_tensor,
|
||||||
ret.DATA_PTR<float>(),
|
norm_type, alpha, beta);
|
||||||
out.DATA_PTR<float>(),
|
|
||||||
true,
|
|
||||||
max_chunks_per_tensor,
|
|
||||||
norm_type,
|
|
||||||
alpha,
|
|
||||||
beta);
|
|
||||||
|
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
Loading…
Reference in New Issue