ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-07 15:47:05 +00:00
Code Issues Packages Projects Releases Wiki Activity

remove debugging statements

Browse Source
This commit is contained in:
khuagarw
2025-11-11 22:33:06 +00:00
parent 869bc5b77b
commit 075c36b5f9
10 changed files with 78 additions and 415 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
example/ck_tile/38_block_scale_gemm/README.md
View File

@@ -64,6 +64,7 @@ args:
-rotating_count Rotating count (default:1000)
-quant_mode Choose aquant, bquant, tensor or rowcol (default:bquant)
-preshuffleb Enable preshuffle of tensor B (default:false)
-preshufflequant Enable preshuffle of quant tensor (defualt:false)
-group_size Quantization group size as MxNxK, e.g., 1x1x128, 1x32x128, 1x64x128 (default:1x1x128)
```

85
example/ck_tile/38_block_scale_gemm/run_gemm_quant_example.inc
View File

@@ -438,14 +438,6 @@ int run_gemm_example_with_layouts(const ck_tile::ArgParser& arg_parser,
ck_tile::HostTensor<CDataType> c_m_n_dev_result(
ck_tile::host_tensor_descriptor(M, N, stride_C, is_row_major(CLayout{})));
// printf("M = %d, N = %d, K = %d, AQK = %d, BQK = %d, stride_AQ = %d, stride_BQ = %d\n",
// M,
// N,
// K,
// AQK,
// BQK,
// stride_AQ,
// stride_BQ);
// Create AQ tensor with appropriate shape
std::unique_ptr<ck_tile::HostTensor<AQDataType>> aq_tensor_ptr = nullptr;
if constexpr(QuantMode == ck_tile::QuantType::AQuantGrouped ||
@@ -531,52 +523,11 @@ int run_gemm_example_with_layouts(const ck_tile::ArgParser& arg_parser,
ck_tile::FillConstant<ADataType>{static_cast<ADataType>(0x38)}(a_m_k);
ck_tile::FillConstant<BDataType>{static_cast<BDataType>(0x22)}(b_k_n);
ck_tile::FillConstant<BQDataType>{static_cast<BQDataType>(0.5f)}(*bq_tensor_ptr);
// if(bq_tensor_ptr)
// {
// BQDataType value = 0;
// for(int i = 0; i < BQK; i++)
// {
// for(int j = 0; j < N; j++)
// {
// (*bq_tensor_ptr)(i, j) = value;
// value += static_cast<BQDataType>(0.1f);
// }
// }
// }
// for(int i = 0; i < BQK; i++)
// {
// for(int j = 0; j < N; j++)
// {
// printf("%.2f ", (*bq_tensor_ptr)(i, j));
// }
// printf("\n");
// }
}
else
{
ck_tile::FillConstant<ADataType>{static_cast<ADataType>(0x22)}(a_m_k);
ck_tile::FillConstant<AQDataType>{static_cast<AQDataType>(0.5f)}(*aq_tensor_ptr);
// if(aq_tensor_ptr)
// {
// AQDataType value = 0;
// for(int i = 0; i < M; i++)
// {
// for(int j = 0; j < AQK; j++)
// {
// (*aq_tensor_ptr)(i, j) = value;
// value += static_cast<AQDataType>(0.1f);
// }
// }
// }
// for(int i = 0; i < M; i++)
// {
// for(int j = 0; j < AQK; j++)
// {
// printf("%.2f ", (*aq_tensor_ptr)(i, j));
// }
// printf("\n");
// }
ck_tile::FillConstant<BDataType>{static_cast<BDataType>(0x38)}(b_k_n);
if constexpr(QuantMode == ck_tile::QuantType::RowColQuant)
@@ -622,24 +573,6 @@ int run_gemm_example_with_layouts(const ck_tile::ArgParser& arg_parser,
{
ck_tile::HostTensor<AQDataType> aq_shuffle_host =
ck_tile::shuffle_aq(aq_tensor_ptr.get(), GemmConfig::K_Tile / QuantGroupSize::kK);
// printf("aq_shuffle_host.get_length(0): %zu, aq_shuffle_host.get_length(1): %zu, "
// "aq_shuffle_host.get_length(2): %zu\n",
// aq_shuffle_host.get_length(0),
// aq_shuffle_host.get_length(1),
// aq_shuffle_host.get_length(2));
// printf("Preshuffle AQ \n");
// for(size_t i = 0; i < aq_shuffle_host.get_length(0); ++i)
// {
// for(size_t j = 0; j < aq_shuffle_host.get_length(1); ++j)
// {
// for(size_t k = 0; k < aq_shuffle_host.get_length(2); ++k)
// {
// printf("%.2f ", aq_shuffle_host(i, j, k));
// }
// printf("\n");
// }
// printf("\n");
// }
aq_dev_buf_ptr->ToDevice(aq_shuffle_host.data());
}
else
@@ -699,24 +632,6 @@ int run_gemm_example_with_layouts(const ck_tile::ArgParser& arg_parser,
{
ck_tile::HostTensor<BQDataType> bq_shuffle_host =
ck_tile::shuffle_bq(bq_tensor_ptr.get(), GemmConfig::K_Tile / QuantGroupSize::kK);
// printf("bq_shuffle_host.get_length(0): %zu, bq_shuffle_host.get_length(1): %zu, "
// "bq_shuffle_host.get_length(2): %zu\n",
// bq_shuffle_host.get_length(0),
// bq_shuffle_host.get_length(1),
// bq_shuffle_host.get_length(2));
// printf("Preshuffle BQ \n");
// for(size_t i = 0; i < bq_shuffle_host.get_length(0); ++i)
// {
// for(size_t j = 0; j < bq_shuffle_host.get_length(1); ++j)
// {
// for(size_t k = 0; k < bq_shuffle_host.get_length(2); ++k)
// {
// printf("%.2f ", bq_shuffle_host(i, j, k));
// }
// printf("\n");
// }
// printf("\n");
// }
bq_dev_buf_ptr->ToDevice(bq_shuffle_host.data());
}
else

4
include/ck_tile/host/tensor_shuffle_utils.hpp
View File

@@ -11,7 +11,7 @@ auto shuffle_aq(const ck_tile::HostTensor<T>* t, int block_aq_k)
}
int m_ = t->get_lengths()[0];
int aqk_ = t->get_lengths()[1];
// printf("m_: %d, aqk_: %d, block_aq_k: %d\n", m_, aqk_, block_aq_k);
if(aqk_ % block_aq_k != 0)
{
throw std::runtime_error("shuffle_aq needs a aqk of multiple times of block_aq_k.");
@@ -30,7 +30,7 @@ auto shuffle_bq(const ck_tile::HostTensor<T>* t, int block_bq_k)
}
int bqk_ = t->get_lengths()[0];
int n_ = t->get_lengths()[1];
// printf("bqk_: %d, n_: %d, block_bq_k: %d\n", bqk_, n_, block_bq_k);
if(bqk_ % block_bq_k != 0)
{
throw std::runtime_error("shuffle_bq needs a bqk of multiple times of block_bq_k.");

25
include/ck_tile/ops/gemm_quant/block/block_universal_gemm_ar_flatbr_bquant_cr.hpp
View File

@@ -207,18 +207,7 @@ struct BlockGemmWeightPreshuffleBQuantARegBRegCReg
pull_from_lane << 2, __builtin_bit_cast(int, scale_reg_dword));
float scale_reg_f = cvt_scale_to_fp32(gathered_scale_reg);
// if(get_block_id() == 0 && get_warp_id() == 0 && get_thread_id() == 0)
// {
// printf("scale_reg_f: %f, reg_offset: %d, MIterPerWarp: %d, "
// "NIterPerWarp: %d, mIter: %d, nIter:%d, kQScale: %d\n",
// scale_reg_f,
// reg_offset,
// MIterPerWarp,
// NIterPerWarp,
// static_cast<int>(mIter),
// static_cast<int>(nIter),
// static_cast<int>(kQScale));
// }
static_for<0, WG::kM * WG::kN / warp_size, 1>{}([&](auto c_row) {
auto& c_ref = c_block_tensor.get_thread_buffer()[tbuf_offset + c_row];
const auto acc_val = c_acc(mIter)(nIter).get_thread_buffer()[c_row];
@@ -230,18 +219,6 @@ struct BlockGemmWeightPreshuffleBQuantARegBRegCReg
constexpr index_t reg_offset = nIter * KPerBlockBQ + kQScale;
auto& scale_reg = bq_block_tensor.get_thread_buffer()[reg_offset];
float scale_reg_f = cvt_scale_to_fp32(scale_reg);
// if(get_block_id() == 0 && get_thread_id() == 0)
// {
// printf("scale_reg_f: %f, reg_offset: %d, MIterPerWarp: %d, "
// "NIterPerWarp: %d, mIter: %d, nIter:%d, kQScale: %d\n",
// scale_reg_f,
// reg_offset,
// MIterPerWarp,
// NIterPerWarp,
// static_cast<int>(mIter),
// static_cast<int>(nIter),
// static_cast<int>(kQScale));
// }
static_for<0, WG::kM * WG::kN / warp_size, 1>{}([&](auto c_row) {
auto& c_ref = c_block_tensor.get_thread_buffer()[tbuf_offset + c_row];

41
include/ck_tile/ops/gemm_quant/block/block_universal_gemm_as_aquant_bs_cr.hpp
View File

@@ -268,10 +268,6 @@ struct AQuantBlockUniversalGemmAsBsCr : public BlockGemmAQuantBase<Problem_>
if constexpr(std::is_same_v<AQDataType, float>)
{
// if(get_block_id() == 0 && get_thread_id() == 0)
// {
// printf("Here exchange_quant_value_across_lanes float\n");
// }
scale_reg_dword = ck_tile::bit_cast<uint32_t>(scale_reg);
}
else
@@ -281,13 +277,6 @@ struct AQuantBlockUniversalGemmAsBsCr : public BlockGemmAQuantBase<Problem_>
int gathered_scale_reg = __builtin_amdgcn_ds_bpermute(
pull_from_lane << 2, __builtin_bit_cast(int, scale_reg_dword));
// if(get_block_id() == 0 && get_thread_id() == 0)
// {
// printf("pull_from_lane: %d, scale_reg_dword: 0x%08x, gathered_scale_reg: %d\n",
// pull_from_lane,
// scale_reg_dword,
// gathered_scale_reg);
// }
return Base::cvt_scale_to_fp32(gathered_scale_reg);
}
@@ -362,17 +351,6 @@ struct AQuantBlockUniversalGemmAsBsCr : public BlockGemmAQuantBase<Problem_>
decltype(threadIdx.x) pull_from_lane = 0;
if constexpr(WarpGemm::kM == 16)
{
// if(get_block_id() == 0 && get_thread_id() == 0)
// {
// printf("\nHere WarpGemm::kM == 16\n");
// printf("lane_id(): %u, Traits::WarpGemm::kN: %d, c_row: %u, "
// "Traits::QScalesPerBlockRow: %d, kQScale: %d\n",
// __lane_id(),
// Traits::WarpGemm::kN,
// c_row,
// Traits::QScalesPerBlockRow,
// kQScale);
// }
pull_from_lane =
(__lane_id() / Traits::WarpGemm::kN * kTileRowsOfCPerThread +
c_row) *
@@ -392,14 +370,6 @@ struct AQuantBlockUniversalGemmAsBsCr : public BlockGemmAQuantBase<Problem_>
static_assert(false, "WarpGemm::kM is not 16 nor 32.");
}
auto& scale_reg = aq_block_tensor.get_thread_buffer()[mIter];
// if(get_block_id() == 0 && get_thread_id() == 0)
// {
// printf("mIter: %d, kQScale: %d, pull_from_lane: %u, scale_reg: %f\n",
// mIter,
// kQScale,
// pull_from_lane,
// scale_reg);
// }
return exchange_quant_value_across_lanes(scale_reg, pull_from_lane);
}
else
@@ -559,16 +529,7 @@ struct AQuantBlockUniversalGemmAsBsCr : public BlockGemmAQuantBase<Problem_>
static_for<0, WarpGemm::kM * WarpGemm::kN / warp_size, 1>{}(
[&](auto c_row) {
float scale_reg_f = aq_picker.template pick<c_row>();
// if(get_block_id() == 0 && get_thread_id() == 0)
// {
// printf("mIter: %d, nIter: %d, kQScale: %d, c_row: %d, "
// "scale_reg_f: %f\n",
// static_cast<int>(mIter),
// static_cast<int>(nIter),
// static_cast<int>(kQScale),
// static_cast<int>(c_row),
// scale_reg_f);
// }
c_block_tensor.get_thread_buffer()[tbuf_offset + c_row] +=
(c_warp_tensor.get_thread_buffer()[c_row] * scale_reg_f);
});

13
include/ck_tile/ops/gemm_quant/block/block_universal_gemm_as_bs_bquant_cr.hpp
View File

@@ -374,18 +374,7 @@ struct BQuantBlockUniversalGemmAsBsCr : public BlockGemmBQuantBase<Problem_>
pull_from_lane << 2, __builtin_bit_cast(int, scale_reg_dword));
float scale_reg_f = Base::cvt_scale_to_fp32(gathered_scale_reg);
// if(get_block_id() == 0 && get_warp_id() == 0 && get_thread_id() == 0)
// {
// printf("scale_reg_f: %f, reg_offset: %d, MIterPerWarp: %d, "
// "NIterPerWarp: %d, mIter: %d, nIter:%d, kQScale: %d\n",
// scale_reg_f,
// reg_offset,
// MIterPerWarp,
// NIterPerWarp,
// static_cast<int>(mIter),
// static_cast<int>(nIter),
// static_cast<int>(kQScale));
// }
static_for<0, WarpGemm::kM * WarpGemm::kN / warp_size, 1>{}(
[&](auto c_row) {
c_block_tensor.get_thread_buffer()[tbuf_offset + c_row] +=

217
include/ck_tile/ops/gemm_quant/kernel/gemm_quant_kernel.hpp
View File

@@ -514,18 +514,8 @@ struct QuantGemmKernel
if constexpr(kQuantType == QuantType::AQuantGrouped && PreshuffleQuant)
{
static_assert(std::is_same_v<AQLayout, tensor_layout::gemm::RowMajor>);
const auto aq_x = kargs.M * GemmPipeline::KPerBlockAQ; // 16*2 =32
const auto aq_y = kargs.QK_A / GemmPipeline::KPerBlockAQ; // 4/2 =2
// if(get_block_id() == 0 && get_thread_id() == 0)
// {
// printf("For aq_desc: \n");
// printf("aq_x: %d, aq_y: %d, GemmPipeline::KPerBlockAQ: %d, "
// "GemmPipeline::GetVectorSizeAQ(): %d\n\n", // 32, 2, 2, 2
// aq_x,
// aq_y,
// GemmPipeline::KPerBlockAQ,
// GemmPipeline::GetVectorSizeAQ());
// }
const auto aq_x = kargs.M * GemmPipeline::KPerBlockAQ;
const auto aq_y = kargs.QK_A / GemmPipeline::KPerBlockAQ;
const auto aq_desc =
make_naive_tensor_descriptor(make_tuple(aq_y, aq_x),
make_tuple(aq_x, 1),
@@ -533,19 +523,7 @@ struct QuantGemmKernel
number<1>{});
const auto block_tile_size = GemmPipeline::MPerBlock * GemmPipeline::KPerBlockAQ;
// if(get_block_id() == 0 && get_thread_id() == 0)
// {
// printf("For aq_pad0_desc: \n");
// printf("GemmPipeline::MPerBlock: %d, GemmPipeline::KPerBlockAQ: %d\n", // 16,
// 2
// GemmPipeline::MPerBlock,
// GemmPipeline::KPerBlockAQ);
// printf("ck_tile::integer_least_multiple(length, alignment) : %d\n",
// ck_tile::integer_least_multiple(aq_x, block_tile_size));
// printf("get_padding_size(aq_x, block_tile_size): %d\n\n",
// get_padding_size(aq_x, block_tile_size));
// }
const auto aq_pad0_desc = transform_tensor_descriptor(
const auto aq_pad0_desc = transform_tensor_descriptor(
aq_desc,
make_tuple(
make_pass_through_transform(aq_y),
@@ -553,22 +531,12 @@ struct QuantGemmKernel
make_tuple(sequence<0>{}, sequence<1>{}),
make_tuple(sequence<0>{}, sequence<1>{}));
const auto pad_aq_x =
aq_pad0_desc.get_lengths()[I1]; // 32 (as no padding needed here)
const auto wave_tile_size = // 16*2 = 32
const auto pad_aq_x = aq_pad0_desc.get_lengths()[I1];
const auto wave_tile_size =
TilePartitioner::BlockGemmShape::WarpTile::at(I0) * GemmPipeline::KPerBlockAQ;
const auto wave_tile_count_x = // 32/32 =1
const auto wave_tile_count_x =
ck_tile::integer_divide_ceil(pad_aq_x, wave_tile_size);
// if(get_block_id() == 0 && get_thread_id() == 0)
// {
// printf("For aq_unmerge_pad0_desc: \n");
// printf("pad_aq_x: %d\n", pad_aq_x);
// printf("wave_tile_size: %d, GemmPipeline::KPerBlockAQ: %d\n",
// wave_tile_size,
// GemmPipeline::KPerBlockAQ);
// printf("wave_tile_count_x: %d\n\n", wave_tile_count_x);
// }
const auto aq_unmerge_pad0_desc = transform_tensor_descriptor(
aq_pad0_desc,
make_tuple(
@@ -577,15 +545,6 @@ struct QuantGemmKernel
make_tuple(sequence<0>{}, sequence<1>{}),
make_tuple(sequence<0>{}, sequence<1, 2>{}));
// if(get_block_id() == 0 && get_thread_id() == 0)
// {
// printf("For aq_pad1_desc: \n");
// printf("aq_y: %d\n", aq_y);
// printf("get_warp_size(): %d\n", get_warp_size());
// printf("get_padding_size(wave_tile_size, get_warp_size()): %d\n\n",
// get_padding_size(wave_tile_size, get_warp_size()));
// }
const auto aq_pad1_desc = transform_tensor_descriptor(
aq_unmerge_pad0_desc,
make_tuple(
@@ -598,15 +557,10 @@ struct QuantGemmKernel
const auto pad_wave_size =
ck_tile::integer_least_multiple(wave_tile_size, get_warp_size());
// if(get_block_id() == 0 && get_thread_id() == 0)
// {
// printf("For aq_merge_pad1_desc: \n");
// printf("pad_wave_size: %d\n\n", pad_wave_size);
// }
const auto aq_merge_pad1_desc = transform_tensor_descriptor(
aq_pad1_desc,
make_tuple(make_merge_transform(make_tuple(aq_y, wave_tile_count_x)), //(2,1)
make_pass_through_transform(pad_wave_size)), //(64)
make_tuple(make_merge_transform(make_tuple(aq_y, wave_tile_count_x)),
make_pass_through_transform(pad_wave_size)),
make_tuple(sequence<0, 1>{}, sequence<2>{}),
make_tuple(sequence<0>{}, sequence<1>{}));
@@ -733,18 +687,8 @@ struct QuantGemmKernel
if constexpr(PreshuffleQuant)
{
static_assert(std::is_same_v<BQLayout, tensor_layout::gemm::ColumnMajor>);
const auto bq_x = kargs.N * GemmPipeline::KPerBlockBQ; // 64*2 =128
const auto bq_y = kargs.QK_B / GemmPipeline::KPerBlockBQ; // 4/2 =2
// if(get_block_id() == 0 && get_thread_id() == 0)
// {
// printf("For bq_desc: \n");
// printf("bq_x: %d, bq_y: %d, GemmPipeline::KPerBlockBQ: %d, "
// "GemmPipeline::GetVectorSizeBQ(): %d\n\n", // 32, 2, 2, 2
// bq_x,
// bq_y,
// GemmPipeline::KPerBlockBQ,
// GemmPipeline::GetVectorSizeBQ());
// }
const auto bq_x = kargs.N * GemmPipeline::KPerBlockBQ;
const auto bq_y = kargs.QK_B / GemmPipeline::KPerBlockBQ;
const auto bq_desc =
make_naive_tensor_descriptor(make_tuple(bq_y, bq_x),
make_tuple(bq_x, 1),
@@ -752,19 +696,7 @@ struct QuantGemmKernel
number<1>{});
const auto block_tile_size =
GemmPipeline::NPerBlock * GemmPipeline::KPerBlockBQ; // 64*2=128
// if(get_block_id() == 0 && get_thread_id() == 0)
// {
// printf("For bq_pad0_desc: \n");
// printf(
// "GemmPipeline::NPerBlock: %d, GemmPipeline::KPerBlockBQ: %d\n",
// GemmPipeline::NPerBlock,
// GemmPipeline::KPerBlockBQ);
// printf("ck_tile::integer_least_multiple(length, alignment) : %d\n",
// ck_tile::integer_least_multiple(bq_x, block_tile_size));
// printf("get_padding_size(bq_x, block_tile_size): %d\n\n",
// get_padding_size(bq_x, block_tile_size));
// }
GemmPipeline::NPerBlock * GemmPipeline::KPerBlockBQ;
const auto bq_pad0_desc = transform_tensor_descriptor(
bq_desc,
make_tuple(make_pass_through_transform(bq_y),
@@ -773,59 +705,34 @@ struct QuantGemmKernel
make_tuple(sequence<0>{}, sequence<1>{}),
make_tuple(sequence<0>{}, sequence<1>{}));
const auto pad_bq_x =
bq_pad0_desc.get_lengths()[I1]; // 128 (as no padding needed here)
const auto wave_tile_size = // 16 * 2 = 32
TilePartitioner::BlockGemmShape::WarpTile::at(I1) *
GemmPipeline::KPerBlockBQ;
const auto wave_tile_count_x = // 128/ 32 = 4
const auto pad_bq_x = bq_pad0_desc.get_lengths()[I1];
const auto wave_tile_size = TilePartitioner::BlockGemmShape::WarpTile::at(I1) *
GemmPipeline::KPerBlockBQ;
const auto wave_tile_count_x =
ck_tile::integer_divide_ceil(pad_bq_x, wave_tile_size);
// if(get_block_id() == 0 && get_thread_id() == 0)
// {
// printf("For bq_unmerge_pad0_desc: \n");
// printf("pad_bq_x: %d\n", pad_bq_x);
// printf("wave_tile_size: %d, GemmPipeline::KPerBlockBQ: %d\n",
// wave_tile_size,
// GemmPipeline::KPerBlockBQ);
// printf("wave_tile_count_x: %d\n\n", wave_tile_count_x);
// }
const auto bq_unmerge_pad0_desc = transform_tensor_descriptor(
bq_pad0_desc,
make_tuple(make_pass_through_transform(bq_y), // 2
make_unmerge_transform(
make_tuple(wave_tile_count_x, wave_tile_size))), //(4, 32)
make_tuple(
make_pass_through_transform(bq_y),
make_unmerge_transform(make_tuple(wave_tile_count_x, wave_tile_size))),
make_tuple(sequence<0>{}, sequence<1>{}),
make_tuple(sequence<0>{}, sequence<1, 2>{}));
// if(get_block_id() == 0 && get_thread_id() == 0)
// {
// printf("For bq_pad1_desc: \n");
// printf("bq_y: %d\n", bq_y); // 2
// printf("get_warp_size(): %d\n", get_warp_size()); // 64
// printf("get_padding_size(wave_tile_size, get_warp_size()): %d\n\n",
// get_padding_size(wave_tile_size, get_warp_size())); // 32
// }
const auto bq_pad1_desc = transform_tensor_descriptor(
bq_unmerge_pad0_desc,
make_tuple(make_pass_through_transform(bq_y), // 2
make_pass_through_transform(wave_tile_count_x), // 4
make_right_pad_transform(
wave_tile_size,
get_padding_size(wave_tile_size, get_warp_size()))), // 64
make_tuple(
make_pass_through_transform(bq_y),
make_pass_through_transform(wave_tile_count_x),
make_right_pad_transform(
wave_tile_size, get_padding_size(wave_tile_size, get_warp_size()))),
make_tuple(sequence<0>{}, sequence<1>{}, sequence<2>{}),
make_tuple(sequence<0>{}, sequence<1>{}, sequence<2>{}));
const auto pad_wave_size =
ck_tile::integer_least_multiple(wave_tile_size, get_warp_size()); // 64
// if(get_block_id() == 0 && get_thread_id() == 0)
// {
// printf("For bq_merge_pad1_desc: \n");
// printf("pad_wave_size: %d\n\n", pad_wave_size);
// }
ck_tile::integer_least_multiple(wave_tile_size, get_warp_size());
const auto bq_merge_pad1_desc = transform_tensor_descriptor(
bq_pad1_desc,
make_tuple(
make_merge_transform(make_tuple(bq_y, wave_tile_count_x)), //(2,4)
make_pass_through_transform(pad_wave_size)), // 64
make_tuple(make_merge_transform(make_tuple(bq_y, wave_tile_count_x)),
make_pass_through_transform(pad_wave_size)),
make_tuple(sequence<0, 1>{}, sequence<2>{}),
make_tuple(sequence<0>{}, sequence<1>{}));
@@ -990,21 +897,9 @@ struct QuantGemmKernel
ck_tile::integer_least_multiple(warp_m * aqk_per_block, get_warp_size());
constexpr auto tile_window_height = block_m / warp_m;
auto block_m_idx = i_m / block_m;
// if(get_block_id() == 0 && get_thread_id() == 0)
// {
// printf("For aq_block_window: \n");
// printf("block_m: %d, warp_m: %d, aqk_per_block: %d\n",
// block_m,
// static_cast<int>(warp_m),
// aqk_per_block);
// printf("tile_window_width: %d, tile_window_height: %d\n",
// tile_window_width,
// tile_window_height);
// printf("i_m: %d, block_m_idx: %d\n\n", i_m, block_m_idx);
// }
return make_tile_window(
aq_pad_view,
make_tuple(number<tile_window_height>{}, number<tile_window_width>{}), // 1, 64
make_tuple(number<tile_window_height>{}, number<tile_window_width>{}),
{block_m_idx * tile_window_height, 0});
}
else if constexpr(kQuantType == QuantType::AQuantGrouped && !PreshuffleQuant)
@@ -1074,45 +969,23 @@ struct QuantGemmKernel
{
static_assert(std::is_same_v<BQLayout, tensor_layout::gemm::ColumnMajor>);
using QuantGroupSize = remove_cvref_t<typename GemmPipeline::QuantGroupSize>;
constexpr auto block_n = TilePartitioner::NPerBlock / QuantGroupSize::kN; // 64
constexpr auto warp_n = TilePartitioner::BlockGemmShape::WarpTile::at(I1); // 16
constexpr auto bqk_per_block =
TilePartitioner::KPerBlock / QuantGroupSize::kK; // 256/128=2
constexpr auto tile_window_width = ck_tile::integer_least_multiple(
warp_n * bqk_per_block, get_warp_size()); //(32, 64) = 64
constexpr auto block_n = TilePartitioner::NPerBlock / QuantGroupSize::kN;
constexpr auto warp_n = TilePartitioner::BlockGemmShape::WarpTile::at(I1);
constexpr auto bqk_per_block = TilePartitioner::KPerBlock / QuantGroupSize::kK;
constexpr auto tile_window_width =
ck_tile::integer_least_multiple(warp_n * bqk_per_block, get_warp_size());
constexpr auto tile_window_height = block_n / warp_n;
auto block_n_idx = i_n / block_n;
// if(get_block_id() == 0 && get_thread_id() == 0)
// {
// printf("For bq_block_window: \n");
// printf("block_n: %d, warp_n: %d, bqk_per_block: %d\n",
// block_n,
// static_cast<int>(warp_n),
// bqk_per_block);
// printf("tile_window_width: %d, tile_window_height: %d\n",
// tile_window_width,
// tile_window_height);
// printf("i_n: %d, block_n_idx: %d\n\n", i_n, block_n_idx);
// }
return make_tile_window(bq_pad_view,
make_tuple(number<tile_window_height>{},
number<tile_window_width>{}), // 4, 64
{block_n_idx * tile_window_height, 0});
return make_tile_window(
bq_pad_view,
make_tuple(number<tile_window_height>{}, number<tile_window_width>{}),
{block_n_idx * tile_window_height, 0});
}
else
{
static_assert(std::is_same_v<BQLayout, tensor_layout::gemm::ColumnMajor>);
using QuantGroupSize = remove_cvref_t<typename GemmPipeline::QuantGroupSize>;
// if(get_block_id() == 0 && get_thread_id() == 0)
// {
// printf("In bq_block_window without preshuffle\n");
// printf("TilePartitioner::NPerBlock: %d, TilePartitioner::KPerBlock %d, "
// "GemmPipeline::QuantGroupSize: %d\n",
// TilePartitioner::NPerBlock,
// TilePartitioner::KPerBlock,
// GemmPipeline::QuantGroupSize);
// }
return make_tile_window(
bq_pad_view,
make_tuple(number<TilePartitioner::KPerBlock / QuantGroupSize::kK>{},
@@ -1181,24 +1054,12 @@ struct QuantGemmKernel
if constexpr(kQuantType == QuantType::AQuantGrouped)
{
const auto& aq_block_window = gemm_tile_windows.at(I1);
// if(get_block_id() == 0 && get_thread_id() == 33)
// {
// printf("In RunGemm, before GemmPipeline call for AQuantGrouped\n");
// aq_block_window.template print_tile_window_range<AQDataType>(
// 0, 16, 0, 64, "aq block window");
// }
return GemmPipeline{}.template operator()(
a_block_window, b_block_window, aq_block_window, kargs.M, num_loop, smem_ptr_0);
}
else if constexpr(kQuantType == QuantType::BQuantGrouped)
{
const auto& bq_block_window = gemm_tile_windows.at(I3);
// if(get_block_id() == 0 && get_thread_id() == 0)
// {
// printf("In RunGemm, before GemmPipeline call for BQuantGrouped\n");
// bq_block_window.template print_tile_window_range<BQDataType>(
// 0, 8, 0, 64, "bq block window");
// }
return GemmPipeline{}.template operator()(
a_block_window, b_block_window, bq_block_window, kargs.N, num_loop, smem_ptr_0);
}
@@ -1288,12 +1149,6 @@ struct QuantGemmKernel
if constexpr(kQuantType == QuantType::BQuantGrouped)
{
const auto& bq_block_window = gemm_tile_windows.at(I3);
// if(get_block_id() == 0 && get_thread_id() == 0)
// {
// printf("In RunGemm2LDS, before GemmPipeline call for BQuantGrouped\n");
// bq_block_window.template print_tile_window_range<BQDataType>(
// 0, 16, 0, 64, /*0, 2, 0, 128,*/ "bq block window");
// }
return GemmPipeline{}.template operator()(a_block_window,
b_block_window,
bq_block_window,

32
include/ck_tile/ops/gemm_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_policy.hpp
View File

@@ -52,25 +52,15 @@ struct GemmAQuantPipelineAgBgCrDefaultPolicy : public UniversalGemmPipelineAgBgC
static_assert(std::is_same_v<AQLayout, tensor_layout::gemm::RowMajor>);
if constexpr(PreshuffleQuant)
{
// if(get_block_id() == 0 && get_thread_id() == 0)
// {
// printf("BlockSize: %d, MPerBlock: %d, KPerBlockAQ: %d, VecLoadSize: %d\n",
// BlockSize,
// MPerBlock,
// KPerBlockAQ,
// VecLoadSize);
// }
using TileEncodingPattern =
tile_distribution_encoding_pattern_aq<BlockGemmShape,
WarpGemm,
BlockSize,
MPerBlock / WarpGemm::kM, // 16/16 = 1
ck_tile::integer_least_multiple(
WarpGemm::kM * KPerBlockAQ,
get_warp_size()), //(32, 64) = 64
KPerBlockAQ,
VecLoadSize,
PreshuffleQuant>;
using TileEncodingPattern = tile_distribution_encoding_pattern_aq<
BlockGemmShape,
WarpGemm,
BlockSize,
MPerBlock / WarpGemm::kM,
ck_tile::integer_least_multiple(WarpGemm::kM * KPerBlockAQ, get_warp_size()),
KPerBlockAQ,
VecLoadSize,
PreshuffleQuant>;
return TileEncodingPattern::make_2d_static_tile_distribution();
}
@@ -92,8 +82,8 @@ struct GemmAQuantPipelineAgBgCrDefaultPolicy : public UniversalGemmPipelineAgBgC
using TileEncodingPattern = tile_distribution_encoding_pattern_aq<BlockGemmShape,
WarpGemm,
BlockSize,
MPerBlock, // 16
KPerBlockAQ, // 2
MPerBlock,
KPerBlockAQ,
KPerBlockAQ,
VecLoadSize,
PreshuffleQuant>;

27
include/ck_tile/ops/gemm_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_policy.hpp
View File

@@ -55,25 +55,14 @@ struct GemmBQuantPipelineAgBgCrDefaultPolicy : public UniversalGemmPipelineAgBgC
static_assert(std::is_same_v<BQLayout, tensor_layout::gemm::ColumnMajor>);
if constexpr(PreshuffleQuant)
{
// if(get_block_id() == 0 && get_thread_id() == 0)
// {
// printf("Inside PreshuffleQuant\n BlockSize: %d, YPerTile: %d, XPerTile: %d, "
// "VecLoadSize: %d\n",
// BlockSize,
// NPerBlock / WarpGemm::kN,
// ck_tile::integer_least_multiple(WarpGemm::kN * KPerBlockBQ,
// get_warp_size()), VecLoadSize);
// }
using TileEncodingPattern =
tile_distribution_encoding_pattern_bq<BlockGemmShape,
WarpGemm,
BlockSize,
NPerBlock / WarpGemm::kN, // 64/16 = 4
ck_tile::integer_least_multiple(
WarpGemm::kN * KPerBlockBQ,
get_warp_size()), //(32, 64) = 64
VecLoadSize,
PreshuffleQuant>;
using TileEncodingPattern = tile_distribution_encoding_pattern_bq<
BlockGemmShape,
WarpGemm,
BlockSize,
NPerBlock / WarpGemm::kN,
ck_tile::integer_least_multiple(WarpGemm::kN * KPerBlockBQ, get_warp_size()),
VecLoadSize,
PreshuffleQuant>;
return TileEncodingPattern::make_2d_static_tile_distribution();
}
else

48
include/ck_tile/ops/gemm_quant/pipeline/gemm_group_quant_utils.hpp
View File

@@ -74,17 +74,6 @@ struct tile_distribution_encoding_pattern_aq : public tile_distribution_encoding
{
if constexpr(PreshuffleQuant)
{
// if(get_block_id() == 0 && get_thread_id() == 0)
// {
// printf("YperTile: %d, XPerTile: %d, MWarps: %d, NWarps: %d, MIterPerWarp: %d, "
// "warp_size: %d\n",
// YPerTile, // 1
// XPerTile, // 64
// MWarps, // 1
// NWarps, // 4
// MIterPerWarp, // 1
// warp_size); // 64
// }
// # of elements per thread
static_assert(XPerTile >= warp_size && XPerTile % warp_size == 0);
constexpr index_t X1 = warp_size;
@@ -95,21 +84,20 @@ struct tile_distribution_encoding_pattern_aq : public tile_distribution_encoding
return make_static_tile_distribution(
tile_distribution_encoding<sequence<NWarps>,
tuple<sequence<Y0, Y1>, sequence<X0, X1>>,
tuple<sequence<1, 0>, sequence<2>>, //(MWarp, NWarp),
//(X1)
tuple<sequence<1, 0>, sequence<1>>, //(1, 4), (64)
sequence<1, 2>, // (1), (64/64) = 1
tuple<sequence<1, 0>, sequence<2>>,
tuple<sequence<1, 0>, sequence<1>>,
sequence<1, 2>,
sequence<0, 0>>{});
}
else
{
// # of elements per thread
constexpr index_t X = XPerTile; // 2
constexpr index_t X = XPerTile;
constexpr index_t Y0 = 1;
constexpr index_t Y1 = MIterPerWarp ? MIterPerWarp : 1; // 1
constexpr index_t Y2 = MWarps; // 1
constexpr index_t Y3 = WarpGemm::kM; // 16
constexpr index_t Y1 = MIterPerWarp ? MIterPerWarp : 1;
constexpr index_t Y2 = MWarps;
constexpr index_t Y3 = WarpGemm::kM;
static_assert(Y3 >= WarpGemm::kM,
"Scales for all rows must be available within the warp.");
static_assert(Y0 * Y1 * Y2 * Y3 == YPerTile,
@@ -117,10 +105,9 @@ struct tile_distribution_encoding_pattern_aq : public tile_distribution_encoding
return make_static_tile_distribution(
tile_distribution_encoding<sequence<NWarps>,
tuple<sequence<Y0, Y1, Y2, Y3>, sequence<X>>,
tuple<sequence<1, 0>, sequence<1, 1>>, //(MWarp, NWarp),
//(Y0, Y3)
tuple<sequence<2, 0>, sequence<0, 3>>, //(1, 4), (1, 16)
sequence<1, 2>, //(1, 2(in X direction))
tuple<sequence<1, 0>, sequence<1, 1>>,
tuple<sequence<2, 0>, sequence<0, 3>>,
sequence<1, 2>,
sequence<1, 0>>{});
}
}
@@ -229,18 +216,17 @@ struct tile_distribution_encoding_pattern_bq : public tile_distribution_encoding
{
if constexpr(PreshuffleQuant)
{
constexpr index_t X1 = warp_size; // 64
constexpr index_t X0 = XPerTile / warp_size; // 64/64 = 1
constexpr index_t Y1 = NWarps; // 4
constexpr index_t Y0 = YPerTile / Y1; // 4/4 = 1
constexpr index_t X1 = warp_size;
constexpr index_t X0 = XPerTile / warp_size;
constexpr index_t Y1 = NWarps;
constexpr index_t Y0 = YPerTile / Y1;
return make_static_tile_distribution(
tile_distribution_encoding<sequence<MWarps>,
tuple<sequence<Y0, Y1>, sequence<X0, X1>>,
tuple<sequence<0, 1>, sequence<2>>, //(MWarp, NWarp),
//(warp_size)
tuple<sequence<0, 1>, sequence<1>>, //(1, 4), (64)
sequence<1, 2>, // (1), (64/64) = 1
tuple<sequence<0, 1>, sequence<2>>,
tuple<sequence<0, 1>, sequence<1>>,
sequence<1, 2>,
sequence<0, 0>>{});
}
else