diff --git a/onnxruntime/contrib_ops/rocm/bert/fast_gelu_impl.cu b/onnxruntime/contrib_ops/rocm/bert/fast_gelu_impl.cu index fc8ff3dcc6..7982fdd497 100644 --- a/onnxruntime/contrib_ops/rocm/bert/fast_gelu_impl.cu +++ b/onnxruntime/contrib_ops/rocm/bert/fast_gelu_impl.cu @@ -41,23 +41,26 @@ namespace rocm { template Status LaunchFastGeluKernel(hipStream_t stream, int input_length, int bias_length, - const T* input, const T* bias, T* output, bool tuning) { - static FastGeluTunableOp op; + const T* input, const T* bias, T* output, bool tuning) { + FastGeluParams params(stream, input, bias, output, input_length, bias_length); + if (tuning) { + static FastGeluTunableOp op; op.EnableTuning(); + return op(¶ms); } - FastGeluParams op_params(stream, input, bias, output, input_length, bias_length); - return op(&op_params); + + return FastGeluStaticSelection(¶ms); } template Status LaunchFastGeluKernel(hipStream_t stream, int input_length, int bias_length, - const float* input, const float* bias, float* output, bool tuning); + const float* input, const float* bias, float* output, bool tuning); template Status LaunchFastGeluKernel(hipStream_t stream, int input_length, int bias_length, - const BFloat16* input, const BFloat16* bias, BFloat16* output, bool tuning); + const BFloat16* input, const BFloat16* bias, BFloat16* output, bool tuning); template Status LaunchFastGeluKernel(hipStream_t stream, int input_length, int bias_length, - const half* input, const half* bias, half* output, bool tuning); + const half* input, const half* bias, half* output, bool tuning); } // namespace rocm } // namespace contrib diff --git a/onnxruntime/contrib_ops/rocm/bert/fast_gelu_impl.h b/onnxruntime/contrib_ops/rocm/bert/fast_gelu_impl.h index 63eb6b052f..5c97066c77 100644 --- a/onnxruntime/contrib_ops/rocm/bert/fast_gelu_impl.h +++ b/onnxruntime/contrib_ops/rocm/bert/fast_gelu_impl.h @@ -14,7 +14,7 @@ namespace rocm { template Status LaunchFastGeluKernel(hipStream_t stream, int input_length, int bias_length, - const T* input, const T* bias, T* output, bool use_half2); + const T* input, const T* bias, T* output, bool tuning); } // namespace rocm } // namespace contrib diff --git a/onnxruntime/contrib_ops/rocm/bert/fast_gelu_impl_kernel.h b/onnxruntime/contrib_ops/rocm/bert/fast_gelu_impl_kernel.h index 6c2b981f60..6e1bfc75fe 100644 --- a/onnxruntime/contrib_ops/rocm/bert/fast_gelu_impl_kernel.h +++ b/onnxruntime/contrib_ops/rocm/bert/fast_gelu_impl_kernel.h @@ -15,18 +15,18 @@ __global__ void FastGeluKernel(int input_length, int bias_length, const T* input const int idx = blockIdx.x * TPB + threadIdx.x; // constants for approximating the normal cdf const T a = T(0.5f); - const T b = T(0.7978845608028654f); // sqrt(2.0/M_PI) + const T b = T(0.7978845608028654f); // sqrt(2.0/M_PI) const T c = T(0.035677408136300125f); // 0.044715 * sqrt(2.0/M_PI) const T oneT = T(1.0f); const T twoT = T(2.0f); if (idx < input_length) { const T x = input[idx]; - const T in = (bias == nullptr) ? x : (x + bias[idx % bias_length]); + const T in = (bias == nullptr) ? x : (T)(x + bias[idx % bias_length]); // const T cdf = a + a * _Tanh(in * (c * in * in + b)); const T u = twoT * in * (c * in * in + b); const T emu = __expf(-u); - const T cdf = a + a * (twoT/(oneT + emu) - oneT); + const T cdf = a + a * (twoT / (oneT + emu) - oneT); output[idx] = in * cdf; } @@ -51,8 +51,8 @@ __global__ void FastGeluKernelVec(int input_length, int bias_length, const T* in VecT* output_val = reinterpret_cast(&output_v); T bias_v[ILP]; if (bias != nullptr) { - VecT* bias_val = reinterpret_cast(&bias_v); - *bias_val = *reinterpret_cast(&bias[idx % bias_length]); + VecT* bias_val = reinterpret_cast(&bias_v); + *bias_val = *reinterpret_cast(&bias[idx % bias_length]); } #pragma unroll @@ -60,7 +60,7 @@ __global__ void FastGeluKernelVec(int input_length, int bias_length, const T* in const T x = (bias == nullptr) ? input_v[i] : (T)(input_v[i] + bias_v[i]); const T u = twoT * x * (c * x * x + b); const T emu = __expf(-u); - const T cdf = a + a * (twoT/(oneT + emu) - oneT); + const T cdf = a + a * (twoT / (oneT + emu) - oneT); output_v[i] = x * cdf; } *(reinterpret_cast(&output[idx])) = *output_val; diff --git a/onnxruntime/contrib_ops/rocm/bert/fast_gelu_tunable_op.h b/onnxruntime/contrib_ops/rocm/bert/fast_gelu_tunable_op.h index ee5e1edb54..5691d15eaa 100644 --- a/onnxruntime/contrib_ops/rocm/bert/fast_gelu_tunable_op.h +++ b/onnxruntime/contrib_ops/rocm/bert/fast_gelu_tunable_op.h @@ -39,14 +39,64 @@ Status FastGeluOp(const FastGeluParams* params) { !((params->bias_length > 0 && params->bias_length % VecSize == 0 && params->input_length % VecSize == 0) || (params->bias_length == 0 && params->input_length % VecSize == 0))); - hipLaunchKernelGGL((FastGeluKernelVec), - dim3(onnxruntime::rocm::CeilDiv(params->input_length, ThreadsPerBlock * VecSize)), - dim3(ThreadsPerBlock), - 0, params->stream, - params->input_length, params->bias_length, params->input, params->bias, params->output); - auto status = hipGetLastError(); - ORT_RETURN_IF(status != hipSuccess, ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, hipGetErrorName(status))); - return Status::OK(); + FastGeluKernelVec + <<input_length, ThreadsPerBlock * VecSize)), + dim3(ThreadsPerBlock), + 0, params->stream>>>( + params->input_length, params->bias_length, params->input, params->bias, params->output); + return HIP_CALL(hipGetLastError()); +} + +template +Status FastGeluStaticSelection(const FastGeluParams* params) { + constexpr int block_size = 256; + const int grid_size = (params->input_length + block_size - 1) / block_size; + FastGeluKernel<<stream>>>( + params->input_length, params->bias_length, params->input, params->bias, params->output); + return HIP_CALL(hipGetLastError()); +} + +template <> +Status FastGeluStaticSelection(const FastGeluParams* params) { + constexpr int block_size = 256; + if (params->bias != nullptr) { + if (0 == (params->bias_length % 8) && (params->input_length >= 3145728)) { // 3145728=8*128*3072 + const int grid_size = (params->input_length / 8 + block_size - 1) / block_size; + FastGeluKernelVec<<stream>>>( + params->input_length, params->bias_length, params->input, params->bias, params->output); + } else if (0 == (params->bias_length % 4)) { + const int grid_size = (params->input_length / 4 + block_size - 1) / block_size; + FastGeluKernelVec<<stream>>>( + params->input_length, params->bias_length, params->input, params->bias, params->output); + } else if (0 == (params->bias_length % 2)) { + const int grid_size = (params->input_length / 2 + block_size - 1) / block_size; + FastGeluKernelVec<<stream>>>( + params->input_length, params->bias_length, params->input, params->bias, params->output); + } else { + const int grid_size = (params->input_length + block_size - 1) / block_size; + FastGeluKernel<<stream>>>( + params->input_length, params->bias_length, params->input, params->bias, params->output); + } + } else { + if (0 == (params->input_length % 8) && (params->input_length >= 3145728)) { // 3145728=8*128*3072 + const int grid_size = (params->input_length / 8 + block_size - 1) / block_size; + FastGeluKernelVec<<stream>>>( + params->input_length, params->bias_length, params->input, params->bias, params->output); + } else if (0 == (params->input_length % 4)) { + const int grid_size = (params->input_length / 4 + block_size - 1) / block_size; + FastGeluKernelVec<<stream>>>( + params->input_length, params->bias_length, params->input, params->bias, params->output); + } else if (0 == (params->input_length % 2)) { + const int grid_size = (params->input_length / 2 + block_size - 1) / block_size; + FastGeluKernelVec<<stream>>>( + params->input_length, params->bias_length, params->input, params->bias, params->output); + } else { + const int grid_size = (params->input_length + block_size - 1) / block_size; + FastGeluKernel<<stream>>>( + params->input_length, params->bias_length, params->input, params->bias, params->output); + } + } + return HIP_CALL(hipGetLastError()); } #define ADD_OP(threads_per_block) \ @@ -60,6 +110,7 @@ template class FastGeluTunableOp : public onnxruntime::rocm::tunable::TunableOp> { public: FastGeluTunableOp() { + this->ops_.emplace_back(FastGeluStaticSelection); ADD_OP(64); ADD_OP(128); ADD_OP(192); @@ -69,8 +120,8 @@ class FastGeluTunableOp : public onnxruntime::rocm::tunable::TunableOpSetDefaultId(15); + // NOTE: the 1st kernel is FastGelu Original implementation. + this->SetDefaultId(0); } }; diff --git a/onnxruntime/contrib_ops/rocm/bert/skip_layer_norm_impl.cu b/onnxruntime/contrib_ops/rocm/bert/skip_layer_norm_impl.cu index f324caf919..d54166c2c1 100644 --- a/onnxruntime/contrib_ops/rocm/bert/skip_layer_norm_impl.cu +++ b/onnxruntime/contrib_ops/rocm/bert/skip_layer_norm_impl.cu @@ -46,15 +46,15 @@ Status LaunchSkipLayerNormKernel( // this must be true because element_count is the total size of the tensor assert(element_count % ld == 0); - SkipLayerNormParams op_params(stream, output, input, skip, gamma, beta, bias, epsilon, ld, element_count); - static SkipLayerNormTunableOp op; + SkipLayerNormParams params(stream, output, input, skip, gamma, beta, bias, epsilon, ld, element_count); - // If disable tuning, the default implementation is SkipLayerNormStaticSelection. if (tuning) { + static SkipLayerNormTunableOp op; op.EnableTuning(); + return op(¶ms); } - return op(&op_params); + return SkipLayerNormStaticSelection(¶ms); } template Status LaunchSkipLayerNormKernel(hipStream_t stream, float* output, const float* input,