saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-25 19:48:11 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

[ROCm] Fix static TunableOp (#13668)

Browse source 
### Description
<!-- Describe your changes. -->

1. Re-add staticSelectionOp for FastGelu.
2. Call TunableOp when enable tuning. Call StaticSelectionOp when
disable tuning.


### Motivation and Context
<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->

Co-authored-by: peixuanzuo <peixuanzuo@linmif39a000004.zvflicr54joexhdgnhvmxrxygg.phxx.internal.cloudapp.net>
This commit is contained in:
PeixuanZuo 2022-11-22 10:51:54 +08:00 committed by GitHub
parent f1b5e4f1c9
commit 45a895cdc3
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
5 changed files with 82 additions and 28 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -41,23 +41,26 @@ namespace rocm {
template <typename T>
Status LaunchFastGeluKernel(hipStream_t stream, int input_length, int bias_length,
const T* input, const T* bias, T* output, bool tuning) {
static FastGeluTunableOp<T> op;
const T* input, const T* bias, T* output, bool tuning) {
FastGeluParams<T> params(stream, input, bias, output, input_length, bias_length);
if (tuning) {
static FastGeluTunableOp<T> op;
op.EnableTuning();
return op(&params);
}
FastGeluParams<T> op_params(stream, input, bias, output, input_length, bias_length);
return op(&op_params);
return FastGeluStaticSelection<T>(&params);
}
template Status LaunchFastGeluKernel<float>(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<BFloat16>(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<half>(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

2
onnxruntime/contrib_ops/rocm/bert/fast_gelu_impl.h
View file

@ -14,7 +14,7 @@ namespace rocm {
template <typename T>
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

12
onnxruntime/contrib_ops/rocm/bert/fast_gelu_impl_kernel.h
View file

@ -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<VecT*>(&output_v);
T bias_v[ILP];
if (bias != nullptr) {
VecT* bias_val = reinterpret_cast<VecT*>(&bias_v);
*bias_val = *reinterpret_cast<const VecT*>(&bias[idx % bias_length]);
VecT* bias_val = reinterpret_cast<VecT*>(&bias_v);
*bias_val = *reinterpret_cast<const VecT*>(&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<VecT*>(&output[idx])) = *output_val;

71
onnxruntime/contrib_ops/rocm/bert/fast_gelu_tunable_op.h
View file

@ -39,14 +39,64 @@ Status FastGeluOp(const FastGeluParams<T>* 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<T, ThreadsPerBlock, VecSize>),
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<T, ThreadsPerBlock, VecSize>
<<<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);
return HIP_CALL(hipGetLastError());
}
template <typename T>
Status FastGeluStaticSelection(const FastGeluParams<T>* params) {
constexpr int block_size = 256;
const int grid_size = (params->input_length + block_size - 1) / block_size;
FastGeluKernel<T, block_size><<<dim3(grid_size), dim3(block_size), 0, params->stream>>>(
params->input_length, params->bias_length, params->input, params->bias, params->output);
return HIP_CALL(hipGetLastError());
}
template <>
Status FastGeluStaticSelection(const FastGeluParams<half>* 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<half, block_size, 8><<<dim3(grid_size), dim3(block_size), 0, params->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<half, block_size, 4><<<dim3(grid_size), dim3(block_size), 0, params->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<half, block_size, 2><<<dim3(grid_size), dim3(block_size), 0, params->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<half, block_size><<<dim3(grid_size), dim3(block_size), 0, params->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<half, block_size, 8><<<dim3(grid_size), dim3(block_size), 0, params->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<half, block_size, 4><<<dim3(grid_size), dim3(block_size), 0, params->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<half, block_size, 2><<<dim3(grid_size), dim3(block_size), 0, params->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<half, block_size><<<dim3(grid_size), dim3(block_size), 0, params->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 <typename T>
class FastGeluTunableOp : public onnxruntime::rocm::tunable::TunableOp<FastGeluParams<T>> {
public:
FastGeluTunableOp() {
this->ops_.emplace_back(FastGeluStaticSelection<T>);
ADD_OP(64);
ADD_OP(128);
ADD_OP(192);
@ -69,8 +120,8 @@ class FastGeluTunableOp : public onnxruntime::rocm::tunable::TunableOp<FastGeluP
ADD_OP(448);
ADD_OP(512);
// NOTE: the 15-th kernel seems to be better in gerenal case, so set it as default one
this->SetDefaultId(15);
// NOTE: the 1st kernel is FastGelu Original implementation.
this->SetDefaultId(0);
}
};

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

@ -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<T> op_params(stream, output, input, skip, gamma, beta, bias, epsilon, ld, element_count);
static SkipLayerNormTunableOp<T> op;
SkipLayerNormParams<T> params(stream, output, input, skip, gamma, beta, bias, epsilon, ld, element_count);
// If disable tuning, the default implementation is SkipLayerNormStaticSelection.
if (tuning) {
static SkipLayerNormTunableOp<T> op;
op.EnableTuning();
return op(&params);
}
return op(&op_params);
return SkipLayerNormStaticSelection<T>(&params);
}
template Status LaunchSkipLayerNormKernel<float>(hipStream_t stream, float* output, const float* input,