From dab036ee8e2b29695a38dda0da629b51e878e4b0 Mon Sep 17 00:00:00 2001 From: Molly Smith Date: Wed, 7 Sep 2022 16:07:53 -0700 Subject: [PATCH 1/4] modified relu kernels --- csrc/transformer/inference/csrc/relu.cu | 71 ++++++++++++++----------- 1 file changed, 41 insertions(+), 30 deletions(-) diff --git a/csrc/transformer/inference/csrc/relu.cu b/csrc/transformer/inference/csrc/relu.cu index 0472d0db3490..1bf5fda750af 100644 --- a/csrc/transformer/inference/csrc/relu.cu +++ b/csrc/transformer/inference/csrc/relu.cu @@ -36,45 +36,56 @@ __global__ void fused_bias_relu(float* input, } } -__global__ void fused_bias_relu(__half* input, - const __half* bias, +__global__ void fused_bias_relu(float* input, + const float* bias, int total_count, int intermediate_size) { -#ifdef HALF_PRECISION_AVAILABLE - - float2* input_cast = reinterpret_cast(input); - const float2* bias_cast = reinterpret_cast(bias); - - int offset = blockIdx.x * blockDim.x + threadIdx.x; + // Input restriction: intermediate_size % vals_per_access == 0 + constexpr int granularity = 16; + constexpr int vals_per_access = granularity / sizeof(float); + const int offset = (blockIdx.x * blockDim.x + threadIdx.x) * vals_per_access; if (offset < total_count) { - float2 vals_vec = input_cast[offset]; - float2 bias_vec = bias_cast[offset % intermediate_size]; - - __half2* vals_half = reinterpret_cast<__half2*>(&vals_vec); - __half2* bias_half = reinterpret_cast<__half2*>(&bias_vec); - - float2 low_data = __half22float2(vals_half[0]); - float2 high_data = __half22float2(vals_half[1]); + float data[vals_per_access]; + float data_bias[vals_per_access]; + mem_access::load_global(data, input + offset); + mem_access::load_global(data_bias, bias + (offset % intermediate_size)); - float2 low_bias = __half22float2(bias_half[0]); - float2 high_bias = __half22float2(bias_half[1]); +#pragma unroll + for (int i = 0; i < vals_per_access; i++) { data[i] = relu(data[i] + data_bias[i]); } - low_data.x += low_bias.x; - low_data.y += low_bias.y; - high_data.x += high_bias.x; - high_data.y += high_bias.y; - - low_data.x = relu(low_data.x); - low_data.y = relu(low_data.y); - high_data.x = relu(high_data.x); - high_data.y = relu(high_data.y); + mem_access::store_global(input + offset, data); + } +} - vals_half[0] = __float22half2_rn(low_data); - vals_half[1] = __float22half2_rn(high_data); +__global__ void fused_bias_relu(__half* input, + const __half* bias, + int total_count, + int intermediate_size) +{ + // Input restriction: intermediate_size % vals_per_access == 0 + // This kernel doubles the per-thread ALU workload as compared to the float implementation +#ifdef HALF_PRECISION_AVAILABLE + constexpr int granularity = 16; + constexpr int vals_per_access = granularity / sizeof(__half); + int offset = (blockIdx.x * blockDim.x + threadIdx.x) * vals_per_access; - input_cast[offset] = vals_vec; + if (offset < total_count) { + // Divide by 2 since we store two values per __half2 + __half2 data[vals_per_access / 2]; + __half2 bias_data[vals_per_access / 2]; + mem_access::load_global(data, input + offset); + mem_access::load_global(bias_data, bias + (offset % intermediate_size)); + +#pragma unroll + for (int i = 0; i < vals_per_access / 2; i++) { + float2 data_f = __half22float2(data[i]); + float2 bias_f = __half22float2(bias_data[i]); + data[i] = __floats2half2_rn(relu(data_f.x + bias_f.x), relu(data_f.y + bias_f.y)); + } + + mem_access::store_global(input + offset, data); } #endif } From 3fdbb8dc2b045584cc71bf74d99910f94b39cd1a Mon Sep 17 00:00:00 2001 From: Molly Smith Date: Wed, 7 Sep 2022 16:13:42 -0700 Subject: [PATCH 2/4] Remove old float fused_bias_relu --- csrc/transformer/inference/csrc/relu.cu | 27 ------------------------- 1 file changed, 27 deletions(-) diff --git a/csrc/transformer/inference/csrc/relu.cu b/csrc/transformer/inference/csrc/relu.cu index 1bf5fda750af..dbd5c7c8153a 100644 --- a/csrc/transformer/inference/csrc/relu.cu +++ b/csrc/transformer/inference/csrc/relu.cu @@ -9,33 +9,6 @@ Copyright 2022 The Microsoft DeepSpeed Team inline __device__ float relu(const float x) { return x < 0 ? 0 : x; } -__global__ void fused_bias_relu(float* input, - const float* bias, - int total_count, - int intermediate_size) -{ - float4* input_cast = reinterpret_cast(input); - const float4* bias_cast = reinterpret_cast(bias); - int offset = blockIdx.x * blockDim.x + threadIdx.x; - - if (offset < total_count) { - float4 data = input_cast[offset]; - float4 bias_data = bias_cast[offset % intermediate_size]; - - data.x += bias_data.x; - data.y += bias_data.y; - data.z += bias_data.z; - data.w += bias_data.w; - - data.x = relu(data.x); - data.y = relu(data.y); - data.z = relu(data.z); - data.w = relu(data.w); - - input_cast[offset] = data; - } -} - __global__ void fused_bias_relu(float* input, const float* bias, int total_count, From 75a31296e69b92025b2e470673f06a3d315d1fd0 Mon Sep 17 00:00:00 2001 From: Molly Smith Date: Thu, 8 Sep 2022 16:29:41 -0700 Subject: [PATCH 3/4] Include header --- csrc/transformer/inference/csrc/relu.cu | 2 ++ 1 file changed, 2 insertions(+) diff --git a/csrc/transformer/inference/csrc/relu.cu b/csrc/transformer/inference/csrc/relu.cu index dbd5c7c8153a..42ccd89c4814 100644 --- a/csrc/transformer/inference/csrc/relu.cu +++ b/csrc/transformer/inference/csrc/relu.cu @@ -3,7 +3,9 @@ Copyright 2022 The Microsoft DeepSpeed Team */ #include "inference_cuda_layers.h" +#include "memory_access_utils.h" +namespace cg = cooperative_groups; #define MAX_CAP 4 #define MAX_SEQ 2048 From fc745349d514f12ba0bdbfcfb331234c047bafc9 Mon Sep 17 00:00:00 2001 From: Molly Smith Date: Thu, 8 Sep 2022 16:49:36 -0700 Subject: [PATCH 4/4] Update launch_bias_relu --- csrc/transformer/inference/csrc/relu.cu | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/csrc/transformer/inference/csrc/relu.cu b/csrc/transformer/inference/csrc/relu.cu index 42ccd89c4814..26445b74e87c 100644 --- a/csrc/transformer/inference/csrc/relu.cu +++ b/csrc/transformer/inference/csrc/relu.cu @@ -72,13 +72,16 @@ void launch_bias_relu(T* input, int batch_size, cudaStream_t stream) { - int total_count = batch_size * (intermediate_size / 4); - int threads = 1024; // intermediate_size / iterations / 4; + constexpr int threads = 1024; + constexpr int granularity = 16; + + const int total_count = batch_size * intermediate_size; + const int elems_per_block = threads * (granularity / sizeof(T)); dim3 block_dims(threads); - dim3 grid_dims(((total_count - 1) / 1024 + 1)); // (batch_size); + dim3 grid_dims((total_count + elems_per_block - 1) / elems_per_block); fused_bias_relu<<>>( - input, bias, total_count, intermediate_size / 4); + input, bias, total_count, intermediate_size); } template void launch_bias_relu(float*, const float*, int, int, cudaStream_t);