From 7eb6dbe7d862117b865cbf51a0329db0dae37779 Mon Sep 17 00:00:00 2001 From: PeixuanZuo <94887879+PeixuanZuo@users.noreply.github.com> Date: Fri, 24 Mar 2023 19:31:14 +0800 Subject: [PATCH] [ROCm] Add compute type for Skiplayernorm to fix ROCm CI (#15192) - Add compute type for Skiplayernorm to fix ROCm CI and get more accurate results. SkipLayerNorm: type T: input, skip, bias type U: epsilon, compute result type V: output, beta, gamma - refactor the usage of aligned_vector, reduce the usage of `reinterpret_cast`. --- .../contrib_ops/rocm/bert/layer_norm.cuh | 89 ++++++-------- .../contrib_ops/rocm/bert/skip_layer_norm.cc | 2 +- .../rocm/bert/skip_layer_norm_impl.cu | 17 +-- .../rocm/bert/skip_layer_norm_impl.h | 8 +- .../rocm/bert/skip_layer_norm_impl_kernel.h | 112 ++++++++---------- .../rocm/bert/skip_layer_norm_tunable_op.h | 81 +++++++------ .../kernels/rocm/skip_layer_norm.cu | 22 ++-- .../kernels/skip_layer_norm_test.py | 2 +- 8 files changed, 157 insertions(+), 176 deletions(-) diff --git a/onnxruntime/contrib_ops/rocm/bert/layer_norm.cuh b/onnxruntime/contrib_ops/rocm/bert/layer_norm.cuh index 169d0ed3b0..9b7dbd5291 100644 --- a/onnxruntime/contrib_ops/rocm/bert/layer_norm.cuh +++ b/onnxruntime/contrib_ops/rocm/bert/layer_norm.cuh @@ -80,16 +80,16 @@ struct KeyValuePairSum { } }; -template +template __device__ inline void LayerNorm( - const hipcub::KeyValuePair& thread_data, const int ld, const int offset, const T* beta, - const T* gamma, const T epsilon, T* output) { + const hipcub::KeyValuePair& thread_data, const int ld, const int offset, const V* beta, + const V* gamma, const U epsilon, V* output) { // Assuming thread_data is already divided by ld - using BlockReduce = hipcub::BlockReduce, TPB>; + using BlockReduce = hipcub::BlockReduce, TPB>; __shared__ typename BlockReduce::TempStorage temp_storage; - __shared__ T mu; // mean - __shared__ T rsigma; // 1 / std.dev. + __shared__ U mu; // mean + __shared__ U rsigma; // 1 / std.dev. KeyValuePairSum pair_sum; const auto sum_kv = BlockReduce(temp_storage).Reduce(thread_data, pair_sum); @@ -102,23 +102,23 @@ __device__ inline void LayerNorm( for (int i = threadIdx.x; i < ld; i += TPB) { const int idx = offset + i; - const T val = output[idx]; - const T g(gamma[i]); - const T b = (nullptr == beta) ? (T)0 : beta[i]; - output[idx] = g * (val - mu) * rsigma + b; + const U val = static_cast(output[idx]); + const U g = static_cast(gamma[i]); + const U b = (nullptr == beta) ? U(0.f) : static_cast(beta[i]); + output[idx] = static_cast(g * (val - mu) * rsigma + b); } } -template +template __device__ inline void LayerNormVec( - const hipcub::KeyValuePair& thread_data, const int ld, const int offset, const T* beta, - const T* gamma, const T epsilon, T* output) { + const hipcub::KeyValuePair& thread_data, const int ld, const int offset, const V* beta, + const V* gamma, const U epsilon, V* output) { // Assuming thread_data is already divided by ld - using VecT = aligned_vector; - using BlockReduce = hipcub::BlockReduce, TPB>; + using VecV = aligned_vector; + using BlockReduce = hipcub::BlockReduce, TPB>; __shared__ typename BlockReduce::TempStorage temp_storage; - __shared__ T mu; // mean - __shared__ T rsigma; // 1 / std.dev. + __shared__ U mu; // mean + __shared__ U rsigma; // 1 / std.dev. KeyValuePairSum pair_sum; const auto sum_kv = BlockReduce(temp_storage).Reduce(thread_data, pair_sum); @@ -130,44 +130,37 @@ __device__ inline void LayerNormVec( __syncthreads(); if (ILP * threadIdx.x < ld) { - T beta_v[ILP], gamma_v[ILP], output_v[ILP]; - VecT* gamma_val = reinterpret_cast(&gamma_v); - VecT* output_val = reinterpret_cast(&output_v); - for (int i = threadIdx.x * ILP; i < ld; i += TPB * ILP) { int idx = offset + i; - if (beta != nullptr) { - VecT* beta_val = reinterpret_cast(&beta_v); - *beta_val = *reinterpret_cast(&beta[i]); - } - *gamma_val = *reinterpret_cast(&gamma[i]); - *output_val = *reinterpret_cast(&output[idx]); + const VecV beta_v = (beta != nullptr) ? *reinterpret_cast(beta + i) : VecV(); + const VecV gamma_v = *reinterpret_cast(gamma + i); + VecV output_v = *reinterpret_cast(output + idx); + #pragma unroll for (int k = 0; k < ILP; k++) { - output_v[k] = (beta != nullptr) ? gamma_v[k] * (output_v[k] - mu) * rsigma + beta_v[k] : - gamma_v[k] * (output_v[k] - mu) * rsigma; + output_v.val[k] = (beta != nullptr) ? U(gamma_v.val[k]) * (U(output_v.val[k]) - mu) * rsigma + U(beta_v.val[k]) : + U(gamma_v.val[k]) * (U(output_v.val[k]) - mu) * rsigma; } - *(reinterpret_cast(&output[idx])) = *reinterpret_cast(&output_v[0]); + *(reinterpret_cast(output + idx)) = output_v; } } } -template -__device__ inline void LayerNormSmall(const T* input_v, const hipcub::KeyValuePair& thread_data, - const int ld, const int idx, const T* beta, const T* gamma, - const T epsilon, T* output) { +template +__device__ inline void LayerNormSmall(const T* input_v, const hipcub::KeyValuePair& thread_data, + const int ld, const int idx, const V* beta, const V* gamma, + const U epsilon, V* output) { // Assuming thread_data is already divided by ld // Small settings: the block covers the leading dimension TPB >= ld. The input // value is available in a register - using VecT = aligned_vector; - using BlockReduce = hipcub::BlockReduce, TPB>; + using VecV = aligned_vector; + using BlockReduce = hipcub::BlockReduce, TPB>; __shared__ typename BlockReduce::TempStorage temp_storage; - __shared__ T mu; // mean - __shared__ T rsigma; // 1 / std.dev. - T beta_v[ILP], gamma_v[ILP], output_v[ILP]; + __shared__ U mu; // mean + __shared__ U rsigma; // 1 / std.dev. KeyValuePairSum pair_sum; - const hipcub::KeyValuePair sum_kv = BlockReduce(temp_storage).Reduce(thread_data, pair_sum); + const hipcub::KeyValuePair sum_kv = BlockReduce(temp_storage).Reduce(thread_data, pair_sum); if (threadIdx.x == 0) { mu = sum_kv.key; @@ -176,20 +169,16 @@ __device__ inline void LayerNormSmall(const T* input_v, const hipcub::KeyValuePa __syncthreads(); if (ILP * threadIdx.x < ld) { - if (beta != nullptr) { - VecT* beta_val = reinterpret_cast(&beta_v); - *beta_val = *reinterpret_cast(&beta[threadIdx.x * ILP]); - } - - VecT* gamma_val = reinterpret_cast(&gamma_v); - *gamma_val = *reinterpret_cast(&gamma[threadIdx.x * ILP]); + const VecV beta_v = (beta != nullptr) ? *reinterpret_cast(beta + threadIdx.x * ILP) : VecV(); + const VecV gamma_v = *reinterpret_cast(gamma + threadIdx.x * ILP); + VecV output_v; #pragma unroll for (int i = 0; i < ILP; i++) { - output_v[i] = (beta != nullptr) ? gamma_v[i] * (input_v[i] - mu) * rsigma + beta_v[i] : - gamma_v[i] * (input_v[i] - mu) * rsigma; + output_v.val[i] = (beta != nullptr) ? U(gamma_v.val[i]) * (U(input_v[i]) - mu) * rsigma + U(beta_v.val[i]) : + U(gamma_v.val[i]) * (U(input_v[i]) - mu) * rsigma; } - *(reinterpret_cast(&output[idx])) = *reinterpret_cast(&output_v[0]); + *(reinterpret_cast(output + idx)) = output_v; } } diff --git a/onnxruntime/contrib_ops/rocm/bert/skip_layer_norm.cc b/onnxruntime/contrib_ops/rocm/bert/skip_layer_norm.cc index a254a8c04f..24dbb87b50 100644 --- a/onnxruntime/contrib_ops/rocm/bert/skip_layer_norm.cc +++ b/onnxruntime/contrib_ops/rocm/bert/skip_layer_norm.cc @@ -101,7 +101,7 @@ Status SkipLayerNorm::ComputeInternal(OpKernelContext* ctx) const { int64_t element_count = input_dims[0] * sequence_length * hidden_size; typedef typename ToHipType::MappedType HipT; - return LaunchSkipLayerNormKernel( + return LaunchSkipLayerNormKernel( GetTuningContext(), Stream(ctx), reinterpret_cast(output->MutableData()), 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 bf33f940b3..c6ac1196ba 100644 --- a/onnxruntime/contrib_ops/rocm/bert/skip_layer_norm_impl.cu +++ b/onnxruntime/contrib_ops/rocm/bert/skip_layer_norm_impl.cu @@ -39,30 +39,31 @@ namespace onnxruntime { namespace contrib { namespace rocm { -template +template Status LaunchSkipLayerNormKernel( - RocmTuningContext* tuning_ctx, hipStream_t stream, T* output, T* skip_input_bias_add_output, const T* input, - const T* skip, const T* gamma, const T* beta, const T* bias, float epsilon, int ld, int element_count) { + RocmTuningContext* tuning_ctx, hipStream_t stream, V* output, T* skip_input_bias_add_output, const T* input, + const T* skip, const V* gamma, const V* beta, const T* bias, float epsilon, int ld, int element_count) { // this must be true because element_count is the total size of the tensor assert(element_count % ld == 0); - SkipLayerNormParams params(tuning_ctx, stream, output, skip_input_bias_add_output, input, skip, gamma, beta, bias, epsilon, ld, element_count); + SkipLayerNormParams params(tuning_ctx, stream, output, skip_input_bias_add_output, input, skip, + gamma, beta, bias, epsilon, ld, element_count); if (tuning_ctx->IsTunableOpEnabled()) { - static SkipLayerNormTunableOp op; + static SkipLayerNormTunableOp op; return op(¶ms); } - return SkipLayerNormStaticSelection(¶ms); + return SkipLayerNormStaticSelection(¶ms); } -template Status LaunchSkipLayerNormKernel( +template Status LaunchSkipLayerNormKernel( RocmTuningContext* tuning_ctx, hipStream_t stream, float* output, float* skip_input_bias_add_output, const float* input, const float* skip, const float* gamma, const float* beta, const float* bias, float epsilon, int ld, int element_count); -template Status LaunchSkipLayerNormKernel( +template Status LaunchSkipLayerNormKernel( RocmTuningContext* tuning_ctx, hipStream_t stream, half* output, half* skip_input_bias_add_output, const half* input, const half* skip, const half* gamma, const half* beta, const half* bias, float epsilon, int ld, diff --git a/onnxruntime/contrib_ops/rocm/bert/skip_layer_norm_impl.h b/onnxruntime/contrib_ops/rocm/bert/skip_layer_norm_impl.h index 911164af92..a1c09142fe 100644 --- a/onnxruntime/contrib_ops/rocm/bert/skip_layer_norm_impl.h +++ b/onnxruntime/contrib_ops/rocm/bert/skip_layer_norm_impl.h @@ -10,16 +10,16 @@ namespace onnxruntime { namespace contrib { namespace rocm { -template +template Status LaunchSkipLayerNormKernel( RocmTuningContext* tuning, hipStream_t stream, - T* output, // output tensor + V* output, // output tensor T* skip_input_bias_add_output, // optional output tensor const T* input, // input tensor const T* skip, // skip tensor - const T* gamma, // Layer normalization gamma tensor - const T* beta, // Layer normalization beta tensor + const V* gamma, // Layer normalization gamma tensor + const V* beta, // Layer normalization beta tensor const T* bias, // Layer normalization beta tensor float epsilon, // Layer normalization epsilon int hidden_size, // hidden size, it is the leading dimension (ld) diff --git a/onnxruntime/contrib_ops/rocm/bert/skip_layer_norm_impl_kernel.h b/onnxruntime/contrib_ops/rocm/bert/skip_layer_norm_impl_kernel.h index aeb954de09..ee8959458b 100644 --- a/onnxruntime/contrib_ops/rocm/bert/skip_layer_norm_impl_kernel.h +++ b/onnxruntime/contrib_ops/rocm/bert/skip_layer_norm_impl_kernel.h @@ -23,133 +23,125 @@ half maybe2half(float x) { return __float2half_rn(x); } -template +template __global__ void SkipLayerNormKernel( - const int ld, const T* input, const T* skip, const T* beta, const T* gamma, const T* bias, - const T epsilon, T* output, T* skip_input_bias_add_output) { - const T reverse_ld = T(1.f / ld); + const int ld, const T* input, const T* skip, const V* beta, const V* gamma, const T* bias, + const U epsilon, V* output, T* skip_input_bias_add_output) { + const U reverse_ld = U(1.f / ld); const int offset = blockIdx.x * ld; KeyValuePairSum pair_sum; // reduce x and x^2 - hipcub::KeyValuePair thread_data(0, 0); + hipcub::KeyValuePair thread_data(U(0.f), U(0.f)); for (int i = threadIdx.x; i < ld; i += TPB) { const int idx = offset + i; - const T val = (bias == nullptr) ? input[idx] + skip[idx] : input[idx] + skip[idx] + bias[i]; - const T rldval = reverse_ld * val; - thread_data = pair_sum(thread_data, hipcub::KeyValuePair(rldval, rldval * val)); + const U val = (bias == nullptr) ? static_cast(input[idx]) + static_cast(skip[idx]) : + static_cast(input[idx]) + static_cast(skip[idx]) + static_cast(bias[i]); + const U rldval = reverse_ld * val; + thread_data = pair_sum(thread_data, hipcub::KeyValuePair(rldval, rldval * val)); if (skip_input_bias_add_output != nullptr) { - skip_input_bias_add_output[idx] = val; + skip_input_bias_add_output[idx] = static_cast(val); } - output[idx] = val; + output[idx] = static_cast(val); } - LayerNorm(thread_data, ld, offset, beta, gamma, epsilon, output); + LayerNorm(thread_data, ld, offset, beta, gamma, epsilon, output); } // Vectorized kernel -template +template __global__ void SkipLayerNormKernelVec( - const int ld, const T* input, const T* skip, const T* beta, const T* gamma, - const T* bias, const T epsilon, T* output, T* skip_input_bias_add_output, + const int ld, const T* input, const T* skip, const V* beta, const V* gamma, + const T* bias, const U epsilon, V* output, T* skip_input_bias_add_output, bool hasBias, bool hasSkipInputBiasAdditionOutput) { - const T reverse_ld = T(1.f / ld); + const U reverse_ld = U(1.f / ld); const int offset = blockIdx.x * ld; KeyValuePairSum pair_sum; // reduce x and x^2 - hipcub::KeyValuePair thread_data(0, 0); + hipcub::KeyValuePair thread_data(U(0.f), U(0.f)); using VecT = aligned_vector; - T input_v[ILP], skip_v[ILP], bias_v[ILP], skip_input_bias_add_output_v[ILP]; + using VecV = aligned_vector; if (threadIdx.x * ILP < ld) { - VecT* input_val = reinterpret_cast(&input_v); - VecT* skip_val = reinterpret_cast(&skip_v); - for (int i = threadIdx.x * ILP; i < ld; i += TPB * ILP) { int idx = offset + i; - *input_val = *reinterpret_cast(&input[idx]); - *skip_val = *reinterpret_cast(&skip[idx]); - if (hasBias) { - VecT* bias_val = reinterpret_cast(&bias_v); - *bias_val = *reinterpret_cast(&bias[i]); - } + const VecT input_v = *reinterpret_cast(input + idx); + const VecT skip_v = *reinterpret_cast(skip + idx); + const VecT bias_v = hasBias ? *reinterpret_cast(bias + i) : VecT(); + VecT skip_input_bias_add_output_v, output_v; #pragma unroll for (int k = 0; k < ILP; k++) { - input_v[k] += hasBias ? skip_v[k] + bias_v[k] : skip_v[k]; + const U val = hasBias ? static_cast(input_v.val[k]) + static_cast(skip_v.val[k]) + static_cast(bias_v.val[k]) : + static_cast(input_v.val[k]) + static_cast(skip_v.val[k]); + const U rldval = reverse_ld * val; if (hasSkipInputBiasAdditionOutput) { - skip_input_bias_add_output_v[k] = input_v[k]; + skip_input_bias_add_output_v.val[k] = static_cast(val); } - - const T rldval = reverse_ld * input_v[k]; - thread_data = pair_sum(thread_data, hipcub::KeyValuePair(rldval, rldval * input_v[k])); + thread_data = pair_sum(thread_data, hipcub::KeyValuePair(rldval, rldval * val)); + output_v.val[k] = static_cast(val); } if (hasSkipInputBiasAdditionOutput) { - *(reinterpret_cast(&skip_input_bias_add_output[idx])) = *reinterpret_cast(&skip_input_bias_add_output_v); + *(reinterpret_cast(skip_input_bias_add_output + idx)) = skip_input_bias_add_output_v; } - *(reinterpret_cast(&output[idx])) = *reinterpret_cast(&input_v[0]); + *(reinterpret_cast(output + idx)) = output_v; } } - LayerNormVec(thread_data, ld, offset, beta, gamma, epsilon, output); + LayerNormVec(thread_data, ld, offset, beta, gamma, epsilon, output); } // Vectorized kernel -template +template __global__ void SkipLayerNormKernelSmall( - const int ld, const T* input, const T* skip, const T* beta, const T* gamma, - const T* bias, const T epsilon, T* output, T* skip_input_bias_add_output, + const int ld, const T* input, const T* skip, const V* beta, const V* gamma, + const T* bias, const U epsilon, V* output, T* skip_input_bias_add_output, bool hasBias, bool hasSkipInputBiasAdditionOutput) { - const T rld = T(1.f / ld); + const U rld = U(1.f / ld); const int idx = blockIdx.x * ld + threadIdx.x * ILP; // grid_size = n / ld using VecT = aligned_vector; - T input_v[ILP], skip_v[ILP], bias_v[ILP], skip_input_bias_add_output_v[ILP]; - - hipcub::KeyValuePair thread_data(T(0.f), T(0.f)); + hipcub::KeyValuePair thread_data(U(0.f), U(0.f)); + VecT input_v; if (ILP * threadIdx.x < ld) { - VecT* input_val = reinterpret_cast(&input_v); - *input_val = *reinterpret_cast(&input[idx]); + input_v = *reinterpret_cast(input + idx); + const VecT skip_v = *reinterpret_cast(skip + idx); + const VecT bias_v = hasBias ? *reinterpret_cast(bias + threadIdx.x * ILP) : VecT(); + VecT skip_input_bias_add_output_v; - VecT* skip_val = reinterpret_cast(&skip_v); - *skip_val = *reinterpret_cast(&skip[idx]); - - if (hasBias) { - VecT* bias_val = reinterpret_cast(&bias_v); - *bias_val = *reinterpret_cast(&bias[threadIdx.x * ILP]); - } - - T rldval_sum = T(0.f); - T rldvalsq_sum = T(0.f); + U rldval_sum = U(0.f); + U rldvalsq_sum = U(0.f); #pragma unroll for (int i = 0; i < ILP; i++) { - input_v[i] += hasBias ? skip_v[i] + bias_v[i] : skip_v[i]; + const U val = hasBias ? static_cast(input_v.val[i]) + static_cast(skip_v.val[i]) + static_cast(bias_v.val[i]) : + static_cast(input_v.val[i]) + static_cast(skip_v.val[i]); if (hasSkipInputBiasAdditionOutput) { - skip_input_bias_add_output_v[i] = input_v[i]; + skip_input_bias_add_output_v.val[i] = static_cast(val); } - const T rldval = rld * input_v[i]; + const U rldval = rld * val; rldval_sum += rldval; - rldvalsq_sum += rldval * input_v[i]; + rldvalsq_sum += rldval * val; + input_v.val[i] = static_cast(val); } if (hasSkipInputBiasAdditionOutput) { - *(reinterpret_cast(&skip_input_bias_add_output[idx])) = *reinterpret_cast(&skip_input_bias_add_output_v); + *(reinterpret_cast(skip_input_bias_add_output + idx)) = skip_input_bias_add_output_v; } - thread_data = hipcub::KeyValuePair(rldval_sum, rldvalsq_sum); + thread_data = hipcub::KeyValuePair(rldval_sum, rldvalsq_sum); } - LayerNormSmall(input_v, thread_data, ld, idx, beta, gamma, epsilon, output); + LayerNormSmall(input_v.val, thread_data, ld, idx, beta, gamma, epsilon, output); } } // namespace rocm diff --git a/onnxruntime/contrib_ops/rocm/bert/skip_layer_norm_tunable_op.h b/onnxruntime/contrib_ops/rocm/bert/skip_layer_norm_tunable_op.h index 4354b79414..a0b2507220 100644 --- a/onnxruntime/contrib_ops/rocm/bert/skip_layer_norm_tunable_op.h +++ b/onnxruntime/contrib_ops/rocm/bert/skip_layer_norm_tunable_op.h @@ -18,76 +18,75 @@ namespace onnxruntime { namespace contrib { namespace rocm { -template +template struct SkipLayerNormParams : OpParams { - SkipLayerNormParams(RocmTuningContext* tuning_ctx, hipStream_t stream, T* output, T* skip_input_bias_add_output, const T* input, - const T* skip, const T* gamma, const T* beta, + SkipLayerNormParams(RocmTuningContext* tuning_ctx, hipStream_t stream, V* output, T* skip_input_bias_add_output, const T* input, + const T* skip, const V* gamma, const V* beta, const T* bias, float epsilon, int ld, int element_count) - : OpParams(tuning_ctx, stream), output(output), skip_input_bias_add_output(skip_input_bias_add_output), input(input), skip(skip), - gamma(gamma), beta(beta), bias(bias), epsilon(epsilon), ld(ld), element_count(element_count) {} + : OpParams(tuning_ctx, stream), output(output), skip_input_bias_add_output(skip_input_bias_add_output), input(input), skip(skip), gamma(gamma), beta(beta), bias(bias), epsilon(epsilon), ld(ld), element_count(element_count) {} std::string Signature() const override { std::string sig = std::to_string(ld) + "_" + std::to_string(element_count); return sig; } - T* output; + V* output; T* skip_input_bias_add_output; const T* input; const T* skip; - const T* gamma; - const T* beta; + const V* gamma; + const V* beta; const T* bias; float epsilon; int ld; int element_count; }; -template -Status SkipLayerNormSmallOp(const SkipLayerNormParams* params) { +template +Status SkipLayerNormSmallOp(const SkipLayerNormParams* params) { // Loosen the hard constraint for ld (hidden_size) to include more possible *Small kernels, // which could offer better performance in some combinations of ThreadsPerBlock and VecSize. TUNABLE_OP_RETURN_UNSUPPORTED_ARGUMENT_IF( !((params->ld <= 8192 && params->ld % VecSize == 0 && params->ld <= ThreadsPerBlock * VecSize && params->ld > (ThreadsPerBlock - GPU_WARP_SIZE) * VecSize))); - SkipLayerNormKernelSmall<<element_count, params->ld)), - dim3(ThreadsPerBlock), - 0, params->stream>>>( + SkipLayerNormKernelSmall<<element_count, params->ld)), + dim3(ThreadsPerBlock), + 0, params->stream>>>( params->ld, params->input, params->skip, - params->beta, params->gamma, params->bias, maybe2half(params->epsilon), params->output, params->skip_input_bias_add_output, + params->beta, params->gamma, params->bias, static_cast(params->epsilon), params->output, params->skip_input_bias_add_output, (params->bias == nullptr) ? false : true, (params->skip_input_bias_add_output == nullptr) ? false : true); return HIP_CALL(hipGetLastError()); } -template -Status SkipLayerNormRegularOp(const SkipLayerNormParams* params) { +template +Status SkipLayerNormRegularOp(const SkipLayerNormParams* params) { TUNABLE_OP_RETURN_UNSUPPORTED_ARGUMENT_IF( !((params->ld > 0 && params->ld % VecSize == 0 && (params->ld >= ThreadsPerBlock * VecSize || (params->ld < GPU_WARP_SIZE && params->ld > (ThreadsPerBlock - GPU_WARP_SIZE) * VecSize))))); - SkipLayerNormKernelVec<<element_count, params->ld)), - dim3(ThreadsPerBlock), - 0, params->stream>>>( + SkipLayerNormKernelVec<<element_count, params->ld)), + dim3(ThreadsPerBlock), + 0, params->stream>>>( params->ld, params->input, params->skip, - params->beta, params->gamma, params->bias, maybe2half(params->epsilon), params->output, params->skip_input_bias_add_output, + params->beta, params->gamma, params->bias, static_cast(params->epsilon), params->output, params->skip_input_bias_add_output, (params->bias == nullptr) ? false : true, (params->skip_input_bias_add_output == nullptr) ? false : true); return HIP_CALL(hipGetLastError()); } -template -Status SkipLayerNormStaticSelection(const SkipLayerNormParams* params) { +template +Status SkipLayerNormStaticSelection(const SkipLayerNormParams* params) { bool hasBias = (params->bias == nullptr) ? false : true; bool hasSkipInputBiasAdditionOutput = (params->skip_input_bias_add_output == nullptr) ? false : true; const int grid_size = params->element_count / params->ld; const int block_size = 256; -#define LAUNCH_SKIPLAYERNORM_SMALL_FORWARD(ELEMENTS, TPB, ILP) \ - if (params->ld <= ELEMENTS) { \ - SkipLayerNormKernelSmall<<stream>>>( \ - params->ld, params->input, params->skip, params->beta, params->gamma, params->bias, \ - maybe2half(params->epsilon), params->output, params->skip_input_bias_add_output, \ - hasBias, hasSkipInputBiasAdditionOutput); \ - break; \ +#define LAUNCH_SKIPLAYERNORM_SMALL_FORWARD(ELEMENTS, TPB, ILP) \ + if (params->ld <= ELEMENTS) { \ + SkipLayerNormKernelSmall<<stream>>>( \ + params->ld, params->input, params->skip, params->beta, params->gamma, params->bias, \ + static_cast(params->epsilon), params->output, params->skip_input_bias_add_output, \ + hasBias, hasSkipInputBiasAdditionOutput); \ + break; \ } if (0 == (params->ld % 4)) { do { @@ -98,9 +97,9 @@ Status SkipLayerNormStaticSelection(const SkipLayerNormParams* params) { LAUNCH_SKIPLAYERNORM_SMALL_FORWARD(768, 192, 4) LAUNCH_SKIPLAYERNORM_SMALL_FORWARD(1024, 256, 4) - SkipLayerNormKernel<<stream>>>( + SkipLayerNormKernel<<stream>>>( params->ld, params->input, params->skip, params->beta, params->gamma, params->bias, - maybe2half(params->epsilon), params->output, params->skip_input_bias_add_output); + static_cast(params->epsilon), params->output, params->skip_input_bias_add_output); } while (0); } else { do { @@ -109,20 +108,20 @@ Status SkipLayerNormStaticSelection(const SkipLayerNormParams* params) { LAUNCH_SKIPLAYERNORM_SMALL_FORWARD(128, 128, 1) LAUNCH_SKIPLAYERNORM_SMALL_FORWARD(384, 384, 1) - SkipLayerNormKernel<<stream>>>( + SkipLayerNormKernel<<stream>>>( params->ld, params->input, params->skip, params->beta, params->gamma, params->bias, - maybe2half(params->epsilon), params->output, params->skip_input_bias_add_output); + static_cast(params->epsilon), params->output, params->skip_input_bias_add_output); } while (0); } return HIP_CALL(hipPeekAtLastError()); } // namespace rocm #define ADD_OP_FOR_ALL_VEC_SIZE(name, threads_per_block) \ - this->RegisterOp(name); \ - this->RegisterOp(name); \ - this->RegisterOp(name); \ - this->RegisterOp(name); \ - this->RegisterOp(name); + this->RegisterOp(name); \ + this->RegisterOp(name); \ + this->RegisterOp(name); \ + this->RegisterOp(name); \ + this->RegisterOp(name); #define ADD_OP_FOR_ALL_THREADS_PER_BLOCK_ALL_VEC_SIZE(name) \ ADD_OP_FOR_ALL_VEC_SIZE(name, 64) \ @@ -141,11 +140,11 @@ Status SkipLayerNormStaticSelection(const SkipLayerNormParams* params) { ADD_OP_FOR_ALL_VEC_SIZE(name, 896) \ ADD_OP_FOR_ALL_VEC_SIZE(name, 1024) -template -class SkipLayerNormTunableOp : public TunableOp> { +template +class SkipLayerNormTunableOp : public TunableOp> { public: SkipLayerNormTunableOp() { - this->RegisterOp(SkipLayerNormStaticSelection); + this->RegisterOp(SkipLayerNormStaticSelection); ADD_OP_FOR_ALL_THREADS_PER_BLOCK_ALL_VEC_SIZE(SkipLayerNormSmallOp) ADD_OP_FOR_ALL_THREADS_PER_BLOCK_ALL_VEC_SIZE(SkipLayerNormRegularOp) diff --git a/onnxruntime/python/tools/kernel_explorer/kernels/rocm/skip_layer_norm.cu b/onnxruntime/python/tools/kernel_explorer/kernels/rocm/skip_layer_norm.cu index ac5ec602f8..37a9f14769 100644 --- a/onnxruntime/python/tools/kernel_explorer/kernels/rocm/skip_layer_norm.cu +++ b/onnxruntime/python/tools/kernel_explorer/kernels/rocm/skip_layer_norm.cu @@ -23,16 +23,16 @@ class SkipLayerNormSmall : public IKernelExplorer { static_cast(beta.ptr()), static_cast(bias.ptr()), epsilon, hidden_size, element_count) {} void Run() override { - ORT_THROW_IF_ERROR((contrib::rocm::SkipLayerNormSmallOp(¶ms_))); + ORT_THROW_IF_ERROR((contrib::rocm::SkipLayerNormSmallOp(¶ms_))); } bool IsSupported() { - Status status = contrib::rocm::SkipLayerNormSmallOp(¶ms_); + Status status = contrib::rocm::SkipLayerNormSmallOp(¶ms_); return status.IsOK(); } private: - using ParamsT = contrib::rocm::SkipLayerNormParams; + using ParamsT = contrib::rocm::SkipLayerNormParams; ParamsT params_{}; }; @@ -47,16 +47,16 @@ class SkipLayerNormRegular : public IKernelExplorer { static_cast(beta.ptr()), static_cast(bias.ptr()), epsilon, hidden_size, element_count) {} void Run() override { - ORT_THROW_IF_ERROR((contrib::rocm::SkipLayerNormRegularOp(¶ms_))); + ORT_THROW_IF_ERROR((contrib::rocm::SkipLayerNormRegularOp(¶ms_))); } bool IsSupported() { - Status status = contrib::rocm::SkipLayerNormRegularOp(¶ms_); + Status status = contrib::rocm::SkipLayerNormRegularOp(¶ms_); return status.IsOK(); } private: - using ParamsT = contrib::rocm::SkipLayerNormParams; + using ParamsT = contrib::rocm::SkipLayerNormParams; ParamsT params_{}; }; @@ -71,16 +71,16 @@ class SkipLayerNormStaticSelection : public IKernelExplorer { static_cast(beta.ptr()), static_cast(bias.ptr()), epsilon, hidden_size, element_count) {} void Run() override { - ORT_THROW_IF_ERROR((contrib::rocm::SkipLayerNormStaticSelection(¶ms_))); + ORT_THROW_IF_ERROR((contrib::rocm::SkipLayerNormStaticSelection(¶ms_))); } bool IsSupported() { - Status status = contrib::rocm::SkipLayerNormStaticSelection(¶ms_); + Status status = contrib::rocm::SkipLayerNormStaticSelection(¶ms_); return status.IsOK(); } private: - using ParamsT = contrib::rocm::SkipLayerNormParams; + using ParamsT = contrib::rocm::SkipLayerNormParams; ParamsT params_{}; }; @@ -105,9 +105,9 @@ class SkipLayerNormTunable : public IKernelExplorer { } private: - using ParamsT = contrib::rocm::SkipLayerNormParams; + using ParamsT = contrib::rocm::SkipLayerNormParams; ParamsT params_{}; - contrib::rocm::SkipLayerNormTunableOp op_{}; + contrib::rocm::SkipLayerNormTunableOp op_{}; }; #define REGISTER_OP(name, type, threads_per_block, vec_size) \ diff --git a/onnxruntime/python/tools/kernel_explorer/kernels/skip_layer_norm_test.py b/onnxruntime/python/tools/kernel_explorer/kernels/skip_layer_norm_test.py index 2a653f92a4..006e563ed8 100644 --- a/onnxruntime/python/tools/kernel_explorer/kernels/skip_layer_norm_test.py +++ b/onnxruntime/python/tools/kernel_explorer/kernels/skip_layer_norm_test.py @@ -92,7 +92,7 @@ def run_skip_layer_norm(batch_size: int, seq_len: int, hidden_size: int, dtype: y_ref, y_optional = skip_layer_norm(input_x, skip, bias, gamma, beta, epsilon) np.testing.assert_almost_equal(y_ref, output_y, decimal=1) if has_optional_output: - np.testing.assert_almost_equal(y_optional, output_optional, decimal=1) + np.testing.assert_almost_equal(y_optional, output_optional, decimal=3) dtypes = ["float32", "float16"]