mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-25 19:48:11 +00:00
[ROCm] Add Tunable GroupNorm (#15298)
refactor GroupNorm and Add Tunable GroupNorm
This commit is contained in:
parent
db4fc12318
commit
ce1eb6d629
6 changed files with 886 additions and 383 deletions
118
onnxruntime/contrib_ops/rocm/diffusion/group_norm_common.h
Normal file
118
onnxruntime/contrib_ops/rocm/diffusion/group_norm_common.h
Normal 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
|
||||
|
|
@ -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(¶ms.redBuffer[(2 * ni + 0) * params.groups + gj], sums.x);
|
||||
atomicAdd(¶ms.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(¶ms);
|
||||
}
|
||||
|
||||
template Status LaunchGroupNormKernel<half>(hipStream_t stream, half* output,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
191
onnxruntime/contrib_ops/rocm/diffusion/group_norm_tunable_op.h
Normal file
191
onnxruntime/contrib_ops/rocm/diffusion/group_norm_tunable_op.h
Normal 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
|
||||
|
|
@ -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
|
||||
)
|
||||
|
|
@ -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(¶ms_);
|
||||
return status.IsOK();
|
||||
}
|
||||
|
||||
void Run() override {
|
||||
ORT_THROW_IF_ERROR(op_(¶ms_));
|
||||
}
|
||||
|
||||
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>(¶ms_)));
|
||||
}
|
||||
|
||||
bool IsSupported() {
|
||||
Status status = contrib::rocm::GroupNormNHWCStaticSelection<T>(¶ms_);
|
||||
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_(¶ms_));
|
||||
}
|
||||
|
||||
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
|
||||
Loading…
Reference in a new issue