From ce1eb6d6299b1b71a3b1f6b3647ec6aabfb6df94 Mon Sep 17 00:00:00 2001 From: PeixuanZuo <94887879+PeixuanZuo@users.noreply.github.com> Date: Wed, 12 Apr 2023 10:55:42 +0800 Subject: [PATCH] [ROCm] Add Tunable GroupNorm (#15298) refactor GroupNorm and Add Tunable GroupNorm --- .../rocm/diffusion/group_norm_common.h | 118 ++++++ .../rocm/diffusion/group_norm_impl.cu | 392 +----------------- .../rocm/diffusion/group_norm_impl_kernel.cuh | 213 ++++++++++ .../rocm/diffusion/group_norm_tunable_op.h | 191 +++++++++ .../kernel_explorer/kernels/groupnorm_test.py | 220 ++++++++++ .../kernel_explorer/kernels/rocm/groupnorm.cu | 135 ++++++ 6 files changed, 886 insertions(+), 383 deletions(-) create mode 100644 onnxruntime/contrib_ops/rocm/diffusion/group_norm_common.h create mode 100644 onnxruntime/contrib_ops/rocm/diffusion/group_norm_impl_kernel.cuh create mode 100644 onnxruntime/contrib_ops/rocm/diffusion/group_norm_tunable_op.h create mode 100644 onnxruntime/python/tools/kernel_explorer/kernels/groupnorm_test.py create mode 100644 onnxruntime/python/tools/kernel_explorer/kernels/rocm/groupnorm.cu diff --git a/onnxruntime/contrib_ops/rocm/diffusion/group_norm_common.h b/onnxruntime/contrib_ops/rocm/diffusion/group_norm_common.h new file mode 100644 index 0000000000..3ccc527526 --- /dev/null +++ b/onnxruntime/contrib_ops/rocm/diffusion/group_norm_common.h @@ -0,0 +1,118 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include + +#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 +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 diff --git a/onnxruntime/contrib_ops/rocm/diffusion/group_norm_impl.cu b/onnxruntime/contrib_ops/rocm/diffusion/group_norm_impl.cu index 1b8a55cfa2..1cec66a9a3 100644 --- a/onnxruntime/contrib_ops/rocm/diffusion/group_norm_impl.cu +++ b/onnxruntime/contrib_ops/rocm/diffusion/group_norm_impl.cu @@ -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 -#include -#include -#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 +#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 -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 -inline __device__ void UpdateSum(const T* src, int64_t offset, U& sum, U& sumSq) { - using VecT = onnxruntime::rocm::aligned_vector; - const VecT input_v = *reinterpret_cast(src + offset); - - #pragma unroll - for (int i = 0; i < ILP; i++) { - const U val = static_cast(input_v.val[i]); - sum += val; - sumSq += val * val; - } -} - -template -__global__ void groupNormNHWCSumKernel(GroupNormNHWCParams params) { - // The object in charge of doing the sums for the different blocks. - typedef hipcub::BlockScan 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(ni) * params.hwc + static_cast(hwi) * params.c + ci; - UpdateSum(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(¶ms.redBuffer[(2 * ni + 0) * params.groups + gj], sums.x); - atomicAdd(¶ms.redBuffer[(2 * ni + 1) * params.groups + gj], sums.y); -} - -template -void groupNormNHWCSum(GroupNormNHWCParams 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<<>>(params); - break; - case 480: - groupNormNHWCSumKernel<<>>(params); - break; - case 256: - groupNormNHWCSumKernel<<>>(params); - break; - case 128: - groupNormNHWCSumKernel<<>>(params); - break; - default: - ORT_NOT_IMPLEMENTED("Not implemented"); - } -} - -template -__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; - const VecT input_v = *reinterpret_cast(src + offset); - VecT output_v; - - #pragma unroll - for (int i = 0; i < ILP; i++) { - U val = static_cast(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(val); - } - *(reinterpret_cast(dst + offset)) = output_v; -} - -template -__global__ void groupNormNHWCScaleKernel(GroupNormNHWCParams 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; - - const VecF gamma_v = *reinterpret_cast(params.gamma + ci); - const VecF beta_v = *reinterpret_cast(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(params.src, params.dst, offset, mean, invStdDev, gamma_v.val, beta_v.val, params.withSwish); - } -} - -template -void groupNormNHWCScale(GroupNormNHWCParams 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<<>>(params); - break; - case 480: - groupNormNHWCScaleKernel<<>>(params); - break; - case 256: - groupNormNHWCScaleKernel<<>>(params); - break; - case 128: - groupNormNHWCScaleKernel<<>>(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 Status LaunchGroupNormKernel( hipStream_t stream, @@ -364,57 +37,10 @@ Status LaunchGroupNormKernel( "only num_groups=32 is supported. Got", num_groups); } - GroupNormNHWCParams 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 params(nullptr, stream, output, reinterpret_cast(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(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(params, stream); - DUMP_TENSOR("workspace", params.redBuffer, batch_size, num_groups, 2); - HIP_RETURN_IF_ERROR(hipGetLastError()); - groupNormNHWCScale(params, stream); - HIP_RETURN_IF_ERROR(hipGetLastError()); - DUMP_TENSOR("output", output, batch_size, num_channels, height * width); - return Status::OK(); + return GroupNormNHWCStaticSelection(¶ms); } template Status LaunchGroupNormKernel(hipStream_t stream, half* output, diff --git a/onnxruntime/contrib_ops/rocm/diffusion/group_norm_impl_kernel.cuh b/onnxruntime/contrib_ops/rocm/diffusion/group_norm_impl_kernel.cuh new file mode 100644 index 0000000000..d6322a12a9 --- /dev/null +++ b/onnxruntime/contrib_ops/rocm/diffusion/group_norm_impl_kernel.cuh @@ -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 +#include +#include +#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 +inline __device__ void UpdateSum(const T* src, int64_t offset, U& sum, U& sumSq) { + using VecT = onnxruntime::rocm::aligned_vector; + const VecT input_v = *reinterpret_cast(src + offset); + +#pragma unroll + for (int i = 0; i < ILP; i++) { + const U val = static_cast(input_v.val[i]); + sum += val; + sumSq += val * val; + } +} + +template +__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 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(ni) * hwc + static_cast(hwi) * c + ci; + UpdateSum(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 +__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; + const VecT input_v = *reinterpret_cast(src + offset); + VecT output_v; + +#pragma unroll + for (int i = 0; i < ILP; i++) { + U val = static_cast(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(val); + } + *(reinterpret_cast(dst + offset)) = output_v; +} + +template +__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; + + const VecF gamma_v = *reinterpret_cast(gamma + ci); + const VecF beta_v = *reinterpret_cast(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(src, dst, offset, mean, invStdDev, gamma_v.val, beta_v.val, withSwish); + } +} + +} // namespace rocm +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/rocm/diffusion/group_norm_tunable_op.h b/onnxruntime/contrib_ops/rocm/diffusion/group_norm_tunable_op.h new file mode 100644 index 0000000000..adeaf2c126 --- /dev/null +++ b/onnxruntime/contrib_ops/rocm/diffusion/group_norm_tunable_op.h @@ -0,0 +1,191 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include +#include +#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 +void groupNormNHWCSum(const GroupNormNHWCParams* 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 \ + <<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 +Status GroupNormNHWCSumOp(const GroupNormNHWCParams* params) { + dim3 grid; + grid.x = params->c / params->cPerBlock; + grid.y = CeilDiv(params->hw, params->hwPerBlock); + grid.z = params->n; + + groupNormNHWCSumKernel + <<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 +void groupNormNHWCScale(const GroupNormNHWCParams* 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 \ + <<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 +Status GroupNormNHWCScaleOp(const GroupNormNHWCParams* params) { + dim3 grid; + grid.x = params->c / params->cPerBlock; + grid.y = CeilDiv(params->hw, params->hwPerBlock); + grid.z = params->n; + + groupNormNHWCScaleKernel + <<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 +class GroupNormNHWCOp { + public: + Status operator()(const GroupNormNHWCParams* params) { + HIP_RETURN_IF_ERROR(hipMemsetAsync(params->redBuffer, 0, GetGroupNormWorkspaceSizeInBytes(), params->stream)); + auto status = GroupNormNHWCSumOp(params); + ORT_RETURN_IF_ERROR(status); + HIP_RETURN_IF_ERROR(hipGetLastError()); + status = GroupNormNHWCScaleOp(params); + ORT_RETURN_IF_ERROR(status); + HIP_RETURN_IF_ERROR(hipGetLastError()); + return Status::OK(); + } + + Status IsSupported(const GroupNormNHWCParams* 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 +Status GroupNormNHWCStaticSelection(const GroupNormNHWCParams* params) { + HIP_RETURN_IF_ERROR(hipMemsetAsync(params->redBuffer, 0, GetGroupNormWorkspaceSizeInBytes(), params->stream)); + groupNormNHWCSum(params); + HIP_RETURN_IF_ERROR(hipGetLastError()); + groupNormNHWCScale(params); + HIP_RETURN_IF_ERROR(hipGetLastError()); + return Status::OK(); +} + +#define ADD_OP_FOR_ALL_VEC_SIZE(name, threads_per_block) \ + this->RegisterOp(name{}); \ + this->RegisterOp(name{}); \ + this->RegisterOp(name{}); \ + this->RegisterOp(name{}); \ + this->RegisterOp(name{}); + +#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 +class GroupNormNHWCTunableOp : public TunableOp> { + public: + GroupNormNHWCTunableOp() { + this->RegisterOp(GroupNormNHWCStaticSelection); + 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 diff --git a/onnxruntime/python/tools/kernel_explorer/kernels/groupnorm_test.py b/onnxruntime/python/tools/kernel_explorer/kernels/groupnorm_test.py new file mode 100644 index 0000000000..f454ebf60b --- /dev/null +++ b/onnxruntime/python/tools/kernel_explorer/kernels/groupnorm_test.py @@ -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 + ) diff --git a/onnxruntime/python/tools/kernel_explorer/kernels/rocm/groupnorm.cu b/onnxruntime/python/tools/kernel_explorer/kernels/rocm/groupnorm.cu new file mode 100644 index 0000000000..b878d692e7 --- /dev/null +++ b/onnxruntime/python/tools/kernel_explorer/kernels/rocm/groupnorm.cu @@ -0,0 +1,135 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include +#include +#include + +#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 +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(output.ptr()), static_cast(workspace.ptr()), + static_cast(input.ptr()), static_cast(gamma.ptr()), static_cast(beta.ptr()), + batch_size, height, width, num_channels, num_groups, epsilon, use_swish) {} + + bool IsSupported() { + Status status = op_.IsSupported(¶ms_); + return status.IsOK(); + } + + void Run() override { + ORT_THROW_IF_ERROR(op_(¶ms_)); + } + + private: + using ParamsT = contrib::rocm::GroupNormNHWCParams; + ParamsT params_{}; + contrib::rocm::GroupNormNHWCOp op_{}; +}; + +template +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(output.ptr()), static_cast(workspace.ptr()), + static_cast(input.ptr()), static_cast(gamma.ptr()), static_cast(beta.ptr()), + batch_size, height, width, num_channels, num_groups, epsilon, use_swish) {} + + void Run() override { + ORT_THROW_IF_ERROR((contrib::rocm::GroupNormNHWCStaticSelection(¶ms_))); + } + + bool IsSupported() { + Status status = contrib::rocm::GroupNormNHWCStaticSelection(¶ms_); + return status.IsOK(); + } + + private: + using ParamsT = contrib::rocm::GroupNormNHWCParams; + ParamsT params_{}; +}; + +template +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(output.ptr()), static_cast(workspace.ptr()), + static_cast(input.ptr()), static_cast(gamma.ptr()), static_cast(beta.ptr()), + batch_size, height, width, num_channels, num_groups, epsilon, use_swish) { + params_.TuningContext()->EnableTunableOp(); + } + + void Run() override { + ORT_THROW_IF_ERROR(op_(¶ms_)); + } + + bool IsSupported() { + return true; + } + + private: + using ParamsT = contrib::rocm::GroupNormNHWCParams; + ParamsT params_{}; + contrib::rocm::GroupNormNHWCTunableOp op_{}; +}; + +#define REGISTER_OP(name, type, threads_per_block, vec_size) \ + py::class_>(m, #name "_" #type "_" #threads_per_block "_" #vec_size) \ + .def(py::init()) \ + .def("SetRepeats", &name::SetRepeats) \ + .def("Profile", &name::Profile) \ + .def("Run", &name::Run) \ + .def("IsSupported", &name::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_>(m, #name "_" #type) \ + .def(py::init()) \ + .def("SetRepeats", &name::SetRepeats) \ + .def("Profile", &name::Profile) \ + .def("Run", &name::Run) \ + .def("IsSupported", &name::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