saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-12 17:57:38 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

[ROCm] Add compute type for Skiplayernorm to fix ROCm CI (#15192)

Browse source 
- 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`.
This commit is contained in:
PeixuanZuo 2023-03-24 19:31:14 +08:00 committed by GitHub
parent 3a4c895765
commit 7eb6dbe7d8
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
8 changed files with 157 additions and 176 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

89
onnxruntime/contrib_ops/rocm/bert/layer_norm.cuh
View file

@ -80,16 +80,16 @@ struct KeyValuePairSum {
}
};
template <typename T, int TPB>
template <typename U, typename V, int TPB>
__device__ inline void LayerNorm(
const hipcub::KeyValuePair<T, T>& thread_data, const int ld, const int offset, const T* beta,
const T* gamma, const T epsilon, T* output) {
const hipcub::KeyValuePair<U, U>& 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<hipcub::KeyValuePair<T, T>, TPB>;
using BlockReduce = hipcub::BlockReduce<hipcub::KeyValuePair<U, U>, 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<U>(output[idx]);
const U g = static_cast<U>(gamma[i]);
const U b = (nullptr == beta) ? U(0.f) : static_cast<U>(beta[i]);
output[idx] = static_cast<V>(g * (val - mu) * rsigma + b);
}
}
template <typename T, int TPB, int ILP>
template <typename U, typename V, int TPB, int ILP>
__device__ inline void LayerNormVec(
const hipcub::KeyValuePair<T, T>& thread_data, const int ld, const int offset, const T* beta,
const T* gamma, const T epsilon, T* output) {
const hipcub::KeyValuePair<U, U>& 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<T, ILP>;
using BlockReduce = hipcub::BlockReduce<hipcub::KeyValuePair<T, T>, TPB>;
using VecV = aligned_vector<V, ILP>;
using BlockReduce = hipcub::BlockReduce<hipcub::KeyValuePair<U, U>, 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<VecT*>(&gamma_v);
VecT* output_val = reinterpret_cast<VecT*>(&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<VecT*>(&beta_v);
*beta_val = *reinterpret_cast<const VecT*>(&beta[i]);
}
*gamma_val = *reinterpret_cast<const VecT*>(&gamma[i]);
*output_val = *reinterpret_cast<const VecT*>(&output[idx]);
const VecV beta_v = (beta != nullptr) ? *reinterpret_cast<const VecV*>(beta + i) : VecV();
const VecV gamma_v = *reinterpret_cast<const VecV*>(gamma + i);
VecV output_v = *reinterpret_cast<const VecV*>(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<VecT*>(&output[idx])) = *reinterpret_cast<VecT*>(&output_v[0]);
*(reinterpret_cast<VecV*>(output + idx)) = output_v;
}
}
}
template <typename T, int TPB, int ILP>
__device__ inline void LayerNormSmall(const T* input_v, const hipcub::KeyValuePair<T, T>& thread_data,
const int ld, const int idx, const T* beta, const T* gamma,
const T epsilon, T* output) {
template <typename T, typename U, typename V, int TPB, int ILP>
__device__ inline void LayerNormSmall(const T* input_v, const hipcub::KeyValuePair<U, U>& 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<T, ILP>;
using BlockReduce = hipcub::BlockReduce<hipcub::KeyValuePair<T, T>, TPB>;
using VecV = aligned_vector<V, ILP>;
using BlockReduce = hipcub::BlockReduce<hipcub::KeyValuePair<U, U>, 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<T, T> sum_kv = BlockReduce(temp_storage).Reduce(thread_data, pair_sum);
const hipcub::KeyValuePair<U, U> 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<VecT*>(&beta_v);
*beta_val = *reinterpret_cast<const VecT*>(&beta[threadIdx.x * ILP]);
}
VecT* gamma_val = reinterpret_cast<VecT*>(&gamma_v);
*gamma_val = *reinterpret_cast<const VecT*>(&gamma[threadIdx.x * ILP]);
const VecV beta_v = (beta != nullptr) ? *reinterpret_cast<const VecV*>(beta + threadIdx.x * ILP) : VecV();
const VecV gamma_v = *reinterpret_cast<const VecV*>(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<VecT*>(&output[idx])) = *reinterpret_cast<VecT*>(&output_v[0]);
*(reinterpret_cast<VecV*>(output + idx)) = output_v;
}
}

2
onnxruntime/contrib_ops/rocm/bert/skip_layer_norm.cc
View file

@ -101,7 +101,7 @@ Status SkipLayerNorm<T>::ComputeInternal(OpKernelContext* ctx) const {
int64_t element_count = input_dims[0] * sequence_length * hidden_size;
typedef typename ToHipType<T>::MappedType HipT;
return LaunchSkipLayerNormKernel<HipT>(
return LaunchSkipLayerNormKernel<HipT, float, HipT>(
GetTuningContext(),
Stream(ctx),
reinterpret_cast<HipT*>(output->MutableData<T>()),

17
onnxruntime/contrib_ops/rocm/bert/skip_layer_norm_impl.cu
View file

@ -39,30 +39,31 @@ namespace onnxruntime {
namespace contrib {
namespace rocm {
template <typename T>
template <typename T, typename U, typename V>
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<T> params(tuning_ctx, stream, output, skip_input_bias_add_output, input, skip, gamma, beta, bias, epsilon, ld, element_count);
SkipLayerNormParams<T, V> 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<T> op;
static SkipLayerNormTunableOp<T, U, V> op;
return op(&params);
}
return SkipLayerNormStaticSelection<T>(&params);
return SkipLayerNormStaticSelection<T, U, V>(&params);
}
template Status LaunchSkipLayerNormKernel<float>(
template Status LaunchSkipLayerNormKernel<float, float, float>(
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<half>(
template Status LaunchSkipLayerNormKernel<half, float, half>(
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,

8
onnxruntime/contrib_ops/rocm/bert/skip_layer_norm_impl.h
View file

@ -10,16 +10,16 @@ namespace onnxruntime {
namespace contrib {
namespace rocm {
template <typename T>
template <typename T, typename U, typename V>
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)

112
onnxruntime/contrib_ops/rocm/bert/skip_layer_norm_impl_kernel.h
View file

@ -23,133 +23,125 @@ half maybe2half(float x) {
return __float2half_rn(x);
}
template <typename T, unsigned TPB>
template <typename T, typename U, typename V, unsigned TPB>
__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<T, T> thread_data(0, 0);
hipcub::KeyValuePair<U, U> 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<T, T>(rldval, rldval * val));
const U val = (bias == nullptr) ? static_cast<U>(input[idx]) + static_cast<U>(skip[idx]) :
static_cast<U>(input[idx]) + static_cast<U>(skip[idx]) + static_cast<U>(bias[i]);
const U rldval = reverse_ld * val;
thread_data = pair_sum(thread_data, hipcub::KeyValuePair<U, U>(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<T>(val);
}
output[idx] = val;
output[idx] = static_cast<V>(val);
}
LayerNorm<T, TPB>(thread_data, ld, offset, beta, gamma, epsilon, output);
LayerNorm<U, V, TPB>(thread_data, ld, offset, beta, gamma, epsilon, output);
}
// Vectorized kernel
template <typename T, unsigned TPB, int ILP>
template <typename T, typename U, typename V, unsigned TPB, int ILP>
__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<T, T> thread_data(0, 0);
hipcub::KeyValuePair<U, U> thread_data(U(0.f), U(0.f));
using VecT = aligned_vector<T, ILP>;
T input_v[ILP], skip_v[ILP], bias_v[ILP], skip_input_bias_add_output_v[ILP];
using VecV = aligned_vector<V, ILP>;
if (threadIdx.x * ILP < ld) {
VecT* input_val = reinterpret_cast<VecT*>(&input_v);
VecT* skip_val = reinterpret_cast<VecT*>(&skip_v);
for (int i = threadIdx.x * ILP; i < ld; i += TPB * ILP) {
int idx = offset + i;
*input_val = *reinterpret_cast<const VecT*>(&input[idx]);
*skip_val = *reinterpret_cast<const VecT*>(&skip[idx]);
if (hasBias) {
VecT* bias_val = reinterpret_cast<VecT*>(&bias_v);
*bias_val = *reinterpret_cast<const VecT*>(&bias[i]);
}
const VecT input_v = *reinterpret_cast<const VecT*>(input + idx);
const VecT skip_v = *reinterpret_cast<const VecT*>(skip + idx);
const VecT bias_v = hasBias ? *reinterpret_cast<const VecT*>(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<U>(input_v.val[k]) + static_cast<U>(skip_v.val[k]) + static_cast<U>(bias_v.val[k]) :
static_cast<U>(input_v.val[k]) + static_cast<U>(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<T>(val);
}
const T rldval = reverse_ld * input_v[k];
thread_data = pair_sum(thread_data, hipcub::KeyValuePair<T, T>(rldval, rldval * input_v[k]));
thread_data = pair_sum(thread_data, hipcub::KeyValuePair<U, U>(rldval, rldval * val));
output_v.val[k] = static_cast<V>(val);
}
if (hasSkipInputBiasAdditionOutput) {
*(reinterpret_cast<VecT*>(&skip_input_bias_add_output[idx])) = *reinterpret_cast<VecT*>(&skip_input_bias_add_output_v);
*(reinterpret_cast<VecT*>(skip_input_bias_add_output + idx)) = skip_input_bias_add_output_v;
}
*(reinterpret_cast<VecT*>(&output[idx])) = *reinterpret_cast<VecT*>(&input_v[0]);
*(reinterpret_cast<VecV*>(output + idx)) = output_v;
}
}
LayerNormVec<T, TPB, ILP>(thread_data, ld, offset, beta, gamma, epsilon, output);
LayerNormVec<U, V, TPB, ILP>(thread_data, ld, offset, beta, gamma, epsilon, output);
}
// Vectorized kernel
template <typename T, unsigned TPB, int ILP>
template <typename T, typename U, typename V, unsigned TPB, int ILP>
__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, ILP>;
T input_v[ILP], skip_v[ILP], bias_v[ILP], skip_input_bias_add_output_v[ILP];
hipcub::KeyValuePair<T, T> thread_data(T(0.f), T(0.f));
hipcub::KeyValuePair<U, U> thread_data(U(0.f), U(0.f));
VecT input_v;
if (ILP * threadIdx.x < ld) {
VecT* input_val = reinterpret_cast<VecT*>(&input_v);
*input_val = *reinterpret_cast<const VecT*>(&input[idx]);
input_v = *reinterpret_cast<const VecT*>(input + idx);
const VecT skip_v = *reinterpret_cast<const VecT*>(skip + idx);
const VecT bias_v = hasBias ? *reinterpret_cast<const VecT*>(bias + threadIdx.x * ILP) : VecT();
VecT skip_input_bias_add_output_v;
VecT* skip_val = reinterpret_cast<VecT*>(&skip_v);
*skip_val = *reinterpret_cast<const VecT*>(&skip[idx]);
if (hasBias) {
VecT* bias_val = reinterpret_cast<VecT*>(&bias_v);
*bias_val = *reinterpret_cast<const VecT*>(&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<U>(input_v.val[i]) + static_cast<U>(skip_v.val[i]) + static_cast<U>(bias_v.val[i]) :
static_cast<U>(input_v.val[i]) + static_cast<U>(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<T>(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<T>(val);
}
if (hasSkipInputBiasAdditionOutput) {
*(reinterpret_cast<VecT*>(&skip_input_bias_add_output[idx])) = *reinterpret_cast<VecT*>(&skip_input_bias_add_output_v);
*(reinterpret_cast<VecT*>(skip_input_bias_add_output + idx)) = skip_input_bias_add_output_v;
}
thread_data = hipcub::KeyValuePair<T, T>(rldval_sum, rldvalsq_sum);
thread_data = hipcub::KeyValuePair<U, U>(rldval_sum, rldvalsq_sum);
}
LayerNormSmall<T, TPB, ILP>(input_v, thread_data, ld, idx, beta, gamma, epsilon, output);
LayerNormSmall<T, U, V, TPB, ILP>(input_v.val, thread_data, ld, idx, beta, gamma, epsilon, output);
}
} // namespace rocm

81
onnxruntime/contrib_ops/rocm/bert/skip_layer_norm_tunable_op.h
View file

@ -18,76 +18,75 @@ namespace onnxruntime {
namespace contrib {
namespace rocm {
template <typename T>
template <typename T, typename V>
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 <typename T, int ThreadsPerBlock, int VecSize>
Status SkipLayerNormSmallOp(const SkipLayerNormParams<T>* params) {
template <typename T, typename U, typename V, int ThreadsPerBlock, int VecSize>
Status SkipLayerNormSmallOp(const SkipLayerNormParams<T, V>* 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<T, ThreadsPerBlock, VecSize><<<dim3(CeilDiv(params->element_count, params->ld)),
dim3(ThreadsPerBlock),
0, params->stream>>>(
SkipLayerNormKernelSmall<T, U, V, ThreadsPerBlock, VecSize><<<dim3(CeilDiv(params->element_count, params->ld)),
dim3(ThreadsPerBlock),
0, params->stream>>>(
params->ld, params->input, params->skip,
params->beta, params->gamma, params->bias, maybe2half<T>(params->epsilon), params->output, params->skip_input_bias_add_output,
params->beta, params->gamma, params->bias, static_cast<U>(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 <typename T, int ThreadsPerBlock, int VecSize>
Status SkipLayerNormRegularOp(const SkipLayerNormParams<T>* params) {
template <typename T, typename U, typename V, int ThreadsPerBlock, int VecSize>
Status SkipLayerNormRegularOp(const SkipLayerNormParams<T, V>* 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<T, ThreadsPerBlock, VecSize><<<dim3(CeilDiv(params->element_count, params->ld)),
dim3(ThreadsPerBlock),
0, params->stream>>>(
SkipLayerNormKernelVec<T, U, V, ThreadsPerBlock, VecSize><<<dim3(CeilDiv(params->element_count, params->ld)),
dim3(ThreadsPerBlock),
0, params->stream>>>(
params->ld, params->input, params->skip,
params->beta, params->gamma, params->bias, maybe2half<T>(params->epsilon), params->output, params->skip_input_bias_add_output,
params->beta, params->gamma, params->bias, static_cast<U>(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 <typename T>
Status SkipLayerNormStaticSelection(const SkipLayerNormParams<T>* params) {
template <typename T, typename U, typename V>
Status SkipLayerNormStaticSelection(const SkipLayerNormParams<T, V>* 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<T, TPB, ILP><<<grid_size, TPB, 0, params->stream>>>( \
params->ld, params->input, params->skip, params->beta, params->gamma, params->bias, \
maybe2half<T>(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<T, U, V, TPB, ILP><<<grid_size, TPB, 0, params->stream>>>( \
params->ld, params->input, params->skip, params->beta, params->gamma, params->bias, \
static_cast<U>(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<T>* params) {
LAUNCH_SKIPLAYERNORM_SMALL_FORWARD(768, 192, 4)
LAUNCH_SKIPLAYERNORM_SMALL_FORWARD(1024, 256, 4)
SkipLayerNormKernel<T, block_size><<<grid_size, block_size, 0, params->stream>>>(
SkipLayerNormKernel<T, U, V, block_size><<<grid_size, block_size, 0, params->stream>>>(
params->ld, params->input, params->skip, params->beta, params->gamma, params->bias,
maybe2half<T>(params->epsilon), params->output, params->skip_input_bias_add_output);
static_cast<U>(params->epsilon), params->output, params->skip_input_bias_add_output);
} while (0);
} else {
do {
@ -109,20 +108,20 @@ Status SkipLayerNormStaticSelection(const SkipLayerNormParams<T>* params) {
LAUNCH_SKIPLAYERNORM_SMALL_FORWARD(128, 128, 1)
LAUNCH_SKIPLAYERNORM_SMALL_FORWARD(384, 384, 1)
SkipLayerNormKernel<T, block_size><<<grid_size, block_size, 0, params->stream>>>(
SkipLayerNormKernel<T, U, V, block_size><<<grid_size, block_size, 0, params->stream>>>(
params->ld, params->input, params->skip, params->beta, params->gamma, params->bias,
maybe2half<T>(params->epsilon), params->output, params->skip_input_bias_add_output);
static_cast<U>(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<T, threads_per_block, 1>); \
this->RegisterOp(name<T, threads_per_block, 2>); \
this->RegisterOp(name<T, threads_per_block, 4>); \
this->RegisterOp(name<T, threads_per_block, 8>); \
this->RegisterOp(name<T, threads_per_block, 16>);
this->RegisterOp(name<T, U, V, threads_per_block, 1>); \
this->RegisterOp(name<T, U, V, threads_per_block, 2>); \
this->RegisterOp(name<T, U, V, threads_per_block, 4>); \
this->RegisterOp(name<T, U, V, threads_per_block, 8>); \
this->RegisterOp(name<T, U, V, threads_per_block, 16>);
#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<T>* params) {
ADD_OP_FOR_ALL_VEC_SIZE(name, 896) \
ADD_OP_FOR_ALL_VEC_SIZE(name, 1024)
template <typename T>
class SkipLayerNormTunableOp : public TunableOp<SkipLayerNormParams<T>> {
template <typename T, typename U, typename V>
class SkipLayerNormTunableOp : public TunableOp<SkipLayerNormParams<T, V>> {
public:
SkipLayerNormTunableOp() {
this->RegisterOp(SkipLayerNormStaticSelection<T>);
this->RegisterOp(SkipLayerNormStaticSelection<T, U, V>);
ADD_OP_FOR_ALL_THREADS_PER_BLOCK_ALL_VEC_SIZE(SkipLayerNormSmallOp)
ADD_OP_FOR_ALL_THREADS_PER_BLOCK_ALL_VEC_SIZE(SkipLayerNormRegularOp)

22
onnxruntime/python/tools/kernel_explorer/kernels/rocm/skip_layer_norm.cu
View file

@ -23,16 +23,16 @@ class SkipLayerNormSmall : public IKernelExplorer {
static_cast<T*>(beta.ptr()), static_cast<T*>(bias.ptr()), epsilon, hidden_size, element_count) {}
void Run() override {
ORT_THROW_IF_ERROR((contrib::rocm::SkipLayerNormSmallOp<T, ThreadsPerBlock, VecSize>(&params_)));
ORT_THROW_IF_ERROR((contrib::rocm::SkipLayerNormSmallOp<T, float, T, ThreadsPerBlock, VecSize>(&params_)));
}
bool IsSupported() {
Status status = contrib::rocm::SkipLayerNormSmallOp<T, ThreadsPerBlock, VecSize>(&params_);
Status status = contrib::rocm::SkipLayerNormSmallOp<T, float, T, ThreadsPerBlock, VecSize>(&params_);
return status.IsOK();
}
private:
using ParamsT = contrib::rocm::SkipLayerNormParams<T>;
using ParamsT = contrib::rocm::SkipLayerNormParams<T, T>;
ParamsT params_{};
};
@ -47,16 +47,16 @@ class SkipLayerNormRegular : public IKernelExplorer {
static_cast<T*>(beta.ptr()), static_cast<T*>(bias.ptr()), epsilon, hidden_size, element_count) {}
void Run() override {
ORT_THROW_IF_ERROR((contrib::rocm::SkipLayerNormRegularOp<T, ThreadsPerBlock, VecSize>(&params_)));
ORT_THROW_IF_ERROR((contrib::rocm::SkipLayerNormRegularOp<T, float, T, ThreadsPerBlock, VecSize>(&params_)));
}
bool IsSupported() {
Status status = contrib::rocm::SkipLayerNormRegularOp<T, ThreadsPerBlock, VecSize>(&params_);
Status status = contrib::rocm::SkipLayerNormRegularOp<T, float, T, ThreadsPerBlock, VecSize>(&params_);
return status.IsOK();
}
private:
using ParamsT = contrib::rocm::SkipLayerNormParams<T>;
using ParamsT = contrib::rocm::SkipLayerNormParams<T, T>;
ParamsT params_{};
};
@ -71,16 +71,16 @@ class SkipLayerNormStaticSelection : public IKernelExplorer {
static_cast<T*>(beta.ptr()), static_cast<T*>(bias.ptr()), epsilon, hidden_size, element_count) {}
void Run() override {
ORT_THROW_IF_ERROR((contrib::rocm::SkipLayerNormStaticSelection<T>(&params_)));
ORT_THROW_IF_ERROR((contrib::rocm::SkipLayerNormStaticSelection<T, float, T>(&params_)));
}
bool IsSupported() {
Status status = contrib::rocm::SkipLayerNormStaticSelection<T>(&params_);
Status status = contrib::rocm::SkipLayerNormStaticSelection<T, float, T>(&params_);
return status.IsOK();
}
private:
using ParamsT = contrib::rocm::SkipLayerNormParams<T>;
using ParamsT = contrib::rocm::SkipLayerNormParams<T, T>;
ParamsT params_{};
};
@ -105,9 +105,9 @@ class SkipLayerNormTunable : public IKernelExplorer {
}
private:
using ParamsT = contrib::rocm::SkipLayerNormParams<T>;
using ParamsT = contrib::rocm::SkipLayerNormParams<T, T>;
ParamsT params_{};
contrib::rocm::SkipLayerNormTunableOp<T> op_{};
contrib::rocm::SkipLayerNormTunableOp<T, float, T> op_{};
};
#define REGISTER_OP(name, type, threads_per_block, vec_size) \

2
onnxruntime/python/tools/kernel_explorer/kernels/skip_layer_norm_test.py
View file

@ -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"]