mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-07-12 02:05:50 +00:00
This commit is contained in:
Ding, Yi
4
example/45_elementwise_normalization/CMakeLists.txt
Normal file
4
example/45_elementwise_normalization/CMakeLists.txt
Normal file
@@ -0,0 +1,4 @@
|
||||
# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
add_example_executable(example_elementwise_layernorm_blockwise elementwise_layernorm_blockwise.cpp)
|
||||
86
example/45_elementwise_normalization/README.md
Normal file
86
example/45_elementwise_normalization/README.md
Normal file
@@ -0,0 +1,86 @@
|
||||
# Elementwise Normalization
|
||||
|
||||
This example demonstrates a fused **elementwise operation followed by normalization**. This pattern combines elementwise tensor arithmetic with a normalization operation in a single kernel, which is particularly useful for implementing custom normalization layers or fused activation-normalization blocks.
|
||||
|
||||
## Mathematical Formulation
|
||||
|
||||
The operation performs an elementwise computation followed by a normalization operation.
|
||||
|
||||
1. **Elementwise Stage**: An elementwise operation is applied to one or more input tensors.
|
||||
$C_{temp} = f(A, B, \dots)$
|
||||
Where `f` is a user-defined elementwise function that operates on corresponding elements of the input tensors.
|
||||
|
||||
2. **Normalization Stage**: The result is then normalized. The normalization can be performed along specified dimensions.
|
||||
- **Compute Statistics**: For each normalization group, compute the mean and variance.
|
||||
$\mu = \frac{1}{N} \sum C_{temp}$
|
||||
$\sigma^2 = \frac{1}{N} \sum (C_{temp} - \mu)^2$
|
||||
- **Normalize**: Apply the normalization formula.
|
||||
$\hat{C} = \frac{C_{temp} - \mu}{\sqrt{\sigma^2 + \epsilon}}$
|
||||
- **Scale and Shift**: Apply learnable parameters.
|
||||
$D = \gamma \cdot \hat{C} + \beta$
|
||||
|
||||
The key optimization is that the intermediate tensor `C_temp` is **never written to global memory**. The elementwise computation feeds directly into the normalization calculation.
|
||||
|
||||
## Algorithmic Strategy: Fused Elementwise with Online Normalization
|
||||
|
||||
The implementation combines elementwise computation with an online normalization algorithm.
|
||||
|
||||
1. **Grid Scheduling**: The normalization groups are distributed among thread blocks. Each block handles one or more normalization groups.
|
||||
|
||||
2. **Fused Two-Pass Algorithm**:
|
||||
- **Pass 1 - Compute Elementwise and Moments**:
|
||||
- Threads cooperatively load input tensors and apply the elementwise function `f`.
|
||||
- The elementwise results are kept in registers/shared memory.
|
||||
- **Welford's Algorithm**: Threads use Welford's online algorithm to compute the mean and variance of the elementwise results within their normalization group.
|
||||
- **Intra-Block Reduction**: A parallel reduction in shared memory computes the final statistics for the group.
|
||||
- **Pass 2 - Normalize and Store**:
|
||||
- Using the computed statistics, threads apply the normalization formula to their elementwise results.
|
||||
- The final normalized result is written to the output tensor `D`.
|
||||
|
||||
This approach ensures that the elementwise computation is performed only once, and the results are immediately consumed by the normalization process without requiring additional memory bandwidth.
|
||||
|
||||
## Source Code Organization
|
||||
|
||||
- [`elementwise_normalization_xdl.cpp`](./elementwise_normalization_xdl.cpp): The main example file. It sets up the input tensors, defines the elementwise operation and normalization parameters, and instantiates the `DeviceElementwiseNormalization` operation.
|
||||
- [`../../include/ck/tensor_operation/gpu/device/device_elementwise_normalization.hpp`](../../include/ck/tensor_operation/gpu/device/device_elementwise_normalization.hpp): The high-level device interface for the fused elementwise normalization operation.
|
||||
- The underlying grid-wise kernel implements the complex fusion of elementwise operations with the two-pass normalization algorithm.
|
||||
|
||||
## Build and Run
|
||||
|
||||
### Prerequisites
|
||||
Ensure the Composable Kernel library is built and installed.
|
||||
```bash
|
||||
cd /path/to/composable_kernel/build
|
||||
make -j install
|
||||
```
|
||||
|
||||
### Build the Example
|
||||
```bash
|
||||
cd /path/to/composable_kernel/example/45_elementwise_normalization
|
||||
mkdir build && cd build
|
||||
|
||||
cmake \
|
||||
-DCMAKE_CXX_COMPILER=/opt/rocm/bin/hipcc \
|
||||
-DCMAKE_PREFIX_PATH="/opt/rocm;${CK_INSTALL_PATH}" \
|
||||
..
|
||||
|
||||
make -j
|
||||
```
|
||||
|
||||
### Run the Example
|
||||
```bash
|
||||
# Run the example with default settings
|
||||
./elementwise_normalization_xdl
|
||||
|
||||
# Run with verification, data initialization, and timing
|
||||
./elementwise_normalization_xdl 1 2 1
|
||||
```
|
||||
|
||||
## Applications
|
||||
|
||||
This fused operation is valuable for implementing custom normalization layers and optimizing activation-normalization sequences.
|
||||
|
||||
- **Custom Activation-Normalization Blocks**: Some architectures use non-standard activation functions followed by normalization. For example, a Swish activation followed by layer normalization can be fused into a single kernel using this pattern.
|
||||
- **Residual Connection with Normalization**: In some variants of residual networks, the residual addition is immediately followed by normalization. This can be expressed as an elementwise addition (residual) followed by normalization.
|
||||
- **Preprocessing Pipelines**: In data preprocessing, tensors might need elementwise transformations (e.g., color space conversion) followed by normalization (e.g., standardization). This kernel can fuse these operations.
|
||||
- **Research Architectures**: Novel normalization techniques often involve custom elementwise operations before the normalization step. This kernel provides a flexible foundation for implementing such research ideas efficiently.
|
||||
@@ -0,0 +1,235 @@
|
||||
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
|
||||
// SPDX-License-Identifier: MIT
|
||||
|
||||
#include <iostream>
|
||||
#include <numeric>
|
||||
#include <initializer_list>
|
||||
#include <cstdlib>
|
||||
#include <getopt.h>
|
||||
|
||||
#include "ck/ck.hpp"
|
||||
#include "ck/utility/reduction_enums.hpp"
|
||||
#include "ck/tensor_operation/gpu/device/impl/device_elementwise_normalization_impl.hpp"
|
||||
#include "ck/tensor_operation/gpu/device/reduction_operator_mapping.hpp"
|
||||
|
||||
#include "ck/library/utility/check_err.hpp"
|
||||
#include "ck/library/utility/device_memory.hpp"
|
||||
#include "ck/library/utility/host_common_util.hpp"
|
||||
#include "ck/library/utility/host_tensor.hpp"
|
||||
#include "ck/library/utility/host_tensor_generator.hpp"
|
||||
#include "ck/library/reference_tensor_operation/cpu/reference_layernorm.hpp"
|
||||
|
||||
using ::ck::DeviceMem;
|
||||
using ::ck::HostTensorDescriptor;
|
||||
using ::ck::Tensor;
|
||||
|
||||
using ADataType = ck::half_t; // Input 1
|
||||
using BDataType = ck::half_t; // Input 2
|
||||
using XDataType = ck::half_t;
|
||||
using GammaDataType = ck::half_t;
|
||||
using BetaDataType = ck::half_t;
|
||||
using YDataType = ck::half_t;
|
||||
using AccDataType = float;
|
||||
using XElementwiseOperation = ck::tensor_operation::element_wise::Add;
|
||||
using YElementwiseOperation = ck::tensor_operation::element_wise::PassThrough;
|
||||
|
||||
constexpr int Rank = 2;
|
||||
constexpr int NumReduceDim = 1;
|
||||
|
||||
// X = Elementwise(input1, input2, input3, ...)
|
||||
// Y = Layernorm(X, beta, gamma)
|
||||
using DeviceInstance = ck::tensor_operation::device::DeviceElementwiseNormalizationImpl<
|
||||
ck::Tuple<ADataType, BDataType>,
|
||||
GammaDataType,
|
||||
BetaDataType,
|
||||
AccDataType,
|
||||
YDataType,
|
||||
XElementwiseOperation,
|
||||
YElementwiseOperation,
|
||||
Rank,
|
||||
NumReduceDim,
|
||||
256, // BlockSize
|
||||
8, // ClusterM
|
||||
32, // ClusterK
|
||||
1, // SliceM
|
||||
32, // SliceK
|
||||
1, // SrcVecDim (0=M, 1=K)
|
||||
8, // SrcScalarPerVector
|
||||
1, // GammaVecDim (0=M, 1=K)
|
||||
8, // GammaScalarPerVector
|
||||
1, // BetaVecDim (0=M, 1=K)
|
||||
8, // BetaScalarPerVector
|
||||
8>; // OutScalarPerVector
|
||||
|
||||
template <typename HostTensorA, typename HostTensorB, typename HostTensorC, typename Functor>
|
||||
void host_elementwise2D(HostTensorC& C,
|
||||
const HostTensorA& A,
|
||||
const HostTensorB& B,
|
||||
const std::vector<std::size_t>& shape,
|
||||
Functor functor)
|
||||
{
|
||||
using ctype = ck::remove_reference_t<decltype(C(0, 0))>;
|
||||
|
||||
for(std::size_t m = 0; m < shape[0]; ++m)
|
||||
for(std::size_t n = 0; n < shape[1]; ++n)
|
||||
{
|
||||
auto a_val = A(m, n);
|
||||
auto b_val = B(m, n);
|
||||
ctype c_val = 0;
|
||||
functor(c_val, a_val, b_val);
|
||||
C(m, n) = c_val;
|
||||
}
|
||||
}
|
||||
|
||||
int main(int argc, char* argv[])
|
||||
{
|
||||
bool do_verification = true;
|
||||
bool time_kernel = false;
|
||||
|
||||
ck::index_t M = 48 * 256;
|
||||
ck::index_t N = 1024;
|
||||
|
||||
if(argc == 1)
|
||||
{
|
||||
// use default
|
||||
}
|
||||
else if(argc == 3 || argc == 5)
|
||||
{
|
||||
do_verification = std::stoi(argv[1]);
|
||||
time_kernel = std::stoi(argv[2]);
|
||||
if(argc == 5)
|
||||
{
|
||||
M = std::stoi(argv[3]);
|
||||
N = std::stoi(argv[4]);
|
||||
}
|
||||
}
|
||||
else
|
||||
{
|
||||
printf("arg1: verification (0=no, 1=yes)\n");
|
||||
printf("arg2: time kernel (0=no, 1=yes)\n");
|
||||
printf("arg3-4: M, N\n");
|
||||
exit(1);
|
||||
}
|
||||
|
||||
ck::index_t Stride = N;
|
||||
|
||||
auto f_host_tensor_descriptor1d = [](std::size_t len, std::size_t stride) {
|
||||
return HostTensorDescriptor(std::vector<std::size_t>({len}),
|
||||
std::vector<std::size_t>({stride}));
|
||||
};
|
||||
|
||||
auto f_host_tensor_descriptor2d = [](std::size_t row, std::size_t col, std::size_t stride) {
|
||||
return HostTensorDescriptor(std::vector<std::size_t>({row, col}),
|
||||
std::vector<std::size_t>({stride, 1}));
|
||||
};
|
||||
|
||||
Tensor<ADataType> a(f_host_tensor_descriptor2d(M, N, Stride));
|
||||
Tensor<BDataType> b(f_host_tensor_descriptor2d(M, N, Stride));
|
||||
Tensor<GammaDataType> gamma(f_host_tensor_descriptor1d(N, 1));
|
||||
Tensor<BetaDataType> beta(f_host_tensor_descriptor1d(N, 1));
|
||||
Tensor<YDataType> y(f_host_tensor_descriptor2d(M, N, Stride));
|
||||
|
||||
a.GenerateTensorValue(GeneratorTensor_2<ADataType>{-5, 5});
|
||||
b.GenerateTensorValue(GeneratorTensor_2<BDataType>{-5, 5});
|
||||
gamma.GenerateTensorValue(GeneratorTensor_2<GammaDataType>{-5, 5});
|
||||
beta.GenerateTensorValue(GeneratorTensor_2<BetaDataType>{-5, 5});
|
||||
|
||||
DeviceMem a_dev(sizeof(ADataType) * a.mDesc.GetElementSpaceSize());
|
||||
DeviceMem b_dev(sizeof(BDataType) * b.mDesc.GetElementSpaceSize());
|
||||
DeviceMem gamma_dev(sizeof(GammaDataType) * gamma.mDesc.GetElementSpaceSize());
|
||||
DeviceMem beta_dev(sizeof(BetaDataType) * beta.mDesc.GetElementSpaceSize());
|
||||
DeviceMem y_dev(sizeof(YDataType) * y.mDesc.GetElementSpaceSize());
|
||||
|
||||
a_dev.ToDevice(a.mData.data());
|
||||
b_dev.ToDevice(b.mData.data());
|
||||
gamma_dev.ToDevice(gamma.mData.data());
|
||||
beta_dev.ToDevice(beta.mData.data());
|
||||
|
||||
std::array<const void*, 2> input = {a_dev.GetDeviceBuffer(), b_dev.GetDeviceBuffer()};
|
||||
|
||||
auto device_instance = DeviceInstance{};
|
||||
auto argument_ptr = device_instance.MakeArgumentPointer(
|
||||
{M, N},
|
||||
{
|
||||
std::vector<ck::index_t>{a.mDesc.GetStrides().begin(), a.mDesc.GetStrides().end()},
|
||||
std::vector<ck::index_t>{b.mDesc.GetStrides().begin(), b.mDesc.GetStrides().end()},
|
||||
},
|
||||
{0, 1},
|
||||
{0, 1},
|
||||
std::vector<ck::index_t>{y.mDesc.GetStrides().begin(), y.mDesc.GetStrides().end()},
|
||||
{1},
|
||||
1e-4,
|
||||
input,
|
||||
gamma_dev.GetDeviceBuffer(),
|
||||
beta_dev.GetDeviceBuffer(),
|
||||
y_dev.GetDeviceBuffer(),
|
||||
XElementwiseOperation{},
|
||||
YElementwiseOperation{});
|
||||
|
||||
if(!device_instance.IsSupportedArgument(argument_ptr.get()))
|
||||
{
|
||||
std::cout << "The runtime parameters are not supported" << std::endl;
|
||||
return 1;
|
||||
};
|
||||
|
||||
auto invoker_ptr = device_instance.MakeInvokerPointer();
|
||||
float ela_time = 0;
|
||||
ela_time = invoker_ptr->Run(argument_ptr.get(), StreamConfig{nullptr, time_kernel});
|
||||
|
||||
float data_mem_size = M * N * sizeof(ADataType) + M * N * sizeof(BDataType) +
|
||||
M * N * sizeof(YDataType) + N * sizeof(GammaDataType) +
|
||||
N * sizeof(BetaDataType);
|
||||
float bandwidth = data_mem_size * 1000 / ela_time / 1024 / 1024 / 1024;
|
||||
|
||||
std::cout << "Bandwidth is : " << bandwidth << "GB/s . " << std::endl;
|
||||
std::cout << "Time elapase is : " << ela_time << " ms . " << std::endl;
|
||||
|
||||
bool pass = true;
|
||||
if(do_verification)
|
||||
{
|
||||
std::vector<std::size_t> mn = {static_cast<unsigned long>(M),
|
||||
static_cast<unsigned long>(N)};
|
||||
Tensor<XDataType> x(f_host_tensor_descriptor2d(M, N, Stride));
|
||||
host_elementwise2D<Tensor<ADataType>,
|
||||
Tensor<BDataType>,
|
||||
Tensor<XDataType>,
|
||||
XElementwiseOperation>(x, a, b, mn, XElementwiseOperation{});
|
||||
|
||||
Tensor<YDataType> host_y(f_host_tensor_descriptor2d(M, N, Stride));
|
||||
Tensor<AccDataType> host_save_mean({M});
|
||||
Tensor<AccDataType> host_save_inv_std({M});
|
||||
using ReferenceInstance =
|
||||
ck::tensor_operation::host::ReferenceLayernorm<XDataType,
|
||||
GammaDataType,
|
||||
BetaDataType,
|
||||
YDataType,
|
||||
AccDataType,
|
||||
AccDataType,
|
||||
YElementwiseOperation,
|
||||
Rank,
|
||||
NumReduceDim>;
|
||||
|
||||
ReferenceInstance ref;
|
||||
auto ref_argument = ref.MakeArgument(x,
|
||||
gamma,
|
||||
beta,
|
||||
host_y,
|
||||
host_save_mean,
|
||||
host_save_inv_std,
|
||||
YElementwiseOperation{},
|
||||
{M, N},
|
||||
{1},
|
||||
1e-4);
|
||||
auto ref_invoker = ref.MakeInvoker();
|
||||
ref_invoker.Run(ref_argument);
|
||||
|
||||
y_dev.FromDevice(y.mData.data());
|
||||
pass &=
|
||||
ck::utils::check_err(y.mData, host_y.mData, "Error: Incorrect results d1", 1e-3, 1e-3);
|
||||
if(!(pass))
|
||||
{
|
||||
std::cout << "layernorm wrong" << std::endl;
|
||||
}
|
||||
}
|
||||
return (pass ? 0 : 1);
|
||||
}
|
||||
Reference in New Issue
Block a user