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

[ROCm] refact SkipLayernorm long if-elseif statements (#14795)

Browse source 
Refact SkipLayernorm long if-elseif statements.
This commit is contained in:
PeixuanZuo 2023-03-14 23:04:55 +08:00 committed by GitHub
parent 0fa00429d5
commit ff2850029b
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 28 additions and 59 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -78,75 +78,44 @@ template <typename T>
Status SkipLayerNormStaticSelection(const SkipLayerNormParams<T>* 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; \
}
if (0 == (params->ld % 4)) {
const int grid_size = params->element_count / params->ld;
if (params->ld <= 32) {
constexpr int block_size = 32;
SkipLayerNormKernelSmall<T, block_size, 1><<<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, hasBias, hasSkipInputBiasAdditionOutput);
} else if (params->ld <= 64) {
constexpr int block_size = 64 / 2;
SkipLayerNormKernelSmall<T, block_size, 2><<<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, hasBias, hasSkipInputBiasAdditionOutput);
} else if (params->ld <= 128) {
constexpr int block_size = 128 / 4;
SkipLayerNormKernelSmall<T, block_size, 4><<<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, hasBias, hasSkipInputBiasAdditionOutput);
} else if (params->ld <= 384) {
constexpr int block_size = 384 / 4;
SkipLayerNormKernelSmall<T, block_size, 4><<<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, hasBias, hasSkipInputBiasAdditionOutput);
} else if (params->ld <= 768) {
constexpr int block_size = 768 / 4;
SkipLayerNormKernelSmall<T, block_size, 4><<<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, hasBias, hasSkipInputBiasAdditionOutput);
} else if (params->ld <= 1024) {
constexpr int block_size = 1024 / 4;
SkipLayerNormKernelSmall<T, block_size, 4><<<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, hasBias, hasSkipInputBiasAdditionOutput);
} else {
constexpr int block_size = 256;
do {
LAUNCH_SKIPLAYERNORM_SMALL_FORWARD(32, 32, 1)
LAUNCH_SKIPLAYERNORM_SMALL_FORWARD(64, 32, 2)
LAUNCH_SKIPLAYERNORM_SMALL_FORWARD(128, 32, 4)
LAUNCH_SKIPLAYERNORM_SMALL_FORWARD(384, 96, 4)
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>>>(
params->ld, params->input, params->skip, params->beta, params->gamma, params->bias,
maybe2half<T>(params->epsilon), params->output, params->skip_input_bias_add_output);
}
} while (0);
} else {
const int grid_size = params->element_count / params->ld;
if (params->ld <= 32) {
constexpr int block_size = 32;
SkipLayerNormKernelSmall<T, block_size, 1><<<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, hasBias, hasSkipInputBiasAdditionOutput);
} else if (params->ld <= 64) {
constexpr int block_size = 64;
SkipLayerNormKernelSmall<T, block_size, 1><<<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, hasBias, hasSkipInputBiasAdditionOutput);
} else if (params->ld <= 128) {
constexpr int block_size = 128;
SkipLayerNormKernelSmall<T, block_size, 1><<<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, hasBias, hasSkipInputBiasAdditionOutput);
} else if (params->ld == 384) {
constexpr int block_size = 384;
SkipLayerNormKernelSmall<T, block_size, 1><<<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, hasBias, hasSkipInputBiasAdditionOutput);
} else {
constexpr int block_size = 256;
do {
LAUNCH_SKIPLAYERNORM_SMALL_FORWARD(32, 32, 1)
LAUNCH_SKIPLAYERNORM_SMALL_FORWARD(64, 64, 1)
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>>>(
params->ld, params->input, params->skip, params->beta, params->gamma, params->bias,
maybe2half<T>(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>); \