mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-07-15 03:30:11 +00:00
Add codegen test example
This commit is contained in:
Clement Lin
committed by
Philip Maybank
Philip Maybank
parent
5b1c397806
commit
0cc5130818
19
example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/CMakeLists.txt
Executable file
19
example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/CMakeLists.txt
Executable file
@@ -0,0 +1,19 @@
|
||||
set(EXAMPLE_REDUCE "codegen_basic_flash_attention_fwd")
|
||||
# not using add_example_executable() to add this target, since we don't want this to have
|
||||
# to be included in "make all/install/check"
|
||||
message("adding example ${EXAMPLE_REDUCE}")
|
||||
|
||||
add_executable(${EXAMPLE_REDUCE} EXCLUDE_FROM_ALL flash_attention_fwd.cpp)
|
||||
target_include_directories(${EXAMPLE_REDUCE} PRIVATE ${CMAKE_CURRENT_LIST_DIR})
|
||||
set(EXAMPLE_REDUCE_COMPILE_OPTIONS)
|
||||
|
||||
# NOTE: we turn off undefined-func-template to let source compile without explicit declare function specializations
|
||||
list(APPEND EXAMPLE_REDUCE_COMPILE_OPTIONS -Wno-undefined-func-template -Wno-float-equal)
|
||||
|
||||
target_compile_options(${EXAMPLE_REDUCE} PRIVATE ${EXAMPLE_REDUCE_COMPILE_OPTIONS})
|
||||
|
||||
# TODO: we have to turn off this global prop, otherwise the progress bar generated
|
||||
# by cmake will print too many files, execvp: /bin/sh: Argument list too long
|
||||
# however, this property may affect global
|
||||
# TODO: consider codegen a makefile by us
|
||||
set_property(GLOBAL PROPERTY RULE_MESSAGES OFF)
|
||||
@@ -0,0 +1,24 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.
|
||||
|
||||
#pragma once
|
||||
|
||||
namespace ck_tile {
|
||||
|
||||
// Problem Description for BlockGemmARegBSmemCReg
|
||||
template <typename ADataType_,
|
||||
typename BDataType_,
|
||||
typename CDataType_,
|
||||
index_t kBlockSize_,
|
||||
typename BlockGemmShape_>
|
||||
struct BlockGemmARegBSmemCRegProblem
|
||||
{
|
||||
using ADataType = remove_cvref_t<ADataType_>;
|
||||
using BDataType = remove_cvref_t<BDataType_>;
|
||||
using CDataType = remove_cvref_t<CDataType_>;
|
||||
using BlockGemmShape = remove_cvref_t<BlockGemmShape_>;
|
||||
|
||||
static constexpr index_t kBlockSize = kBlockSize_;
|
||||
};
|
||||
|
||||
} // namespace ck_tile
|
||||
@@ -0,0 +1,300 @@
|
||||
// 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/core/tensor/tile_distribution.hpp"
|
||||
|
||||
#include "block_gemm_areg_bsmem_creg_problem.hpp"
|
||||
#include "block_gemm_areg_bsmem_creg_v1_default_policy.hpp"
|
||||
#include "block_gemm_areg_bsmem_creg_v1_iteratek_policy.hpp"
|
||||
|
||||
namespace ck_tile {
|
||||
|
||||
// A is block distributed tensor
|
||||
// B is block window on shared memory
|
||||
// C is block distributed tensor
|
||||
template <typename Problem, typename Policy = BlockGemmARegBSmemCRegV1DefaultPolicy>
|
||||
struct BlockGemmARegBSmemCRegV1
|
||||
{
|
||||
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 BlockGemmShape = remove_cvref_t<typename Problem::BlockGemmShape>;
|
||||
using BlockGemmPolicy = Policy;
|
||||
|
||||
static constexpr index_t kBlockSize = Problem::kBlockSize;
|
||||
|
||||
// C += A * B
|
||||
template <typename CBlockTensor, typename ABlockTensorTmp, typename BBlockWindowTmp>
|
||||
__device__ void operator()(CBlockTensor& c_block_tensor,
|
||||
const ABlockTensorTmp& a_block_tensor_tmp,
|
||||
const BBlockWindowTmp& b_block_window_tmp) const
|
||||
{
|
||||
static_assert(std::is_same_v<ADataType, remove_cv_t<typename ABlockTensorTmp::DataType>> &&
|
||||
std::is_same_v<BDataType, remove_cv_t<typename BBlockWindowTmp::DataType>> &&
|
||||
std::is_same_v<CDataType, remove_cv_t<typename CBlockTensor::DataType>>,
|
||||
"wrong!");
|
||||
|
||||
constexpr index_t MPerBlock = ABlockTensorTmp{}.get_lengths()[number<0>{}];
|
||||
constexpr index_t NPerBlock = BBlockWindowTmp{}.get_window_lengths()[number<0>{}];
|
||||
constexpr index_t KPerBlock = ABlockTensorTmp{}.get_lengths()[number<1>{}];
|
||||
|
||||
static_assert(MPerBlock == BlockGemmShape::kM && NPerBlock == BlockGemmShape::kN &&
|
||||
KPerBlock == BlockGemmShape::kK, "wrong!");
|
||||
|
||||
constexpr auto config = Policy::template GetWarpGemmMWarpNWarp<Problem>();
|
||||
|
||||
using WG = remove_cvref_t<decltype(config.template get<0>())>;
|
||||
|
||||
constexpr index_t MWarp = config.template get<1>();
|
||||
constexpr index_t NWarp = config.template get<2>();
|
||||
|
||||
constexpr index_t MIterPerWarp = MPerBlock / (MWarp * WG::kM);
|
||||
constexpr index_t NIterPerWarp = NPerBlock / (NWarp * WG::kN);
|
||||
constexpr index_t KIterPerWarp = KPerBlock / WG::kK;
|
||||
|
||||
constexpr index_t NPerBlockPerIter = NPerBlock / NIterPerWarp;
|
||||
constexpr index_t KPerBlockPerIter = KPerBlock / KIterPerWarp;
|
||||
|
||||
const index_t iNWarp = get_warp_id() % NWarp;
|
||||
|
||||
constexpr auto a_block_outer_dstr_encoding = tile_distribution_encoding<
|
||||
sequence<NWarp>,
|
||||
tuple<sequence<MIterPerWarp, MWarp>, sequence<KIterPerWarp>>,
|
||||
tuple<sequence<1, 0>>,
|
||||
tuple<sequence<1, 0>>,
|
||||
sequence<1, 2>,
|
||||
sequence<0, 0>>{};
|
||||
|
||||
constexpr auto c_block_outer_dstr_encoding = tile_distribution_encoding<
|
||||
sequence<>,
|
||||
tuple<sequence<MIterPerWarp, MWarp>, sequence<NIterPerWarp, NWarp>>,
|
||||
tuple<sequence<1, 2>>,
|
||||
tuple<sequence<1, 1>>,
|
||||
sequence<1, 2>,
|
||||
sequence<0, 0>>{};
|
||||
|
||||
constexpr auto a_block_dstr_encode = detail::make_embed_tile_distribution_encoding(
|
||||
a_block_outer_dstr_encoding, typename WG::AWarpDstrEncoding{});
|
||||
|
||||
constexpr auto c_block_dstr_encode = detail::make_embed_tile_distribution_encoding(
|
||||
c_block_outer_dstr_encoding, typename WG::CWarpDstrEncoding{});
|
||||
|
||||
constexpr auto a_block_dstr = make_static_tile_distribution(a_block_dstr_encode);
|
||||
|
||||
// constrcut from A-block-tensor from A-Block-tensor-tmp
|
||||
// FIXME: need method to check a_block_tensor and a_block_tensor_tmp have equivalent
|
||||
// distribution
|
||||
auto a_block_tensor =
|
||||
make_static_distributed_tensor<typename ABlockTensorTmp::DataType>(a_block_dstr);
|
||||
|
||||
a_block_tensor.get_thread_buffer() = a_block_tensor_tmp.get_thread_buffer();
|
||||
|
||||
// construct B-warp-window
|
||||
auto b_warp_window_tmp = make_tile_window(
|
||||
b_block_window_tmp.get_bottom_tensor_view(),
|
||||
make_tuple(number<WG::kN>{}, number<WG::kK>{}),
|
||||
{b_block_window_tmp.get_window_origin().at(number<0>{}) + iNWarp * WG::kN, b_block_window_tmp.get_window_origin().at(number<1>{})},
|
||||
make_static_tile_distribution(typename WG::BWarpDstrEncoding{}));
|
||||
|
||||
statically_indexed_array<statically_indexed_array<decltype(b_warp_window_tmp), KIterPerWarp>,
|
||||
NIterPerWarp> b_warp_windows;
|
||||
|
||||
static_for<0, NIterPerWarp, 1>{}([&](auto nIter) {
|
||||
static_for<0, KIterPerWarp, 1>{}([&](auto kIter) {
|
||||
b_warp_windows(nIter)(kIter) = b_warp_window_tmp;
|
||||
|
||||
move_tile_window(b_warp_windows(nIter)(kIter),
|
||||
{nIter * NPerBlockPerIter, kIter * KPerBlockPerIter});
|
||||
});
|
||||
});
|
||||
|
||||
// check C-block-distribution
|
||||
static_assert(std::is_same_v<remove_cvref_t<decltype(c_block_dstr_encode)>,
|
||||
remove_cvref_t<decltype(CBlockTensor::get_tile_distribution()
|
||||
.get_static_tile_distribution_encoding())>>, "wrong!");
|
||||
|
||||
using AWarpDstr = typename WG::AWarpDstr;
|
||||
using CWarpDstr = typename WG::CWarpDstr;
|
||||
|
||||
using AWarpTensor = typename WG::AWarpTensor;
|
||||
using CWarpTensor = typename WG::CWarpTensor;
|
||||
|
||||
constexpr auto a_warp_y_lengths = to_sequence(AWarpDstr{}.get_ys_to_d_descriptor().get_lengths());
|
||||
constexpr auto c_warp_y_lengths = to_sequence(CWarpDstr{}.get_ys_to_d_descriptor().get_lengths());
|
||||
|
||||
constexpr auto a_warp_y_index_zeros = uniform_sequence_gen_t<AWarpDstr::NDimY, 0>{};
|
||||
constexpr auto c_warp_y_index_zeros = uniform_sequence_gen_t<CWarpDstr::NDimY, 0>{};
|
||||
|
||||
// hot loop:
|
||||
static_for<0, KIterPerWarp, 1>{}([&](auto kIter) {
|
||||
static_for<0, MIterPerWarp, 1>{}([&](auto mIter) {
|
||||
// read A warp tensor from A block tensor
|
||||
AWarpTensor a_warp_tensor;
|
||||
|
||||
a_warp_tensor.get_thread_buffer() = a_block_tensor.get_y_sliced_thread_data(
|
||||
merge_sequences(sequence<mIter, kIter>{}, a_warp_y_index_zeros),
|
||||
merge_sequences(sequence<1, 1>{}, a_warp_y_lengths));
|
||||
|
||||
static_for<0, NIterPerWarp, 1>{}([&](auto nIter) {
|
||||
// read B warp tensor from B Block window
|
||||
const auto b_warp_tensor = load_tile(b_warp_windows(nIter)(kIter));
|
||||
// read C warp tensor from C block tensor
|
||||
CWarpTensor c_warp_tensor;
|
||||
|
||||
c_warp_tensor.get_thread_buffer() = c_block_tensor.get_y_sliced_thread_data(
|
||||
merge_sequences(sequence<mIter, nIter>{}, c_warp_y_index_zeros),
|
||||
merge_sequences(sequence<1, 1>{}, c_warp_y_lengths));
|
||||
|
||||
// warp GEMM
|
||||
WG{}(c_warp_tensor, a_warp_tensor, b_warp_tensor);
|
||||
|
||||
// write C warp tensor into C block tensor
|
||||
c_block_tensor.set_y_sliced_thread_data(
|
||||
merge_sequences(sequence<mIter, nIter>{}, c_warp_y_index_zeros),
|
||||
merge_sequences(sequence<1, 1>{}, c_warp_y_lengths),
|
||||
c_warp_tensor.get_thread_buffer());
|
||||
});
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
// C = A * B
|
||||
template <typename ABlockTensorTmp, typename BBlockWindowTmp>
|
||||
__device__ auto operator()(const ABlockTensorTmp& a_block_tensor_tmp,
|
||||
const BBlockWindowTmp& b_block_window_tmp) const
|
||||
{
|
||||
static_assert(std::is_same_v<ADataType, remove_cv_t<typename ABlockTensorTmp::DataType>> &&
|
||||
std::is_same_v<BDataType, remove_cv_t<typename BBlockWindowTmp::DataType>>,
|
||||
"wrong!");
|
||||
|
||||
constexpr index_t MPerBlock = ABlockTensorTmp{}.get_lengths()[number<0>{}];
|
||||
constexpr index_t NPerBlock = BBlockWindowTmp{}.get_window_lengths()[number<0>{}];
|
||||
constexpr index_t KPerBlock = ABlockTensorTmp{}.get_lengths()[number<1>{}];
|
||||
|
||||
static_assert(MPerBlock == BlockGemmShape::kM && NPerBlock == BlockGemmShape::kN &&
|
||||
KPerBlock == BlockGemmShape::kK, "wrong!");
|
||||
|
||||
constexpr auto config = Policy::template GetWarpGemmMWarpNWarp<Problem>();
|
||||
|
||||
using WG = remove_cvref_t<decltype(config.template get<0>())>;
|
||||
|
||||
constexpr index_t MWarp = config.template get<1>();
|
||||
constexpr index_t NWarp = config.template get<2>();
|
||||
|
||||
constexpr index_t MIterPerWarp = MPerBlock / (MWarp * WG::kM);
|
||||
constexpr index_t NIterPerWarp = NPerBlock / (NWarp * WG::kN);
|
||||
constexpr index_t KIterPerWarp = KPerBlock / WG::kK;
|
||||
|
||||
constexpr index_t NPerBlockPerIter = NPerBlock / NIterPerWarp;
|
||||
constexpr index_t KPerBlockPerIter = KPerBlock / KIterPerWarp;
|
||||
|
||||
const index_t iNWarp = get_warp_id() % NWarp;
|
||||
|
||||
constexpr auto a_block_outer_dstr_encoding = tile_distribution_encoding<
|
||||
sequence<NWarp>,
|
||||
tuple<sequence<MIterPerWarp, MWarp>, sequence<KIterPerWarp>>,
|
||||
tuple<sequence<1, 0>>,
|
||||
tuple<sequence<1, 0>>,
|
||||
sequence<1, 2>,
|
||||
sequence<0, 0>>{};
|
||||
|
||||
constexpr auto c_block_outer_dstr_encoding = tile_distribution_encoding<
|
||||
sequence<>,
|
||||
tuple<sequence<MIterPerWarp, MWarp>, sequence<NIterPerWarp, NWarp>>,
|
||||
tuple<sequence<1, 2>>,
|
||||
tuple<sequence<1, 1>>,
|
||||
sequence<1, 2>,
|
||||
sequence<0, 0>>{};
|
||||
|
||||
constexpr auto a_block_dstr_encode = detail::make_embed_tile_distribution_encoding(
|
||||
a_block_outer_dstr_encoding, typename WG::AWarpDstrEncoding{});
|
||||
|
||||
constexpr auto c_block_dstr_encode = detail::make_embed_tile_distribution_encoding(
|
||||
c_block_outer_dstr_encoding, typename WG::CWarpDstrEncoding{});
|
||||
|
||||
constexpr auto a_block_dstr = make_static_tile_distribution(a_block_dstr_encode);
|
||||
constexpr auto c_block_dstr = make_static_tile_distribution(c_block_dstr_encode);
|
||||
|
||||
// constrcut from A-block-tensor from A-Block-tensor-tmp
|
||||
// FIXME: need method to check a_block_tensor and a_block_tensor_tmp have equivalent
|
||||
// distribution
|
||||
auto a_block_tensor =
|
||||
make_static_distributed_tensor<typename ABlockTensorTmp::DataType>(a_block_dstr);
|
||||
|
||||
a_block_tensor.get_thread_buffer() = a_block_tensor_tmp.get_thread_buffer();
|
||||
|
||||
// construct B-warp-window
|
||||
auto b_warp_window_tmp = make_tile_window(
|
||||
b_block_window_tmp.get_bottom_tensor_view(),
|
||||
make_tuple(number<WG::kN>{}, number<WG::kK>{}),
|
||||
{b_block_window_tmp.get_window_origin().at(number<0>{}) + iNWarp * WG::kN, b_block_window_tmp.get_window_origin().at(number<1>{})},
|
||||
make_static_tile_distribution(typename WG::BWarpDstrEncoding{}));
|
||||
|
||||
statically_indexed_array<statically_indexed_array<decltype(b_warp_window_tmp), KIterPerWarp>,
|
||||
NIterPerWarp> b_warp_windows;
|
||||
|
||||
static_for<0, NIterPerWarp, 1>{}([&](auto nIter) {
|
||||
static_for<0, KIterPerWarp, 1>{}([&](auto kIter) {
|
||||
b_warp_windows(nIter)(kIter) = b_warp_window_tmp;
|
||||
|
||||
move_tile_window(b_warp_windows(nIter)(kIter),
|
||||
{nIter * NPerBlockPerIter, kIter * KPerBlockPerIter});
|
||||
});
|
||||
});
|
||||
|
||||
// Construct C-Block-Tensor
|
||||
auto c_block_tensor = make_static_distributed_tensor<CDataType>(c_block_dstr);
|
||||
|
||||
using AWarpDstr = typename WG::AWarpDstr;
|
||||
using CWarpDstr = typename WG::CWarpDstr;
|
||||
|
||||
using AWarpTensor = typename WG::AWarpTensor;
|
||||
using CWarpTensor = typename WG::CWarpTensor;
|
||||
|
||||
constexpr auto a_warp_y_lengths = to_sequence(AWarpDstr{}.get_ys_to_d_descriptor().get_lengths());
|
||||
constexpr auto c_warp_y_lengths = to_sequence(CWarpDstr{}.get_ys_to_d_descriptor().get_lengths());
|
||||
|
||||
constexpr auto a_warp_y_index_zeros = uniform_sequence_gen_t<AWarpDstr::NDimY, 0>{};
|
||||
constexpr auto c_warp_y_index_zeros = uniform_sequence_gen_t<CWarpDstr::NDimY, 0>{};
|
||||
|
||||
// hot loop:
|
||||
static_for<0, KIterPerWarp, 1>{}([&](auto kIter) {
|
||||
static_for<0, MIterPerWarp, 1>{}([&](auto mIter) {
|
||||
// read A warp tensor from A block tensor
|
||||
AWarpTensor a_warp_tensor;
|
||||
|
||||
a_warp_tensor.get_thread_buffer() = a_block_tensor.get_y_sliced_thread_data(
|
||||
merge_sequences(sequence<mIter, kIter>{}, a_warp_y_index_zeros),
|
||||
merge_sequences(sequence<1, 1>{}, a_warp_y_lengths));
|
||||
|
||||
static_for<0, NIterPerWarp, 1>{}([&](auto nIter) {
|
||||
// read B warp tensor from B Block window
|
||||
const auto b_warp_tensor = load_tile(b_warp_windows(nIter)(kIter));
|
||||
|
||||
// read C warp tensor from C block tensor
|
||||
CWarpTensor c_warp_tensor;
|
||||
|
||||
c_warp_tensor.get_thread_buffer() = c_block_tensor.get_y_sliced_thread_data(
|
||||
merge_sequences(sequence<mIter, nIter>{}, c_warp_y_index_zeros),
|
||||
merge_sequences(sequence<1, 1>{}, c_warp_y_lengths));
|
||||
|
||||
// warp GEMM
|
||||
WG{}(c_warp_tensor, a_warp_tensor, b_warp_tensor);
|
||||
|
||||
// write C warp tensor into C block tensor
|
||||
c_block_tensor.set_y_sliced_thread_data(
|
||||
merge_sequences(sequence<mIter, nIter>{}, c_warp_y_index_zeros),
|
||||
merge_sequences(sequence<1, 1>{}, c_warp_y_lengths),
|
||||
c_warp_tensor.get_thread_buffer());
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
return c_block_tensor;
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace ck_tile
|
||||
@@ -0,0 +1,20 @@
|
||||
// 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/core/tensor/tile_distribution.hpp"
|
||||
|
||||
namespace ck_tile {
|
||||
|
||||
struct BlockGemmARegBSmemCRegV1DefaultPolicy
|
||||
{
|
||||
template <typename Problem>
|
||||
CK_TILE_HOST_DEVICE static constexpr auto GetWarpGemmMWarpNWarp()
|
||||
{
|
||||
return make_tuple(WarpGemmMfmaF16F16F32M32N32K8TransposedCDistribution{}, 4, 1);
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace ck_tile
|
||||
@@ -0,0 +1,20 @@
|
||||
// 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/core/tensor/tile_distribution.hpp"
|
||||
|
||||
namespace ck_tile {
|
||||
|
||||
struct BlockGemmARegBSmemCRegV1K8Policy
|
||||
{
|
||||
template <typename Problem>
|
||||
CK_TILE_HOST_DEVICE static constexpr auto GetWarpGemmMWarpNWarp()
|
||||
{
|
||||
return make_tuple(WarpGemmMfmaF16F16F32M32N32K16TransposedCDistribution{}, 4, 1);
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace ck_tile
|
||||
@@ -0,0 +1,545 @@
|
||||
// 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/core/tensor/tile_distribution.hpp"
|
||||
|
||||
#include "block_gemm_pipeline_agmem_bgmem_creg_v2_askiplds_policy.hpp"
|
||||
|
||||
namespace ck_tile {
|
||||
|
||||
// A Tile Window: global memory
|
||||
// B Tile Window: global memory
|
||||
// C Distributed tensor: register
|
||||
template <typename Problem>
|
||||
struct BlockGemmPipelineAGmemBGmemCReg<Problem, BlockGemmPipelineAGmemBGmemCRegSkipALdsPolicy>
|
||||
{
|
||||
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 BlockGemmShape = remove_cvref_t<typename Problem::BlockGemmShape>;
|
||||
using Policy = BlockGemmPipelineAGmemBGmemCRegSkipALdsPolicy;
|
||||
|
||||
static constexpr index_t kBlockSize = Problem::kBlockSize;
|
||||
|
||||
static constexpr index_t kMPerBlock = BlockGemmShape::kM;
|
||||
static constexpr index_t kNPerBlock = BlockGemmShape::kN;
|
||||
static constexpr index_t kKPerBlock = BlockGemmShape::kK;
|
||||
|
||||
// Move this part into Policy?
|
||||
__host__ __device__ static constexpr index_t GetStaticLdsSize()
|
||||
{
|
||||
return sizeof(BDataType) *
|
||||
Policy::template MakeBLdsBlockDescriptor<Problem>().get_element_space_size();
|
||||
}
|
||||
|
||||
template <typename ADramBlockWindowTmp,
|
||||
typename BDramBlockWindowTmp,
|
||||
typename AElementFunction,
|
||||
typename BElementFunction>
|
||||
__host__ __device__ auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp,
|
||||
const AElementFunction& a_element_func,
|
||||
const BDramBlockWindowTmp& b_dram_block_window_tmp,
|
||||
const BElementFunction& b_element_func,
|
||||
index_t num_loop,
|
||||
void* p_smem) const
|
||||
{
|
||||
static_assert(
|
||||
std::is_same_v<ADataType, remove_cvref_t<typename ADramBlockWindowTmp::DataType>> &&
|
||||
std::is_same_v<BDataType, remove_cvref_t<typename BDramBlockWindowTmp::DataType>>,
|
||||
"wrong!");
|
||||
|
||||
static_assert(kMPerBlock == ADramBlockWindowTmp{}.get_window_lengths()[number<0>{}] &&
|
||||
kNPerBlock == BDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] &&
|
||||
kKPerBlock == ADramBlockWindowTmp{}.get_window_lengths()[number<1>{}],
|
||||
"wrong!");
|
||||
|
||||
// A tile in Reg，blockTensor
|
||||
// This tensor distribution used to construct both distributed tensor for local buffer store
|
||||
// and read. without buffer address info
|
||||
constexpr auto a_reg_block_dstr = Policy::template MakeARegBlockDescriptor<Problem>();
|
||||
|
||||
// B tile in LDS, blockWindow
|
||||
BDataType* p_b_lds =
|
||||
static_cast<BDataType*>(static_cast<void*>(static_cast<char*>(p_smem)));
|
||||
|
||||
constexpr auto b_lds_block_desc = Policy::template MakeBLdsBlockDescriptor<Problem>();
|
||||
|
||||
// This tensor view used to construct both tile window for lds store and read, with buffer
|
||||
// address info
|
||||
auto b_lds_block = make_tensor_view<address_space_enum::lds>(p_b_lds, b_lds_block_desc);
|
||||
|
||||
// A DRAM tile window for load
|
||||
auto a_copy_dram_window =
|
||||
make_tile_window(a_dram_block_window_tmp.get_bottom_tensor_view(),
|
||||
make_tuple(number<kMPerBlock>{}, number<kKPerBlock>{}),
|
||||
a_dram_block_window_tmp.get_window_origin(),
|
||||
Policy::template MakeADramTileDistribution<Problem>());
|
||||
|
||||
// A Reg tensor for store, also used for block GEMM
|
||||
auto a_copy_reg_tensor = make_static_distributed_tensor<ADataType>(a_reg_block_dstr);
|
||||
|
||||
// B DRAM tile window for load
|
||||
auto b_copy_dram_window =
|
||||
make_tile_window(b_dram_block_window_tmp.get_bottom_tensor_view(),
|
||||
make_tuple(number<kNPerBlock>{}, number<kKPerBlock>{}),
|
||||
b_dram_block_window_tmp.get_window_origin(),
|
||||
Policy::template MakeBDramTileDistribution<Problem>());
|
||||
|
||||
// B LDS tile window for store
|
||||
auto b_copy_lds_window =
|
||||
make_tile_window(b_lds_block,
|
||||
make_tuple(number<kNPerBlock>{}, number<kKPerBlock>{}),
|
||||
{0, 0},
|
||||
b_copy_dram_window.get_tile_distribution());
|
||||
|
||||
// B LDS tile for block GEMM
|
||||
auto b_lds_gemm_window = make_tile_window(
|
||||
b_lds_block, make_tuple(number<kNPerBlock>{}, number<kKPerBlock>{}), {0, 0});
|
||||
|
||||
// Block GEMM
|
||||
constexpr auto block_gemm = Policy::template GetBlockGemm<Problem>();
|
||||
|
||||
// Acc register tile
|
||||
auto c_block_tile = decltype(block_gemm(a_copy_reg_tensor, b_lds_gemm_window)){};
|
||||
|
||||
// prefetch
|
||||
// global read 0
|
||||
auto a_block_tile = load_tile(a_copy_dram_window);
|
||||
auto b_block_tile = load_tile(b_copy_dram_window);
|
||||
|
||||
{
|
||||
// move to 1
|
||||
move_tile_window(a_copy_dram_window, {0, kKPerBlock});
|
||||
move_tile_window(b_copy_dram_window, {0, kKPerBlock});
|
||||
|
||||
// Initialize C
|
||||
tile_elementwise_inout([](auto& c) { c = 0; }, c_block_tile);
|
||||
|
||||
// block buffer write 0
|
||||
const auto a_block_tile_tmp = tile_elementwise_in(a_element_func, a_block_tile);
|
||||
// store_tile -> shuffle store tile
|
||||
store_tile(a_copy_reg_tensor, a_block_tile_tmp);
|
||||
// global read 1
|
||||
a_block_tile = load_tile(a_copy_dram_window);
|
||||
|
||||
// LDS write 0
|
||||
const auto b_block_tile_tmp = tile_elementwise_in(b_element_func, b_block_tile);
|
||||
store_tile(b_copy_lds_window, b_block_tile_tmp);
|
||||
// global read 1
|
||||
b_block_tile = load_tile(b_copy_dram_window);
|
||||
}
|
||||
|
||||
index_t iCounter = num_loop - 2;
|
||||
|
||||
do
|
||||
{
|
||||
block_sync_lds();
|
||||
|
||||
// GEMM i
|
||||
block_gemm(c_block_tile, a_copy_reg_tensor, b_lds_gemm_window);
|
||||
|
||||
block_sync_lds();
|
||||
|
||||
// move to i + 2
|
||||
move_tile_window(a_copy_dram_window, {0, kKPerBlock});
|
||||
move_tile_window(b_copy_dram_window, {0, kKPerBlock});
|
||||
|
||||
// LDS write i + 1
|
||||
const auto a_block_tile_tmp = tile_elementwise_in(a_element_func, a_block_tile);
|
||||
store_tile(a_copy_reg_tensor, a_block_tile_tmp);
|
||||
// global read i + 2
|
||||
a_block_tile = load_tile(a_copy_dram_window);
|
||||
|
||||
// LDS write i + 1
|
||||
const auto b_block_tile_tmp = tile_elementwise_in(b_element_func, b_block_tile);
|
||||
store_tile(b_copy_lds_window, b_block_tile_tmp);
|
||||
// global read i + 2
|
||||
b_block_tile = load_tile(b_copy_dram_window);
|
||||
|
||||
iCounter--;
|
||||
|
||||
} while(iCounter > 0);
|
||||
|
||||
// tail
|
||||
{
|
||||
block_sync_lds();
|
||||
|
||||
// GEMM num_loop - 2
|
||||
block_gemm(c_block_tile, a_copy_reg_tensor, b_lds_gemm_window);
|
||||
|
||||
block_sync_lds();
|
||||
|
||||
// LDS write num_loop - 1
|
||||
const auto a_block_tile_tmp = tile_elementwise_in(a_element_func, a_block_tile);
|
||||
store_tile(a_copy_reg_tensor, a_block_tile_tmp);
|
||||
|
||||
const auto b_block_tile_tmp = tile_elementwise_in(b_element_func, b_block_tile);
|
||||
store_tile(b_copy_lds_window, b_block_tile_tmp);
|
||||
|
||||
block_sync_lds();
|
||||
|
||||
// GEMM num_loop - 1
|
||||
block_gemm(c_block_tile, a_copy_reg_tensor, b_lds_gemm_window);
|
||||
}
|
||||
|
||||
return c_block_tile;
|
||||
}
|
||||
|
||||
template <typename ADramBlockWindowTmp, typename BDramBlockWindowTmp>
|
||||
__device__ auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp,
|
||||
const BDramBlockWindowTmp& b_dram_block_window_tmp,
|
||||
index_t num_loop,
|
||||
void* p_smem) const
|
||||
{
|
||||
return operator()(
|
||||
a_dram_block_window_tmp,
|
||||
[](const ADataType& a) { return a; },
|
||||
b_dram_block_window_tmp,
|
||||
[](const BDataType& b) { return b; },
|
||||
num_loop,
|
||||
p_smem);
|
||||
}
|
||||
};
|
||||
|
||||
// A Tile Window: global memory
|
||||
// B Tile Window: global memory
|
||||
// C Distributed tensor: register
|
||||
template <typename Problem, index_t kHeadDim>
|
||||
struct BlockGemmPipelineAGmemBGmemCReg<
|
||||
Problem,
|
||||
BlockGemmPipelineAGmemBGmemCRegSkipALdsPersistentQRegCachePolicy<kHeadDim>>
|
||||
{
|
||||
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 BlockGemmShape = remove_cvref_t<typename Problem::BlockGemmShape>;
|
||||
using Policy = BlockGemmPipelineAGmemBGmemCRegSkipALdsPersistentQRegCachePolicy<kHeadDim>;
|
||||
|
||||
static constexpr index_t kBlockSize = Problem::kBlockSize;
|
||||
|
||||
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 k_loops = Policy::AKDim / kKPerBlock;
|
||||
|
||||
// Move this part into Policy?
|
||||
__host__ __device__ static constexpr index_t GetStaticLdsSize()
|
||||
{
|
||||
return sizeof(BDataType) *
|
||||
Policy::template MakeBLdsBlockDescriptor<Problem>().get_element_space_size();
|
||||
}
|
||||
|
||||
// Cold A Register Cache
|
||||
template <typename ADramBlockWindowTmp,
|
||||
typename BDramBlockWindowTmp,
|
||||
typename AElementFunction,
|
||||
typename BElementFunction,
|
||||
typename ARegBlockTensorTmp>
|
||||
__host__ __device__ auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp,
|
||||
const AElementFunction& a_element_func,
|
||||
const BDramBlockWindowTmp& b_dram_block_window_tmp,
|
||||
const BElementFunction& b_element_func,
|
||||
ARegBlockTensorTmp& a_reg_block_tensor_tmp,
|
||||
void* p_smem) const
|
||||
{
|
||||
static_assert(
|
||||
std::is_same_v<ADataType, remove_cvref_t<typename ADramBlockWindowTmp::DataType>> &&
|
||||
std::is_same_v<BDataType, remove_cvref_t<typename BDramBlockWindowTmp::DataType>>,
|
||||
"wrong!");
|
||||
|
||||
static_assert(kMPerBlock == ADramBlockWindowTmp{}.get_window_lengths()[number<0>{}] &&
|
||||
kNPerBlock == BDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] &&
|
||||
kKPerBlock == ADramBlockWindowTmp{}.get_window_lengths()[number<1>{}],
|
||||
"wrong!");
|
||||
|
||||
ignore = a_element_func;
|
||||
ignore = b_element_func;
|
||||
|
||||
// A tile in Reg，blockTensor
|
||||
// This tensor distribution used to construct both distributed tensor for local buffer store
|
||||
// and read. without buffer address info
|
||||
constexpr auto a_reg_block_dstr = Policy::template MakeARegBlockDescriptor<Problem>();
|
||||
|
||||
// B tile in LDS, blockWindow
|
||||
BDataType* p_b_lds =
|
||||
static_cast<BDataType*>(static_cast<void*>(static_cast<char*>(p_smem)));
|
||||
|
||||
constexpr auto b_lds_block_desc = Policy::template MakeBLdsBlockDescriptor<Problem>();
|
||||
|
||||
// This tensor view used to construct both tile window for lds store and read, with buffer
|
||||
// address info
|
||||
auto b_lds_block = make_tensor_view<address_space_enum::lds>(p_b_lds, b_lds_block_desc);
|
||||
|
||||
// A DRAM tile window for load
|
||||
auto a_copy_dram_window =
|
||||
make_tile_window(a_dram_block_window_tmp.get_bottom_tensor_view(),
|
||||
make_tuple(number<kMPerBlock>{}, number<kKPerBlock>{}),
|
||||
a_dram_block_window_tmp.get_window_origin(),
|
||||
Policy::template MakeADramTileDistribution<Problem>());
|
||||
|
||||
// A Reg tensor for store, also used for block GEMM
|
||||
auto a_copy_reg_tensor = make_static_distributed_tensor<ADataType>(a_reg_block_dstr);
|
||||
|
||||
// B DRAM tile window for load
|
||||
auto b_copy_dram_window =
|
||||
make_tile_window(b_dram_block_window_tmp.get_bottom_tensor_view(),
|
||||
make_tuple(number<kNPerBlock>{}, number<kKPerBlock>{}),
|
||||
b_dram_block_window_tmp.get_window_origin(),
|
||||
Policy::template MakeBDramTileDistribution<Problem>());
|
||||
|
||||
// B LDS tile window for store
|
||||
auto b_copy_lds_window =
|
||||
make_tile_window(b_lds_block,
|
||||
make_tuple(number<kNPerBlock>{}, number<kKPerBlock>{}),
|
||||
{0, 0},
|
||||
b_copy_dram_window.get_tile_distribution());
|
||||
|
||||
// B LDS tile for block GEMM
|
||||
auto b_lds_gemm_window = make_tile_window(
|
||||
b_lds_block, make_tuple(number<kNPerBlock>{}, number<kKPerBlock>{}), {0, 0});
|
||||
|
||||
// Block GEMM
|
||||
constexpr auto block_gemm = Policy::template GetBlockGemm<Problem>();
|
||||
|
||||
// Acc register tile
|
||||
auto c_block_tile = decltype(block_gemm(
|
||||
get_slice_tile(a_copy_reg_tensor, sequence<0, 0>{}, sequence<kMPerBlock, kKPerBlock>{}),
|
||||
b_lds_gemm_window)){};
|
||||
|
||||
auto a_block_tile = load_tile(a_copy_dram_window);
|
||||
auto b_block_tile = load_tile(b_copy_dram_window);
|
||||
{
|
||||
move_tile_window(a_copy_dram_window, {0, kKPerBlock});
|
||||
move_tile_window(b_copy_dram_window, {0, kKPerBlock});
|
||||
|
||||
tile_elementwise_inout([](auto& c) { c = 0; }, c_block_tile);
|
||||
|
||||
set_slice_tile(a_copy_reg_tensor,
|
||||
a_block_tile,
|
||||
sequence<0, 0>{},
|
||||
sequence<kMPerBlock, kKPerBlock>{});
|
||||
a_block_tile = load_tile(a_copy_dram_window);
|
||||
|
||||
store_tile(b_copy_lds_window, b_block_tile);
|
||||
b_block_tile = load_tile(b_copy_dram_window);
|
||||
}
|
||||
if constexpr(k_loops > 2)
|
||||
{
|
||||
static_for<0, k_loops - 2, 1>{}([&](auto i_k0) {
|
||||
block_sync_lds();
|
||||
|
||||
block_gemm(c_block_tile,
|
||||
get_slice_tile(a_copy_reg_tensor,
|
||||
sequence<0, (i_k0)*kKPerBlock>{},
|
||||
sequence<kMPerBlock, (i_k0 + 1) * kKPerBlock>{}),
|
||||
b_copy_lds_window);
|
||||
|
||||
block_sync_lds();
|
||||
|
||||
move_tile_window(a_copy_dram_window, {0, kKPerBlock});
|
||||
move_tile_window(b_copy_dram_window, {0, kKPerBlock});
|
||||
|
||||
set_slice_tile(a_copy_reg_tensor,
|
||||
a_block_tile,
|
||||
sequence<0, (i_k0 + 1) * kKPerBlock>{},
|
||||
sequence<kMPerBlock, (i_k0 + 2) * kKPerBlock>{});
|
||||
a_block_tile = load_tile(a_copy_dram_window);
|
||||
|
||||
store_tile(b_copy_lds_window, b_block_tile);
|
||||
b_block_tile = load_tile(b_copy_dram_window);
|
||||
});
|
||||
}
|
||||
|
||||
// tail
|
||||
{
|
||||
block_sync_lds();
|
||||
|
||||
block_gemm(c_block_tile,
|
||||
get_slice_tile(a_copy_reg_tensor,
|
||||
sequence<0, (k_loops - 2) * kKPerBlock>{},
|
||||
sequence<kMPerBlock, (k_loops - 1) * kKPerBlock>{}),
|
||||
b_copy_lds_window);
|
||||
|
||||
block_sync_lds();
|
||||
|
||||
set_slice_tile(a_copy_reg_tensor,
|
||||
a_block_tile,
|
||||
sequence<0, (k_loops - 1) * kKPerBlock>{},
|
||||
sequence<kMPerBlock, k_loops * kKPerBlock>{});
|
||||
|
||||
store_tile(b_copy_lds_window, b_block_tile);
|
||||
|
||||
block_sync_lds();
|
||||
|
||||
block_gemm(c_block_tile,
|
||||
get_slice_tile(a_copy_reg_tensor,
|
||||
sequence<0, (k_loops - 1) * kKPerBlock>{},
|
||||
sequence<kMPerBlock, (k_loops)*kKPerBlock>{}),
|
||||
b_copy_lds_window);
|
||||
}
|
||||
|
||||
// store_tile(a_reg_block_tensor_tmp, a_copy_reg_tensor);
|
||||
set_slice_tile(a_reg_block_tensor_tmp,
|
||||
a_copy_reg_tensor,
|
||||
sequence<0, 0>{},
|
||||
sequence<kMPerBlock, k_loops * kKPerBlock>{});
|
||||
|
||||
return c_block_tile;
|
||||
}
|
||||
|
||||
// Hot A Register Cache
|
||||
template <typename BDramBlockWindowTmp, typename BElementFunction, typename ARegBlockTensorTmp>
|
||||
__host__ __device__ auto operator()(const BDramBlockWindowTmp& b_dram_block_window_tmp,
|
||||
const BElementFunction& b_element_func,
|
||||
const ARegBlockTensorTmp& a_reg_block_tensor_tmp,
|
||||
void* p_smem) const
|
||||
{
|
||||
static_assert(std::is_same_v<BDataType, remove_cvref_t<typename BDramBlockWindowTmp::DataType>>,
|
||||
"wrong!");
|
||||
|
||||
static_assert(kNPerBlock == BDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] &&
|
||||
kKPerBlock == BDramBlockWindowTmp{}.get_window_lengths()[number<1>{}],
|
||||
"wrong!");
|
||||
|
||||
ignore = b_element_func;
|
||||
|
||||
// A tile in Reg，blockTensor
|
||||
// This tensor distribution used to construct both distributed tensor for local buffer store
|
||||
// and read. without buffer address info
|
||||
constexpr auto a_reg_block_dstr = Policy::template MakeARegBlockDescriptor<Problem>();
|
||||
|
||||
// A Reg tensor for store, also used for block GEMM
|
||||
auto a_copy_reg_tensor = make_static_distributed_tensor<ADataType>(a_reg_block_dstr);
|
||||
// store_tile(a_copy_reg_tensor, a_reg_block_tensor_tmp);
|
||||
|
||||
set_slice_tile(a_copy_reg_tensor,
|
||||
a_reg_block_tensor_tmp,
|
||||
sequence<0, 0>{},
|
||||
sequence<kMPerBlock, k_loops * kKPerBlock>{});
|
||||
|
||||
// B tile in LDS, blockWindow
|
||||
BDataType* p_b_lds =
|
||||
static_cast<BDataType*>(static_cast<void*>(static_cast<char*>(p_smem)));
|
||||
|
||||
constexpr auto b_lds_block_desc = Policy::template MakeBLdsBlockDescriptor<Problem>();
|
||||
|
||||
// This tensor view used to construct both tile window for lds store and read, with buffer
|
||||
// address info
|
||||
auto b_lds_block = make_tensor_view<address_space_enum::lds>(p_b_lds, b_lds_block_desc);
|
||||
|
||||
// B DRAM tile window for load
|
||||
auto b_copy_dram_window =
|
||||
make_tile_window(b_dram_block_window_tmp.get_bottom_tensor_view(),
|
||||
make_tuple(number<kNPerBlock>{}, number<kKPerBlock>{}),
|
||||
b_dram_block_window_tmp.get_window_origin(),
|
||||
Policy::template MakeBDramTileDistribution<Problem>());
|
||||
|
||||
// B LDS tile window for store
|
||||
auto b_copy_lds_window =
|
||||
make_tile_window(b_lds_block,
|
||||
make_tuple(number<kNPerBlock>{}, number<kKPerBlock>{}),
|
||||
{0, 0},
|
||||
b_copy_dram_window.get_tile_distribution());
|
||||
|
||||
// B LDS tile for block GEMM
|
||||
auto b_lds_gemm_window = make_tile_window(
|
||||
b_lds_block, make_tuple(number<kNPerBlock>{}, number<kKPerBlock>{}), {0, 0});
|
||||
|
||||
// Block GEMM
|
||||
constexpr auto block_gemm = Policy::template GetBlockGemm<Problem>();
|
||||
|
||||
// Acc register tile
|
||||
auto c_block_tile = decltype(block_gemm(
|
||||
get_slice_tile(a_copy_reg_tensor, sequence<0, 0>{}, sequence<kMPerBlock, kKPerBlock>{}),
|
||||
b_lds_gemm_window)){};
|
||||
|
||||
auto b_block_tile = load_tile(b_copy_dram_window);
|
||||
{
|
||||
move_tile_window(b_copy_dram_window, {0, kKPerBlock});
|
||||
|
||||
tile_elementwise_inout([](auto& c) { c = 0; }, c_block_tile);
|
||||
|
||||
store_tile(b_copy_lds_window, b_block_tile);
|
||||
b_block_tile = load_tile(b_copy_dram_window);
|
||||
}
|
||||
if constexpr(k_loops > 2)
|
||||
{
|
||||
static_for<0, k_loops - 2, 1>{}([&](auto i_k0) {
|
||||
block_sync_lds();
|
||||
|
||||
block_gemm(c_block_tile,
|
||||
get_slice_tile(a_copy_reg_tensor,
|
||||
sequence<0, (i_k0)*kKPerBlock>{},
|
||||
sequence<kMPerBlock, (i_k0 + 1) * kKPerBlock>{}),
|
||||
b_copy_lds_window);
|
||||
|
||||
block_sync_lds();
|
||||
|
||||
move_tile_window(b_copy_dram_window, {0, kKPerBlock});
|
||||
|
||||
store_tile(b_copy_lds_window, b_block_tile);
|
||||
b_block_tile = load_tile(b_copy_dram_window);
|
||||
});
|
||||
}
|
||||
|
||||
// tail
|
||||
{
|
||||
block_sync_lds();
|
||||
|
||||
block_gemm(c_block_tile,
|
||||
get_slice_tile(a_copy_reg_tensor,
|
||||
sequence<0, (k_loops - 2) * kKPerBlock>{},
|
||||
sequence<kMPerBlock, (k_loops - 1) * kKPerBlock>{}),
|
||||
b_copy_lds_window);
|
||||
|
||||
block_sync_lds();
|
||||
|
||||
store_tile(b_copy_lds_window, b_block_tile);
|
||||
|
||||
block_sync_lds();
|
||||
|
||||
block_gemm(c_block_tile,
|
||||
get_slice_tile(a_copy_reg_tensor,
|
||||
sequence<0, (k_loops - 1) * kKPerBlock>{},
|
||||
sequence<kMPerBlock, (k_loops)*kKPerBlock>{}),
|
||||
b_copy_lds_window);
|
||||
}
|
||||
|
||||
return c_block_tile;
|
||||
}
|
||||
|
||||
template <typename ADramBlockWindowTmp,
|
||||
typename BDramBlockWindowTmp,
|
||||
typename ARegBlockTensorTmp>
|
||||
__device__ auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp,
|
||||
const BDramBlockWindowTmp& b_dram_block_window_tmp,
|
||||
ARegBlockTensorTmp& a_reg_block_tensor_tmp,
|
||||
void* p_smem) const
|
||||
{
|
||||
return operator()(
|
||||
a_dram_block_window_tmp,
|
||||
[](const ADataType& a) { return a; },
|
||||
b_dram_block_window_tmp,
|
||||
[](const BDataType& b) { return b; },
|
||||
a_reg_block_tensor_tmp,
|
||||
p_smem);
|
||||
}
|
||||
|
||||
template <typename BDramBlockWindowTmp, typename ARegBlockTensorTmp>
|
||||
__device__ auto operator()(const BDramBlockWindowTmp& b_dram_block_window_tmp,
|
||||
const ARegBlockTensorTmp& a_reg_block_tensor_tmp,
|
||||
void* p_smem) const
|
||||
{
|
||||
return operator()(
|
||||
b_dram_block_window_tmp,
|
||||
[](const BDataType& b) { return b; },
|
||||
a_reg_block_tensor_tmp,
|
||||
p_smem);
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace ck_tile
|
||||
@@ -0,0 +1,96 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "blockgemm_pipeline_agmem_bgmem_creg_policy_impl.hpp"
|
||||
#include "../../../example/ck_tile/99_toy_example/02_gemm/block_gemm_pipeline_agmem_bgmem_creg.hpp"
|
||||
|
||||
namespace ck_tile {
|
||||
|
||||
// NOTE: Assume A is K-Major
|
||||
struct BlockGemmPipelineAGmemBGmemCRegSkipALdsPolicy
|
||||
{
|
||||
template <typename Problem>
|
||||
__host__ __device__ static constexpr auto MakeARegBlockDescriptor()
|
||||
{
|
||||
constexpr auto blockgemm = GetBlockGemm<Problem>();
|
||||
using BlockGemm = remove_cvref_t<decltype(blockgemm)>;
|
||||
|
||||
return policy_impl::make_a_reg_block_descriptor<Problem, BlockGemm>();
|
||||
}
|
||||
|
||||
template <typename Problem>
|
||||
__host__ __device__ static constexpr auto MakeBLdsBlockDescriptor()
|
||||
{
|
||||
return policy_impl::make_b_lds_block_descriptor_3d_pad<Problem>();
|
||||
}
|
||||
|
||||
template <typename Problem>
|
||||
__host__ __device__ static constexpr auto MakeADramTileDistribution()
|
||||
{
|
||||
constexpr auto blockgemm = GetBlockGemm<Problem>();
|
||||
using BlockGemm = remove_cvref_t<decltype(blockgemm)>;
|
||||
|
||||
return policy_impl::make_a_dram_tile_distribution_skip_lds<Problem, BlockGemm>();
|
||||
}
|
||||
|
||||
template <typename Problem>
|
||||
__host__ __device__ static constexpr auto MakeBDramTileDistribution()
|
||||
{
|
||||
return policy_impl::make_b_dram_tile_distribution<Problem>();
|
||||
}
|
||||
|
||||
template <typename Problem>
|
||||
__host__ __device__ static constexpr auto GetBlockGemm()
|
||||
{
|
||||
using BlockGemmPolicy = BlockGemmARegBSmemCRegV1K8Policy;
|
||||
|
||||
return BlockGemmARegBSmemCRegV1<Problem, BlockGemmPolicy>{};
|
||||
}
|
||||
};
|
||||
|
||||
template <index_t AKDim_>
|
||||
struct BlockGemmPipelineAGmemBGmemCRegSkipALdsPersistentQRegCachePolicy
|
||||
: BlockGemmPipelineAGmemBGmemCRegSkipALdsPolicy
|
||||
{
|
||||
static constexpr index_t AKDim = AKDim_;
|
||||
|
||||
template <typename Problem>
|
||||
__host__ __device__ static constexpr auto MakeARegBlockDescriptor()
|
||||
{
|
||||
constexpr auto blockgemm = GetBlockGemm<Problem>();
|
||||
using BlockGemm = remove_cvref_t<decltype(blockgemm)>;
|
||||
|
||||
constexpr index_t kMPerBlock = Problem::BlockGemmShape::kM;
|
||||
constexpr index_t kKPerBlock = AKDim;
|
||||
|
||||
constexpr auto config =
|
||||
BlockGemm::BlockGemmPolicy::template GetWarpGemmMWarpNWarp<Problem>();
|
||||
|
||||
using WG = remove_cvref_t<decltype(config.template get<0>())>;
|
||||
|
||||
constexpr index_t MWarp = config.template get<1>();
|
||||
constexpr index_t NWarp = config.template get<2>();
|
||||
|
||||
constexpr index_t MIterPerWarp = kMPerBlock / (MWarp * WG::kM);
|
||||
constexpr index_t KIterPerWarp = kKPerBlock / WG::kK;
|
||||
|
||||
constexpr auto a_block_outer_dstr_encoding = tile_distribution_encoding<
|
||||
sequence<NWarp>,
|
||||
tuple<sequence<MIterPerWarp, MWarp>, sequence<KIterPerWarp>>,
|
||||
tuple<sequence<1, 0>>,
|
||||
tuple<sequence<1, 0>>,
|
||||
sequence<1, 2>,
|
||||
sequence<0, 0>>{};
|
||||
|
||||
constexpr auto a_block_dstr_encode = detail::make_embed_tile_distribution_encoding(
|
||||
a_block_outer_dstr_encoding, typename WG::AWarpDstrEncoding{});
|
||||
|
||||
constexpr auto a_block_dstr = make_static_tile_distribution(a_block_dstr_encode);
|
||||
|
||||
return a_block_dstr;
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace ck_tile
|
||||
@@ -0,0 +1,26 @@
|
||||
// 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 {
|
||||
|
||||
template <typename ADataType_,
|
||||
typename BDataType_,
|
||||
typename CDataType_,
|
||||
index_t kBlockSize_,
|
||||
typename BlockGemmShape_>
|
||||
struct BlockGemmPipelineProblem
|
||||
{
|
||||
using ADataType = remove_cvref_t<ADataType_>;
|
||||
using BDataType = remove_cvref_t<BDataType_>;
|
||||
using CDataType = remove_cvref_t<CDataType_>;
|
||||
using BlockGemmShape = remove_cvref_t<BlockGemmShape_>;
|
||||
|
||||
static constexpr index_t kBlockSize = kBlockSize_;
|
||||
};
|
||||
|
||||
} // namespace ck_tile
|
||||
|
||||
@@ -0,0 +1,180 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.
|
||||
|
||||
#include "ck_tile/core.hpp"
|
||||
#include "ck_tile/core/tensor/tile_distribution.hpp"
|
||||
|
||||
namespace ck_tile {
|
||||
namespace policy_impl {
|
||||
|
||||
// 3d + padding
|
||||
template <typename Problem>
|
||||
__host__ __device__ static constexpr auto make_a_lds_block_descriptor_3d_pad()
|
||||
{
|
||||
constexpr index_t kMPerBlock = Problem::BlockGemmShape::kM;
|
||||
constexpr index_t kKPerBlock = Problem::BlockGemmShape::kK;
|
||||
|
||||
constexpr auto a_lds_block_desc_0 = make_naive_tensor_descriptor(
|
||||
make_tuple(number<kKPerBlock / 8>{}, number<kMPerBlock>{}, number<8>{}),
|
||||
make_tuple(number<(kMPerBlock + 1) * 8>{}, number<8>{}, number<1>{}),
|
||||
number<8>{},
|
||||
number<1>{});
|
||||
|
||||
constexpr auto a_lds_block_desc =
|
||||
transform_tensor_descriptor(a_lds_block_desc_0,
|
||||
make_tuple(make_pass_through_transform(kMPerBlock),
|
||||
make_merge_transform(make_tuple(kKPerBlock / 8, 8))),
|
||||
make_tuple(sequence<1>{}, sequence<0, 2>{}),
|
||||
make_tuple(sequence<0>{}, sequence<1>{}));
|
||||
|
||||
return a_lds_block_desc;
|
||||
}
|
||||
|
||||
// 3d + padding
|
||||
template <typename Problem>
|
||||
__host__ __device__ static constexpr auto make_b_lds_block_descriptor_3d_pad()
|
||||
{
|
||||
constexpr index_t kNPerBlock = Problem::BlockGemmShape::kN;
|
||||
constexpr index_t kKPerBlock = Problem::BlockGemmShape::kK;
|
||||
|
||||
constexpr auto b_lds_block_desc_0 = make_naive_tensor_descriptor(
|
||||
make_tuple(number<kKPerBlock / 8>{}, number<kNPerBlock>{}, number<8>{}),
|
||||
make_tuple(number<(kNPerBlock + 1) * 8>{}, number<8>{}, number<1>{}),
|
||||
number<8>{},
|
||||
number<1>{});
|
||||
|
||||
constexpr auto b_lds_block_desc =
|
||||
transform_tensor_descriptor(b_lds_block_desc_0,
|
||||
make_tuple(make_pass_through_transform(kNPerBlock),
|
||||
make_merge_transform(make_tuple(kKPerBlock / 8, 8))),
|
||||
make_tuple(sequence<1>{}, sequence<0, 2>{}),
|
||||
make_tuple(sequence<0>{}, sequence<1>{}));
|
||||
|
||||
return b_lds_block_desc;
|
||||
}
|
||||
|
||||
template <typename Problem, typename BlockGemm>
|
||||
__host__ __device__ static constexpr auto make_a_reg_block_descriptor()
|
||||
{
|
||||
constexpr index_t kMPerBlock = Problem::BlockGemmShape::kM;
|
||||
constexpr index_t kKPerBlock = Problem::BlockGemmShape::kK;
|
||||
|
||||
constexpr auto config = BlockGemm::BlockGemmPolicy::template GetWarpGemmMWarpNWarp<Problem>();
|
||||
|
||||
using WG = remove_cvref_t<decltype(config.template get<0>())>;
|
||||
|
||||
constexpr index_t MWarp = config.template get<1>();
|
||||
constexpr index_t NWarp = config.template get<2>();
|
||||
|
||||
constexpr index_t MIterPerWarp = kMPerBlock / (MWarp * WG::kM);
|
||||
constexpr index_t KIterPerWarp = kKPerBlock / WG::kK;
|
||||
|
||||
constexpr auto a_block_outer_dstr_encoding =
|
||||
tile_distribution_encoding<sequence<NWarp>,
|
||||
tuple<sequence<MIterPerWarp, MWarp>, sequence<KIterPerWarp>>,
|
||||
tuple<sequence<1, 0>>,
|
||||
tuple<sequence<1, 0>>,
|
||||
sequence<1, 2>,
|
||||
sequence<0, 0>>{};
|
||||
|
||||
constexpr auto a_block_dstr_encode = detail::make_embed_tile_distribution_encoding(
|
||||
a_block_outer_dstr_encoding, typename WG::AWarpDstrEncoding{});
|
||||
|
||||
constexpr auto a_block_dstr = make_static_tile_distribution(a_block_dstr_encode);
|
||||
|
||||
return a_block_dstr;
|
||||
}
|
||||
|
||||
template <typename Problem>
|
||||
__host__ __device__ static constexpr auto make_a_dram_tile_distribution()
|
||||
{
|
||||
using ADataType = remove_cvref_t<typename Problem::ADataType>;
|
||||
|
||||
constexpr index_t kBlockSize = Problem::kBlockSize;
|
||||
|
||||
constexpr index_t kMPerBlock = Problem::BlockGemmShape::kM;
|
||||
constexpr index_t kKPerBlock = Problem::BlockGemmShape::kK;
|
||||
|
||||
constexpr index_t K1 = 16 / sizeof(ADataType);
|
||||
constexpr index_t K0 = kKPerBlock / K1;
|
||||
constexpr index_t M2 = get_warp_size() / K0;
|
||||
|
||||
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>>{});
|
||||
}
|
||||
|
||||
template <typename Problem, typename BlockGemm>
|
||||
__host__ __device__ static constexpr auto make_a_dram_tile_distribution_skip_lds()
|
||||
{
|
||||
constexpr auto config = BlockGemm::BlockGemmPolicy::template GetWarpGemmMWarpNWarp<Problem>();
|
||||
|
||||
using WG = remove_cvref_t<decltype(config.template get<0>())>;
|
||||
|
||||
constexpr index_t MWarp = config.template get<1>();
|
||||
|
||||
constexpr index_t kMPerBlock = Problem::BlockGemmShape::kM;
|
||||
constexpr index_t kKPerBlock = Problem::BlockGemmShape::kK;
|
||||
|
||||
constexpr index_t K2 =
|
||||
WG::kK / WG::WarpGemmAttribute::Impl::kABKLane; // WG::WarpGemmAttribute::Impl::kABKPerLane;
|
||||
// // 16 / sizeof(ADataType);
|
||||
constexpr index_t K1 = WG::WarpGemmAttribute::Impl::kABKLane;
|
||||
constexpr index_t K0 = kKPerBlock / (K1 * K2);
|
||||
|
||||
constexpr index_t M2 = WG::WarpGemmAttribute::Impl::kAMLane;
|
||||
constexpr index_t M1 = MWarp;
|
||||
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, K2>>,
|
||||
tuple<sequence<1>, sequence<2, 1>>,
|
||||
tuple<sequence<1>, sequence<1, 2>>,
|
||||
sequence<2, 1, 2>,
|
||||
sequence<0, 0, 2>>{});
|
||||
}
|
||||
|
||||
template <typename Problem>
|
||||
__host__ __device__ static constexpr auto make_b_dram_tile_distribution()
|
||||
{
|
||||
using BDataType = remove_cvref_t<typename Problem::BDataType>;
|
||||
|
||||
constexpr index_t kBlockSize = Problem::kBlockSize;
|
||||
|
||||
constexpr index_t kNPerBlock = Problem::BlockGemmShape::kN;
|
||||
constexpr index_t kKPerBlock = Problem::BlockGemmShape::kK;
|
||||
|
||||
constexpr index_t K1 = 16 / sizeof(BDataType);
|
||||
constexpr index_t K0 = kKPerBlock / K1;
|
||||
constexpr index_t N2 = get_warp_size() / K0;
|
||||
|
||||
constexpr index_t N1 = kBlockSize / get_warp_size();
|
||||
constexpr index_t N0 = kNPerBlock / (N2 * N1);
|
||||
|
||||
return make_static_tile_distribution(
|
||||
tile_distribution_encoding<sequence<1>,
|
||||
tuple<sequence<N0, N1, N2>, sequence<K0, K1>>,
|
||||
tuple<sequence<1>, sequence<1, 2>>,
|
||||
tuple<sequence<1>, sequence<2, 0>>,
|
||||
sequence<1, 2>,
|
||||
sequence<0, 1>>{});
|
||||
}
|
||||
|
||||
template <typename Problem>
|
||||
__host__ __device__ static constexpr auto get_block_gemm()
|
||||
{
|
||||
using BlockGemmPolicy = BlockGemmASmemBSmemCRegDefaultPolicy;
|
||||
|
||||
return BlockGemmASmemBSmemCReg<Problem, BlockGemmPolicy>{};
|
||||
}
|
||||
|
||||
} // namespace policy_impl
|
||||
} // namespace ck_tile
|
||||
@@ -0,0 +1,202 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.
|
||||
|
||||
#include <cstring>
|
||||
|
||||
#include "ck_tile/host.hpp"
|
||||
|
||||
#include "reference_batched_gemm.hpp"
|
||||
#include "reference_batched_softmax.hpp"
|
||||
#include "flash_attention_fwd.hpp"
|
||||
|
||||
/*
|
||||
* Toy code of flash attention forward pass
|
||||
* Assume simplest case.
|
||||
* Q [Batch, HeadNum, SeqenceLengthQ, HeadDim]
|
||||
* K [Batch, HeadNum, SeqenceLengthK, HeadDim]
|
||||
* V [Batch, HeadNum, HeadDim, SeqenceLengthK]
|
||||
* O [Batch, HeadNum, SeqenceLengthQ, HeadDim]
|
||||
*/
|
||||
|
||||
int main(int argc, char* argv[])
|
||||
{
|
||||
using QDataType = ck_tile::half_t;
|
||||
using KDataType = ck_tile::half_t;
|
||||
using VDataType = ck_tile::half_t;
|
||||
using SaccDataType = float;
|
||||
using SMPLComputeDataType = float;
|
||||
using PDataType = ck_tile::half_t;
|
||||
using OaccDataType = float;
|
||||
using ODataType = ck_tile::half_t;
|
||||
|
||||
ck_tile::index_t Batch = 64; // Batch Number * Head Number
|
||||
ck_tile::index_t M0 = 4096; // SequenceLengthQ
|
||||
ck_tile::index_t N0 = 4096; // SequencelengthK
|
||||
ck_tile::index_t K0 = 128; // HeadDim
|
||||
ck_tile::index_t N1 = 128; // HeadDim
|
||||
ck_tile::index_t verification = 0;
|
||||
ck_tile::index_t init_method = 1;
|
||||
[[maybe_unused]] ck_tile::index_t time_kernel = 0;
|
||||
|
||||
if(argc == 4)
|
||||
{
|
||||
init_method = std::stoi(argv[1]);
|
||||
time_kernel = std::stoi(argv[2]);
|
||||
verification = std::stoi(argv[3]);
|
||||
}
|
||||
|
||||
if(argc == 9)
|
||||
{
|
||||
init_method = std::stoi(argv[1]);
|
||||
time_kernel = std::stoi(argv[2]);
|
||||
verification = std::stoi(argv[3]);
|
||||
Batch = std::stoi(argv[4]);
|
||||
M0 = std::stoi(argv[5]);
|
||||
N0 = std::stoi(argv[6]);
|
||||
K0 = std::stoi(argv[7]);
|
||||
N1 = std::stoi(argv[8]);
|
||||
}
|
||||
|
||||
std::array<ck_tile::index_t, 3> q_lengths{Batch, M0, K0};
|
||||
std::array<ck_tile::index_t, 3> q_strides{M0 * K0, K0, 1};
|
||||
|
||||
std::array<ck_tile::index_t, 3> k_lengths{Batch, N0, K0};
|
||||
std::array<ck_tile::index_t, 3> k_strides{N0 * K0, K0, 1};
|
||||
|
||||
std::array<ck_tile::index_t, 3> v_lengths{Batch, N1, N0};
|
||||
std::array<ck_tile::index_t, 3> v_strides{N1 * N0, N0, 1};
|
||||
|
||||
std::array<ck_tile::index_t, 3> s_lengths{Batch, M0, N0};
|
||||
std::array<ck_tile::index_t, 3> s_strides{M0 * N0, N0, 1};
|
||||
|
||||
std::array<ck_tile::index_t, 3> p_lengths{Batch, M0, N0};
|
||||
std::array<ck_tile::index_t, 3> p_strides{M0 * N0, N0, 1};
|
||||
|
||||
std::array<ck_tile::index_t, 3> o_lengths{Batch, M0, N1};
|
||||
std::array<ck_tile::index_t, 3> o_strides{M0 * N1, N1, 1};
|
||||
|
||||
// host verify
|
||||
ck_tile::HostTensor<QDataType> q_host(q_lengths, q_strides);
|
||||
ck_tile::HostTensor<KDataType> k_host(k_lengths, k_strides);
|
||||
ck_tile::HostTensor<VDataType> v_host(v_lengths, v_strides);
|
||||
ck_tile::HostTensor<ODataType> o_host_dev(o_lengths, o_strides);
|
||||
|
||||
switch(init_method)
|
||||
{
|
||||
case 0: break;
|
||||
case 1:
|
||||
ck_tile::FillUniformDistributionIntegerValue<QDataType>{-3.f, 3.f}(q_host);
|
||||
ck_tile::FillUniformDistributionIntegerValue<KDataType>{-3.f, 3.f}(k_host);
|
||||
ck_tile::FillUniformDistributionIntegerValue<VDataType>{-3.f, 3.f}(v_host);
|
||||
break;
|
||||
case 2:
|
||||
ck_tile::FillUniformDistribution<QDataType>{-3.f, 3.f}(q_host);
|
||||
ck_tile::FillUniformDistribution<KDataType>{-3.f, 3.f}(k_host);
|
||||
ck_tile::FillUniformDistribution<VDataType>{-3.f, 3.f}(v_host);
|
||||
break;
|
||||
default:
|
||||
ck_tile::FillUniformDistributionIntegerValue<QDataType>{-2.f, 2.f}(q_host);
|
||||
ck_tile::FillUniformDistributionIntegerValue<KDataType>{-2.f, 2.f}(k_host);
|
||||
ck_tile::FillUniformDistributionIntegerValue<VDataType>{-2.f, 2.f}(v_host);
|
||||
}
|
||||
ck_tile::DeviceMem q_buf(q_host.get_element_space_size_in_bytes());
|
||||
ck_tile::DeviceMem k_buf(k_host.get_element_space_size_in_bytes());
|
||||
ck_tile::DeviceMem v_buf(v_host.get_element_space_size_in_bytes());
|
||||
ck_tile::DeviceMem o_buf(o_host_dev.get_element_space_size_in_bytes());
|
||||
|
||||
q_buf.ToDevice(q_host.mData.data());
|
||||
k_buf.ToDevice(k_host.mData.data());
|
||||
v_buf.ToDevice(v_host.mData.data());
|
||||
|
||||
constexpr ck_tile::index_t kM0PerBlock = 128;
|
||||
constexpr ck_tile::index_t kN0PerBlock = 128;
|
||||
constexpr ck_tile::index_t kK0PerBlock = 32;
|
||||
constexpr ck_tile::index_t kN1PerBlock = 128;
|
||||
constexpr ck_tile::index_t kK1PerBlock = 32;
|
||||
|
||||
constexpr ck_tile::index_t kBlockSize = 256;
|
||||
constexpr ck_tile::index_t kHeadDim = 128;
|
||||
|
||||
ck_tile::index_t kGridSize = Batch * (M0 / kM0PerBlock) * (N1 / kN1PerBlock);
|
||||
|
||||
std::cout << "grid size " << kGridSize << std::endl;
|
||||
|
||||
constexpr ck_tile::index_t kWarpPerCu = 8; // 2 warps per SIMD
|
||||
constexpr ck_tile::index_t kWarpPerBlock = kBlockSize / warpSize;
|
||||
constexpr ck_tile::index_t kBlockPerCu = kWarpPerCu / kWarpPerBlock;
|
||||
|
||||
float ave_time = ck_tile::launch_kernel(ck_tile::stream_config{nullptr, true},
|
||||
ck_tile::make_kernel<kBlockSize, kBlockPerCu>(
|
||||
ck_tile::FlashAttentionFwd<QDataType,
|
||||
KDataType,
|
||||
VDataType,
|
||||
SaccDataType,
|
||||
SMPLComputeDataType,
|
||||
PDataType,
|
||||
OaccDataType,
|
||||
ODataType,
|
||||
kBlockSize,
|
||||
kHeadDim,
|
||||
kM0PerBlock,
|
||||
kN0PerBlock,
|
||||
kK0PerBlock,
|
||||
kN1PerBlock,
|
||||
kK1PerBlock>{},
|
||||
kGridSize,
|
||||
kBlockSize,
|
||||
0,
|
||||
static_cast<QDataType*>(q_buf.GetDeviceBuffer()),
|
||||
static_cast<KDataType*>(k_buf.GetDeviceBuffer()),
|
||||
static_cast<VDataType*>(v_buf.GetDeviceBuffer()),
|
||||
static_cast<ODataType*>(o_buf.GetDeviceBuffer()),
|
||||
M0,
|
||||
N0,
|
||||
K0,
|
||||
N1,
|
||||
Batch,
|
||||
K0, // StrideQ
|
||||
K0, // StrideK
|
||||
N0, // StrideV
|
||||
N1, // StrideO
|
||||
M0 * K0, // BatchStrideQ
|
||||
N0 * K0, // BatchStrideK
|
||||
N1 * N0, // BatchStrideV
|
||||
M0 * N1)); // BatchStrideO
|
||||
|
||||
// reference
|
||||
auto pass = true;
|
||||
if(verification)
|
||||
{
|
||||
o_buf.FromDevice(o_host_dev.mData.data());
|
||||
|
||||
ck_tile::HostTensor<SMPLComputeDataType> s_host_ref(s_lengths, s_strides);
|
||||
ck_tile::HostTensor<PDataType> p_host_ref(p_lengths, p_strides);
|
||||
ck_tile::HostTensor<ODataType> o_host_ref(o_lengths, o_strides);
|
||||
|
||||
ck_tile::reference_batched_gemm<QDataType, KDataType, SaccDataType, SMPLComputeDataType>(
|
||||
q_host, k_host, s_host_ref);
|
||||
ck_tile::reference_batched_softmax<SMPLComputeDataType, SMPLComputeDataType, PDataType>(s_host_ref,
|
||||
p_host_ref);
|
||||
ck_tile::reference_batched_gemm<PDataType, VDataType, OaccDataType, ODataType>(
|
||||
p_host_ref, v_host, o_host_ref);
|
||||
|
||||
pass &= ck_tile::check_err(o_host_dev, o_host_ref);
|
||||
std::cout << "valid:" << (pass ? "y" : "n") << std::endl;
|
||||
}
|
||||
|
||||
std::size_t flop =
|
||||
std::size_t(2) * Batch * M0 * N0 * K0 + std::size_t(2) * Batch * M0 * N1 * N0;
|
||||
std::size_t num_btype =
|
||||
sizeof(QDataType) * Batch * M0 * K0 + sizeof(KDataType) * Batch * N0 * K0 +
|
||||
sizeof(VDataType) * Batch * N1 * N0 + sizeof(ODataType) * Batch * M0 * N1;
|
||||
|
||||
float tflops = static_cast<float>(flop) / 1.E9 / ave_time;
|
||||
|
||||
float gb_per_sec = num_btype / 1.E6 / ave_time;
|
||||
|
||||
std::cout << "Perf: " << ave_time << " ms, " << tflops << " TFlops, " << gb_per_sec << " GB/s"
|
||||
<< std::endl;
|
||||
|
||||
return !pass;
|
||||
}
|
||||
|
||||
@@ -0,0 +1,109 @@
|
||||
// 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/common.hpp"
|
||||
#include "ck_tile/ops/gemm/warp/warp_gemm.hpp"
|
||||
#include "ck_tile/core/tensor/tile_distribution.hpp"
|
||||
|
||||
#include "../../../example/ck_tile/99_toy_example/02_gemm/block_gemm_pipeline_agmem_bgmem_creg.hpp"
|
||||
#include "block_gemm_pipeline_problem.hpp"
|
||||
#include "block_gemm_areg_bsmem_creg_v1.hpp"
|
||||
#include "flash_attention_fwd_impl.hpp"
|
||||
|
||||
namespace ck_tile {
|
||||
|
||||
// S[M0, N0] = Q[M0, K0] * K[N0, K0]
|
||||
// P[M0, N0] = Softmax(S[M0, N0])
|
||||
// O[M0, N1] = P[M0, N0] * V[N1, N0]
|
||||
template <typename QDataType,
|
||||
typename KDataType,
|
||||
typename VDataType,
|
||||
typename SaccDataType,
|
||||
typename SMPLComputeDataType,
|
||||
typename PDataType,
|
||||
typename OaccDataType,
|
||||
typename ODataType,
|
||||
index_t kBlockSize,
|
||||
index_t kHeadDim,
|
||||
index_t kM0PerBlock,
|
||||
index_t kN0PerBlock,
|
||||
index_t kK0PerBlock,
|
||||
index_t kN1PerBlock,
|
||||
index_t kK1PerBlock>
|
||||
struct FlashAttentionFwd
|
||||
{
|
||||
__device__ void operator()(const QDataType* q_ptr,
|
||||
const KDataType* k_ptr,
|
||||
const VDataType* v_ptr,
|
||||
ODataType* o_ptr,
|
||||
const index_t M0,
|
||||
const index_t N0,
|
||||
const index_t K0,
|
||||
const index_t N1,
|
||||
const index_t /* Batch */,
|
||||
const index_t StrideQ,
|
||||
const index_t StrideK,
|
||||
const index_t StrideV,
|
||||
const index_t StrideO,
|
||||
const index_t BatchStrideQ,
|
||||
const index_t BatchStrideK,
|
||||
const index_t BatchStrideV,
|
||||
const index_t BatchStrideO) const
|
||||
{
|
||||
// divide problem
|
||||
const index_t num_tile_m0 = M0 / kM0PerBlock;
|
||||
const index_t num_tile_n1 = N1 / kN1PerBlock;
|
||||
|
||||
const index_t id_block = get_block_id();
|
||||
|
||||
const auto f = [](index_t dividend, index_t divisor) {
|
||||
index_t quotient = dividend / divisor;
|
||||
index_t modulus = dividend - quotient * divisor;
|
||||
|
||||
return make_tuple(quotient, modulus);
|
||||
};
|
||||
|
||||
const auto [itmp, id_tile_n] = f(id_block, num_tile_n1);
|
||||
const auto [id_tile_batch, id_tile_m] = f(itmp, num_tile_m0);
|
||||
|
||||
const index_t iBatch = __builtin_amdgcn_readfirstlane(id_tile_batch);
|
||||
const index_t iM0 = __builtin_amdgcn_readfirstlane(id_tile_m * kM0PerBlock);
|
||||
const index_t iN1 = __builtin_amdgcn_readfirstlane(id_tile_n * kN1PerBlock);
|
||||
|
||||
const auto kernel_impl = FlashAttentionFwdImpl<QDataType,
|
||||
KDataType,
|
||||
VDataType,
|
||||
SaccDataType,
|
||||
SMPLComputeDataType,
|
||||
PDataType,
|
||||
OaccDataType,
|
||||
ODataType,
|
||||
kBlockSize,
|
||||
kHeadDim,
|
||||
kM0PerBlock,
|
||||
kN0PerBlock,
|
||||
kK0PerBlock,
|
||||
kN1PerBlock,
|
||||
kK1PerBlock>{};
|
||||
|
||||
kernel_impl(q_ptr + iBatch * BatchStrideQ,
|
||||
k_ptr + iBatch * BatchStrideK,
|
||||
v_ptr + iBatch * BatchStrideV,
|
||||
o_ptr + iBatch * BatchStrideO,
|
||||
M0,
|
||||
N0,
|
||||
K0,
|
||||
N1,
|
||||
StrideQ,
|
||||
StrideK,
|
||||
StrideV,
|
||||
StrideO,
|
||||
iM0,
|
||||
iN1);
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace ck_tile
|
||||
@@ -0,0 +1,355 @@
|
||||
// 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/common.hpp"
|
||||
#include "ck_tile/ops/gemm/warp/warp_gemm.hpp"
|
||||
#include "ck_tile/core/tensor/tile_distribution.hpp"
|
||||
|
||||
#include "tile_gemm_shape.hpp"
|
||||
#include "../../../example/ck_tile/99_toy_example/02_gemm/block_gemm_pipeline_agmem_bgmem_creg.hpp"
|
||||
|
||||
#include "block_gemm_pipeline_agmem_bgmem_creg_v2_askiplds.hpp"
|
||||
#include "block_gemm_pipeline_problem.hpp"
|
||||
#include "block_gemm_areg_bsmem_creg_v1.hpp"
|
||||
#include "ck_tile/ops/reduce.hpp"
|
||||
|
||||
|
||||
namespace ck_tile {
|
||||
|
||||
// S[M0, N0] = Q[M0, K0] * K[N0, K0]
|
||||
// P[M0, N0] = Softmax(S[M0, N0])
|
||||
// O[M0, N1] = P[M0, N0] * V[N1, N0]
|
||||
template <typename QDataType,
|
||||
typename KDataType,
|
||||
typename VDataType,
|
||||
typename SaccDataType,
|
||||
typename SMPLComputeDataType,
|
||||
typename PDataType,
|
||||
typename OaccDataType,
|
||||
typename ODataType,
|
||||
index_t kBlockSize,
|
||||
index_t kHeadDim,
|
||||
index_t kM0PerBlock,
|
||||
index_t kN0PerBlock,
|
||||
index_t kK0PerBlock,
|
||||
index_t kN1PerBlock,
|
||||
index_t kK1PerBlock>
|
||||
struct FlashAttentionFwdImpl
|
||||
{
|
||||
// block gemm0 pipeline
|
||||
using BlockGemm0Problem = BlockGemmPipelineProblem<
|
||||
QDataType,
|
||||
KDataType,
|
||||
SaccDataType,
|
||||
kBlockSize,
|
||||
TileGemmShape<kM0PerBlock, kN0PerBlock, kK0PerBlock>>;
|
||||
|
||||
using BlockGemm0Policy =
|
||||
BlockGemmPipelineAGmemBGmemCRegSkipALdsPersistentQRegCachePolicy<kHeadDim>;
|
||||
|
||||
using BlockGemm0Pipeline =
|
||||
BlockGemmPipelineAGmemBGmemCReg<BlockGemm0Problem, BlockGemm0Policy>;
|
||||
|
||||
// block gemm1
|
||||
using BlockGemm1 = BlockGemmARegBSmemCRegV1<
|
||||
BlockGemmARegBSmemCRegProblem<
|
||||
PDataType,
|
||||
VDataType,
|
||||
OaccDataType,
|
||||
kBlockSize,
|
||||
TileGemmShape<kM0PerBlock, kN1PerBlock, kK1PerBlock>>,
|
||||
BlockGemmARegBSmemCRegV1DefaultPolicy>;
|
||||
|
||||
// 3d, with padding
|
||||
__device__ static constexpr auto MakeVLdsBlockDescriptor()
|
||||
{
|
||||
constexpr index_t kNPerBlock = kN1PerBlock;
|
||||
constexpr index_t kKPerBlock = kK1PerBlock;
|
||||
constexpr index_t kPad = 1;
|
||||
// 2% faster than use kK1 = 8
|
||||
constexpr index_t kK1 = 4;
|
||||
|
||||
constexpr auto b_lds_block_desc_0 = make_naive_tensor_descriptor(
|
||||
make_tuple(number<kKPerBlock / kK1>{}, number<kNPerBlock>{}, number<kK1>{}),
|
||||
make_tuple(number<(kNPerBlock + kPad) * kK1>{}, number<kK1>{}, number<1>{}),
|
||||
number<kK1>{},
|
||||
number<1>{});
|
||||
|
||||
constexpr auto b_lds_block_desc = transform_tensor_descriptor(
|
||||
b_lds_block_desc_0,
|
||||
make_tuple(make_pass_through_transform(kNPerBlock),
|
||||
make_merge_transform(make_tuple(number<kKPerBlock / kK1>{}, number<kK1>{}))),
|
||||
make_tuple(sequence<1>{}, sequence<0, 2>{}),
|
||||
make_tuple(sequence<0>{}, sequence<1>{}));
|
||||
|
||||
return b_lds_block_desc;
|
||||
}
|
||||
|
||||
__device__ static constexpr auto MakeVDramTileDistribution()
|
||||
{
|
||||
using BDataType = VDataType;
|
||||
|
||||
constexpr index_t kNPerBlock = kN1PerBlock;
|
||||
constexpr index_t kKPerBlock = kK1PerBlock;
|
||||
|
||||
constexpr index_t K1 = 16 / sizeof(BDataType);
|
||||
constexpr index_t K0 = kKPerBlock / K1;
|
||||
constexpr index_t N2 = get_warp_size() / K0;
|
||||
constexpr index_t N1 = kBlockSize / get_warp_size();
|
||||
constexpr index_t N0 = kNPerBlock / (N2 * N1);
|
||||
|
||||
return make_static_tile_distribution(
|
||||
tile_distribution_encoding<sequence<1>,
|
||||
tuple<sequence<N0, N1, N2>, sequence<K0, K1>>,
|
||||
tuple<sequence<1>, sequence<1, 2>>,
|
||||
tuple<sequence<1>, sequence<2, 0>>,
|
||||
sequence<1, 2>,
|
||||
sequence<0, 1>>{});
|
||||
}
|
||||
|
||||
__device__ static constexpr index_t GetStaticLdsSize()
|
||||
{
|
||||
return max(BlockGemm0Pipeline::GetStaticLdsSize(),
|
||||
static_cast<index_t>(MakeVLdsBlockDescriptor().get_element_space_size() *
|
||||
sizeof(VDataType)));
|
||||
}
|
||||
|
||||
__device__ void operator()(const QDataType* q_ptr,
|
||||
const KDataType* k_ptr,
|
||||
const VDataType* v_ptr,
|
||||
ODataType* o_ptr,
|
||||
const index_t M0,
|
||||
const index_t N0,
|
||||
const index_t K0,
|
||||
const index_t N1,
|
||||
const index_t StrideQ,
|
||||
const index_t StrideK,
|
||||
const index_t StrideV,
|
||||
const index_t StrideO,
|
||||
const index_t iM0,
|
||||
const index_t iN1) const
|
||||
{
|
||||
constexpr auto I0 = number<0>{};
|
||||
constexpr auto I1 = number<1>{};
|
||||
|
||||
// allocate LDS
|
||||
__shared__ char smem_ptr[GetStaticLdsSize()];
|
||||
|
||||
// Q/K/V DRAM and DRAM window
|
||||
const auto q_dram = make_naive_tensor_view<address_space_enum::global>(
|
||||
q_ptr, make_tuple(M0, K0), make_tuple(StrideQ, 1), number<32>{}, number<1>{});
|
||||
|
||||
const auto k_dram = make_naive_tensor_view<address_space_enum::global>(
|
||||
k_ptr, make_tuple(N0, K0), make_tuple(StrideK, 1), number<32>{}, number<1>{});
|
||||
|
||||
const auto v_dram = make_naive_tensor_view<address_space_enum::global>(
|
||||
v_ptr, make_tuple(N1, N0), make_tuple(StrideV, 1), number<32>{}, number<1>{});
|
||||
|
||||
auto q_dram_window = make_tile_window(
|
||||
q_dram, make_tuple(number<kM0PerBlock>{}, number<kK0PerBlock>{}), {iM0, 0});
|
||||
|
||||
auto k_dram_window = make_tile_window(
|
||||
k_dram, make_tuple(number<kN0PerBlock>{}, number<kK0PerBlock>{}), {0, 0});
|
||||
|
||||
auto v_dram_window =
|
||||
make_tile_window(v_dram,
|
||||
make_tuple(number<kN1PerBlock>{}, number<kK1PerBlock>{}),
|
||||
{iN1, 0},
|
||||
MakeVDramTileDistribution());
|
||||
|
||||
// Q in Register
|
||||
auto q_reg_tensor = make_static_distributed_tensor<QDataType>(
|
||||
BlockGemm0Policy::template MakeARegBlockDescriptor<BlockGemm0Problem>());
|
||||
|
||||
// V LDS and LDS window
|
||||
// V LDS occupies the same LDS allocation Q/K LDS
|
||||
auto v_lds = make_tensor_view<address_space_enum::lds>(reinterpret_cast<VDataType*>(smem_ptr),
|
||||
MakeVLdsBlockDescriptor());
|
||||
|
||||
auto v_lds_window = make_tile_window(
|
||||
v_lds, make_tuple(number<kN1PerBlock>{}, number<kK1PerBlock>{}), {0, 0});
|
||||
|
||||
// Block GEMM0 pipeline and Block GEMM1
|
||||
constexpr auto gemm0_pipeline = BlockGemm0Pipeline{};
|
||||
constexpr auto gemm1 = BlockGemm1{};
|
||||
|
||||
// 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; };
|
||||
|
||||
// infer Sacc, S, P, M, L, Oacc type
|
||||
using SaccBlockTileType =
|
||||
decltype(gemm0_pipeline(q_dram_window, k_dram_window, q_reg_tensor, nullptr));
|
||||
|
||||
using SBlockTileType = decltype(tile_elementwise_in(
|
||||
type_convert<SMPLComputeDataType, SaccDataType>, SaccBlockTileType{}));
|
||||
|
||||
using PBlockTileType = decltype(tile_elementwise_in(type_convert<PDataType, SaccDataType>,
|
||||
SaccBlockTileType{}));
|
||||
|
||||
using MLBlockTileType = decltype(block_tile_reduce<SMPLComputeDataType>(
|
||||
SBlockTileType{}, sequence<1>{}, f_max, SMPLComputeDataType{0}));
|
||||
|
||||
using OaccBlockTileType = decltype(gemm1(
|
||||
get_slice_tile(
|
||||
PBlockTileType{}, sequence<0, 0>{}, sequence<kM0PerBlock, kK1PerBlock>{}),
|
||||
v_dram_window));
|
||||
|
||||
// init Sacc, Oacc, M, L
|
||||
auto s_acc = SaccBlockTileType{};
|
||||
auto o_acc = OaccBlockTileType{};
|
||||
auto m = MLBlockTileType{};
|
||||
auto l = MLBlockTileType{};
|
||||
|
||||
tile_elementwise_inout([](auto& e) { e = 0; }, o_acc);
|
||||
tile_elementwise_inout([](auto& e) { e = std::numeric_limits<SMPLComputeDataType>::lowest(); },
|
||||
m);
|
||||
tile_elementwise_inout([](auto& e) { e = 0; }, l);
|
||||
|
||||
// loop over Column of S (J loop)
|
||||
index_t iN0 = 0;
|
||||
|
||||
// Cold Q_Reg_Cache
|
||||
s_acc = gemm0_pipeline(q_dram_window, k_dram_window, q_reg_tensor, smem_ptr);
|
||||
do
|
||||
{
|
||||
// Hot Q_Reg_Cache
|
||||
if(iN0 > 0)
|
||||
{
|
||||
s_acc = gemm0_pipeline(k_dram_window, q_reg_tensor, smem_ptr);
|
||||
}
|
||||
// S{j}
|
||||
const auto s =
|
||||
tile_elementwise_in(type_convert<SMPLComputeDataType, SaccDataType>, s_acc);
|
||||
|
||||
// prefetch load v tile
|
||||
const auto v_prefetch = load_tile(v_dram_window);
|
||||
|
||||
// m_local = rowmax(S{j})
|
||||
auto m_local = block_tile_reduce<SMPLComputeDataType>(
|
||||
s, sequence<1>{}, f_max, std::numeric_limits<SMPLComputeDataType>::lowest());
|
||||
|
||||
block_tile_reduce_sync(m_local, f_max);
|
||||
|
||||
// m{j-1}
|
||||
const auto m_old = m;
|
||||
|
||||
// m{j}
|
||||
tile_elementwise_inout(
|
||||
[](auto& e0, auto e1, auto e2) { e0 = max(e1, e2); }, m, m_old, m_local);
|
||||
|
||||
// Pcompute{j}
|
||||
auto p_compute =
|
||||
make_static_distributed_tensor<SMPLComputeDataType>(s.get_tile_distribution());
|
||||
|
||||
constexpr auto p_spans = decltype(p_compute)::get_distributed_spans();
|
||||
|
||||
sweep_tile_span(p_spans[I0], [&](auto idx0) {
|
||||
constexpr auto i_idx = make_tuple(idx0);
|
||||
|
||||
sweep_tile_span(p_spans[I1], [&](auto idx1) {
|
||||
constexpr auto i_j_idx = make_tuple(idx0, idx1);
|
||||
|
||||
p_compute(i_j_idx) = exp(s[i_j_idx] - m[i_idx]);
|
||||
});
|
||||
});
|
||||
|
||||
// rowsum(Pcompute{j})
|
||||
auto rowsum_p = block_tile_reduce<SMPLComputeDataType>(
|
||||
p_compute, sequence<1>{}, f_sum, SMPLComputeDataType{0});
|
||||
|
||||
block_tile_reduce_sync(rowsum_p, f_sum);
|
||||
|
||||
// l{j}, Oacc{j}
|
||||
sweep_tile_span(p_spans[I0], [&](auto idx0) {
|
||||
constexpr auto i_idx = make_tuple(idx0);
|
||||
|
||||
const auto tmp = exp(m_old[i_idx] - m[i_idx]);
|
||||
|
||||
l(i_idx) = tmp * l[i_idx] + rowsum_p[i_idx];
|
||||
|
||||
sweep_tile_span(p_spans[I1], [&](auto idx1) {
|
||||
constexpr auto i_j_idx = make_tuple(idx0, idx1);
|
||||
|
||||
o_acc(i_j_idx) *= tmp;
|
||||
});
|
||||
});
|
||||
|
||||
block_sync_lds();
|
||||
store_tile(v_lds_window, v_prefetch);
|
||||
move_tile_window(v_dram_window, {0, kK1PerBlock});
|
||||
|
||||
// type cast Pcompute{j} into P{j}
|
||||
const auto p =
|
||||
tile_elementwise_in(type_convert<PDataType, SMPLComputeDataType>, p_compute);
|
||||
|
||||
// Oacc{j}
|
||||
constexpr index_t k1_loops = kN0PerBlock / kK1PerBlock;
|
||||
|
||||
if constexpr(k1_loops > 1)
|
||||
{
|
||||
static_for<0, k1_loops - 1, 1>{}([&](auto i_k1) {
|
||||
const auto v = load_tile(v_dram_window); // load next v
|
||||
block_sync_lds();
|
||||
gemm1(o_acc,
|
||||
get_slice_tile(p,
|
||||
sequence<0, i_k1 * kK1PerBlock>{},
|
||||
sequence<kM0PerBlock, (i_k1 + 1) * kK1PerBlock>{}),
|
||||
v_lds_window);
|
||||
block_sync_lds();
|
||||
store_tile(v_lds_window, v);
|
||||
move_tile_window(v_dram_window, {0, kK1PerBlock});
|
||||
});
|
||||
}
|
||||
// tail
|
||||
{
|
||||
block_sync_lds();
|
||||
gemm1(o_acc,
|
||||
get_slice_tile(p,
|
||||
sequence<0, (k1_loops - 1) * kK1PerBlock>{},
|
||||
sequence<kM0PerBlock, kN0PerBlock>{}),
|
||||
v_lds_window);
|
||||
block_sync_lds();
|
||||
}
|
||||
// move tile windows
|
||||
move_tile_window(k_dram_window, {kN0PerBlock, 0});
|
||||
iN0 += kN0PerBlock;
|
||||
} while(iN0 < N0);
|
||||
|
||||
// Oacc
|
||||
constexpr auto o_spans = decltype(o_acc)::get_distributed_spans();
|
||||
|
||||
sweep_tile_span(o_spans[I0], [&](auto idx0) {
|
||||
constexpr auto i_idx = make_tuple(idx0);
|
||||
|
||||
const auto tmp = 1 / l[i_idx];
|
||||
|
||||
sweep_tile_span(o_spans[I1], [&](auto idx1) {
|
||||
constexpr auto i_j_idx = make_tuple(idx0, idx1);
|
||||
|
||||
o_acc(i_j_idx) *= tmp;
|
||||
});
|
||||
});
|
||||
|
||||
// type cast Oacc into O
|
||||
const auto o = tile_elementwise_in(type_convert<ODataType, OaccDataType>, o_acc);
|
||||
|
||||
// O DRAM and O DRAM window
|
||||
auto o_dram = make_naive_tensor_view<address_space_enum::global>(
|
||||
o_ptr, make_tuple(M0, N1), make_tuple(StrideO, 1), number<32>{}, number<1>{});
|
||||
|
||||
auto o_dram_window =
|
||||
make_tile_window(o_dram,
|
||||
make_tuple(number<kM0PerBlock>{}, number<kN1PerBlock>{}),
|
||||
{iM0, iN1},
|
||||
o.get_tile_distribution());
|
||||
|
||||
// store O
|
||||
store_tile(o_dram_window, o);
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace ck_tile
|
||||
@@ -0,0 +1,37 @@
|
||||
// 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/host/host_tensor.hpp"
|
||||
|
||||
template <typename ADataType, typename BDataType, typename AccDataType, typename CDataType>
|
||||
void reference_batched_gemm(const ck_tile::HostTensor<ADataType>& a_b_m_k,
|
||||
const ck_tile::HostTensor<BDataType>& b_b_n_k,
|
||||
ck_tile::HostTensor<CDataType>& c_b_m_n)
|
||||
{
|
||||
const int N = b_b_n_k.mDesc.get_lengths()[1];
|
||||
const int K = b_b_n_k.mDesc.get_lengths()[2];
|
||||
|
||||
auto f = [&](auto batch, auto m) {
|
||||
for(int n = 0; n < N; ++n)
|
||||
{
|
||||
AccDataType v_acc = 0;
|
||||
|
||||
for(int k = 0; k < K; ++k)
|
||||
{
|
||||
ADataType v_a = a_b_m_k(batch, m, k);
|
||||
BDataType v_b = b_b_n_k(batch, n, k);
|
||||
|
||||
v_acc += ck_tile::type_convert<AccDataType>(v_a) * ck_tile::type_convert<AccDataType>(v_b);
|
||||
}
|
||||
|
||||
c_b_m_n(batch, m, n) = ck_tile::type_convert<CDataType>(v_acc);
|
||||
}
|
||||
};
|
||||
|
||||
ck_tile::make_ParallelTensorFunctor(f, c_b_m_n.mDesc.get_lengths()[0], c_b_m_n.mDesc.get_lengths()[1])(
|
||||
std::thread::hardware_concurrency());
|
||||
}
|
||||
|
||||
@@ -0,0 +1,47 @@
|
||||
// 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/host/host_tensor.hpp"
|
||||
|
||||
template <typename ADataType, typename AccDataType, typename BDataType>
|
||||
void reference_batched_softmax(const ck_tile::HostTensor<ADataType>& a_b_m_n, ck_tile::HostTensor<BDataType>& b_b_m_n)
|
||||
{
|
||||
const int N = a_b_m_n.mDesc.get_lengths()[2];
|
||||
|
||||
auto f = [&](auto batch, auto m) {
|
||||
AccDataType v_max = std::numeric_limits<ADataType>::lowest();
|
||||
|
||||
// max
|
||||
for(int n = 0; n < N; ++n)
|
||||
{
|
||||
const ADataType v_a = a_b_m_n(batch, m, n);
|
||||
|
||||
v_max = v_max < v_a ? v_a : v_max;
|
||||
}
|
||||
|
||||
AccDataType v_exp_sum = 0;
|
||||
|
||||
// sum
|
||||
for(int n = 0; n < N; ++n)
|
||||
{
|
||||
const ADataType v_a = a_b_m_n(batch, m, n);
|
||||
|
||||
v_exp_sum += ck_tile::exp(v_a - v_max);
|
||||
}
|
||||
|
||||
// elementwise
|
||||
for(int n = 0; n < N; ++n)
|
||||
{
|
||||
const ADataType v_a = a_b_m_n(batch, m, n);
|
||||
|
||||
b_b_m_n(batch, m, n) = ck_tile::exp(v_a - v_max) / v_exp_sum;
|
||||
}
|
||||
};
|
||||
|
||||
ck_tile::make_ParallelTensorFunctor(f, b_b_m_n.mDesc.get_lengths()[0], b_b_m_n.mDesc.get_lengths()[1])(
|
||||
std::thread::hardware_concurrency());
|
||||
}
|
||||
|
||||
@@ -0,0 +1,18 @@
|
||||
// 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 {
|
||||
|
||||
template <index_t kMPerTile, index_t kNPerTile, index_t kKPerTile>
|
||||
struct TileGemmShape
|
||||
{
|
||||
static constexpr index_t kM = kMPerTile;
|
||||
static constexpr index_t kN = kNPerTile;
|
||||
static constexpr index_t kK = kKPerTile;
|
||||
};
|
||||
|
||||
} // namespace ck_tile
|
||||
@@ -5,3 +5,4 @@ include_directories(AFTER
|
||||
add_subdirectory(01_add)
|
||||
add_subdirectory(02_gemm)
|
||||
add_subdirectory(03_flash_attention_fwd)
|
||||
add_subdirectory(04_codegen_flash_attention_fwd)
|
||||
|
||||
Reference in New Issue
Block a user