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

ROCm optimized layernorm for MI100 (#7682)

Browse source 
* layernorm optimizations

* Changed HIP flag from HIP_VERSION to __HIP_PLATFORM_HCC__
This commit is contained in:
Aswin John Mathews 2021-05-13 17:54:06 -05:00 committed by GitHub
parent d90a99aad5
commit 4afdc19958
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 43 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
onnxruntime/contrib_ops/cuda/layer_norm_impl.cu
View file

@ -364,7 +364,11 @@ void HostApplyLayerNorm(
const int warp_size = prop.warpSize;
ORT_ENFORCE(warp_size == GPU_WARP_SIZE);
const dim3 threads(warp_size, 4, 1);
dim3 threads(warp_size, 4, 1);
#ifdef __HIP_PLATFORM_HCC__
// Optimization for ROCm MI100
threads.y = 1;
#endif
const dim3 blocks(1, std::min<unsigned int>(n1, maxGridY), 1);
int nshared =
threads.y > 1 ? threads.y * sizeof(U) + (threads.y / 2) * sizeof(U) : 0;

39
orttraining/orttraining/training_ops/cuda/nn/layer_norm_impl.cu
View file

@ -306,6 +306,7 @@ __global__ void cuComputeGradInput(
const int numx = blockDim.x * blockDim.y;
const int thrx = threadIdx.x + threadIdx.y * blockDim.x;
if (use_gamma) {
#ifndef __HIP_PLATFORM_HCC__
int l = 4 * thrx;
for (; l + 3 < n2; l += 4 * numx) {
for (int k = 0; k < 4; ++k) {
@ -331,7 +332,24 @@ __global__ void cuComputeGradInput(
sum_loss2 += c_loss * (c_output - U(beta[l]));
}
}
#else
// Optimization for ROCm MI100
for( int l = 0; l < n2 ; l += numx) {
int idx = l + thrx;
T gamma_idx = (idx<n2)?gamma[ idx ]:T(0);
const U c_loss = static_cast<U>( (idx<n2)?k_dout[ idx ]:T(0) );
sum_loss1 += c_loss * U( gamma_idx );
if (use_mean) {
const U c_h = static_cast<U>( (idx<n2)?k_input[ idx ]:T(0) );
sum_loss2 += c_loss * U(gamma_idx) * (c_h - c_mean) * c_invvar;
} else {
const U c_output = static_cast<U>( (idx<n2)?k_output[idx]:T(0) );
sum_loss2 += c_loss * (c_output - U( (idx<n2)?beta[idx]:T(0) ));
}
}
#endif
} else {
#ifndef __HIP_PLATFORM_HCC__
int l = 4 * thrx;
for (; l + 3 < n2; l += 4 * numx) {
for (int k = 0; k < 4; ++k) {
@ -357,6 +375,21 @@ __global__ void cuComputeGradInput(
sum_loss2 += c_loss * c_output;
}
}
#else
// Optimization for ROCm MI100
for( int l = 0; l < n2 ; l += numx) {
int idx = l + thrx;
const U c_loss = static_cast<U>((idx<n2)?k_dout[idx]:T(0));
sum_loss1 += c_loss;
if (use_mean) {
const U c_h = static_cast<U>((idx<n2)?k_input[idx]:T(0));
sum_loss2 += c_loss * (c_h - c_mean) * c_invvar;
} else {
const U c_output = static_cast<U>((idx<n2)?k_output[idx]:T(0));
sum_loss2 += c_loss * c_output;
}
}
#endif
}
// intra-warp reductions
for (int mask = blockDim.x / 2; mask > 0; mask /= 2) {
@ -504,7 +537,11 @@ void HostLayerNormGradient(
// compute grad_input
const uint64_t maxGridY = prop.maxGridSize[1];
const dim3 blocks1(1, std::min<unsigned int>(static_cast<unsigned int>(n1), static_cast<unsigned int>(maxGridY)), 1);
const dim3 threads1(warp_size, 4, 1);
dim3 threads1(warp_size, 4, 1);
#ifdef __HIP_PLATFORM_HCC__
// Optimization for ROCm MI100
threads1.y = 2;
#endif
int nshared =
threads1.y > 1 ? threads1.y * threads1.x * sizeof(U) : 0;
if (mean == nullptr && !simplified) {