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

basic gemm softmax topk

Browse Source
This commit is contained in:
huizzhan
2025-08-05 09:04:26 +00:00
parent edf79c7064
commit e298f8194a
6 changed files with 306 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files

81
example/ck_tile/39_gemm_softmax/block_gemm_pipeline_agmem_bgmem_creg.hpp
View File

@@ -23,6 +23,8 @@ struct BlockGemmSoftmaxPipelineAGmemBGmemCReg
using ADataType = remove_cvref_t<typename Problem::ADataType>;
using BDataType = remove_cvref_t<typename Problem::BDataType>;
using CDataType = remove_cvref_t<typename Problem::CDataType>;
using WeightType = remove_cvref_t<typename Problem::WeightType>;
using IndexType = remove_cvref_t<typename Problem::IndexType>;
using ComputeDataType = float;
using BlockGemmShape = remove_cvref_t<typename Problem::BlockGemmShape>;
@@ -31,6 +33,14 @@ struct BlockGemmSoftmaxPipelineAGmemBGmemCReg
static constexpr index_t kMPerBlock = BlockGemmShape::kM;
static constexpr index_t kNPerBlock = BlockGemmShape::kN;
static constexpr index_t kKPerBlock = BlockGemmShape::kK;
static constexpr index_t topk = BlockGemmShape::kTopK;
// for topk computing
struct ArgmaxPacket
{
WeightType arg;
IndexType value;
};
using BlockGemm = remove_cvref_t<decltype(Policy::template GetBlockGemm<Problem>())>;
@@ -198,9 +208,11 @@ struct BlockGemmSoftmaxPipelineAGmemBGmemCReg
}
#endif
template <typename ADramBlockWindowTmp, typename BDramBlockWindowTmp>
CK_TILE_HOST_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp,
template <typename ADramBlockWindowTmp, typename BDramBlockWindowTmp, typename ValueBlockTile, typename IndexBlockTile>
CK_TILE_HOST_DEVICE void operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp,
const BDramBlockWindowTmp& b_dram_block_window_tmp,
ValueBlockTile& value_block_tile,
IndexBlockTile& index_block_tile,
index_t num_loop,
void* p_smem) const
{
@@ -453,11 +465,70 @@ struct BlockGemmSoftmaxPipelineAGmemBGmemCReg
p_compute(i_j_idx) = p_compute[i_j_idx] / rowsum_p[i_idx];
});
});
// CDramBlockWindowTmp c_dram_block_window_tmp = c_dram_block_window;
// apply topk for softmax output
auto x_tmp = p_compute;
// constexpr auto dst_dist = BlockGemmPipelineAGmemBGmemCRegDefaultPolicy::MakeOutputDistribution();
// store_tile(c_dram_block_window_tmp, type_convert<CDataType>(p_compute));
// argmax for topk
const auto f_argmax = [](ArgmaxPacket e0, ArgmaxPacket e1) {
return e0.arg > e1.arg ? e0 : e1;
};
return p_compute;
for(index_t i_k = 0; i_k < topk; i_k++)
{
constexpr auto span_2d = decltype(p_compute)::get_distributed_spans();
auto packet = [&]() {
auto tmp = make_static_distributed_tensor<ArgmaxPacket>(p_compute.get_tile_distribution());
sweep_tile_span(span_2d[number<0>{}], [&](auto idx0) {
sweep_tile_span(span_2d[number<1>{}], [&](auto idx1) {
const auto tile_idx = get_x_indices_from_distributed_indices(
tmp.get_tile_distribution(), make_tuple(idx0, idx1));
constexpr auto i_j_idx = make_tuple(idx0, idx1);
ArgmaxPacket t;
t.arg = x_tmp(i_j_idx); // !!! we reference p_compute here
t.value = tile_idx.at(number<1>{});
tmp(i_j_idx) = t;
});
});
return tmp;
}();
auto argmax_init = ArgmaxPacket{-numeric<WeightType>::infinity(), 0};
auto r = block_tile_reduce<ArgmaxPacket>(packet, sequence<1>{}, f_argmax, argmax_init);
block_tile_reduce_xor_sync(r, f_argmax);
// auto value_block_tile = make_static_distributed_tensor<WeightType>(dst_dist);
// auto index_block_tile = make_static_distributed_tensor<IndexType>(dst_dist);
// Initialize value_block_tile and index_block_tile
tile_elementwise_inout([](auto& value) { value = 0; }, value_block_tile);
tile_elementwise_inout([](auto& index) { index = 0; }, index_block_tile);
sweep_tile_span(span_2d[number<0>{}], [&](auto idx0) {
sweep_tile_span(span_2d[number<1>{}], [&](auto idx1) {
constexpr auto i_j_idx = make_tuple(idx0, idx1);
ArgmaxPacket tmp = r(i_j_idx);
value_block_tile(i_j_idx) = tmp.arg;
index_block_tile(i_j_idx) = tmp.value;
});
});
// update value
sweep_tile_span(span_2d[number<0>{}], [&](auto idx0) {
sweep_tile_span(span_2d[number<1>{}], [&](auto idx1) {
const auto tile_idx = get_x_indices_from_distributed_indices(
p_compute.get_tile_distribution(), make_tuple(idx0, idx1));
auto col_id = tile_idx.at(number<1>{});
constexpr auto i_j_idx = make_tuple(idx0, idx1);
x_tmp(i_j_idx) = (col_id == r(i_j_idx).value) ? -numeric<WeightType>::infinity()
: x_tmp(i_j_idx);
});
});
}
}
};

33
example/ck_tile/39_gemm_softmax/block_gemm_pipeline_agmem_bgmem_creg_default_policy.hpp
View File

@@ -261,6 +261,39 @@ struct BlockGemmPipelineAGmemBGmemCRegDefaultPolicy
sequence<0, 1>>{});
}
// template <typename Problem>
// CK_TILE_HOST_DEVICE static constexpr auto MakeOutputDistribution()
// {
// using WeightType = remove_cvref_t<typename Problem::WeightType>;
// constexpr index_t kBlockSize = Problem::kBlockSize;
// constexpr index_t kMPerBlock = Problem::BlockGemmShape::kM;
// constexpr index_t topk = Problem::BlockGemmShape::kTopK;
// constexpr index_t K1 = 16 / sizeof(WeightType);
// constexpr index_t K0 = topk / K1;
// constexpr index_t M2 = get_warp_size() / K0;
// // coalesce reading for each blocks
// constexpr index_t M1 = kBlockSize / get_warp_size();
// constexpr index_t M0 = kMPerBlock / (M2 * M1);
// // return 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>>{});
// return make_static_tile_distribution(
// tile_distribution_encoding<sequence<1>,
// tuple<sequence<2, 4, 16>, sequence<4, 4>>,
// tuple<sequence<1>, sequence<1, 2>>,
// tuple<sequence<1>, sequence<2, 0>>,
// sequence<1, 2>,
// sequence<0, 1>>{});
// }
#if defined(ENABLE_INSTRUCTION_SCH)
static constexpr auto I0 = number<0>{};
static constexpr auto I1 = number<1>{};

42
example/ck_tile/39_gemm_softmax/gemm.hpp
View File

@@ -18,6 +18,8 @@ template <typename ADataType_,
typename BDataType_,
typename AccDataType_,
typename CDataType_,
typename WeightType_,
typename IndexType_,
typename CElementFunction_>
struct GridGemmProblem
{
@@ -25,21 +27,26 @@ struct GridGemmProblem
using BDataType = BDataType_;
using AccDataType = AccDataType_;
using CDataType = CDataType_;
using WeightType = WeightType_;
using IndexType = IndexType_;
using CElementFunction = CElementFunction_;
};
template <index_t kMPerTile, index_t kNPerTile, index_t kKPerTile>
template <index_t kMPerTile, index_t kNPerTile, index_t kKPerTile, index_t TopkKPerTile>
struct TileGemmShape
{
static constexpr index_t kM = kMPerTile;
static constexpr index_t kN = kNPerTile;
static constexpr index_t kK = kKPerTile;
static constexpr index_t kTopK = TopkKPerTile;
};
template <typename ADataType_,
typename BDataType_,
typename CDataType_,
typename WeightType_,
typename IndexType_,
index_t kBlockSize_,
typename BlockGemmShape_>
struct BlockGemmPipelineProblem
@@ -47,6 +54,8 @@ struct BlockGemmPipelineProblem
using ADataType = remove_cvref_t<ADataType_>;
using BDataType = remove_cvref_t<BDataType_>;
using CDataType = remove_cvref_t<CDataType_>;
using WeightType = remove_cvref_t<WeightType_>;
using IndexType = remove_cvref_t<IndexType_>;
using BlockGemmShape = remove_cvref_t<BlockGemmShape_>;
static constexpr index_t kBlockSize = kBlockSize_;
@@ -57,18 +66,21 @@ template <typename ADataType,
typename BDataType,
typename AccDataType,
typename CDataType,
typename WeightType,
typename IndexType,
typename CElementFunction,
index_t kAAlignment,
index_t kBAlignment,
index_t kCAlignment,
index_t kOutAlignment,
index_t kBlockSize_,
index_t kMPerBlock_,
index_t kNPerBlock_,
index_t kKPerBlock_>
index_t kKPerBlock_,
index_t kTopKPerBlock_>
struct Gemm
{
using GridGemmProblem =
GridGemmProblem<ADataType, BDataType, AccDataType, CDataType, CElementFunction>;
GridGemmProblem<ADataType, BDataType, AccDataType, CDataType, WeightType, IndexType, CElementFunction>;
struct GridGemmPolicy
{
@@ -76,6 +88,7 @@ struct Gemm
static constexpr index_t kMPerBlock = kMPerBlock_;
static constexpr index_t kNPerBlock = kNPerBlock_;
static constexpr index_t kKPerBlock = kKPerBlock_;
static constexpr index_t kTopKPerBlock = kTopKPerBlock_;
template <typename Problem>
CK_TILE_HOST_DEVICE static constexpr auto MakeBlock2TileMap(index_t M0, index_t N0)
@@ -154,8 +167,10 @@ struct Gemm
BlockGemmPipelineProblem<ADataType,
BDataType,
AccDataType,
WeightType,
IndexType,
kBlockSize,
TileGemmShape<kMPerBlock, kNPerBlock, kKPerBlock>>;
TileGemmShape<kMPerBlock, kNPerBlock, kKPerBlock, kTopKPerBlock>>;
return BlockGemmSoftmaxPipelineAGmemBGmemCReg<BlockGemmPipelineProblem_>{};
}
};
@@ -164,13 +179,15 @@ struct Gemm
CK_TILE_DEVICE void operator()(const ADataType* p_a,
const BDataType* p_b,
CDataType* p_c,
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 Lda,
const index_t Ldb,
const index_t Ldc,
const index_t Ldout,
const CElementFunction& c_element_func) const
{
const auto a_dram = [&] {
@@ -183,12 +200,17 @@ struct Gemm
p_b, make_tuple(N, K), make_tuple(Ldb, 1), number<kBAlignment>{}, number<1>{});
}();
const auto c_dram = [&] {
const auto value_dram = [&] {
return make_naive_tensor_view<address_space_enum::global>(
p_c, make_tuple(M, N), make_tuple(Ldc, 1), number<kCAlignment>{}, number<1>{});
p_value, make_tuple(M, topK), make_tuple(Ldout, 1), number<kOutAlignment>{}, number<1>{});
}();
GridGemm{}(a_dram, b_dram, c_dram, c_element_func);
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>{});
}();
GridGemm{}(a_dram, b_dram, value_dram, index_dram, c_element_func);
}
};

128
example/ck_tile/39_gemm_softmax/gemm_softmax.cpp
View File

@@ -26,28 +26,66 @@ struct CElementFunction
}
};
// different threshold for different dtype
template <typename DataType>
auto get_elimit(std::string /*init_method*/)
{
double rtol = 1e-3;
double atol = 1e-3;
return ck_tile::make_tuple(rtol, atol);
}
template <>
auto get_elimit<ck_tile::bf16_t>(std::string /*init_method*/)
{
double rtol = 1e-2;
double atol = 1e-2;
return ck_tile::make_tuple(rtol, atol);
}
template <>
auto get_elimit<ck_tile::fp8_t>(std::string init_method)
{
if(init_method == "ui" || init_method == "ni")
{
unsigned max_rounding_point_distance = 0;
double atol = 2e-3;
return ck_tile::make_tuple(max_rounding_point_distance, atol);
}
else
{
unsigned max_rounding_point_distance = 1;
double atol = 0.0625;
return ck_tile::make_tuple(max_rounding_point_distance, atol);
}
}
int main(int argc, char* argv[])
{
using ADataType = ck_tile::half_t;
using BDataType = ck_tile::half_t;
using AccDataType = float;
using CDataType = ck_tile::half_t;
using WeightType = float;
using IndexType = ck_tile::index_t;
ck_tile::index_t verification = 0;
ck_tile::index_t M = 3328;
ck_tile::index_t N = 4096;
ck_tile::index_t K = 4096;
ck_tile::index_t topk = 16;
if(argc == 2)
{
verification = std::stoi(argv[1]);
}
if(argc == 5)
if(argc == 6)
{
verification = std::stoi(argv[1]);
M = std::stoi(argv[2]);
N = std::stoi(argv[3]);
K = std::stoi(argv[4]);
topk = std::stoi(argv[5]);
}
#if defined(KERNEL_A)
@@ -95,7 +133,8 @@ int main(int argc, char* argv[])
const ck_tile::index_t Lda = K;
const ck_tile::index_t Ldb = K;
const ck_tile::index_t Ldc = N;
// const ck_tile::index_t Ldc = N;
const ck_tile::index_t Ldout = topk;
const auto a_lengths = std::array<ck_tile::index_t, 2>{M, K};
const auto a_strides = std::array<ck_tile::index_t, 2>{Lda, 1};
@@ -103,20 +142,27 @@ 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 c_lengths = std::array<ck_tile::index_t, 2>{M, N};
const auto c_strides = std::array<ck_tile::index_t, 2>{Ldc, 1};
// const auto c_lengths = std::array<ck_tile::index_t, 2>{M, N};
// const auto c_strides = std::array<ck_tile::index_t, 2>{Ldc, 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};
// host verify
ck_tile::HostTensor<ADataType> a_host(a_lengths, a_strides);
ck_tile::HostTensor<BDataType> b_host(b_lengths, b_strides);
ck_tile::HostTensor<CDataType> c_host_dev(c_lengths, c_strides);
// ck_tile::HostTensor<CDataType> c_host_dev(c_lengths, c_strides);
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::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 c_buf(c_host_dev.get_element_space_size_in_bytes());
// ck_tile::DeviceMem c_buf(c_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());
a_buf.ToDevice(a_host.mData.data());
b_buf.ToDevice(b_host.mData.data());
@@ -124,7 +170,8 @@ 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 kCAlignment = 8;
constexpr ck_tile::index_t kOutAlignment = 8;
constexpr ck_tile::index_t kBlockSize = 256;
@@ -136,6 +183,7 @@ int main(int argc, char* argv[])
constexpr ck_tile::index_t kGemmKPerBlock = 16;
#endif
constexpr ck_tile::index_t kGemmNPerBlock = 256;
constexpr ck_tile::index_t kGemmTopKPerBlock = 16;
ck_tile::index_t kGridSize = (M / kGemmMPerBlock) * (N / kGemmNPerBlock);
@@ -149,14 +197,17 @@ int main(int argc, char* argv[])
BDataType,
AccDataType,
CDataType,
WeightType,
IndexType,
CElementFunction,
kAAlignment,
kBAlignment,
kCAlignment,
kOutAlignment,
kBlockSize,
kGemmMPerBlock,
kGemmNPerBlock,
kGemmKPerBlock>;
kGemmKPerBlock,
kGemmTopKPerBlock>;
float ave_time = ck_tile::launch_kernel(ck_tile::stream_config{nullptr, true, 0, 5, 1000},
ck_tile::make_kernel<kBlockSize, kBlockPerCu>(
@@ -166,25 +217,61 @@ int main(int argc, char* argv[])
0,
static_cast<ADataType*>(a_buf.GetDeviceBuffer()),
static_cast<BDataType*>(b_buf.GetDeviceBuffer()),
static_cast<CDataType*>(c_buf.GetDeviceBuffer()),
static_cast<WeightType*>(value_buf.GetDeviceBuffer()),
static_cast<IndexType*>(index_buf.GetDeviceBuffer()),
M,
N,
K,
topk,
Lda,
Ldb,
Ldc,
Ldout,
CElementFunction{}));
auto pass = true;
// auto pass = true;
bool rtn = true;
if(verification)
{
// reference gemm
ck_tile::HostTensor<CDataType> c_host_ref(c_lengths, c_strides);
reference_basic_gemm_softmax<ADataType, ADataType, AccDataType, CDataType>(
a_host, b_host, c_host_ref);
c_buf.FromDevice(c_host_dev.mData.data());
pass &= ck_tile::check_err(c_host_dev, c_host_ref);
std::cout << "valid:" << (pass ? "y" : "n") << std::endl;
// ck_tile::HostTensor<CDataType> c_host_ref(c_lengths, c_strides);
// reference_basic_gemm_softmax<ADataType, ADataType, AccDataType, CDataType>(
// a_host, b_host, c_host_ref);
// c_buf.FromDevice(c_host_dev.mData.data());
// pass &= ck_tile::check_err(c_host_dev, c_host_ref);
// std::cout << "valid:" << (pass ? "y" : "n") << std::endl;
ck_tile::HostTensor<WeightType> value_ref(out_lengths, out_strides);
ck_tile::HostTensor<IndexType> index_ref(out_lengths, out_strides);
reference_basic_gemm_softmax_topk<ADataType, ADataType, AccDataType, WeightType, IndexType>(
a_host, b_host, value_ref, index_ref, topk);
value_buf.FromDevice(value_host_dev.mData.data());
index_buf.FromDevice(index_host_dev.mData.data());
// pass &= ck_tile::check_err(c_host_dev, c_host_ref);
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);
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);
}
std::cout << "valid:" << (rtn ? "y" : "n") << std::endl;
}
std::size_t flop = std::size_t(2) * M * N * K;
@@ -198,5 +285,6 @@ int main(int argc, char* argv[])
std::cout << "Perf: " << ave_time << " ms, " << tflops << " TFlops, " << gb_per_sec << " GB/s"
<< std::endl;
return !pass;
// return !pass;
return rtn;
}

62
example/ck_tile/39_gemm_softmax/grid_gemm.hpp
View File

@@ -11,6 +11,8 @@ struct GridGemm
using ADataType = typename Problem::ADataType;
using BDataType = typename Problem::BDataType;
using CDataType = typename Problem::CDataType;
using WeightType = typename Problem::WeightType;
using IndexType = typename Problem::IndexType;
using AccDataType = typename Problem::AccDataType;
using ComputeDataType = float;
using CElementFunction = typename Problem::CElementFunction;
@@ -18,15 +20,17 @@ struct GridGemm
static constexpr auto kMPerBlock = Policy::kMPerBlock;
static constexpr auto kNPerBlock = Policy::kNPerBlock;
static constexpr auto kKPerBlock = Policy::kKPerBlock;
static constexpr auto topk = Policy::kTopKPerBlock;
template <typename AGridTensorView, typename BGridTensorView, typename CGridTensorView>
template <typename AGridTensorView, typename BGridTensorView, typename ValueGridTensorView, typename IndexGridTensorView>
CK_TILE_DEVICE void operator()(const AGridTensorView& a_grid,
const BGridTensorView& b_grid,
CGridTensorView& c_grid,
ValueGridTensorView& value_grid,
IndexGridTensorView& index_grid,
const CElementFunction& c_element_func) const
{
const auto M = a_grid.get_tensor_descriptor().get_length(number<0>{});
const auto N = c_grid.get_tensor_descriptor().get_length(number<1>{});
const auto N = b_grid.get_tensor_descriptor().get_length(number<0>{});
const auto K = a_grid.get_tensor_descriptor().get_length(number<1>{});
// divide problem
@@ -56,22 +60,52 @@ 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 C
// auto c_window = make_tile_window(
// c_grid, make_tuple(number<kMPerBlock>{}, number<kNPerBlock>{}), {iM, iN});
// block_gemm_pipeline(a_block_window, b_block_window, c_window, K / kKPerBlock, p_smem_char, c_element_func);
// store value and index
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<2, 4, 16>, sequence<4, 4>>,
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(
tile_distribution_encoding<sequence<1>,
tuple<sequence<2, 4, 16>, sequence<4, 4>>,
tuple<sequence<1>, sequence<1, 2>>,
tuple<sequence<1>, sequence<2, 0>>,
sequence<1, 2>,
sequence<0, 1>>{}));
const auto acc_block_tile =
block_gemm_pipeline(a_block_window, b_block_window, K / kKPerBlock, p_smem_char);
using ValueBlockTileDistr = decltype(value_window.get_tile_distribution());
using IndexBlockTileDistr = decltype(index_window.get_tile_distribution());
// cast to CDataType and apply CElementFunction
const auto c_block_tile = tile_elementwise_in(
[&](const auto& acc) { return c_element_func(type_convert<CDataType>(acc)); },
acc_block_tile);
using ValueBlockTile = decltype(make_static_distributed_tensor<WeightType>(ValueBlockTileDistr{}));
using IndexBlockTile = decltype(make_static_distributed_tensor<IndexType>(IndexBlockTileDistr{}));
ValueBlockTile value_block_tile;
IndexBlockTile index_block_tile;
store_tile(c_window, c_block_tile);
block_gemm_pipeline(a_block_window, b_block_window, value_block_tile, index_block_tile, K / kKPerBlock, p_smem_char);
// 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);
}
};

12
example/ck_tile/39_gemm_softmax/reference_gemm.hpp
View File

@@ -6,13 +6,17 @@
#include "ck_tile/core.hpp"
#include "ck_tile/host/host_tensor.hpp"
template <typename ADataType, typename BDataType, typename AccDataType, typename CDataType>
void reference_basic_gemm_softmax(const ck_tile::HostTensor<ADataType>& a_m_k,
template <typename ADataType, typename BDataType, typename AccDataType, typename WeightType, typename IndexType>
void reference_basic_gemm_softmax_topk(const ck_tile::HostTensor<ADataType>& a_m_k,
const ck_tile::HostTensor<BDataType>& b_n_k,
ck_tile::HostTensor<CDataType>& c_m_n)
ck_tile::HostTensor<WeightType>& y_values,
ck_tile::HostTensor<IndexType>& y_indices,
ck_tile::index_t topk)
{
const int M = a_m_k.mDesc.get_lengths()[0];
const int N = b_n_k.mDesc.get_lengths()[0];
const int K = b_n_k.mDesc.get_lengths()[1];
ck_tile::HostTensor<AccDataType> c_m_n({M, N}, {N, 1});
auto f = [&](auto m) {
for(int n = 0; n < N; ++n)
@@ -62,4 +66,6 @@ void reference_basic_gemm_softmax(const ck_tile::HostTensor<ADataType>& a_m_k,
ck_tile::make_ParallelTensorFunctor(f, c_m_n.mDesc.get_lengths()[0])(
std::thread::hardware_concurrency());
reference_topk(c_m_n, y_values, y_indices, topk);
}