github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

[NFC] polish colossalai/kernel/cuda_native/csrc/multi_tensor_l2norm_kernel.cu code style (#635)

pull/673/head
Wangbo Zhao 2022-04-02 10:45:04 +08:00 committed by binmakeswell
parent 8a5d526e95
commit 6fcb381801
1 changed files with 304 additions and 372 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

322
colossalai/kernel/cuda_native/csrc/multi_tensor_l2norm_kernel.cu
View File

@ -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,37 +10,29 @@
#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
typename std::aligned_storage<ILP * sizeof(T), ILP * alignof(T)>::type LT;
((LT *)dst)[dst_offset] = ((LT *)src)[src_offset]; ((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,
TensorListMetadata<1> &tl,
float *output,
float *output_per_tensor,
bool per_tensor,
int max_chunks_per_tensor)
{
// I'd like this kernel to propagate infs/nans. // I'd like this kernel to propagate infs/nans.
// if(*noop_gmem == 1) // if(*noop_gmem == 1)
// return; // return;
@ -55,39 +48,34 @@ struct L2NormFunctor
__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
vals[ILP]; // = {0}; // this probably works too but I want to be sure...
x_t r_x[ILP]; x_t r_x[ILP];
for (int i = 0; i < ILP; i++) 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 // to make things simple, we put aligned case in a different code path
if (n % ILP == 0 && chunk_size % ILP == 0 && is_aligned(x)) if (n % ILP == 0 && chunk_size % ILP == 0 && is_aligned(x)) {
{ for (int i_start = threadIdx.x;
for (int i_start = threadIdx.x; i_start * ILP < n && i_start * ILP < chunk_size; i_start += blockDim.x) i_start * ILP < n && i_start * ILP < chunk_size;
{ i_start += blockDim.x) {
// load // load
load_store(r_x, x, 0, i_start); load_store(r_x, x, 0, i_start);
#pragma unroll #pragma unroll
for (int ii = 0; ii < ILP; ii++) for (int ii = 0; ii < ILP; ii++) {
{
float next = static_cast<float>(r_x[ii]); float next = static_cast<float>(r_x[ii]);
vals[ii] += next * next; vals[ii] += next * next;
} }
} }
} } else {
else for (int i_start = 0; i_start < n && i_start < chunk_size;
{ i_start += blockDim.x * ILP) {
for (int i_start = 0; i_start < n && i_start < chunk_size; i_start += blockDim.x * ILP)
{
#pragma unroll #pragma unroll
for (int ii = 0; ii < ILP; ii++) for (int ii = 0; ii < ILP; ii++) {
{
int i = i_start + threadIdx.x + ii * blockDim.x; int i = i_start + threadIdx.x + ii * blockDim.x;
if (i < n && i < chunk_size) if (i < n && i < chunk_size) {
{
float next = static_cast<float>(x[i]); float next = static_cast<float>(x[i]);
vals[ii] += next * next; vals[ii] += next * next;
} }
@ -101,30 +89,26 @@ struct L2NormFunctor
float final = reduce_block_into_lanes(s_vals, val); float final = reduce_block_into_lanes(s_vals, val);
if (threadIdx.x == 0) if (threadIdx.x == 0) {
{
if (!isfinite(final)) if (!isfinite(final))
*noop_gmem = 1; // Blindly fire off a write. These will race but that's ok. *noop_gmem =
1; // Blindly fire off a write. These will race but that's ok.
output[blockIdx.x] += final; output[blockIdx.x] += final;
if (per_tensor) if (per_tensor)
output_per_tensor[(tl.start_tensor_this_launch + tensor_loc) * max_chunks_per_tensor + chunk_idx] = final; 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,
float *output,
float *output_per_tensor,
bool per_tensor,
int max_chunks_per_tensor)
{
// I'd like this kernel to propagate infs/nans. // I'd like this kernel to propagate infs/nans.
// if(*noop_gmem == 1) // if(*noop_gmem == 1)
// return; // return;
@ -140,39 +124,34 @@ struct MaxNormFunctor
__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
vals[ILP]; // = {0}; // this probably works too but I want to be sure...
x_t r_x[ILP]; x_t r_x[ILP];
for (int i = 0; i < ILP; i++) 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 // to make things simple, we put aligned case in a different code path
if (n % ILP == 0 && chunk_size % ILP == 0 && is_aligned(x)) if (n % ILP == 0 && chunk_size % ILP == 0 && is_aligned(x)) {
{ for (int i_start = threadIdx.x;
for (int i_start = threadIdx.x; i_start * ILP < n && i_start * ILP < chunk_size; i_start += blockDim.x) i_start * ILP < n && i_start * ILP < chunk_size;
{ i_start += blockDim.x) {
// load // load
load_store(r_x, x, 0, i_start); load_store(r_x, x, 0, i_start);
#pragma unroll #pragma unroll
for (int ii = 0; ii < ILP; ii++) for (int ii = 0; ii < ILP; ii++) {
{
float next = static_cast<float>(r_x[ii]); float next = static_cast<float>(r_x[ii]);
vals[ii] = fmaxf(fabsf(vals[ii]), fabsf(next)); vals[ii] = fmaxf(fabsf(vals[ii]), fabsf(next));
} }
} }
} } else {
else for (int i_start = 0; i_start < n && i_start < chunk_size;
{ i_start += blockDim.x * ILP) {
for (int i_start = 0; i_start < n && i_start < chunk_size; i_start += blockDim.x * ILP)
{
#pragma unroll #pragma unroll
for (int ii = 0; ii < ILP; ii++) for (int ii = 0; ii < ILP; ii++) {
{
int i = i_start + threadIdx.x + ii * blockDim.x; int i = i_start + threadIdx.x + ii * blockDim.x;
if (i < n && i < chunk_size) if (i < n && i < chunk_size) {
{
float next = static_cast<float>(x[i]); float next = static_cast<float>(x[i]);
vals[ii] = fmaxf(fabsf(vals[ii]), fabsf(next)); vals[ii] = fmaxf(fabsf(vals[ii]), fabsf(next));
} }
@ -186,29 +165,25 @@ struct MaxNormFunctor
float final = reduce_block_into_lanes_max_op(s_vals, val); float final = reduce_block_into_lanes_max_op(s_vals, val);
if (threadIdx.x == 0) if (threadIdx.x == 0) {
{
if (!isfinite(final)) if (!isfinite(final))
*noop_gmem = 1; // Blindly fire off a write. These will race but that's ok. *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)); output[blockIdx.x] = fmaxf(fabsf(output[blockIdx.x]), fabsf(final));
if (per_tensor) if (per_tensor)
output_per_tensor[(tl.start_tensor_this_launch + tensor_loc) * max_chunks_per_tensor + chunk_idx] = final; 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,
float *ret_per_tensor,
bool per_tensor,
int max_chunks_per_tensor)
{
__shared__ float vals[512]; __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];
@ -219,9 +194,9 @@ __global__ void cleanup(
*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)
@ -234,45 +209,33 @@ __global__ void cleanup(
} }
} }
__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,
bool per_tensor,
int max_chunks_per_tensor,
int norm_type,
float alpha,
float beta)
{
__shared__ float vals[512]; __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 {
else
{
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(alpha * (*ret) * (*ret) + beta * final); *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]));
@ -280,10 +243,9 @@ __global__ void cleanup_v2(
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];
@ -291,18 +253,19 @@ __global__ void cleanup_v2(
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 =
bool per_tensor = per_tensor_python.has_value() ? per_tensor_python.value() : false; 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);
@ -313,40 +276,34 @@ std::tuple<at::Tensor, at::Tensor> multi_tensor_l2norm_cuda(
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); output_per_tensor =
at::zeros({ntensors * max_chunks_per_tensor}, float_options);
ret_per_tensor = at::empty({ntensors}, float_options); ret_per_tensor = at::empty({ntensors}, float_options);
} } else {
else
{
ret_per_tensor = at::empty({0}, float_options); 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(
tensor_lists[0][0].scalar_type(), 0, "multi_tensor_l2norm_cuda",
multi_tensor_apply<1>( multi_tensor_apply<1>(
BLOCK_SIZE, BLOCK_SIZE, chunk_size, noop_flag, tensor_lists,
chunk_size, L2NormFunctor<scalar_t_0>(), output.DATA_PTR<float>(),
noop_flag,
tensor_lists,
L2NormFunctor<scalar_t_0>(),
output.DATA_PTR<float>(),
per_tensor ? output_per_tensor.DATA_PTR<float>() : nullptr, per_tensor ? output_per_tensor.DATA_PTR<float>() : nullptr,
per_tensor, per_tensor, max_chunks_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();
@ -354,8 +311,7 @@ std::tuple<at::Tensor, at::Tensor> multi_tensor_l2norm_cuda(
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);
@ -366,16 +322,12 @@ 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,
const float alpha,
const float beta,
const int norm_type)
{
auto float_options = tensor_lists[0][0].options().dtype(at::kFloat); 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"); 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 // we don't need global thus uses empty here
auto output = at::empty({320}, float_options); auto output = at::empty({320}, float_options);
@ -385,54 +337,40 @@ void multi_tensor_norm_out_cuda(
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, BLOCK_SIZE, chunk_size, noop_flag, tensor_lists,
chunk_size, MaxNormFunctor<scalar_t_0>(), output.DATA_PTR<float>(),
noop_flag, output_per_tensor.DATA_PTR<float>(), true, max_chunks_per_tensor);)
tensor_lists, } else {
MaxNormFunctor<scalar_t_0>(),
output.DATA_PTR<float>(),
output_per_tensor.DATA_PTR<float>(),
true,
max_chunks_per_tensor);)
}
else
{
DISPATCH_FLOAT_AND_HALF( DISPATCH_FLOAT_AND_HALF(
tensor_lists[0][0].scalar_type(), 0, "multi_tensor_l2norm_cuda", tensor_lists[0][0].scalar_type(), 0, "multi_tensor_l2norm_cuda",
multi_tensor_apply<1>( multi_tensor_apply<1>(
BLOCK_SIZE, BLOCK_SIZE, chunk_size, noop_flag, tensor_lists,
chunk_size, L2NormFunctor<scalar_t_0>(), output.DATA_PTR<float>(),
noop_flag, output_per_tensor.DATA_PTR<float>(), true, max_chunks_per_tensor);)
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(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
@ -441,15 +379,9 @@ void multi_tensor_norm_out_cuda(
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;
} }