From 2a90922f01e2fa9861dda0c7a769cfed09658167 Mon Sep 17 00:00:00 2001 From: Hubert Lu <55214931+hubertlu-tw@users.noreply.github.com> Date: Thu, 5 May 2022 14:19:21 -0700 Subject: [PATCH] Using vectorized loads (float2) for fp16 to improve performance (#11390) --- .../contrib_ops/rocm/bert/fast_gelu_impl.cu | 139 ++++++++++++++---- 1 file changed, 113 insertions(+), 26 deletions(-) diff --git a/onnxruntime/contrib_ops/rocm/bert/fast_gelu_impl.cu b/onnxruntime/contrib_ops/rocm/bert/fast_gelu_impl.cu index f2c2abf950..b973073bae 100644 --- a/onnxruntime/contrib_ops/rocm/bert/fast_gelu_impl.cu +++ b/onnxruntime/contrib_ops/rocm/bert/fast_gelu_impl.cu @@ -1,7 +1,7 @@ /* The implementation of this file is based on gelu plugin in TensorRT demo: https://github.com/NVIDIA/TensorRT/tree/release/5.1/demo/BERT/ - + Copyright 2019 NVIDIA Corporation Licensed under the Apache License, Version 2.0 (the "License"); @@ -43,12 +43,10 @@ constexpr float one = 1.0; constexpr float two = 2.0; template -__global__ void FastGeluKernel(const T a, const T b, const T c, int input_length, int bias_length, const T* input, const T* bias, T* output) { +__global__ void FastGeluKernel(const T a, const T b, const T c, const T oneT, const T twoT, + int input_length, int bias_length, const T* input, const T* bias, T* output) { const int idx = blockIdx.x * TPB + threadIdx.x; - const T twoT = T(two); - const T oneT = T(one); - if (idx < input_length) { const T x = input[idx]; const T in = (bias == nullptr) ? x : (x + bias[idx % bias_length]); @@ -63,12 +61,11 @@ __global__ void FastGeluKernel(const T a, const T b, const T c, int input_length } template -__global__ void FastGeluKernel2(const half2 a, const half2 b, const half2 c, int input_length, int bias_length, const half2* input, const half2* bias, half2* output) { +__global__ void FastGeluKernel2(const half2 a, const half2 b, const half2 c, const half2 one2, const half2 two2, + int input_length, int bias_length, const half2* input, const half2* bias, + half2* output) { const int idx = blockIdx.x * TPB + threadIdx.x; - const half2 two2 = __floats2half2_rn(two, two); - const half2 one2 = __floats2half2_rn(one, one); - if (idx < input_length) { const half2 x = input[idx]; const half2 in = (bias == nullptr) ? x : (x + bias[idx % bias_length]); @@ -82,32 +79,120 @@ __global__ void FastGeluKernel2(const half2 a, const half2 b, const half2 c, int } } +template +__global__ void FastGeluKernel4Bias(const half2 a, const half2 b, const half2 c, const half2 one2, const half2 two2, + int input_length, int bias_length, const float2* input, const float2* bias, + float2* output) { + const int idx = blockIdx.x * TPB + threadIdx.x; + + if (idx < input_length) { + float2 input_vec = input[idx]; + float2 bias_vec = bias[idx % bias_length]; + float2 output_vec = output[idx]; + + half2* input_half = reinterpret_cast(&input_vec); + half2* bias_half = reinterpret_cast(&bias_vec); + half2* output_half = reinterpret_cast(&output_vec); + + half2 lo_data = input_half[0]; + half2 hi_data = input_half[1]; + half2 lo_bias = bias_half[0]; + half2 hi_bias = bias_half[1]; + + lo_data += lo_bias; + hi_data += hi_bias; + + const half2 lo_u = two2 * lo_data * (c * lo_data * lo_data + b); + const half2 hi_u = two2 * hi_data * (c * hi_data * hi_data + b); + const half2 lo_emu = h2exp(-lo_u); + const half2 hi_emu = h2exp(-hi_u); + const half2 lo_cdf = a + a * (two2/(one2 + lo_emu) - one2); + const half2 hi_cdf = a + a * (two2/(one2 + hi_emu) - one2); + + output_half[0] = lo_data * lo_cdf; + output_half[1] = hi_data * hi_cdf; + + output[idx] = output_vec; + } +} + +template +__global__ void FastGeluKernel4(const half2 a, const half2 b, const half2 c, const half2 one2, const half2 two2, + int input_length, const float2* input, float2* output) { + const int idx = blockIdx.x * TPB + threadIdx.x; + + if (idx < input_length) { + float2 input_vec = input[idx]; + float2 output_vec = output[idx]; + + half2* input_half = reinterpret_cast(&input_vec); + half2* output_half = reinterpret_cast(&output_vec); + + half2 lo_data = input_half[0]; + half2 hi_data = input_half[1]; + + const half2 lo_u = two2 * lo_data * (c * lo_data * lo_data + b); + const half2 hi_u = two2 * hi_data * (c * hi_data * hi_data + b); + const half2 lo_emu = h2exp(-lo_u); + const half2 hi_emu = h2exp(-hi_u); + const half2 lo_cdf = a + a * (two2/(one2 + lo_emu) - one2); + const half2 hi_cdf = a + a * (two2/(one2 + hi_emu) - one2); + + output_half[0] = lo_data * lo_cdf; + output_half[1] = hi_data * hi_cdf; + + output[idx] = output_vec; + } +} + template <> -bool LaunchFastGeluKernel(const hipDeviceProp_t& prop, hipStream_t stream, int input_length, int bias_length, const float* input, const float* bias, float* output, bool /*use_half2*/) { +bool LaunchFastGeluKernel(const hipDeviceProp_t& prop, hipStream_t stream, int input_length, int bias_length, + const float* input, const float* bias, float* output, bool /*use_half2*/) { constexpr int blockSize = 256; const int gridSize = (input_length + blockSize - 1) / blockSize; - hipLaunchKernelGGL(HIP_KERNEL_NAME(FastGeluKernel), dim3(gridSize), dim3(blockSize), 0, stream, A, B, C, input_length, bias_length, input, bias, output); + hipLaunchKernelGGL(HIP_KERNEL_NAME(FastGeluKernel), dim3(gridSize), dim3(blockSize), 0, + stream, A, B, C, one, two, input_length, bias_length, input, bias, output); return HIP_CALL(hipPeekAtLastError()); } template <> -bool LaunchFastGeluKernel(const hipDeviceProp_t& prop, hipStream_t stream, int input_length, int bias_length, const half* input, const half* bias, half* output, bool use_half2) { +bool LaunchFastGeluKernel(const hipDeviceProp_t& prop, hipStream_t stream, int input_length, int bias_length, + const half* input, const half* bias, half* output, bool use_half2) { constexpr int blockSize = 256; - - if (use_half2 && 0 == (bias_length & 1) && prop.major >= 7) { - const int n = input_length / 2; - const int gridSize = (n + blockSize - 1) / blockSize; - const half2 A2 = __floats2half2_rn(A, A); - const half2 B2 = __floats2half2_rn(B, B); - const half2 C2 = __floats2half2_rn(C, C); - const half2* input2 = reinterpret_cast(input); - const half2* bias2 = reinterpret_cast(bias); - half2* output2 = reinterpret_cast(output); - hipLaunchKernelGGL(HIP_KERNEL_NAME(FastGeluKernel2), dim3(gridSize), dim3(blockSize), 0, stream, A2, B2, C2, n, bias_length / 2, input2, bias2, output2); + if (use_half2 && prop.major >= 7 && (0 == (bias_length % 4) || 0 == (bias_length & 1))) { + const half2 A2 = __float2half2_rn(A); + const half2 B2 = __float2half2_rn(B); + const half2 C2 = __float2half2_rn(C); + const half2 one2 = __float2half2_rn(one); + const half2 two2 = __float2half2_rn(two); + if (0 == (bias_length % 4)) { + const int n = input_length / 4; + const int gridSize = (n + blockSize - 1) / blockSize; + const float2* input4 = reinterpret_cast(input); + const float2* bias4 = reinterpret_cast(bias); + float2* output4 = reinterpret_cast(output); + if (bias == nullptr) + hipLaunchKernelGGL(HIP_KERNEL_NAME(FastGeluKernel4), dim3(gridSize), dim3(blockSize), 0, + stream, A2, B2, C2, one2, two2, n, input4, output4); + else + hipLaunchKernelGGL(HIP_KERNEL_NAME(FastGeluKernel4Bias), dim3(gridSize), dim3(blockSize), 0, + stream, A2, B2, C2, one2, two2, n, bias_length / 4, input4, bias4, output4); + } else { + const int n = input_length / 2; + const int gridSize = (n + blockSize - 1) / blockSize; + const half2* input2 = reinterpret_cast(input); + const half2* bias2 = reinterpret_cast(bias); + half2* output2 = reinterpret_cast(output); + hipLaunchKernelGGL(HIP_KERNEL_NAME(FastGeluKernel2), dim3(gridSize), dim3(blockSize), 0, + stream, A2, B2, C2, one2, two2, n, bias_length / 2, input2, bias2, output2); + } } else { const int gridSize = (input_length + blockSize - 1) / blockSize; - hipLaunchKernelGGL(HIP_KERNEL_NAME(FastGeluKernel), dim3(gridSize), dim3(blockSize), 0, stream, A, B, C, input_length, bias_length, input, bias, output); + const half oneT = half(one); + const half twoT = half(two); + hipLaunchKernelGGL(HIP_KERNEL_NAME(FastGeluKernel), dim3(gridSize), dim3(blockSize), 0, + stream, A, B, C, oneT, twoT, input_length, bias_length, input, bias, output); } return HIP_CALL(hipPeekAtLastError()); @@ -118,8 +203,10 @@ bool LaunchFastGeluKernel(const hipDeviceProp_t& prop, hipStream_t stream, int i const BFloat16* input, const BFloat16* bias, BFloat16* output, bool /*use_half2*/) { constexpr int blockSize = 256; const int gridSize = (input_length + blockSize - 1) / blockSize; - hipLaunchKernelGGL(HIP_KERNEL_NAME(FastGeluKernel), dim3(gridSize), dim3(blockSize), 0, stream, - A, B, C, input_length, bias_length, input, bias, output); + const BFloat16 oneT = BFloat16(one); + const BFloat16 twoT = BFloat16(two); + hipLaunchKernelGGL(HIP_KERNEL_NAME(FastGeluKernel), dim3(gridSize), dim3(blockSize), 0, + stream, A, B, C, oneT, twoT, input_length, bias_length, input, bias, output); return HIP_CALL(hipPeekAtLastError()); }