[ROCm] Add Tunable GroupNorm (#15298)

refactor GroupNorm and Add Tunable GroupNorm
This commit is contained in:
PeixuanZuo 2023-04-12 10:55:42 +08:00 committed by GitHub
parent db4fc12318
commit ce1eb6d629
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
6 changed files with 886 additions and 383 deletions

View file

@ -0,0 +1,118 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
#include <string>
#include "core/providers/rocm/cu_inc/common.cuh"
#include "core/providers/rocm/rocm_common.h"
#include "core/providers/rocm/tunable/rocm_tunable.h"
namespace onnxruntime {
namespace contrib {
namespace rocm {
using onnxruntime::rocm::CeilDiv;
int32_t findMaxDivisor(int32_t n, int32_t maxAllowedDivisor) {
int32_t maxDivisor = -1;
for (int32_t i = 1; i <= std::sqrt(n); i++) {
if (n % i == 0) {
int32_t divisor1 = n / i;
int32_t divisor2 = i;
if (divisor1 > maxDivisor && divisor1 < maxAllowedDivisor) {
maxDivisor = divisor1;
}
if (divisor2 > maxDivisor && divisor2 < maxAllowedDivisor) {
maxDivisor = divisor2;
}
}
}
return maxDivisor;
}
template <typename T>
struct GroupNormNHWCParams : OpParams {
GroupNormNHWCParams(RocmTuningContext* tuning_ctx, hipStream_t stream, T* dst, float* redBuffer, const T* src, const float* gamma,
const float* beta, int32_t n, int32_t h, int32_t w, int32_t c, int32_t groups, float epsilon, bool withSwish)
: OpParams(tuning_ctx, stream), dst(dst), src(src), gamma(gamma), beta(beta), redBuffer(redBuffer), epsilon(epsilon), n(n), h(h), w(w), c(c), groups(groups), withSwish(withSwish) {
int32_t maxBlocksPerHW = 1024;
switch (c) {
case 960:
case 1920:
cPerBlock = 480;
break;
case 512:
case 256:
cPerBlock = 256;
break;
case 128:
cPerBlock = 128;
break;
default:
cPerBlock = 320;
}
hw = h * w;
const int32_t blocksPerHW = findMaxDivisor(hw, maxBlocksPerHW);
hwPerBlock = CeilDiv(hw, blocksPerHW);
cPerGroup = c / groups;
hwc = hw * c;
invHWC = 1.F / (float)(hw * cPerGroup);
groupsPerBlock = cPerBlock / cPerGroup;
}
std::string Signature() const override {
std::string sig = std::to_string(n) + "_" + std::to_string(h * w) + "_" + std::to_string(c) + "_" + std::to_string(groups);
return sig;
}
// The output buffer. Layout NHWC.
T* dst;
// The input buffer. Layout NHWC.
T const* src;
// The gamma scaling factor.
float const* gamma;
// The beta term to add in GN.
float const* beta;
// The temporary buffer to do the global parallel reduction. Size:
// BLOCKS_PER_BATCH x C x 2.
float* redBuffer;
float epsilon;
// The number of instances in the batch.
int32_t n;
// The height and width of each activation map.
int32_t h;
int32_t w;
// The number of channels.
int32_t c;
// The number of groups.
int32_t groups;
// Do we apply the Swish activation function?
bool withSwish;
// Precomputed values and parameters to control the execution of the kernels.
// The number of activations per instance (h * w) and the number of
// activations per block.
int32_t hw;
int32_t hwPerBlock;
// The number of channels per group and blocks per activation in the C
// dimension.
int32_t cPerBlock;
int32_t cPerGroup;
// The precomputed stride between instances.
int32_t hwc;
// The inverse of hwc in floats (to compute mean/var).
float invHWC;
// The precomputed number of groups per block.
int32_t groupsPerBlock;
};
} // namespace rocm
} // namespace contrib
} // namespace onnxruntime

View file

@ -1,344 +1,17 @@
#include "hip/hip_runtime.h"
/*
* SPDX-FileCopyrightText: Copyright (c) 1993-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: Apache-2.0
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
// The ROCM kernel is hipified from CUDA kernel.
#include <hip/hip_fp16.h>
#include <hip/hip_runtime_api.h>
#include <hipcub/hipcub.hpp>
#include "core/providers/rocm/cu_inc/common.cuh"
#include "core/providers/rocm/rocm_common.h"
#include "contrib_ops/rocm/diffusion/group_norm_impl.h"
#include "contrib_ops/rocm/transformers/dump_rocm_tensor.h"
#include <hip/hip_fp16.h>
#include "contrib_ops/rocm/diffusion/group_norm_common.h"
#include "contrib_ops/rocm/diffusion/group_norm_tunable_op.h"
namespace onnxruntime {
namespace contrib {
namespace rocm {
static inline int32_t divUp(int32_t m, int32_t n) {
return (m + n - 1) / n;
}
static inline __device__ __host__ float sigmoid(float x) {
return 1.F / (1.F + expf(-x));
}
struct GroupSums {
// Is it the 1st element of the group?
int32_t flag;
// The sum.
float sum;
// The sum of squares.
float sumSq;
};
struct GroupSumsOp {
inline __device__ GroupSums operator()(GroupSums const& a, GroupSums const& b) {
GroupSums dst;
dst.sum = b.flag ? b.sum : (a.sum + b.sum);
dst.sumSq = b.flag ? b.sumSq : (a.sumSq + b.sumSq);
dst.flag = a.flag + b.flag;
return dst;
}
};
template <typename T>
struct GroupNormNHWCParams {
// The output buffer. Layout NHWC.
T* dst;
// The input buffer. Layout NHWC.
T const* src;
// The gamma scaling factor.
float const* gamma;
// The beta term to add in GN.
float const* beta;
// The temporary buffer to do the global parallel reduction. Size:
// BLOCKS_PER_BATCH x C x 2.
float* redBuffer;
// The number of instances in the batch.
int32_t n;
// The height and width of each activation map.
int32_t h;
int32_t w;
// The number of channels.
int32_t c;
// The number of groups.
int32_t groups;
// Do we apply the Swish activation function?
bool withSwish;
// Precomputed values and parameters to control the execution of the kernels.
// The number of activations per instance (h * w) and the number of
// activations per block.
int32_t hw;
int32_t hwPerBlock;
// The number of channels per group and blocks per activation in the C
// dimension.
int32_t cPerBlock;
int32_t cPerGroup;
// The precomputed stride between instances.
int32_t hwc;
// The inverse of hwc in floats (to compute mean/var).
float invHWC;
// The precomputed number of groups per block.
int32_t groupsPerBlock;
};
template <typename T, typename U, int32_t ILP>
inline __device__ void UpdateSum(const T* src, int64_t offset, U& sum, U& sumSq) {
using VecT = onnxruntime::rocm::aligned_vector<T, ILP>;
const VecT input_v = *reinterpret_cast<const VecT*>(src + offset);
#pragma unroll
for (int i = 0; i < ILP; i++) {
const U val = static_cast<U>(input_v.val[i]);
sum += val;
sumSq += val * val;
}
}
template <typename T, int32_t tTHREADS_PER_BLOCK, int32_t ILP>
__global__ void groupNormNHWCSumKernel(GroupNormNHWCParams<T> params) {
// The object in charge of doing the sums for the different blocks.
typedef hipcub::BlockScan<GroupSums, tTHREADS_PER_BLOCK> BlockScan;
// Allocate shared memory for BlockScan.
__shared__ typename BlockScan::TempStorage tempStorage;
// Allocate shared memory for the groups. We could reduce the amount of shared
// memory reserved.
__shared__ float2 smem[tTHREADS_PER_BLOCK];
// The instance in the batch.
int32_t ni = blockIdx.z;
// The channel loaded by that thread (ILP channels per thread).
int32_t ci = blockIdx.x * params.cPerBlock + threadIdx.x * ILP;
// The first activation loaded by that block.
int32_t hwBegin = blockIdx.y * params.hwPerBlock;
// The last activation loaded by that block.
int32_t hwEnd = min(hwBegin + params.hwPerBlock, params.hw);
// The sums.
float sum = 0.F;
float sumSq = 0.F;
// Iterate over the activations to compute the sums.
if (ci < params.c) {
for (int32_t hwi = hwBegin; hwi < hwEnd; ++hwi) {
// The offset.
int64_t offset = static_cast<int64_t>(ni) * params.hwc + static_cast<int64_t>(hwi) * params.c + ci;
UpdateSum<T, float, ILP>(params.src, offset, sum, sumSq);
}
}
// The group that thread works on and the channel in the group (modulus).
int32_t gi = threadIdx.x * ILP / params.cPerGroup;
int32_t cj = threadIdx.x * ILP - params.cPerGroup * gi;
// The data for the summations.
GroupSums inp{cj == 0 ? 1 : 0, sum, sumSq};
// Do the segmented scan.
GroupSums out;
BlockScan(tempStorage).InclusiveScan(inp, out, GroupSumsOp());
// Store the results for the groups in shared memory (to produce coalesced
// stores later).
if (cj == params.cPerGroup - ILP) { // ILP channels per thread
smem[gi] = make_float2(out.sum, out.sumSq);
}
// Make sure the data is in shared memory.
__syncthreads();
// The global group index.
int32_t gj = blockIdx.x * params.groupsPerBlock + threadIdx.x;
// Threads that have nothing left to do, exit.
if (threadIdx.x >= params.groupsPerBlock || gj >= params.groups) {
return;
}
// The first threads (those storing to global memory, load the values).
float2 sums = smem[threadIdx.x];
// Store to global memory.
atomicAdd(&params.redBuffer[(2 * ni + 0) * params.groups + gj], sums.x);
atomicAdd(&params.redBuffer[(2 * ni + 1) * params.groups + gj], sums.y);
}
template <typename T>
void groupNormNHWCSum(GroupNormNHWCParams<T> const& params, hipStream_t stream) {
// Make sure the values are as we expect.
ORT_ENFORCE(params.c % params.cPerBlock == 0 && params.hw % params.hwPerBlock == 0);
// Make sure a group does not span multiple blocks.
ORT_ENFORCE(params.cPerBlock % params.cPerGroup == 0);
dim3 grid;
// The number of blocks to compute all the channels.
grid.x = params.c / params.cPerBlock;
// The number of blocks to compute all the activations in a given instance.
grid.y = divUp(params.hw, params.hwPerBlock);
// The number of instances.
grid.z = params.n;
switch (params.cPerBlock) {
case 320:
groupNormNHWCSumKernel<T, 256, 2><<<grid, 256, 0, stream>>>(params);
break;
case 480:
groupNormNHWCSumKernel<T, 256, 2><<<grid, 256, 0, stream>>>(params);
break;
case 256:
groupNormNHWCSumKernel<T, 128, 2><<<grid, 128, 0, stream>>>(params);
break;
case 128:
groupNormNHWCSumKernel<T, 64, 2><<<grid, 64, 0, stream>>>(params);
break;
default:
ORT_NOT_IMPLEMENTED("Not implemented");
}
}
template <typename T, typename U, int32_t ILP>
__device__ void computeGroupNorm(const T* src, T* dst, int64_t offset, U mean, U invStdDev,
const U* gamma_v, const U* beta_v, bool swish) {
using VecT = onnxruntime::rocm::aligned_vector<T, ILP>;
const VecT input_v = *reinterpret_cast<const VecT*>(src + offset);
VecT output_v;
#pragma unroll
for (int i = 0; i < ILP; i++) {
U val = static_cast<U>(input_v.val[i]);
val = (val - mean) * invStdDev;
val = gamma_v[i] * val + beta_v[i];
if (swish) {
val = val * sigmoid(val);
}
output_v.val[i] = static_cast<T>(val);
}
*(reinterpret_cast<VecT*>(dst + offset)) = output_v;
}
template <typename T, int32_t tTHREADS_PER_BLOCK, int32_t ILP>
__global__ void groupNormNHWCScaleKernel(GroupNormNHWCParams<T> params) {
// The channel loaded by that thread (ILP channels per thread for F16x2).
int32_t ci = blockIdx.x * params.cPerBlock + threadIdx.x * ILP;
if (ci >= params.c) {
return;
}
// The instance in the batch.
int32_t ni = blockIdx.z;
// The group that thread works on and the channel in the group (modulus).
int32_t gi = ci / params.cPerGroup;
// Load the sum and sum of squares for the group.
float sum = 0.F, sumSq = 0.F;
if (gi < params.groups) {
sum = params.redBuffer[(2 * ni + 0) * params.groups + gi];
sumSq = params.redBuffer[(2 * ni + 1) * params.groups + gi];
}
using VecF = onnxruntime::rocm::aligned_vector<float, ILP>;
const VecF gamma_v = *reinterpret_cast<const VecF*>(params.gamma + ci);
const VecF beta_v = *reinterpret_cast<const VecF*>(params.beta + ci);
// Compute the mean.
float mean = sum * params.invHWC;
// Compute the variance.
float var = sumSq * params.invHWC - (mean * mean);
// Compute the inverse of the stddev.
float invStdDev = var <= 0.F ? 1.F : rsqrtf(var);
// The first activation loaded by that block.
int32_t hwBegin = blockIdx.y * params.hwPerBlock;
// The last activation loaded by that block.
int32_t hwEnd = min(hwBegin + params.hwPerBlock, params.hw);
// Iterate over the activations to compute the sums.
for (int32_t hwi = hwBegin; hwi < hwEnd; ++hwi) {
// The src/dst offset.
int64_t offset = (int64_t)ni * params.hwc + hwi * params.c + ci;
// Fetch ILP channels per thread.
computeGroupNorm<T, float, ILP>(params.src, params.dst, offset, mean, invStdDev, gamma_v.val, beta_v.val, params.withSwish);
}
}
template <typename T>
void groupNormNHWCScale(GroupNormNHWCParams<T> const& params, hipStream_t stream) {
// Make sure the dimensions are aligned with what we expect.
ORT_ENFORCE(params.c % params.cPerBlock == 0);
// Make sure a group does not span multiple blocks.
ORT_ENFORCE(params.cPerBlock % params.cPerGroup == 0);
dim3 grid;
// The number of blocks to compute all the channels.
grid.x = params.c / params.cPerBlock;
// The number of blocks to compute all the activations in a given instance.
grid.y = divUp(params.hw, params.hwPerBlock);
// The number of instances.
grid.z = params.n;
switch (params.cPerBlock) {
case 320:
groupNormNHWCScaleKernel<T, 256, 2><<<grid, 256, 0, stream>>>(params);
break;
case 480:
groupNormNHWCScaleKernel<T, 256, 2><<<grid, 256, 0, stream>>>(params);
break;
case 256:
groupNormNHWCScaleKernel<T, 128, 2><<<grid, 128, 0, stream>>>(params);
break;
case 128:
groupNormNHWCScaleKernel<T, 64, 2><<<grid, 64, 0, stream>>>(params);
break;
default:
ORT_NOT_IMPLEMENTED("Not implemented");
}
}
int32_t findMaxDivisor(int32_t n, int32_t maxAllowedDivisor) {
int32_t maxDivisor = -1;
for (int32_t i = 1; i <= std::sqrt(n); i++) {
if (n % i == 0) {
int32_t divisor1 = n / i;
int32_t divisor2 = i;
if (divisor1 > maxDivisor && divisor1 < maxAllowedDivisor) {
maxDivisor = divisor1;
}
if (divisor2 > maxDivisor && divisor2 < maxAllowedDivisor) {
maxDivisor = divisor2;
}
}
}
return maxDivisor;
}
template <typename T>
Status LaunchGroupNormKernel(
hipStream_t stream,
@ -364,57 +37,10 @@ Status LaunchGroupNormKernel(
"only num_groups=32 is supported. Got", num_groups);
}
GroupNormNHWCParams<T> params;
int32_t cPerBlock = 320;
int32_t maxBlocksPerHW = 1024;
switch (num_channels) {
case 960:
case 1920:
cPerBlock = 480;
break;
case 512:
case 256:
cPerBlock = 256;
break;
case 128:
cPerBlock = 128;
break;
default:
cPerBlock = 320;
}
GroupNormNHWCParams<T> params(nullptr, stream, output, reinterpret_cast<float*>(workspace), input, gamma, beta,
batch_size, height, width, num_channels, num_groups, epsilon, use_swish_activation);
params.withSwish = use_swish_activation;
params.dst = output;
params.src = input;
params.gamma = gamma;
params.beta = beta;
params.redBuffer = reinterpret_cast<float*>(workspace);
params.n = batch_size;
params.h = height;
params.w = width;
params.c = num_channels;
params.groups = num_groups;
params.hw = params.h * params.w;
const int32_t blocksPerHW = findMaxDivisor(params.hw, maxBlocksPerHW);
params.hwPerBlock = divUp(params.hw, blocksPerHW);
params.cPerBlock = cPerBlock;
params.cPerGroup = params.c / params.groups;
params.hwc = params.hw * params.c;
params.invHWC = 1.F / (float)(params.hw * params.cPerGroup);
params.groupsPerBlock = cPerBlock / params.cPerGroup;
DUMP_TENSOR_INIT();
DUMP_TENSOR("input", input, batch_size, num_channels, height * width);
DUMP_TENSOR("gamma", gamma, 1, num_channels);
DUMP_TENSOR("beta", beta, 1, num_channels);
HIP_RETURN_IF_ERROR(hipMemsetAsync(params.redBuffer, 0, GetGroupNormWorkspaceSizeInBytes(), stream));
groupNormNHWCSum<T>(params, stream);
DUMP_TENSOR("workspace", params.redBuffer, batch_size, num_groups, 2);
HIP_RETURN_IF_ERROR(hipGetLastError());
groupNormNHWCScale<T>(params, stream);
HIP_RETURN_IF_ERROR(hipGetLastError());
DUMP_TENSOR("output", output, batch_size, num_channels, height * width);
return Status::OK();
return GroupNormNHWCStaticSelection(&params);
}
template Status LaunchGroupNormKernel<half>(hipStream_t stream, half* output,

View file

@ -0,0 +1,213 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
// The ROCm kernel is modified from TensorRT 8.5.
/*
* SPDX-FileCopyrightText: Copyright (c) 1993-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: Apache-2.0
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#pragma once
#include <hip/hip_fp16.h>
#include <hip/hip_runtime_api.h>
#include <hipcub/hipcub.hpp>
#include "core/providers/rocm/cu_inc/common.cuh"
#include "core/providers/rocm/rocm_common.h"
namespace onnxruntime {
namespace contrib {
namespace rocm {
static inline __device__ __host__ float sigmoid(float x) {
return 1.F / (1.F + expf(-x));
}
struct GroupSums {
// Is it the 1st element of the group?
int32_t flag;
// The sum.
float sum;
// The sum of squares.
float sumSq;
};
struct GroupSumsOp {
inline __device__ GroupSums operator()(GroupSums const& a, GroupSums const& b) {
GroupSums dst;
dst.sum = b.flag ? b.sum : (a.sum + b.sum);
dst.sumSq = b.flag ? b.sumSq : (a.sumSq + b.sumSq);
dst.flag = a.flag + b.flag;
return dst;
}
};
template <typename T, typename U, int ILP>
inline __device__ void UpdateSum(const T* src, int64_t offset, U& sum, U& sumSq) {
using VecT = onnxruntime::rocm::aligned_vector<T, ILP>;
const VecT input_v = *reinterpret_cast<const VecT*>(src + offset);
#pragma unroll
for (int i = 0; i < ILP; i++) {
const U val = static_cast<U>(input_v.val[i]);
sum += val;
sumSq += val * val;
}
}
template <typename T, int ThreadsPerBlock, int ILP>
__global__ void groupNormNHWCSumKernel(const T* src, float* redBuffer, int32_t cPerBlock, int32_t hwPerBlock, int32_t hw,
int32_t hwc, int32_t c, int32_t cPerGroup, int32_t groups, int32_t groupsPerBlock) {
// The object in charge of doing the sums for the different blocks.
typedef hipcub::BlockScan<GroupSums, ThreadsPerBlock> BlockScan;
// Allocate shared memory for BlockScan.
__shared__ typename BlockScan::TempStorage tempStorage;
// Allocate shared memory for the groups. We could reduce the amount of shared
// memory reserved.
__shared__ float2 smem[ThreadsPerBlock];
// The instance in the batch.
int32_t ni = blockIdx.z;
// The channel loaded by that thread (ILP channels per thread).
int32_t ci = blockIdx.x * cPerBlock + threadIdx.x * ILP;
// The first activation loaded by that block.
int32_t hwBegin = blockIdx.y * hwPerBlock;
// The last activation loaded by that block.
int32_t hwEnd = min(hwBegin + hwPerBlock, hw);
// The sums.
float sum = 0.F;
float sumSq = 0.F;
// Iterate over the activations to compute the sums.
if (ci < c) {
for (int32_t hwi = hwBegin; hwi < hwEnd; ++hwi) {
// The offset.
int64_t offset = static_cast<int64_t>(ni) * hwc + static_cast<int64_t>(hwi) * c + ci;
UpdateSum<T, float, ILP>(src, offset, sum, sumSq);
}
}
// The group that thread works on and the channel in the group (modulus).
int32_t gi = threadIdx.x * ILP / cPerGroup;
int32_t cj = threadIdx.x * ILP - cPerGroup * gi;
// The data for the summations.
GroupSums inp{cj == 0 ? 1 : 0, sum, sumSq};
// Do the segmented scan.
GroupSums out;
BlockScan(tempStorage).InclusiveScan(inp, out, GroupSumsOp());
// Store the results for the groups in shared memory (to produce coalesced
// stores later).
if (cj == cPerGroup - ILP) { // ILP channels per thread
smem[gi] = make_float2(out.sum, out.sumSq);
}
// Make sure the data is in shared memory.
__syncthreads();
// The global group index.
int32_t gj = blockIdx.x * groupsPerBlock + threadIdx.x;
// Threads that have nothing left to do, exit.
if (threadIdx.x >= groupsPerBlock || gj >= groups) {
return;
}
// The first threads (those storing to global memory, load the values).
float2 sums = smem[threadIdx.x];
// Store to global memory.
atomicAdd(&redBuffer[(2 * ni + 0) * groups + gj], sums.x);
atomicAdd(&redBuffer[(2 * ni + 1) * groups + gj], sums.y);
}
template <typename T, typename U, int32_t ILP>
__device__ void computeGroupNorm(const T* src, T* dst, int64_t offset, U mean, U invStdDev,
const U* gamma_v, const U* beta_v, bool swish) {
using VecT = onnxruntime::rocm::aligned_vector<T, ILP>;
const VecT input_v = *reinterpret_cast<const VecT*>(src + offset);
VecT output_v;
#pragma unroll
for (int i = 0; i < ILP; i++) {
U val = static_cast<U>(input_v.val[i]);
val = (val - mean) * invStdDev;
val = gamma_v[i] * val + beta_v[i];
if (swish) {
val = val * sigmoid(val);
}
output_v.val[i] = static_cast<T>(val);
}
*(reinterpret_cast<VecT*>(dst + offset)) = output_v;
}
template <typename T, int ThreadsPerBlock, int ILP>
__global__ void groupNormNHWCScaleKernel(T* dst, const T* src, const float* gamma, const float* beta, const float* redBuffer, float epsilon, int32_t c, int32_t cPerBlock,
int32_t cPerGroup, int32_t groups, int32_t hwc, float invHWC, int32_t hw, int32_t hwPerBlock, bool withSwish) {
// The channel loaded by that thread (ILP channels per thread for F16x2).
int32_t ci = blockIdx.x * cPerBlock + threadIdx.x * ILP;
if (ci >= c) {
return;
}
// The instance in the batch.
int32_t ni = blockIdx.z;
// The group that thread works on and the channel in the group (modulus).
int32_t gi = ci / cPerGroup;
// Load the sum and sum of squares for the group.
float sum = 0.F, sumSq = 0.F;
if (gi < groups) {
sum = redBuffer[(2 * ni + 0) * groups + gi];
sumSq = redBuffer[(2 * ni + 1) * groups + gi];
}
using VecF = onnxruntime::rocm::aligned_vector<float, ILP>;
const VecF gamma_v = *reinterpret_cast<const VecF*>(gamma + ci);
const VecF beta_v = *reinterpret_cast<const VecF*>(beta + ci);
// Compute the mean.
float mean = sum * invHWC;
// Compute the variance.
float var = sumSq * invHWC - (mean * mean);
// Compute the inverse of the stddev.
float invStdDev = var <= 0.F ? 1.F : rsqrtf(var + epsilon);
// The first activation loaded by that block.
int32_t hwBegin = blockIdx.y * hwPerBlock;
// The last activation loaded by that block.
int32_t hwEnd = min(hwBegin + hwPerBlock, hw);
// Iterate over the activations to compute the sums.
for (int32_t hwi = hwBegin; hwi < hwEnd; ++hwi) {
// The src/dst offset.
int64_t offset = (int64_t)ni * hwc + hwi * c + ci;
// Fetch ILP channels per thread.
computeGroupNorm<T, float, ILP>(src, dst, offset, mean, invStdDev, gamma_v.val, beta_v.val, withSwish);
}
}
} // namespace rocm
} // namespace contrib
} // namespace onnxruntime

View file

@ -0,0 +1,191 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
#include <hip/hip_fp16.h>
#include <hip/hip_runtime_api.h>
#include "core/providers/rocm/cu_inc/common.cuh"
#include "core/providers/rocm/rocm_common.h"
#include "contrib_ops/rocm/diffusion/group_norm_common.h"
#include "contrib_ops/rocm/diffusion/group_norm_impl.h"
#include "contrib_ops/rocm/diffusion/group_norm_impl_kernel.cuh"
namespace onnxruntime {
namespace contrib {
namespace rocm {
using onnxruntime::rocm::GPU_WARP_SIZE;
template <typename T>
void groupNormNHWCSum(const GroupNormNHWCParams<T>* params) {
// Make sure the values are as we expect.
ORT_ENFORCE(params->c % params->cPerBlock == 0 && params->hw % params->hwPerBlock == 0);
// Make sure a group does not span multiple blocks.
ORT_ENFORCE(params->cPerBlock % params->cPerGroup == 0);
dim3 grid;
// The number of blocks to compute all the channels.
grid.x = params->c / params->cPerBlock;
// The number of blocks to compute all the activations in a given instance.
grid.y = CeilDiv(params->hw, params->hwPerBlock);
// The number of instances.
grid.z = params->n;
#define LAUNCH_GROUPNORM_SUM(ThreadsPerBlock, VecSize) \
groupNormNHWCSumKernel<T, ThreadsPerBlock, VecSize> \
<<<grid, ThreadsPerBlock, 0, params->stream>>>(params->src, params->redBuffer, params->cPerBlock, \
params->hwPerBlock, params->hw, params->hwc, params->c, \
params->cPerGroup, params->groups, params->groupsPerBlock); \
break;
switch (params->cPerBlock) {
case 320:
LAUNCH_GROUPNORM_SUM(256, 2)
case 480:
LAUNCH_GROUPNORM_SUM(256, 2)
case 256:
LAUNCH_GROUPNORM_SUM(128, 2)
case 128:
LAUNCH_GROUPNORM_SUM(64, 2)
default:
ORT_NOT_IMPLEMENTED("Not implemented");
}
}
template <typename T, int ThreadsPerBlock, int VecSize>
Status GroupNormNHWCSumOp(const GroupNormNHWCParams<T>* params) {
dim3 grid;
grid.x = params->c / params->cPerBlock;
grid.y = CeilDiv(params->hw, params->hwPerBlock);
grid.z = params->n;
groupNormNHWCSumKernel<T, ThreadsPerBlock, VecSize>
<<<grid, ThreadsPerBlock, 0, params->stream>>>(
params->src, params->redBuffer, params->cPerBlock, params->hwPerBlock,
params->hw, params->hwc, params->c, params->cPerGroup, params->groups, params->groupsPerBlock);
return HIP_CALL(hipGetLastError());
}
template <typename T>
void groupNormNHWCScale(const GroupNormNHWCParams<T>* params) {
// Make sure the dimensions are aligned with what we expect.
ORT_ENFORCE(params->c % params->cPerBlock == 0);
// Make sure a group does not span multiple blocks.
ORT_ENFORCE(params->cPerBlock % params->cPerGroup == 0);
dim3 grid;
// The number of blocks to compute all the channels.
grid.x = params->c / params->cPerBlock;
// The number of blocks to compute all the activations in a given instance.
grid.y = CeilDiv(params->hw, params->hwPerBlock);
// The number of instances.
grid.z = params->n;
#define LAUNCH_GROUPNORM_SCALE(ThreadsPerBlock, VecSize) \
groupNormNHWCScaleKernel<T, ThreadsPerBlock, VecSize> \
<<<grid, ThreadsPerBlock, 0, params->stream>>>(params->dst, params->src, params->gamma, params->beta, \
params->redBuffer, params->epsilon, params->c, params->cPerBlock, \
params->cPerGroup, params->groups, params->hwc, params->invHWC, \
params->hw, params->hwPerBlock, params->withSwish); \
break;
switch (params->cPerBlock) {
case 320:
LAUNCH_GROUPNORM_SCALE(256, 2)
case 480:
LAUNCH_GROUPNORM_SCALE(256, 2)
case 256:
LAUNCH_GROUPNORM_SCALE(128, 2)
case 128:
LAUNCH_GROUPNORM_SCALE(64, 2)
default:
ORT_NOT_IMPLEMENTED("Not implemented");
}
}
template <typename T, int ThreadsPerBlock, int VecSize>
Status GroupNormNHWCScaleOp(const GroupNormNHWCParams<T>* params) {
dim3 grid;
grid.x = params->c / params->cPerBlock;
grid.y = CeilDiv(params->hw, params->hwPerBlock);
grid.z = params->n;
groupNormNHWCScaleKernel<T, ThreadsPerBlock, VecSize>
<<<grid, ThreadsPerBlock, 0, params->stream>>>(
params->dst, params->src, params->gamma, params->beta, params->redBuffer, params->epsilon, params->c, params->cPerBlock,
params->cPerGroup, params->groups, params->hwc, params->invHWC, params->hw, params->hwPerBlock, params->withSwish);
return HIP_CALL(hipGetLastError());
}
template <typename T, int ThreadsPerBlock, int VecSize>
class GroupNormNHWCOp {
public:
Status operator()(const GroupNormNHWCParams<T>* params) {
HIP_RETURN_IF_ERROR(hipMemsetAsync(params->redBuffer, 0, GetGroupNormWorkspaceSizeInBytes(), params->stream));
auto status = GroupNormNHWCSumOp<T, ThreadsPerBlock, VecSize>(params);
ORT_RETURN_IF_ERROR(status);
HIP_RETURN_IF_ERROR(hipGetLastError());
status = GroupNormNHWCScaleOp<T, ThreadsPerBlock, VecSize>(params);
ORT_RETURN_IF_ERROR(status);
HIP_RETURN_IF_ERROR(hipGetLastError());
return Status::OK();
}
Status IsSupported(const GroupNormNHWCParams<T>* params) {
TUNABLE_OP_RETURN_UNSUPPORTED_ARGUMENT_IF(
!(params->c % VecSize == 0 && params->cPerGroup % VecSize == 0));
TUNABLE_OP_RETURN_UNSUPPORTED_ARGUMENT_IF(!(params->cPerBlock % params->cPerGroup == 0 &&
params->c % params->cPerBlock == 0 &&
params->hw % params->hwPerBlock == 0));
TUNABLE_OP_RETURN_UNSUPPORTED_ARGUMENT_IF(!(params->cPerBlock <= ThreadsPerBlock * VecSize &&
params->cPerBlock > (ThreadsPerBlock - GPU_WARP_SIZE) * VecSize));
return Status::OK();
}
};
template <typename T>
Status GroupNormNHWCStaticSelection(const GroupNormNHWCParams<T>* params) {
HIP_RETURN_IF_ERROR(hipMemsetAsync(params->redBuffer, 0, GetGroupNormWorkspaceSizeInBytes(), params->stream));
groupNormNHWCSum<T>(params);
HIP_RETURN_IF_ERROR(hipGetLastError());
groupNormNHWCScale<T>(params);
HIP_RETURN_IF_ERROR(hipGetLastError());
return Status::OK();
}
#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>{});
#define ADD_OP_FOR_ALL_THREADS_PER_BLOCK_ALL_VEC_SIZE(name) \
ADD_OP_FOR_ALL_VEC_SIZE(name, 64) \
ADD_OP_FOR_ALL_VEC_SIZE(name, 128) \
ADD_OP_FOR_ALL_VEC_SIZE(name, 192) \
ADD_OP_FOR_ALL_VEC_SIZE(name, 256) \
ADD_OP_FOR_ALL_VEC_SIZE(name, 320) \
ADD_OP_FOR_ALL_VEC_SIZE(name, 384) \
ADD_OP_FOR_ALL_VEC_SIZE(name, 448) \
ADD_OP_FOR_ALL_VEC_SIZE(name, 512)
template <typename T>
class GroupNormNHWCTunableOp : public TunableOp<GroupNormNHWCParams<T>> {
public:
GroupNormNHWCTunableOp() {
this->RegisterOp(GroupNormNHWCStaticSelection<T>);
ADD_OP_FOR_ALL_THREADS_PER_BLOCK_ALL_VEC_SIZE(GroupNormNHWCOp)
}
};
#undef ADD_OP_FOR_ALL_VEC_SIZE
#undef ADD_OP_FOR_ALL_THREADS_PER_BLOCK_ALL_VEC_SIZE
} // namespace rocm
} // namespace contrib
} // namespace onnxruntime

View file

@ -0,0 +1,220 @@
# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
# --------------------------------------------------------------------------
import re
import sys
from dataclasses import dataclass
from itertools import product
import kernel_explorer as ke
import numpy as np
import pytest
from utils import dtype_to_bytes
def get_sd_sizes():
batch_sizes = [1, 2]
height = [8, 16, 32, 64]
num_channels = [320, 640, 1280, 1920, 2560]
num_groups = [32]
return product(batch_sizes, height, num_channels, num_groups)
def dtype_to_funcs(dtype):
type_map = {
"float16": list(filter(lambda x: re.search("GroupNormNHWC.*_half", x), dir(ke))),
"float32": list(filter(lambda x: re.search("GroupNormNHWC.*_float", x), dir(ke))),
}
return type_map[dtype]
def sigmoid_function(x):
return 1.0 / (1.0 + np.exp(-x))
def group_norm(input_x, gamma, beta, num_groups, epsilon, with_swish):
n, h, w, c = input_x.shape
input_x = input_x.transpose([0, 3, 1, 2])
assert c % num_groups == 0
x = input_x.reshape((n, num_groups, -1))
mean = np.mean(x, axis=-1, keepdims=True)
var = np.var(x, axis=-1, keepdims=True)
x = (x - mean) / np.sqrt(var + epsilon)
x = x.reshape((n, c, h, w))
x = x.transpose([0, 2, 3, 1])
x = x * gamma + beta
if with_swish:
x = x * sigmoid_function(x)
return x
def run_group_norm(batch_size: int, height: int, num_channels: int, num_groups: int, dtype: str, func):
np.random.seed(0)
width = height
input_x = np.random.rand(batch_size, height, width, num_channels).astype(np.float32)
gamma = np.random.rand(num_channels).astype(np.float32)
beta = np.random.rand(num_channels).astype(np.float32)
# the size of workspace is defined in onnxruntime/contrib_ops/cuda/diffusion/group_norm_impl.h L18
workspace = np.random.rand((np.dtype(np.float32).itemsize * 2) * 32 * 32).astype(np.float32)
epsilon = 1e-05
output_y = np.random.rand(batch_size, height, width, num_channels).astype(dtype)
use_swish = True
host_x = input_x.astype(dtype)
input_d = ke.DeviceArray(host_x)
gamma_d = ke.DeviceArray(gamma)
beta_d = ke.DeviceArray(beta)
workspace_d = ke.DeviceArray(workspace)
y_d = ke.DeviceArray(output_y)
f = getattr(ke, func)
my_op = f(
y_d,
workspace_d,
input_d,
gamma_d,
beta_d,
batch_size,
height,
width,
num_channels,
num_groups,
epsilon,
use_swish,
)
if my_op.IsSupported():
my_op.Run()
y_d.UpdateHostNumpyArray()
y_ref = group_norm(input_x, gamma, beta, num_groups, epsilon, use_swish).astype(dtype)
np.testing.assert_allclose(y_ref, output_y, atol=1e-02)
dtypes = ["float32", "float16"]
@pytest.mark.parametrize("sd_sizes", get_sd_sizes())
@pytest.mark.parametrize("dtype", dtypes)
def test_skip_layer_norm(sd_sizes, dtype):
for func in dtype_to_funcs(dtype):
run_group_norm(*sd_sizes, dtype, func)
@dataclass
class GroupNormNHWCMetric(ke.BandwidthMetric):
batch_size: int
height: int
width: int
num_channels: int
groups: int
def report(self):
common = (
f"{self.dtype:<4} batch={self.batch_size:<4} height={self.height:<4} width={self.width:<4}"
+ f"num_channels={self.num_channels:<6} groups={self.groups:<4} {self.name}"
)
if self.duration > 0:
return f"{self.duration:.2f} us, {self.gbps:.2f} GB/s " + common
return "not supported " + common
def profile_group_norm_func(
batch_size: int, height: int, width: int, num_channels: int, num_groups: int, dtype: str, func
):
np.random.seed(0)
input_x = np.random.rand(batch_size, height, width, num_channels).astype(dtype)
gamma = np.random.rand(num_channels).astype(np.float32)
beta = np.random.rand(num_channels).astype(np.float32)
workspace = np.random.rand(np.dtype(np.float32).itemsize * 2 * 32 * 32).astype(np.float32)
epsilon = 0.05
output_y = np.random.rand(batch_size, height, width, num_channels).astype(dtype)
use_swish = True
input_d = ke.DeviceArray(input_x)
gamma_d = ke.DeviceArray(gamma)
beta_d = ke.DeviceArray(beta)
workspace_d = ke.DeviceArray(workspace)
y_d = ke.DeviceArray(output_y)
f = getattr(ke, func)
my_op = f(
y_d,
workspace_d,
input_d,
gamma_d,
beta_d,
batch_size,
height,
width,
num_channels,
num_groups,
epsilon,
use_swish,
)
duration_ms = -1
if my_op.IsSupported():
duration_ms = my_op.Profile()
total_bytes = (input_x.size * 2 + gamma.size * 2) * dtype_to_bytes(dtype)
ke.report(
GroupNormNHWCMetric(func, dtype, duration_ms, total_bytes, batch_size, height, width, num_channels, num_groups)
)
def profile_with_args(batch_size, height, width, num_channels, num_groups, dtype, sort=True):
with ke.benchmark(sort):
for func in dtype_to_funcs(dtype):
profile_group_norm_func(batch_size, height, width, num_channels, num_groups, dtype, func)
sd_profile_sizes = [
(2, 64, 64, 320, 32),
(2, 32, 32, 640, 32),
(2, 16, 16, 1280, 32),
(2, 64, 64, 640, 32),
(2, 16, 16, 2560, 32),
(2, 32, 32, 1280, 32),
(2, 32, 32, 1920, 32),
(2, 8, 8, 1280, 32),
(2, 64, 64, 960, 32),
(2, 32, 32, 960, 32),
(2, 32, 32, 320, 32),
(2, 16, 16, 640, 32),
(2, 16, 16, 1920, 32),
(2, 8, 8, 2560, 32),
]
def profile():
for dtype in dtypes:
for sd_size in sd_profile_sizes:
profile_with_args(*sd_size, dtype)
print()
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser()
group = parser.add_argument_group("profile with args")
group.add_argument("batch_size", type=int)
group.add_argument("height", type=int)
group.add_argument("width", type=int)
group.add_argument("num_channels", type=int)
group.add_argument("num_groups", type=int)
group.add_argument("dtype", choices=dtypes)
group.add_argument("--sort", action="store_true")
if len(sys.argv) == 1:
profile()
else:
args = parser.parse_args()
profile_with_args(
args.batch_size, args.height, args.width, args.num_channels, args.num_groups, args.dtype, args.sort
)

View file

@ -0,0 +1,135 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include <stdio.h>
#include <hip/hip_fp16.h>
#include <pybind11/pybind11.h>
#include "contrib_ops/rocm/diffusion/group_norm_common.h"
#include "contrib_ops/rocm/diffusion/group_norm_tunable_op.h"
#include "python/tools/kernel_explorer/device_array.h"
#include "python/tools/kernel_explorer/kernel_explorer_interface.h"
namespace py = pybind11;
namespace onnxruntime {
template <typename T, int ThreadsPerBlock, int VecSize>
class GroupNormNHWC : public IKernelExplorer {
public:
GroupNormNHWC(DeviceArray& output, DeviceArray& workspace, DeviceArray& input, DeviceArray& gamma, DeviceArray& beta,
int batch_size, int height, int width, int num_channels, int num_groups, float epsilon, bool use_swish)
: params_(TuningContext(), Stream(), static_cast<T*>(output.ptr()), static_cast<float*>(workspace.ptr()),
static_cast<T*>(input.ptr()), static_cast<float*>(gamma.ptr()), static_cast<float*>(beta.ptr()),
batch_size, height, width, num_channels, num_groups, epsilon, use_swish) {}
bool IsSupported() {
Status status = op_.IsSupported(&params_);
return status.IsOK();
}
void Run() override {
ORT_THROW_IF_ERROR(op_(&params_));
}
private:
using ParamsT = contrib::rocm::GroupNormNHWCParams<T>;
ParamsT params_{};
contrib::rocm::GroupNormNHWCOp<T, ThreadsPerBlock, VecSize> op_{};
};
template <typename T>
class GroupNormNHWCStaticSelection : public IKernelExplorer {
public:
GroupNormNHWCStaticSelection(DeviceArray& output, DeviceArray& workspace, DeviceArray& input, DeviceArray& gamma, DeviceArray& beta,
int batch_size, int height, int width, int num_channels, int num_groups, float epsilon, bool use_swish)
: params_(TuningContext(), Stream(), static_cast<T*>(output.ptr()), static_cast<float*>(workspace.ptr()),
static_cast<T*>(input.ptr()), static_cast<float*>(gamma.ptr()), static_cast<float*>(beta.ptr()),
batch_size, height, width, num_channels, num_groups, epsilon, use_swish) {}
void Run() override {
ORT_THROW_IF_ERROR((contrib::rocm::GroupNormNHWCStaticSelection<T>(&params_)));
}
bool IsSupported() {
Status status = contrib::rocm::GroupNormNHWCStaticSelection<T>(&params_);
return status.IsOK();
}
private:
using ParamsT = contrib::rocm::GroupNormNHWCParams<T>;
ParamsT params_{};
};
template <typename T>
class GroupNormNHWCTunable : public IKernelExplorer {
public:
GroupNormNHWCTunable(DeviceArray& output, DeviceArray& workspace, DeviceArray& input, DeviceArray& gamma, DeviceArray& beta,
int batch_size, int height, int width, int num_channels, int num_groups, float epsilon, bool use_swish)
: params_(TuningContext(), Stream(), static_cast<T*>(output.ptr()), static_cast<float*>(workspace.ptr()),
static_cast<T*>(input.ptr()), static_cast<float*>(gamma.ptr()), static_cast<float*>(beta.ptr()),
batch_size, height, width, num_channels, num_groups, epsilon, use_swish) {
params_.TuningContext()->EnableTunableOp();
}
void Run() override {
ORT_THROW_IF_ERROR(op_(&params_));
}
bool IsSupported() {
return true;
}
private:
using ParamsT = contrib::rocm::GroupNormNHWCParams<T>;
ParamsT params_{};
contrib::rocm::GroupNormNHWCTunableOp<T> op_{};
};
#define REGISTER_OP(name, type, threads_per_block, vec_size) \
py::class_<name<type, threads_per_block, vec_size>>(m, #name "_" #type "_" #threads_per_block "_" #vec_size) \
.def(py::init<DeviceArray&, DeviceArray&, DeviceArray&, DeviceArray&, DeviceArray&, \
int, int, int, int, int, float, bool>()) \
.def("SetRepeats", &name<type, threads_per_block, vec_size>::SetRepeats) \
.def("Profile", &name<type, threads_per_block, vec_size>::Profile) \
.def("Run", &name<type, threads_per_block, vec_size>::Run) \
.def("IsSupported", &name<type, threads_per_block, vec_size>::IsSupported);
#define REGISTER_OP_FOR_ALL_VEC_SIZE(name, type, threads_per_block) \
REGISTER_OP(name, type, threads_per_block, 1) \
REGISTER_OP(name, type, threads_per_block, 2) \
REGISTER_OP(name, type, threads_per_block, 4) \
REGISTER_OP(name, type, threads_per_block, 8) \
REGISTER_OP(name, type, threads_per_block, 16)
#define REGISTER_OP_FOR_ALL_THREADS_PER_BLOCK_ALL_VEC_SIZE(name, type) \
REGISTER_OP_FOR_ALL_VEC_SIZE(name, type, 64) \
REGISTER_OP_FOR_ALL_VEC_SIZE(name, type, 128) \
REGISTER_OP_FOR_ALL_VEC_SIZE(name, type, 192) \
REGISTER_OP_FOR_ALL_VEC_SIZE(name, type, 256) \
REGISTER_OP_FOR_ALL_VEC_SIZE(name, type, 320) \
REGISTER_OP_FOR_ALL_VEC_SIZE(name, type, 384) \
REGISTER_OP_FOR_ALL_VEC_SIZE(name, type, 448) \
REGISTER_OP_FOR_ALL_VEC_SIZE(name, type, 512)
#define REGISTER_OP_TYPED(name, type) \
py::class_<name<type>>(m, #name "_" #type) \
.def(py::init<DeviceArray&, DeviceArray&, DeviceArray&, DeviceArray&, DeviceArray&, \
int, int, int, int, int, float, bool>()) \
.def("SetRepeats", &name<type>::SetRepeats) \
.def("Profile", &name<type>::Profile) \
.def("Run", &name<type>::Run) \
.def("IsSupported", &name<type>::IsSupported);
KE_REGISTER(m) {
REGISTER_OP_FOR_ALL_THREADS_PER_BLOCK_ALL_VEC_SIZE(GroupNormNHWC, half);
REGISTER_OP_FOR_ALL_THREADS_PER_BLOCK_ALL_VEC_SIZE(GroupNormNHWC, float);
REGISTER_OP_TYPED(GroupNormNHWCTunable, half);
REGISTER_OP_TYPED(GroupNormNHWCTunable, float);
REGISTER_OP_TYPED(GroupNormNHWCStaticSelection, half);
REGISTER_OP_TYPED(GroupNormNHWCStaticSelection, float);
}
} // namespace onnxruntime