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

refine code

Browse Source
This commit is contained in:
huizzhan
2025-08-14 10:34:02 +00:00
parent 71cea3dd70
commit d86f895180
4 changed files with 84 additions and 229 deletions
Download Patch File Download Diff File Expand all files Collapse all files

65
example/ck_tile/39_gemm_softmax_grouped_topk/block_gemm_softmax_grouped_topk_pipeline_agmem_bgmem_creg.hpp
View File

@@ -210,15 +210,15 @@ struct BlockGemmSoftmaxGroupedTopkPipelineAGmemBGmemCReg
}
#endif
// template <typename ADramBlockWindowTmp, typename BDramBlockWindowTmp, typename DebugBlockTile, typename ValueBlockTile, typename IndexBlockTile>
template <typename ADramBlockWindowTmp, typename BDramBlockWindowTmp>
// template <typename ADramBlockWindowTmp, typename BDramBlockWindowTmp, typename ValueBlockWindow, typename IndexBlockWindow, typename CElementFunction>
CK_TILE_HOST_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp,
const BDramBlockWindowTmp& b_dram_block_window_tmp,
// ValueBlockWindow& value_window,
// IndexBlockWindow& index_window,
// const CElementFunction& c_element_func,
index_t num_loop,
void* p_smem) const
// DebugBlockTile& debug_block_tile,
// ValueBlockTile& value_block_tile,
// IndexBlockTile& index_block_tile,
{
static_assert(
std::is_same_v<ADataType, remove_cvref_t<typename ADramBlockWindowTmp::DataType>> &&
@@ -426,7 +426,8 @@ struct BlockGemmSoftmaxGroupedTopkPipelineAGmemBGmemCReg
}
#endif
// apply softmax for c_block_tile
// -------------------------------------------------------------------------------------
// softmax part
// reduction function for softmax
const auto f_max = [](auto e0, auto e1) { return max(e0, e1); };
const auto f_sum = [](auto e0, auto e1) { return e0 + e1; };
@@ -441,9 +442,6 @@ struct BlockGemmSoftmaxGroupedTopkPipelineAGmemBGmemCReg
auto p_compute =
make_static_distributed_tensor<ComputeDataType>(c_block_tile.get_tile_distribution());
auto debug_block_tile =
make_static_distributed_tensor<WeightType>(p_compute.get_tile_distribution());
constexpr auto p_spans = decltype(p_compute)::get_distributed_spans();
sweep_tile_span(p_spans[number<0>{}], [&](auto idx0) {
@@ -472,8 +470,16 @@ struct BlockGemmSoftmaxGroupedTopkPipelineAGmemBGmemCReg
p_compute(i_j_idx) = p_compute[i_j_idx] / rowsum_p[i_idx];
});
});
// apply topk for softmax output
// -------------------------------------------------------------------------------------
// grouped topk part
auto value_block_tile =
make_static_distributed_tensor<WeightType>(p_compute.get_tile_distribution());
auto index_block_tile =
make_static_distributed_tensor<IndexType>(p_compute.get_tile_distribution());
// auto debug_block_tile =
// make_static_distributed_tensor<WeightType>(p_compute.get_tile_distribution());
auto x_tmp = p_compute;
// calculate group score, need to creat group scores tensor
int num_expert_group = 16;
@@ -518,7 +524,7 @@ struct BlockGemmSoftmaxGroupedTopkPipelineAGmemBGmemCReg
});
}
// // another scheme to reshape x_tmp from 2d to 3d
// // To Do - another scheme to reshape x_tmp from 2d to 3d
// const auto x_tmp_3d = x_tmp.block_tile_reshape((kM, num_expert_group, kN/num_expert_group));
// auto group_scores = block_tile_reduce<ComputeDataType>(
// x_tmp_3d, sequence<2>{}, f_max, std::numeric_limits<ComputeDataType>::lowest());
@@ -530,16 +536,6 @@ struct BlockGemmSoftmaxGroupedTopkPipelineAGmemBGmemCReg
return e0.value > e1.value ? e0 : e1;
};
// // topk_group_index = topk(group_scores, topk_group)
// auto topk_group_index = x_tmp;
// // init topk_group_index to -inf
// sweep_tile_span(p_compute_spans[number<0>{}], [&](auto idx0) {
// sweep_tile_span(p_compute_spans[number<1>{}], [&](auto idx1) {
// constexpr auto i_j_idx = make_tuple(idx0, idx1);
// topk_group_index(i_j_idx) = -numeric<WeightType>::infinity();
// });
// });
// topk_group_mask(1 for selected group scores, -inf for other group scores)
auto topk_group_scores_mask = x_tmp;
// init topk_group_scores_mask to -inf
@@ -560,7 +556,7 @@ struct BlockGemmSoftmaxGroupedTopkPipelineAGmemBGmemCReg
tmp.get_tile_distribution(), make_tuple(idx0, idx1));
constexpr auto i_j_idx = make_tuple(idx0, idx1);
ArgmaxPacket t;
t.value = group_scores(i_j_idx); // !!! we reference p_compute here
t.value = group_scores(i_j_idx);
t.arg = tile_idx.at(number<1>{});
tmp(i_j_idx) = t;
});
@@ -572,8 +568,6 @@ struct BlockGemmSoftmaxGroupedTopkPipelineAGmemBGmemCReg
auto group_r = block_tile_reduce<ArgmaxPacket>(group_packet, sequence<1>{}, f_argmax, argmax_init);
block_tile_reduce_xor_sync(group_r, f_argmax);
// constexpr auto value_spans = decltype(value_block_tile)::get_distributed_spans();
sweep_tile_span(p_compute_spans[number<0>{}], [&](auto idx0) {
constexpr auto i_idx = make_tuple(idx0);
@@ -584,8 +578,6 @@ struct BlockGemmSoftmaxGroupedTopkPipelineAGmemBGmemCReg
auto col_id = tile_idx.at(number<1>{});
constexpr auto i_j_idx = make_tuple(idx0, idx1);
auto k_group_idx = group_r(i_idx).arg;
// topk_group_index(i_j_idx) = (col_id == k_group) ? tmp.value: topk_group_index(i_j_idx);
// topk_group_index(i_j_idx) = (col_id == k_group) ? k_group_idx: topk_group_index(i_j_idx);
topk_group_scores_mask(i_j_idx) = ((col_id >= (k_group_idx * expert_per_group)) && (col_id < ((k_group_idx + 1) * expert_per_group))) ? 1 : topk_group_scores_mask(i_j_idx);
});
});
@@ -638,8 +630,6 @@ struct BlockGemmSoftmaxGroupedTopkPipelineAGmemBGmemCReg
auto r = block_tile_reduce<ArgmaxPacket>(packet, sequence<1>{}, f_argmax, argmax_init);
block_tile_reduce_xor_sync(r, f_argmax);
// constexpr auto value_spans = decltype(value_block_tile)::get_distributed_spans();
sweep_tile_span(p_compute_spans[number<0>{}], [&](auto idx0) {
constexpr auto i_idx = make_tuple(idx0);
@@ -650,10 +640,8 @@ struct BlockGemmSoftmaxGroupedTopkPipelineAGmemBGmemCReg
auto col_id = tile_idx.at(number<1>{});
constexpr auto i_j_idx = make_tuple(idx0, idx1);
ArgmaxPacket tmp = r(i_idx);
// debug_block_tile(i_j_idx) = (col_id == i_k) ? tmp.value: debug_block_tile(i_j_idx);
debug_block_tile(i_j_idx) = (col_id == i_k) ? tmp.arg: debug_block_tile(i_j_idx);
// value_block_tile(i_j_idx) = tmp.value;
// index_block_tile(i_j_idx) = tmp.arg;
value_block_tile(i_j_idx) = (col_id == i_k) ? tmp.value: value_block_tile(i_j_idx);
index_block_tile(i_j_idx) = (col_id == i_k) ? tmp.arg: index_block_tile(i_j_idx);
});
});
@@ -672,8 +660,17 @@ struct BlockGemmSoftmaxGroupedTopkPipelineAGmemBGmemCReg
});
});
}
return debug_block_tile;
// return x_tmp_masked;
// // cast DataType and apply CElementFunction
// const auto value_cast_block_tile = tile_elementwise_in(
// [&](const auto& value) { return c_element_func(type_convert<WeightType>(value)); },
// value_block_tile);
// const auto index_cast_block_tile = tile_elementwise_in(
// [&](const auto& index) { return c_element_func(type_convert<IndexType>(index)); },
// index_block_tile);
// store_tile(value_window, value_cast_block_tile);
// store_tile(index_window, index_cast_block_tile);
return value_block_tile;
}
};

155
example/ck_tile/39_gemm_softmax_grouped_topk/gemm_softmax_grouped_topk.cpp
View File

@@ -69,9 +69,7 @@ int main(int argc, char* argv[])
using AccDataType = float;
using CDataType = ck_tile::half_t;
using WeightType = float;
using DebugType = float;
// using IndexType = ck_tile::index_t;
using IndexType = float;
using IndexType = ck_tile::index_t;
ck_tile::index_t verification = 0;
ck_tile::index_t M = 3328;
@@ -146,9 +144,6 @@ int main(int argc, char* argv[])
const auto b_lengths = std::array<ck_tile::index_t, 2>{N, K};
const auto b_strides = std::array<ck_tile::index_t, 2>{Ldb, 1};
const auto debug_lengths = std::array<ck_tile::index_t, 2>{M, N};
const auto debug_strides = std::array<ck_tile::index_t, 2>{N, 1};
const auto out_lengths = std::array<ck_tile::index_t, 2>{M, topk};
const auto out_strides = std::array<ck_tile::index_t, 2>{Ldout, 1};
@@ -158,57 +153,11 @@ int main(int argc, char* argv[])
ck_tile::HostTensor<WeightType> value_host_dev(out_lengths, out_strides);
ck_tile::HostTensor<IndexType> index_host_dev(out_lengths, out_strides);
// ck_tile::FillUniformDistributionIntegerValue<ADataType>{-5.f, 5.f}(a_host);
// ck_tile::FillUniformDistributionIntegerValue<BDataType>{-5.f, 5.f}(b_host);
ck_tile::FillUniformDistribution<ADataType>{0.01f, 0.05f}(a_host);
ck_tile::FillUniformDistribution<BDataType>{0.01f, 0.05f}(b_host);
// ck_tile::HostTensor<WeightType> debug_host_input(debug_lengths, debug_strides);
ck_tile::HostTensor<DebugType> debug_host_dev(debug_lengths, debug_strides);
// // std::random_device rd;
// // std::mt19937 gen(rd());
// // std::uniform_real_distribution<> dist_b(0.01f, 0.05f);
// printf("===============debug input=====================\n");
// // std::mt19937 rng(123);
// // std::uniform_int_distribution<int> dist_debug_input(1, 100);
// for(int m = 0; m < M; ++m) {
// printf("m: %d ", m);
// for(int n = 0; n < N; ++n) {
// // debug_host_input(m, n) = float(dist_debug_input(rng));
// debug_host_input(m, n) = sin(float(m + n)) * 100;
// printf("[%d]:%.4f ", n, debug_host_input(m, n));
// }
// printf("/n");
// }
// // std::random_device rd;
// std::mt19937 gen(12345);
// std::uniform_real_distribution<> dist_a(0.f, 0.001f);
// std::uniform_real_distribution<> dist_b(0.001f, 0.005f);
// for(int m = 0; m < M; ++m) {
// for(int k = 0; k < K; ++k) {
// a_host(m, k) = dist_a(gen);
// }
// }
// for(int n = 0; n < K; ++n) {
// for(int k = 0; k < K; ++k) {
// b_host(n, k) = dist_b(gen);
// }
// }
for(std::size_t i = 0; i < 20; ++i) {
const double a = *std::next(std::begin(a_host), i);
const double b = *std::next(std::begin(b_host), i);
std::cout << " a[" << i << "]: " << a << " " << "b[" << i << "]: " << b << std::endl;
}
ck_tile::DeviceMem a_buf(a_host.get_element_space_size_in_bytes());
ck_tile::DeviceMem b_buf(b_host.get_element_space_size_in_bytes());
ck_tile::DeviceMem debug_buf(debug_host_dev.get_element_space_size_in_bytes());
ck_tile::DeviceMem value_buf(value_host_dev.get_element_space_size_in_bytes());
ck_tile::DeviceMem index_buf(index_host_dev.get_element_space_size_in_bytes());
@@ -218,7 +167,6 @@ int main(int argc, char* argv[])
// Alignment
constexpr ck_tile::index_t kAAlignment = 8;
constexpr ck_tile::index_t kBAlignment = 8;
// constexpr ck_tile::index_t kCAlignment = 8;
constexpr ck_tile::index_t kOutAlignment = 8;
constexpr ck_tile::index_t kBlockSize = 256;
@@ -265,13 +213,12 @@ int main(int argc, char* argv[])
0,
static_cast<ADataType*>(a_buf.GetDeviceBuffer()),
static_cast<BDataType*>(b_buf.GetDeviceBuffer()),
static_cast<DebugType*>(debug_buf.GetDeviceBuffer()),
static_cast<WeightType*>(value_buf.GetDeviceBuffer()),
static_cast<IndexType*>(index_buf.GetDeviceBuffer()),
M,
N,
K,
topk,
// topk,
Lda,
Ldb,
Ldout,
@@ -280,27 +227,14 @@ int main(int argc, char* argv[])
bool rtn = true;
if(verification)
{
ck_tile::HostTensor<DebugType> debug_ref({M, N}, {N, 1});
ck_tile::HostTensor<WeightType> value_ref(out_lengths, out_strides);
ck_tile::HostTensor<IndexType> index_ref(out_lengths, out_strides);
// reference_topk(debug_host_input, value_ref, index_ref, topk);
// debug_ref = reference_basic_gemm<ADataType, ADataType, AccDataType>(a_host, b_host);
// debug_ref = reference_basic_gemm_softmax<ADataType, ADataType, AccDataType>(a_host, b_host);
reference_basic_gemm_softmax_grouped_topk<ADataType, ADataType, AccDataType, WeightType, IndexType>(
a_host, b_host, value_ref, index_ref, topk);
debug_buf.FromDevice(debug_host_dev.mData.data());
value_buf.FromDevice(value_host_dev.mData.data());
index_buf.FromDevice(index_host_dev.mData.data());
// rtn &= ck_tile::check_err(debug_host_dev, debug_ref);
// for(std::size_t i = 0; i < debug_ref.size(); ++i) {
// const double o = *std::next(std::begin(debug_host_dev), i);
// const double r = *std::next(std::begin(debug_ref), i);
// std::cout << " out[" << i << "] != ref[" << i << "]: " << o << " != " << r << std::endl;
// }
// std::cout << "valid:" << (rtn ? "y" : "n") << std::endl;
const ck_tile::index_t tokens = M;
auto [rtol, atol] = get_elimit<WeightType>("");
for(int i_t = 0; i_t < tokens; i_t++)
@@ -308,75 +242,40 @@ int main(int argc, char* argv[])
auto s_begin = std::vector<size_t>{static_cast<size_t>(i_t), static_cast<size_t>(0)};
auto s_end =
std::vector<size_t>{static_cast<size_t>(i_t + 1), static_cast<size_t>(topk)};
// auto s_end =
// std::vector<size_t>{static_cast<size_t>(i_t + 1), static_cast<size_t>(N)};
auto s_debug_host = debug_host_dev.slice(s_begin, s_end);
// auto s_debug_ref = debug_ref.slice(s_begin, s_end);
// auto s_debug_ref = value_ref.slice(s_begin, s_end);
auto s_debug_ref = index_ref.slice(s_begin, s_end);
rtn &= ck_tile::check_err(s_debug_host,
s_debug_ref,
auto s_value_host = value_host_dev.slice(s_begin, s_end);
auto s_value_ref = value_ref.slice(s_begin, s_end);
rtn &= ck_tile::check_err(s_value_host,
s_value_ref,
std::string("[") + std::to_string(i_t) +
std::string("] Value Error:"),
rtol,
atol);
printf("row [%d]\n", i_t);
for(std::size_t i = 0; i < s_debug_ref.size(); ++i) {
// const double o = *std::next(std::begin(s_debug_host), i);
const double r = *std::next(std::begin(s_debug_ref), i);
printf("ref[%zu]:%.8f ", i, r);
// std::cout << i_t << " out[" << i << "] != ref[" << i << "]: " << o << " != " << r << std::endl;
}
printf("\n");
for(std::size_t i = 0; i < s_debug_ref.size(); ++i) {
const double o = *std::next(std::begin(s_debug_host), i);
// const double r = *std::next(std::begin(s_debug_ref), i);
printf("out[%zu]:%.8f ", i, o);
// std::cout << i_t << " out[" << i << "] != ref[" << i << "]: " << o << " != " << r << std::endl;
}
printf("\n");
// auto s_index_host = index_host_dev.slice(s_begin, s_end);
// auto s_index_ref = index_ref.slice(s_begin, s_end);
// rtn &= ck_tile::check_err(s_index_host,
// s_index_ref,
// std::string("[") + std::to_string(i_t) +
// std::string("] Index Error:"),
// rtol,
// atol);
// printf("row [%d]\n", i_t);
// for(std::size_t i = 0; i < s_debug_ref.size(); ++i) {
// const double r = *std::next(std::begin(s_index_ref), i);
// printf("ref[%zu]:%.8f ", i, r);
// }
// printf("\n");
// for(std::size_t i = 0; i < s_debug_ref.size(); ++i) {
// const double o = *std::next(std::begin(s_index_host), i);
// printf("out[%zu]:%.8f ", i, o);
// }
// printf("\n");
}
std::cout << "valid:" << (rtn ? "y" : "n") << std::endl;
}
// const ck_tile::index_t tokens = M;
// auto [rtol, atol] = get_elimit<ADataType>("");
// for(int i_t = 0; i_t < tokens; i_t++)
// {
// auto s_begin = std::vector<size_t>{static_cast<size_t>(i_t), static_cast<size_t>(0)};
// auto s_end =
// std::vector<size_t>{static_cast<size_t>(i_t + 1), static_cast<size_t>(topk)};
// auto s_value_host = value_host_dev.slice(s_begin, s_end);
// auto s_value_ref = value_ref.slice(s_begin, s_end);
// rtn &= ck_tile::check_err(s_value_host,
// s_value_ref,
// std::string("[") + std::to_string(i_t) +
// std::string("] Value Error:"),
// rtol,
// atol);
// // for(std::size_t i = 0; i < s_value_ref.size(); ++i) {
// // const double o = *std::next(std::begin(s_value_host), i);
// // const double r = *std::next(std::begin(s_value_ref), i);
// // std::cout << " out[" << i << "] != ref[" << i << "]: " << o << " != " << r << std::endl;
// // }
// auto s_index_host = index_host_dev.slice(s_begin, s_end);
// auto s_index_ref = index_ref.slice(s_begin, s_end);
// rtn &= ck_tile::check_err(s_index_host,
// s_index_ref,
// std::string("[") + std::to_string(i_t) +
// std::string("] Index Error:"),
// rtol,
// atol);
// // for(std::size_t i = 0; i < s_index_ref.size(); ++i) {
// // const double o = *std::next(std::begin(s_index_host), i);
// // const double r = *std::next(std::begin(s_index_ref), i);
// // std::cout << " out[" << i << "] != ref[" << i << "]: " << o << " != " << r << std::endl;
// // }
// }
// std::cout << "valid:" << (rtn ? "y" : "n") << std::endl;
// }
std::cout << "Perf: " << ave_time << " ms, " << std::endl;
return rtn;
// return !rtn;
}

14
example/ck_tile/39_gemm_softmax_grouped_topk/gemm_softmax_grouped_topk.hpp
View File

@@ -179,13 +179,12 @@ struct Gemm
CK_TILE_DEVICE void operator()(const ADataType* p_a,
const BDataType* p_b,
WeightType* p_debug,
WeightType* p_value,
IndexType* p_index,
const index_t M,
const index_t N,
const index_t K,
const index_t topK,
// const index_t topK,
const index_t Lda,
const index_t Ldb,
const index_t Ldout,
@@ -201,22 +200,17 @@ struct Gemm
p_b, make_tuple(N, K), make_tuple(Ldb, 1), number<kBAlignment>{}, number<1>{});
}();
const auto debug_dram = [&] {
return make_naive_tensor_view<address_space_enum::global>(
p_debug, make_tuple(M, N), make_tuple(N, 1), number<kOutAlignment>{}, number<1>{});
}();
const auto value_dram = [&] {
return make_naive_tensor_view<address_space_enum::global>(
p_value, make_tuple(M, topK), make_tuple(Ldout, 1), number<kOutAlignment>{}, number<1>{});
p_value, make_tuple(M, N), make_tuple(Ldout, 1), number<kOutAlignment>{}, number<1>{});
}();
const auto index_dram = [&] {
return make_naive_tensor_view<address_space_enum::global>(
p_index, make_tuple(M, topK), make_tuple(Ldout, 1), number<kOutAlignment>{}, number<1>{});
p_index, make_tuple(M, N), make_tuple(Ldout, 1), number<kOutAlignment>{}, number<1>{});
}();
GridGemm{}(a_dram, b_dram, debug_dram, value_dram, index_dram, c_element_func);
GridGemm{}(a_dram, b_dram, value_dram, index_dram, c_element_func);
}
};

79
example/ck_tile/39_gemm_softmax_grouped_topk/grid_gemm_softmax_grouped_topk.hpp
View File

@@ -23,10 +23,9 @@ struct GridGemm
static constexpr auto topk = Policy::kTopKPerBlock;
static constexpr auto kBlockSize = Policy::kBlockSize;
template <typename AGridTensorView, typename BGridTensorView, typename DebugGridTensorView, typename ValueGridTensorView, typename IndexGridTensorView>
template <typename AGridTensorView, typename BGridTensorView, typename ValueGridTensorView, typename IndexGridTensorView>
CK_TILE_DEVICE void operator()(const AGridTensorView& a_grid,
const BGridTensorView& b_grid,
DebugGridTensorView& debug_grid,
ValueGridTensorView& value_grid,
IndexGridTensorView& index_grid,
const CElementFunction& c_element_func) const
@@ -62,15 +61,13 @@ struct GridGemm
__shared__ char p_smem_char[block_gemm_pipeline.GetStaticLdsSize()];
// // store C
// auto c_window = make_tile_window(
// c_grid, make_tuple(number<kMPerBlock>{}, number<kNPerBlock>{}), {iM, iN});
// store value and index
// constexpr index_t kBlockSize = Problem::kBlockSize;
// constexpr index_t kMPerBlock = Problem::BlockGemmShape::kM;
// constexpr index_t kKPerBlock = Problem::BlockGemmShape::kK;
// store value and index window
auto value_window = make_tile_window(
value_grid, make_tuple(number<kMPerBlock>{}, number<kNPerBlock>{}), {iM, iN});
// auto index_window = make_tile_window(
// index_grid, make_tuple(number<kMPerBlock>{}, number<kNPerBlock>{}), {iM, iN});
// create tile_distribution for value and index window
constexpr index_t K1 = 16 / sizeof(WeightType);
constexpr index_t K0 = topk / K1;
constexpr index_t M2 = get_warp_size() / K0;
@@ -78,15 +75,15 @@ struct GridGemm
constexpr index_t M1 = kBlockSize / get_warp_size();
constexpr index_t M0 = kMPerBlock / (M2 * M1);
auto value_window = make_tile_window(
value_grid, make_tuple(number<kMPerBlock>{}, number<topk>{}), {iM, iN},
make_static_tile_distribution(
tile_distribution_encoding<sequence<1>,
tuple<sequence<M0, M1, M2>, sequence<K0, K1>>,
tuple<sequence<1>, sequence<1, 2>>,
tuple<sequence<1>, sequence<2, 0>>,
sequence<1, 2>,
sequence<0, 1>>{}));
// auto value_window = make_tile_window(
// value_grid, make_tuple(number<kMPerBlock>{}, number<topk>{}), {iM, iN},
// make_static_tile_distribution(
// tile_distribution_encoding<sequence<1>,
// tuple<sequence<M0, M1, M2>, sequence<K0, K1>>,
// tuple<sequence<1>, sequence<1, 2>>,
// tuple<sequence<1>, sequence<2, 0>>,
// sequence<1, 2>,
// sequence<0, 1>>{}));
auto index_window = make_tile_window(
index_grid, make_tuple(number<kMPerBlock>{}, number<topk>{}), {iM, iN},
make_static_tile_distribution(
@@ -97,53 +94,23 @@ struct GridGemm
sequence<1, 2>,
sequence<0, 1>>{}));
using ValueBlockTileDistr = decltype(value_window.get_tile_distribution());
// using ValueBlockTileDistr = decltype(value_window.get_tile_distribution());
using IndexBlockTileDistr = decltype(index_window.get_tile_distribution());
using ValueBlockTile = decltype(make_static_distributed_tensor<WeightType>(ValueBlockTileDistr{}));
// using ValueBlockTile = decltype(make_static_distributed_tensor<WeightType>(ValueBlockTileDistr{}));
using IndexBlockTile = decltype(make_static_distributed_tensor<IndexType>(IndexBlockTileDistr{}));
ValueBlockTile value_block_tile;
// ValueBlockTile value_block_tile;
IndexBlockTile index_block_tile;
// Initialize value_block_tile and index_block_tile
tile_elementwise_inout([](auto& value) { value = 0; }, value_block_tile);
// tile_elementwise_inout([](auto& value) { value = 0; }, value_block_tile);
tile_elementwise_inout([](auto& index) { index = 0; }, index_block_tile);
// constexpr index_t debugK1 = 16 / sizeof(WeightType);
// constexpr index_t debugK0 = kNPerBlock / debugK1;
// constexpr index_t debugM2 = get_warp_size() / debugK0;
// // coalesce reading for each blocks
// constexpr index_t debugM1 = kBlockSize / get_warp_size();
// constexpr index_t debugM0 = kMPerBlock / (debugM2 * debugM1);
auto debug_window = make_tile_window(
debug_grid, make_tuple(number<kMPerBlock>{}, number<kNPerBlock>{}), {iM, iN});
// auto debug_window = make_tile_window(
// debug_grid, make_tuple(number<kMPerBlock>{}, number<kNPerBlock>{}), {iM, iN},
// make_static_tile_distribution(
// tile_distribution_encoding<sequence<1>,
// tuple<sequence<debugM0, debugM1, debugM2>, sequence<debugK0, debugK1>>,
// tuple<sequence<1>, sequence<1, 2>>,
// tuple<sequence<1>, sequence<2, 0>>,
// sequence<1, 2>,
// sequence<0, 1>>{}));
// using DebugBlockTileDistr = decltype(debug_window.get_tile_distribution());
// using DebugBlockTile = decltype(make_static_distributed_tensor<WeightType>(DebugBlockTileDistr{}));
// DebugBlockTile debug_block_tile;
// tile_elementwise_inout([](auto& debug) { debug = 0; }, debug_block_tile);
// block_gemm_pipeline(a_block_window, b_block_window, debug_block_tile, value_block_tile, index_block_tile, K / kKPerBlock, p_smem_char);
const auto debug_block_tile = block_gemm_pipeline(a_block_window, b_block_window, K / kKPerBlock, p_smem_char);
// block_gemm_pipeline(a_block_window, b_block_window, debug_block_tile, K / kKPerBlock, p_smem_char);
// block_gemm_pipeline(a_block_window, b_block_window, value_window, index_window, K / kKPerBlock, p_smem_char, c_element_func);
const auto value_block_tile = block_gemm_pipeline(a_block_window, b_block_window, K / kKPerBlock, p_smem_char);
// cast DataType and apply CElementFunction
const auto debug_cast_block_tile = tile_elementwise_in(
[&](const auto& debug) { return c_element_func(type_convert<WeightType>(debug)); },
debug_block_tile);
const auto value_cast_block_tile = tile_elementwise_in(
[&](const auto& value) { return c_element_func(type_convert<WeightType>(value)); },
value_block_tile);
@@ -152,8 +119,6 @@ struct GridGemm
[&](const auto& index) { return c_element_func(type_convert<IndexType>(index)); },
index_block_tile);
store_tile(debug_window, debug_cast_block_tile);
store_tile(value_window, value_cast_block_tile);
store_tile(index_window, index_cast_block_tile);
}