ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-16 08:44:55 +00:00
Code Issues Packages Projects Releases Wiki Activity

add tiling, pipeline, follow layernorm2d

Browse Source
This commit is contained in:
Mohsen Saffari
2025-11-28 18:29:10 +00:00
parent c4fbbaaa55
commit 3194b653f7
6 changed files with 347 additions and 151 deletions
Download Patch File Download Diff File Expand all files Collapse all files

22
example/ck_tile/42_batchnorm/batchnorm_simple.cpp
View File

@@ -139,8 +139,13 @@ bool run(const ck_tile::ArgParser& arg_parser)
// Fill input with random data
ck_tile::FillUniformDistribution<XDataType>{-5.f, 5.f}(x_host);
ck_tile::FillUniformDistribution<ComputeDataType>{0.8f, 1.2f}(gamma_host); // Scale around 1.0
ck_tile::FillUniformDistribution<ComputeDataType>{-0.5f, 0.5f}(beta_host); // Bias around 0.0
// Set gamma=1.0 and beta=0.0 for testing (identity transform)
for(ck_tile::index_t c = 0; c < C; ++c)
{
gamma_host.mData[c] = static_cast<ComputeDataType>(1.0);
beta_host.mData[c] = static_cast<ComputeDataType>(0.0);
}
// Allocate device memory
ck_tile::DeviceMem x_buf(x_host.get_element_space_size_in_bytes());
@@ -152,13 +157,14 @@ bool run(const ck_tile::ArgParser& arg_parser)
gamma_buf.ToDevice(gamma_host.data());
beta_buf.ToDevice(beta_host.data());
// Define kernel configuration
using BlockWarps = ck_tile::sequence<4, 1>;
using BlockTile = ck_tile::sequence<1, 256>; // Simplified for POC
using WarpTile = ck_tile::sequence<1, 256>;
using Vector = ck_tile::sequence<1, 1>;
// Define kernel configuration using Generic2dBlockShape
// For N=2, H=8, W=8: per-channel elements = 2×8×8 = 128
// Use Repeat_N=2: Block_N = Repeat_N × ThreadPerBlock_N × Vector_N = 2×64×1 = 128 ✓
using BlockTile = ck_tile::sequence<1, 128>; // Block size: 1 channel, 128 spatial
using ThreadPerBlock = ck_tile::sequence<1, 128>; // 64 threads (must be warp size)
using Vector = ck_tile::sequence<1, 1>; // Vector size (start with 1)
using Shape = ck_tile::BatchnormShape<BlockWarps, BlockTile, WarpTile, Vector>;
using Shape = ck_tile::BatchnormShape<BlockTile, ThreadPerBlock, Vector>;
// Define traits (compile-time configuration)
using Traits = ck_tile::BatchnormFwdTraits<false, false>; // No save, no update

2
include/ck_tile/ops/batchnorm.hpp
View File

@@ -5,6 +5,8 @@
#include "ck_tile/ops/batchnorm/block/block_welford.hpp"
#include "ck_tile/ops/batchnorm/kernel/batchnorm_fwd_kernel.hpp"
#include "ck_tile/ops/batchnorm/pipeline/batchnorm_fwd_pipeline.hpp"
#include "ck_tile/ops/batchnorm/pipeline/batchnorm_fwd_policy.hpp"
#include "ck_tile/ops/batchnorm/pipeline/batchnorm_fwd_traits.hpp"
#include "ck_tile/ops/batchnorm/pipeline/batchnorm_problem.hpp"
#include "ck_tile/ops/batchnorm/pipeline/batchnorm_shape.hpp"

194
include/ck_tile/ops/batchnorm/kernel/batchnorm_fwd_kernel.hpp
View File

@@ -5,6 +5,7 @@
#include "ck_tile/core.hpp"
#include "ck_tile/ops/batchnorm/block/block_welford.hpp"
#include "ck_tile/ops/batchnorm/pipeline/batchnorm_fwd_pipeline.hpp"
#include "ck_tile/ops/batchnorm/pipeline/batchnorm_problem.hpp"
#include "ck_tile/ops/batchnorm/pipeline/batchnorm_shape.hpp"
@@ -38,11 +39,14 @@ struct BatchnormFwd
{
using Problem = remove_cvref_t<Problem_>;
using XDataType = typename Problem::XDataType;
using GammaDataType = typename Problem::GammaDataType;
using BetaDataType = typename Problem::BetaDataType;
using ComputeDataType = typename Problem::ComputeDataType;
using YDataType = typename Problem::YDataType;
using MeanVarDataType = typename Problem::MeanVarDataType;
using BlockShape = typename Problem::BlockShape;
static constexpr index_t kBlockSize = BlockShape::kBlockSize;
static constexpr index_t kBlockSize = BlockShape::BlockSize;
// Kernel arguments
struct BatchnormFwdKargs
@@ -98,130 +102,110 @@ struct BatchnormFwd
return kBlockSize;
}
// Shared memory size (must be constexpr for __shared__ allocation)
CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSize()
{
return BatchnormFwdPipeline<Problem>::GetSmemSize();
}
CK_TILE_DEVICE void operator()(Kargs kargs) const
{
// Cast pointers to typed pointers
const XDataType* p_x = static_cast<const XDataType*>(kargs.p_x);
YDataType* p_y = static_cast<YDataType*>(kargs.p_y);
using Pipeline = BatchnormFwdPipeline<Problem>;
const index_t N = kargs.N;
const index_t C = kargs.C;
const index_t H = kargs.H;
const index_t W = kargs.W;
const ComputeDataType epsilon = static_cast<ComputeDataType>(kargs.epsilon);
const index_t spatial_size = H * W;
const index_t per_channel_size = N * spatial_size; // Reduce over N×H×W
const index_t thread_id = get_thread_id();
const index_t block_id = get_block_id();
// Each block handles one channel
const index_t c = block_id;
const index_t c = block_id; // Channel index
if(c >= C)
return;
// Thread-local Welford statistics
ComputeDataType thread_mean = type_convert<ComputeDataType>(0);
ComputeDataType thread_m2 = type_convert<ComputeDataType>(0);
index_t thread_count = 0;
const index_t spatial_size = H * W;
const index_t per_channel_size = N * spatial_size;
// Each thread processes elements across ALL samples (N) and spatial positions (H×W)
// for this channel
for(index_t idx = thread_id; idx < per_channel_size; idx += kBlockSize)
{
// Calculate which sample (n) and spatial position (hw) this is
const index_t n = idx / spatial_size;
const index_t hw = idx % spatial_size;
// Use block dimensions from BlockShape (like layernorm2d)
static constexpr index_t Block_M = BlockShape::Block_M;
static constexpr index_t Block_N = BlockShape::Block_N;
// Create tensor views WITHOUT distributions (will be applied in pipeline)
const auto x_window = [&]() {
const XDataType* p_x = static_cast<const XDataType*>(kargs.p_x);
const auto tmp_ = make_naive_tensor_view<address_space_enum::global>(
p_x + c * spatial_size,
make_tuple(N, spatial_size),
make_tuple(C * spatial_size, 1),
number<1>{},
number<1>{});
// Memory layout: [N, C, H, W] with C-major for H×W
const index_t offset = n * C * H * W + c * H * W + hw;
ComputeDataType val = type_convert<ComputeDataType>(p_x[offset]);
const auto tmp2_ = pad_tensor_view(
tmp_, make_tuple(number<Block_M>{}, number<Block_N>{}), sequence<false, false>{});
// Online Welford update
thread_count++;
ComputeDataType delta = val - thread_mean;
thread_mean += delta / type_convert<ComputeDataType>(thread_count);
ComputeDataType delta2 = val - thread_mean;
thread_m2 += delta * delta2;
}
return make_tile_window(tmp2_, make_tuple(number<Block_M>{}, number<Block_N>{}), {0, 0});
}();
// Allocate shared memory for block-level reduction
__shared__ char smem[BlockWelford<ComputeDataType>::template GetSmemSize<index_t, kBlockSize>()];
// Block-level Welford reduction
ComputeDataType block_mean = thread_mean;
ComputeDataType block_var = thread_m2;
index_t block_count = thread_count;
BlockWelford<ComputeDataType>::template Run<index_t, kBlockSize>(
block_mean, block_var, block_count, smem);
// Load scale (gamma) and bias (beta) for this channel
// Following old CK pattern: gamma/beta are ALWAYS provided (no nullptr checks)
// All threads load (efficient, no branching)
const ComputeDataType* p_gamma = static_cast<const ComputeDataType*>(kargs.p_gamma);
const ComputeDataType* p_beta = static_cast<const ComputeDataType*>(kargs.p_beta);
const ComputeDataType gamma = p_gamma[c];
const ComputeDataType beta = p_beta[c];
// Compute inverse standard deviation
ComputeDataType inv_std = type_convert<ComputeDataType>(1) /
ck_tile::sqrt(block_var + epsilon);
// Normalize and write output with scale and bias
// Formula: y = gamma * (x - mean) / std + beta
// = gamma * (x - mean) * inv_std + beta
for(index_t idx = thread_id; idx < per_channel_size; idx += kBlockSize)
{
const index_t n = idx / spatial_size;
const index_t hw = idx % spatial_size;
const auto gamma_window = [&]() {
const GammaDataType* p_gamma = static_cast<const GammaDataType*>(kargs.p_gamma);
const auto tmp_ = make_naive_tensor_view_packed<address_space_enum::global>(
p_gamma + c,
make_tuple(1),
number<1>{});
const index_t offset = (n * C * H * W) + (c * H * W) + hw;
ComputeDataType val = type_convert<ComputeDataType>(p_x[offset]);
// Apply batch normalization with scale and bias
ComputeDataType normalized = gamma * ((val - block_mean) * inv_std) + beta;
p_y[offset] = type_convert<YDataType>(normalized);
}
const auto tmp2_ = pad_tensor_view(tmp_, make_tuple(number<Block_M>{}), sequence<false>{});
return make_tile_window(tmp2_, make_tuple(number<Block_M>{}), {0});
}();
// Save mean and inverse std for backward pass (compile-time check)
if constexpr(Problem::Traits::kSaveMeanInvStd)
{
if(thread_id == 0)
{
using MeanVarDataType = typename Problem::MeanVarDataType;
MeanVarDataType* p_save_mean = static_cast<MeanVarDataType*>(kargs.p_save_mean);
MeanVarDataType* p_save_inv_std = static_cast<MeanVarDataType*>(kargs.p_save_inv_std);
p_save_mean[c] = type_convert<MeanVarDataType>(block_mean);
p_save_inv_std[c] = type_convert<MeanVarDataType>(inv_std);
}
}
const auto beta_window = [&]() {
const BetaDataType* p_beta = static_cast<const BetaDataType*>(kargs.p_beta);
const auto tmp_ = make_naive_tensor_view_packed<address_space_enum::global>(
p_beta + c,
make_tuple(1),
number<1>{});
const auto tmp2_ = pad_tensor_view(tmp_, make_tuple(number<Block_M>{}), sequence<false>{});
return make_tile_window(tmp2_, make_tuple(number<Block_M>{}), {0});
}();
// Update running mean and variance (compile-time check)
if constexpr(Problem::Traits::kUpdateMovingAverage)
{
if(thread_id == 0)
{
using MeanVarDataType = typename Problem::MeanVarDataType;
MeanVarDataType* p_running_mean = static_cast<MeanVarDataType*>(kargs.p_running_mean);
MeanVarDataType* p_running_var = static_cast<MeanVarDataType*>(kargs.p_running_var);
const ComputeDataType momentum = static_cast<ComputeDataType>(kargs.momentum);
const ComputeDataType one_minus_momentum = type_convert<ComputeDataType>(1) - momentum;
// Exponential moving average: new = (1-momentum)*old + momentum*current
ComputeDataType old_mean = type_convert<ComputeDataType>(p_running_mean[c]);
ComputeDataType old_var = type_convert<ComputeDataType>(p_running_var[c]);
p_running_mean[c] = type_convert<MeanVarDataType>(one_minus_momentum * old_mean + momentum * block_mean);
p_running_var[c] = type_convert<MeanVarDataType>(one_minus_momentum * old_var + momentum * block_var);
}
}
auto y_window = [&]() {
YDataType* p_y = static_cast<YDataType*>(kargs.p_y);
const auto tmp_ = make_naive_tensor_view<address_space_enum::global>(
p_y + c * spatial_size,
make_tuple(N, spatial_size),
make_tuple(C * spatial_size, 1),
number<1>{},
number<1>{});
const auto tmp2_ = pad_tensor_view(
tmp_, make_tuple(number<Block_M>{}, number<Block_N>{}), sequence<false, false>{});
return make_tile_window(tmp2_, make_tuple(number<Block_M>{}, number<Block_N>{}), {0, 0});
}();
// Allocate shared memory (use kernel's constexpr GetSmemSize)
__shared__ char smem[GetSmemSize()];
// Cast pointers for optional features
MeanVarDataType* p_running_mean = static_cast<MeanVarDataType*>(kargs.p_running_mean);
MeanVarDataType* p_running_var = static_cast<MeanVarDataType*>(kargs.p_running_var);
MeanVarDataType* p_save_mean = static_cast<MeanVarDataType*>(kargs.p_save_mean);
MeanVarDataType* p_save_inv_std = static_cast<MeanVarDataType*>(kargs.p_save_inv_std);
// Call pipeline with properly distributed tile windows
Pipeline{}(x_window,
gamma_window,
beta_window,
y_window,
p_running_mean,
p_running_var,
p_save_mean,
p_save_inv_std,
static_cast<ComputeDataType>(kargs.epsilon),
static_cast<ComputeDataType>(kargs.momentum),
per_channel_size,
c,
smem);
}
// Validate arguments

177
include/ck_tile/ops/batchnorm/pipeline/batchnorm_fwd_pipeline.hpp Normal file
View File

@@ -0,0 +1,177 @@
// SPDX-License-Identifier: MIT
// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.
#pragma once
#include "ck_tile/core.hpp"
#include "ck_tile/ops/batchnorm/block/block_welford.hpp"
#include "ck_tile/ops/batchnorm/pipeline/batchnorm_fwd_policy.hpp"
namespace ck_tile {
// BatchnormFwdPipeline: Computation logic for batch normalization
// Takes tile windows and performs Welford reduction + normalization
template <typename Problem_, typename Policy_ = BatchnormFwdPipelineDefaultPolicy>
struct BatchnormFwdPipeline
{
using Problem = remove_cvref_t<Problem_>;
using Policy = remove_cvref_t<Policy_>;
using XDataType = typename Problem::XDataType;
using GammaDataType = typename Problem::GammaDataType;
using BetaDataType = typename Problem::BetaDataType;
using ComputeDataType = typename Problem::ComputeDataType;
using YDataType = typename Problem::YDataType;
using MeanVarDataType = typename Problem::MeanVarDataType;
static constexpr index_t kBlockSize = Problem::BlockShape::BlockSize;
CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSize()
{
return Policy::template GetSmemSize<Problem>();
}
template <typename XWindow,
typename GammaWindow,
typename BetaWindow,
typename YWindow>
CK_TILE_DEVICE void operator()(const XWindow& x_window_,
const GammaWindow& gamma_window_,
const BetaWindow& beta_window_,
YWindow& y_window_,
MeanVarDataType* p_running_mean,
MeanVarDataType* p_running_var,
MeanVarDataType* p_save_mean,
MeanVarDataType* p_save_inv_std,
ComputeDataType epsilon,
ComputeDataType momentum,
[[maybe_unused]] index_t per_channel_size,
index_t channel_idx,
void* smem) const
{
const index_t thread_id = get_thread_id();
// Apply tile distributions (like layernorm2d does)
// Note: x_window and y_window are NOT const (need to move them)
auto x_window =
make_tile_window(x_window_, Policy::template MakeXBlockTileDistribution<Problem>());
const auto gamma_window = make_tile_window(
gamma_window_, Policy::template MakeGammaBetaBlockTileDistribution<Problem>());
const auto beta_window = make_tile_window(
beta_window_, Policy::template MakeGammaBetaBlockTileDistribution<Problem>());
auto y_window =
make_tile_window(y_window_, Policy::template MakeXBlockTileDistribution<Problem>());
// Load gamma/beta once (constant per channel)
[[maybe_unused]]const auto gamma = load_tile(gamma_window);
[[maybe_unused]]const auto beta = load_tile(beta_window);
// Calculate how many tiles needed (like layernorm2d two-pass)
constexpr index_t Block_N = Problem::BlockShape::Block_N;
index_t num_tile_iteration = integer_divide_ceil(per_channel_size, Block_N);
// ==========================================
// PHASE 1: WELFORD REDUCTION OVER ALL TILES
// ==========================================
ComputeDataType thread_mean = type_convert<ComputeDataType>(0);
ComputeDataType thread_m2 = type_convert<ComputeDataType>(0);
index_t thread_count = 0;
// Iterate over tiles for Welford accumulation
for(index_t tile_idx = 0; tile_idx < num_tile_iteration; ++tile_idx)
{
auto x = load_tile(x_window);
sweep_tile(x, [&](auto idx) {
ComputeDataType val = type_convert<ComputeDataType>(x[idx]);
thread_count++;
ComputeDataType delta = val - thread_mean;
thread_mean += delta / type_convert<ComputeDataType>(thread_count);
ComputeDataType delta2 = val - thread_mean;
thread_m2 += delta * delta2;
});
// Move to next tile
if(tile_idx < num_tile_iteration - 1)
{
move_tile_window(x_window, {0, Block_N});
}
}
// Move x_window back to start
move_tile_window(x_window, {0, -static_cast<int>(Block_N * (num_tile_iteration - 1))});
// Block-level reduction
ComputeDataType block_mean = thread_mean;
ComputeDataType block_var = thread_m2;
index_t block_count = thread_count;
BlockWelford<ComputeDataType>::template Run<index_t, kBlockSize>(
block_mean, block_var, block_count, smem);
// ==========================================
// PHASE 2: COMPUTE INVERSE STD
// ==========================================
ComputeDataType inv_std = type_convert<ComputeDataType>(1) /
ck_tile::sqrt(block_var + epsilon);
// ==========================================
// PHASE 3: NORMALIZE AND STORE (ITERATE OVER TILES)
// ==========================================
for(index_t tile_idx = 0; tile_idx < num_tile_iteration; ++tile_idx)
{
auto x = load_tile(x_window);
auto y = make_static_distributed_tensor<YDataType>(x.get_tile_distribution());
sweep_tile(y, [&](auto idx) {
ComputeDataType x_val = type_convert<ComputeDataType>(x[idx]);
// y = (x - mean) / std (no gamma/beta for now)
ComputeDataType normalized = (x_val - block_mean) * inv_std;
y(idx) = type_convert<YDataType>(normalized);
});
store_tile(y_window, y);
// Move to next tile
if(tile_idx < num_tile_iteration - 1)
{
move_tile_window(x_window, {0, Block_N});
move_tile_window(y_window, {0, Block_N});
}
}
// ==========================================
// PHASE 6: SAVE STATISTICS (Optional)
// ==========================================
if constexpr(Problem::Traits::kSaveMeanInvStd)
{
if(thread_id == 0)
{
p_save_mean[channel_idx] = type_convert<MeanVarDataType>(block_mean);
p_save_inv_std[channel_idx] = type_convert<MeanVarDataType>(inv_std);
}
}
// ==========================================
// PHASE 7: UPDATE RUNNING STATISTICS (Optional)
// ==========================================
if constexpr(Problem::Traits::kUpdateMovingAverage)
{
if(thread_id == 0)
{
ComputeDataType one_minus_momentum = type_convert<ComputeDataType>(1) - momentum;
ComputeDataType old_mean = type_convert<ComputeDataType>(p_running_mean[channel_idx]);
ComputeDataType old_var = type_convert<ComputeDataType>(p_running_var[channel_idx]);
p_running_mean[channel_idx] = type_convert<MeanVarDataType>(
one_minus_momentum * old_mean + momentum * block_mean);
p_running_var[channel_idx] = type_convert<MeanVarDataType>(
one_minus_momentum * old_var + momentum * block_var);
}
}
}
};
} // namespace ck_tile

59
include/ck_tile/ops/batchnorm/pipeline/batchnorm_fwd_policy.hpp Normal file
View File

@@ -0,0 +1,59 @@
// SPDX-License-Identifier: MIT
// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.
#pragma once
#include "ck_tile/core.hpp"
namespace ck_tile {
// Default policy for batchnorm forward pipeline
// Defines tile distributions and helper functions
struct BatchnormFwdPipelineDefaultPolicy
{
// Tile distribution for input data (following layernorm2d pattern exactly)
template <typename Problem>
CK_TILE_DEVICE static constexpr auto MakeXBlockTileDistribution()
{
using S = typename Problem::BlockShape;
return make_static_tile_distribution(
tile_distribution_encoding<
sequence<>,
tuple<sequence<S::Repeat_M, S::WarpPerBlock_M, S::ThreadPerWarp_M, S::Vector_M>,
sequence<S::Repeat_N, S::WarpPerBlock_N, S::ThreadPerWarp_N, S::Vector_N>>,
tuple<sequence<1, 2>, sequence<1, 2>>,
tuple<sequence<1, 1>, sequence<2, 2>>,
sequence<1, 1, 2, 2>,
sequence<0, 3, 0, 3>>{});
}
// Tile distribution for gamma/beta parameters (following layernorm2d pattern)
template <typename Problem>
CK_TILE_DEVICE static constexpr auto MakeGammaBetaBlockTileDistribution()
{
using S = typename Problem::BlockShape;
return make_static_tile_distribution(
tile_distribution_encoding<
sequence<S::WarpPerBlock_M, S::ThreadPerWarp_M>,
tuple<sequence<S::Repeat_N, S::WarpPerBlock_N, S::ThreadPerWarp_N, S::Vector_N>>,
tuple<sequence<0, 1>, sequence<0, 1>>,
tuple<sequence<0, 1>, sequence<1, 2>>,
sequence<1, 1>,
sequence<0, 3>>{});
}
// Calculate shared memory size
template <typename Problem>
CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSize()
{
// For POC, use BlockWelford's smem requirement
using ComputeDataType = typename Problem::ComputeDataType;
constexpr index_t kBlockSize = Problem::BlockShape::BlockSize;
return BlockWelford<ComputeDataType>::template GetSmemSize<index_t, kBlockSize>();
}
};
} // namespace ck_tile

44
include/ck_tile/ops/batchnorm/pipeline/batchnorm_shape.hpp
View File

@@ -4,46 +4,14 @@
#pragma once
#include "ck_tile/core.hpp"
#include "ck_tile/ops/common/generic_2d_block_shape.hpp"
namespace ck_tile {
// BatchnormShape defines the block/warp/thread tile configuration
// Similar to reduce2d_shape but adapted for batchnorm
template <typename BlockWarps_,
typename BlockTile_,
typename WarpTile_,
typename Vector_>
struct BatchnormShape
{
using BlockWarps = remove_cvref_t<BlockWarps_>;
using BlockTile = remove_cvref_t<BlockTile_>;
using WarpTile = remove_cvref_t<WarpTile_>;
using Vector = remove_cvref_t<Vector_>;
static constexpr index_t kBlockWarps_M = BlockWarps::at(number<0>{});
static constexpr index_t kBlockWarps_N = BlockWarps::at(number<1>{});
static constexpr index_t kBlockSize = kBlockWarps_M * kBlockWarps_N * get_warp_size();
static constexpr index_t Block_M = BlockTile::at(number<0>{});
static constexpr index_t Block_N = BlockTile::at(number<1>{});
static constexpr index_t WarpTile_M = WarpTile::at(number<0>{});
static constexpr index_t WarpTile_N = WarpTile::at(number<1>{});
static constexpr index_t VectorSize_M = Vector::at(number<0>{});
static constexpr index_t VectorSize_N = Vector::at(number<1>{});
// Thread tile sizes
static constexpr index_t ThreadTile_M = WarpTile_M / kBlockWarps_M;
static constexpr index_t ThreadTile_N = WarpTile_N / kBlockWarps_N;
// For batchnorm:
// - M dimension represents the batch*channel (merged N*C)
// - N dimension represents the spatial (merged H*W)
// We reduce over N (spatial) dimension per M (batch*channel)
static constexpr auto BlockSize = number<kBlockSize>{};
};
// BatchnormShape using Generic2dBlockShape for proper tile distribution support
template <typename BlockTile_, // Block size, sequence<M, N>
typename ThreadPerBlock_, // Threads along sequence<M, N>
typename Vector_> // Vector size along sequence<M, N>
using BatchnormShape = Generic2dBlockShape<BlockTile_, ThreadPerBlock_, Vector_>;
} // namespace ck_tile