From 0cc513081884518e4ba2e2bf3ce6809b698918a0 Mon Sep 17 00:00:00 2001 From: Clement Lin Date: Thu, 10 Apr 2025 22:30:10 +0800 Subject: [PATCH] Add codegen test example --- .../CMakeLists.txt | 19 + .../block_gemm_areg_bsmem_creg_problem.hpp | 24 + .../block_gemm_areg_bsmem_creg_v1.hpp | 300 ++++++++++ ...gemm_areg_bsmem_creg_v1_default_policy.hpp | 20 + ...emm_areg_bsmem_creg_v1_iteratek_policy.hpp | 20 + ..._pipeline_agmem_bgmem_creg_v2_askiplds.hpp | 545 ++++++++++++++++++ ...ne_agmem_bgmem_creg_v2_askiplds_policy.hpp | 96 +++ .../block_gemm_pipeline_problem.hpp | 26 + ..._pipeline_agmem_bgmem_creg_policy_impl.hpp | 180 ++++++ .../flash_attention_fwd.cpp | 202 +++++++ .../flash_attention_fwd.hpp | 109 ++++ .../flash_attention_fwd_impl.hpp | 355 ++++++++++++ .../reference_batched_gemm.hpp | 37 ++ .../reference_batched_softmax.hpp | 47 ++ .../tile_gemm_shape.hpp | 18 + example/ck_tile/99_toy_example/CMakeLists.txt | 1 + 16 files changed, 1999 insertions(+) create mode 100755 example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/CMakeLists.txt create mode 100644 example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/block_gemm_areg_bsmem_creg_problem.hpp create mode 100644 example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/block_gemm_areg_bsmem_creg_v1.hpp create mode 100644 example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/block_gemm_areg_bsmem_creg_v1_default_policy.hpp create mode 100644 example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/block_gemm_areg_bsmem_creg_v1_iteratek_policy.hpp create mode 100644 example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/block_gemm_pipeline_agmem_bgmem_creg_v2_askiplds.hpp create mode 100644 example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/block_gemm_pipeline_agmem_bgmem_creg_v2_askiplds_policy.hpp create mode 100644 example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/block_gemm_pipeline_problem.hpp create mode 100644 example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/blockgemm_pipeline_agmem_bgmem_creg_policy_impl.hpp create mode 100644 example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/flash_attention_fwd.cpp create mode 100644 example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/flash_attention_fwd.hpp create mode 100644 example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/flash_attention_fwd_impl.hpp create mode 100644 example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/reference_batched_gemm.hpp create mode 100644 example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/reference_batched_softmax.hpp create mode 100644 example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/tile_gemm_shape.hpp diff --git a/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/CMakeLists.txt b/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/CMakeLists.txt new file mode 100755 index 0000000000..93a1364b37 --- /dev/null +++ b/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/CMakeLists.txt @@ -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) diff --git a/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/block_gemm_areg_bsmem_creg_problem.hpp b/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/block_gemm_areg_bsmem_creg_problem.hpp new file mode 100644 index 0000000000..a0a74c55aa --- /dev/null +++ b/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/block_gemm_areg_bsmem_creg_problem.hpp @@ -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 +struct BlockGemmARegBSmemCRegProblem +{ + using ADataType = remove_cvref_t; + using BDataType = remove_cvref_t; + using CDataType = remove_cvref_t; + using BlockGemmShape = remove_cvref_t; + + static constexpr index_t kBlockSize = kBlockSize_; +}; + +} // namespace ck_tile diff --git a/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/block_gemm_areg_bsmem_creg_v1.hpp b/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/block_gemm_areg_bsmem_creg_v1.hpp new file mode 100644 index 0000000000..740c540d6c --- /dev/null +++ b/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/block_gemm_areg_bsmem_creg_v1.hpp @@ -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 +struct BlockGemmARegBSmemCRegV1 +{ + using ADataType = remove_cvref_t; + using BDataType = remove_cvref_t; + using CDataType = remove_cvref_t; + using BlockGemmShape = remove_cvref_t; + using BlockGemmPolicy = Policy; + + static constexpr index_t kBlockSize = Problem::kBlockSize; + + // C += A * B + template + __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> && + std::is_same_v> && + std::is_same_v>, + "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(); + + using WG = remove_cvref_t())>; + + 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, + tuple, sequence>, + tuple>, + tuple>, + sequence<1, 2>, + sequence<0, 0>>{}; + + constexpr auto c_block_outer_dstr_encoding = tile_distribution_encoding< + sequence<>, + tuple, sequence>, + tuple>, + tuple>, + 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(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{}, number{}), + {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, + 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>, "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{}; + constexpr auto c_warp_y_index_zeros = uniform_sequence_gen_t{}; + + // 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{}, 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{}, 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{}, c_warp_y_index_zeros), + merge_sequences(sequence<1, 1>{}, c_warp_y_lengths), + c_warp_tensor.get_thread_buffer()); + }); + }); + }); + } + + // C = A * B + template + __device__ auto operator()(const ABlockTensorTmp& a_block_tensor_tmp, + const BBlockWindowTmp& b_block_window_tmp) const + { + static_assert(std::is_same_v> && + std::is_same_v>, + "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(); + + using WG = remove_cvref_t())>; + + 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, + tuple, sequence>, + tuple>, + tuple>, + sequence<1, 2>, + sequence<0, 0>>{}; + + constexpr auto c_block_outer_dstr_encoding = tile_distribution_encoding< + sequence<>, + tuple, sequence>, + tuple>, + tuple>, + 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(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{}, number{}), + {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, + 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(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{}; + constexpr auto c_warp_y_index_zeros = uniform_sequence_gen_t{}; + + // 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{}, 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{}, 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{}, 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 diff --git a/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/block_gemm_areg_bsmem_creg_v1_default_policy.hpp b/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/block_gemm_areg_bsmem_creg_v1_default_policy.hpp new file mode 100644 index 0000000000..fb1516eb52 --- /dev/null +++ b/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/block_gemm_areg_bsmem_creg_v1_default_policy.hpp @@ -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 + CK_TILE_HOST_DEVICE static constexpr auto GetWarpGemmMWarpNWarp() + { + return make_tuple(WarpGemmMfmaF16F16F32M32N32K8TransposedCDistribution{}, 4, 1); + } +}; + +} // namespace ck_tile diff --git a/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/block_gemm_areg_bsmem_creg_v1_iteratek_policy.hpp b/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/block_gemm_areg_bsmem_creg_v1_iteratek_policy.hpp new file mode 100644 index 0000000000..32dc09f95e --- /dev/null +++ b/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/block_gemm_areg_bsmem_creg_v1_iteratek_policy.hpp @@ -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 + CK_TILE_HOST_DEVICE static constexpr auto GetWarpGemmMWarpNWarp() + { + return make_tuple(WarpGemmMfmaF16F16F32M32N32K16TransposedCDistribution{}, 4, 1); + } +}; + +} // namespace ck_tile diff --git a/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/block_gemm_pipeline_agmem_bgmem_creg_v2_askiplds.hpp b/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/block_gemm_pipeline_agmem_bgmem_creg_v2_askiplds.hpp new file mode 100644 index 0000000000..fc98792615 --- /dev/null +++ b/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/block_gemm_pipeline_agmem_bgmem_creg_v2_askiplds.hpp @@ -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 +struct BlockGemmPipelineAGmemBGmemCReg +{ + using ADataType = remove_cvref_t; + using BDataType = remove_cvref_t; + using CDataType = remove_cvref_t; + using BlockGemmShape = remove_cvref_t; + 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().get_element_space_size(); + } + + template + __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> && + std::is_same_v>, + "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(); + + // B tile in LDS, blockWindow + BDataType* p_b_lds = + static_cast(static_cast(static_cast(p_smem))); + + constexpr auto b_lds_block_desc = Policy::template MakeBLdsBlockDescriptor(); + + // 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(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{}, number{}), + a_dram_block_window_tmp.get_window_origin(), + Policy::template MakeADramTileDistribution()); + + // A Reg tensor for store, also used for block GEMM + auto a_copy_reg_tensor = make_static_distributed_tensor(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{}, number{}), + b_dram_block_window_tmp.get_window_origin(), + Policy::template MakeBDramTileDistribution()); + + // B LDS tile window for store + auto b_copy_lds_window = + make_tile_window(b_lds_block, + make_tuple(number{}, number{}), + {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{}, number{}), {0, 0}); + + // Block GEMM + constexpr auto block_gemm = Policy::template GetBlockGemm(); + + // 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 + __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 +struct BlockGemmPipelineAGmemBGmemCReg< + Problem, + BlockGemmPipelineAGmemBGmemCRegSkipALdsPersistentQRegCachePolicy> +{ + using ADataType = remove_cvref_t; + using BDataType = remove_cvref_t; + using CDataType = remove_cvref_t; + using BlockGemmShape = remove_cvref_t; + using Policy = BlockGemmPipelineAGmemBGmemCRegSkipALdsPersistentQRegCachePolicy; + + 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().get_element_space_size(); + } + + // Cold A Register Cache + template + __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> && + std::is_same_v>, + "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(); + + // B tile in LDS, blockWindow + BDataType* p_b_lds = + static_cast(static_cast(static_cast(p_smem))); + + constexpr auto b_lds_block_desc = Policy::template MakeBLdsBlockDescriptor(); + + // 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(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{}, number{}), + a_dram_block_window_tmp.get_window_origin(), + Policy::template MakeADramTileDistribution()); + + // A Reg tensor for store, also used for block GEMM + auto a_copy_reg_tensor = make_static_distributed_tensor(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{}, number{}), + b_dram_block_window_tmp.get_window_origin(), + Policy::template MakeBDramTileDistribution()); + + // B LDS tile window for store + auto b_copy_lds_window = + make_tile_window(b_lds_block, + make_tuple(number{}, number{}), + {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{}, number{}), {0, 0}); + + // Block GEMM + constexpr auto block_gemm = Policy::template GetBlockGemm(); + + // Acc register tile + auto c_block_tile = decltype(block_gemm( + get_slice_tile(a_copy_reg_tensor, sequence<0, 0>{}, sequence{}), + 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{}); + 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{}), + 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{}); + 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{}), + b_copy_lds_window); + + block_sync_lds(); + + set_slice_tile(a_copy_reg_tensor, + a_block_tile, + sequence<0, (k_loops - 1) * kKPerBlock>{}, + sequence{}); + + 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{}), + 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{}); + + return c_block_tile; + } + + // Hot A Register Cache + template + __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>, + "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(); + + // A Reg tensor for store, also used for block GEMM + auto a_copy_reg_tensor = make_static_distributed_tensor(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{}); + + // B tile in LDS, blockWindow + BDataType* p_b_lds = + static_cast(static_cast(static_cast(p_smem))); + + constexpr auto b_lds_block_desc = Policy::template MakeBLdsBlockDescriptor(); + + // 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(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{}, number{}), + b_dram_block_window_tmp.get_window_origin(), + Policy::template MakeBDramTileDistribution()); + + // B LDS tile window for store + auto b_copy_lds_window = + make_tile_window(b_lds_block, + make_tuple(number{}, number{}), + {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{}, number{}), {0, 0}); + + // Block GEMM + constexpr auto block_gemm = Policy::template GetBlockGemm(); + + // Acc register tile + auto c_block_tile = decltype(block_gemm( + get_slice_tile(a_copy_reg_tensor, sequence<0, 0>{}, sequence{}), + 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{}), + 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{}), + 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{}), + b_copy_lds_window); + } + + return c_block_tile; + } + + template + __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 + __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 diff --git a/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/block_gemm_pipeline_agmem_bgmem_creg_v2_askiplds_policy.hpp b/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/block_gemm_pipeline_agmem_bgmem_creg_v2_askiplds_policy.hpp new file mode 100644 index 0000000000..cdce1b1f31 --- /dev/null +++ b/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/block_gemm_pipeline_agmem_bgmem_creg_v2_askiplds_policy.hpp @@ -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 + __host__ __device__ static constexpr auto MakeARegBlockDescriptor() + { + constexpr auto blockgemm = GetBlockGemm(); + using BlockGemm = remove_cvref_t; + + return policy_impl::make_a_reg_block_descriptor(); + } + + template + __host__ __device__ static constexpr auto MakeBLdsBlockDescriptor() + { + return policy_impl::make_b_lds_block_descriptor_3d_pad(); + } + + template + __host__ __device__ static constexpr auto MakeADramTileDistribution() + { + constexpr auto blockgemm = GetBlockGemm(); + using BlockGemm = remove_cvref_t; + + return policy_impl::make_a_dram_tile_distribution_skip_lds(); + } + + template + __host__ __device__ static constexpr auto MakeBDramTileDistribution() + { + return policy_impl::make_b_dram_tile_distribution(); + } + + template + __host__ __device__ static constexpr auto GetBlockGemm() + { + using BlockGemmPolicy = BlockGemmARegBSmemCRegV1K8Policy; + + return BlockGemmARegBSmemCRegV1{}; + } +}; + +template +struct BlockGemmPipelineAGmemBGmemCRegSkipALdsPersistentQRegCachePolicy + : BlockGemmPipelineAGmemBGmemCRegSkipALdsPolicy +{ + static constexpr index_t AKDim = AKDim_; + + template + __host__ __device__ static constexpr auto MakeARegBlockDescriptor() + { + constexpr auto blockgemm = GetBlockGemm(); + using BlockGemm = remove_cvref_t; + + constexpr index_t kMPerBlock = Problem::BlockGemmShape::kM; + constexpr index_t kKPerBlock = AKDim; + + constexpr auto config = + BlockGemm::BlockGemmPolicy::template GetWarpGemmMWarpNWarp(); + + using WG = remove_cvref_t())>; + + 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, + tuple, sequence>, + tuple>, + tuple>, + 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 diff --git a/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/block_gemm_pipeline_problem.hpp b/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/block_gemm_pipeline_problem.hpp new file mode 100644 index 0000000000..1a620ba54b --- /dev/null +++ b/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/block_gemm_pipeline_problem.hpp @@ -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 +struct BlockGemmPipelineProblem +{ + using ADataType = remove_cvref_t; + using BDataType = remove_cvref_t; + using CDataType = remove_cvref_t; + using BlockGemmShape = remove_cvref_t; + + static constexpr index_t kBlockSize = kBlockSize_; +}; + +} // namespace ck_tile + diff --git a/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/blockgemm_pipeline_agmem_bgmem_creg_policy_impl.hpp b/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/blockgemm_pipeline_agmem_bgmem_creg_policy_impl.hpp new file mode 100644 index 0000000000..ad6d6d3996 --- /dev/null +++ b/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/blockgemm_pipeline_agmem_bgmem_creg_policy_impl.hpp @@ -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 +__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{}, number{}, 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 +__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{}, number{}, 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 +__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(); + + using WG = remove_cvref_t())>; + + 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, + tuple, sequence>, + tuple>, + tuple>, + 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 +__host__ __device__ static constexpr auto make_a_dram_tile_distribution() +{ + using ADataType = remove_cvref_t; + + 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, + tuple, sequence>, + tuple, sequence<1, 2>>, + tuple, sequence<2, 0>>, + sequence<1, 2>, + sequence<0, 1>>{}); +} + +template +__host__ __device__ static constexpr auto make_a_dram_tile_distribution_skip_lds() +{ + constexpr auto config = BlockGemm::BlockGemmPolicy::template GetWarpGemmMWarpNWarp(); + + using WG = remove_cvref_t())>; + + 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, + tuple, sequence>, + tuple, sequence<2, 1>>, + tuple, sequence<1, 2>>, + sequence<2, 1, 2>, + sequence<0, 0, 2>>{}); +} + +template +__host__ __device__ static constexpr auto make_b_dram_tile_distribution() +{ + using BDataType = remove_cvref_t; + + 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, + tuple, sequence>, + tuple, sequence<1, 2>>, + tuple, sequence<2, 0>>, + sequence<1, 2>, + sequence<0, 1>>{}); +} + +template +__host__ __device__ static constexpr auto get_block_gemm() +{ + using BlockGemmPolicy = BlockGemmASmemBSmemCRegDefaultPolicy; + + return BlockGemmASmemBSmemCReg{}; +} + +} // namespace policy_impl +} // namespace ck_tile diff --git a/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/flash_attention_fwd.cpp b/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/flash_attention_fwd.cpp new file mode 100644 index 0000000000..8ce1d6c6c7 --- /dev/null +++ b/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/flash_attention_fwd.cpp @@ -0,0 +1,202 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +#include + +#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 q_lengths{Batch, M0, K0}; + std::array q_strides{M0 * K0, K0, 1}; + + std::array k_lengths{Batch, N0, K0}; + std::array k_strides{N0 * K0, K0, 1}; + + std::array v_lengths{Batch, N1, N0}; + std::array v_strides{N1 * N0, N0, 1}; + + std::array s_lengths{Batch, M0, N0}; + std::array s_strides{M0 * N0, N0, 1}; + + std::array p_lengths{Batch, M0, N0}; + std::array p_strides{M0 * N0, N0, 1}; + + std::array o_lengths{Batch, M0, N1}; + std::array o_strides{M0 * N1, N1, 1}; + + // host verify + ck_tile::HostTensor q_host(q_lengths, q_strides); + ck_tile::HostTensor k_host(k_lengths, k_strides); + ck_tile::HostTensor v_host(v_lengths, v_strides); + ck_tile::HostTensor o_host_dev(o_lengths, o_strides); + + switch(init_method) + { + case 0: break; + case 1: + ck_tile::FillUniformDistributionIntegerValue{-3.f, 3.f}(q_host); + ck_tile::FillUniformDistributionIntegerValue{-3.f, 3.f}(k_host); + ck_tile::FillUniformDistributionIntegerValue{-3.f, 3.f}(v_host); + break; + case 2: + ck_tile::FillUniformDistribution{-3.f, 3.f}(q_host); + ck_tile::FillUniformDistribution{-3.f, 3.f}(k_host); + ck_tile::FillUniformDistribution{-3.f, 3.f}(v_host); + break; + default: + ck_tile::FillUniformDistributionIntegerValue{-2.f, 2.f}(q_host); + ck_tile::FillUniformDistributionIntegerValue{-2.f, 2.f}(k_host); + ck_tile::FillUniformDistributionIntegerValue{-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( + ck_tile::FlashAttentionFwd{}, + kGridSize, + kBlockSize, + 0, + static_cast(q_buf.GetDeviceBuffer()), + static_cast(k_buf.GetDeviceBuffer()), + static_cast(v_buf.GetDeviceBuffer()), + static_cast(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 s_host_ref(s_lengths, s_strides); + ck_tile::HostTensor p_host_ref(p_lengths, p_strides); + ck_tile::HostTensor o_host_ref(o_lengths, o_strides); + + ck_tile::reference_batched_gemm( + q_host, k_host, s_host_ref); + ck_tile::reference_batched_softmax(s_host_ref, + p_host_ref); + ck_tile::reference_batched_gemm( + 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(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; +} + diff --git a/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/flash_attention_fwd.hpp b/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/flash_attention_fwd.hpp new file mode 100644 index 0000000000..caeeece8e9 --- /dev/null +++ b/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/flash_attention_fwd.hpp @@ -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 +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{}; + + 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 diff --git a/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/flash_attention_fwd_impl.hpp b/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/flash_attention_fwd_impl.hpp new file mode 100644 index 0000000000..4229db5250 --- /dev/null +++ b/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/flash_attention_fwd_impl.hpp @@ -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 +struct FlashAttentionFwdImpl +{ + // block gemm0 pipeline + using BlockGemm0Problem = BlockGemmPipelineProblem< + QDataType, + KDataType, + SaccDataType, + kBlockSize, + TileGemmShape>; + + using BlockGemm0Policy = + BlockGemmPipelineAGmemBGmemCRegSkipALdsPersistentQRegCachePolicy; + + using BlockGemm0Pipeline = + BlockGemmPipelineAGmemBGmemCReg; + + // block gemm1 + using BlockGemm1 = BlockGemmARegBSmemCRegV1< + BlockGemmARegBSmemCRegProblem< + PDataType, + VDataType, + OaccDataType, + kBlockSize, + TileGemmShape>, + 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{}, number{}, number{}), + make_tuple(number<(kNPerBlock + kPad) * kK1>{}, number{}, number<1>{}), + number{}, + 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{}, number{}))), + 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, + tuple, sequence>, + tuple, sequence<1, 2>>, + tuple, sequence<2, 0>>, + sequence<1, 2>, + sequence<0, 1>>{}); + } + + __device__ static constexpr index_t GetStaticLdsSize() + { + return max(BlockGemm0Pipeline::GetStaticLdsSize(), + static_cast(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( + q_ptr, make_tuple(M0, K0), make_tuple(StrideQ, 1), number<32>{}, number<1>{}); + + const auto k_dram = make_naive_tensor_view( + k_ptr, make_tuple(N0, K0), make_tuple(StrideK, 1), number<32>{}, number<1>{}); + + const auto v_dram = make_naive_tensor_view( + 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{}, number{}), {iM0, 0}); + + auto k_dram_window = make_tile_window( + k_dram, make_tuple(number{}, number{}), {0, 0}); + + auto v_dram_window = + make_tile_window(v_dram, + make_tuple(number{}, number{}), + {iN1, 0}, + MakeVDramTileDistribution()); + + // Q in Register + auto q_reg_tensor = make_static_distributed_tensor( + BlockGemm0Policy::template MakeARegBlockDescriptor()); + + // V LDS and LDS window + // V LDS occupies the same LDS allocation Q/K LDS + auto v_lds = make_tensor_view(reinterpret_cast(smem_ptr), + MakeVLdsBlockDescriptor()); + + auto v_lds_window = make_tile_window( + v_lds, make_tuple(number{}, number{}), {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, SaccBlockTileType{})); + + using PBlockTileType = decltype(tile_elementwise_in(type_convert, + SaccBlockTileType{})); + + using MLBlockTileType = decltype(block_tile_reduce( + SBlockTileType{}, sequence<1>{}, f_max, SMPLComputeDataType{0})); + + using OaccBlockTileType = decltype(gemm1( + get_slice_tile( + PBlockTileType{}, sequence<0, 0>{}, sequence{}), + 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::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, 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( + s, sequence<1>{}, f_max, std::numeric_limits::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(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( + 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, 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{}), + 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{}), + 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, o_acc); + + // O DRAM and O DRAM window + auto o_dram = make_naive_tensor_view( + 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{}, number{}), + {iM0, iN1}, + o.get_tile_distribution()); + + // store O + store_tile(o_dram_window, o); + } +}; + +} // namespace ck_tile diff --git a/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/reference_batched_gemm.hpp b/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/reference_batched_gemm.hpp new file mode 100644 index 0000000000..2762e66464 --- /dev/null +++ b/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/reference_batched_gemm.hpp @@ -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 +void reference_batched_gemm(const ck_tile::HostTensor& a_b_m_k, + const ck_tile::HostTensor& b_b_n_k, + ck_tile::HostTensor& 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(v_a) * ck_tile::type_convert(v_b); + } + + c_b_m_n(batch, m, n) = ck_tile::type_convert(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()); +} + diff --git a/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/reference_batched_softmax.hpp b/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/reference_batched_softmax.hpp new file mode 100644 index 0000000000..3713a22c6a --- /dev/null +++ b/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/reference_batched_softmax.hpp @@ -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 +void reference_batched_softmax(const ck_tile::HostTensor& a_b_m_n, ck_tile::HostTensor& 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::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()); +} + diff --git a/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/tile_gemm_shape.hpp b/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/tile_gemm_shape.hpp new file mode 100644 index 0000000000..02a7106eb7 --- /dev/null +++ b/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/tile_gemm_shape.hpp @@ -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 +struct TileGemmShape +{ + static constexpr index_t kM = kMPerTile; + static constexpr index_t kN = kNPerTile; + static constexpr index_t kK = kKPerTile; +}; + +} // namespace ck_tile diff --git a/example/ck_tile/99_toy_example/CMakeLists.txt b/example/ck_tile/99_toy_example/CMakeLists.txt index adecf1a1c1..cadbb1c06f 100644 --- a/example/ck_tile/99_toy_example/CMakeLists.txt +++ b/example/ck_tile/99_toy_example/CMakeLists.txt @@ -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)