From 4d8fce33dddfc003432ae06848f6416a9d5d5e2f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Bart=C5=82omiej=20Kocot?= Date: Fri, 13 Dec 2024 21:08:35 +0100 Subject: [PATCH 01/25] Add SplitK support into Batched GEMM V3 (#1729) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * add bmm api * add bf16 multi_d * add ckProfiler for bf16 * add ckProfiler files * add more instance; fixed 64bit index issue * fixed naming * enabled batched Ds * use long_index for ds offsets * clean * add bmm fp8 ckProfiler * Update example/24_batched_gemm/batched_gemm_xdl_bf16_v3.cpp Co-authored-by: Bartłomiej Kocot * Update example/24_batched_gemm/batched_gemm_xdl_fp8_rowwise_v3.cpp Co-authored-by: Bartłomiej Kocot * Update example/24_batched_gemm/run_batched_gemm_example_rowwise.inc Co-authored-by: Bartłomiej Kocot * Update library/src/tensor_operation_instance/gpu/gemm_universal_batched/device_batched_gemm_xdl_universal_bf16_bf16_bf16/device_batched_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn.hpp Co-authored-by: Bartłomiej Kocot * Update library/src/tensor_operation_instance/gpu/gemm_universal_batched/device_batched_gemm_xdl_universal_bf16_bf16_bf16/device_batched_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_mem_v1_default_instance.cpp Co-authored-by: Bartłomiej Kocot * Update library/src/tensor_operation_instance/gpu/gemm_universal_batched/device_batched_gemm_xdl_universal_bf16_bf16_bf16/device_batched_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_mem_v2_default_instance.cpp Co-authored-by: Bartłomiej Kocot * Update profiler/src/profile_gemm_universal_batched.cpp Co-authored-by: Bartłomiej Kocot * Update profiler/include/profiler/profile_gemm_universal_batched_impl.hpp Co-authored-by: Bartłomiej Kocot * clean * Update include/ck/tensor_operation/gpu/device/impl/device_batched_gemm_multiple_d_xdl_cshuffle_v3.hpp * Update include/ck/tensor_operation/gpu/device/impl/device_batched_gemm_multiple_d_xdl_cshuffle_v3.hpp * Update library/src/tensor_operation_instance/gpu/gemm_universal_batched/device_batched_gemm_xdl_universal_bf16_bf16_bf16/device_batched_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_comp_default_instance.cpp * Update include/ck/tensor_operation/gpu/device/impl/device_batched_gemm_multiple_d_xdl_cshuffle_v3.hpp * Update include/ck/tensor_operation/gpu/device/impl/device_batched_gemm_multiple_d_xdl_cshuffle_v3.hpp * Update include/ck/tensor_operation/gpu/device/impl/device_batched_gemm_multiple_d_xdl_cshuffle_v3.hpp * refactor batch offset func * add splitk suppport into bmm_v3 * clean * clean * format * fixed * fix --------- Co-authored-by: Jing Zhang Co-authored-by: zjing14 --- .../batched_gemm_xdl_bf16_v3.cpp | 4 +- .../device/device_batched_gemm_multi_d.hpp | 3 +- ...atched_gemm_multiple_d_xdl_cshuffle_v3.hpp | 45 ++++-- .../gridwise_gemm_xdl_cshuffle_v3_multi_d.hpp | 16 +- ..._xdl_universal_bf16_bf16_bf16_mk_nk_mn.hpp | 3 + ...gemm_xdl_universal_f8_f8_bf16_mk_nk_mn.hpp | 2 + .../profile_gemm_universal_batched_impl.hpp | 148 ++++++++++-------- .../src/profile_gemm_universal_batched.cpp | 20 +-- 8 files changed, 137 insertions(+), 104 deletions(-) diff --git a/example/24_batched_gemm/batched_gemm_xdl_bf16_v3.cpp b/example/24_batched_gemm/batched_gemm_xdl_bf16_v3.cpp index fa8b752185..548500518f 100644 --- a/example/24_batched_gemm/batched_gemm_xdl_bf16_v3.cpp +++ b/example/24_batched_gemm/batched_gemm_xdl_bf16_v3.cpp @@ -78,14 +78,14 @@ using DeviceGemmInstance = ck::tensor_operation::device::DeviceBatchedGemmMultiD 2, // ABlockTransferSrcVectorDim 8, // ABlockTransferSrcScalarPerVector 8, // ABlockTransferDstScalarPerVector_AK1 - 1, // ABlockLdsExtraM + 0, // ABlockLdsExtraM S<4, 64, 1>, // BBlockTransferThreadClusterLengths_BK0_N_BK1 S<1, 0, 2>, // BBlockTransferThreadClusterArrangeOrder S<1, 0, 2>, // BBlockTransferSrcAccessOrder 2, // BBlockTransferSrcVectorDim 8, // BBlockTransferSrcScalarPerVector 8, // BBlockTransferDstScalarPerVector_BK1 - 1, // BBlockLdsExtraN + 0, // BBlockLdsExtraN 1, // CShuffleMXdlPerWavePerShuffle 1, // CShuffleNXdlPerWavePerShuffle S<1, 32, 1, 8>, // CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock diff --git a/include/ck/tensor_operation/gpu/device/device_batched_gemm_multi_d.hpp b/include/ck/tensor_operation/gpu/device/device_batched_gemm_multi_d.hpp index 58c0288e8f..8fb4a71f55 100644 --- a/include/ck/tensor_operation/gpu/device/device_batched_gemm_multi_d.hpp +++ b/include/ck/tensor_operation/gpu/device/device_batched_gemm_multi_d.hpp @@ -89,7 +89,8 @@ struct DeviceBatchedGemmV2MultiD : public BaseOperator index_t BatchStrideE, AElementwiseOperation a_element_op, BElementwiseOperation b_element_op, - CDEElementwiseOperation cde_element_op) = 0; + CDEElementwiseOperation cde_element_op, + index_t KBatch) = 0; virtual std::unique_ptr MakeInvokerPointer() = 0; }; diff --git a/include/ck/tensor_operation/gpu/device/impl/device_batched_gemm_multiple_d_xdl_cshuffle_v3.hpp b/include/ck/tensor_operation/gpu/device/impl/device_batched_gemm_multiple_d_xdl_cshuffle_v3.hpp index 314ecdf76e..5f5bea4f86 100644 --- a/include/ck/tensor_operation/gpu/device/impl/device_batched_gemm_multiple_d_xdl_cshuffle_v3.hpp +++ b/include/ck/tensor_operation/gpu/device/impl/device_batched_gemm_multiple_d_xdl_cshuffle_v3.hpp @@ -41,12 +41,15 @@ __global__ void __shared__ char p_shared[GridwiseGemm::GetSharedMemoryNumberOfByte()]; const index_t g_idx = blockIdx.z % karg.Batch; + const index_t k_idx = blockIdx.z / karg.Batch; const auto a_batch_offset = karg.compute_ptr_offset_of_batch.GetAPtrOffset(g_idx); const auto b_batch_offset = karg.compute_ptr_offset_of_batch.GetBPtrOffset(g_idx); const auto ds_batch_offset = karg.compute_ptr_offset_of_batch.GetDsPtrOffset(g_idx); const auto c_batch_offset = karg.compute_ptr_offset_of_batch.GetCPtrOffset(g_idx); + auto splitk_batch_offset = typename GridwiseGemm::SplitKBatchOffset(karg, k_idx); + // populate pointer, desc for Ds static_for<0, GridwiseGemm::NumDTensor, 1>{}([&](auto i) { // D pointer @@ -54,8 +57,8 @@ __global__ void }); GridwiseGemm::template Run( - karg.p_a_grid + a_batch_offset, - karg.p_b_grid + b_batch_offset, + karg.p_a_grid + a_batch_offset + splitk_batch_offset.a_k_split_offset, + karg.p_b_grid + b_batch_offset + splitk_batch_offset.b_k_split_offset, karg.p_ds_grid, karg.p_c_grid + c_batch_offset, p_shared, @@ -87,12 +90,15 @@ __global__ void __shared__ char p_shared_1[GridwiseGemm::GetSharedMemoryNumberOfByte()]; const index_t g_idx = blockIdx.z % karg.Batch; + const index_t k_idx = blockIdx.z / karg.Batch; const auto a_batch_offset = karg.compute_ptr_offset_of_batch.GetAPtrOffset(g_idx); const auto b_batch_offset = karg.compute_ptr_offset_of_batch.GetBPtrOffset(g_idx); const auto ds_batch_offset = karg.compute_ptr_offset_of_batch.GetDsPtrOffset(g_idx); const auto c_batch_offset = karg.compute_ptr_offset_of_batch.GetCPtrOffset(g_idx); + auto splitk_batch_offset = typename GridwiseGemm::SplitKBatchOffset(karg, k_idx); + // populate pointer, desc for Ds static_for<0, GridwiseGemm::NumDTensor, 1>{}([&](auto i) { // D pointer @@ -100,8 +106,8 @@ __global__ void }); GridwiseGemm::template Run_2Lds( - karg.p_a_grid + a_batch_offset, - karg.p_b_grid + b_batch_offset, + karg.p_a_grid + a_batch_offset + splitk_batch_offset.a_k_split_offset, + karg.p_b_grid + b_batch_offset + splitk_batch_offset.b_k_split_offset, karg.p_ds_grid, karg.p_c_grid + c_batch_offset, p_shared_0, @@ -303,7 +309,8 @@ struct DeviceBatchedGemmMultiD_Xdl_CShuffle_V3 index_t Batch_, AElementwiseOperation a_element_op_, BElementwiseOperation b_element_op_, - CElementwiseOperation c_element_op_) + CElementwiseOperation c_element_op_, + index_t KBatch_) : GridwiseGemm::Argument{p_a_grid_, p_b_grid_, p_ds_grid_, @@ -315,7 +322,7 @@ struct DeviceBatchedGemmMultiD_Xdl_CShuffle_V3 StrideB_, StrideDs_, StrideE_, - 1, + KBatch_, a_element_op_, b_element_op_, c_element_op_}, @@ -336,13 +343,14 @@ struct DeviceBatchedGemmMultiD_Xdl_CShuffle_V3 arg.Print(); } - if(!GridwiseGemm::CheckValidity(arg) || arg.KBatch > 1) + if(!GridwiseGemm::CheckValidity(arg)) { throw std::runtime_error("wrong! GridwiseGemm has invalid setting"); } index_t gdx, gdy, gdz; - std::tie(gdx, gdy, gdz) = GridwiseGemm::CalculateGridSize(arg.M, arg.N, arg.Batch); + std::tie(gdx, gdy, gdz) = + GridwiseGemm::CalculateGridSize(arg.M, arg.N, arg.Batch * arg.KBatch); float ave_time = 0; @@ -387,10 +395,11 @@ struct DeviceBatchedGemmMultiD_Xdl_CShuffle_V3 rotating_mem.Next(); // clear c mem if(arg_.KBatch > 1) - hipGetErrorString(hipMemsetAsync(arg_.p_c_grid, - 0, - arg_.M * arg_.N * sizeof(CDataType), - stream_config.stream_id_)); + hipGetErrorString( + hipMemsetAsync(arg_.p_c_grid, + 0, + arg.Batch * arg_.M * arg_.N * sizeof(CDataType), + stream_config.stream_id_)); }; ave_time = ck::utility::launch_and_time_kernel_with_preprocess( @@ -889,7 +898,8 @@ struct DeviceBatchedGemmMultiD_Xdl_CShuffle_V3 index_t BatchStrideE, AElementwiseOperation a_element_op, BElementwiseOperation b_element_op, - CElementwiseOperation c_element_op) + CElementwiseOperation c_element_op, + index_t KBatch = 1) { return Argument{static_cast(p_a), static_cast(p_b), @@ -909,7 +919,8 @@ struct DeviceBatchedGemmMultiD_Xdl_CShuffle_V3 Batch, a_element_op, b_element_op, - c_element_op}; + c_element_op, + KBatch}; } static auto MakeInvoker() { return Invoker{}; } @@ -934,7 +945,8 @@ struct DeviceBatchedGemmMultiD_Xdl_CShuffle_V3 index_t BatchStrideE, AElementwiseOperation a_element_op, BElementwiseOperation b_element_op, - CElementwiseOperation c_element_op) override + CElementwiseOperation c_element_op, + index_t KBatch = 1) override { return std::make_unique(static_cast(p_a), static_cast(p_b), @@ -954,7 +966,8 @@ struct DeviceBatchedGemmMultiD_Xdl_CShuffle_V3 Batch, a_element_op, b_element_op, - c_element_op); + c_element_op, + KBatch); } // polymorphic diff --git a/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_v3_multi_d.hpp b/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_v3_multi_d.hpp index c7038ed4fa..e5a31f8d1f 100644 --- a/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_v3_multi_d.hpp +++ b/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_v3_multi_d.hpp @@ -41,7 +41,7 @@ __global__ void #if(!defined(__HIP_DEVICE_COMPILE__) || defined(__gfx9__)) __shared__ char p_shared[GridwiseGemm::GetSharedMemoryNumberOfByte()]; - auto splitk_batch_offset = typename GridwiseGemm::SplitKBatchOffset(karg); + auto splitk_batch_offset = typename GridwiseGemm::SplitKBatchOffset(karg, blockIdx.z); GridwiseGemm::template Run( karg.p_a_grid + splitk_batch_offset.a_k_split_offset, @@ -76,7 +76,7 @@ __global__ void __shared__ char p_shared_0[GridwiseGemm::GetSharedMemoryNumberOfByte()]; __shared__ char p_shared_1[GridwiseGemm::GetSharedMemoryNumberOfByte()]; - auto splitk_batch_offset = typename GridwiseGemm::SplitKBatchOffset(karg); + auto splitk_batch_offset = typename GridwiseGemm::SplitKBatchOffset(karg, blockIdx.z); GridwiseGemm::template Run_2Lds( karg.p_a_grid + splitk_batch_offset.a_k_split_offset, @@ -639,27 +639,27 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3 struct SplitKBatchOffset { - __device__ SplitKBatchOffset(Argument& karg) + __device__ SplitKBatchOffset(Argument& karg, index_t k_id) { if constexpr(is_same_v) { - a_k_split_offset = blockIdx.z * karg.KRead; + a_k_split_offset = k_id * karg.KRead; } else if constexpr(is_same_v) { - a_k_split_offset = blockIdx.z * karg.KRead * karg.StrideA; + a_k_split_offset = k_id * karg.KRead * karg.StrideA; } if constexpr(is_same_v) { - b_k_split_offset = blockIdx.z * karg.KRead * karg.StrideB; + b_k_split_offset = k_id * karg.KRead * karg.StrideB; } else if constexpr(is_same_v) { - b_k_split_offset = blockIdx.z * karg.KRead; + b_k_split_offset = k_id * karg.KRead; } - if(blockIdx.z < static_cast(karg.KBatch - 1)) + if(k_id < karg.KBatch - 1) { karg.K = karg.KRead; } diff --git a/library/src/tensor_operation_instance/gpu/gemm_universal_batched/device_batched_gemm_xdl_universal_bf16_bf16_bf16/device_batched_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn.hpp b/library/src/tensor_operation_instance/gpu/gemm_universal_batched/device_batched_gemm_xdl_universal_bf16_bf16_bf16/device_batched_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn.hpp index 5db041de09..21cef335c5 100644 --- a/library/src/tensor_operation_instance/gpu/gemm_universal_batched/device_batched_gemm_xdl_universal_bf16_bf16_bf16/device_batched_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn.hpp +++ b/library/src/tensor_operation_instance/gpu/gemm_universal_batched/device_batched_gemm_xdl_universal_bf16_bf16_bf16/device_batched_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn.hpp @@ -52,6 +52,9 @@ using device_batched_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_comp_instances = DeviceBatchedGemmMultiD_Xdl_CShuffle_V3< Row, Col, DsLayout, Row, BF16, BF16, DsDataType, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 256, 256, 32, 8, 8, 32, 32, 4, 4, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, 1, 1, S<1, 16, 1, 16>, S<4>, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v5>, DeviceBatchedGemmMultiD_Xdl_CShuffle_V3< Row, Col, DsLayout, Row, BF16, BF16, DsDataType, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 224, 256, 64, 8, 8, 16, 16, 7, 8, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, 1, 2, S<1, 16, 1, 16>, S<4>, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v3>, DeviceBatchedGemmMultiD_Xdl_CShuffle_V3< Row, Col, DsLayout, Row, BF16, BF16, DsDataType, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 256, 224, 64, 8, 8, 16, 16, 8, 7, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, 2, 1, S<1, 32, 1, 8>, S<4>, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v3>, + DeviceBatchedGemmMultiD_Xdl_CShuffle_V3< Row, Col, DsLayout, Row, BF16, BF16, DsDataType, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 256, 160, 64, 8, 8, 16, 16, 8, 5, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, 2, 1, S<1, 32, 1, 8>, S<4>, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v3>, + DeviceBatchedGemmMultiD_Xdl_CShuffle_V3< Row, Col, DsLayout, Row, BF16, BF16, DsDataType, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 128, 160, 64, 8, 8, 32, 32, 1, 5, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, 1, 1, S<1, 64, 1, 4>, S<8>, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v3>, + DeviceBatchedGemmMultiD_Xdl_CShuffle_V3< Row, Col, DsLayout, Row, BF16, BF16, DsDataType, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 160, 128, 64, 8, 8, 32, 32, 5, 1, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, 1, 1, S<1, 16, 1, 16>, S<4>, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v3>, DeviceBatchedGemmMultiD_Xdl_CShuffle_V3< Row, Col, DsLayout, Row, BF16, BF16, DsDataType, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 128, 128, 64, 8, 8, 32, 32, 2, 2, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, 1, 1, S<1, 16, 1, 16>, S<4>, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v3>, DeviceBatchedGemmMultiD_Xdl_CShuffle_V3< Row, Col, DsLayout, Row, BF16, BF16, DsDataType, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 128, 128, 64, 8, 8, 32, 32, 2, 2, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, 1, 1, S<1, 16, 1, 16>, S<4>, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v5>, DeviceBatchedGemmMultiD_Xdl_CShuffle_V3< Row, Col, DsLayout, Row, BF16, BF16, DsDataType, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 128, 128, 64, 8, 8, 32, 32, 2, 2, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, 1, 1, S<1, 16, 1, 16>, S<4>, BlockGemmPipelineScheduler::Interwave, BlockGemmPipelineVersion::v1> diff --git a/library/src/tensor_operation_instance/gpu/gemm_universal_batched/device_batched_gemm_xdl_universal_f8_f8_bf16/device_batched_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn.hpp b/library/src/tensor_operation_instance/gpu/gemm_universal_batched/device_batched_gemm_xdl_universal_f8_f8_bf16/device_batched_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn.hpp index 355dc3212b..552ac3cd00 100644 --- a/library/src/tensor_operation_instance/gpu/gemm_universal_batched/device_batched_gemm_xdl_universal_f8_f8_bf16/device_batched_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn.hpp +++ b/library/src/tensor_operation_instance/gpu/gemm_universal_batched/device_batched_gemm_xdl_universal_f8_f8_bf16/device_batched_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn.hpp @@ -42,6 +42,7 @@ using device_batched_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_comp_instances = std //##################################| | | | | Type| Type| Type| Type| Type| Type| Elementwise| Elementwise| Elementwise|Specialization| Size| Block| Block| Block| | | XDL| XDL| Per| Per| ThreadCluster| ThreadCluster| SrcAccessOrder| SrcVectorDim| SrcScalar| DstScalar| AddExtraM| ThreadCluster| ThreadCluster| SrcAccessOrder| SrcVectorDim| SrcScalar| DstScalar| AddExtraN| MXdlPerWave| NXdlPerWave| _MBlock_MXdlPerWave_MWaveMPerXdl| ScalarPerVector| Pipeline| Pipeline| //##################################| | | | | | | | | | | Operation| Operation| Operation| | | | | | | | | | Wave| Wave| Lengths_K0_M_K1| ArrangeOrder| | | PerVector| PerVector_K1| | Lengths_K0_N_K1| ArrangeOrder| | | PerVector| PerVector_K1| | PerShuffle| PerShuffle| _NBlock_NXdlPerWave_NWaveNPerXdl| _NWaveNPerXdl| Scheduler| Verision| //##################################| | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | + #ifdef __gfx94__ // Compute friendly DeviceBatchedGemmMultiD_Xdl_CShuffle_V3< Row, Col, DsLayout, Row, F8, F8, DsDataType, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 256, 256, 64, 16, 16, 32, 32, 4, 4, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, 1, 1, S<1, 32, 1, 8>, S<8>, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v4, F8>, @@ -72,6 +73,7 @@ using device_batched_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_mem_instances = std: //##################################| | | | | Type| Type| Type| Type| Type| Type| Elementwise| Elementwise| Elementwise|Specialization| Size| Block| Block| Block| | | XDL| XDL| Per| Per| ThreadCluster| ThreadCluster| SrcAccessOrder| SrcVectorDim| SrcScalar| DstScalar| AddExtraM| ThreadCluster| ThreadCluster| SrcAccessOrder| SrcVectorDim| SrcScalar| DstScalar| AddExtraN| MXdlPerWave| NXdlPerWave| _MBlock_MXdlPerWave_MWaveMPerXdl| ScalarPerVector| Pipeline| Pipeline| //##################################| | | | | | | | | | | Operation| Operation| Operation| | | | | | | | | | Wave| Wave| Lengths_K0_M_K1| ArrangeOrder| | | PerVector| PerVector_K1| | Lengths_K0_N_K1| ArrangeOrder| | | PerVector| PerVector_K1| | PerShuffle| PerShuffle| _NBlock_NXdlPerWave_NWaveNPerXdl| _NWaveNPerXdl| Scheduler| Verision| //##################################| | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | + #if defined(__gfx94__) || defined(CK_USE_FP8_ON_UNSUPPORTED_ARCH) DeviceBatchedGemmMultiD_Xdl_CShuffle_V3< Row, Col, DsLayout, Row, F8, F8, DsDataType, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 128, 32, 16, 128, 16, 16, 16, 16, 1, 1, S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, 1, 1, S<1, 16, 1, 8>, S<2>, BlkGemmPipeSched, BlockGemmPipelineVersion::v1, F8>, DeviceBatchedGemmMultiD_Xdl_CShuffle_V3< Row, Col, DsLayout, Row, F8, F8, DsDataType, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 64, 16, 16, 128, 16, 16, 16, 16, 1, 1, S<8, 8, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, S<8, 8, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, 1, 1, S<1, 16, 1, 4>, S<4>, BlkGemmPipeSched, BlockGemmPipelineVersion::v1, F8>, diff --git a/profiler/include/profiler/profile_gemm_universal_batched_impl.hpp b/profiler/include/profiler/profile_gemm_universal_batched_impl.hpp index 53f81162ac..f4300af8d8 100644 --- a/profiler/include/profiler/profile_gemm_universal_batched_impl.hpp +++ b/profiler/include/profiler/profile_gemm_universal_batched_impl.hpp @@ -48,6 +48,7 @@ bool profile_gemm_universal_batched_impl(int do_verification, int StrideB, int StrideC, int BatchCount, + int KBatch, int n_warmup, int n_iter, uint64_t rotating = 0) @@ -147,89 +148,100 @@ bool profile_gemm_universal_batched_impl(int do_verification, float best_ave_time = 0; float best_tflops = 0; float best_gb_per_sec = 0; + float best_kbatch = 0; // profile device op instances for(auto& op_ptr : op_ptrs) { - std::unique_ptr argument_ptr; - // false branch for multi d dl kernel + std::vector kbatch_list = {1, 2, 4, 8, 16, 19, 32, 38}; - argument_ptr = - op_ptr->MakeArgumentPointer(static_cast(a_device_buf.GetDeviceBuffer()), - static_cast(b_device_buf.GetDeviceBuffer()), - {}, - static_cast(c_device_buf.GetDeviceBuffer()), - M, - N, - K, - BatchCount, - StrideA, - StrideB, - {}, - StrideC, - BatchStrideA, - BatchStrideB, - {}, - BatchStrideC, - ck::tensor_operation::element_wise::PassThrough{}, - ck::tensor_operation::element_wise::PassThrough{}, - ck::tensor_operation::element_wise::PassThrough{}); - - auto invoker_ptr = op_ptr->MakeInvokerPointer(); - - if(op_ptr->IsSupportedArgument(argument_ptr.get())) + if(KBatch > 0) { - // re-init C to zero before profiling next kernel - c_device_buf.SetZero(); + kbatch_list = {KBatch}; + } - std::string op_name = op_ptr->GetTypeString(); + for(std::size_t i = 0; i < kbatch_list.size(); i++) + { + auto kbatch_curr = kbatch_list[i]; - float ave_time = invoker_ptr->Run( - argument_ptr.get(), - StreamConfig{nullptr, time_kernel, 0, n_warmup, n_iter, true, rotating_count}); + auto argument_ptr = + op_ptr->MakeArgumentPointer(static_cast(a_device_buf.GetDeviceBuffer()), + static_cast(b_device_buf.GetDeviceBuffer()), + {}, + static_cast(c_device_buf.GetDeviceBuffer()), + M, + N, + K, + BatchCount, + StrideA, + StrideB, + {}, + StrideC, + BatchStrideA, + BatchStrideB, + {}, + BatchStrideC, + ck::tensor_operation::element_wise::PassThrough{}, + ck::tensor_operation::element_wise::PassThrough{}, + ck::tensor_operation::element_wise::PassThrough{}, + kbatch_curr); - std::size_t flop = std::size_t(2) * BatchCount * M * N * K; + auto invoker_ptr = op_ptr->MakeInvokerPointer(); - std::size_t num_btype = (sizeof(ADataType) * M * K + sizeof(BDataType) * K * N + - sizeof(CDataType) * M * N) * - BatchCount; - - 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, " << op_name << std::endl; - - if(tflops > best_tflops) + if(op_ptr->IsSupportedArgument(argument_ptr.get())) { - best_op_name = op_name; - best_tflops = tflops; - best_ave_time = ave_time; - best_gb_per_sec = gb_per_sec; - } + std::string op_name = op_ptr->GetTypeString(); - if(do_verification) - { - c_device_buf.FromDevice(c_g_m_n_device_result.mData.data()); + float ave_time = invoker_ptr->Run( + argument_ptr.get(), + StreamConfig{nullptr, time_kernel, 0, n_warmup, n_iter, true, rotating_count}); - pass = pass & ck::utils::check_err(c_g_m_n_device_result, c_g_m_n_host_result); + std::size_t flop = std::size_t(2) * BatchCount * M * N * K; - if(do_log) + std::size_t num_btype = (sizeof(ADataType) * M * K + sizeof(BDataType) * K * N + + sizeof(CDataType) * M * N) * + BatchCount; + + 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, " << op_name << ", KBatch " << kbatch_curr << std::endl; + + if(tflops > best_tflops) { - LogRangeAsType(std::cout << "a : ", a_g_m_k.mData, ",") << std::endl; - LogRangeAsType(std::cout << "b: ", b_g_k_n.mData, ",") << std::endl; - LogRangeAsType(std::cout << "c_host: ", c_g_m_n_host_result.mData, ",") - << std::endl; - LogRangeAsType( - std::cout << "c_device: ", c_g_m_n_device_result.mData, ",") - << std::endl; + best_op_name = op_name; + best_tflops = tflops; + best_ave_time = ave_time; + best_gb_per_sec = gb_per_sec; + best_kbatch = kbatch_curr; + } + + if(do_verification) + { + c_device_buf.FromDevice(c_g_m_n_device_result.mData.data()); + + pass = pass & ck::utils::check_err(c_g_m_n_device_result, c_g_m_n_host_result); + + if(do_log) + { + LogRangeAsType(std::cout << "a : ", a_g_m_k.mData, ",") << std::endl; + LogRangeAsType(std::cout << "b: ", b_g_k_n.mData, ",") << std::endl; + LogRangeAsType( + std::cout << "c_host: ", c_g_m_n_host_result.mData, ",") + << std::endl; + LogRangeAsType( + std::cout << "c_device: ", c_g_m_n_device_result.mData, ",") + << std::endl; + } } } - } - else - { - std::cout << op_ptr->GetTypeString() << " does not support this problem" << std::endl; + else + { + std::cout << op_ptr->GetTypeString() << " does not support this problem" + << std::endl; + } } } @@ -270,8 +282,8 @@ bool profile_gemm_universal_batched_impl(int do_verification, std::cout << " B = " << BatchCount << " M = " << M << " N = " << N << " K = " << K << " StrideA = " << StrideA << " StrideB = " << StrideB << " StrideC = " << StrideC - << ": " << best_ave_time << " ms, " << best_tflops << " TFlops, " << best_gb_per_sec - << " GB/s, " << best_op_name << std::endl; + << " KBatch = " << best_kbatch << ": " << best_ave_time << " ms, " << best_tflops + << " TFlops, " << best_gb_per_sec << " GB/s, " << best_op_name << std::endl; return pass; } diff --git a/profiler/src/profile_gemm_universal_batched.cpp b/profiler/src/profile_gemm_universal_batched.cpp index 4afef8e55f..d57511fbfc 100644 --- a/profiler/src/profile_gemm_universal_batched.cpp +++ b/profiler/src/profile_gemm_universal_batched.cpp @@ -31,7 +31,7 @@ enum struct GemmDataType int profile_batched_gemm_universal(int argc, char* argv[]) { - if(argc != 18 && argc != 21) + if(argc != 19 && argc != 22) { // clang-format off printf("arg1: tensor operation (" OP_NAME ": " OP_DESC ")\n"); @@ -44,11 +44,11 @@ int profile_batched_gemm_universal(int argc, char* argv[]) printf("arg5: initialization (0: no init; 1: integer value; 2: decimal value)\n"); printf("arg6: print tensor value (0: no; 1: yes)\n"); printf("arg7: time kernel (0=n0, 1=yes)\n"); - printf("arg8 to 17: M, N, K, StrideA, StrideB, StrideC, BatchStrideA, BatchStrideB, BatchStrideC, BatchCount\n"); + printf("arg8 to 18: M, N, K, StrideA, StrideB, StrideC, BatchStrideA, BatchStrideB, BatchStrideC, BatchCount, KBatch\n"); printf("optional:\n"); - printf("arg18: number of warm-up cycles (default 1)\n"); - printf("arg19: number of iterations (default 10)\n"); - printf("arg20: memory for rotating buffer (default 0, size in MB)\n"); + printf("arg19: number of warm-up cycles (default 1)\n"); + printf("arg20: number of iterations (default 10)\n"); + printf("arg21: memory for rotating buffer (default 0, size in MB)\n"); // clang-format on exit(1); } @@ -56,11 +56,11 @@ int profile_batched_gemm_universal(int argc, char* argv[]) int n_warmup = 1; int n_iter = 10; uint64_t rotating = 0; - if(argc == 21) + if(argc == 22) { - n_warmup = std::stoi(argv[18]); - n_iter = std::stoi(argv[19]); - rotating = std::stoull(argv[20]) * 1024 * 1024; + n_warmup = std::stoi(argv[19]); + n_iter = std::stoi(argv[20]); + rotating = std::stoull(argv[21]) * 1024 * 1024; } const auto data_type = static_cast(std::stoi(argv[2])); @@ -83,6 +83,7 @@ int profile_batched_gemm_universal(int argc, char* argv[]) const int BatchStrideC = std::stoi(argv[16]); const int BatchCount = std::stoi(argv[17]); + const int KBatch = std::stoi(argv[18]); #if defined(CK_USE_FP8_ON_UNSUPPORTED_ARCH) || defined(CK_USE_GFX94) using F8 = ck::f8_t; @@ -159,6 +160,7 @@ int profile_batched_gemm_universal(int argc, char* argv[]) StrideB_, StrideC_, BatchCount, + KBatch, n_warmup, n_iter, rotating); From 41ebf117a5927654a504803c19d18749babdeddd Mon Sep 17 00:00:00 2001 From: Illia Silin <98187287+illsilin@users.noreply.github.com> Date: Fri, 13 Dec 2024 16:30:22 -0800 Subject: [PATCH 02/25] Add zstd lib for building hipTensor. (#1745) * add zstd library to CI docker * fix the libzstd name --- Dockerfile | 1 + 1 file changed, 1 insertion(+) diff --git a/Dockerfile b/Dockerfile index 8ce158a200..4329c54c17 100644 --- a/Dockerfile +++ b/Dockerfile @@ -64,6 +64,7 @@ RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --allow- nano \ zlib1g-dev \ zip \ + libzstd-dev \ openssh-server \ clang-format-12 \ kmod && \ From d68974a5c68bd25bb8433302886213d7f5ff0d88 Mon Sep 17 00:00:00 2001 From: Illia Silin <98187287+illsilin@users.noreply.github.com> Date: Fri, 13 Dec 2024 16:30:39 -0800 Subject: [PATCH 03/25] upgrade pandas package (#1746) --- Dockerfile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Dockerfile b/Dockerfile index 4329c54c17..83edbfb8ee 100644 --- a/Dockerfile +++ b/Dockerfile @@ -94,7 +94,7 @@ RUN pip install --upgrade cmake==3.27.5 && \ dpkg -i dumb-init_*.deb && rm dumb-init_*.deb && \ # Install packages for processing the performance results pip3 install --upgrade pip && \ - pip3 install sqlalchemy==1.4.46 pymysql pandas==2.0.3 setuptools-rust sshtunnel==0.4.0 && \ + pip3 install sqlalchemy==1.4.46 pymysql pandas==2.2.3 setuptools-rust sshtunnel==0.4.0 && \ # Add render group groupadd -f render && \ # Install the new rocm-cmake version From f57d720c67123b43cb6f18f4b8b5aa0c7c9f51ba Mon Sep 17 00:00:00 2001 From: "Xu, Shengnan" <117875955+shengnxu@users.noreply.github.com> Date: Sun, 15 Dec 2024 20:13:10 +0800 Subject: [PATCH 04/25] added moe interleaving pipeline (#1712) * added moe interleaving pipeline * remove redundant code * formater --------- Co-authored-by: root --- include/ck_tile/ops/flatmm.hpp | 1 + ...latmm_sn_32x128x512_1x4x1_16x16x32_itl.hpp | 510 +++++++++++++ ..._uk_gfx9_32x128x512_1x4x1_16x16x16_itl.inc | 708 ++++++++++++++++++ .../fused_moegemm_pipeline_flatmm_policy.hpp | 29 +- .../pipeline/fused_moegemm_traits.hpp | 4 +- 5 files changed, 1249 insertions(+), 3 deletions(-) create mode 100644 include/ck_tile/ops/flatmm/block/flatmm_sn_32x128x512_1x4x1_16x16x32_itl.hpp create mode 100644 include/ck_tile/ops/flatmm/block/uk/flatmm_sn_uk_gfx9_32x128x512_1x4x1_16x16x16_itl.inc diff --git a/include/ck_tile/ops/flatmm.hpp b/include/ck_tile/ops/flatmm.hpp index eee80cda4a..ba76e3070d 100644 --- a/include/ck_tile/ops/flatmm.hpp +++ b/include/ck_tile/ops/flatmm.hpp @@ -5,6 +5,7 @@ #include "ck_tile/ops/flatmm/block/flatmm_32x512x128_1x4x1_16x16x32.hpp" #include "ck_tile/ops/flatmm/block/flatmm_sn_32x128x512_1x4x1_16x16x32.hpp" +#include "ck_tile/ops/flatmm/block/flatmm_sn_32x128x512_1x4x1_16x16x32_itl.hpp" #include "ck_tile/ops/flatmm/block/flatmm_uk_config.hpp" #include "ck_tile/ops/common/generic_2d_block_shape.hpp" #include "ck_tile/ops/common/tensor_layout.hpp" diff --git a/include/ck_tile/ops/flatmm/block/flatmm_sn_32x128x512_1x4x1_16x16x32_itl.hpp b/include/ck_tile/ops/flatmm/block/flatmm_sn_32x128x512_1x4x1_16x16x32_itl.hpp new file mode 100644 index 0000000000..681a696036 --- /dev/null +++ b/include/ck_tile/ops/flatmm/block/flatmm_sn_32x128x512_1x4x1_16x16x32_itl.hpp @@ -0,0 +1,510 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved. + +#pragma once + +#include "ck_tile/core.hpp" +#include "ck_tile/ops/gemm/warp/warp_gemm.hpp" +#include "ck_tile/ops/flatmm/block/flatmm_uk_config.hpp" +#include "ck_tile/ops/flatmm/block/flatmm_sn_32x128x512_1x4x1_16x16x32.hpp" + +namespace ck_tile { + +// "S"tream update output along "N" +// A in smem, B load from global +// require 4 wave, occupancy=1c + +struct FlatmmSn_32x128x512_1x4x1_16x16x32_BF16_itl : public FlatmmSn_32x128x512_1x4x1_16x16x32_Base +{ + using BDataType = bf16_t; + using ODataType = bf16_t; + + // TODO: need paired with tile_window_linear! + // TODO: need call init_raw() before call this function! + // template + template + CK_TILE_DEVICE auto + operator()(const BRes& res_b, + const BCoords& cached_coords_b, + const ORes& res_o, + const OCoords& cached_coords_o, + const OFlags& o_flags, // this should be in sgpr + CK_TILE_LDS_ADDR void* smem, + index_t n, // loop along n dim + const ScaleTensor& scale_, + index_t tile_offset_b, // stride b is fixed to blockKr * blockW, but still can adjust + index_t tile_offset_o) + { + static_assert(BCoords::size() == 8); // 8 + static_assert(OCoords::size() == 8); + + const index_t tile_stride_b_bytes = tile_offset_b * sizeof(BDataType); + const index_t tile_stride_o_bytes = tile_offset_o * sizeof(ODataType); + + static_assert(ScaleTensor::size() == 2); + float s0 = scale_[number<0>{}]; + float s1 = scale_[number<1>{}]; + + // index_t loop_cnt = n / Block_N; + + register float v_c0 asm("v64"); + register float v_c1 asm("v65"); + register float v_c2 asm("v66"); + register float v_c3 asm("v67"); + register float v_c4 asm("v68"); + register float v_c5 asm("v69"); + register float v_c6 asm("v70"); + register float v_c7 asm("v71"); + register float v_c8 asm("v72"); + register float v_c9 asm("v73"); + register float v_c10 asm("v74"); + register float v_c11 asm("v75"); + register float v_c12 asm("v76"); + register float v_c13 asm("v77"); + register float v_c14 asm("v78"); + register float v_c15 asm("v79"); + register float v_c16 asm("v80"); + register float v_c17 asm("v81"); + register float v_c18 asm("v82"); + register float v_c19 asm("v83"); + register float v_c20 asm("v84"); + register float v_c21 asm("v85"); + register float v_c22 asm("v86"); + register float v_c23 asm("v87"); + register float v_c24 asm("v88"); + register float v_c25 asm("v89"); + register float v_c26 asm("v90"); + register float v_c27 asm("v91"); + register float v_c28 asm("v92"); + register float v_c29 asm("v93"); + register float v_c30 asm("v94"); + register float v_c31 asm("v95"); + int32_t nan_hi = 0x7fff0000; + int32_t nan_lo = 0x00007fff; + + // in smem, the layout is M0(2)*K0(128)*M1(16)*K1(4) + // every threads need 8xK in contiguous register + // ... and every wave need the same data + int lane_id = threadIdx.x % 64; + int sld_y_os = (lane_id % 16) * 4 + (lane_id / 16) * 128; + sld_y_os *= 2; + + // y y p p p y + // reg before shfl M0(2)*N0(2)*Nl(4)*Nw(4)*Mw(16)*Nv(4) + // but order is N0*M0*Nv + // in LDS we need store as + // M0(2)* N0(2) * Nl(4) * Nw(4) * (Mw(16)*Nv(4) + 4) + // y y wave-id lid/16 lid%16 v + // sst(v3) = (v0/16*34 + v0%16 * 2 + wid*136) * 4 + int sfl_sst = (threadIdx.x % 16 * 4) + (threadIdx.x / 16) * (64 + 4); + sfl_sst *= 2; + + // from LDS we need load as + // M0(2)* N0(2) * Nl(4) * Nw(4) * (Mw(16) * Nv(4) + 4) + // ( 2 issue) (rem 32-lane) (4 wave*4issue) 2lane*1ussue(pk2) + // sld(v4) = v0/2 *34*4 + v0 % 2 *4 + wid*2 *4 + int sfl_sld = (lane_id % 2) * 2 + (lane_id / 2) * (64 + 4) + (threadIdx.x / 64) * 4; + sfl_sld *= 2; + + // B nr->kr + // clang-format off +#pragma clang diagnostic push +#pragma clang diagnostic ignored "-Winline-asm" + asm volatile( +#define CK_TILE_FLATMM_UK_MFMA CK_TILE_FLATMM_UK_MFMA_BF16 +#include "uk/flatmm_sn_uk_gfx9_32x128x512_1x4x1_16x16x16_itl.inc" +#undef CK_TILE_FLATMM_UK_MFMA + :[smem_]"+r"(smem), + // [s_loop_cnt]"+s"(loop_cnt), + [s_loop_cnt]"+s"(n), + [c0]"+v" (v_c0), + [c1]"+v" (v_c1), + [c2]"+v" (v_c2), + [c3]"+v" (v_c3), + [c4]"+v" (v_c4), + [c5]"+v" (v_c5), + [c6]"+v" (v_c6), + [c7]"+v" (v_c7), + [c8]"+v" (v_c8), + [c9]"+v" (v_c9), + [c10]"+v"(v_c10), + [c11]"+v"(v_c11), + [c12]"+v"(v_c12), + [c13]"+v"(v_c13), + [c14]"+v"(v_c14), + [c15]"+v"(v_c15), + [c16]"+v"(v_c16), + [c17]"+v"(v_c17), + [c18]"+v"(v_c18), + [c19]"+v"(v_c19), + [c20]"+v"(v_c20), + [c21]"+v"(v_c21), + [c22]"+v"(v_c22), + [c23]"+v"(v_c23), + [c24]"+v"(v_c24), + [c25]"+v"(v_c25), + [c26]"+v"(v_c26), + [c27]"+v"(v_c27), + [c28]"+v"(v_c28), + [c29]"+v"(v_c29), + [c30]"+v"(v_c30), + [c31]"+v"(v_c31) + : + [sld_a_base]"n"(0), + [shfl_base]"n"(0), + [v_sld_y_os]"v"(sld_y_os), + [v_sfl_sld]"v"(sfl_sld), + [v_sfl_sst]"v"(sfl_sst), + [s_res_o0]"s"(res_o[0]), + [s_res_o1]"s"(res_o[1]), + //[s_res_o2]"s"(res_o[2]), + //[s_res_o3]"s"(res_o[3]), + [s_res_b0]"s"(res_b[0]), + [s_res_b1]"s"(res_b[1]), + [s_res_b2]"s"(res_b[2]), + [s_res_b3]"s"(res_b[3]), + [v_os_o0]"v"(static_cast(cached_coords_o[number<0>{}] * sizeof(ODataType))), + [v_os_o1]"v"(static_cast(cached_coords_o[number<1>{}] * sizeof(ODataType))), + [v_os_o2]"v"(static_cast(cached_coords_o[number<2>{}] * sizeof(ODataType))), + [v_os_o3]"v"(static_cast(cached_coords_o[number<3>{}] * sizeof(ODataType))), + [v_os_o4]"v"(static_cast(cached_coords_o[number<4>{}] * sizeof(ODataType))), + [v_os_o5]"v"(static_cast(cached_coords_o[number<5>{}] * sizeof(ODataType))), + [v_os_o6]"v"(static_cast(cached_coords_o[number<6>{}] * sizeof(ODataType))), + [v_os_o7]"v"(static_cast(cached_coords_o[number<7>{}] * sizeof(ODataType))), + [v_os_b0]"v"(static_cast(cached_coords_b[number<0>{}] * sizeof(BDataType))), + [v_os_b1]"v"(static_cast(cached_coords_b[number<1>{}] * sizeof(BDataType))), + [v_os_b2]"v"(static_cast(cached_coords_b[number<2>{}] * sizeof(BDataType))), + [v_os_b3]"v"(static_cast(cached_coords_b[number<3>{}] * sizeof(BDataType))), + [v_os_b4]"v"(static_cast(cached_coords_b[number<4>{}] * sizeof(BDataType))), + [v_os_b5]"v"(static_cast(cached_coords_b[number<5>{}] * sizeof(BDataType))), + [v_os_b6]"v"(static_cast(cached_coords_b[number<6>{}] * sizeof(BDataType))), + [v_os_b7]"v"(static_cast(cached_coords_b[number<7>{}] * sizeof(BDataType))), + + [s_tile_os_o]"s"(tile_stride_o_bytes), + [s_tile_os_b]"s"(tile_stride_b_bytes), + [scale_0]"v"(s0), + [scale_1]"v"(s1), + [v_nan_lo]"v"(nan_lo), + [v_nan_hi]"v"(nan_hi), + [s_execflag_0]"s"(o_flags[number<0>{}]), + [s_execflag_1]"s"(o_flags[number<1>{}]), + [s_execflag_2]"s"(o_flags[number<2>{}]), + [s_execflag_3]"s"(o_flags[number<3>{}]), + [s_execflag_4]"s"(o_flags[number<4>{}]), + [s_execflag_5]"s"(o_flags[number<5>{}]), + [s_execflag_6]"s"(o_flags[number<6>{}]), + [s_execflag_7]"s"(o_flags[number<7>{}]) + : + "memory", "a0", "a1", "a2", "a3", "a4", "a5", "a6", "a7", "a8", "a9", + "a10", "a11", "a12", "a13", "a14", "a15", "a16", "a17", "a18", "a19", + "a20", "a21", "a22", "a23", "a24", "a25", "a26", "a27", "a28", "a29", + "a30", "a31", "a32", "a33", "a34", "a35", "a36", "a37", "a38", "a39", + "a40", "a41", "a42", "a43", "a44", "a45", "a46", "a47", "a48", "a49", + "a50", "a51", "a52", "a53", "a54", "a55", "a56", "a57", "a58", "a59", + "a60", "a61", "a62", "a63", "a64", "a65", "a66", "a67", "a68", "a69", + "a70", "a71", "a72", "a73", "a74", "a75", "a76", "a77", "a78", "a79", + "a80", "a81", "a82", "a83", "a84", "a85", "a86", "a87", "a88", "a89", + "a90", "a91", "a92", "a93", "a94", "a95", "a96", "a97", "a98", "a99", + "a100", "a101", "a102", "a103", "a104", "a105", "a106", "a107", + "a108", "a109", "a110", "a111", "a112", "a113", "a114", "a115", + "a116", "a117", "a118", "a119", "a120", "a121", "a122", "a123", + "a124", "a125", "a126", "a127", "a128", "a129", "a130", "a131", + "a132", "a133", "a134", "a135", "a136", "a137", "a138", "a139", + "a140", "a141", "a142", "a143", "a144", "a145", "a146", "a147", + "a148", "a149", "a150", "a151", "a152", "a153", "a154", "a155", + "a156", "a157", "a158", "a159", "a160", "a161", "a162", "a163", + "a164", "a165", "a166", "a167", "a168", "a169", "a170", "a171", + "a172", "a173", "a174", "a175", "a176", "a177", "a178", "a179", + "a180", "a181", "a182", "a183", "a184", "a185", "a186", "a187", + "a188", "a189", "a190", "a191", "a192", "a193", "a194", "a195", + "a196", "a197", "a198", "a199", "a200", "a201", "a202", "a203", + "a204", "a205", "a206", "a207", "a208", "a209", "a210", "a211", + "a212", "a213", "a214", "a215", "a216", "a217", "a218", "a219", + "a220", "a221", "a222", "a223", "a224", "a225", "a226", "a227", + "a228", "a229", "a230", "a231", "a232", "a233", "a234", "a235", + "a236", "a237", "a238", "a239", "a240", "a241", "a242", "a243", + "a244", "a245", "a246", "a247", "a248", "a249", "a250", "a251", + "a252", "a253", "a254", "a255", + "s8", "s9", "s12", "s13", "s14", "s15", "s38", "s39", "s52", "s86", + "s36", "s37","s59","s80", + "v10", "v11", "v12", "v13", "v14", "v15", "v16", "v17", + "v50", "v54", "v55", + "v64","v65","v66","v67","v68","v69","v70","v71", + "v72","v73","v74","v75","v76","v77","v78","v79", + "v80","v81","v82","v83","v84","v85","v86","v87", + "v88","v89","v90","v91","v92","v93","v94","v95", + "v128", "v129", "v130", "v131", + "v132", "v133", "v134", "v135", "v136", "v137", "v138", "v139", + "v140", "v141", "v142", "v143", "v144", "v145", "v146", "v147", + "v148", "v149", "v150", "v151", "v152", "v153", "v154", "v155", + "v156", "v157", "v158", "v159", "v160", "v161", "v162", "v163", + "v164", "v165", "v166", "v167", "v168", "v169", "v170", "v171", + "v172", "v173", "v174", "v175", "v176", "v177", "v178", "v179", + "v180", "v181", "v182", "v183", "v184", "v185", "v186", "v187", + "v188", "v189", "v190", "v191", "v192", "v193", "v194", "v195", + "v196", "v197", "v198", "v199", "v200", "v201", "v202", "v203", + "v204", "v205", "v206", "v207", "v208", "v209", "v210", "v211", + "v212", "v213", "v214", "v215", "v216", "v217", "v218", "v219", + "v220", "v221", "v222", "v223", "v224", "v225", "v226", "v227", + "v228", "v229", "v230", "v231", "v232", "v233", "v234", "v235", + "v236", "v237", "v238", "v239", "v240", "v241", "v242", "v243", + "v244", "v245", "v246", "v247", "v248", "v249", "v250", "v251", + "v252", "v253", "v254", "v255" + ); +#pragma clang diagnostic pop + // clang-format on + } +}; + +struct FlatmmSn_32x128x512_1x4x1_16x16x32_FP16_itl : public FlatmmSn_32x128x512_1x4x1_16x16x32_Base +{ + using BDataType = bf16_t; + using ODataType = bf16_t; + + // TODO: need paired with tile_window_linear! + // TODO: need call init_raw() before call this function! + // template + template + CK_TILE_DEVICE auto + operator()(const BRes& res_b, + const BCoords& cached_coords_b, + const ORes& res_o, + const OCoords& cached_coords_o, + const OFlags& o_flags, // this should be in sgpr + CK_TILE_LDS_ADDR void* smem, + index_t n, // loop along n dim + const ScaleTensor& scale_, + index_t tile_offset_b, // stride b is fixed to blockKr * blockW, but still can adjust + index_t tile_offset_o) + { + static_assert(BCoords::size() == 8); // 8 + static_assert(OCoords::size() == 8); + + const index_t tile_stride_b_bytes = tile_offset_b * sizeof(BDataType); + const index_t tile_stride_o_bytes = tile_offset_o * sizeof(ODataType); + + static_assert(ScaleTensor::size() == 2); + float s0 = scale_[number<0>{}]; + float s1 = scale_[number<1>{}]; + + // index_t loop_cnt = n / Block_N; + + register float v_c0 asm("v64"); + register float v_c1 asm("v65"); + register float v_c2 asm("v66"); + register float v_c3 asm("v67"); + register float v_c4 asm("v68"); + register float v_c5 asm("v69"); + register float v_c6 asm("v70"); + register float v_c7 asm("v71"); + register float v_c8 asm("v72"); + register float v_c9 asm("v73"); + register float v_c10 asm("v74"); + register float v_c11 asm("v75"); + register float v_c12 asm("v76"); + register float v_c13 asm("v77"); + register float v_c14 asm("v78"); + register float v_c15 asm("v79"); + register float v_c16 asm("v80"); + register float v_c17 asm("v81"); + register float v_c18 asm("v82"); + register float v_c19 asm("v83"); + register float v_c20 asm("v84"); + register float v_c21 asm("v85"); + register float v_c22 asm("v86"); + register float v_c23 asm("v87"); + register float v_c24 asm("v88"); + register float v_c25 asm("v89"); + register float v_c26 asm("v90"); + register float v_c27 asm("v91"); + register float v_c28 asm("v92"); + register float v_c29 asm("v93"); + register float v_c30 asm("v94"); + register float v_c31 asm("v95"); + int32_t nan_hi = 0x7fff0000; + int32_t nan_lo = 0x00007fff; + + // in smem, the layout is M0(2)*K0(128)*M1(16)*K1(4) + // every threads need 8xK in contiguous register + // ... and every wave need the same data + int lane_id = threadIdx.x % 64; + int sld_y_os = (lane_id % 16) * 4 + (lane_id / 16) * 128; + sld_y_os *= 2; + + // y y p p p y + // reg before shfl M0(2)*N0(2)*Nl(4)*Nw(4)*Mw(16)*Nv(4) + // but order is N0*M0*Nv + // in LDS we need store as + // M0(2)* N0(2) * Nl(4) * Nw(4) * (Mw(16)*Nv(4) + 4) + // y y wave-id lid/16 lid%16 v + // sst(v3) = (v0/16*34 + v0%16 * 2 + wid*136) * 4 + int sfl_sst = (threadIdx.x % 16 * 4) + (threadIdx.x / 16) * (64 + 4); + sfl_sst *= 2; + + // from LDS we need load as + // M0(2)* N0(2) * Nl(4) * Nw(4) * (Mw(16) * Nv(4) + 4) + // ( 2 issue) (rem 32-lane) (4 wave*4issue) 2lane*1ussue(pk2) + // sld(v4) = v0/2 *34*4 + v0 % 2 *4 + wid*2 *4 + int sfl_sld = (lane_id % 2) * 2 + (lane_id / 2) * (64 + 4) + (threadIdx.x / 64) * 4; + sfl_sld *= 2; + + // B nr->kr + // clang-format off +#pragma clang diagnostic push +#pragma clang diagnostic ignored "-Winline-asm" + asm volatile( +#define CK_TILE_FLATMM_UK_MFMA CK_TILE_FLATMM_UK_MFMA_FP16 +#include "uk/flatmm_sn_uk_gfx9_32x128x512_1x4x1_16x16x16_itl.inc" +#undef CK_TILE_FLATMM_UK_MFMA + :[smem_]"+r"(smem), + [s_loop_cnt]"+s"(n), + [c0]"+v" (v_c0), + [c1]"+v" (v_c1), + [c2]"+v" (v_c2), + [c3]"+v" (v_c3), + [c4]"+v" (v_c4), + [c5]"+v" (v_c5), + [c6]"+v" (v_c6), + [c7]"+v" (v_c7), + [c8]"+v" (v_c8), + [c9]"+v" (v_c9), + [c10]"+v"(v_c10), + [c11]"+v"(v_c11), + [c12]"+v"(v_c12), + [c13]"+v"(v_c13), + [c14]"+v"(v_c14), + [c15]"+v"(v_c15), + [c16]"+v"(v_c16), + [c17]"+v"(v_c17), + [c18]"+v"(v_c18), + [c19]"+v"(v_c19), + [c20]"+v"(v_c20), + [c21]"+v"(v_c21), + [c22]"+v"(v_c22), + [c23]"+v"(v_c23), + [c24]"+v"(v_c24), + [c25]"+v"(v_c25), + [c26]"+v"(v_c26), + [c27]"+v"(v_c27), + [c28]"+v"(v_c28), + [c29]"+v"(v_c29), + [c30]"+v"(v_c30), + [c31]"+v"(v_c31) + : + [sld_a_base]"n"(0), + [shfl_base]"n"(0), + [v_sld_y_os]"v"(sld_y_os), + [v_sfl_sld]"v"(sfl_sld), + [v_sfl_sst]"v"(sfl_sst), + [s_res_o0]"s"(res_o[0]), + [s_res_o1]"s"(res_o[1]), + //[s_res_o2]"s"(res_o[2]), + //[s_res_o3]"s"(res_o[3]), + [s_res_b0]"s"(res_b[0]), + [s_res_b1]"s"(res_b[1]), + [s_res_b2]"s"(res_b[2]), + [s_res_b3]"s"(res_b[3]), + [v_os_o0]"v"(static_cast(cached_coords_o[number<0>{}] * sizeof(ODataType))), + [v_os_o1]"v"(static_cast(cached_coords_o[number<1>{}] * sizeof(ODataType))), + [v_os_o2]"v"(static_cast(cached_coords_o[number<2>{}] * sizeof(ODataType))), + [v_os_o3]"v"(static_cast(cached_coords_o[number<3>{}] * sizeof(ODataType))), + [v_os_o4]"v"(static_cast(cached_coords_o[number<4>{}] * sizeof(ODataType))), + [v_os_o5]"v"(static_cast(cached_coords_o[number<5>{}] * sizeof(ODataType))), + [v_os_o6]"v"(static_cast(cached_coords_o[number<6>{}] * sizeof(ODataType))), + [v_os_o7]"v"(static_cast(cached_coords_o[number<7>{}] * sizeof(ODataType))), + [v_os_b0]"v"(static_cast(cached_coords_b[number<0>{}] * sizeof(BDataType))), + [v_os_b1]"v"(static_cast(cached_coords_b[number<1>{}] * sizeof(BDataType))), + [v_os_b2]"v"(static_cast(cached_coords_b[number<2>{}] * sizeof(BDataType))), + [v_os_b3]"v"(static_cast(cached_coords_b[number<3>{}] * sizeof(BDataType))), + [v_os_b4]"v"(static_cast(cached_coords_b[number<4>{}] * sizeof(BDataType))), + [v_os_b5]"v"(static_cast(cached_coords_b[number<5>{}] * sizeof(BDataType))), + [v_os_b6]"v"(static_cast(cached_coords_b[number<6>{}] * sizeof(BDataType))), + [v_os_b7]"v"(static_cast(cached_coords_b[number<7>{}] * sizeof(BDataType))), + + [s_tile_os_o]"s"(tile_stride_o_bytes), + [s_tile_os_b]"s"(tile_stride_b_bytes), + [scale_0]"v"(s0), + [scale_1]"v"(s1), + [v_nan_lo]"v"(nan_lo), + [v_nan_hi]"v"(nan_hi), + [s_execflag_0]"s"(o_flags[number<0>{}]), + [s_execflag_1]"s"(o_flags[number<1>{}]), + [s_execflag_2]"s"(o_flags[number<2>{}]), + [s_execflag_3]"s"(o_flags[number<3>{}]), + [s_execflag_4]"s"(o_flags[number<4>{}]), + [s_execflag_5]"s"(o_flags[number<5>{}]), + [s_execflag_6]"s"(o_flags[number<6>{}]), + [s_execflag_7]"s"(o_flags[number<7>{}]) + : + "memory", "a0", "a1", "a2", "a3", "a4", "a5", "a6", "a7", "a8", "a9", + "a10", "a11", "a12", "a13", "a14", "a15", "a16", "a17", "a18", "a19", + "a20", "a21", "a22", "a23", "a24", "a25", "a26", "a27", "a28", "a29", + "a30", "a31", "a32", "a33", "a34", "a35", "a36", "a37", "a38", "a39", + "a40", "a41", "a42", "a43", "a44", "a45", "a46", "a47", "a48", "a49", + "a50", "a51", "a52", "a53", "a54", "a55", "a56", "a57", "a58", "a59", + "a60", "a61", "a62", "a63", "a64", "a65", "a66", "a67", "a68", "a69", + "a70", "a71", "a72", "a73", "a74", "a75", "a76", "a77", "a78", "a79", + "a80", "a81", "a82", "a83", "a84", "a85", "a86", "a87", "a88", "a89", + "a90", "a91", "a92", "a93", "a94", "a95", "a96", "a97", "a98", "a99", + "a100", "a101", "a102", "a103", "a104", "a105", "a106", "a107", + "a108", "a109", "a110", "a111", "a112", "a113", "a114", "a115", + "a116", "a117", "a118", "a119", "a120", "a121", "a122", "a123", + "a124", "a125", "a126", "a127", "a128", "a129", "a130", "a131", + "a132", "a133", "a134", "a135", "a136", "a137", "a138", "a139", + "a140", "a141", "a142", "a143", "a144", "a145", "a146", "a147", + "a148", "a149", "a150", "a151", "a152", "a153", "a154", "a155", + "a156", "a157", "a158", "a159", "a160", "a161", "a162", "a163", + "a164", "a165", "a166", "a167", "a168", "a169", "a170", "a171", + "a172", "a173", "a174", "a175", "a176", "a177", "a178", "a179", + "a180", "a181", "a182", "a183", "a184", "a185", "a186", "a187", + "a188", "a189", "a190", "a191", "a192", "a193", "a194", "a195", + "a196", "a197", "a198", "a199", "a200", "a201", "a202", "a203", + "a204", "a205", "a206", "a207", "a208", "a209", "a210", "a211", + "a212", "a213", "a214", "a215", "a216", "a217", "a218", "a219", + "a220", "a221", "a222", "a223", "a224", "a225", "a226", "a227", + "a228", "a229", "a230", "a231", "a232", "a233", "a234", "a235", + "a236", "a237", "a238", "a239", "a240", "a241", "a242", "a243", + "a244", "a245", "a246", "a247", "a248", "a249", "a250", "a251", + "a252", "a253", "a254", "a255", + "s8", "s9", "s12", "s13", "s14", "s15", "s38", "s39", "s52", "s86", + "s36", "s37","s59","s80", + "v10", "v11", "v12", "v13", "v14", "v15", "v16", "v17", + "v50", "v54", "v55", + "v64","v65","v66","v67","v68","v69","v70","v71", + "v72","v73","v74","v75","v76","v77","v78","v79", + "v80","v81","v82","v83","v84","v85","v86","v87", + "v88","v89","v90","v91","v92","v93","v94","v95", + "v128", "v129", "v130", "v131", + "v132", "v133", "v134", "v135", "v136", "v137", "v138", "v139", + "v140", "v141", "v142", "v143", "v144", "v145", "v146", "v147", + "v148", "v149", "v150", "v151", "v152", "v153", "v154", "v155", + "v156", "v157", "v158", "v159", "v160", "v161", "v162", "v163", + "v164", "v165", "v166", "v167", "v168", "v169", "v170", "v171", + "v172", "v173", "v174", "v175", "v176", "v177", "v178", "v179", + "v180", "v181", "v182", "v183", "v184", "v185", "v186", "v187", + "v188", "v189", "v190", "v191", "v192", "v193", "v194", "v195", + "v196", "v197", "v198", "v199", "v200", "v201", "v202", "v203", + "v204", "v205", "v206", "v207", "v208", "v209", "v210", "v211", + "v212", "v213", "v214", "v215", "v216", "v217", "v218", "v219", + "v220", "v221", "v222", "v223", "v224", "v225", "v226", "v227", + "v228", "v229", "v230", "v231", "v232", "v233", "v234", "v235", + "v236", "v237", "v238", "v239", "v240", "v241", "v242", "v243", + "v244", "v245", "v246", "v247", "v248", "v249", "v250", "v251", + "v252", "v253", "v254", "v255" + ); +#pragma clang diagnostic pop + // clang-format on + } +}; + +} // namespace ck_tile diff --git a/include/ck_tile/ops/flatmm/block/uk/flatmm_sn_uk_gfx9_32x128x512_1x4x1_16x16x16_itl.inc b/include/ck_tile/ops/flatmm/block/uk/flatmm_sn_uk_gfx9_32x128x512_1x4x1_16x16x16_itl.inc new file mode 100644 index 0000000000..b8c6d20021 --- /dev/null +++ b/include/ck_tile/ops/flatmm/block/uk/flatmm_sn_uk_gfx9_32x128x512_1x4x1_16x16x16_itl.inc @@ -0,0 +1,708 @@ +#ifndef CK_TILE_FLATMM_UK_MFMA +#define CK_TILE_FLATMM_UK_MFMA CK_TILE_FLATMM_UK_MFMA_BF16 +#endif + +#if CK_TILE_FLATMM_UK_MFMA == CK_TILE_FLATMM_UK_MFMA_BF16 +# define _UK_MFMA_ "v_mfma_f32_16x16x16_bf16" + +# define _UK_PK_CVT_(x0_, x1_, y_) \ + " v_cmp_u_f32 s[36:37], " x0_ ", " x0_ " \n" \ + " v_add3_u32 v50, " x0_ ", %[v_nan_lo], 1 \n" \ + " v_cndmask_b32 v54, v50, %[v_nan_hi], s[36:37] \n" \ + " v_cmp_u_f32 s[36:37], " x1_ ", " x1_ " \n" \ + " v_add3_u32 v50, " x1_ ", %[v_nan_lo], 1 \n" \ + " v_cndmask_b32 v55, v50, %[v_nan_hi], s[36:37] \n" \ + " v_perm_b32 " y_ ", v55, v54, s52 \n" + +# define _UK_ATOMIC_ADD_ "global_atomic_pk_add_bf16" + +#elif CK_TILE_FLATMM_UK_MFMA == CK_TILE_FLATMM_UK_MFMA_FP16 +#define _UK_MFMA_ "v_mfma_f32_16x16x16_f16" + +# define _UK_PK_CVT_(x0_, x1_, y_) \ + " v_cvt_f16_f32 v54, " x0_ " \n" \ + " v_cvt_f16_f32 v55, " x1_ " \n" \ + " v_pack_b32_f16 " y_ ", v54, v55 \n" + +# define _UK_ATOMIC_ADD_ "global_atomic_pk_add_f16" + +#endif + + +";-------------------------------------------------------------\n" +" s_mov_b32 s52, 0x07060302 ; v_perm\n" +" s_mov_b64 s[38:39], exec ; save current exec\n" +" s_mov_b32 s8, %[s_res_o0] \n" +" s_mov_b32 s9, %[s_res_o1] \n" +" s_mov_b32 s12, %[s_res_b0] \n" +" s_mov_b32 s13, %[s_res_b1] \n" +" s_mov_b32 s14, %[s_res_b2] \n" +" s_mov_b32 s15, %[s_res_b3] \n" +" s_mov_b32 s59, 0 \n" +" ds_read_b64 v[128:129], %[v_sld_y_os] offset:0 + %[sld_a_base] \n" +" ds_read_b64 v[130:131], %[v_sld_y_os] offset:128 + %[sld_a_base] \n" +" ds_read_b64 v[132:133], %[v_sld_y_os] offset:1024 + %[sld_a_base] \n" +" ds_read_b64 v[134:135], %[v_sld_y_os] offset:1152 + %[sld_a_base] \n" +" ds_read_b64 v[136:137], %[v_sld_y_os] offset:2048 + %[sld_a_base] \n" +" ds_read_b64 v[138:139], %[v_sld_y_os] offset:2176 + %[sld_a_base] \n" +" ds_read_b64 v[140:141], %[v_sld_y_os] offset:3072 + %[sld_a_base] \n" +" ds_read_b64 v[142:143], %[v_sld_y_os] offset:3200 + %[sld_a_base] \n" +" ds_read_b64 v[144:145], %[v_sld_y_os] offset:4096 + %[sld_a_base] \n" +" ds_read_b64 v[146:147], %[v_sld_y_os] offset:4224 + %[sld_a_base] \n" +" ds_read_b64 v[148:149], %[v_sld_y_os] offset:5120 + %[sld_a_base] \n" +" ds_read_b64 v[150:151], %[v_sld_y_os] offset:5248 + %[sld_a_base] \n" +" ds_read_b64 v[152:153], %[v_sld_y_os] offset:6144 + %[sld_a_base] \n" +" ds_read_b64 v[154:155], %[v_sld_y_os] offset:6272 + %[sld_a_base] \n" +" ds_read_b64 v[156:157], %[v_sld_y_os] offset:7168 + %[sld_a_base] \n" +" ds_read_b64 v[158:159], %[v_sld_y_os] offset:7296 + %[sld_a_base] \n" +" ds_read_b64 v[160:161], %[v_sld_y_os] offset:8192 + %[sld_a_base] \n" +" ds_read_b64 v[162:163], %[v_sld_y_os] offset:8320 + %[sld_a_base] \n" +" ds_read_b64 v[164:165], %[v_sld_y_os] offset:9216 + %[sld_a_base] \n" +" ds_read_b64 v[166:167], %[v_sld_y_os] offset:9344 + %[sld_a_base] \n" +" ds_read_b64 v[168:169], %[v_sld_y_os] offset:10240 + %[sld_a_base] \n" +" ds_read_b64 v[170:171], %[v_sld_y_os] offset:10368 + %[sld_a_base] \n" +" ds_read_b64 v[172:173], %[v_sld_y_os] offset:11264 + %[sld_a_base] \n" +" ds_read_b64 v[174:175], %[v_sld_y_os] offset:11392 + %[sld_a_base] \n" +" ds_read_b64 v[176:177], %[v_sld_y_os] offset:12288 + %[sld_a_base] \n" +" ds_read_b64 v[178:179], %[v_sld_y_os] offset:12416 + %[sld_a_base] \n" +" ds_read_b64 v[180:181], %[v_sld_y_os] offset:13312 + %[sld_a_base] \n" +" ds_read_b64 v[182:183], %[v_sld_y_os] offset:13440 + %[sld_a_base] \n" +" ds_read_b64 v[184:185], %[v_sld_y_os] offset:14336 + %[sld_a_base] \n" +" ds_read_b64 v[186:187], %[v_sld_y_os] offset:14464 + %[sld_a_base] \n" +" ds_read_b64 v[188:189], %[v_sld_y_os] offset:15360 + %[sld_a_base] \n" +" ds_read_b64 v[190:191], %[v_sld_y_os] offset:15488 + %[sld_a_base] \n" +" ds_read_b64 v[192:193], %[v_sld_y_os] offset:16384 + %[sld_a_base] \n" +" ds_read_b64 v[194:195], %[v_sld_y_os] offset:16512 + %[sld_a_base] \n" +" ds_read_b64 v[196:197], %[v_sld_y_os] offset:17408 + %[sld_a_base] \n" +" ds_read_b64 v[198:199], %[v_sld_y_os] offset:17536 + %[sld_a_base] \n" +" ds_read_b64 v[200:201], %[v_sld_y_os] offset:18432 + %[sld_a_base] \n" +" ds_read_b64 v[202:203], %[v_sld_y_os] offset:18560 + %[sld_a_base] \n" +" ds_read_b64 v[204:205], %[v_sld_y_os] offset:19456 + %[sld_a_base] \n" +" ds_read_b64 v[206:207], %[v_sld_y_os] offset:19584 + %[sld_a_base] \n" +" ds_read_b64 v[208:209], %[v_sld_y_os] offset:20480 + %[sld_a_base] \n" +" ds_read_b64 v[210:211], %[v_sld_y_os] offset:20608 + %[sld_a_base] \n" +" ds_read_b64 v[212:213], %[v_sld_y_os] offset:21504 + %[sld_a_base] \n" +" ds_read_b64 v[214:215], %[v_sld_y_os] offset:21632 + %[sld_a_base] \n" +" ds_read_b64 v[216:217], %[v_sld_y_os] offset:22528 + %[sld_a_base] \n" +" ds_read_b64 v[218:219], %[v_sld_y_os] offset:22656 + %[sld_a_base] \n" +" ds_read_b64 v[220:221], %[v_sld_y_os] offset:23552 + %[sld_a_base] \n" +" ds_read_b64 v[222:223], %[v_sld_y_os] offset:23680 + %[sld_a_base] \n" +" ds_read_b64 v[224:225], %[v_sld_y_os] offset:24576 + %[sld_a_base] \n" +" ds_read_b64 v[226:227], %[v_sld_y_os] offset:24704 + %[sld_a_base] \n" +" ds_read_b64 v[228:229], %[v_sld_y_os] offset:25600 + %[sld_a_base] \n" +" ds_read_b64 v[230:231], %[v_sld_y_os] offset:25728 + %[sld_a_base] \n" +" ds_read_b64 v[232:233], %[v_sld_y_os] offset:26624 + %[sld_a_base] \n" +" ds_read_b64 v[234:235], %[v_sld_y_os] offset:26752 + %[sld_a_base] \n" +" ds_read_b64 v[236:237], %[v_sld_y_os] offset:27648 + %[sld_a_base] \n" +" ds_read_b64 v[238:239], %[v_sld_y_os] offset:27776 + %[sld_a_base] \n" +" ds_read_b64 v[240:241], %[v_sld_y_os] offset:28672 + %[sld_a_base] \n" +" ds_read_b64 v[242:243], %[v_sld_y_os] offset:28800 + %[sld_a_base] \n" +" ds_read_b64 v[244:245], %[v_sld_y_os] offset:29696 + %[sld_a_base] \n" +" ds_read_b64 v[246:247], %[v_sld_y_os] offset:29824 + %[sld_a_base] \n" +" ds_read_b64 v[248:249], %[v_sld_y_os] offset:30720 + %[sld_a_base] \n" +" ds_read_b64 v[250:251], %[v_sld_y_os] offset:30848 + %[sld_a_base] \n" +" ds_read_b64 v[252:253], %[v_sld_y_os] offset:31744 + %[sld_a_base] \n" +" ds_read_b64 v[254:255], %[v_sld_y_os] offset:31872 + %[sld_a_base] \n" +" s_waitcnt 0 \n" +" buffer_load_dwordx4 acc[0:3], %[v_os_b0], s[12:15], 0 offen \n" +" buffer_load_dwordx4 acc[4:7], %[v_os_b0], s[12:15], 0 offen offset:1024 \n" +" buffer_load_dwordx4 acc[8:11], %[v_os_b0], s[12:15], 0 offen offset:2048 \n" +" buffer_load_dwordx4 acc[12:15], %[v_os_b0], s[12:15], 0 offen offset:3072 \n" +" buffer_load_dwordx4 acc[16:19], %[v_os_b1], s[12:15], 0 offen \n" +" buffer_load_dwordx4 acc[20:23], %[v_os_b1], s[12:15], 0 offen offset:1024 \n" +" buffer_load_dwordx4 acc[24:27], %[v_os_b1], s[12:15], 0 offen offset:2048 \n" +" buffer_load_dwordx4 acc[28:31], %[v_os_b1], s[12:15], 0 offen offset:3072 \n" +" buffer_load_dwordx4 acc[32:35], %[v_os_b2], s[12:15], 0 offen \n" +" buffer_load_dwordx4 acc[36:39], %[v_os_b2], s[12:15], 0 offen offset:1024 \n" +" buffer_load_dwordx4 acc[40:43], %[v_os_b2], s[12:15], 0 offen offset:2048 \n" +" buffer_load_dwordx4 acc[44:47], %[v_os_b2], s[12:15], 0 offen offset:3072 \n" +" buffer_load_dwordx4 acc[48:51], %[v_os_b3], s[12:15], 0 offen \n" +" buffer_load_dwordx4 acc[52:55], %[v_os_b3], s[12:15], 0 offen offset:1024 \n" +" buffer_load_dwordx4 acc[56:59], %[v_os_b3], s[12:15], 0 offen offset:2048 \n" +" buffer_load_dwordx4 acc[60:63], %[v_os_b3], s[12:15], 0 offen offset:3072 \n" +" buffer_load_dwordx4 acc[64:67], %[v_os_b4], s[12:15], 0 offen \n" +" buffer_load_dwordx4 acc[68:71], %[v_os_b4], s[12:15], 0 offen offset:1024 \n" +" buffer_load_dwordx4 acc[72:75], %[v_os_b4], s[12:15], 0 offen offset:2048 \n" +" buffer_load_dwordx4 acc[76:79], %[v_os_b4], s[12:15], 0 offen offset:3072 \n" +" buffer_load_dwordx4 acc[80:83], %[v_os_b5], s[12:15], 0 offen \n" +" buffer_load_dwordx4 acc[84:87], %[v_os_b5], s[12:15], 0 offen offset:1024 \n" +" buffer_load_dwordx4 acc[88:91], %[v_os_b5], s[12:15], 0 offen offset:2048 \n" +" buffer_load_dwordx4 acc[92:95], %[v_os_b5], s[12:15], 0 offen offset:3072 \n" +" buffer_load_dwordx4 acc[96:99], %[v_os_b6], s[12:15], 0 offen \n" +" buffer_load_dwordx4 acc[100:103], %[v_os_b6], s[12:15], 0 offen offset:1024 \n" +" buffer_load_dwordx4 acc[104:107], %[v_os_b6], s[12:15], 0 offen offset:2048 \n" +" buffer_load_dwordx4 acc[108:111], %[v_os_b6], s[12:15], 0 offen offset:3072 \n" +" buffer_load_dwordx4 acc[112:115], %[v_os_b7], s[12:15], 0 offen \n" +" buffer_load_dwordx4 acc[116:119], %[v_os_b7], s[12:15], 0 offen offset:1024 \n" +" buffer_load_dwordx4 acc[120:123], %[v_os_b7], s[12:15], 0 offen offset:2048 \n" +" buffer_load_dwordx4 acc[124:127], %[v_os_b7], s[12:15], 0 offen offset:3072 \n" +" s_add_u32 s12, %[s_tile_os_b], s12 \n" +" s_addc_u32 s13, 0, s13 \n" +" v_mov_b32 v64, 0 \n" +" v_mov_b32 v80, 0 \n" +" v_mov_b32 v65, 0 \n" +" v_mov_b32 v81, 0 \n" +" v_mov_b32 v66, 0 \n" +" v_mov_b32 v82, 0 \n" +" v_mov_b32 v67, 0 \n" +" v_mov_b32 v83, 0 \n" +" v_mov_b32 v68, 0 \n" +" v_mov_b32 v84, 0 \n" +" v_mov_b32 v69, 0 \n" +" v_mov_b32 v85, 0 \n" +" v_mov_b32 v70, 0 \n" +" v_mov_b32 v86, 0 \n" +" v_mov_b32 v71, 0 \n" +" v_mov_b32 v87, 0 \n" +" ds_write_b64 %[v_sfl_sst], [%[c0],%[c1]] offset:16640 \n" +" ds_write_b64 %[v_sfl_sst], [%[c2],%[c3]] offset:20992 \n" +" ds_write_b64 %[v_sfl_sst], [%[c4],%[c5]] offset:18816 \n" +" ds_write_b64 %[v_sfl_sst], [%[c6],%[c7]] offset:23168 \n" +" s_mov_b32 s80, 0 \n" +" s_waitcnt vmcnt(24) \n" +"label_0AA6: \n" +" s_waitcnt vmcnt(30) & lgkmcnt(0) \n" +" s_barrier \n" + _UK_MFMA_ " [%[c0], %[c1], %[c2], %[c3]], acc[0:1], v[128:129], 0 \n" +" ds_read_b32 v10, %[v_sfl_sld] offset:16640 \n" +" ds_read_b32 v11, %[v_sfl_sld] offset:16672 \n" +" ds_write_b64 %[v_sfl_sst], [%[c16],%[c17]] offset:25344 \n" + _UK_MFMA_ " [%[c0], %[c1], %[c2], %[c3]], acc[2:3], v[130:131], v[64:67] \n" + " buffer_load_dwordx4 acc[128:131], %[v_os_b0], s[12:15], 0 offen \n" +" ds_write_b64 %[v_sfl_sst], [%[c18],%[c19]] offset:29696 \n" + _UK_MFMA_ " [%[c0], %[c1], %[c2], %[c3]], acc[4:5], v[132:133], v[64:67] \n" +" ds_read_b32 v12, %[v_sfl_sld] offset:16704 \n" +" ds_read_b32 v13, %[v_sfl_sld] offset:16736 \n" +" ds_write_b64 %[v_sfl_sst], [%[c20],%[c21]] offset:27520 \n" + _UK_MFMA_ " [%[c0], %[c1], %[c2], %[c3]], acc[6:7], v[134:135], v[64:67] \n" +" ds_write_b64 %[v_sfl_sst], [%[c22],%[c23]] offset:31872 \n" + _UK_MFMA_ " [%[c0], %[c1], %[c2], %[c3]], acc[8:9], v[136:137], v[64:67] \n" +" ds_read_b32 v14, %[v_sfl_sld] offset:20992 \n" +" ds_read_b32 v15, %[v_sfl_sld] offset:21024 \n" + _UK_MFMA_ " [%[c0], %[c1], %[c2], %[c3]], acc[10:11], v[138:139], v[64:67] \n" + " buffer_load_dwordx4 acc[132:135], %[v_os_b0], s[12:15], 0 offen offset:1024 \n" + _UK_MFMA_ " [%[c0], %[c1], %[c2], %[c3]], acc[12:13], v[140:141], v[64:67] \n" +" ds_read_b32 v16, %[v_sfl_sld] offset:21056 \n" +" ds_read_b32 v17, %[v_sfl_sld] offset:21088 \n" + _UK_MFMA_ " [%[c0], %[c1], %[c2], %[c3]], acc[14:15], v[142:143], v[64:67] \n" + _UK_MFMA_ " [%[c4], %[c5], %[c6], %[c7]], acc[0:1], v[192:193], 0 \n" + _UK_MFMA_ " [%[c4], %[c5], %[c6], %[c7]], acc[2:3], v[194:195], v[68:71] \n" + " buffer_load_dwordx4 acc[136:139], %[v_os_b0], s[12:15], 0 offen offset:2048 \n" + _UK_MFMA_ " [%[c4], %[c5], %[c6], %[c7]], acc[4:5], v[196:197], v[68:71] \n" + _UK_MFMA_ " [%[c4], %[c5], %[c6], %[c7]], acc[6:7], v[198:199], v[68:71] \n" + _UK_MFMA_ " [%[c4], %[c5], %[c6], %[c7]], acc[8:9], v[200:201], v[68:71] \n" + _UK_MFMA_ " [%[c4], %[c5], %[c6], %[c7]], acc[10:11], v[202:203], v[68:71] \n" + " buffer_load_dwordx4 acc[140:143], %[v_os_b0], s[12:15], 0 offen offset:3072 \n" + _UK_MFMA_ " [%[c4], %[c5], %[c6], %[c7]], acc[12:13], v[204:205], v[68:71] \n" + _UK_MFMA_ " [%[c4], %[c5], %[c6], %[c7]], acc[14:15], v[206:207], v[68:71] \n" + " s_waitcnt lgkmcnt(0) \n" + " s_mov_b64 exec, %[s_execflag_0] \n" +_UK_ATOMIC_ADD_ " %[v_os_o0], v10, s[8:9] \n" +" s_mov_b64 exec, s[38:39] \n" + _UK_MFMA_ " [%[c8], %[c9], %[c10], %[c11]], acc[16:17], v[128:129], 0 \n" + _UK_MFMA_ " [%[c8], %[c9], %[c10], %[c11]], acc[18:19], v[130:131], v[72:75] \n" + " buffer_load_dwordx4 acc[144:147], %[v_os_b1], s[12:15], 0 offen \n" + _UK_MFMA_ " [%[c8], %[c9], %[c10], %[c11]], acc[20:21], v[132:133], v[72:75] \n" + _UK_MFMA_ " [%[c8], %[c9], %[c10], %[c11]], acc[22:23], v[134:135], v[72:75] \n" + _UK_MFMA_ " [%[c8], %[c9], %[c10], %[c11]], acc[24:25], v[136:137], v[72:75] \n" + _UK_MFMA_ " [%[c8], %[c9], %[c10], %[c11]], acc[26:27], v[138:139], v[72:75] \n" + " buffer_load_dwordx4 acc[148:151], %[v_os_b1], s[12:15], 0 offen offset:1024 \n" + _UK_MFMA_ " [%[c8], %[c9], %[c10], %[c11]], acc[28:29], v[140:141], v[72:75] \n" + _UK_MFMA_ " [%[c8], %[c9], %[c10], %[c11]], acc[30:31], v[142:143], v[72:75] \n" + _UK_MFMA_ " [%[c12], %[c13], %[c14], %[c15]], acc[16:17], v[192:193], 0 \n" + _UK_MFMA_ " [%[c12], %[c13], %[c14], %[c15]], acc[18:19], v[194:195], v[76:79] \n" + " buffer_load_dwordx4 acc[152:155], %[v_os_b1], s[12:15], 0 offen offset:2048 \n" + _UK_MFMA_ " [%[c12], %[c13], %[c14], %[c15]], acc[20:21], v[196:197], v[76:79] \n" + _UK_MFMA_ " [%[c12], %[c13], %[c14], %[c15]], acc[22:23], v[198:199], v[76:79] \n" + _UK_MFMA_ " [%[c12], %[c13], %[c14], %[c15]], acc[24:25], v[200:201], v[76:79] \n" + _UK_MFMA_ " [%[c12], %[c13], %[c14], %[c15]], acc[26:27], v[202:203], v[76:79] \n" + " buffer_load_dwordx4 acc[156:159], %[v_os_b1], s[12:15], 0 offen offset:3072 \n" + _UK_MFMA_ " [%[c12], %[c13], %[c14], %[c15]], acc[28:29], v[204:205], v[76:79] \n" + _UK_MFMA_ " [%[c12], %[c13], %[c14], %[c15]], acc[30:31], v[206:207], v[76:79] \n" + " s_mov_b64 exec, %[s_execflag_1] \n" +_UK_ATOMIC_ADD_ " %[v_os_o1], v11, s[8:9] \n" +" s_mov_b64 exec, s[38:39] \n" +" s_waitcnt vmcnt(30) \n" + _UK_MFMA_ " [%[c0], %[c1], %[c2], %[c3]], acc[32:33], v[144:145], v[64:67] \n" + _UK_MFMA_ " [%[c0], %[c1], %[c2], %[c3]], acc[34:35], v[146:147], v[64:67] \n" + " buffer_load_dwordx4 acc[160:163], %[v_os_b2], s[12:15], 0 offen \n" + _UK_MFMA_ " [%[c0], %[c1], %[c2], %[c3]], acc[36:37], v[148:149], v[64:67] \n" + _UK_MFMA_ " [%[c0], %[c1], %[c2], %[c3]], acc[38:39], v[150:151], v[64:67] \n" + _UK_MFMA_ " [%[c0], %[c1], %[c2], %[c3]], acc[40:41], v[152:153], v[64:67] \n" + _UK_MFMA_ " [%[c0], %[c1], %[c2], %[c3]], acc[42:43], v[154:155], v[64:67] \n" + " buffer_load_dwordx4 acc[164:167], %[v_os_b2], s[12:15], 0 offen offset:1024 \n" + _UK_MFMA_ " [%[c0], %[c1], %[c2], %[c3]], acc[44:45], v[156:157], v[64:67] \n" + _UK_MFMA_ " [%[c0], %[c1], %[c2], %[c3]], acc[46:47], v[158:159], v[64:67] \n" + _UK_MFMA_ " [%[c4], %[c5], %[c6], %[c7]], acc[32:33], v[208:209], v[68:71] \n" + _UK_MFMA_ " [%[c4], %[c5], %[c6], %[c7]], acc[34:35], v[210:211], v[68:71] \n" + " buffer_load_dwordx4 acc[168:171], %[v_os_b2], s[12:15], 0 offen offset:2048 \n" + _UK_MFMA_ " [%[c4], %[c5], %[c6], %[c7]], acc[36:37], v[212:213], v[68:71] \n" + _UK_MFMA_ " [%[c4], %[c5], %[c6], %[c7]], acc[38:39], v[214:215], v[68:71] \n" + _UK_MFMA_ " [%[c4], %[c5], %[c6], %[c7]], acc[40:41], v[216:217], v[68:71] \n" + _UK_MFMA_ " [%[c4], %[c5], %[c6], %[c7]], acc[42:43], v[218:219], v[68:71] \n" + " buffer_load_dwordx4 acc[172:175], %[v_os_b2], s[12:15], 0 offen offset:3072 \n" + _UK_MFMA_ " [%[c4], %[c5], %[c6], %[c7]], acc[44:45], v[220:221], v[68:71] \n" + _UK_MFMA_ " [%[c4], %[c5], %[c6], %[c7]], acc[46:47], v[222:223], v[68:71] \n" + " s_mov_b64 exec, %[s_execflag_2] \n" +_UK_ATOMIC_ADD_ " %[v_os_o2], v12, s[8:9] \n" +" s_mov_b64 exec, s[38:39] \n" + _UK_MFMA_ " [%[c8], %[c9], %[c10], %[c11]], acc[48:49], v[144:145], v[72:75] \n" + _UK_MFMA_ " [%[c8], %[c9], %[c10], %[c11]], acc[50:51], v[146:147], v[72:75] \n" + " buffer_load_dwordx4 acc[176:179], %[v_os_b3], s[12:15], 0 offen \n" + _UK_MFMA_ " [%[c8], %[c9], %[c10], %[c11]], acc[52:53], v[148:149], v[72:75] \n" + _UK_MFMA_ " [%[c8], %[c9], %[c10], %[c11]], acc[54:55], v[150:151], v[72:75] \n" + _UK_MFMA_ " [%[c8], %[c9], %[c10], %[c11]], acc[56:57], v[152:153], v[72:75] \n" + _UK_MFMA_ " [%[c8], %[c9], %[c10], %[c11]], acc[58:59], v[154:155], v[72:75] \n" + " buffer_load_dwordx4 acc[180:183], %[v_os_b3], s[12:15], 0 offen offset:1024 \n" + _UK_MFMA_ " [%[c8], %[c9], %[c10], %[c11]], acc[60:61], v[156:157], v[72:75] \n" + _UK_MFMA_ " [%[c8], %[c9], %[c10], %[c11]], acc[62:63], v[158:159], v[72:75] \n" + _UK_MFMA_ " [%[c12], %[c13], %[c14], %[c15]], acc[48:49], v[208:209], v[76:79] \n" + _UK_MFMA_ " [%[c12], %[c13], %[c14], %[c15]], acc[50:51], v[210:211], v[76:79] \n" + " buffer_load_dwordx4 acc[184:187], %[v_os_b3], s[12:15], 0 offen offset:2048 \n" + _UK_MFMA_ " [%[c12], %[c13], %[c14], %[c15]], acc[52:53], v[212:213], v[76:79] \n" + _UK_MFMA_ " [%[c12], %[c13], %[c14], %[c15]], acc[54:55], v[214:215], v[76:79] \n" + _UK_MFMA_ " [%[c12], %[c13], %[c14], %[c15]], acc[56:57], v[216:217], v[76:79] \n" + _UK_MFMA_ " [%[c12], %[c13], %[c14], %[c15]], acc[58:59], v[218:219], v[76:79] \n" + " buffer_load_dwordx4 acc[188:191], %[v_os_b3], s[12:15], 0 offen offset:3072 \n" + _UK_MFMA_ " [%[c12], %[c13], %[c14], %[c15]], acc[60:61], v[220:221], v[76:79] \n" + _UK_MFMA_ " [%[c12], %[c13], %[c14], %[c15]], acc[62:63], v[222:223], v[76:79] \n" + " s_mov_b64 exec, %[s_execflag_3] \n" +_UK_ATOMIC_ADD_ " %[v_os_o3], v13, s[8:9] \n" +" s_mov_b64 exec, s[38:39] \n" +" s_waitcnt vmcnt(30) \n" + _UK_MFMA_ " [%[c0], %[c1], %[c2], %[c3]], acc[64:65], v[160:161], v[64:67] \n" + _UK_MFMA_ " [%[c0], %[c1], %[c2], %[c3]], acc[66:67], v[162:163], v[64:67] \n" + " buffer_load_dwordx4 acc[192:195], %[v_os_b4], s[12:15], 0 offen \n" + _UK_MFMA_ " [%[c0], %[c1], %[c2], %[c3]], acc[68:69], v[164:165], v[64:67] \n" + _UK_MFMA_ " [%[c0], %[c1], %[c2], %[c3]], acc[70:71], v[166:167], v[64:67] \n" + _UK_MFMA_ " [%[c0], %[c1], %[c2], %[c3]], acc[72:73], v[168:169], v[64:67] \n" + _UK_MFMA_ " [%[c0], %[c1], %[c2], %[c3]], acc[74:75], v[170:171], v[64:67] \n" + " buffer_load_dwordx4 acc[196:199], %[v_os_b4], s[12:15], 0 offen offset:1024 \n" + _UK_MFMA_ " [%[c0], %[c1], %[c2], %[c3]], acc[76:77], v[172:173], v[64:67] \n" + _UK_MFMA_ " [%[c0], %[c1], %[c2], %[c3]], acc[78:79], v[174:175], v[64:67] \n" + _UK_MFMA_ " [%[c4], %[c5], %[c6], %[c7]], acc[64:65], v[224:225], v[68:71] \n" + _UK_MFMA_ " [%[c4], %[c5], %[c6], %[c7]], acc[66:67], v[226:227], v[68:71] \n" + " buffer_load_dwordx4 acc[200:203], %[v_os_b4], s[12:15], 0 offen offset:2048 \n" + _UK_MFMA_ " [%[c4], %[c5], %[c6], %[c7]], acc[68:69], v[228:229], v[68:71] \n" + _UK_MFMA_ " [%[c4], %[c5], %[c6], %[c7]], acc[70:71], v[230:231], v[68:71] \n" + _UK_MFMA_ " [%[c4], %[c5], %[c6], %[c7]], acc[72:73], v[232:233], v[68:71] \n" + _UK_MFMA_ " [%[c4], %[c5], %[c6], %[c7]], acc[74:75], v[234:235], v[68:71] \n" + " buffer_load_dwordx4 acc[204:207], %[v_os_b4], s[12:15], 0 offen offset:3072 \n" + _UK_MFMA_ " [%[c4], %[c5], %[c6], %[c7]], acc[76:77], v[236:237], v[68:71] \n" + _UK_MFMA_ " [%[c4], %[c5], %[c6], %[c7]], acc[78:79], v[238:239], v[68:71] \n" + " s_mov_b64 exec, %[s_execflag_4] \n" +_UK_ATOMIC_ADD_ " %[v_os_o4], v14, s[8:9] \n" +" s_mov_b64 exec, s[38:39] \n" + _UK_MFMA_ " [%[c8], %[c9], %[c10], %[c11]], acc[80:81], v[160:161], v[72:75] \n" + _UK_MFMA_ " [%[c8], %[c9], %[c10], %[c11]], acc[82:83], v[162:163], v[72:75] \n" + " buffer_load_dwordx4 acc[208:211], %[v_os_b5], s[12:15], 0 offen \n" + _UK_MFMA_ " [%[c8], %[c9], %[c10], %[c11]], acc[84:85], v[164:165], v[72:75] \n" + _UK_MFMA_ " [%[c8], %[c9], %[c10], %[c11]], acc[86:87], v[166:167], v[72:75] \n" + _UK_MFMA_ " [%[c8], %[c9], %[c10], %[c11]], acc[88:89], v[168:169], v[72:75] \n" + _UK_MFMA_ " [%[c8], %[c9], %[c10], %[c11]], acc[90:91], v[170:171], v[72:75] \n" + " buffer_load_dwordx4 acc[212:215], %[v_os_b5], s[12:15], 0 offen offset:1024 \n" + _UK_MFMA_ " [%[c8], %[c9], %[c10], %[c11]], acc[92:93], v[172:173], v[72:75] \n" + _UK_MFMA_ " [%[c8], %[c9], %[c10], %[c11]], acc[94:95], v[174:175], v[72:75] \n" + _UK_MFMA_ " [%[c12], %[c13], %[c14], %[c15]], acc[80:81], v[224:225], v[76:79] \n" + _UK_MFMA_ " [%[c12], %[c13], %[c14], %[c15]], acc[82:83], v[226:227], v[76:79] \n" + " buffer_load_dwordx4 acc[216:219], %[v_os_b5], s[12:15], 0 offen offset:2048 \n" + _UK_MFMA_ " [%[c12], %[c13], %[c14], %[c15]], acc[84:85], v[228:229], v[76:79] \n" + _UK_MFMA_ " [%[c12], %[c13], %[c14], %[c15]], acc[86:87], v[230:231], v[76:79] \n" + _UK_MFMA_ " [%[c12], %[c13], %[c14], %[c15]], acc[88:89], v[232:233], v[76:79] \n" + _UK_MFMA_ " [%[c12], %[c13], %[c14], %[c15]], acc[90:91], v[234:235], v[76:79] \n" + " buffer_load_dwordx4 acc[220:223], %[v_os_b5], s[12:15], 0 offen offset:3072 \n" + _UK_MFMA_ " [%[c12], %[c13], %[c14], %[c15]], acc[92:93], v[236:237], v[76:79] \n" + _UK_MFMA_ " [%[c12], %[c13], %[c14], %[c15]], acc[94:95], v[238:239], v[76:79] \n" + " s_mov_b64 exec, %[s_execflag_5] \n" +_UK_ATOMIC_ADD_ " %[v_os_o5], v15, s[8:9] \n" +" s_mov_b64 exec, s[38:39] \n" +" s_waitcnt vmcnt(30) \n" + _UK_MFMA_ " [%[c0], %[c1], %[c2], %[c3]], acc[96:97], v[176:177], v[64:67] \n" + _UK_MFMA_ " [%[c0], %[c1], %[c2], %[c3]], acc[98:99], v[178:179], v[64:67] \n" + " buffer_load_dwordx4 acc[224:227], %[v_os_b6], s[12:15], 0 offen \n" + _UK_MFMA_ " [%[c0], %[c1], %[c2], %[c3]], acc[100:101], v[180:181], v[64:67] \n" + _UK_MFMA_ " [%[c0], %[c1], %[c2], %[c3]], acc[102:103], v[182:183], v[64:67] \n" + _UK_MFMA_ " [%[c0], %[c1], %[c2], %[c3]], acc[104:105], v[184:185], v[64:67] \n" + _UK_MFMA_ " [%[c0], %[c1], %[c2], %[c3]], acc[106:107], v[186:187], v[64:67] \n" + " buffer_load_dwordx4 acc[228:231], %[v_os_b6], s[12:15], 0 offen offset:1024 \n" + _UK_MFMA_ " [%[c0], %[c1], %[c2], %[c3]], acc[108:109], v[188:189], v[64:67] \n" + _UK_MFMA_ " [%[c0], %[c1], %[c2], %[c3]], acc[110:111], v[190:191], v[64:67] \n" + _UK_MFMA_ " [%[c4], %[c5], %[c6], %[c7]], acc[96:97], v[240:241], v[68:71] \n" + _UK_MFMA_ " [%[c4], %[c5], %[c6], %[c7]], acc[98:99], v[242:243], v[68:71] \n" + " buffer_load_dwordx4 acc[232:235], %[v_os_b6], s[12:15], 0 offen offset:2048 \n" + _UK_MFMA_ " [%[c4], %[c5], %[c6], %[c7]], acc[100:101], v[244:245], v[68:71] \n" + _UK_MFMA_ " [%[c4], %[c5], %[c6], %[c7]], acc[102:103], v[246:247], v[68:71] \n" + _UK_MFMA_ " [%[c4], %[c5], %[c6], %[c7]], acc[104:105], v[248:249], v[68:71] \n" + _UK_MFMA_ " [%[c4], %[c5], %[c6], %[c7]], acc[106:107], v[250:251], v[68:71] \n" + " buffer_load_dwordx4 acc[236:239], %[v_os_b6], s[12:15], 0 offen offset:3072 \n" + _UK_MFMA_ " [%[c4], %[c5], %[c6], %[c7]], acc[108:109], v[252:253], v[68:71] \n" + _UK_MFMA_ " [%[c4], %[c5], %[c6], %[c7]], acc[110:111], v[254:255], v[68:71] \n" + " s_mov_b64 exec, %[s_execflag_6] \n" +_UK_ATOMIC_ADD_ " %[v_os_o6], v16, s[8:9] \n" +" s_mov_b64 exec, s[38:39] \n" + _UK_MFMA_ " [%[c8], %[c9], %[c10], %[c11]], acc[112:113], v[176:177], v[72:75] \n" + _UK_MFMA_ " [%[c8], %[c9], %[c10], %[c11]], acc[114:115], v[178:179], v[72:75] \n" + " buffer_load_dwordx4 acc[240:243], %[v_os_b7], s[12:15], 0 offen \n" + _UK_MFMA_ " [%[c8], %[c9], %[c10], %[c11]], acc[116:117], v[180:181], v[72:75] \n" + _UK_MFMA_ " [%[c8], %[c9], %[c10], %[c11]], acc[118:119], v[182:183], v[72:75] \n" + _UK_MFMA_ " [%[c8], %[c9], %[c10], %[c11]], acc[120:121], v[184:185], v[72:75] \n" + _UK_MFMA_ " [%[c8], %[c9], %[c10], %[c11]], acc[122:123], v[186:187], v[72:75] \n" + " buffer_load_dwordx4 acc[244:247], %[v_os_b7], s[12:15], 0 offen offset:1024 \n" + _UK_MFMA_ " [%[c8], %[c9], %[c10], %[c11]], acc[124:125], v[188:189], v[72:75] \n" + _UK_MFMA_ " [%[c8], %[c9], %[c10], %[c11]], acc[126:127], v[190:191], v[72:75] \n" + _UK_MFMA_ " [%[c12], %[c13], %[c14], %[c15]], acc[112:113], v[240:241], v[76:79] \n" + _UK_MFMA_ " [%[c12], %[c13], %[c14], %[c15]], acc[114:115], v[242:243], v[76:79] \n" + " buffer_load_dwordx4 acc[248:251], %[v_os_b7], s[12:15], 0 offen offset:2048 \n" + _UK_MFMA_ " [%[c12], %[c13], %[c14], %[c15]], acc[116:117], v[244:245], v[76:79] \n" + _UK_MFMA_ " [%[c12], %[c13], %[c14], %[c15]], acc[118:119], v[246:247], v[76:79] \n" + _UK_MFMA_ " [%[c12], %[c13], %[c14], %[c15]], acc[120:121], v[248:249], v[76:79] \n" + _UK_MFMA_ " [%[c12], %[c13], %[c14], %[c15]], acc[122:123], v[250:251], v[76:79] \n" + " buffer_load_dwordx4 acc[252:255], %[v_os_b7], s[12:15], 0 offen offset:3072 \n" + _UK_MFMA_ " [%[c12], %[c13], %[c14], %[c15]], acc[124:125], v[252:253], v[76:79] \n" + _UK_MFMA_ " [%[c12], %[c13], %[c14], %[c15]], acc[126:127], v[254:255], v[76:79] \n" + " s_mov_b64 exec, %[s_execflag_7] \n" +_UK_ATOMIC_ADD_ " %[v_os_o7], v17, s[8:9] \n" +" s_mov_b64 exec, s[38:39] \n" +" s_add_u32 s60, 0x00000100, s80 \n" +" s_cmp_lt_u32 s60, %[s_loop_cnt] \n" +" s_cselect_b32 s56, %[s_tile_os_b], 0 \n" +" s_add_u32 s12, s56, s12 \n" +" s_addc_u32 s13, 0, s13 \n" +" s_cmp_ge_u32 s80, 0x00000100 \n" +" s_cselect_b32 s59, %[s_tile_os_o], s59 \n" +" s_add_u32 s8, s59, s8 \n" +" s_addc_u32 s9, 0, s9 \n" +" v_mul_f32 %[c0], %[scale_0], %[c0] \n" +" v_mul_f32 %[c1], %[scale_0], %[c1] \n" +" v_mul_f32 %[c2], %[scale_0], %[c2] \n" +" v_mul_f32 %[c3], %[scale_0], %[c3] \n" +" v_mul_f32 %[c4], %[scale_1], %[c4] \n" +" v_mul_f32 %[c5], %[scale_1], %[c5] \n" +" v_mul_f32 %[c6], %[scale_1], %[c6] \n" +" v_mul_f32 %[c7], %[scale_1], %[c7] \n" +" v_mul_f32 %[c8], %[scale_0], %[c8] \n" +" v_mul_f32 %[c9], %[scale_0], %[c9] \n" +" v_mul_f32 %[c10], %[scale_0], %[c10] \n" +" v_mul_f32 %[c11], %[scale_0], %[c11] \n" +" v_mul_f32 %[c12], %[scale_1], %[c12] \n" +" v_mul_f32 %[c13], %[scale_1], %[c13] \n" +" v_mul_f32 %[c14], %[scale_1], %[c14] \n" +" v_mul_f32 %[c15], %[scale_1], %[c15] \n" +_UK_PK_CVT_("%[c0]", "%[c1]", "%[c0]") +_UK_PK_CVT_("%[c2]", "%[c3]", "%[c1]") +_UK_PK_CVT_("%[c4]", "%[c5]", "%[c2]") +_UK_PK_CVT_("%[c6]", "%[c7]", "%[c3]") +_UK_PK_CVT_("%[c8]", "%[c9]", "%[c4]") +_UK_PK_CVT_("%[c10]", "%[c11]", "%[c5]") +_UK_PK_CVT_("%[c12]", "%[c13]", "%[c6]") +_UK_PK_CVT_("%[c14]", "%[c15]", "%[c7]") +" s_addk_i32 s80, 0x0080 \n" +" s_cmp_lt_i32 s80, %[s_loop_cnt] \n" +" s_cbranch_scc0 label_0EC1 \n" +" s_waitcnt vmcnt(30) & lgkmcnt(0) \n" +" s_barrier \n" + _UK_MFMA_ " [%[c16], %[c17], %[c18], %[c19]], acc[128:129], v[128:129], 0 \n" +" ds_read_b32 v10, %[v_sfl_sld] offset:25344 \n" +" ds_read_b32 v11, %[v_sfl_sld] offset:25376 \n" +" ds_write_b64 v3, v[64:65] offset:16640 \n" + _UK_MFMA_ " [%[c16], %[c17], %[c18], %[c19]], acc[130:131], v[130:131], v[80:83] \n" + " buffer_load_dwordx4 acc[0:3], %[v_os_b0], s[12:15], 0 offen \n" +" ds_write_b64 v3, v[66:67] offset:20992 \n" + _UK_MFMA_ " [%[c16], %[c17], %[c18], %[c19]], acc[132:133], v[132:133], v[80:83] \n" +" ds_read_b32 v12, %[v_sfl_sld] offset:25408 \n" +" ds_read_b32 v13, %[v_sfl_sld] offset:25440 \n" +" ds_write_b64 v3, v[68:69] offset:18816 \n" + _UK_MFMA_ " [%[c16], %[c17], %[c18], %[c19]], acc[134:135], v[134:135], v[80:83] \n" +" ds_write_b64 v3, v[70:71] offset:23168 \n" + _UK_MFMA_ " [%[c16], %[c17], %[c18], %[c19]], acc[136:137], v[136:137], v[80:83] \n" +" ds_read_b32 v14, %[v_sfl_sld] offset:29696 \n" +" ds_read_b32 v15, %[v_sfl_sld] offset:29728 \n" + _UK_MFMA_ " [%[c16], %[c17], %[c18], %[c19]], acc[138:139], v[138:139], v[80:83] \n" + " buffer_load_dwordx4 acc[4:7], %[v_os_b0], s[12:15], 0 offen offset:1024 \n" + _UK_MFMA_ " [%[c16], %[c17], %[c18], %[c19]], acc[140:141], v[140:141], v[80:83] \n" +" ds_read_b32 v16, %[v_sfl_sld] offset:29760 \n" +" ds_read_b32 v17, %[v_sfl_sld] offset:29792 \n" + _UK_MFMA_ " [%[c16], %[c17], %[c18], %[c19]], acc[142:143], v[142:143], v[80:83] \n" + _UK_MFMA_ " [%[c20], %[c21], %[c22], %[c23]], acc[128:129], v[192:193], 0 \n" + _UK_MFMA_ " [%[c20], %[c21], %[c22], %[c23]], acc[130:131], v[194:195], v[84:87] \n" + " buffer_load_dwordx4 acc[8:11], %[v_os_b0], s[12:15], 0 offen offset:2048 \n" + _UK_MFMA_ " [%[c20], %[c21], %[c22], %[c23]], acc[132:133], v[196:197], v[84:87] \n" + _UK_MFMA_ " [%[c20], %[c21], %[c22], %[c23]], acc[134:135], v[198:199], v[84:87] \n" + _UK_MFMA_ " [%[c20], %[c21], %[c22], %[c23]], acc[136:137], v[200:201], v[84:87] \n" + _UK_MFMA_ " [%[c20], %[c21], %[c22], %[c23]], acc[138:139], v[202:203], v[84:87] \n" + " buffer_load_dwordx4 acc[12:15], %[v_os_b0], s[12:15], 0 offen offset:3072 \n" + _UK_MFMA_ " [%[c20], %[c21], %[c22], %[c23]], acc[140:141], v[204:205], v[84:87] \n" + _UK_MFMA_ " [%[c20], %[c21], %[c22], %[c23]], acc[142:143], v[206:207], v[84:87] \n" + " s_waitcnt lgkmcnt(0) \n" + " s_mov_b64 exec, %[s_execflag_0] \n" +_UK_ATOMIC_ADD_ " %[v_os_o0], v10, s[8:9] \n" +" s_mov_b64 exec, s[38:39] \n" + _UK_MFMA_ " [%[c24], %[c25], %[c26], %[c27]], acc[144:145], v[128:129], 0 \n" + _UK_MFMA_ " [%[c24], %[c25], %[c26], %[c27]], acc[146:147], v[130:131], v[88:91] \n" + " buffer_load_dwordx4 acc[16:19], %[v_os_b1], s[12:15], 0 offen \n" + _UK_MFMA_ " [%[c24], %[c25], %[c26], %[c27]], acc[148:149], v[132:133], v[88:91] \n" + _UK_MFMA_ " [%[c24], %[c25], %[c26], %[c27]], acc[150:151], v[134:135], v[88:91] \n" + _UK_MFMA_ " [%[c24], %[c25], %[c26], %[c27]], acc[152:153], v[136:137], v[88:91] \n" + _UK_MFMA_ " [%[c24], %[c25], %[c26], %[c27]], acc[154:155], v[138:139], v[88:91] \n" + " buffer_load_dwordx4 acc[20:23], %[v_os_b1], s[12:15], 0 offen offset:1024 \n" + _UK_MFMA_ " [%[c24], %[c25], %[c26], %[c27]], acc[156:157], v[140:141], v[88:91] \n" + _UK_MFMA_ " [%[c24], %[c25], %[c26], %[c27]], acc[158:159], v[142:143], v[88:91] \n" + _UK_MFMA_ " [%[c28], %[c29], %[c30], %[c31]], acc[144:145], v[192:193], 0 \n" + _UK_MFMA_ " [%[c28], %[c29], %[c30], %[c31]], acc[146:147], v[194:195], v[92:95] \n" + " buffer_load_dwordx4 acc[24:27], %[v_os_b1], s[12:15], 0 offen offset:2048 \n" + _UK_MFMA_ " [%[c28], %[c29], %[c30], %[c31]], acc[148:149], v[196:197], v[92:95] \n" + _UK_MFMA_ " [%[c28], %[c29], %[c30], %[c31]], acc[150:151], v[198:199], v[92:95] \n" + _UK_MFMA_ " [%[c28], %[c29], %[c30], %[c31]], acc[152:153], v[200:201], v[92:95] \n" + _UK_MFMA_ " [%[c28], %[c29], %[c30], %[c31]], acc[154:155], v[202:203], v[92:95] \n" + " buffer_load_dwordx4 acc[28:31], %[v_os_b1], s[12:15], 0 offen offset:3072 \n" + _UK_MFMA_ " [%[c28], %[c29], %[c30], %[c31]], acc[156:157], v[204:205], v[92:95] \n" + _UK_MFMA_ " [%[c28], %[c29], %[c30], %[c31]], acc[158:159], v[206:207], v[92:95] \n" + " s_mov_b64 exec, %[s_execflag_1] \n" +_UK_ATOMIC_ADD_ " %[v_os_o1], v11, s[8:9] \n" +" s_mov_b64 exec, s[38:39] \n" +" s_waitcnt vmcnt(30) \n" + _UK_MFMA_ " [%[c16], %[c17], %[c18], %[c19]], acc[160:161], v[144:145], v[80:83] \n" + _UK_MFMA_ " [%[c16], %[c17], %[c18], %[c19]], acc[162:163], v[146:147], v[80:83] \n" + " buffer_load_dwordx4 acc[32:35], %[v_os_b2], s[12:15], 0 offen \n" + _UK_MFMA_ " [%[c16], %[c17], %[c18], %[c19]], acc[164:165], v[148:149], v[80:83] \n" + _UK_MFMA_ " [%[c16], %[c17], %[c18], %[c19]], acc[166:167], v[150:151], v[80:83] \n" + _UK_MFMA_ " [%[c16], %[c17], %[c18], %[c19]], acc[168:169], v[152:153], v[80:83] \n" + _UK_MFMA_ " [%[c16], %[c17], %[c18], %[c19]], acc[170:171], v[154:155], v[80:83] \n" + " buffer_load_dwordx4 acc[36:39], %[v_os_b2], s[12:15], 0 offen offset:1024 \n" + _UK_MFMA_ " [%[c16], %[c17], %[c18], %[c19]], acc[172:173], v[156:157], v[80:83] \n" + _UK_MFMA_ " [%[c16], %[c17], %[c18], %[c19]], acc[174:175], v[158:159], v[80:83] \n" + _UK_MFMA_ " [%[c20], %[c21], %[c22], %[c23]], acc[160:161], v[208:209], v[84:87] \n" + _UK_MFMA_ " [%[c20], %[c21], %[c22], %[c23]], acc[162:163], v[210:211], v[84:87] \n" + " buffer_load_dwordx4 acc[40:43], %[v_os_b2], s[12:15], 0 offen offset:2048 \n" + _UK_MFMA_ " [%[c20], %[c21], %[c22], %[c23]], acc[164:165], v[212:213], v[84:87] \n" + _UK_MFMA_ " [%[c20], %[c21], %[c22], %[c23]], acc[166:167], v[214:215], v[84:87] \n" + _UK_MFMA_ " [%[c20], %[c21], %[c22], %[c23]], acc[168:169], v[216:217], v[84:87] \n" + _UK_MFMA_ " [%[c20], %[c21], %[c22], %[c23]], acc[170:171], v[218:219], v[84:87] \n" + " buffer_load_dwordx4 acc[44:47], %[v_os_b2], s[12:15], 0 offen offset:3072 \n" + _UK_MFMA_ " [%[c20], %[c21], %[c22], %[c23]], acc[172:173], v[220:221], v[84:87] \n" + _UK_MFMA_ " [%[c20], %[c21], %[c22], %[c23]], acc[174:175], v[222:223], v[84:87] \n" + " s_mov_b64 exec, %[s_execflag_2] \n" +_UK_ATOMIC_ADD_ " %[v_os_o2], v12, s[8:9] \n" +" s_mov_b64 exec, s[38:39] \n" + _UK_MFMA_ " [%[c24], %[c25], %[c26], %[c27]], acc[176:177], v[144:145], v[88:91] \n" + _UK_MFMA_ " [%[c24], %[c25], %[c26], %[c27]], acc[178:179], v[146:147], v[88:91] \n" + " buffer_load_dwordx4 acc[48:51], %[v_os_b3], s[12:15], 0 offen \n" + _UK_MFMA_ " [%[c24], %[c25], %[c26], %[c27]], acc[180:181], v[148:149], v[88:91] \n" + _UK_MFMA_ " [%[c24], %[c25], %[c26], %[c27]], acc[182:183], v[150:151], v[88:91] \n" + _UK_MFMA_ " [%[c24], %[c25], %[c26], %[c27]], acc[184:185], v[152:153], v[88:91] \n" + _UK_MFMA_ " [%[c24], %[c25], %[c26], %[c27]], acc[186:187], v[154:155], v[88:91] \n" + " buffer_load_dwordx4 acc[52:55], %[v_os_b3], s[12:15], 0 offen offset:1024 \n" + _UK_MFMA_ " [%[c24], %[c25], %[c26], %[c27]], acc[188:189], v[156:157], v[88:91] \n" + _UK_MFMA_ " [%[c24], %[c25], %[c26], %[c27]], acc[190:191], v[158:159], v[88:91] \n" + _UK_MFMA_ " [%[c28], %[c29], %[c30], %[c31]], acc[176:177], v[208:209], v[92:95] \n" + _UK_MFMA_ " [%[c28], %[c29], %[c30], %[c31]], acc[178:179], v[210:211], v[92:95] \n" + " buffer_load_dwordx4 acc[56:59], %[v_os_b3], s[12:15], 0 offen offset:2048 \n" + _UK_MFMA_ " [%[c28], %[c29], %[c30], %[c31]], acc[180:181], v[212:213], v[92:95] \n" + _UK_MFMA_ " [%[c28], %[c29], %[c30], %[c31]], acc[182:183], v[214:215], v[92:95] \n" + _UK_MFMA_ " [%[c28], %[c29], %[c30], %[c31]], acc[184:185], v[216:217], v[92:95] \n" + _UK_MFMA_ " [%[c28], %[c29], %[c30], %[c31]], acc[186:187], v[218:219], v[92:95] \n" + " buffer_load_dwordx4 acc[60:63], %[v_os_b3], s[12:15], 0 offen offset:3072 \n" + _UK_MFMA_ " [%[c28], %[c29], %[c30], %[c31]], acc[188:189], v[220:221], v[92:95] \n" + _UK_MFMA_ " [%[c28], %[c29], %[c30], %[c31]], acc[190:191], v[222:223], v[92:95] \n" + " s_mov_b64 exec, %[s_execflag_3] \n" +_UK_ATOMIC_ADD_ " %[v_os_o3], v13, s[8:9] \n" +" s_mov_b64 exec, s[38:39] \n" +" s_waitcnt vmcnt(30) \n" + _UK_MFMA_ " [%[c16], %[c17], %[c18], %[c19]], acc[192:193], v[160:161], v[80:83] \n" + _UK_MFMA_ " [%[c16], %[c17], %[c18], %[c19]], acc[194:195], v[162:163], v[80:83] \n" + " buffer_load_dwordx4 acc[64:67], %[v_os_b4], s[12:15], 0 offen \n" + _UK_MFMA_ " [%[c16], %[c17], %[c18], %[c19]], acc[196:197], v[164:165], v[80:83] \n" + _UK_MFMA_ " [%[c16], %[c17], %[c18], %[c19]], acc[198:199], v[166:167], v[80:83] \n" + _UK_MFMA_ " [%[c16], %[c17], %[c18], %[c19]], acc[200:201], v[168:169], v[80:83] \n" + _UK_MFMA_ " [%[c16], %[c17], %[c18], %[c19]], acc[202:203], v[170:171], v[80:83] \n" + " buffer_load_dwordx4 acc[68:71], %[v_os_b4], s[12:15], 0 offen offset:1024 \n" + _UK_MFMA_ " [%[c16], %[c17], %[c18], %[c19]], acc[204:205], v[172:173], v[80:83] \n" + _UK_MFMA_ " [%[c16], %[c17], %[c18], %[c19]], acc[206:207], v[174:175], v[80:83] \n" + _UK_MFMA_ " [%[c20], %[c21], %[c22], %[c23]], acc[192:193], v[224:225], v[84:87] \n" + _UK_MFMA_ " [%[c20], %[c21], %[c22], %[c23]], acc[194:195], v[226:227], v[84:87] \n" + " buffer_load_dwordx4 acc[72:75], %[v_os_b4], s[12:15], 0 offen offset:2048 \n" + _UK_MFMA_ " [%[c20], %[c21], %[c22], %[c23]], acc[196:197], v[228:229], v[84:87] \n" + _UK_MFMA_ " [%[c20], %[c21], %[c22], %[c23]], acc[198:199], v[230:231], v[84:87] \n" + _UK_MFMA_ " [%[c20], %[c21], %[c22], %[c23]], acc[200:201], v[232:233], v[84:87] \n" + _UK_MFMA_ " [%[c20], %[c21], %[c22], %[c23]], acc[202:203], v[234:235], v[84:87] \n" + " buffer_load_dwordx4 acc[76:79], %[v_os_b4], s[12:15], 0 offen offset:3072 \n" + _UK_MFMA_ " [%[c20], %[c21], %[c22], %[c23]], acc[204:205], v[236:237], v[84:87] \n" + _UK_MFMA_ " [%[c20], %[c21], %[c22], %[c23]], acc[206:207], v[238:239], v[84:87] \n" + " s_mov_b64 exec, %[s_execflag_4] \n" +_UK_ATOMIC_ADD_ " %[v_os_o4], v14, s[8:9] \n" +" s_mov_b64 exec, s[38:39] \n" + _UK_MFMA_ " [%[c24], %[c25], %[c26], %[c27]], acc[208:209], v[160:161], v[88:91] \n" + _UK_MFMA_ " [%[c24], %[c25], %[c26], %[c27]], acc[210:211], v[162:163], v[88:91] \n" + " buffer_load_dwordx4 acc[80:83], %[v_os_b5], s[12:15], 0 offen \n" + _UK_MFMA_ " [%[c24], %[c25], %[c26], %[c27]], acc[212:213], v[164:165], v[88:91] \n" + _UK_MFMA_ " [%[c24], %[c25], %[c26], %[c27]], acc[214:215], v[166:167], v[88:91] \n" + _UK_MFMA_ " [%[c24], %[c25], %[c26], %[c27]], acc[216:217], v[168:169], v[88:91] \n" + _UK_MFMA_ " [%[c24], %[c25], %[c26], %[c27]], acc[218:219], v[170:171], v[88:91] \n" + " buffer_load_dwordx4 acc[84:87], %[v_os_b5], s[12:15], 0 offen offset:1024 \n" + _UK_MFMA_ " [%[c24], %[c25], %[c26], %[c27]], acc[220:221], v[172:173], v[88:91] \n" + _UK_MFMA_ " [%[c24], %[c25], %[c26], %[c27]], acc[222:223], v[174:175], v[88:91] \n" + _UK_MFMA_ " [%[c28], %[c29], %[c30], %[c31]], acc[208:209], v[224:225], v[92:95] \n" + _UK_MFMA_ " [%[c28], %[c29], %[c30], %[c31]], acc[210:211], v[226:227], v[92:95] \n" + " buffer_load_dwordx4 acc[88:91], %[v_os_b5], s[12:15], 0 offen offset:2048 \n" + _UK_MFMA_ " [%[c28], %[c29], %[c30], %[c31]], acc[212:213], v[228:229], v[92:95] \n" + _UK_MFMA_ " [%[c28], %[c29], %[c30], %[c31]], acc[214:215], v[230:231], v[92:95] \n" + _UK_MFMA_ " [%[c28], %[c29], %[c30], %[c31]], acc[216:217], v[232:233], v[92:95] \n" + _UK_MFMA_ " [%[c28], %[c29], %[c30], %[c31]], acc[218:219], v[234:235], v[92:95] \n" + " buffer_load_dwordx4 acc[92:95], %[v_os_b5], s[12:15], 0 offen offset:3072 \n" + _UK_MFMA_ " [%[c28], %[c29], %[c30], %[c31]], acc[220:221], v[236:237], v[92:95] \n" + _UK_MFMA_ " [%[c28], %[c29], %[c30], %[c31]], acc[222:223], v[238:239], v[92:95] \n" + " s_mov_b64 exec, %[s_execflag_5] \n" +_UK_ATOMIC_ADD_ " %[v_os_o5], v15, s[8:9] \n" +" s_mov_b64 exec, s[38:39] \n" +" s_waitcnt vmcnt(30) \n" + _UK_MFMA_ " [%[c16], %[c17], %[c18], %[c19]], acc[224:225], v[176:177], v[80:83] \n" + _UK_MFMA_ " [%[c16], %[c17], %[c18], %[c19]], acc[226:227], v[178:179], v[80:83] \n" + " buffer_load_dwordx4 acc[96:99], %[v_os_b6], s[12:15], 0 offen \n" + _UK_MFMA_ " [%[c16], %[c17], %[c18], %[c19]], acc[228:229], v[180:181], v[80:83] \n" + _UK_MFMA_ " [%[c16], %[c17], %[c18], %[c19]], acc[230:231], v[182:183], v[80:83] \n" + _UK_MFMA_ " [%[c16], %[c17], %[c18], %[c19]], acc[232:233], v[184:185], v[80:83] \n" + _UK_MFMA_ " [%[c16], %[c17], %[c18], %[c19]], acc[234:235], v[186:187], v[80:83] \n" + " buffer_load_dwordx4 acc[100:103], %[v_os_b6], s[12:15], 0 offen offset:1024 \n" + _UK_MFMA_ " [%[c16], %[c17], %[c18], %[c19]], acc[236:237], v[188:189], v[80:83] \n" + _UK_MFMA_ " [%[c16], %[c17], %[c18], %[c19]], acc[238:239], v[190:191], v[80:83] \n" + _UK_MFMA_ " [%[c20], %[c21], %[c22], %[c23]], acc[224:225], v[240:241], v[84:87] \n" + _UK_MFMA_ " [%[c20], %[c21], %[c22], %[c23]], acc[226:227], v[242:243], v[84:87] \n" + " buffer_load_dwordx4 acc[104:107], %[v_os_b6], s[12:15], 0 offen offset:2048 \n" + _UK_MFMA_ " [%[c20], %[c21], %[c22], %[c23]], acc[228:229], v[244:245], v[84:87] \n" + _UK_MFMA_ " [%[c20], %[c21], %[c22], %[c23]], acc[230:231], v[246:247], v[84:87] \n" + _UK_MFMA_ " [%[c20], %[c21], %[c22], %[c23]], acc[232:233], v[248:249], v[84:87] \n" + _UK_MFMA_ " [%[c20], %[c21], %[c22], %[c23]], acc[234:235], v[250:251], v[84:87] \n" + " buffer_load_dwordx4 acc[108:111], %[v_os_b6], s[12:15], 0 offen offset:3072 \n" + _UK_MFMA_ " [%[c20], %[c21], %[c22], %[c23]], acc[236:237], v[252:253], v[84:87] \n" + _UK_MFMA_ " [%[c20], %[c21], %[c22], %[c23]], acc[238:239], v[254:255], v[84:87] \n" + " s_mov_b64 exec, %[s_execflag_6] \n" +_UK_ATOMIC_ADD_ " %[v_os_o6], v16, s[8:9] \n" +" s_mov_b64 exec, s[38:39] \n" + _UK_MFMA_ " [%[c24], %[c25], %[c26], %[c27]], acc[240:241], v[176:177], v[88:91] \n" + _UK_MFMA_ " [%[c24], %[c25], %[c26], %[c27]], acc[242:243], v[178:179], v[88:91] \n" + " buffer_load_dwordx4 acc[112:115], %[v_os_b7], s[12:15], 0 offen \n" + _UK_MFMA_ " [%[c24], %[c25], %[c26], %[c27]], acc[244:245], v[180:181], v[88:91] \n" + _UK_MFMA_ " [%[c24], %[c25], %[c26], %[c27]], acc[246:247], v[182:183], v[88:91] \n" + _UK_MFMA_ " [%[c24], %[c25], %[c26], %[c27]], acc[248:249], v[184:185], v[88:91] \n" + _UK_MFMA_ " [%[c24], %[c25], %[c26], %[c27]], acc[250:251], v[186:187], v[88:91] \n" + " buffer_load_dwordx4 acc[116:119], %[v_os_b7], s[12:15], 0 offen offset:1024 \n" + _UK_MFMA_ " [%[c24], %[c25], %[c26], %[c27]], acc[252:253], v[188:189], v[88:91] \n" + _UK_MFMA_ " [%[c24], %[c25], %[c26], %[c27]], acc[254:255], v[190:191], v[88:91] \n" + _UK_MFMA_ " [%[c28], %[c29], %[c30], %[c31]], acc[240:241], v[240:241], v[92:95] \n" + _UK_MFMA_ " [%[c28], %[c29], %[c30], %[c31]], acc[242:243], v[242:243], v[92:95] \n" + " buffer_load_dwordx4 acc[120:123], %[v_os_b7], s[12:15], 0 offen offset:2048 \n" + _UK_MFMA_ " [%[c28], %[c29], %[c30], %[c31]], acc[244:245], v[244:245], v[92:95] \n" + _UK_MFMA_ " [%[c28], %[c29], %[c30], %[c31]], acc[246:247], v[246:247], v[92:95] \n" + _UK_MFMA_ " [%[c28], %[c29], %[c30], %[c31]], acc[248:249], v[248:249], v[92:95] \n" + _UK_MFMA_ " [%[c28], %[c29], %[c30], %[c31]], acc[250:251], v[250:251], v[92:95] \n" + " buffer_load_dwordx4 acc[124:127], %[v_os_b7], s[12:15], 0 offen offset:3072 \n" + _UK_MFMA_ " [%[c28], %[c29], %[c30], %[c31]], acc[252:253], v[252:253], v[92:95] \n" + _UK_MFMA_ " [%[c28], %[c29], %[c30], %[c31]], acc[254:255], v[254:255], v[92:95] \n" + " s_mov_b64 exec, %[s_execflag_7] \n" +_UK_ATOMIC_ADD_ " %[v_os_o7], v17, s[8:9] \n" +" s_mov_b64 exec, s[38:39] \n" +" s_add_u32 s60, 0x00000100, s80 \n" +" s_cmp_lt_u32 s60, %[s_loop_cnt] \n" +" s_cselect_b32 s56, s56, 0 \n" +" s_add_u32 s12, s56, s12 \n" +" s_addc_u32 s13, 0, s13 \n" +" s_cmp_ge_u32 s80, 0x00000100 \n" +" s_cselect_b32 s59, 0x00000100, s59 \n" +" s_add_u32 s8, s59, s8 \n" +" s_addc_u32 s9, 0, s9 \n" +" v_mul_f32 %[c16], %[scale_0], %[c16] \n" +" v_mul_f32 %[c17], %[scale_0], %[c17] \n" +" v_mul_f32 %[c18], %[scale_0], %[c18] \n" +" v_mul_f32 %[c19], %[scale_0], %[c19] \n" +" v_mul_f32 %[c20], %[scale_1], %[c20] \n" +" v_mul_f32 %[c21], %[scale_1], %[c21] \n" +" v_mul_f32 %[c22], %[scale_1], %[c22] \n" +" v_mul_f32 %[c23], %[scale_1], %[c23] \n" +" v_mul_f32 %[c24], %[scale_0], %[c24] \n" +" v_mul_f32 %[c25], %[scale_0], %[c25] \n" +" v_mul_f32 %[c26], %[scale_0], %[c26] \n" +" v_mul_f32 %[c27], %[scale_0], %[c27] \n" +" v_mul_f32 %[c28], %[scale_1], %[c28] \n" +" v_mul_f32 %[c29], %[scale_1], %[c29] \n" +" v_mul_f32 %[c30], %[scale_1], %[c30] \n" +" v_mul_f32 %[c31], %[scale_1], %[c31] \n" +_UK_PK_CVT_("%[c16]", "%[c17]", "%[c16]") +_UK_PK_CVT_("%[c18]", "%[c19]", "%[c17]") +_UK_PK_CVT_("%[c20]", "%[c21]", "%[c18]") +_UK_PK_CVT_("%[c22]", "%[c23]", "%[c19]") +_UK_PK_CVT_("%[c24]", "%[c25]", "%[c20]") +_UK_PK_CVT_("%[c26]", "%[c27]", "%[c21]") +_UK_PK_CVT_("%[c28]", "%[c29]", "%[c22]") +_UK_PK_CVT_("%[c30]", "%[c31]", "%[c23]") +" s_addk_i32 s80, 0x0080 \n" +" s_cmp_lt_i32 s80, %[s_loop_cnt] \n" +" s_cbranch_scc0 label_0EC1 \n" +" s_branch label_0AA6 \n" +" label_0EC1: \n" +" s_waitcnt lgkmcnt(0) \n" +" s_barrier \n" +" ds_read_b32 v10, %[v_sfl_sld] offset:16640 \n" +" ds_read_b32 v11, %[v_sfl_sld] offset:16672 \n" +" ds_read_b32 v12, %[v_sfl_sld] offset:16704 \n" +" ds_read_b32 v13, %[v_sfl_sld] offset:16736 \n" +" ds_read_b32 v14, %[v_sfl_sld] offset:20992 \n" +" ds_read_b32 v15, %[v_sfl_sld] offset:21024 \n" +" ds_read_b32 v16, %[v_sfl_sld] offset:21056 \n" +" ds_read_b32 v17, %[v_sfl_sld] offset:21088 \n" +" s_waitcnt lgkmcnt(0) \n" + " s_mov_b64 exec, %[s_execflag_0] \n" +_UK_ATOMIC_ADD_ " %[v_os_o0], v10, s[8:9] \n" + " s_mov_b64 exec, %[s_execflag_1] \n" +_UK_ATOMIC_ADD_ " %[v_os_o1], v11, s[8:9] \n" + " s_mov_b64 exec, %[s_execflag_2] \n" +_UK_ATOMIC_ADD_ " %[v_os_o2], v12, s[8:9] \n" + " s_mov_b64 exec, %[s_execflag_3] \n" +_UK_ATOMIC_ADD_ " %[v_os_o3], v13, s[8:9] \n" + " s_mov_b64 exec, %[s_execflag_4] \n" +_UK_ATOMIC_ADD_ " %[v_os_o4], v14, s[8:9] \n" + " s_mov_b64 exec, %[s_execflag_5] \n" +_UK_ATOMIC_ADD_ " %[v_os_o5], v15, s[8:9] \n" + " s_mov_b64 exec, %[s_execflag_6] \n" +_UK_ATOMIC_ADD_ " %[v_os_o6], v16, s[8:9] \n" + " s_mov_b64 exec, %[s_execflag_7] \n" +_UK_ATOMIC_ADD_ " %[v_os_o7], v17, s[8:9] \n" +" s_mov_b64 exec, s[38:39] \n" +" s_add_u32 s8, s59, s8 \n" +" s_addc_u32 s9, 0, s9 \n" +" ds_write_b64 %[v_sfl_sst], [%[c16],%[c17]] offset:25344 \n" +" ds_write_b64 %[v_sfl_sst], [%[c18],%[c19]] offset:29696 \n" +" ds_write_b64 %[v_sfl_sst], [%[c20],%[c21]] offset:27520 \n" +" ds_write_b64 %[v_sfl_sst], [%[c22],%[c23]] offset:31872 \n" +" s_waitcnt lgkmcnt(0) \n" +" s_barrier \n" +" ds_read_b32 v10, %[v_sfl_sld] offset:25344 \n" +" ds_read_b32 v11, %[v_sfl_sld] offset:25376 \n" +" ds_read_b32 v12, %[v_sfl_sld] offset:25408 \n" +" ds_read_b32 v13, %[v_sfl_sld] offset:25440 \n" +" ds_read_b32 v14, %[v_sfl_sld] offset:29696 \n" +" ds_read_b32 v15, %[v_sfl_sld] offset:29728 \n" +" ds_read_b32 v16, %[v_sfl_sld] offset:29760 \n" +" ds_read_b32 v17, %[v_sfl_sld] offset:29792 \n" +" s_waitcnt lgkmcnt(0) \n" +" s_mov_b64 exec, %[s_execflag_0] \n" +_UK_ATOMIC_ADD_ " %[v_os_o0], v10, s[8:9] \n" + " s_mov_b64 exec, %[s_execflag_1] \n" +_UK_ATOMIC_ADD_ " %[v_os_o1], v11, s[8:9] \n" + " s_mov_b64 exec, %[s_execflag_2] \n" +_UK_ATOMIC_ADD_ " %[v_os_o2], v12, s[8:9] \n" + " s_mov_b64 exec, %[s_execflag_3] \n" +_UK_ATOMIC_ADD_ " %[v_os_o3], v13, s[8:9] \n" + " s_mov_b64 exec, %[s_execflag_4] \n" +_UK_ATOMIC_ADD_ " %[v_os_o4], v14, s[8:9] \n" + " s_mov_b64 exec, %[s_execflag_5] \n" +_UK_ATOMIC_ADD_ " %[v_os_o5], v15, s[8:9] \n" + " s_mov_b64 exec, %[s_execflag_6] \n" +_UK_ATOMIC_ADD_ " %[v_os_o6], v16, s[8:9] \n" + " s_mov_b64 exec, %[s_execflag_7] \n" +_UK_ATOMIC_ADD_ " %[v_os_o7], v17, s[8:9] \n" +" s_mov_b64 exec, s[38:39] \n" + +#undef _UK_MFMA_ +#undef _UK_PK_CVT_ +#undef _UK_ATOMIC_ADD_ + diff --git a/include/ck_tile/ops/fused_moe/pipeline/fused_moegemm_pipeline_flatmm_policy.hpp b/include/ck_tile/ops/fused_moe/pipeline/fused_moegemm_pipeline_flatmm_policy.hpp index fea30f0297..629f0ee8f1 100644 --- a/include/ck_tile/ops/fused_moe/pipeline/fused_moegemm_pipeline_flatmm_policy.hpp +++ b/include/ck_tile/ops/fused_moe/pipeline/fused_moegemm_pipeline_flatmm_policy.hpp @@ -810,21 +810,46 @@ struct FusedMoeGemmPipelineFlatmmPolicy CK_TILE_HOST_DEVICE static constexpr auto GetUK_1() { using S_ = typename Problem::BlockShape; + using T_ = typename Problem::Traits; if constexpr(std::is_same_v && std::is_same_v && std::is_same_v && S_::Block_M1 == 32 && S_::Block_N1 == 128 && S_::Block_K1 == 512 && - S_::Warp_M0 == 16 && S_::Warp_N0 == 16 && S_::Warp_K0 == 32) + S_::Warp_M0 == 16 && S_::Warp_N0 == 16 && S_::Warp_K0 == 32 && + T_::PipeInterleave == false) { return FlatmmSn_32x128x512_1x4x1_16x16x32_BF16{}; + // return FlatmmSn_32x128x512_1x4x1_16x16x32_BF16_itl{}; } else if constexpr(std::is_same_v && std::is_same_v && std::is_same_v && S_::Block_M1 == 32 && S_::Block_N1 == 128 && S_::Block_K1 == 512 && - S_::Warp_M0 == 16 && S_::Warp_N0 == 16 && S_::Warp_K0 == 32) + S_::Warp_M0 == 16 && S_::Warp_N0 == 16 && S_::Warp_K0 == 32 && + T_::PipeInterleave == false) { return FlatmmSn_32x128x512_1x4x1_16x16x32_FP16{}; + // return FlatmmSn_32x128x512_1x4x1_16x16x32_FP16_itl{}; + } + else if constexpr(std::is_same_v && + std::is_same_v && + std::is_same_v && + S_::Block_M1 == 32 && S_::Block_N1 == 128 && S_::Block_K1 == 512 && + S_::Warp_M0 == 16 && S_::Warp_N0 == 16 && S_::Warp_K0 == 32 && + T_::PipeInterleave == true) + { + // return FlatmmSn_32x128x512_1x4x1_16x16x32_FP16{}; + return FlatmmSn_32x128x512_1x4x1_16x16x32_BF16_itl{}; + } + else if constexpr(std::is_same_v && + std::is_same_v && + std::is_same_v && + S_::Block_M1 == 32 && S_::Block_N1 == 128 && S_::Block_K1 == 512 && + S_::Warp_M0 == 16 && S_::Warp_N0 == 16 && S_::Warp_K0 == 32 && + T_::PipeInterleave == true) + { + // return FlatmmSn_32x128x512_1x4x1_16x16x32_FP16{}; + return FlatmmSn_32x128x512_1x4x1_16x16x32_FP16_itl{}; } } }; diff --git a/include/ck_tile/ops/fused_moe/pipeline/fused_moegemm_traits.hpp b/include/ck_tile/ops/fused_moe/pipeline/fused_moegemm_traits.hpp index d7127b098c..3fb82bc099 100644 --- a/include/ck_tile/ops/fused_moe/pipeline/fused_moegemm_traits.hpp +++ b/include/ck_tile/ops/fused_moe/pipeline/fused_moegemm_traits.hpp @@ -22,7 +22,8 @@ template + bool PadIntermediateSize_ = false, + bool PipeInterleave_ = true> struct FusedMoeGemmTraits { // Gate+Up or Gate only @@ -32,6 +33,7 @@ struct FusedMoeGemmTraits static constexpr FusedMoeGemmWeightPermuteEnum PermuteEnum = PermuteEnum_; static constexpr bool PadHiddenSize = PadHiddenSize_; static constexpr bool PadIntermediateSize = PadIntermediateSize_; + static constexpr bool PipeInterleave = PipeInterleave_; }; // Note: this need to be a bit mask From fdfe2102304f62ec62194706a5f67766ae824dc6 Mon Sep 17 00:00:00 2001 From: Illia Silin <98187287+illsilin@users.noreply.github.com> Date: Sun, 15 Dec 2024 16:25:21 -0800 Subject: [PATCH 05/25] upgrade sqlalchemy version (#1748) * upgrade sqlalchemy version * replace the connection with engine in to_sql call * change the hipTes=nsor ctest syntax --- Dockerfile | 2 +- Jenkinsfile | 4 +--- script/process_perf_data.py | 2 +- 3 files changed, 3 insertions(+), 5 deletions(-) diff --git a/Dockerfile b/Dockerfile index 83edbfb8ee..a3bf3866bf 100644 --- a/Dockerfile +++ b/Dockerfile @@ -94,7 +94,7 @@ RUN pip install --upgrade cmake==3.27.5 && \ dpkg -i dumb-init_*.deb && rm dumb-init_*.deb && \ # Install packages for processing the performance results pip3 install --upgrade pip && \ - pip3 install sqlalchemy==1.4.46 pymysql pandas==2.2.3 setuptools-rust sshtunnel==0.4.0 && \ + pip3 install sqlalchemy==2.0.36 pymysql pandas==2.2.3 setuptools-rust sshtunnel==0.4.0 && \ # Add render group groupadd -f render && \ # Install the new rocm-cmake version diff --git a/Jenkinsfile b/Jenkinsfile index f82c34afa6..87c9457fcb 100644 --- a/Jenkinsfile +++ b/Jenkinsfile @@ -566,11 +566,9 @@ def Build_CK(Map conf=[:]){ ls -ltr CC=hipcc CXX=hipcc cmake -Bbuild . -D CMAKE_PREFIX_PATH="${env.WORKSPACE}/install" cmake --build build -- -j + ctest --test-dir build """ } - dir("hipTensor-${params.hipTensor_branch}/build"){ - sh 'ctest' - } } } } diff --git a/script/process_perf_data.py b/script/process_perf_data.py index fbfec94eef..32e2e15d7a 100644 --- a/script/process_perf_data.py +++ b/script/process_perf_data.py @@ -332,7 +332,7 @@ def main(): table_name="ck_fmha_bwd_tflops" tflops_base = get_baseline(table_name,conn) - store_new_test_result(table_name, results, testlist, branch_name, node_id, gpu_arch, compute_units, rocm_vers, hip_vers, environment, conn) + store_new_test_result(table_name, results, testlist, branch_name, node_id, gpu_arch, compute_units, rocm_vers, hip_vers, environment, sqlEngine) conn.close() #compare the results to the baseline if baseline exists From a8ad7fcce912c8e462ca69d5ca680d99b2ef56dd Mon Sep 17 00:00:00 2001 From: Max Podkorytov <4273004+tenpercent@users.noreply.github.com> Date: Tue, 10 Dec 2024 18:14:52 +0000 Subject: [PATCH 06/25] add template placeholders --- .github/CONTRIBUTING.md | 0 .github/ISSUE_TEMPLATE.md | 14 ++++++++++++++ .github/PULL_REQUEST_TEMPLATE.md | 0 3 files changed, 14 insertions(+) create mode 100644 .github/CONTRIBUTING.md create mode 100644 .github/ISSUE_TEMPLATE.md create mode 100644 .github/PULL_REQUEST_TEMPLATE.md diff --git a/.github/CONTRIBUTING.md b/.github/CONTRIBUTING.md new file mode 100644 index 0000000000..e69de29bb2 diff --git a/.github/ISSUE_TEMPLATE.md b/.github/ISSUE_TEMPLATE.md new file mode 100644 index 0000000000..263cc3480d --- /dev/null +++ b/.github/ISSUE_TEMPLATE.md @@ -0,0 +1,14 @@ +When creating an issue, please check if a similar issue already exists. + +### When reporting a bug, please include: +- [ ] A descriptive title +- [ ] An isolated way to reproduce the behavior (preferably a docker container with a repro) +- [ ] ROCm version, clang version, Composable Kernel commit pin +- [ ] Environment variables +- [ ] The behavior you expect to see, and the behavior you actually see + +### When requesting a feature, please include: +- [ ] A descriptive title +- [ ] A detailed description of the problem you are trying to solve +- [ ] An overview of the suggested solution +- [ ] Explanation why the solution is an improvement \ No newline at end of file diff --git a/.github/PULL_REQUEST_TEMPLATE.md b/.github/PULL_REQUEST_TEMPLATE.md new file mode 100644 index 0000000000..e69de29bb2 From 30a37cac0e76298ef184597b1f7d3ef0d3f4bb60 Mon Sep 17 00:00:00 2001 From: Max Podkorytov <4273004+tenpercent@users.noreply.github.com> Date: Tue, 10 Dec 2024 18:50:27 +0000 Subject: [PATCH 07/25] add pull request template placeholder --- .github/PULL_REQUEST_TEMPLATE.md | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/.github/PULL_REQUEST_TEMPLATE.md b/.github/PULL_REQUEST_TEMPLATE.md index e69de29bb2..c5161f7f8f 100644 --- a/.github/PULL_REQUEST_TEMPLATE.md +++ b/.github/PULL_REQUEST_TEMPLATE.md @@ -0,0 +1,19 @@ +## Proposed changes + +Please describe the motivation behind the pull request, whether it enables a new feature or fixes a bug. If there are associated pull requests or issues, please link them to the pull request. + +## Checklist + +Please put an `x` into the boxes that apply. You can also fill these out after creating the PR. If you're not sure, please don't hesitate to ask. + +- [ ] I have added tests relevant to the introduced functionality, and the unit tests are passing locally +- [ ] I have added inline documentation which enables the maintainers with understanding the motivation +- [ ] I have removed the stale documentation which is no longer relevant after this pull request +- [ ] I have added release notes which provide the end users with a brief summary of the improvement from this pull request +- [ ] I have run `clang-format` on all changed files +- [ ] Any dependent changes have been merged + +## Discussion + +If this is a relatively large or complex change, feel free to start a discussion by explaining why you chose the solution you did and what alternatives you considered + From 1b75c77da41afdfa8cff30a40bbe0fc4bd1d643f Mon Sep 17 00:00:00 2001 From: Max Podkorytov <4273004+tenpercent@users.noreply.github.com> Date: Tue, 10 Dec 2024 19:14:37 +0000 Subject: [PATCH 08/25] add contributing placeholder --- .github/CONTRIBUTING.md | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/.github/CONTRIBUTING.md b/.github/CONTRIBUTING.md index e69de29bb2..56f2acee71 100644 --- a/.github/CONTRIBUTING.md +++ b/.github/CONTRIBUTING.md @@ -0,0 +1,10 @@ +We'd love for you to contribute to our source code! + +Some helpful links: + +- [Code of Conduct guidelines](https://www.contributor-covenant.org/version/2/1/code_of_conduct/code_of_conduct.txt) +- [New issue guidelines](https://github.com/rocm/composable_kernel/blob/develop/.github/ISSUE_TEMPLATE.md) +- [Submitting a pull request guidelines](https://github.com/rocm/composable_kernel/blob/develop/.github/PULL_REQUEST_TEMPLATE.md) +- [Maintainers](https://github.com/rocm/composable_kernel/blob/develop/CONTRIBUTORS.md) +- [General information](https://github.com/rocm/composable_kernel/blob/develop/README.md) +- [ROCm documentation](https://rocm.docs.amd.com/en/latest/how-to/llm-fine-tuning-optimization/optimizing-with-composable-kernel.html) \ No newline at end of file From 0fd6978d2a3c5973d9c0486616b2a71ea7aa5f86 Mon Sep 17 00:00:00 2001 From: Max Podkorytov <4273004+tenpercent@users.noreply.github.com> Date: Tue, 10 Dec 2024 20:29:49 +0000 Subject: [PATCH 09/25] clarify release notes bullet point --- .github/PULL_REQUEST_TEMPLATE.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/PULL_REQUEST_TEMPLATE.md b/.github/PULL_REQUEST_TEMPLATE.md index c5161f7f8f..b3fcabec34 100644 --- a/.github/PULL_REQUEST_TEMPLATE.md +++ b/.github/PULL_REQUEST_TEMPLATE.md @@ -9,7 +9,7 @@ Please put an `x` into the boxes that apply. You can also fill these out after c - [ ] I have added tests relevant to the introduced functionality, and the unit tests are passing locally - [ ] I have added inline documentation which enables the maintainers with understanding the motivation - [ ] I have removed the stale documentation which is no longer relevant after this pull request -- [ ] I have added release notes which provide the end users with a brief summary of the improvement from this pull request +- [ ] (If this change is user-facing) I have added release notes which provide the end users with a brief summary of the improvement from this pull request - [ ] I have run `clang-format` on all changed files - [ ] Any dependent changes have been merged From d46196f291a33539a089d7d09bcbc4d2270733c2 Mon Sep 17 00:00:00 2001 From: Adam Osewski <19374865+aosewski@users.noreply.github.com> Date: Tue, 17 Dec 2024 09:19:44 +0100 Subject: [PATCH 10/25] Enhance printing functionality (#1751) * Added object print with all template parameters * fix clang format --------- Co-authored-by: ravil-mobile Co-authored-by: illsilin --- .../gpu/device/device_base.hpp | 34 + .../impl/device_gemm_xdl_cshuffle_v3.hpp | 1 + ...m_sn_uk_gfx9_32x128x512_1x4x1_16x16x16.inc | 1383 +++++++++------- ..._uk_gfx9_32x128x512_1x4x1_16x16x16_itl.inc | 1439 +++++++++-------- ...atmm_uk_gfx9_32x512x128_1x1x1_16x16x16.inc | 1007 ++++++------ .../profiler/profile_gemm_universal_impl.hpp | 18 +- 6 files changed, 2095 insertions(+), 1787 deletions(-) diff --git a/include/ck/tensor_operation/gpu/device/device_base.hpp b/include/ck/tensor_operation/gpu/device/device_base.hpp index 908ada016d..736e241fdf 100644 --- a/include/ck/tensor_operation/gpu/device/device_base.hpp +++ b/include/ck/tensor_operation/gpu/device/device_base.hpp @@ -5,6 +5,8 @@ #include #include +#include +#include #include "ck/stream_config.hpp" @@ -12,6 +14,34 @@ namespace ck { namespace tensor_operation { namespace device { +#define GET_OBJECT_NAME_IMLP \ + std::optional GetObjectName() const override \ + { \ + std::string str = __PRETTY_FUNCTION__; \ + static std::regex obj_name_expr{" (.*)::GetObjectName"}; \ + std::smatch match; \ + if(!std::regex_search(str, match, obj_name_expr)) \ + { \ + return str; \ + } \ + return std::string(match[1]) + ';'; \ + } + +#define GET_TEMPLATE_INFO_IMPL \ + std::optional GetTemplateInfo() const override \ + { \ + std::string str = __PRETTY_FUNCTION__; \ + static std::regex template_expr{"\\[(.*)\\]"}; \ + std::smatch match; \ + if(!std::regex_search(str, match, template_expr)) \ + { \ + return std::nullopt; \ + } \ + return std::string(match[1]); \ + } + +#define REGISTER_EXTRA_PRINTING_METHODS GET_OBJECT_NAME_IMLP GET_TEMPLATE_INFO_IMPL + struct BaseArgument { BaseArgument() = default; @@ -48,6 +78,10 @@ struct BaseOperator virtual std::string GetTypeIdName() const { return typeid(*this).name(); } + virtual std::optional GetObjectName() const { return std::nullopt; } + + virtual std::optional GetTemplateInfo() const { return std::nullopt; } + virtual std::string GetTypeIdHashCode() const { std::ostringstream oss; diff --git a/include/ck/tensor_operation/gpu/device/impl/device_gemm_xdl_cshuffle_v3.hpp b/include/ck/tensor_operation/gpu/device/impl/device_gemm_xdl_cshuffle_v3.hpp index 4489b2e5ce..ad6aa1e7c3 100644 --- a/include/ck/tensor_operation/gpu/device/impl/device_gemm_xdl_cshuffle_v3.hpp +++ b/include/ck/tensor_operation/gpu/device/impl/device_gemm_xdl_cshuffle_v3.hpp @@ -729,6 +729,7 @@ struct DeviceGemm_Xdl_CShuffleV3 : public DeviceGemmV2 best_op_object_name; float best_ave_time = 0; float best_tflops = 0; float best_gb_per_sec = 0; @@ -225,7 +226,8 @@ bool profile_gemm_universal_impl(int do_verification, } } - std::string op_name = op_ptr->GetTypeString(); + std::string op_name = op_ptr->GetTypeString(); + std::optional op_obj_name = op_ptr->GetObjectName(); float ave_time = invoker_ptr->Run(argument_ptr.get(), StreamConfig{nullptr, @@ -251,11 +253,12 @@ bool profile_gemm_universal_impl(int do_verification, if(tflops > best_tflops && ave_time > 1e-10) { - best_op_name = op_name; - best_tflops = tflops; - best_ave_time = ave_time; - best_gb_per_sec = gb_per_sec; - best_kbatch = kbatch_curr; + best_op_name = op_name; + best_op_object_name = op_obj_name; + best_tflops = tflops; + best_ave_time = ave_time; + best_gb_per_sec = gb_per_sec; + best_kbatch = kbatch_curr; } } else @@ -306,6 +309,9 @@ bool profile_gemm_universal_impl(int do_verification, << " : " << best_ave_time << " ms, " << best_tflops << " TFlops, " << best_gb_per_sec << " GB/s, " << best_op_name << std::endl; + if(best_op_object_name) + std::cout << best_op_object_name.value() << std::endl; + return pass; } From 627a27bda3f38b3d904f844ec0b4d988e50cc262 Mon Sep 17 00:00:00 2001 From: jakpiase Date: Tue, 17 Dec 2024 14:25:22 +0100 Subject: [PATCH 11/25] Added unit tests for CK Tile compute bound gemm pipeline (#1728) --- test/ck_tile/gemm/CMakeLists.txt | 2 +- test/ck_tile/gemm/test_gemm_mem_pipeline.cpp | 36 ----------- test/ck_tile/gemm/test_gemm_pipeline.cpp | 42 +++++++++++++ ...es.inc => test_gemm_pipeline_ut_cases.inc} | 10 +-- ...e_util.hpp => test_gemm_pipeline_util.hpp} | 62 +++++++++++++------ 5 files changed, 90 insertions(+), 62 deletions(-) delete mode 100644 test/ck_tile/gemm/test_gemm_mem_pipeline.cpp create mode 100644 test/ck_tile/gemm/test_gemm_pipeline.cpp rename test/ck_tile/gemm/{test_gemm_mem_pipeline_ut_cases.inc => test_gemm_pipeline_ut_cases.inc} (79%) rename test/ck_tile/gemm/{test_gemm_mem_pipeline_util.hpp => test_gemm_pipeline_util.hpp} (80%) diff --git a/test/ck_tile/gemm/CMakeLists.txt b/test/ck_tile/gemm/CMakeLists.txt index f96ad9c6e0..ecfbd4e55b 100644 --- a/test/ck_tile/gemm/CMakeLists.txt +++ b/test/ck_tile/gemm/CMakeLists.txt @@ -1,4 +1,4 @@ # Currently ck_tile is only built on gfx9 if(GPU_TARGETS MATCHES "gfx9") - add_gtest_executable(test_ck_tile_gemm_mem_pipeline test_gemm_mem_pipeline.cpp) + add_gtest_executable(test_ck_tile_gemm_pipeline test_gemm_pipeline.cpp) endif() diff --git a/test/ck_tile/gemm/test_gemm_mem_pipeline.cpp b/test/ck_tile/gemm/test_gemm_mem_pipeline.cpp deleted file mode 100644 index aeb383c87d..0000000000 --- a/test/ck_tile/gemm/test_gemm_mem_pipeline.cpp +++ /dev/null @@ -1,36 +0,0 @@ -// SPDX-License-Identifier: MIT -// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved. - -#include - -#include "gtest/gtest.h" - -#include "ck_tile/host.hpp" -#include "test_gemm_mem_pipeline_util.hpp" - -using F16 = ck_tile::half_t; -using F32 = float; -using Row = ck_tile::tensor_layout::gemm::RowMajor; -using Col = ck_tile::tensor_layout::gemm::ColumnMajor; -using Intrawave = ck_tile::integral_constant; -using Interwave = ck_tile::integral_constant; - -// clang-format off -using KernelTypes = ::testing::Types< - // ALayout, BLayout, CLayout, ADataType, BDataType, AccDataType, CDataType, GemmPipelineScheduler - std::tuple< Row, Row, Row, F16, F16, F32, F16, Intrawave>, - std::tuple< Row, Row, Row, F16, F16, F32, F16, Interwave>, - std::tuple< Row, Col, Row, F16, F16, F32, F16, Intrawave>, - std::tuple< Row, Col, Row, F16, F16, F32, F16, Interwave>, - std::tuple< Col, Row, Row, F16, F16, F32, F16, Intrawave>, - std::tuple< Col, Row, Row, F16, F16, F32, F16, Interwave>, - std::tuple< Col, Col, Row, F16, F16, F32, F16, Intrawave>, - std::tuple< Col, Col, Row, F16, F16, F32, F16, Interwave> - >; -// clang-format on - -TYPED_TEST_SUITE(TestCkTileGemmMemPipeline, KernelTypes); - -#include "test_gemm_mem_pipeline_ut_cases.inc" diff --git a/test/ck_tile/gemm/test_gemm_pipeline.cpp b/test/ck_tile/gemm/test_gemm_pipeline.cpp new file mode 100644 index 0000000000..48a2b86a63 --- /dev/null +++ b/test/ck_tile/gemm/test_gemm_pipeline.cpp @@ -0,0 +1,42 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved. + +#include + +#include "gtest/gtest.h" + +#include "ck_tile/host.hpp" +#include "test_gemm_pipeline_util.hpp" + +using F16 = ck_tile::half_t; +using F32 = float; +using Row = ck_tile::tensor_layout::gemm::RowMajor; +using Col = ck_tile::tensor_layout::gemm::ColumnMajor; +using Intrawave = ck_tile::integral_constant; +using Interwave = ck_tile::integral_constant; +using Mem = ck_tile::integral_constant; +using Comp = ck_tile::integral_constant; + +// clang-format off +using KernelTypes = ::testing::Types< + // ALayout, BLayout, CLayout, ADataType, BDataType, AccDataType, CDataType, GemmPipelineScheduler, PipelineType + std::tuple< Row, Row, Row, F16, F16, F32, F16, Intrawave, Mem>, + std::tuple< Row, Row, Row, F16, F16, F32, F16, Intrawave, Comp>, + std::tuple< Row, Row, Row, F16, F16, F32, F16, Interwave, Mem>, + std::tuple< Row, Col, Row, F16, F16, F32, F16, Intrawave, Mem>, + std::tuple< Row, Col, Row, F16, F16, F32, F16, Intrawave, Comp>, + std::tuple< Row, Col, Row, F16, F16, F32, F16, Interwave, Mem>, + std::tuple< Col, Row, Row, F16, F16, F32, F16, Intrawave, Mem>, + std::tuple< Col, Row, Row, F16, F16, F32, F16, Intrawave, Comp>, + std::tuple< Col, Row, Row, F16, F16, F32, F16, Interwave, Mem>, + std::tuple< Col, Col, Row, F16, F16, F32, F16, Intrawave, Mem>, + std::tuple< Col, Col, Row, F16, F16, F32, F16, Intrawave, Comp>, + std::tuple< Col, Col, Row, F16, F16, F32, F16, Interwave, Mem> + >; +// clang-format on + +TYPED_TEST_SUITE(TestCkTileGemmPipeline, KernelTypes); + +#include "test_gemm_pipeline_ut_cases.inc" diff --git a/test/ck_tile/gemm/test_gemm_mem_pipeline_ut_cases.inc b/test/ck_tile/gemm/test_gemm_pipeline_ut_cases.inc similarity index 79% rename from test/ck_tile/gemm/test_gemm_mem_pipeline_ut_cases.inc rename to test/ck_tile/gemm/test_gemm_pipeline_ut_cases.inc index af94d68f2c..c78d69601c 100644 --- a/test/ck_tile/gemm/test_gemm_mem_pipeline_ut_cases.inc +++ b/test/ck_tile/gemm/test_gemm_pipeline_ut_cases.inc @@ -3,7 +3,7 @@ #pragma once -TYPED_TEST(TestCkTileGemmMemPipeline, SmallM) +TYPED_TEST(TestCkTileGemmPipeline, SmallM) { std::vector Ms{1, 2, 3, 4, 5, 6}; constexpr int N = 1024; @@ -13,7 +13,7 @@ TYPED_TEST(TestCkTileGemmMemPipeline, SmallM) this->Run(M, N, K); } -TYPED_TEST(TestCkTileGemmMemPipeline, MidLargeM) +TYPED_TEST(TestCkTileGemmPipeline, MidLargeM) { std::vector Ms{127, 255, 312, 799, 1573}; constexpr int N = 1024; @@ -23,7 +23,7 @@ TYPED_TEST(TestCkTileGemmMemPipeline, MidLargeM) this->Run(M, N, K); } -TYPED_TEST(TestCkTileGemmMemPipeline, PaddK) +TYPED_TEST(TestCkTileGemmPipeline, PaddK) { std::vector Ms{127}; constexpr int N = 1024; @@ -33,7 +33,7 @@ TYPED_TEST(TestCkTileGemmMemPipeline, PaddK) this->Run(M, N, K); } -TYPED_TEST(TestCkTileGemmMemPipeline, Regular) +TYPED_TEST(TestCkTileGemmPipeline, Regular) { std::vector Ms{512}; constexpr int N = 1024; @@ -43,7 +43,7 @@ TYPED_TEST(TestCkTileGemmMemPipeline, Regular) this->Run(M, N, K); } -TYPED_TEST(TestCkTileGemmMemPipeline, NotSupportedArgument) +TYPED_TEST(TestCkTileGemmPipeline, NotSupportedArgument) { constexpr int M = 512; constexpr int N = 1025; diff --git a/test/ck_tile/gemm/test_gemm_mem_pipeline_util.hpp b/test/ck_tile/gemm/test_gemm_pipeline_util.hpp similarity index 80% rename from test/ck_tile/gemm/test_gemm_mem_pipeline_util.hpp rename to test/ck_tile/gemm/test_gemm_pipeline_util.hpp index 6941a7596a..a514986024 100644 --- a/test/ck_tile/gemm/test_gemm_mem_pipeline_util.hpp +++ b/test/ck_tile/gemm/test_gemm_pipeline_util.hpp @@ -11,18 +11,24 @@ #include "ck_tile/ops/epilogue.hpp" #include "ck_tile/ops/gemm.hpp" +enum struct GemmPipelineType +{ + Mem, + Comp +}; template -class TestCkTileGemmMemPipeline : public ::testing::Test +class TestCkTileGemmPipeline : public ::testing::Test { protected: - using ALayout = std::tuple_element_t<0, Tuple>; - using BLayout = std::tuple_element_t<1, Tuple>; - using CLayout = std::tuple_element_t<2, Tuple>; - using ADataType = std::tuple_element_t<3, Tuple>; - using BDataType = std::tuple_element_t<4, Tuple>; - using AccDataType = std::tuple_element_t<5, Tuple>; - using CDataType = std::tuple_element_t<6, Tuple>; - static constexpr auto Scheduler = std::tuple_element_t<7, Tuple>::value; + using ALayout = std::tuple_element_t<0, Tuple>; + using BLayout = std::tuple_element_t<1, Tuple>; + using CLayout = std::tuple_element_t<2, Tuple>; + using ADataType = std::tuple_element_t<3, Tuple>; + using BDataType = std::tuple_element_t<4, Tuple>; + using AccDataType = std::tuple_element_t<5, Tuple>; + using CDataType = std::tuple_element_t<6, Tuple>; + static constexpr auto Scheduler = std::tuple_element_t<7, Tuple>::value; + static constexpr auto PipelineType = std::tuple_element_t<8, Tuple>::value; // TODO: expose tile size through test t-param ? struct gemm_args @@ -74,8 +80,13 @@ class TestCkTileGemmMemPipeline : public ::testing::Test using Traits = ck_tile::TileGemmTraits; - using BaseGemmPipeline = ck_tile::BaseGemmPipelineAgBgCrMem< - ck_tile::GemmPipelineProblem>; + using BaseGemmPipeline = std::conditional_t< + PipelineType == GemmPipelineType::Mem, + ck_tile::BaseGemmPipelineAgBgCrMem< + ck_tile::GemmPipelineProblem>, + ck_tile::BaseGemmPipelineAgBgCrCompV3< + ck_tile:: + GemmPipelineProblem>>; const ck_tile::index_t num_loop = TilePartitioner::GetLoopNum(args.K); const bool has_hot_loop = BaseGemmPipeline::BlockHasHotloop(num_loop); @@ -85,15 +96,26 @@ class TestCkTileGemmMemPipeline : public ::testing::Test constexpr bool has_hot_loop_v = has_hot_loop_.value; constexpr auto tail_number_v = tail_number_.value; - using GemmPipeline = ck_tile::GemmPipelineAgBgCrMem< - ck_tile::UniversalGemmPipelineProblem>; + using GemmPipeline = + std::conditional_t>, + ck_tile::GemmPipelineAgBgCrCompV3< + ck_tile::UniversalGemmPipelineProblem>>; using Kernel = ck_tile::GemmKernel; auto kargs = Kernel::MakeKargs(args.p_a, args.p_b, From 0e54d7ae5a638c9c1cbdc478dd12159354cd7e97 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Tue, 17 Dec 2024 06:57:55 -0800 Subject: [PATCH 12/25] Bump rocm-docs-core from 1.11.0 to 1.12.0 in /docs/sphinx (#1753) Bumps [rocm-docs-core](https://github.com/ROCm/rocm-docs-core) from 1.11.0 to 1.12.0. - [Release notes](https://github.com/ROCm/rocm-docs-core/releases) - [Changelog](https://github.com/ROCm/rocm-docs-core/blob/develop/CHANGELOG.md) - [Commits](https://github.com/ROCm/rocm-docs-core/compare/v1.11.0...v1.12.0) --- updated-dependencies: - dependency-name: rocm-docs-core dependency-type: direct:production update-type: version-update:semver-minor ... Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> --- docs/sphinx/requirements.in | 2 +- docs/sphinx/requirements.txt | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/sphinx/requirements.in b/docs/sphinx/requirements.in index d1b3465b9c..46a61a87fc 100644 --- a/docs/sphinx/requirements.in +++ b/docs/sphinx/requirements.in @@ -1,2 +1,2 @@ -rocm-docs-core==1.11.0 +rocm-docs-core==1.12.0 sphinxcontrib-bibtex==2.6.3 diff --git a/docs/sphinx/requirements.txt b/docs/sphinx/requirements.txt index 26d0aa2446..c2e74baae3 100644 --- a/docs/sphinx/requirements.txt +++ b/docs/sphinx/requirements.txt @@ -103,7 +103,7 @@ requests==2.32.3 # via # pygithub # sphinx -rocm-docs-core==1.11.0 +rocm-docs-core==1.12.0 # via -r requirements.in six==1.16.0 # via pybtex From 6ef8d3c295686b872d7e7a86621b68f765d98572 Mon Sep 17 00:00:00 2001 From: Max Podkorytov <4273004+tenpercent@users.noreply.github.com> Date: Thu, 12 Dec 2024 19:47:57 +0000 Subject: [PATCH 13/25] refactor conditional usage; fix build on rocm6.1 where the reference didn't exist --- include/ck/utility/amd_ck_fp8.hpp | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/include/ck/utility/amd_ck_fp8.hpp b/include/ck/utility/amd_ck_fp8.hpp index 7b21ad6464..1bdb1d078e 100644 --- a/include/ck/utility/amd_ck_fp8.hpp +++ b/include/ck/utility/amd_ck_fp8.hpp @@ -18,6 +18,12 @@ #define CK_USE_OCP_FP8 0 #endif +namespace { +// https://en.cppreference.com/w/cpp/types/conditional +template struct conditional { using type = T; }; +template struct conditional { using type = F; }; +} + namespace ck { using f8_fnuz_t = _BitInt(8); @@ -191,10 +197,10 @@ __host__ __device__ static inline T cast_from_f8(fp8_storage_t x) } } - typename __hip_internal::conditional< + typename conditional< sizeof(T) == 2, unsigned short int, - typename __hip_internal::conditional:: + typename conditional:: type>::type retval; if constexpr(we == 5 && is_half && !is_fnuz) @@ -538,10 +544,10 @@ __host__ __device__ static inline fp8_storage_t cast_to_f8(T _x, unsigned int rn constexpr int mfmt = (sizeof(T) == 8) ? 52 : ((sizeof(T) == 4) ? 23 : 10); - using T_bitwise = typename __hip_internal::conditional< + using T_bitwise = typename conditional< sizeof(T) == 2, unsigned short int, - typename __hip_internal::conditional:: + typename conditional:: type>::type; T_bitwise x_bitwise = bit_cast(_x); From 689a5ae45be802f51fc947a9f92208dcfb143f77 Mon Sep 17 00:00:00 2001 From: Illia Silin <98187287+illsilin@users.noreply.github.com> Date: Tue, 17 Dec 2024 10:17:29 -0800 Subject: [PATCH 14/25] Pass build flags to config.h (#1760) * pass the build flags to config.h * fix clang format --- CMakeLists.txt | 4 ++++ include/ck/config.h.in | 16 ++++++++++++++++ include/ck/utility/amd_ck_fp8.hpp | 20 +++++++++++++------- 3 files changed, 33 insertions(+), 7 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 2c86987561..be4efd3dfd 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -183,14 +183,17 @@ message("Building CK for the following targets: ${SUPPORTED_GPU_TARGETS}") if (SUPPORTED_GPU_TARGETS MATCHES "gfx9") message("Enabling XDL instances") add_definitions(-DCK_USE_XDL) + set(CK_USE_XDL "ON") endif() if (SUPPORTED_GPU_TARGETS MATCHES "gfx94") message("Enabling FP8 gemms on native architectures") add_definitions(-DCK_USE_GFX94) + set(CK_USE_GFX94 "ON") endif() if (SUPPORTED_GPU_TARGETS MATCHES "gfx11" OR SUPPORTED_GPU_TARGETS MATCHES "gfx12") message("Enabling WMMA instances") add_definitions(-DCK_USE_WMMA) + set(CK_USE_WMMA "ON") endif() if (SUPPORTED_GPU_TARGETS MATCHES "gfx12") add_definitions(-DCK_USE_OCP_FP8) @@ -204,6 +207,7 @@ endif() option(CK_USE_FP8_ON_UNSUPPORTED_ARCH "Enable FP8 GEMM instances on older architectures" OFF) if(CK_USE_FP8_ON_UNSUPPORTED_ARCH AND (SUPPORTED_GPU_TARGETS MATCHES "gfx90a" OR SUPPORTED_GPU_TARGETS MATCHES "gfx908")) add_definitions(-DCK_USE_FP8_ON_UNSUPPORTED_ARCH) + set(CK_USE_FP8_ON_UNSUPPORTED_ARCH "ON") endif() # CK config file to record supported datatypes, etc. diff --git a/include/ck/config.h.in b/include/ck/config.h.in index 0f0b7bd607..55a498073f 100644 --- a/include/ck/config.h.in +++ b/include/ck/config.h.in @@ -111,6 +111,22 @@ #cmakedefine CK_USE_WMMA @CK_USE_WMMA@ #endif +#ifndef CK_USE_GFX94 +#cmakedefine CK_USE_GFX94 @CK_USE_GFX94@ +#endif + +#ifndef DCK_USE_OCP_FP8 +#cmakedefine DCK_USE_OCP_FP8 @DCK_USE_OCP_FP8@ +#endif + +#ifndef CK_USE_FNUZ_FP8 +#cmakedefine CK_USE_FNUZ_FP8 @CK_USE_FNUZ_FP8@ +#endif + +#ifndef CK_USE_FP8_ON_UNSUPPORTED_ARCH +#cmakedefine CK_USE_FP8_ON_UNSUPPORTED_ARCH @CK_USE_FP8_ON_UNSUPPORTED_ARCH@ +#endif + // clang-format on #endif // CK_CONFIG_H_IN diff --git a/include/ck/utility/amd_ck_fp8.hpp b/include/ck/utility/amd_ck_fp8.hpp index 1bdb1d078e..e9174904c9 100644 --- a/include/ck/utility/amd_ck_fp8.hpp +++ b/include/ck/utility/amd_ck_fp8.hpp @@ -20,9 +20,17 @@ namespace { // https://en.cppreference.com/w/cpp/types/conditional -template struct conditional { using type = T; }; -template struct conditional { using type = F; }; -} +template +struct conditional +{ + using type = T; +}; +template +struct conditional +{ + using type = F; +}; +} // namespace namespace ck { @@ -200,8 +208,7 @@ __host__ __device__ static inline T cast_from_f8(fp8_storage_t x) typename conditional< sizeof(T) == 2, unsigned short int, - typename conditional:: - type>::type retval; + typename conditional::type>::type retval; if constexpr(we == 5 && is_half && !is_fnuz) { @@ -547,8 +554,7 @@ __host__ __device__ static inline fp8_storage_t cast_to_f8(T _x, unsigned int rn using T_bitwise = typename conditional< sizeof(T) == 2, unsigned short int, - typename conditional:: - type>::type; + typename conditional::type>::type; T_bitwise x_bitwise = bit_cast(_x); unsigned long long x{x_bitwise}; From d9e37c6874402023f5fe033f6821bde6869c5da5 Mon Sep 17 00:00:00 2001 From: Harisankar Sadasivan <135730918+hsadasiv@users.noreply.github.com> Date: Tue, 17 Dec 2024 10:31:21 -0800 Subject: [PATCH 15/25] updated fp16 instances to be on parity with universal gemm instances (#1754) * updated fp16 instances to be on parity with universal gemm instances * corrected instance name to streamk instance --- ...universal_streamk_f16_f16_f16_mk_kn_mn.hpp | 18 ++++++++++-- ...universal_streamk_f16_f16_f16_mk_nk_mn.hpp | 29 +++++++++++++++---- 2 files changed, 39 insertions(+), 8 deletions(-) mode change 100644 => 100755 library/src/tensor_operation_instance/gpu/gemm_universal_streamk/device_gemm_xdl_universal_streamk_f16_f16_f16/device_gemm_xdl_universal_streamk_f16_f16_f16_mk_kn_mn.hpp mode change 100644 => 100755 library/src/tensor_operation_instance/gpu/gemm_universal_streamk/device_gemm_xdl_universal_streamk_f16_f16_f16/device_gemm_xdl_universal_streamk_f16_f16_f16_mk_nk_mn.hpp diff --git a/library/src/tensor_operation_instance/gpu/gemm_universal_streamk/device_gemm_xdl_universal_streamk_f16_f16_f16/device_gemm_xdl_universal_streamk_f16_f16_f16_mk_kn_mn.hpp b/library/src/tensor_operation_instance/gpu/gemm_universal_streamk/device_gemm_xdl_universal_streamk_f16_f16_f16/device_gemm_xdl_universal_streamk_f16_f16_f16_mk_kn_mn.hpp old mode 100644 new mode 100755 index 6e8d5c798b..5460f7f857 --- a/library/src/tensor_operation_instance/gpu/gemm_universal_streamk/device_gemm_xdl_universal_streamk_f16_f16_f16/device_gemm_xdl_universal_streamk_f16_f16_f16_mk_kn_mn.hpp +++ b/library/src/tensor_operation_instance/gpu/gemm_universal_streamk/device_gemm_xdl_universal_streamk_f16_f16_f16/device_gemm_xdl_universal_streamk_f16_f16_f16_mk_kn_mn.hpp @@ -41,6 +41,8 @@ using device_gemm_xdl_universal_streamk_f16_f16_f16_mk_kn_mn_comp_instances = st //#########################| | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Row, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 256, 256, 32, 8, 4, 32, 32, 4, 4, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<8, 32, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 8, 4, 0, 1, 1, S<1, 32, 1, 8>, 8, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v4>, + DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Row, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 256, 256, 32, 4, 4, 32, 32, 4, 4, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 4, 4, 0, S<8, 32, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 4, 4, 0, 1, 1, S<1, 32, 1, 8>, 8, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v4>, + DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Row, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 256, 256, 32, 2, 2, 32, 32, 4, 4, S<16,16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 2, 2, 0, S<16,16, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 2, 2, 0, 1, 1, S<1, 32, 1, 8>, 8, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v4>, DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Row, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 128, 128, 64, 8, 4, 32, 32, 2, 2, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<16, 16, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 8, 4, 0, 1, 1, S<1, 32, 1, 8>, 8, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v4>, DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Row, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 256, 256, 32, 8, 4, 32, 32, 4, 4, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<8, 32, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 8, 4, 0, 1, 1, S<1, 32, 1, 8>, 8, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v5>, DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Row, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 256, 256, 32, 8, 4, 32, 32, 4, 4, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<8, 32, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 8, 4, 0, 1, 1, S<1, 32, 1, 8>, 8, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v3>, @@ -49,7 +51,9 @@ using device_gemm_xdl_universal_streamk_f16_f16_f16_mk_kn_mn_comp_instances = st DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Row, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 128, 128, 64, 8, 4, 32, 32, 2, 2, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<16, 16, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 8, 4, 0, 1, 1, S<1, 32, 1, 8>, 8, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v5>, DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Row, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 128, 256, 32, 8, 4, 32, 32, 2, 4, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<8, 32, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 8, 4, 0, 1, 1, S<1, 32, 1, 8>, 8, BlockGemmPipelineScheduler::Interwave, BlockGemmPipelineVersion::v1>, DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Row, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 256, 128, 32, 8, 4, 32, 32, 4, 2, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<8, 32, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 4, 4, 0, 1, 1, S<1, 32, 1, 8>, 8, BlockGemmPipelineScheduler::Interwave, BlockGemmPipelineVersion::v1>, - DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Row, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 128, 128, 64, 8, 4, 32, 32, 2, 2, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<16, 16, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 8, 4, 0, 1, 1, S<1, 32, 1, 8>, 8, BlockGemmPipelineScheduler::Interwave, BlockGemmPipelineVersion::v1> + DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Row, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 128, 128, 64, 8, 4, 32, 32, 2, 2, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<16, 16, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 8, 4, 0, 1, 1, S<1, 32, 1, 8>, 8, BlockGemmPipelineScheduler::Interwave, BlockGemmPipelineVersion::v1>, + DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Row, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 128, 128, 64, 4, 4, 32, 32, 2, 2, S<16,16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 4, 4, 0, S<16, 16, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 4, 4, 0, 1, 1, S<1, 32, 1, 8>, 8, BlockGemmPipelineScheduler::Interwave, BlockGemmPipelineVersion::v1>, + DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Row, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 128, 128, 64, 2, 2, 32, 32, 2, 2, S<32, 8, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 2, 2, 0, S<32, 8, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 2, 2, 0, 1, 1, S<1, 32, 1, 8>, 8, BlockGemmPipelineScheduler::Interwave, BlockGemmPipelineVersion::v1> // clang-format on >; @@ -61,14 +65,21 @@ using device_gemm_xdl_universal_streamk_f16_f16_f16_mk_kn_mn_mem_instances = std //#########################| | | | | | | | | Operation| Operation| Operation| | | | | | | | | | Wave| Wave| Lengths_K0_M_K1| ArrangeOrder| | | PerVector| PerVector_K1| | Lengths_K0_N_K1| ArrangeOrder| | | PerVector| PerVector_K1| | PerShuffle| PerShuffle| _NBlock_NXdlPerWave_NWaveNPerXdl| _NWaveNPerXdl| Scheduler| Verision| //#########################| | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | - // Latency friendly + // Latency friendly DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Row, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 128, 32, 16, 64, 8, 4, 16, 16, 1, 1, S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<16, 8, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 2, 4, 0, 1, 1, S<1, 16, 1, 8>, 2, BlkGemmPipeSched, BlockGemmPipelineVersion::v1>, - DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Row, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 64, 16, 16, 128, 8, 4, 16, 16, 1, 1, S<16, 4, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<16, 4, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 4, 4, 0, 1, 1, S<1, 16, 1, 4>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v1>, + DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Row, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 128, 32, 16, 64, 4, 4, 16, 16, 1, 1, S<16, 8, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 4, 4, 0, S<16, 8, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 2, 4, 0, 1, 1, S<1, 16, 1, 8>, 2, BlkGemmPipeSched, BlockGemmPipelineVersion::v1>, + DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Row, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 128, 32, 16, 64, 2, 2, 16, 16, 1, 1, S<32, 4, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 2, 2, 0, S<32, 4, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 2, 2, 0, 1, 1, S<1, 16, 1, 8>, 2, BlkGemmPipeSched, BlockGemmPipelineVersion::v1>, + DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Row, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 64, 16, 16, 128, 8, 4, 16, 16, 1, 1, S<16, 4, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<32, 2, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 4, 4, 0, 1, 1, S<1, 16, 1, 4>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v1>, + DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Row, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 64, 16, 16, 128, 4, 4, 16, 16, 1, 1, S<32, 2, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 4, 4, 0, S<32, 2, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 4, 4, 0, 1, 1, S<1, 16, 1, 4>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v1>, DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Row, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 64, 16, 16, 64, 8, 4, 16, 16, 1, 1, S<8, 8, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<16, 4, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 4, 4, 0, 1, 1, S<1, 16, 1, 4>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v1>, + DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Row, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 64, 16, 16, 64, 4, 4, 16, 16, 1, 1, S<8, 8, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 4, 4, 0, S<16, 4, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 4, 4, 0, 1, 1, S<1, 16, 1, 4>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v1>, + DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Row, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 64, 16, 16, 64, 2, 2, 16, 16, 1, 1, S<8, 8, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 2, 2, 0, S<16, 4, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 2, 2, 0, 1, 1, S<1, 16, 1, 4>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v1>, DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Row, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 128, 16, 32, 64, 8, 4, 16, 16, 1, 1, S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<16, 8, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 4, 4, 0, 1, 1, S<1, 16, 1, 8>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v1>, // Memory friendly DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Row, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 256, 32, 64, 8, 2, 32, 32, 2, 1, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<32, 8, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 4, 2, 0, 1, 1, S<1, 32, 1, 8>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v2>, + DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Row, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 256, 32, 64, 2, 2, 32, 32, 2, 1, S<32, 8, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 2, 2, 0, S<32, 8, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 2, 2, 0, 1, 1, S<1, 32, 1, 8>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v2>, DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Row, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 256, 16, 64, 8, 2, 16, 16, 4, 1, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<32, 8, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 2, 2, 0, 1, 1, S<1, 32, 1, 8>, 2, BlkGemmPipeSched, BlockGemmPipelineVersion::v2>, + DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Row, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 256, 16, 64, 2, 2, 16, 16, 4, 1, S<32, 8, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 2, 2, 0, S<32, 8, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 2, 2, 0, 1, 1, S<1, 32, 1, 8>, 2, BlkGemmPipeSched, BlockGemmPipelineVersion::v2>, DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Row, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 128, 128, 32, 64, 8, 4, 32, 32, 2, 1, S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<16, 8, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 4, 4, 0, 1, 1, S<1, 16, 1, 8>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v2>, DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Row, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 128, 128, 16, 64, 8, 4, 16, 16, 4, 1, S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<16, 8, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 2, 4, 0, 1, 1, S<1, 16, 1, 8>, 2, BlkGemmPipeSched, BlockGemmPipelineVersion::v2>, DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Row, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 128, 64, 32, 64, 8, 4, 32, 32, 1, 1, S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<16, 8, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 4, 4, 0, 1, 1, S<1, 16, 1, 8>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v2>, @@ -82,6 +93,7 @@ using device_gemm_xdl_universal_streamk_f16_f16_f16_mk_kn_mn_mem_instances = std DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Row, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 128, 16, 128, 64, 8, 4, 16, 16, 1, 4, S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<8, 16, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 8, 4, 0, 1, 1, S<1, 16, 1, 8>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v2>, DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Row, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 128, 32, 128, 64, 8, 4, 32, 32, 1, 2, S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<8, 16, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 8, 4, 0, 1, 1, S<1, 16, 1, 8>, 8, BlkGemmPipeSched, BlockGemmPipelineVersion::v2>, DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Row, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 16, 256, 64, 8, 4, 16, 16, 1, 4, S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<8, 32, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 8, 4, 0, 1, 1, S<1, 16, 1, 16>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v2>, + DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Row, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 16, 256, 64, 4, 4, 16, 16, 1, 4, S<16, 8, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 4, 4, 0, S<8, 32, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 4, 4, 0, 1, 1, S<1, 16, 1, 16>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v2>, DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Row, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 32, 256, 64, 8, 4, 32, 32, 1, 2, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<8, 32, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 8, 4, 0, 1, 1, S<1, 16, 1, 16>, 8, BlkGemmPipeSched, BlockGemmPipelineVersion::v2> // clang-format on >; diff --git a/library/src/tensor_operation_instance/gpu/gemm_universal_streamk/device_gemm_xdl_universal_streamk_f16_f16_f16/device_gemm_xdl_universal_streamk_f16_f16_f16_mk_nk_mn.hpp b/library/src/tensor_operation_instance/gpu/gemm_universal_streamk/device_gemm_xdl_universal_streamk_f16_f16_f16/device_gemm_xdl_universal_streamk_f16_f16_f16_mk_nk_mn.hpp old mode 100644 new mode 100755 index e00c1733e0..e716b3e85c --- a/library/src/tensor_operation_instance/gpu/gemm_universal_streamk/device_gemm_xdl_universal_streamk_f16_f16_f16/device_gemm_xdl_universal_streamk_f16_f16_f16_mk_nk_mn.hpp +++ b/library/src/tensor_operation_instance/gpu/gemm_universal_streamk/device_gemm_xdl_universal_streamk_f16_f16_f16/device_gemm_xdl_universal_streamk_f16_f16_f16_mk_nk_mn.hpp @@ -42,14 +42,21 @@ using device_gemm_xdl_universal_streamk_f16_f16_f16_mk_nk_mn_comp_instances = st // Compute friendly DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Col, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 256, 256, 32, 8, 8, 32, 32, 4, 4, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, 1, 1, S<1, 32, 1, 8>, 8, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v4>, + DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Col, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 256, 256, 32, 4, 4, 32, 32, 4, 4, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 4, 4, 0, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 4, 4, 0, 1, 1, S<1, 32, 1, 8>, 8, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v4>, DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Col, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 128, 128, 64, 8, 8, 32, 32, 2, 2, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, 1, 1, S<1, 32, 1, 8>, 8, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v4>, DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Col, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 128, 128, 32, 8, 8, 32, 32, 2, 2, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, 1, 1, S<1, 32, 1, 8>, 8, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v4>, + DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Col, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 128, 128, 32, 4, 4, 32, 32, 2, 2, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 4, 4, 0, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 4, 4, 0, 1, 1, S<1, 32, 1, 8>, 8, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v4>, + DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Col, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 128, 128, 32, 2, 2, 32, 32, 2, 2, S<16,16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 2, 2, 0, S<16,16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 2, 2, 0, 1, 1, S<1, 32, 1, 8>, 8, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v4>, DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Col, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 256, 256, 32, 8, 8, 32, 32, 4, 4, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, 1, 1, S<1, 32, 1, 8>, 8, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v3>, + DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Col, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 256, 256, 32, 4, 4, 32, 32, 4, 4, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 4, 4, 0, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 4, 4, 0, 1, 1, S<1, 32, 1, 8>, 8, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v3>, + DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Col, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 256, 256, 32, 2, 2, 32, 32, 4, 4, S<16,16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 2, 2, 0, S<16,16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 2, 2, 0, 1, 1, S<1, 32, 1, 8>, 8, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v3>, DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Col, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 256, 256, 32, 8, 8, 32, 32, 4, 4, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, 1, 1, S<1, 32, 1, 8>, 8, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v5>, + DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Col, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 256, 256, 32, 4, 4, 32, 32, 4, 4, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 4, 4, 0, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 4, 4, 0, 1, 1, S<1, 32, 1, 8>, 8, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v5>, + DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Col, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 256, 256, 32, 2, 2, 32, 32, 4, 4, S<16,16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 2, 2, 0, S<16,16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 2, 2, 0, 1, 1, S<1, 32, 1, 8>, 8, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v5>, DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Col, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 256, 256, 32, 8, 8, 16, 16, 8, 8, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, 1, 2, S<1, 32, 1, 8>, 8, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v3>, - // AGPR Spill - // DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Col, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 256, 256, 32, 8, 8, 16, 16, 8, 8, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, 1, 2, S<1, 32, 1, 8>, 8, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v5>, - // AGPR Spill when use permuted lds layout. so, use padding for these two. + DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Col, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 256, 256, 32, 4, 4, 16, 16, 8, 8, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 4, 4, 0, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 4, 4, 0, 1, 2, S<1, 32, 1, 8>, 8, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v3>, + DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Col, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 256, 256, 32, 2, 2, 16, 16, 8, 8, S<16,16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 2, 2, 0, S<16,16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 2, 2, 0, 1, 2, S<1, 32, 1, 8>, 8, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v3>, + DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Col, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 224, 256, 64, 8, 8, 16, 16, 7, 8, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, 1, 2, S<1, 32, 1, 8>, 8, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v3>, DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Col, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 256, 224, 64, 8, 8, 16, 16, 8, 7, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, 2, 1, S<1, 64, 1, 4>, 8, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v3>, DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Col, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 128, 128, 64, 8, 8, 32, 32, 2, 2, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, 1, 1, S<1, 32, 1, 8>, 8, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v3>, @@ -68,15 +75,23 @@ using device_gemm_xdl_universal_streamk_f16_f16_f16_mk_nk_mn_mem_instances = std //#########################| | | | | | | | | Operation| Operation| Operation| | | | | | | | | | Wave| Wave| Lengths_K0_M_K1| ArrangeOrder| | | PerVector| PerVector_K1| | Lengths_K0_N_K1| ArrangeOrder| | | PerVector| PerVector_K1| | PerShuffle| PerShuffle| _NBlock_NXdlPerWave_NWaveNPerXdl| _NWaveNPerXdl| Scheduler| Verision| //#########################| | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | - // Latency friendly + // Latency friendly DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Col, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 128, 32, 16, 64, 8, 8, 16, 16, 1, 1, S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, 1, 1, S<1, 16, 1, 8>, 2, BlkGemmPipeSched, BlockGemmPipelineVersion::v1>, + DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Col, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 128, 32, 16, 64, 4, 4, 16, 16, 1, 1, S<16, 8, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 4, 4, 0, S<16, 8, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 4, 4, 0, 1, 1, S<1, 16, 1, 8>, 2, BlkGemmPipeSched, BlockGemmPipelineVersion::v1>, + DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Col, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 128, 32, 16, 64, 2, 2, 16, 16, 1, 1, S<32, 4, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 2, 2, 0, S<32, 4, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 2, 2, 0, 1, 1, S<1, 16, 1, 8>, 2, BlkGemmPipeSched, BlockGemmPipelineVersion::v1>, DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Col, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 64, 16, 16, 128, 8, 8, 16, 16, 1, 1, S<16, 4, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<16, 4, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, 1, 1, S<1, 16, 1, 4>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v1>, DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Col, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 64, 16, 16, 64, 8, 8, 16, 16, 1, 1, S<8, 8, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<8, 8, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, 1, 1, S<1, 16, 1, 4>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v1>, + DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Col, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 64, 16, 16, 64, 4, 4, 16, 16, 1, 1, S<16, 4, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 4, 4, 0, S<16, 4, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 4, 4, 0, 1, 1, S<1, 16, 1, 4>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v1>, + DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Col, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 64, 16, 16, 64, 2, 2, 16, 16, 1, 1, S<32, 2, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 2, 2, 0, S<32, 2, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 2, 2, 0, 1, 1, S<1, 16, 1, 4>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v1>, DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Col, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 128, 16, 32, 64, 8, 8, 16, 16, 1, 1, S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, 1, 1, S<1, 16, 1, 8>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v1>, // Memory friendly DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Col, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 256, 32, 64, 8, 8, 32, 32, 2, 1, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, 1, 1, S<1, 32, 1, 8>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v2>, + DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Col, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 256, 32, 64, 4, 4, 32, 32, 2, 1, S<16,16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 4, 4, 0, S<16,16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 4, 4, 0, 1, 1, S<1, 32, 1, 8>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v2>, + DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Col, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 256, 32, 64, 2, 2, 32, 32, 2, 1, S<32, 8, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 2, 2, 0, S<32, 8, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 2, 2, 0, 1, 1, S<1, 32, 1, 8>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v2>, DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Col, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 256, 16, 64, 8, 8, 16, 16, 4, 1, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, 1, 1, S<1, 32, 1, 8>, 2, BlkGemmPipeSched, BlockGemmPipelineVersion::v2>, DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Col, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 128, 128, 32, 64, 8, 8, 32, 32, 2, 1, S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, 1, 1, S<1, 16, 1, 8>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v2>, + DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Col, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 128, 128, 32, 64, 4, 4, 32, 32, 2, 1, S<16, 8, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 4, 4, 0, S<16, 8, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 4, 4, 0, 1, 1, S<1, 16, 1, 8>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v2>, + DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Col, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 128, 128, 32, 64, 2, 2, 32, 32, 2, 1, S<32, 4, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 2, 2, 0, S<32, 4, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 2, 2, 0, 1, 1, S<1, 16, 1, 8>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v2>, DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Col, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 128, 128, 16, 64, 8, 8, 16, 16, 4, 1, S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, 1, 1, S<1, 16, 1, 8>, 2, BlkGemmPipeSched, BlockGemmPipelineVersion::v2>, DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Col, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 128, 64, 32, 64, 8, 8, 32, 32, 1, 1, S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, 1, 1, S<1, 16, 1, 8>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v2>, DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Col, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 128, 64, 16, 64, 8, 8, 16, 16, 2, 1, S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, 1, 1, S<1, 16, 1, 8>, 2, BlkGemmPipeSched, BlockGemmPipelineVersion::v2>, @@ -84,12 +99,16 @@ using device_gemm_xdl_universal_streamk_f16_f16_f16_mk_nk_mn_mem_instances = std DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Col, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 64, 16, 16, 128, 8, 8, 16, 16, 1, 1, S<16, 4, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<16, 4, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, 1, 1, S<1, 16, 1, 4>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v2>, DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Col, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 64, 16, 16, 64, 8, 8, 16, 16, 1, 1, S<8, 8, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<8, 8, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, 1, 1, S<1, 16, 1, 4>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v2>, DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Col, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 128, 16, 32, 64, 8, 8, 16, 16, 1, 1, S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, 1, 1, S<1, 16, 1, 8>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v2>, + DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Col, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 128, 16, 32, 64, 4, 4, 16, 16, 1, 1, S<16, 8, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 4, 4, 0, S<16, 8, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 4, 4, 0, 1, 1, S<1, 16, 1, 8>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v2>, + DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Col, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 128, 16, 32, 64, 2, 2, 16, 16, 1, 1, S<32, 4, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 2, 2, 0, S<32, 4, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 2, 2, 0, 1, 1, S<1, 16, 1, 8>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v2>, DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Col, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 128, 16, 64, 64, 8, 8, 16, 16, 1, 2, S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, 1, 1, S<1, 16, 1, 8>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v2>, DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Col, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 128, 32, 64, 64, 8, 8, 32, 32, 1, 1, S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, 1, 1, S<1, 16, 1, 8>, 8, BlkGemmPipeSched, BlockGemmPipelineVersion::v2>, DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Col, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 128, 16, 128, 64, 8, 8, 16, 16, 1, 4, S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, 1, 1, S<1, 16, 1, 8>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v2>, DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Col, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 128, 32, 128, 64, 8, 8, 32, 32, 1, 2, S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, 1, 1, S<1, 16, 1, 8>, 8, BlkGemmPipeSched, BlockGemmPipelineVersion::v2>, DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Col, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 16, 256, 64, 8, 8, 16, 16, 1, 4, S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, 1, 1, S<1, 16, 1, 16>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v2>, - DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Col, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 32, 256, 64, 8, 8, 32, 32, 1, 2, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, 1, 1, S<1, 16, 1, 16>, 8, BlkGemmPipeSched, BlockGemmPipelineVersion::v2> + DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Col, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 32, 256, 64, 8, 8, 32, 32, 1, 2, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, 1, 1, S<1, 16, 1, 16>, 8, BlkGemmPipeSched, BlockGemmPipelineVersion::v2>, + DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Col, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 32, 256, 64, 4, 4, 32, 32, 1, 2, S<16,16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 4, 4, 0, S<16,16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 4, 4, 0, 1, 1, S<1, 16, 1, 16>, 8, BlkGemmPipeSched, BlockGemmPipelineVersion::v2>, + DeviceGemm_Xdl_CShuffle_Streamk_V3< Row, Col, Row, F16, F16, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 32, 256, 64, 2, 2, 32, 32, 1, 2, S<32, 8, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 2, 2, 0, S<32, 8, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 2, 2, 0, 1, 1, S<1, 16, 1, 16>, 8, BlkGemmPipeSched, BlockGemmPipelineVersion::v2> // clang-format on >; } // namespace instance From f6c4d614e35b7424774160a23d8e8bef3b15faad Mon Sep 17 00:00:00 2001 From: aledudek Date: Wed, 18 Dec 2024 09:45:58 +0100 Subject: [PATCH 16/25] [CK_TILE] Move hipmalloc/memcpy calls out of gpu reference gemm (#1743) * [CK_TILE] Move hipmalloc/memcpy calls out of gpu reference gemm * [CK_TILE] Move hipmalloc/memcpy calls out of gpu reference gemm - review changes * [CK_TILE] Move hipmalloc/memcpy calls out of gpu reference gemm - review fix --- example/ck_tile/03_gemm/run_gemm_example.inc | 29 +++- .../run_batched_gemm_example.inc | 33 +++- .../ck_tile/host/reference/reference_gemm.hpp | 162 ++---------------- 3 files changed, 68 insertions(+), 156 deletions(-) diff --git a/example/ck_tile/03_gemm/run_gemm_example.inc b/example/ck_tile/03_gemm/run_gemm_example.inc index a1fc155775..2b7a967bab 100644 --- a/example/ck_tile/03_gemm/run_gemm_example.inc +++ b/example/ck_tile/03_gemm/run_gemm_example.inc @@ -161,14 +161,39 @@ int run_gemm_example_with_layouts(int argc, c_m_n_gpu_ref.SetZero(); c_m_n_gpu_buf_ref.SetZero(); + ADataType* d_A; + BDataType* d_B; + CDataType* d_C; + + ck_tile::hip_check_error(hipMalloc(&d_A, M * K * sizeof(ADataType))); + ck_tile::hip_check_error(hipMalloc(&d_B, N * K * sizeof(BDataType))); + ck_tile::hip_check_error(hipMalloc(&d_C, M * N * sizeof(CDataType))); + + ck_tile::hip_check_error(hipMemcpy(d_A, + a_m_k_dev_buf.GetDeviceBuffer(), + M * K * sizeof(ADataType), + hipMemcpyHostToDevice)); + ck_tile::hip_check_error(hipMemcpy(d_B, + b_k_n_dev_buf.GetDeviceBuffer(), + N * K * sizeof(BDataType), + hipMemcpyHostToDevice)); + ck_tile::reference_gemm_gpu( - a_m_k_dev_buf, b_k_n_dev_buf, c_m_n_gpu_buf_ref, M, N, K, stride_A, stride_B, stride_C); + CLayout>(d_A, d_B, d_C, M, N, K, stride_A, stride_B, stride_C); + + ck_tile::hip_check_error(hipMemcpy(c_m_n_gpu_buf_ref.GetDeviceBuffer(), + d_C, + M * N * sizeof(CDataType), + hipMemcpyDeviceToHost)); + + ck_tile::hip_check_error(hipFree(d_A)); + ck_tile::hip_check_error(hipFree(d_B)); + ck_tile::hip_check_error(hipFree(d_C)); c_m_n_gpu_buf_ref.FromDevice(c_m_n_gpu_ref.data()); pass = ck_tile::check_err(c_m_n_dev_result, c_m_n_gpu_ref); diff --git a/example/ck_tile/16_batched_gemm/run_batched_gemm_example.inc b/example/ck_tile/16_batched_gemm/run_batched_gemm_example.inc index dacca2042e..8345eef95b 100644 --- a/example/ck_tile/16_batched_gemm/run_batched_gemm_example.inc +++ b/example/ck_tile/16_batched_gemm/run_batched_gemm_example.inc @@ -188,15 +188,33 @@ int run_batched_gemm_example_with_layouts(int argc, c_m_n_gpu_ref.SetZero(); c_m_n_gpu_buf_ref.SetZero(); + ADataType* d_A; + BDataType* d_B; + CDataType* d_C; + + ck_tile::hip_check_error(hipMalloc(&d_A, batch_count * M * K * sizeof(ADataType))); + ck_tile::hip_check_error(hipMalloc(&d_B, batch_count * N * K * sizeof(BDataType))); + ck_tile::hip_check_error(hipMalloc(&d_C, batch_count * M * N * sizeof(CDataType))); + + ck_tile::hip_check_error(hipMemcpy(d_A, + a_m_k_dev_buf.GetDeviceBuffer(), + batch_count * M * K * sizeof(ADataType), + hipMemcpyHostToDevice)); + + ck_tile::hip_check_error(hipMemcpy(d_B, + b_k_n_dev_buf.GetDeviceBuffer(), + batch_count * N * K * sizeof(BDataType), + hipMemcpyHostToDevice)); + ck_tile::reference_batched_gemm_gpu(a_m_k_dev_buf, - b_k_n_dev_buf, - c_m_n_gpu_buf_ref, + CLayout>(d_A, + d_B, + d_C, M, N, K, @@ -208,6 +226,15 @@ int run_batched_gemm_example_with_layouts(int argc, batch_stride_C, batch_count); + ck_tile::hip_check_error(hipMemcpy(c_m_n_gpu_buf_ref.GetDeviceBuffer(), + d_C, + batch_count * M * N * sizeof(CDataType), + hipMemcpyDeviceToHost)); + + ck_tile::hip_check_error(hipFree(d_A)); + ck_tile::hip_check_error(hipFree(d_B)); + ck_tile::hip_check_error(hipFree(d_C)); + c_m_n_gpu_buf_ref.FromDevice(c_m_n_gpu_ref.data()); pass = ck_tile::check_err(c_m_n_dev_result, c_m_n_gpu_ref); diff --git a/include/ck_tile/host/reference/reference_gemm.hpp b/include/ck_tile/host/reference/reference_gemm.hpp index 8bd1f5b048..fc412e8831 100644 --- a/include/ck_tile/host/reference/reference_gemm.hpp +++ b/include/ck_tile/host/reference/reference_gemm.hpp @@ -97,9 +97,9 @@ template -void reference_gemm_gpu(DeviceMem& a_device, - DeviceMem& b_device, - DeviceMem& c_device, +void reference_gemm_gpu(ADataType* a_ptr, + BDataType* b_ptr, + CDataType* c_ptr, index_t M, index_t N, index_t K, @@ -107,79 +107,13 @@ void reference_gemm_gpu(DeviceMem& a_device, index_t stride_b, index_t stride_c) { - - ADataType* d_A; - BDataType* d_B; - CDataType* d_C; - - hipError_t errA = hipMalloc(&d_A, M * K * sizeof(ADataType)); - hipError_t errB = hipMalloc(&d_B, N * K * sizeof(BDataType)); - hipError_t errC = hipMalloc(&d_C, M * N * sizeof(CDataType)); - if(errA != hipSuccess) - { - std::cerr << "Error allocating device memory for A: " << hipGetErrorString(errA) - << std::endl; - return; // Early exit on error - } - - if(errB != hipSuccess) - { - std::cerr << "Error allocating device memory for B: " << hipGetErrorString(errB) - << std::endl; - return; // Early exit on error - } - - if(errC != hipSuccess) - { - std::cerr << "Error allocating device memory for C: " << hipGetErrorString(errC) - << std::endl; - return; // Early exit on error - } - - errA = hipMemcpy( - d_A, a_device.GetDeviceBuffer(), M * K * sizeof(ADataType), hipMemcpyHostToDevice); - if(errA != hipSuccess) - { - std::cerr << "Error copying A to device: " << hipGetErrorString(errA) << std::endl; - } - - errB = hipMemcpy( - d_B, b_device.GetDeviceBuffer(), N * K * sizeof(BDataType), hipMemcpyHostToDevice); - if(errB != hipSuccess) - { - std::cerr << "Error copying B to device: " << hipGetErrorString(errB) << std::endl; - } - int totalElements = M * N; int numThreadsPerBlock = 256; // Common choice for threads per block int numBlocks = (totalElements + numThreadsPerBlock - 1) / numThreadsPerBlock; naive_gemm_kernel - <<>>(d_A, d_B, d_C, M, N, K, stride_a, stride_b, stride_c); - errC = hipMemcpy( - c_device.GetDeviceBuffer(), d_C, M * N * sizeof(CDataType), hipMemcpyDeviceToHost); - if(errC != hipSuccess) - { - std::cerr << "Error copying C to device: " << hipGetErrorString(errC) << std::endl; - } - - errA = hipFree(d_A); - if(errA != hipSuccess) - { - std::cerr << "Error free the A memory: " << hipGetErrorString(errA) << std::endl; - } - - errB = hipFree(d_B); - if(errB != hipSuccess) - { - std::cerr << "Error free the B memory: " << hipGetErrorString(errB) << std::endl; - } - - errC = hipFree(d_C); - if(errC != hipSuccess) - { - std::cerr << "Error free the C memory: " << hipGetErrorString(errC) << std::endl; - } + <<>>( + a_ptr, b_ptr, c_ptr, M, N, K, stride_a, stride_b, stride_c); return; } @@ -191,9 +125,9 @@ template -void reference_batched_gemm_gpu(DeviceMem& a_device, - DeviceMem& b_device, - DeviceMem& c_device, +void reference_batched_gemm_gpu(ADataType* a_ptr, + BDataType* b_ptr, + CDataType* c_ptr, index_t M, index_t N, index_t K, @@ -205,94 +139,20 @@ void reference_batched_gemm_gpu(DeviceMem& a_device, index_t batch_stride_C, index_t batch_count) { - - ADataType* d_A; - BDataType* d_B; - CDataType* d_C; - - hipError_t errA = hipMalloc(&d_A, batch_count * M * K * sizeof(ADataType)); - hipError_t errB = hipMalloc(&d_B, batch_count * N * K * sizeof(BDataType)); - hipError_t errC = hipMalloc(&d_C, batch_count * M * N * sizeof(CDataType)); - if(errA != hipSuccess) - { - std::cerr << "Error allocating device memory for A: " << hipGetErrorString(errA) - << std::endl; - return; // Early exit on error - } - - if(errB != hipSuccess) - { - std::cerr << "Error allocating device memory for B: " << hipGetErrorString(errB) - << std::endl; - return; // Early exit on error - } - - if(errC != hipSuccess) - { - std::cerr << "Error allocating device memory for C: " << hipGetErrorString(errC) - << std::endl; - return; // Early exit on error - } - - errA = hipMemcpy(d_A, - a_device.GetDeviceBuffer(), - batch_count * M * K * sizeof(ADataType), - hipMemcpyHostToDevice); - if(errA != hipSuccess) - { - std::cerr << "Error copying A to device: " << hipGetErrorString(errA) << std::endl; - } - - errB = hipMemcpy(d_B, - b_device.GetDeviceBuffer(), - batch_count * N * K * sizeof(BDataType), - hipMemcpyHostToDevice); - if(errB != hipSuccess) - { - std::cerr << "Error copying B to device: " << hipGetErrorString(errB) << std::endl; - } - int totalElements = M * N; int numThreadsPerBlock = 256; // Common choice for threads per block int numBlocks = (totalElements + numThreadsPerBlock - 1) / numThreadsPerBlock; for(index_t batch_id = 0; batch_id < batch_count; ++batch_id) { - ADataType* d_ATemp = d_A + batch_id * batch_stride_A; - BDataType* d_BTemp = d_B + batch_id * batch_stride_B; - CDataType* d_CTemp = d_C + batch_id * batch_stride_C; + ADataType* d_ATemp = a_ptr + batch_id * batch_stride_A; + BDataType* d_BTemp = b_ptr + batch_id * batch_stride_B; + CDataType* d_CTemp = c_ptr + batch_id * batch_stride_C; naive_gemm_kernel <<>>( d_ATemp, d_BTemp, d_CTemp, M, N, K, stride_a, stride_b, stride_c); } - errC = hipMemcpy(c_device.GetDeviceBuffer(), - d_C, - batch_count * M * N * sizeof(CDataType), - hipMemcpyDeviceToHost); - if(errC != hipSuccess) - { - std::cerr << "Error copying C to device: " << hipGetErrorString(errC) << std::endl; - } - - errA = hipFree(d_A); - if(errA != hipSuccess) - { - std::cerr << "Error free the A memory: " << hipGetErrorString(errA) << std::endl; - } - - errB = hipFree(d_B); - if(errB != hipSuccess) - { - std::cerr << "Error free the B memory: " << hipGetErrorString(errB) << std::endl; - } - - errC = hipFree(d_C); - if(errC != hipSuccess) - { - std::cerr << "Error free the C memory: " << hipGetErrorString(errC) << std::endl; - } - return; } } // namespace ck_tile From 1c1b336371e2367fece6b33644b36ab30d92b2d3 Mon Sep 17 00:00:00 2001 From: Xiaodong Wang Date: Wed, 18 Dec 2024 02:32:38 -0800 Subject: [PATCH 17/25] Disambiguate bit_cast (#1749) Adding namespace to disambiguate with std::bit_cast Co-authored-by: Po Yen Chen --- include/ck_tile/core/container/meta_data_buffer.hpp | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/include/ck_tile/core/container/meta_data_buffer.hpp b/include/ck_tile/core/container/meta_data_buffer.hpp index 7493b93d80..eba60fac75 100644 --- a/include/ck_tile/core/container/meta_data_buffer.hpp +++ b/include/ck_tile/core/container/meta_data_buffer.hpp @@ -30,7 +30,7 @@ struct meta_data_buffer { constexpr index_t size = sizeof(T); - auto tmp = bit_cast>(data); + auto tmp = ck_tile::bit_cast>(data); for(int i = 0; i < size; i++) { @@ -66,7 +66,7 @@ struct meta_data_buffer pos++; } - data = bit_cast(tmp); + data = ck_tile::bit_cast(tmp); } return data; @@ -86,7 +86,7 @@ struct meta_data_buffer pos++; } - auto data = bit_cast(tmp); + auto data = ck_tile::bit_cast(tmp); return data; } From 453ca373479e1c3510bff66c03a773a29f1caada Mon Sep 17 00:00:00 2001 From: aledudek Date: Wed, 18 Dec 2024 17:52:46 +0100 Subject: [PATCH 18/25] [CK TILE] Refactor GemmKernel to be reused by other GEMM related operators (#1730) * Gemm Kernel Refactor part1 * Gemm Kernel Refactor common gemm pipeline part2 * [CK TILE] Refactor batched gemm to reuse GemmKernel * [CK TILE] Refactor GemmKernel - review changes part1 * [CK TILE] Refactor GemmKernel - references fix * [CK TILE] Refactor GemmKernel - naming changes, add problem * [CK_TILE] Refactor GemmKernel - update tests * [CK_TILE] Refactor GemmKernel - review changes * [CK_TILE] Refactor GemmKernel - update test * [CK_TILE] Refactor GemmKernel - constness fixes * [CK_TILE] Refactor GemmKernel - update tests --- example/ck_tile/03_gemm/gemm_basic.cpp | 14 +- example/ck_tile/03_gemm/gemm_basic.hpp | 16 +- example/ck_tile/03_gemm/run_gemm_example.inc | 10 +- .../ck_tile/16_batched_gemm/batched_gemm.cpp | 6 +- .../ck_tile/16_batched_gemm/batched_gemm.hpp | 6 +- .../run_batched_gemm_example.inc | 2 +- .../ops/gemm/kernel/batched_gemm_kernel.hpp | 274 +++++------------- .../ck_tile/ops/gemm/kernel/gemm_kernel.hpp | 259 ++++++++++++----- .../batched_gemm/test_batched_gemm_util.hpp | 40 ++- test/ck_tile/gemm/test_gemm_pipeline_util.hpp | 38 +-- 10 files changed, 297 insertions(+), 368 deletions(-) diff --git a/example/ck_tile/03_gemm/gemm_basic.cpp b/example/ck_tile/03_gemm/gemm_basic.cpp index f5260c306e..4c630375f4 100644 --- a/example/ck_tile/03_gemm/gemm_basic.cpp +++ b/example/ck_tile/03_gemm/gemm_basic.cpp @@ -15,7 +15,7 @@ #include "gemm_basic.hpp" template -float gemm_calc(const gemm_basic_args& args, const ck_tile::stream_config& s) +float gemm_calc(const ck_tile::GemmHostArgs& args, const ck_tile::stream_config& s) { // The kPadM, kPadN, kPadK & kBlockPerCu should also come from the Codegen part. constexpr bool kPadM = false; @@ -79,17 +79,9 @@ float gemm_calc(const gemm_basic_args& args, const ck_tile::stream_config& s) // Now we only use the BlockGemmASmemBSmemCRegV1DefaultPolicy. using Kernel = ck_tile::GemmKernel; - auto kargs = Kernel::MakeKargs(args.p_a, - args.p_b, - args.p_c, - args.M, - args.N, - args.K, - args.stride_A, - args.stride_B, - args.stride_C); + auto kargs = Kernel::MakeKernelArgs(args); - const dim3 grids = Kernel::GridSize(args.M, args.N, args.kbatch); + const dim3 grids = Kernel::GridSize(args.M, args.N, args.k_batch); constexpr dim3 blocks = Kernel::BlockSize(); if(!Kernel::IsSupportedArgument(kargs)) diff --git a/example/ck_tile/03_gemm/gemm_basic.hpp b/example/ck_tile/03_gemm/gemm_basic.hpp index 23e99bc2a8..58cdaea7d8 100644 --- a/example/ck_tile/03_gemm/gemm_basic.hpp +++ b/example/ck_tile/03_gemm/gemm_basic.hpp @@ -51,20 +51,6 @@ using BDataType = Types::BDataType; using AccDataType = Types::AccDataType; using CDataType = Types::CDataType; -struct gemm_basic_args -{ - const void* p_a; - const void* p_b; - void* p_c; - ck_tile::index_t kbatch; - ck_tile::index_t M; - ck_tile::index_t N; - ck_tile::index_t K; - ck_tile::index_t stride_A; - ck_tile::index_t stride_B; - ck_tile::index_t stride_C; -}; - auto create_args(int argc, char* argv[]) { ck_tile::ArgParser arg_parser; @@ -89,4 +75,4 @@ auto create_args(int argc, char* argv[]) } // host API -float gemm_calc(gemm_basic_args args, const ck_tile::stream_config& s); +float gemm_calc(const ck_tile::GemmHostArgs& args, const ck_tile::stream_config& s); diff --git a/example/ck_tile/03_gemm/run_gemm_example.inc b/example/ck_tile/03_gemm/run_gemm_example.inc index 2b7a967bab..68df389bfc 100644 --- a/example/ck_tile/03_gemm/run_gemm_example.inc +++ b/example/ck_tile/03_gemm/run_gemm_example.inc @@ -16,11 +16,11 @@ float invoke_gemm(ck_tile::DeviceMem& a_m_k_dev_buf, int n_warmup, int n_repeat) { - gemm_basic_args args; - args.p_a = a_m_k_dev_buf.GetDeviceBuffer(); - args.p_b = b_k_n_dev_buf.GetDeviceBuffer(); - args.p_c = c_m_n_dev_buf.GetDeviceBuffer(); - args.kbatch = kbatch; + ck_tile::GemmHostArgs args; + args.a_ptr = a_m_k_dev_buf.GetDeviceBuffer(); + args.b_ptr = b_k_n_dev_buf.GetDeviceBuffer(); + args.c_ptr = c_m_n_dev_buf.GetDeviceBuffer(); + args.k_batch = kbatch; args.M = M; args.N = N; args.K = K; diff --git a/example/ck_tile/16_batched_gemm/batched_gemm.cpp b/example/ck_tile/16_batched_gemm/batched_gemm.cpp index bfdd74126e..9b4ed9a9e7 100644 --- a/example/ck_tile/16_batched_gemm/batched_gemm.cpp +++ b/example/ck_tile/16_batched_gemm/batched_gemm.cpp @@ -16,7 +16,7 @@ #include "batched_gemm.hpp" template -float batched_gemm(const batched_gemm_kargs& args, const ck_tile::stream_config& s) +float batched_gemm(const ck_tile::BatchedGemmHostArgs& args, const ck_tile::stream_config& s) { // The kPadM, kPadN, kPadK & kBlockPerCu should also come from the Codegen part. constexpr bool kPadM = false; @@ -79,9 +79,9 @@ float batched_gemm(const batched_gemm_kargs& args, const ck_tile::stream_config& // Now we only use the BlockGemmASmemBSmemCRegV1DefaultPolicy. using Kernel = ck_tile::BatchedGemmKernel; - auto kargs = Kernel::MakeKargs(args); + auto kargs = Kernel::MakeKernelArgs(args); - const dim3 grids = Kernel::GridSize(args); + const dim3 grids = Kernel::GridSize(args.M, args.N, args.batch_count); constexpr dim3 blocks = Kernel::BlockSize(); if(s.log_level_ > 0) diff --git a/example/ck_tile/16_batched_gemm/batched_gemm.hpp b/example/ck_tile/16_batched_gemm/batched_gemm.hpp index e252c0f673..f0c0c9efba 100644 --- a/example/ck_tile/16_batched_gemm/batched_gemm.hpp +++ b/example/ck_tile/16_batched_gemm/batched_gemm.hpp @@ -29,10 +29,6 @@ using BDataType = Types::BDataType; using AccDataType = Types::AccDataType; using CDataType = Types::CDataType; -struct batched_gemm_kargs : public ck_tile::BatchedGemmHostArgs -{ -}; - auto create_args(int argc, char* argv[]) { ck_tile::ArgParser arg_parser; @@ -60,4 +56,4 @@ auto create_args(int argc, char* argv[]) } // host API -float batched_gemm(batched_gemm_kargs args, const ck_tile::stream_config& s); +float batched_gemm(const ck_tile::BatchedGemmHostArgs& args, const ck_tile::stream_config& s); diff --git a/example/ck_tile/16_batched_gemm/run_batched_gemm_example.inc b/example/ck_tile/16_batched_gemm/run_batched_gemm_example.inc index 8345eef95b..4e7218b5b1 100644 --- a/example/ck_tile/16_batched_gemm/run_batched_gemm_example.inc +++ b/example/ck_tile/16_batched_gemm/run_batched_gemm_example.inc @@ -20,7 +20,7 @@ float invoke_batched_gemm(ck_tile::DeviceMem& a_m_k_dev_buf, int n_warmup, int n_repeat) { - batched_gemm_kargs args; + ck_tile::BatchedGemmHostArgs args; args.a_ptr = a_m_k_dev_buf.GetDeviceBuffer(); args.b_ptr = b_k_n_dev_buf.GetDeviceBuffer(); args.c_ptr = c_m_n_dev_buf.GetDeviceBuffer(); diff --git a/include/ck_tile/ops/gemm/kernel/batched_gemm_kernel.hpp b/include/ck_tile/ops/gemm/kernel/batched_gemm_kernel.hpp index 07b4af5730..07a4cf8fbe 100644 --- a/include/ck_tile/ops/gemm/kernel/batched_gemm_kernel.hpp +++ b/include/ck_tile/ops/gemm/kernel/batched_gemm_kernel.hpp @@ -3,90 +3,93 @@ #pragma once -#include -#include - -#include "ck_tile/core.hpp" -#include "ck_tile/ops/common.hpp" +#include "ck_tile/ops/gemm/kernel/gemm_kernel.hpp" namespace ck_tile { -struct BatchedGemmHostArgs +struct BatchedGemmHostArgs : public ck_tile::GemmHostArgs { - const void* a_ptr; - const void* b_ptr; - void* c_ptr; - index_t M; - index_t N; - index_t K; - index_t stride_A; - index_t stride_B; - index_t stride_C; - index_t batch_stride_A; - index_t batch_stride_B; - index_t batch_stride_C; - index_t batch_count; + CK_TILE_HOST BatchedGemmHostArgs() = default; + CK_TILE_HOST BatchedGemmHostArgs(const void* a_ptr_, + const void* b_ptr_, + void* c_ptr_, + ck_tile::index_t k_batch_, + ck_tile::index_t M_, + ck_tile::index_t N_, + ck_tile::index_t K_, + ck_tile::index_t stride_A_, + ck_tile::index_t stride_B_, + ck_tile::index_t stride_C_, + ck_tile::index_t batch_stride_A_, + ck_tile::index_t batch_stride_B_, + ck_tile::index_t batch_stride_C_, + ck_tile::index_t batch_count_) + : GemmHostArgs( + a_ptr_, b_ptr_, c_ptr_, k_batch_, M_, N_, K_, stride_A_, stride_B_, stride_C_), + batch_stride_A(batch_stride_A_), + batch_stride_B(batch_stride_B_), + batch_stride_C(batch_stride_C_), + batch_count(batch_count_) + { + } + + ck_tile::index_t batch_stride_A; + ck_tile::index_t batch_stride_B; + ck_tile::index_t batch_stride_C; + ck_tile::index_t batch_count; }; template -struct BatchedGemmKernel +struct BatchedGemmKernel : public GemmKernel { - using TilePartitioner = remove_cvref_t; - using GemmPipeline = remove_cvref_t; - using EpiloguePipeline = remove_cvref_t; - using ALayout = remove_cvref_t; - using BLayout = remove_cvref_t; - using CLayout = remove_cvref_t; - static constexpr index_t KernelBlockSize = GemmPipeline::BlockSize; + using Base = GemmKernel; - using ADataType = remove_cvref_t; - using BDataType = remove_cvref_t; - using CDataType = remove_cvref_t; + using GemmKernelArgs = typename Base::GemmKernelArgs; - struct BatchedGemmKargs + using ADataType = typename Base::ADataType; + using BDataType = typename Base::BDataType; + using CDataType = typename Base::CDataType; + + using TilePartitioner = typename Base::TilePartitioner; + using GemmPipeline = typename Base::GemmPipeline; + using EpiloguePipeline = typename Base::EpiloguePipeline; + using ALayout = typename Base::ALayout; + using BLayout = typename Base::BLayout; + using CLayout = typename Base::CLayout; + + struct BatchedGemmKernelArgs : GemmKernelArgs { - const void* a_ptr; - const void* b_ptr; - void* c_ptr; - index_t M; - index_t N; - index_t K; - index_t stride_A; - index_t stride_B; - index_t stride_C; index_t batch_stride_A; index_t batch_stride_B; index_t batch_stride_C; index_t batch_count; }; - using Kargs = BatchedGemmKargs; - using Hargs = BatchedGemmHostArgs; + using KernelArgs = BatchedGemmKernelArgs; - __host__ static constexpr auto GridSize(const Hargs& h) + __host__ static constexpr auto GridSize(index_t M, index_t N, index_t batch_count) { - return TilePartitioner::GridSize(h.M, h.N, h.batch_count); + return TilePartitioner::GridSize(M, N, batch_count); } - __host__ static constexpr auto BlockSize() { return dim3(KernelBlockSize); } + __host__ static constexpr auto BlockSize() { return dim3(Base::KernelBlockSize); } - CK_TILE_HOST static constexpr BatchedGemmKargs MakeKargs(const Hargs& h) + CK_TILE_HOST static constexpr BatchedGemmKernelArgs + MakeKernelArgs(const BatchedGemmHostArgs& hostArgs) { - Kargs k; - k.a_ptr = h.a_ptr; - k.b_ptr = h.b_ptr; - k.c_ptr = h.c_ptr; - k.M = h.M; - k.N = h.N; - k.K = h.K; - k.stride_A = h.stride_A; - k.stride_B = h.stride_B; - k.stride_C = h.stride_C; - k.batch_stride_A = h.batch_stride_A; - k.batch_stride_B = h.batch_stride_B; - k.batch_stride_C = h.batch_stride_C; - k.batch_count = h.batch_count; - return k; + return BatchedGemmKernelArgs{{hostArgs.a_ptr, + hostArgs.b_ptr, + hostArgs.c_ptr, + hostArgs.M, + hostArgs.N, + hostArgs.K, + hostArgs.stride_A, + hostArgs.stride_B, + hostArgs.stride_C}, + hostArgs.batch_stride_A, + hostArgs.batch_stride_B, + hostArgs.batch_stride_C, + hostArgs.batch_count}; } CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSize() @@ -94,7 +97,7 @@ struct BatchedGemmKernel return max(GemmPipeline::GetSmemSize(), EpiloguePipeline::GetSmemSize()); } - CK_TILE_DEVICE void operator()(Kargs kargs) const + CK_TILE_DEVICE void operator()(BatchedGemmKernelArgs kargs) const { const auto [i_m, i_n] = TilePartitioner{}(); const auto i_batch = __builtin_amdgcn_readfirstlane(blockIdx.z); @@ -102,156 +105,17 @@ struct BatchedGemmKernel // options const auto batch_stride_A = __builtin_amdgcn_readfirstlane(kargs.batch_stride_A); const auto batch_offset_A = __builtin_amdgcn_readfirstlane(i_batch * batch_stride_A); - const ADataType* a_start = static_cast(kargs.a_ptr); + const ADataType* a_ptr = static_cast(kargs.a_ptr) + batch_offset_A; const auto batch_stride_B = __builtin_amdgcn_readfirstlane(kargs.batch_stride_B); const auto batch_offset_B = __builtin_amdgcn_readfirstlane(i_batch * batch_stride_B); - const BDataType* b_start = static_cast(kargs.b_ptr); - - // Convert pointers to tensor views - auto a_tensor_view = [&]() { - if constexpr(std::is_same_v) - { - return make_naive_tensor_view( - a_start + batch_offset_A, - make_tuple(kargs.M, kargs.K), - make_tuple(kargs.stride_A, 1), - number{}, - number<1>{}); - } - else - { - return make_naive_tensor_view( - a_start + batch_offset_A, - make_tuple(kargs.M, kargs.K), - make_tuple(1, kargs.stride_A), - number<1>{}, - number<1>{}); - } - }(); - - auto b_tensor_view = [&]() { - if constexpr(std::is_same_v) - { - return make_naive_tensor_view( - b_start + batch_offset_B, - make_tuple(kargs.N, kargs.K), - make_tuple(1, kargs.stride_B), - number<1>{}, - number<1>{}); - } - else - { - return make_naive_tensor_view( - b_start + batch_offset_B, - make_tuple(kargs.N, kargs.K), - make_tuple(kargs.stride_B, 1), - number{}, - number<1>{}); - } - }(); - - auto a_pad_view = [&]() { - if constexpr(std::is_same_v) - { - return pad_tensor_view( - a_tensor_view, - make_tuple(number{}, number{}), - sequence{}); - } - else - { - return pad_tensor_view( - a_tensor_view, - make_tuple(number{}, number{}), - sequence{}); - } - }(); - // clang-format on - - auto a_block_window = make_tile_window( - a_pad_view, - make_tuple(number{}, number{}), - {i_m, 0}); - - auto b_pad_view = [&]() { - if constexpr(std::is_same_v) - { - return pad_tensor_view( - b_tensor_view, - make_tuple(number{}, number{}), - sequence{}); - } - else - { - return pad_tensor_view( - b_tensor_view, - make_tuple(number{}, number{}), - sequence{}); - } - }(); - // clang-format on - - auto b_block_window = make_tile_window( - b_pad_view, - make_tuple(number{}, number{}), - {i_n, 0}); - - // allocate LDS - __shared__ char smem_ptr[GetSmemSize()]; - - const index_t num_loop = TilePartitioner::GetLoopNum(kargs.K); - - // Run GEMM cooperatively by whole wokrgroup. - auto c_block_tile = - GemmPipeline{}.template operator()(a_block_window, b_block_window, num_loop, smem_ptr); + const BDataType* b_ptr = static_cast(kargs.b_ptr) + batch_offset_B; const auto batch_stride_C = __builtin_amdgcn_readfirstlane(kargs.batch_stride_C); const auto batch_offset_C = __builtin_amdgcn_readfirstlane(i_batch * batch_stride_C); - CDataType* c_start = static_cast(kargs.c_ptr); - auto c_tensor_view = [&]() { - if constexpr(std::is_same_v) - { - return make_naive_tensor_view( - c_start + batch_offset_C, - make_tuple(kargs.M, kargs.N), - make_tuple(kargs.stride_C, 1), - number{}, - number<1>{}); - } - else - { - return make_naive_tensor_view( - c_start + batch_offset_C, - make_tuple(kargs.M, kargs.N), - make_tuple(1, kargs.stride_C), - number<1>{}, - number<1>{}); - } - }(); + CDataType* c_ptr = static_cast(kargs.c_ptr) + batch_offset_C; - auto c_pad_view = [&]() { - if constexpr(std::is_same_v) - { - return pad_tensor_view( - c_tensor_view, - make_tuple(number{}, number{}), - sequence{}); - } - else - { - return pad_tensor_view( - c_tensor_view, - make_tuple(number{}, number{}), - sequence{}); - } - }(); - auto c_block_window = make_tile_window( - c_pad_view, - make_tuple(number{}, number{}), - {i_m, i_n}); - - EpiloguePipeline{}(c_block_window, c_block_tile); + this->RunGemm(a_ptr, b_ptr, c_ptr, kargs, i_m, i_n); } }; diff --git a/include/ck_tile/ops/gemm/kernel/gemm_kernel.hpp b/include/ck_tile/ops/gemm/kernel/gemm_kernel.hpp index 763d8cad9c..925648a886 100644 --- a/include/ck_tile/ops/gemm/kernel/gemm_kernel.hpp +++ b/include/ck_tile/ops/gemm/kernel/gemm_kernel.hpp @@ -12,6 +12,50 @@ namespace ck_tile { +struct GemmProblem +{ + CK_TILE_HOST GemmProblem() = default; + CK_TILE_HOST GemmProblem( + index_t M_, index_t N_, index_t K_, index_t stride_A_, index_t stride_B_, index_t stride_C_) + : M(M_), N(N_), K(K_), stride_A(stride_A_), stride_B(stride_B_), stride_C(stride_C_) + { + } + + index_t M; + index_t N; + index_t K; + index_t stride_A; + index_t stride_B; + index_t stride_C; +}; + +struct GemmHostArgs : public GemmProblem +{ + CK_TILE_HOST GemmHostArgs() = default; + CK_TILE_HOST GemmHostArgs(const void* a_ptr_, + const void* b_ptr_, + void* c_ptr_, + index_t k_batch_, + index_t M_, + index_t N_, + index_t K_, + index_t stride_A_, + index_t stride_B_, + index_t stride_C_) + : GemmProblem(M_, N_, K_, stride_A_, stride_B_, stride_C_), + a_ptr(a_ptr_), + b_ptr(b_ptr_), + c_ptr(c_ptr_), + k_batch(k_batch_) + { + } + + const void* a_ptr; + const void* b_ptr; + void* c_ptr; + index_t k_batch; +}; + template struct GemmKernel { @@ -25,9 +69,12 @@ struct GemmKernel using ADataType = remove_cvref_t; using BDataType = remove_cvref_t; - // using CAccDataType = remove_cvref_t; using CDataType = remove_cvref_t; + static constexpr auto I0 = number<0>(); + static constexpr auto I1 = number<1>(); + static constexpr auto I2 = number<2>(); + __host__ static constexpr auto GridSize(index_t M, index_t N, index_t KBatch) { return TilePartitioner::GridSize(M, N, KBatch); @@ -35,7 +82,7 @@ struct GemmKernel __host__ static constexpr auto BlockSize() { return dim3(KernelBlockSize); } - struct GemmCommonKargs + struct GemmKernelArgs { const void* a_ptr; const void* b_ptr; @@ -48,25 +95,37 @@ struct GemmKernel index_t stride_C; }; - CK_TILE_HOST static constexpr GemmCommonKargs MakeKargs(const void* a_ptr, - const void* b_ptr, - void* c_ptr, - index_t M, - index_t N, - index_t K, - index_t stride_A, - index_t stride_B, - index_t stride_C) + CK_TILE_HOST static constexpr GemmKernelArgs MakeKernelArgs(const GemmHostArgs& hostArgs) { - return GemmCommonKargs{a_ptr, b_ptr, c_ptr, M, N, K, stride_A, stride_B, stride_C}; + return GemmKernelArgs{hostArgs.a_ptr, + hostArgs.b_ptr, + hostArgs.c_ptr, + hostArgs.M, + hostArgs.N, + hostArgs.K, + hostArgs.stride_A, + hostArgs.stride_B, + hostArgs.stride_C}; } + // CK_TILE_HOST static constexpr GemmKernelArgs MakeKernelArgs(const void* a_ptr, + // const void* b_ptr, + // void* c_ptr, + // index_t M, + // index_t N, + // index_t K, + // index_t stride_A, + // index_t stride_B, + // index_t stride_C) + // { + // return GemmKernelArgs{a_ptr, b_ptr, c_ptr, M, N, K, stride_A, stride_B, stride_C}; + // } CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSize() { return max(GemmPipeline::GetSmemSize(), EpiloguePipeline::GetSmemSize()); } - CK_TILE_HOST static bool IsSupportedArgument(const GemmCommonKargs& kargs) + CK_TILE_HOST static bool IsSupportedArgument(const GemmKernelArgs& kargs) { if constexpr(std::is_same_v) { @@ -139,18 +198,16 @@ struct GemmKernel return true; } - CK_TILE_DEVICE void operator()(GemmCommonKargs kargs) const + CK_TILE_DEVICE auto MakeGemmTensorViews(const ADataType* a_ptr, + const BDataType* b_ptr, + CDataType* c_ptr, + const GemmKernelArgs& kargs) const { - const auto [i_m, i_n] = TilePartitioner{}(); - // options - const ADataType* a_start = static_cast(kargs.a_ptr); - const BDataType* b_start = static_cast(kargs.b_ptr); - // Convert pointers to tensor views - auto a_tensor_view = [&]() { + const auto& a_tensor_view = [&]() { if constexpr(std::is_same_v) { return make_naive_tensor_view( - a_start, + a_ptr, make_tuple(kargs.M, kargs.K), make_tuple(kargs.stride_A, 1), number{}, @@ -159,7 +216,7 @@ struct GemmKernel else { return make_naive_tensor_view( - a_start, + a_ptr, make_tuple(kargs.M, kargs.K), make_tuple(1, kargs.stride_A), number<1>{}, @@ -167,11 +224,11 @@ struct GemmKernel } }(); - auto b_tensor_view = [&]() { + const auto& b_tensor_view = [&]() { if constexpr(std::is_same_v) { return make_naive_tensor_view( - b_start, + b_ptr, make_tuple(kargs.N, kargs.K), make_tuple(1, kargs.stride_B), number<1>{}, @@ -180,7 +237,7 @@ struct GemmKernel else { return make_naive_tensor_view( - b_start, + b_ptr, make_tuple(kargs.N, kargs.K), make_tuple(kargs.stride_B, 1), number{}, @@ -188,7 +245,35 @@ struct GemmKernel } }(); - auto a_pad_view = [&]() { + const auto& c_tensor_view = [&]() { + if constexpr(std::is_same_v) + { + return make_naive_tensor_view( + c_ptr, + make_tuple(kargs.M, kargs.N), + make_tuple(kargs.stride_C, 1), + number{}, + number<1>{}); + } + else + { + return make_naive_tensor_view( + c_ptr, + make_tuple(kargs.M, kargs.N), + make_tuple(1, kargs.stride_C), + number<1>{}, + number<1>{}); + } + }(); + + return make_tuple(a_tensor_view, b_tensor_view, c_tensor_view); + } + + template + CK_TILE_DEVICE auto MakeGemmPadViews(const TensorView& views) const + { + const auto& a_pad_view = [&]() { + const auto& a_tensor_view = views.at(I0); if constexpr(std::is_same_v) { return pad_tensor_view( @@ -204,14 +289,9 @@ struct GemmKernel sequence{}); } }(); - // clang-format on - auto a_block_window = make_tile_window( - a_pad_view, - make_tuple(number{}, number{}), - {i_m, 0}); - - auto b_pad_view = [&]() { + const auto& b_pad_view = [&]() { + const auto& b_tensor_view = views.at(I1); if constexpr(std::is_same_v) { return pad_tensor_view( @@ -228,43 +308,8 @@ struct GemmKernel } }(); - auto b_block_window = make_tile_window( - b_pad_view, - make_tuple(number{}, number{}), - {i_n, 0}); - - // allocate LDS - __shared__ char smem_ptr[GetSmemSize()]; - - const index_t num_loop = TilePartitioner::GetLoopNum(kargs.K); - - // Run GEMM cooperatively by whole wokrgroup. - auto c_block_tile = - GemmPipeline{}.template operator()(a_block_window, b_block_window, num_loop, smem_ptr); - - CDataType* c_start = static_cast(kargs.c_ptr); - auto c_tensor_view = [&]() { - if constexpr(std::is_same_v) - { - return make_naive_tensor_view( - c_start, - make_tuple(kargs.M, kargs.N), - make_tuple(kargs.stride_C, 1), - number{}, - number<1>{}); - } - else - { - return make_naive_tensor_view( - c_start, - make_tuple(kargs.M, kargs.N), - make_tuple(1, kargs.stride_C), - number<1>{}, - number<1>{}); - } - }(); - - auto c_pad_view = [&]() { + const auto& c_pad_view = [&]() { + const auto& c_tensor_view = views.at(I2); if constexpr(std::is_same_v) { return pad_tensor_view( @@ -280,12 +325,82 @@ struct GemmKernel sequence{}); } }(); - auto CBlockWindow_pad = make_tile_window( + + return make_tuple(a_pad_view, b_pad_view, c_pad_view); + } + + template + CK_TILE_DEVICE auto + MakeGemmTileWindows(const PadView& views, const index_t i_m, const index_t i_n) const + { + const auto& a_pad_view = views.at(I0); + const auto& a_block_window = make_tile_window( + a_pad_view, + make_tuple(number{}, number{}), + {i_m, 0}); + + const auto& b_pad_view = views.at(I1); + const auto& b_block_window = make_tile_window( + b_pad_view, + make_tuple(number{}, number{}), + {i_n, 0}); + + const auto& c_pad_view = views.at(I2); + auto c_block_window = make_tile_window( c_pad_view, make_tuple(number{}, number{}), {i_m, i_n}); - EpiloguePipeline{}(CBlockWindow_pad, c_block_tile); + return make_tuple(a_block_window, b_block_window, c_block_window); + } + + /** + * @brief Runs single GEMM problem cooperatively by whole workgroup. + * + * @param a_ptr input A pointer + * @param b_ptr input B pointer + * @param c_ptr output C pointer + * @param kargs GEMM kernel arguments + * @param block_idx_m The GEMM's output M dimension tile index processed by this workgroup. + * @param block_idx_n The GEMM's output N dimension tile index processed by this workgroup. + */ + CK_TILE_DEVICE void RunGemm(const ADataType* a_ptr, + const BDataType* b_ptr, + CDataType* c_ptr, + const GemmKernelArgs& kargs, + const index_t block_idx_m, + const index_t block_idx_n) const + { + // Create Gemm tensor views, pad views and tile windows + const auto& gemm_tensor_views_tuple = MakeGemmTensorViews(a_ptr, b_ptr, c_ptr, kargs); + const auto& gemm_pad_views = MakeGemmPadViews(gemm_tensor_views_tuple); + auto gemm_tile_windows = MakeGemmTileWindows(gemm_pad_views, block_idx_m, block_idx_n); + + // allocate LDS + __shared__ char smem_ptr[GetSmemSize()]; + + const index_t num_loop = TilePartitioner::GetLoopNum(kargs.K); + + // Run GEMM cooperatively by whole workgroup. + const auto& a_block_window = gemm_tile_windows.at(I0); + const auto& b_block_window = gemm_tile_windows.at(I1); + const auto& c_block_tile = + GemmPipeline{}.template operator()(a_block_window, b_block_window, num_loop, smem_ptr); + + // Run Epilogue Pipeline + auto& c_block_window = gemm_tile_windows.at(I2); + EpiloguePipeline{}(c_block_window, c_block_tile); + } + + CK_TILE_DEVICE void operator()(GemmKernelArgs kargs) const + { + const auto [i_m, i_n] = TilePartitioner{}(); + // options + const ADataType* a_ptr = static_cast(kargs.a_ptr); + const BDataType* b_ptr = static_cast(kargs.b_ptr); + CDataType* c_ptr = static_cast(kargs.c_ptr); + + RunGemm(a_ptr, b_ptr, c_ptr, kargs, i_m, i_n); } }; diff --git a/test/ck_tile/batched_gemm/test_batched_gemm_util.hpp b/test/ck_tile/batched_gemm/test_batched_gemm_util.hpp index 88145b987b..d3f3077870 100644 --- a/test/ck_tile/batched_gemm/test_batched_gemm_util.hpp +++ b/test/ck_tile/batched_gemm/test_batched_gemm_util.hpp @@ -24,12 +24,9 @@ class TestCkTileBatchedGemm : public ::testing::Test using AccDataType = std::tuple_element_t<5, Tuple>; using CDataType = std::tuple_element_t<6, Tuple>; - struct batched_gemm_kargs : public ck_tile::BatchedGemmHostArgs - { - }; - template - void invoke_batched_gemm(const batched_gemm_kargs& args, const ck_tile::stream_config& s) + void invoke_batched_gemm(const ck_tile::BatchedGemmHostArgs& args, + const ck_tile::stream_config& s) { // The kPadM, kPadN, kPadK & kBlockPerCu should also come from the Codegen part. constexpr bool kPadM = false; @@ -94,9 +91,9 @@ class TestCkTileBatchedGemm : public ::testing::Test using Kernel = ck_tile::BatchedGemmKernel; - auto kargs = Kernel::MakeKargs(args); + auto kargs = Kernel::MakeKernelArgs(args); - const dim3 grids = Kernel::GridSize(args); + const dim3 grids = Kernel::GridSize(args.M, args.N, args.batch_count); constexpr dim3 blocks = Kernel::BlockSize(); if(s.log_level_ > 0) @@ -185,21 +182,22 @@ class TestCkTileBatchedGemm : public ::testing::Test c_m_n_dev_buf.SetZero(); c_m_n_dev_result.SetZero(); - batched_gemm_kargs kargs{a_m_k_dev_buf.GetDeviceBuffer(), - b_k_n_dev_buf.GetDeviceBuffer(), - c_m_n_dev_buf.GetDeviceBuffer(), - M, - N, - K, - StrideA, - StrideB, - StrideC, - BatchStrideA, - BatchStrideB, - BatchStrideC, - BatchCount}; + ck_tile::BatchedGemmHostArgs args; + args.a_ptr = a_m_k_dev_buf.GetDeviceBuffer(); + args.b_ptr = b_k_n_dev_buf.GetDeviceBuffer(); + args.c_ptr = c_m_n_dev_buf.GetDeviceBuffer(); + args.M = M; + args.N = N; + args.K = K; + args.stride_A = StrideA; + args.stride_B = StrideB; + args.stride_C = StrideC; + args.batch_stride_A = BatchStrideA; + args.batch_stride_B = BatchStrideB; + args.batch_stride_C = BatchStrideC; + args.batch_count = BatchCount; - invoke_batched_gemm(kargs, + invoke_batched_gemm(args, ck_tile::stream_config{nullptr, false}); std::cout << "Run kernel with M =" << M << " N =" << N << " K =" << K diff --git a/test/ck_tile/gemm/test_gemm_pipeline_util.hpp b/test/ck_tile/gemm/test_gemm_pipeline_util.hpp index a514986024..53ead4d8d6 100644 --- a/test/ck_tile/gemm/test_gemm_pipeline_util.hpp +++ b/test/ck_tile/gemm/test_gemm_pipeline_util.hpp @@ -31,22 +31,8 @@ class TestCkTileGemmPipeline : public ::testing::Test static constexpr auto PipelineType = std::tuple_element_t<8, Tuple>::value; // TODO: expose tile size through test t-param ? - struct gemm_args - { - const void* p_a; - const void* p_b; - void* p_c; - ck_tile::index_t kbatch; - ck_tile::index_t M; - ck_tile::index_t N; - ck_tile::index_t K; - ck_tile::index_t stride_A; - ck_tile::index_t stride_B; - ck_tile::index_t stride_C; - }; - template - void invoke_gemm(const gemm_args& args, const ck_tile::stream_config& s) + void invoke_gemm(const ck_tile::GemmHostArgs& args, const ck_tile::stream_config& s) { // TODO: This should be parameterized in tests constexpr ck_tile::index_t M_Tile = 128; @@ -117,17 +103,9 @@ class TestCkTileGemmPipeline : public ::testing::Test has_hot_loop_v, tail_number_v>>>; using Kernel = ck_tile::GemmKernel; - auto kargs = Kernel::MakeKargs(args.p_a, - args.p_b, - args.p_c, - args.M, - args.N, - args.K, - args.stride_A, - args.stride_B, - args.stride_C); + auto kargs = Kernel::MakeKernelArgs(args); - const dim3 grids = Kernel::GridSize(args.M, args.N, args.kbatch); + const dim3 grids = Kernel::GridSize(args.M, args.N, args.k_batch); constexpr dim3 blocks = Kernel::BlockSize(); if(!Kernel::IsSupportedArgument(kargs)) @@ -319,11 +297,11 @@ class TestCkTileGemmPipeline : public ::testing::Test c_m_n_dev_buf.SetZero(); c_m_n_dev_result.SetZero(); - gemm_args args; - args.p_a = a_m_k_dev_buf.GetDeviceBuffer(); - args.p_b = b_k_n_dev_buf.GetDeviceBuffer(); - args.p_c = c_m_n_dev_buf.GetDeviceBuffer(); - args.kbatch = kbatch; + ck_tile::GemmHostArgs args; + args.a_ptr = a_m_k_dev_buf.GetDeviceBuffer(); + args.b_ptr = b_k_n_dev_buf.GetDeviceBuffer(); + args.c_ptr = c_m_n_dev_buf.GetDeviceBuffer(); + args.k_batch = kbatch; args.M = M; args.N = N; args.K = K; From e758d006a55dd45ee9aae009b5ab554d42736dfb Mon Sep 17 00:00:00 2001 From: Mateusz Ozga <110818320+mozga-amd@users.noreply.github.com> Date: Thu, 19 Dec 2024 17:55:35 +0100 Subject: [PATCH 19/25] Apply Ck-tile argument parser for vectors [I/O] (#1758) * Parser for a vector was added. Additionaly we valid correctnes of numbers * Remove unnecessary comments * Review part 1 * Review part 2 * Add const to variadic lambda * Rename C->K --- .../ck_tile/17_grouped_gemm/grouped_gemm.hpp | 20 +++++--- .../run_grouped_gemm_example.inc | 34 ++++++++------ include/ck_tile/host/arg_parser.hpp | 46 ++++++++++++++++++- 3 files changed, 78 insertions(+), 22 deletions(-) diff --git a/example/ck_tile/17_grouped_gemm/grouped_gemm.hpp b/example/ck_tile/17_grouped_gemm/grouped_gemm.hpp index 94af4711d1..20ba740884 100644 --- a/example/ck_tile/17_grouped_gemm/grouped_gemm.hpp +++ b/example/ck_tile/17_grouped_gemm/grouped_gemm.hpp @@ -34,13 +34,19 @@ using grouped_gemm_kargs = ck_tile::GroupedGemmHostArgs; auto create_args(int argc, char* argv[]) { ck_tile::ArgParser arg_parser; - arg_parser.insert("a_layout", "R", "A tensor data layout - Row by default") - .insert("b_layout", "R", "B tensor data layout - Row by default") - .insert("c_layout", "R", "C tensor data layout - Row by default") - .insert("validate", "1", "0. No validation, 1. Validation on CPU") - .insert("warmup", "10", "number of iterations before benchmark the kernel") - .insert("repeat", "100", "number of iterations to benchmark the kernel") - .insert("group_count", "16", "group count"); + arg_parser.insert("Ms", "", "M dimensions - empty by default.") + .insert("Ns", "", "N dimensions - empty by default.") + .insert("Ks", "", "K dimensions - empty by default.") + .insert("stride_As", "", "Tensor A strides - it is empty by default.") + .insert("stride_Bs", "", "Tensor B strides - it is empty by default.") + .insert("stride_Cs", "", "Tensor C strides - it is empty by default.") + .insert("a_layout", "R", "A tensor data layout - Row by default.") + .insert("b_layout", "R", "B tensor data layout - Row by default.") + .insert("c_layout", "R", "C tensor data layout - Row by default.") + .insert("validate", "1", "0. No validation, 1. Validation on CPU.") + .insert("warmup", "10", "number of iterations before benchmark the kernel.") + .insert("repeat", "100", "number of iterations to benchmark the kernel.") + .insert("group_count", "16", "group count."); bool result = arg_parser.parse(argc, argv); return std::make_tuple(result, arg_parser); diff --git a/example/ck_tile/17_grouped_gemm/run_grouped_gemm_example.inc b/example/ck_tile/17_grouped_gemm/run_grouped_gemm_example.inc index cd5b1c2864..11faa6642c 100644 --- a/example/ck_tile/17_grouped_gemm/run_grouped_gemm_example.inc +++ b/example/ck_tile/17_grouped_gemm/run_grouped_gemm_example.inc @@ -53,26 +53,34 @@ int run_grouped_gemm_example_with_layouts(int argc, return -1; }; + auto valid_input_data = [&](int group_count, const auto&... args) { + return !(args.empty() || ...) && group_count == (args.size() == ...); + }; + const int group_count = arg_parser.get_int("group_count"); const int repeat = arg_parser.get_int("repeat"); const int warmup = arg_parser.get_int("warmup"); - std::vector Ms; - std::vector Ns; - std::vector Ks; - std::vector stride_As; - std::vector stride_Bs; - std::vector stride_Cs; + std::vector Ms = arg_parser.get_int_vec("Ms"); + std::vector Ns = arg_parser.get_int_vec("Ns"); + std::vector Ks = arg_parser.get_int_vec("Ks"); + std::vector stride_As = arg_parser.get_int_vec("stride_As"); + std::vector stride_Bs = arg_parser.get_int_vec("stride_Bs"); + std::vector stride_Cs = arg_parser.get_int_vec("stride_Cs"); - for(int i = 0; i < group_count; i++) + if(!valid_input_data(group_count, Ms, Ns, Ks, stride_As, stride_Bs, stride_Cs)) { - Ms.push_back(256 + 256 * i); - Ns.push_back(128 + 128 * i); - Ks.push_back(128 + 64 * i); + std::cout << "Please check the input data. Default values will be used." << std::endl; + for(int i = 0; i < group_count; i++) + { + Ms.push_back(256 + 256 * i); + Ns.push_back(128 + 128 * i); + Ks.push_back(128 + 64 * i); - stride_As.push_back(Ks[i]); - stride_Bs.push_back(Ks[i]); - stride_Cs.push_back(Ns[i]); + stride_As.push_back(Ks[i]); + stride_Bs.push_back(Ks[i]); + stride_Cs.push_back(Ns[i]); + } } std::vector> a_m_k_tensors; diff --git a/include/ck_tile/host/arg_parser.hpp b/include/ck_tile/host/arg_parser.hpp index 3765156df0..df309f312a 100644 --- a/include/ck_tile/host/arg_parser.hpp +++ b/include/ck_tile/host/arg_parser.hpp @@ -15,11 +15,14 @@ namespace ck_tile { /* - * a host side utility, arg parser for - * -[key0]=[value0] -[key1]=[value1] ... + * a host side utility, arg parser for, either + * -[key0] = [value0, value1, value2] + * or + * -[key0]=[value0] -[key1]=[value1] ... */ class ArgParser { + public: class Arg { @@ -187,6 +190,45 @@ class ArgParser return value; } + std::vector get_string_vec(const std::string& name, + const std::string& delimiter = ",") const + { + if(get_str(name).empty()) + { + return {}; + } + std::string s = get_str(name); + std::vector tokens; + size_t pos = 0; + std::string token; + while((pos = s.find(delimiter)) != std::string::npos) + { + token = s.substr(0, pos); + tokens.push_back(token); + s.erase(0, pos + delimiter.length()); + } + tokens.push_back(s); + + return tokens; + } + + std::vector get_int_vec(const std::string& name, const std::string& delimiter = ",") const + { + if(get_str(name).empty()) + { + return {}; + } + const std::vector args = get_string_vec(name, delimiter); + std::vector tokens; + tokens.reserve(static_cast(args.size())); + for(const std::string& token : args) + { + int value = atoi(token.c_str()); + tokens.push_back(value); + } + return tokens; + } + private: std::unordered_map input_map; std::vector keys; From 2944c508941055a0cf36d5a96092d6c739f53c36 Mon Sep 17 00:00:00 2001 From: Illia Silin <98187287+illsilin@users.noreply.github.com> Date: Thu, 19 Dec 2024 17:24:05 -0800 Subject: [PATCH 20/25] fix profiler_grouped_gemm (#1766) --- profiler/include/profiler/profile_grouped_gemm_impl.hpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/profiler/include/profiler/profile_grouped_gemm_impl.hpp b/profiler/include/profiler/profile_grouped_gemm_impl.hpp index c10cd0ea9f..367e94de11 100644 --- a/profiler/include/profiler/profile_grouped_gemm_impl.hpp +++ b/profiler/include/profiler/profile_grouped_gemm_impl.hpp @@ -77,7 +77,7 @@ bool profile_grouped_gemm_impl(int do_verification, std::vector> c_m_n_host_results; std::vector> c_m_n_device_results; - ComputeDataType max_abs_in_val = 0.f; + double max_abs_in_val = 0.f; for(std::size_t i = 0; i < group_count; i++) { a_m_k.push_back( From 37cdbf4f0ec88ba5064f46c3370633b5950bc7ae Mon Sep 17 00:00:00 2001 From: Po Yen Chen Date: Fri, 20 Dec 2024 14:41:01 +0800 Subject: [PATCH 21/25] [CK_TILE] Add fmha fwd N-Warp S-Shuffle pipeline (fmha fwd splitkv pipeline variant) (#1705) * Add check for zero values * Add static assertions * Remove invalid option '-e' in smoke_test.sh * Use correct path of smoke_test.sh * Avoid zero-sized shared memory array * Add warning comment * Replace expr by integer_divide_ceil() call * Use more readable constant names * Write down assumption as static assertion * Add more diagnostic error messages * Fix wrong BlockWarps when using default pipeline policy * Add more static assertions for A LDS desc * Allow using vector size < 8 for data type fp16/bf16 * Align vector size between DRAM dist & LDS desc * Remove no-longer used func decl * Fix wrong displayed piepline name * Undo policy template changes for tile_example_gemm_basic * Add missing space and make error message stands out * Unify print precision * Add missing include directive * Replace constant 64 by get_warp_size() call * Replace constant 128 by named variable: BankLength * Add kAMBlock/kBNBlock attributes * Allow usig different A/B warp dist for multiple blocks * Add helper function to get warp dist encodings * Add 4x64x4 fp16 warp gemm attribute impl * Complete the A/B warp dist encoding logic * Fix wrong thread mapping for C matrix * Use smaller vector size for small tile * Add static assert to block unsupported warp gemm impl * Extract common code out as helper method * Add 4x64x16 fp16 warp gemm type alias * Add comment to warning developers * Undo WarpGemmAtrributeMfma<> changes * Use more clear static assertion error message * Add trivial wrapper to get warp dstr encodings * Only transpose warp gemm result if it's square * Fix compilation error * Support multi-block warp gemm (on N direction) * Remove duplicated code * Fix output encoding of warp gemm * Fix wrong shape of WarpGemmAtrributeMfmaIterateK<> * Remove unused code * Fix wrong shape of WarpGemmAttributeMfmaImplF16F16F32M4N64K4 * Add type config for bf16_t * Add 4x64x16 bf16 warp gemm * Update WarpGemmAtrributeMfmaIterateKAndTransposedCDistribution * Add 64x4x4 fp16/bf16 warp gemm impl * Add 64x4x16 fp16/bf16 warp gemm * Add static assertion for better error diagnostic * Get Q dram dstr directly form block gemm * Add missing header: fused_moe.hpp * Allow specifying different warp-gemm for gemm0 & gemm1 * Store P matrix into LDS before gemm1 * Fix inconsistant kernel name * Remove constraint on gemm0 & gemm1 block warps * Remove unsupported vector size from checking list * Allow using 4x64x16 warp gemm for gemm0 * Finish policy customization * Finish pipeline modification F# * Use block warps in codegen * Fix wrong rank of m_lds_window origin * Use better distributed tensor * Make P-store earlier * Remove duplicated experssions * Remove unnecessary tile window * Create new files for new splitkv pipeline * Separate old/new pipeline codegen logic * Sync changes form develop * Undo gemm kernel/pipeline changes * Undo gemm example changes * Remove blank lines * Fix typo * Use new warp gemm interface * Fix link error * Fix wrong pipeline tag * Fix more link error * Avoid unnecessary padding * Always use vector load for K * Padding on fastest dimension when necessary * Force padding Q on hdim_q * Set high dimension padding flag to false * Re-format headers * Use warps=<1, 4, 1> for both gemm0 & gemm1 * Fix complilation errors * Remove m/l shuffle logics * Ignore duplicate data when write lse_acc * Use gemm0 block warps as lds tile width * Remove hard-coded numbers * Fix wrong distribution width * Remove unnecessary code * Add s_barrier before writing to LDS * Store Q into LDS before gemm0 * Fix wrong Q tile size * Use simple Q lds descriptor for debuging * Use more realistic Q lds descriptor * Add comment & use better variable name * Make Q lds space not overlapped with others * Remove unnecessary block_tile_reduce_sync() call * Move Q load statements * Move block_sync_lds() right before use * Re-order instructions * Remove necessary lambda expression * Use 8 threads on kMaxSplits direction while doing reduction * Tiny correction for using 8 threads on kMaxSplits direction for combine kernel * Padding num_split direction of o_acc tile window to 4x * Update splitkv combine pipeline design * Add kN1 back to splitkv combine pipeline problem * Fix compilation errors * Add missing template parameter * Fix wrong splitkv combine kernel name * Fix wrong origin * Fix wrong LDS descriptor shape * Fix sync & reduction logics * Remove unnecessary static assertions * Extract tile size computation logics * Make sure we can reuse padding flags in combine kernels * Rename variables * Use OaccDataType in BlockFmhaSplitKVCombinePipelineTileSizes<> * Remove unnecessary static assertion * Fix function name typo * Add constraint on kN1 template parameter * Hide K tile loading latency in earlier iteration * Fix wrong splitkv kernel name * Use s_shuffling to replace p_shuffling which removes the needs of cross-warp reduction * Rename pipeline * Fix wrong pipeline name attribute * Add GetAlignmentQ() for NWarpSShuffle pipeline * Separate Q tile into dram tile & register tile concepts * Remove non-squre warp gemm transpose c type alias * Fallback tile size changes for fmha fwd splitkv * Remove redundant change * Refine naming for the S tile * Use better naming of the S tile dstr (read from lds) * Share Q lds with K lds * Tiny change * Fix with using static_for for passing CI checking --------- Co-authored-by: Qianfeng Zhang --- .../ck_tile/01_fmha/codegen/cpp_symbol_map.py | 1 + .../ck_tile/01_fmha/codegen/ops/fmha_fwd.py | 42 +- .../01_fmha/codegen/ops/fmha_fwd_splitkv.py | 85 +- example/ck_tile/01_fmha/fmha_fwd.hpp | 2 - .../core/arch/amd_buffer_addressing.hpp | 4 +- .../core/tensor/static_distributed_tensor.hpp | 1 + include/ck_tile/ops/fmha.hpp | 2 + .../ops/fmha/kernel/fmha_fwd_kernel.hpp | 6 +- .../fmha_fwd_splitkv_combine_kernel.hpp | 56 +- .../fmha/kernel/fmha_fwd_splitkv_kernel.hpp | 9 +- ...lock_fmha_fwd_splitkv_combine_pipeline.hpp | 89 +- ...plitkv_combine_pipeline_default_policy.hpp | 173 ++-- ...litkv_pipeline_nwarp_sshuffle_qr_ks_vs.hpp | 794 ++++++++++++++++++ ...nwarp_sshuffle_qr_ks_vs_default_policy.hpp | 226 +++++ .../pipeline/block_fmha_pipeline_problem.hpp | 36 +- ...k_fmha_pipeline_qx_ks_vs_custom_policy.hpp | 55 +- .../ops/fmha/pipeline/tile_fmha_shape.hpp | 2 - ...block_gemm_areg_bsmem_creg_one_warp_v1.hpp | 44 +- .../block/block_gemm_areg_bsmem_creg_v2.hpp | 44 +- include/ck_tile/ops/gemm/warp/warp_gemm.hpp | 16 + .../gemm/warp/warp_gemm_attribute_mfma.hpp | 303 ++++++- .../warp/warp_gemm_attribute_mfma_impl.hpp | 271 ++++++ .../ops/gemm/warp/warp_gemm_dispatcher.hpp | 4 + 23 files changed, 1990 insertions(+), 275 deletions(-) create mode 100644 include/ck_tile/ops/fmha/pipeline/block_fmha_fwd_splitkv_pipeline_nwarp_sshuffle_qr_ks_vs.hpp create mode 100644 include/ck_tile/ops/fmha/pipeline/block_fmha_fwd_splitkv_pipeline_nwarp_sshuffle_qr_ks_vs_default_policy.hpp diff --git a/example/ck_tile/01_fmha/codegen/cpp_symbol_map.py b/example/ck_tile/01_fmha/codegen/cpp_symbol_map.py index f6df44a318..332707eafd 100644 --- a/example/ck_tile/01_fmha/codegen/cpp_symbol_map.py +++ b/example/ck_tile/01_fmha/codegen/cpp_symbol_map.py @@ -119,6 +119,7 @@ PIPELINE_MAP = { PIPELINE_ENUM_MAP = { "qr" : "ck_tile::BlockFmhaPipelineEnum::QRKSVS", "qr_async" : "ck_tile::BlockFmhaPipelineEnum::QRKSVS_ASYNC", + "qr_nwarp_sshuffle" : "ck_tile::BlockFmhaPipelineEnum::QRKSVS", } BOOL_MAP = { diff --git a/example/ck_tile/01_fmha/codegen/ops/fmha_fwd.py b/example/ck_tile/01_fmha/codegen/ops/fmha_fwd.py index eca638784d..66814f5a16 100644 --- a/example/ck_tile/01_fmha/codegen/ops/fmha_fwd.py +++ b/example/ck_tile/01_fmha/codegen/ops/fmha_fwd.py @@ -44,13 +44,12 @@ FMHA_FWD_KERNEL_BODY=""" using fmha_dtype_{F_idx} = {F_dtype}; using fmha_block_tile_{F_idx} = ck_tile::sequence<{F_bm0}, {F_bn0}, {F_bk0}, {F_bn1}, {F_bk1}, {F_bk0max}>; -using fmha_warp_tile_{F_idx} = ck_tile::sequence<{F_wm}, {F_wn}, {F_wk}>; using fmha_shape_{F_idx} = ck_tile::TileFmhaShape, - fmha_warp_tile_{F_idx}, + ck_tile::sequence<{F_wm0}, {F_wn0}, {F_wk0}>, ck_tile::sequence<{F_rm1}, {F_rn1}, {F_rk1}>, - fmha_warp_tile_{F_idx}, + ck_tile::sequence<{F_wm1}, {F_wn1}, {F_wk1}>, {F_vlayout}>; using fmha_trait_{F_idx} = ck_tile::TileFmhaTraits<{F_spad}, @@ -306,15 +305,19 @@ class FmhaFwdTileSize: F_rm1 : int # number of warps for gemm1 along q seqlen F_rn1 : int # number of warps for gemm1 along head dim v F_rk1 : int # number of warps for gemm1 along k seqlen (not used) - F_wm : int # warp size along m (warp size) - F_wn : int # warp size along n - F_wk : int # warp size along k + F_wm0 : int # gemm0 warp size along m + F_wn0 : int # gemm0 warp size along n + F_wk0 : int # gemm0 warp size along k + F_wm1 : int # gemm1 warp size along m + F_wn1 : int # gemm1 warp size along n + F_wk1 : int # gemm1 warp size along k F_occupancy : int # occupancy, -1 will let pipeline decide the occupancy, other value will overwrite occupancy @property def name(self) -> str: return f"b{self.F_bm0}x{self.F_bn0}x{self.F_bk0}x{self.F_bn1}x{self.F_bk1}x{self.F_bk0max}" +\ f"_r{self.F_rm0}x{self.F_rn0}x{self.F_rk0}_r{self.F_rm1}x{self.F_rn1}x{self.F_rk1}" +\ - f"_w{self.F_wm}x{self.F_wn}x{self.F_wk}" + ("" if self.F_occupancy == -1 else f"_o{self.F_occupancy}") + f"_w{self.F_wm0}x{self.F_wn0}x{self.F_wk0}_w{self.F_wm1}x{self.F_wn1}x{self.F_wk1}" +\ + ("" if self.F_occupancy == -1 else f"_o{self.F_occupancy}") @dataclass class FmhaFwdKernel: @@ -352,9 +355,12 @@ class FmhaFwdKernel: F_rm1 = self.F_tile.F_rm1, F_rn1 = self.F_tile.F_rn1, F_rk1 = self.F_tile.F_rk1, - F_wm = self.F_tile.F_wm, - F_wn = self.F_tile.F_wn, - F_wk = self.F_tile.F_wk, + F_wm0 = self.F_tile.F_wm0, + F_wn0 = self.F_tile.F_wn0, + F_wk0 = self.F_tile.F_wk0, + F_wm1 = self.F_tile.F_wm1, + F_wn1 = self.F_tile.F_wn1, + F_wk1 = self.F_tile.F_wk1, F_vlayout = LAYOUT_MAP[self.F_pipeline.F_vlayout], F_spad = BOOL_MAP[self.F_pipeline.F_spad], F_skpad = BOOL_MAP[self.F_pipeline.F_skpad], @@ -409,17 +415,17 @@ class FmhaFwdKernel: def get_fmha_fwd_tile_dict_from_dtype(dtype : str) -> Optional[dict]: if dtype == 'fp16' or dtype == 'bf16': return { - '32' : FmhaFwdTileSize(128, 64, 16, 32, 32, 32, 2, 1, 1, 2, 1, 1, 32, 32, 16, -1), - '64' : FmhaFwdTileSize(128, 64, 32, 64, 32, 64, 4, 1, 1, 4, 1, 1, 32, 32, 16, -1), - ## '96' : FmhaFwdTileSize(128, 128, 32, 128, 32, 96, 4, 1, 1, 4, 1, 1, 32, 32, 16, -1), - '128' : FmhaFwdTileSize(128, 128, 32, 128, 32, 128, 4, 1, 1, 4, 1, 1, 32, 32, 16, -1), - '256' : FmhaFwdTileSize(128, 128, 32, 256, 32, 256, 4, 1, 1, 4, 1, 1, 32, 32, 16, -1), + '32' : FmhaFwdTileSize(128, 64, 16, 32, 32, 32, 2, 1, 1, 2, 1, 1, 32, 32, 16, 32, 32, 16, -1), + '64' : FmhaFwdTileSize(128, 64, 32, 64, 32, 64, 4, 1, 1, 4, 1, 1, 32, 32, 16, 32, 32, 16, -1), + ### '96' : FmhaFwdTileSize(128, 128, 32, 128, 32, 96, 4, 1, 1, 4, 1, 1, 32, 32, 16, 32, 32, 16, -1), + '128' : FmhaFwdTileSize(128, 128, 32, 128, 32, 128, 4, 1, 1, 4, 1, 1, 32, 32, 16, 32, 32, 16, -1), + '256' : FmhaFwdTileSize(128, 128, 32, 256, 32, 256, 4, 1, 1, 4, 1, 1, 32, 32, 16, 32, 32, 16, -1), } elif dtype == 'fp8' or dtype == 'bf8': return { - '64' : FmhaFwdTileSize(128, 64, 32, 64, 32, 64, 2, 1, 1, 2, 1, 1, 32, 32, 32, -1), - '128' : FmhaFwdTileSize(128, 128, 32, 128, 32, 128, 4, 1, 1, 4, 1, 1, 32, 32, 32, -1), - '256' : FmhaFwdTileSize(128, 128, 32, 256, 32, 256, 4, 1, 1, 4, 1, 1, 32, 32, 32, -1) + '64' : FmhaFwdTileSize(128, 64, 32, 64, 32, 64, 2, 1, 1, 2, 1, 1, 32, 32, 32, 32, 32, 32, -1), + '128' : FmhaFwdTileSize(128, 128, 32, 128, 32, 128, 4, 1, 1, 4, 1, 1, 32, 32, 32, 32, 32, 32, -1), + '256' : FmhaFwdTileSize(128, 128, 32, 256, 32, 256, 4, 1, 1, 4, 1, 1, 32, 32, 32, 32, 32, 32, -1), } else: return None diff --git a/example/ck_tile/01_fmha/codegen/ops/fmha_fwd_splitkv.py b/example/ck_tile/01_fmha/codegen/ops/fmha_fwd_splitkv.py index e448902cf8..df5b9cecc6 100644 --- a/example/ck_tile/01_fmha/codegen/ops/fmha_fwd_splitkv.py +++ b/example/ck_tile/01_fmha/codegen/ops/fmha_fwd_splitkv.py @@ -39,6 +39,7 @@ K0_MAX_SUBMAX_MAP = { FMHA_FWD_SPLITKV_PIPELINE_MAP = { "qr" : "ck_tile::BlockFmhaFwdSplitKVPipelineQRKSVS", + "qr_nwarp_sshuffle" : "ck_tile::BlockFmhaFwdSplitKVPipelineNWarpSShuffleQRKSVS", "qr_async" : "ck_tile::BlockFmhaFwdSplitKVPipelineQRKSVSAsync", } @@ -50,13 +51,12 @@ namespace {{ template struct kernel_runner {{ using fmha_block_tile = ck_tile::sequence<{F_bm0}, {F_bn0}, {F_bk0}, {F_bn1}, {F_bk1}, {F_bk0max}>; -using fmha_warp_tile = ck_tile::sequence<{F_wm}, {F_wn}, {F_wk}>; using fmha_shape = ck_tile::TileFmhaShape, - fmha_warp_tile, + ck_tile::sequence<{F_wm0}, {F_wn0}, {F_wk0}>, ck_tile::sequence<{F_rm1}, {F_rn1}, {F_rk1}>, - fmha_warp_tile, + ck_tile::sequence<{F_wm1}, {F_wn1}, {F_wk1}>, {F_vlayout}>; using fmha_trait = ck_tile::TileFmhaFwdSplitKVTraits<{F_spad}, @@ -161,9 +161,8 @@ using fmha_pipeline_problem = ck_tile::BlockFmhaSplitKVCombinePipelineProblem< typename FmhaFwdTypeConfig::OaccDataType, typename FmhaFwdTypeConfig::ODataType, {F_hdim}, - {F_bm0}, - {F_bn1}, {F_mode}, + {F_bn1}, fmha_trait>; using fmha_pipeline = ck_tile::BlockFmhaFwdSplitKVCombinePipeline< @@ -177,9 +176,11 @@ using fmha_epilogue = false, false>>; using fmha_kernel = - ck_tile::FmhaFwdSplitKVCombineKernel, - fmha_pipeline, - fmha_epilogue>; + ck_tile::FmhaFwdSplitKVCombineKernel< + ck_tile::FmhaFwdSplitKVCombineTilePartitioner< + fmha_pipeline_problem::kM0, fmha_pipeline_problem::kN1>, + fmha_pipeline, + fmha_epilogue>; static void run(const ck_tile::stream_config& s, fmha_fwd_splitkv_args a) {{ @@ -192,7 +193,7 @@ static void run(const ck_tile::stream_config& s, fmha_fwd_splitkv_args a) }}; }} -using trait_{F_idx} = fmha_fwd_splitkv_combine_traits_<{F_hdim}, {F_dtype}, {F_mode}, {F_bm0}, {F_bn1}, +using trait_{F_idx} = fmha_fwd_splitkv_combine_traits_<{F_hdim}, {F_dtype}, {F_mode}, {F_bn1}, {F_lse}, {F_squant}, {F_spad}, {F_dvpad}>; #include @@ -250,16 +251,25 @@ float fmha_fwd_splitkv(fmha_fwd_splitkv_traits t, fmha_fwd_splitkv_args a, const FMHA_FWD_SPLITKV_API_INNER_DISPATCH=""" {F_if}((t.is_group_mode == {F_mode}) && (t.is_v_rowmajor == {F_vlayout}) && ({F_mask_check}) && (t.bias_type == {F_bias_check}) && (t.do_fp8_static_quant == {F_squant}) && ((a.block_table_ptr != nullptr) == {F_pagedkv}) && ({F_scheck}) && ({F_skcheck}) && ({F_dcheck}) && ({F_dvcheck})) {{ using traits_ = fmha_fwd_splitkv_traits_<{F_hdim}, {F_dtype}, {F_mode}, {F_bm0}, {F_bn0}, {F_bk0}, {F_bn1}, {F_bk1}, {F_bk0max}, {F_vlayout}, {F_pipeline_enum}, {F_mask}, {F_bias}, true, {F_squant}, {F_pagedkv}, {F_spad}, {F_skpad}, {F_dpad}, {F_dvpad}>; + + // get combine kernel tile sizes + using OaccDataType = typename FmhaFwdTypeConfig<{F_dtype}>::OaccDataType; + constexpr ck_tile::index_t kM0 = ck_tile::BlockFmhaSplitKVCombinePipelineTileSizes::kM0; + + // make sure we can reuse the padding flags in combine kernels + static_assert({F_bm0} % kM0 == 0); + static_assert({F_bn1} % 32 == 0); + if (t.has_lse) {{ if constexpr (std::is_same_v<{F_dtype}, ck_tile::fp8_t>) {{ return -1; }} else {{ - using traits2_ = fmha_fwd_splitkv_combine_traits_<{F_hdim}, {F_dtype}, {F_mode}, {F_bm0}/2, {F_bn1}/2, true, {F_squant}, {F_spad}, {F_dvpad}>; + using traits2_ = fmha_fwd_splitkv_combine_traits_<{F_hdim}, {F_dtype}, {F_mode}, /*F_bn1=*/32, true, {F_squant}, {F_spad}, {F_dvpad}>; return fmha_fwd_splitkv_(s, a); }} }} else {{ - using traits2_ = fmha_fwd_splitkv_combine_traits_<{F_hdim}, {F_dtype}, {F_mode}, {F_bm0}/2, {F_bn1}/2, false, {F_squant}, {F_spad}, {F_dvpad}>; + using traits2_ = fmha_fwd_splitkv_combine_traits_<{F_hdim}, {F_dtype}, {F_mode}, /*F_bn1=*/32, false, {F_squant}, {F_spad}, {F_dvpad}>; return fmha_fwd_splitkv_(s, a); }} @@ -302,7 +312,7 @@ class FmhaFwdSplitKVApiTrait: if self.pipeline_tag == 'qr_async': if self.spad == 't' : return 'true' # always support else : return 'true' - elif self.pipeline_tag in ['qr']: + elif self.pipeline_tag in ['qr', 'qr_nwarp_sshuffle']: if self.spad == 't' : return f'true /*a.seqlen_q % {self.bm0} != 0*/' # TODO: order of get_pipelines() matters! (ugly) else : return f'a.seqlen_q % {self.bm0} == 0' else: assert False @@ -313,7 +323,7 @@ class FmhaFwdSplitKVApiTrait: if self.pipeline_tag == 'qr_async': if self.skpad == 't' : return f'a.seqlen_k == 0 || a.seqlen_k % {self.bn0} != 0' else : return f'a.seqlen_k != 0 && a.seqlen_k % {self.bn0} == 0' - elif self.pipeline_tag in ['qr', 'qr_fp8']: + elif self.pipeline_tag in ['qr', 'qr_nwarp_sshuffle']: if self.skpad == 't' : return f'true /*a.seqlen_k % {self.bn0} != 0*/' # TODO: order of get_pipelines() matters! (ugly) else : return f'a.seqlen_k % {self.bn0} == 0' else: assert False @@ -324,7 +334,7 @@ class FmhaFwdSplitKVApiTrait: vec = int((32 * 4) / DTYPE_BITS[self.dtype]) if self.dpad == 't': return f'a.hdim_q % {vec} == 0' else : assert False - elif self.pipeline_tag in ['qr']: + elif self.pipeline_tag in ['qr', 'qr_nwarp_sshuffle']: bk0submax = K0_MAX_SUBMAX_MAP[self.bk0max] if self.dpad == 't': return f'true /*a.hdim_q % {bk0submax} != 0*/' # TODO: order of get_pipelines() matters! (ugly) else : return f'a.hdim_q % {bk0submax} == 0' @@ -336,7 +346,7 @@ class FmhaFwdSplitKVApiTrait: vec = int((32 * 4) / DTYPE_BITS[self.dtype]) if self.dvpad == 't': return f'a.hdim_v % {vec} == 0' else : assert False - elif self.pipeline_tag in ['qr']: + elif self.pipeline_tag in ['qr', 'qr_nwarp_sshuffle']: bk0submax = K0_MAX_SUBMAX_MAP[self.bk0max] if self.dvpad == 't': return f'true /*a.hdim_v % {bk0submax} != 0*/' # TODO: order of get_pipelines() matters! (ugly) else : return f'a.hdim_v % {bk0submax} == 0' @@ -447,12 +457,11 @@ class FmhaFwdSplitKVApiPool: @dataclass class FmhaFwdSplitKVCombineTileSize: - F_bm0 : int # tile size along q seqlen F_bn1 : int # tile size along v head_dim F_occupancy : int # occupancy, -1 will let pipeline decide the occupancy, other value will overwrite occupancy @property def name(self) -> str: - return f"b{self.F_bm0}x{self.F_bn1}" +\ + return f"b{self.F_bn1}" +\ ("" if self.F_occupancy == -1 else f"_o{self.F_occupancy}") @dataclass @@ -485,9 +494,12 @@ class FmhaFwdSplitKVKernel: F_rm1 = self.F_tile.F_rm1, F_rn1 = self.F_tile.F_rn1, F_rk1 = self.F_tile.F_rk1, - F_wm = self.F_tile.F_wm, - F_wn = self.F_tile.F_wn, - F_wk = self.F_tile.F_wk, + F_wm0 = self.F_tile.F_wm0, + F_wn0 = self.F_tile.F_wn0, + F_wk0 = self.F_tile.F_wk0, + F_wm1 = self.F_tile.F_wm1, + F_wn1 = self.F_tile.F_wn1, + F_wk1 = self.F_tile.F_wk1, F_vlayout = LAYOUT_MAP[self.F_pipeline.F_vlayout], F_spad = BOOL_MAP[self.F_pipeline.F_spad], F_skpad = BOOL_MAP[self.F_pipeline.F_skpad], @@ -553,7 +565,6 @@ class FmhaFwdSplitKVCombineKernel: F_idx = self.F_idx, F_hdim = self.F_hdim, F_dtype = FWD_DTYPE_MAP[self.F_dtype], - F_bm0 = self.F_tile.F_bm0, F_bn1 = self.F_tile.F_bn1, F_spad = BOOL_MAP[self.F_pipeline.F_spad], F_dvpad = BOOL_MAP[self.F_pipeline.F_dvpad], @@ -577,17 +588,17 @@ class FmhaFwdSplitKVCombineKernel: def get_fmha_fwd_tile_dict_from_dtype(dtype : str) -> Optional[dict]: if dtype == 'fp16' or dtype == 'bf16': return { - '32' : FmhaFwdTileSize(32, 64, 16, 32, 32, 32, 2, 1, 1, 2, 1, 1, 16, 16, 16, -1), - '64' : FmhaFwdTileSize(64, 64, 32, 64, 32, 64, 4, 1, 1, 4, 1, 1, 16, 16, 16, -1), - ## '96' : FmhaFwdTileSize(64, 128, 32, 128, 32, 96, 4, 1, 1, 4, 1, 1, 16, 16, 16, -1), - '128' : FmhaFwdTileSize(64, 128, 32, 128, 32, 128, 4, 1, 1, 4, 1, 1, 16, 16, 16, -1), - '256' : FmhaFwdTileSize(64, 128, 32, 256, 32, 256, 4, 1, 1, 4, 1, 1, 16, 16, 16, -1), + '32' : FmhaFwdTileSize(32, 64, 16, 32, 32, 32, 2, 1, 1, 2, 1, 1, 16, 16, 16, 16, 16, 16, -1), + '64' : FmhaFwdTileSize(64, 64, 32, 64, 32, 64, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1), + ### '96' : FmhaFwdTileSize(64, 128, 32, 128, 32, 96, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1), + '128' : FmhaFwdTileSize(64, 128, 32, 128, 32, 128, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1), + '256' : FmhaFwdTileSize(64, 128, 32, 256, 32, 256, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1), } elif dtype == 'fp8' or dtype == 'bf8': return { - '64' : FmhaFwdTileSize(128, 64, 32, 64, 32, 64, 2, 1, 1, 2, 1, 1, 32, 32, 32, -1), - '128' : FmhaFwdTileSize(128, 128, 32, 128, 32, 128, 4, 1, 1, 4, 1, 1, 32, 32, 32, -1), - '256' : FmhaFwdTileSize(128, 128, 32, 256, 32, 256, 4, 1, 1, 4, 1, 1, 32, 32, 32, -1) + '64' : FmhaFwdTileSize(128, 64, 32, 64, 32, 64, 2, 1, 1, 2, 1, 1, 32, 32, 32, 32, 32, 32, -1), + '128' : FmhaFwdTileSize(128, 128, 32, 128, 32, 128, 4, 1, 1, 4, 1, 1, 32, 32, 32, 32, 32, 32, -1), + '256' : FmhaFwdTileSize(128, 128, 32, 256, 32, 256, 4, 1, 1, 4, 1, 1, 32, 32, 32, 32, 32, 32, -1), } else: return None @@ -595,17 +606,17 @@ def get_fmha_fwd_tile_dict_from_dtype(dtype : str) -> Optional[dict]: def get_fmha_fwd_splitkv_combine_tile_dict_from_dtype(dtype : str) -> Optional[dict]: if dtype == 'fp16' or dtype == 'bf16': return { - '32' : FmhaFwdSplitKVCombineTileSize(16, 16, -1), - '64' : FmhaFwdSplitKVCombineTileSize(32, 32, -1), - ## '96' : FmhaFwdSplitKVCombineTileSize(32, 64, -1), - '128' : FmhaFwdSplitKVCombineTileSize(32, 64, -1), - '256' : FmhaFwdSplitKVCombineTileSize(32, 128, -1), + '32' : FmhaFwdSplitKVCombineTileSize(32, -1), + '64' : FmhaFwdSplitKVCombineTileSize(32, -1), + ### '96' : FmhaFwdSplitKVCombineTileSize(32, -1), + '128' : FmhaFwdSplitKVCombineTileSize(32, -1), + '256' : FmhaFwdSplitKVCombineTileSize(32, -1), } elif dtype == 'fp8' or dtype == 'bf8': return { - '64' : FmhaFwdSplitKVCombineTileSize(64, 32, -1), - '128' : FmhaFwdSplitKVCombineTileSize(64, 64, -1), - '256' : FmhaFwdSplitKVCombineTileSize(64, 128, -1), + '64' : FmhaFwdSplitKVCombineTileSize(32, -1), + '128' : FmhaFwdSplitKVCombineTileSize(32, -1), + '256' : FmhaFwdSplitKVCombineTileSize(32, -1), } else: return None diff --git a/example/ck_tile/01_fmha/fmha_fwd.hpp b/example/ck_tile/01_fmha/fmha_fwd.hpp index aee54b4758..0e821ed5d9 100644 --- a/example/ck_tile/01_fmha/fmha_fwd.hpp +++ b/example/ck_tile/01_fmha/fmha_fwd.hpp @@ -709,7 +709,6 @@ std::string fmha_fwd_splitkv_get_name_(); template ; static constexpr bool kIsGroupMode = kIsGroupMode_; - static constexpr ck_tile::index_t kM0 = kM0_; static constexpr ck_tile::index_t kN1 = kN1_; static constexpr bool kStoreLse = kStoreLse_; static constexpr bool kDoFp8StaticQuant = kDoFp8StaticQuant_; diff --git a/include/ck_tile/core/arch/amd_buffer_addressing.hpp b/include/ck_tile/core/arch/amd_buffer_addressing.hpp index bebf035e9c..107aae5516 100644 --- a/include/ck_tile/core/arch/amd_buffer_addressing.hpp +++ b/include/ck_tile/core/arch/amd_buffer_addressing.hpp @@ -1303,8 +1303,8 @@ CK_TILE_DEVICE thread_buffer amd_buffer_load_impl(int32x4_t src_wave_buffe static_assert( (std::is_same::value && (N == 1 || N == 2 || N == 4 || N == 8)) || (std::is_same::value && (N == 1 || N == 2 || N == 4 || N == 8 || N == 16)) || - (std::is_same::value && (N == 1 || N == 2 || N == 4 || N == 8 || N == 16)) || - (std::is_same::value && (N == 1 || N == 2 || N == 4 || N == 8 || N == 16)) || + (std::is_same::value && (N == 1 || N == 2 || N == 4 || N == 8)) || + (std::is_same::value && (N == 1 || N == 2 || N == 4 || N == 8)) || (std::is_same::value && (N == 1 || N == 2 || N == 4 || N == 8 || N == 16)) || (std::is_same::value && (N == 1 || N == 2 || N == 4 || N == 8 || N == 16)) || diff --git a/include/ck_tile/core/tensor/static_distributed_tensor.hpp b/include/ck_tile/core/tensor/static_distributed_tensor.hpp index 568d618ec2..8d2f88af39 100644 --- a/include/ck_tile/core/tensor/static_distributed_tensor.hpp +++ b/include/ck_tile/core/tensor/static_distributed_tensor.hpp @@ -29,6 +29,7 @@ struct static_distributed_tensor remove_cvref_t; static constexpr index_t kThreadElementSpaceSize = ThreadTensorDesc{}.get_element_space_size(); + static_assert(0 < kThreadElementSpaceSize, "Make sure tile distribution is valid"); CK_TILE_HOST_DEVICE static constexpr auto get_num_of_dimension() { diff --git a/include/ck_tile/ops/fmha.hpp b/include/ck_tile/ops/fmha.hpp index e106264cef..7a09e4622d 100644 --- a/include/ck_tile/ops/fmha.hpp +++ b/include/ck_tile/ops/fmha.hpp @@ -29,6 +29,8 @@ #include "ck_tile/ops/fmha/pipeline/block_fmha_fwd_appendkv_pipeline_default_policy.hpp" #include "ck_tile/ops/fmha/pipeline/block_fmha_fwd_splitkv_combine_pipeline.hpp" #include "ck_tile/ops/fmha/pipeline/block_fmha_fwd_splitkv_combine_pipeline_default_policy.hpp" +#include "ck_tile/ops/fmha/pipeline/block_fmha_fwd_splitkv_pipeline_nwarp_sshuffle_qr_ks_vs.hpp" +#include "ck_tile/ops/fmha/pipeline/block_fmha_fwd_splitkv_pipeline_nwarp_sshuffle_qr_ks_vs_default_policy.hpp" #include "ck_tile/ops/fmha/pipeline/block_fmha_fwd_splitkv_pipeline_qr_ks_vs.hpp" #include "ck_tile/ops/fmha/pipeline/block_fmha_fwd_splitkv_pipeline_qr_ks_vs_default_policy.hpp" #include "ck_tile/ops/fmha/pipeline/block_fmha_pipeline_enum.hpp" diff --git a/include/ck_tile/ops/fmha/kernel/fmha_fwd_kernel.hpp b/include/ck_tile/ops/fmha/kernel/fmha_fwd_kernel.hpp index 3de433d6a7..90102a6c6f 100644 --- a/include/ck_tile/ops/fmha/kernel/fmha_fwd_kernel.hpp +++ b/include/ck_tile/ops/fmha/kernel/fmha_fwd_kernel.hpp @@ -71,7 +71,8 @@ struct FmhaFwdKernel using bfs = typename FmhaPipeline::BlockFmhaShape; using g0br = typename bfs::Gemm0BlockWarps; using g1br = typename bfs::Gemm1BlockWarps; - using gwt = typename bfs::Gemm0WarpTile; + using g0wt = typename bfs::Gemm0WarpTile; + using g1wt = typename bfs::Gemm1WarpTile; #define _SS_ std::string #define _TS_ std::to_string auto pn = [&] () { @@ -88,7 +89,8 @@ struct FmhaFwdKernel _TS_(bfs::kN1) + "x" + _TS_(bfs::kK1) + "x" + _TS_(bfs::kQKHeaddim) + "_" + "r" + _TS_(g0br::at(ck_tile::number<0>{})) + "x" + _TS_(g0br::at(ck_tile::number<1>{})) + "x" + _TS_(g0br::at(ck_tile::number<2>{})) + "_" + "r" + _TS_(g1br::at(ck_tile::number<0>{})) + "x" + _TS_(g1br::at(ck_tile::number<1>{})) + "x" + _TS_(g1br::at(ck_tile::number<2>{})) + "_" + - "w" + _TS_(gwt::at(ck_tile::number<0>{})) + "x" + _TS_(gwt::at(ck_tile::number<1>{})) + "x" + _TS_(gwt::at(ck_tile::number<2>{})) + "_" + + "w" + _TS_(g0wt::at(ck_tile::number<0>{})) + "x" + _TS_(g0wt::at(ck_tile::number<1>{})) + "x" + _TS_(g0wt::at(ck_tile::number<2>{})) + "_" + + "w" + _TS_(g1wt::at(ck_tile::number<0>{})) + "x" + _TS_(g1wt::at(ck_tile::number<1>{})) + "x" + _TS_(g1wt::at(ck_tile::number<2>{})) + "_" + (kBlockPerCuInput == -1 ? "" : ("o" + _TS_(kBlockPerCu) + "_")) + _SS_(FmhaPipeline::name) + "_" + "v" + (std::is_same_v ? "r" : "c") + (pn.empty() ? "" : "_" + pn) + (BiasEnum == BlockAttentionBiasEnum::NO_BIAS ? _SS_("") : (_SS_("_") + BlockAttentionBiasEnumToStr::name)) + diff --git a/include/ck_tile/ops/fmha/kernel/fmha_fwd_splitkv_combine_kernel.hpp b/include/ck_tile/ops/fmha/kernel/fmha_fwd_splitkv_combine_kernel.hpp index 0bccabdd2f..a0adfdc127 100644 --- a/include/ck_tile/ops/fmha/kernel/fmha_fwd_splitkv_combine_kernel.hpp +++ b/include/ck_tile/ops/fmha/kernel/fmha_fwd_splitkv_combine_kernel.hpp @@ -8,9 +8,11 @@ namespace ck_tile { template struct FmhaFwdSplitKVCombineKernel { - using TilePartitioner = remove_cvref_t; - using FmhaPipeline = remove_cvref_t; - using EpiloguePipeline = remove_cvref_t; + using TilePartitioner = remove_cvref_t; + using FmhaPipeline = remove_cvref_t; + using EpiloguePipeline = remove_cvref_t; + + static constexpr index_t kNumWarps = FmhaPipeline::kNumWarps; static constexpr index_t kBlockSize = FmhaPipeline::kBlockSize; static constexpr index_t kBlockPerCu = FmhaPipeline::kBlockPerCu; static_assert(kBlockPerCu > 0); @@ -50,8 +52,7 @@ struct FmhaFwdSplitKVCombineKernel return _SS_("fmha_fwd_splitkv_combine_d") + _TS_(FmhaPipeline::kHeadDimV) + "_" + _SS_(t2s::name) + "_" + (kIsGroupMode ? "group" : "batch") + "_" - "b" + _TS_(FmhaPipeline::kM0) + "x" + - _TS_(FmhaPipeline::kN1) + "_" + + "b" + _TS_(FmhaPipeline::kN1) + "_" + (kBlockPerCuInput == -1 ? "" : ("o" + _TS_(kBlockPerCu) + "_")) + _SS_(FmhaPipeline::name) + (pn.empty() ? "" : "_" + pn) + @@ -339,37 +340,56 @@ struct FmhaFwdSplitKVCombineKernel number{}, number<1>{}); + // read 4 * (kM0, kN1) o_acc tiles simultaneously by 4 warps const auto o_acc_dram_view = pad_tensor_view( o_acc_dram_naive, - make_tuple(number<1>{}, number{}, number{}), - sequence{}); + make_tuple( + number{}, number{}, number{}), + sequence{}); + const index_t padded_num_splits = + o_acc_dram_view.get_tensor_descriptor().get_lengths()[number<0>{}]; const index_t padded_seqlen_q = o_acc_dram_view.get_tensor_descriptor().get_lengths()[number<1>{}]; const index_t padded_hdim_v = o_acc_dram_view.get_tensor_descriptor().get_lengths()[number<2>{}]; - return transform_tensor_view( + const index_t num_m_tiles = integer_divide_floor(padded_seqlen_q, FmhaPipeline::kM0); + + // transform tensor view by following steps, given shape: (padded_num_splits, + // padded_seqlen_q, padded_hdim_v) + // 1. unmerge to (padded_num_splits, num_m_tiles, kM0, padded_hdim_v) + // 2. transpose to (num_m_tiles, padded_num_splits, kM0, padded_hdim_v) + // 3. merge to (num_m_tiles * padded_num_splits * kM0, padded_hdim_v) + auto transposed = transform_tensor_view( o_acc_dram_view, - make_tuple(make_merge_transform(make_tuple(kargs.num_splits, padded_seqlen_q)), + make_tuple(make_pass_through_transform(padded_num_splits), + make_unmerge_transform(make_tuple(num_m_tiles, FmhaPipeline::kM0)), make_pass_through_transform(padded_hdim_v)), - make_tuple(sequence<0, 1>{}, sequence<2>{}), + make_tuple(sequence<0>{}, sequence<1>{}, sequence<2>{}), + make_tuple(sequence<1>{}, sequence<0, 2>{}, sequence<3>{})); + + return transform_tensor_view( + transposed, + make_tuple(make_merge_transform( + make_tuple(num_m_tiles, padded_num_splits, FmhaPipeline::kM0)), + make_pass_through_transform(padded_hdim_v)), + make_tuple(sequence<0, 1, 2>{}, sequence<3>{}), make_tuple(sequence<0>{}, sequence<1>{})); }(); auto lse_acc_dram_window = make_tile_window( lse_acc_dram, - [&]() { - return make_tuple(number{}, number{}); - }(), + make_tuple(number{}, number{}), {0, i_m0}); + const index_t padded_num_splits = + integer_divide_ceil(kargs.num_splits, kNumWarps) * kNumWarps; + auto o_acc_dram_window = make_tile_window( o_acc_dram, - [&]() { - return make_tuple(number{}, number{}); - }(), - {i_m0, i_n1}); + make_tuple(number{}, number{}), + {i_tile_m * padded_num_splits * FmhaPipeline::kM0, i_n1}); // LSE DRAM window auto lse_dram_window = [&, i_nhead_ = i_nhead]() { @@ -410,7 +430,6 @@ struct FmhaFwdSplitKVCombineKernel identity{}, // lse_element_func composes(saturates{}, scales{kargs.scale_o}), // o_acc_element_func kargs.num_splits, - kargs.seqlen_q, smem_ptr); } else @@ -419,7 +438,6 @@ struct FmhaFwdSplitKVCombineKernel o_acc_dram_window, lse_dram_window, kargs.num_splits, - kargs.seqlen_q, smem_ptr); } }(); diff --git a/include/ck_tile/ops/fmha/kernel/fmha_fwd_splitkv_kernel.hpp b/include/ck_tile/ops/fmha/kernel/fmha_fwd_splitkv_kernel.hpp index f37e676da0..dc17487262 100644 --- a/include/ck_tile/ops/fmha/kernel/fmha_fwd_splitkv_kernel.hpp +++ b/include/ck_tile/ops/fmha/kernel/fmha_fwd_splitkv_kernel.hpp @@ -45,6 +45,7 @@ struct FmhaFwdSplitKVKernel static constexpr bool kPadHeadDimQ = FmhaPipeline::kPadHeadDimQ; static constexpr bool kPadHeadDimV = FmhaPipeline::kPadHeadDimV; static constexpr auto BiasEnum = FmhaPipeline::BiasEnum; + static constexpr bool kStoreLSE = FmhaPipeline::kStoreLSE; static constexpr bool kDoFp8StaticQuant = FmhaPipeline::Problem::kDoFp8StaticQuant; static constexpr bool kIsPagedKV = FmhaPipeline::Problem::kIsPagedKV; @@ -67,7 +68,8 @@ struct FmhaFwdSplitKVKernel using bfs = typename FmhaPipeline::BlockFmhaShape; using g0br = typename bfs::Gemm0BlockWarps; using g1br = typename bfs::Gemm1BlockWarps; - using gwt = typename bfs::Gemm0WarpTile; + using g0wt = typename bfs::Gemm0WarpTile; + using g1wt = typename bfs::Gemm1WarpTile; #define _SS_ std::string #define _TS_ std::to_string auto pn = [&] () { @@ -84,11 +86,12 @@ struct FmhaFwdSplitKVKernel _TS_(bfs::kN1) + "x" + _TS_(bfs::kK1) + "x" + _TS_(bfs::kQKHeaddim) + "_" + "r" + _TS_(g0br::at(ck_tile::number<0>{})) + "x" + _TS_(g0br::at(ck_tile::number<1>{})) + "x" + _TS_(g0br::at(ck_tile::number<2>{})) + "_" + "r" + _TS_(g1br::at(ck_tile::number<0>{})) + "x" + _TS_(g1br::at(ck_tile::number<1>{})) + "x" + _TS_(g1br::at(ck_tile::number<2>{})) + "_" + - "w" + _TS_(gwt::at(ck_tile::number<0>{})) + "x" + _TS_(gwt::at(ck_tile::number<1>{})) + "x" + _TS_(gwt::at(ck_tile::number<2>{})) + "_" + + "w" + _TS_(g0wt::at(ck_tile::number<0>{})) + "x" + _TS_(g0wt::at(ck_tile::number<1>{})) + "x" + _TS_(g0wt::at(ck_tile::number<2>{})) + "_" + + "w" + _TS_(g1wt::at(ck_tile::number<0>{})) + "x" + _TS_(g1wt::at(ck_tile::number<1>{})) + "x" + _TS_(g1wt::at(ck_tile::number<2>{})) + "_" + (kBlockPerCuInput == -1 ? "" : ("o" + _TS_(kBlockPerCu) + "_")) + _SS_(FmhaPipeline::name) + "_" + "v" + (std::is_same_v ? "r" : "c") + (pn.empty() ? "" : "_" + pn) + (BiasEnum == BlockAttentionBiasEnum::NO_BIAS ? _SS_("") : (_SS_("_") + BlockAttentionBiasEnumToStr::name)) + - (kHasMask ? "_" + _SS_(FmhaMask::name) : "") + (kDoFp8StaticQuant ? "_squant" : "") + (kIsPagedKV ? "_pagedkv" : "" ); + (kHasMask ? "_" + _SS_(FmhaMask::name) : "") + (kStoreLSE ? "_lse" : "" ) + (kDoFp8StaticQuant ? "_squant" : "") + (kIsPagedKV ? "_pagedkv" : "" ); #undef _SS_ #undef _TS_ // clang-format on diff --git a/include/ck_tile/ops/fmha/pipeline/block_fmha_fwd_splitkv_combine_pipeline.hpp b/include/ck_tile/ops/fmha/pipeline/block_fmha_fwd_splitkv_combine_pipeline.hpp index 7c49fce99a..7ac86e6d12 100644 --- a/include/ck_tile/ops/fmha/pipeline/block_fmha_fwd_splitkv_combine_pipeline.hpp +++ b/include/ck_tile/ops/fmha/pipeline/block_fmha_fwd_splitkv_combine_pipeline.hpp @@ -53,6 +53,7 @@ struct BlockFmhaFwdSplitKVCombinePipeline using OaccDataType = remove_cvref_t; using ODataType = remove_cvref_t; + static constexpr index_t kNumWarps = Problem::kNumWarps; static constexpr index_t kBlockSize = Problem::kBlockSize; static constexpr index_t kHeadDimV = Problem::kHeadDimV; @@ -117,7 +118,6 @@ struct BlockFmhaFwdSplitKVCombinePipeline const LSEElementFunction& lse_element_func, const OaccElementFunction& o_acc_element_func, index_t num_splits, - index_t seqlen_q, void* smem_ptr) const { // lse_acc tile in LDS @@ -143,11 +143,12 @@ struct BlockFmhaFwdSplitKVCombinePipeline // copy lse_acc tile (shape=[kMaxSplits, kM0]) to LDS (shape=[kMaxSplits, kM0]). auto lse_acc_tile = load_tile(lse_acc_dram_window); store_tile(lse_acc_lds_write_window, lse_acc_tile); - block_sync_lds(); auto lse_accum = make_static_distributed_tensor( Policy::template MakeLSEaccRegTileDistribution()); + __builtin_amdgcn_sched_barrier(0); + block_sync_lds(); // copy LDS (shape=[kM0, kMaxSplits]) to lse_accum (shape=[kM0, kMaxSplits]) // and fill up -INF values outside the [kM0, num_splits] region. { @@ -264,45 +265,93 @@ struct BlockFmhaFwdSplitKVCombinePipeline } }); } - block_sync_lds(); if constexpr(kStoreLSE) { store_tile(lse_dram_window_tmp, tile_elementwise_in(lse_element_func, lse_logsum)); } - auto o_acc_dist = Policy::template MakeOaccDramTileDistribution(); - auto o_acc_dram_window = + auto o_acc_4_dist = Policy::template MakeOacc4DramTileDistribution(); + auto o_acc_4_dram_window = make_tile_window(o_acc_dram_block_window_tmp.get_bottom_tensor_view(), o_acc_dram_block_window_tmp.get_window_lengths(), o_acc_dram_block_window_tmp.get_window_origin(), - o_acc_dist); + o_acc_4_dist); + + // shape=[4 * KM0, kN1] + auto o_acc_4 = make_static_distributed_tensor(o_acc_4_dist); + clear_tile(o_acc_4); + + const index_t padded_num_splits = integer_divide_ceil(num_splits, kNumWarps) * kNumWarps; + + __builtin_amdgcn_sched_barrier(0); + block_sync_lds(); + // each warp handles a [KM0, kN1] tile + for(index_t split_start = 0; split_start < padded_num_splits; split_start += kNumWarps) + { + auto o_tile = load_tile(o_acc_4_dram_window); + const index_t i_split = split_start + get_warp_id(); + const index_t row_start = kM0 * get_warp_id(); + { + constexpr auto spans = decltype(o_acc_4)::get_distributed_spans(); + sweep_tile_span(spans[number<0>{}], [&](auto idx0) { + sweep_tile_span(spans[number<1>{}], [&](auto idx1) { + constexpr auto i_j_idx = make_tuple(idx0, idx1); + const auto x_indices = get_x_indices_from_distributed_indices( + o_acc_4.get_tile_distribution(), i_j_idx); + + const auto row = x_indices.at(number<0>{}); + + const LSEDataType lse_scale = lse_acc_lds(row - row_start, i_split); + o_acc_4(i_j_idx) += lse_scale * o_tile(i_j_idx); + }); + }); + } + + move_tile_window(o_acc_4_dram_window, {kNumWarps * kM0, 0}); + } + + // 4 o_acc tiles in LDS. shape=[4 * kM0, kN1] + OaccDataType* o_acc_4_lds_ptr = static_cast(static_cast( + static_cast(smem_ptr) + Policy::template GetSmemSizeLSEacc())); + + { + auto o_acc_4_lds_window = [&]() { + auto desc = Policy::template MakeOacc4LdsBlockDescriptor(); + auto view = make_tensor_view(o_acc_4_lds_ptr, desc); + return make_tile_window(view, desc.get_lengths(), {0, 0}); + }(); + store_tile(o_acc_4_lds_window, o_acc_4); + } + + auto o_acc_dist = Policy::template MakeOaccDramTileDistribution(); + + auto o_acc_4_lds_window = [&]() { + auto desc = Policy::template MakeOacc4LdsBlockDescriptor(); + auto view = make_tensor_view(o_acc_4_lds_ptr, desc); + return make_tile_window(view, desc.get_lengths(), {0, 0}, o_acc_dist); + }(); + auto o_acc = make_static_distributed_tensor(o_acc_dist); clear_tile(o_acc); - const index_t padded_seqlen_q = integer_divide_ceil(seqlen_q, kM0) * kM0; + __builtin_amdgcn_sched_barrier(0); + block_sync_lds(); + static_for<0, kNumWarps, 1>{}([&](auto) { + auto o_acc_in = load_tile(o_acc_4_lds_window); - for(index_t i_split = 0; i_split < num_splits; ++i_split) - { - auto o_tile = load_tile(o_acc_dram_window); { constexpr auto spans = decltype(o_acc)::get_distributed_spans(); sweep_tile_span(spans[number<0>{}], [&](auto idx0) { sweep_tile_span(spans[number<1>{}], [&](auto idx1) { constexpr auto i_j_idx = make_tuple(idx0, idx1); - const auto x_indices = get_x_indices_from_distributed_indices( - o_acc.get_tile_distribution(), i_j_idx); - - const auto row = x_indices.at(number<0>{}); - - const LSEDataType lse_scale = lse_acc_lds(row, i_split); - o_acc(i_j_idx) += lse_scale * o_tile(i_j_idx); + o_acc(i_j_idx) += o_acc_in(i_j_idx); }); }); } - move_tile_window(o_acc_dram_window, {padded_seqlen_q, 0}); - } + move_tile_window(o_acc_4_lds_window, {kM0, 0}); + }); o_acc = tile_elementwise_in(o_acc_element_func, o_acc); @@ -316,7 +365,6 @@ struct BlockFmhaFwdSplitKVCombinePipeline const OaccDramBlockWindow& o_acc_dram_block_window, LSEDramBlockWindow& lse_dram_block_window, index_t num_splits, - index_t seqlen_q, void* smem_ptr) const { return operator()(lse_acc_dram_block_window, @@ -325,7 +373,6 @@ struct BlockFmhaFwdSplitKVCombinePipeline identity{}, identity{}, num_splits, - seqlen_q, smem_ptr); } }; diff --git a/include/ck_tile/ops/fmha/pipeline/block_fmha_fwd_splitkv_combine_pipeline_default_policy.hpp b/include/ck_tile/ops/fmha/pipeline/block_fmha_fwd_splitkv_combine_pipeline_default_policy.hpp index ebd69c0cf8..2d4abb3888 100644 --- a/include/ck_tile/ops/fmha/pipeline/block_fmha_fwd_splitkv_combine_pipeline_default_policy.hpp +++ b/include/ck_tile/ops/fmha/pipeline/block_fmha_fwd_splitkv_combine_pipeline_default_policy.hpp @@ -10,23 +10,38 @@ namespace ck_tile { struct BlockFmhaFwdSplitKVCombinePipelineDefaultPolicy { - template + template + CK_TILE_HOST_DEVICE static constexpr auto GetMaxNumWarpsForTile() + { + static_assert(NumWarps == 1 || NumWarps == 2 || NumWarps == 4); + + constexpr index_t ElemPerThread = (M * N) / (NumWarps * get_warp_size()); + if constexpr(0 < ElemPerThread) + { + return NumWarps; + } + else + { // try dividing tile by smaller # of warps + return GetMaxNumWarpsForTile(); + } + } + + template CK_TILE_HOST_DEVICE static constexpr auto GetVectorSizeForTile() { - constexpr index_t PixelsPerThread = (M * N) / BlockSize; - static_assert(0 < PixelsPerThread); + constexpr index_t MaxNumWarps = GetMaxNumWarpsForTile(); + + constexpr index_t ElemPerThread = (M * N) / (MaxNumWarps * get_warp_size()); constexpr index_t MaxNPerThread = 16 / sizeof(DataType); - constexpr index_t NPerThread = min(MaxNPerThread, PixelsPerThread); - - return NPerThread; + return min(MaxNPerThread, ElemPerThread); } // alignment for dram lse tile (shape=[kMaxSplits, kM0]) template CK_TILE_HOST_DEVICE static constexpr auto GetAlignmentLSE() { - return GetVectorSizeForTile(); @@ -56,40 +71,54 @@ struct BlockFmhaFwdSplitKVCombinePipelineDefaultPolicy } template - CK_TILE_HOST_DEVICE static constexpr ck_tile::index_t GetSmemSize() + CK_TILE_HOST_DEVICE static constexpr ck_tile::index_t GetSmemSizeLSEacc() { return sizeof(typename Problem::LSEDataType) * MakeLSEaccLdsBlockDescriptor().get_element_space_size(); } + template + CK_TILE_HOST_DEVICE static constexpr ck_tile::index_t GetSmemSizeOacc4() + { + return sizeof(typename Problem::OaccDataType) * + MakeOacc4LdsBlockDescriptor().get_element_space_size(); + } + + template + CK_TILE_HOST_DEVICE static constexpr ck_tile::index_t GetSmemSize() + { + return GetSmemSizeLSEacc() + GetSmemSizeOacc4(); + } + // shape=[kMaxSplits, kM0] template CK_TILE_HOST_DEVICE static constexpr auto MakeLSEaccDramTileDistribution() { using LSEDataType = remove_cvref_t; - constexpr index_t kBlockSize = Problem::kBlockSize; - constexpr index_t kNumWarps = Problem::kNumWarps; - - constexpr index_t kNPerBlock = Problem::kM0; constexpr index_t kMPerBlock = Problem::kMaxSplits; + constexpr index_t kNPerBlock = Problem::kM0; + + constexpr index_t MaxNumWarps = + GetMaxNumWarpsForTile(); + constexpr index_t Replicate = Problem::kNumWarps / MaxNumWarps; constexpr index_t NPerThread = - GetVectorSizeForTile(); + GetVectorSizeForTile(); constexpr index_t NThreads = kNPerBlock / NPerThread; constexpr index_t MThreadsPerWarp = get_warp_size() / NThreads; - constexpr index_t MPerThread = kMPerBlock / (kNumWarps * MThreadsPerWarp); + constexpr index_t MPerThread = kMPerBlock / (MaxNumWarps * MThreadsPerWarp); + static_assert(MPerThread * MaxNumWarps * MThreadsPerWarp == kMPerBlock); static_assert(NThreads * NPerThread == kNPerBlock); - static_assert(MPerThread * kNumWarps * MThreadsPerWarp == kMPerBlock); return make_static_tile_distribution( - tile_distribution_encoding, - tuple, + tile_distribution_encoding, + tuple, sequence>, - tuple, sequence<1, 2>>, - tuple, sequence<2, 0>>, + tuple, sequence<1, 2>>, + tuple, sequence<2, 0>>, sequence<1, 2>, sequence<0, 1>>{}); } @@ -100,17 +129,15 @@ struct BlockFmhaFwdSplitKVCombinePipelineDefaultPolicy { using LSEDataType = remove_cvref_t; - constexpr index_t kBlockSize = Problem::kBlockSize; - - constexpr index_t kMPerBlock = Problem::kMaxSplits; - constexpr index_t kNPerBlock = Problem::kM0; + constexpr index_t kMPerBlock = Problem::kM0; + constexpr index_t kNPerBlock = Problem::kMaxSplits; constexpr index_t NPack = - GetVectorSizeForTile(); + GetVectorSizeForTile(); constexpr auto lse_acc_lds_block_desc_0 = make_naive_tensor_descriptor( make_tuple(number{}, number{}, number{}), make_tuple(number<(kMPerBlock + 1) * NPack>{}, number{}, number<1>{}), - number<8>{}, + number{}, number<1>{}); constexpr auto lse_acc_lds_block_desc = transform_tensor_descriptor( @@ -129,17 +156,15 @@ struct BlockFmhaFwdSplitKVCombinePipelineDefaultPolicy { using LSEDataType = remove_cvref_t; - constexpr index_t kBlockSize = Problem::kBlockSize; - - constexpr index_t kMPerBlock = Problem::kMaxSplits; - constexpr index_t kNPerBlock = Problem::kM0; + constexpr index_t kMPerBlock = Problem::kM0; + constexpr index_t kNPerBlock = Problem::kMaxSplits; constexpr index_t NPack = - GetVectorSizeForTile(); + GetVectorSizeForTile(); constexpr auto lse_acc_lds_block_desc_0 = make_naive_tensor_descriptor( make_tuple(number{}, number{}, number{}), make_tuple(number<(kMPerBlock + 1) * NPack>{}, number{}, number<1>{}), - number<8>{}, + number{}, number<1>{}); constexpr auto lse_acc_t_lds_block_desc = transform_tensor_descriptor( @@ -152,33 +177,86 @@ struct BlockFmhaFwdSplitKVCombinePipelineDefaultPolicy return lse_acc_t_lds_block_desc; } + // 3d + padding, shape=[4 * kM0, kN1] + template + CK_TILE_HOST_DEVICE static constexpr auto MakeOacc4LdsBlockDescriptor() + { + using LSEDataType = remove_cvref_t; + + constexpr index_t kMPerBlock = 4 * Problem::kM0; + constexpr index_t kNPerBlock = Problem::kN1; + constexpr index_t NPack = + GetVectorSizeForTile(); + + constexpr auto o_acc_lds_block_desc_0 = make_naive_tensor_descriptor( + make_tuple(number{}, number{}, number{}), + make_tuple(number<(kMPerBlock + 1) * NPack>{}, number{}, number<1>{}), + number<8>{}, + number<1>{}); + + constexpr auto o_acc_t_lds_block_desc = transform_tensor_descriptor( + o_acc_lds_block_desc_0, + make_tuple(make_pass_through_transform(kMPerBlock), + make_merge_transform(make_tuple(kNPerBlock / NPack, NPack))), + make_tuple(sequence<1>{}, sequence<0, 2>{}), + make_tuple(sequence<1>{}, sequence<0>{})); + + return o_acc_t_lds_block_desc; + } + + // shape=[kM0, kMaxSplits] template CK_TILE_HOST_DEVICE static constexpr auto MakeLSEaccRegTileDistribution() { - constexpr index_t kBlockSize = Problem::kBlockSize; - - constexpr index_t kNPerBlock = Problem::kMaxSplits; constexpr index_t kMPerBlock = Problem::kM0; + constexpr index_t kNPerBlock = Problem::kMaxSplits; - constexpr index_t NThreads = 4; - constexpr index_t NPerThread = kNPerBlock / NThreads; + constexpr index_t MaxNThreads = 8; + constexpr index_t NThreads = min(kNPerBlock, MaxNThreads); + constexpr index_t NPerThread = kNPerBlock / NThreads; - constexpr index_t MThreads = kBlockSize / NThreads; - constexpr index_t MPerThread = kMPerBlock / MThreads; - constexpr index_t MWarps = kBlockSize / get_warp_size(); + constexpr index_t MPerThread = 1; + constexpr index_t MThreads = kMPerBlock / MPerThread; constexpr index_t MThreadPerWarp = get_warp_size() / NThreads; + constexpr index_t MaxNumWarps = (MThreads * NThreads) / get_warp_size(); + constexpr index_t Replicate = Problem::kNumWarps / MaxNumWarps; + + static_assert(MaxNumWarps * MThreadPerWarp * MPerThread == kMPerBlock); static_assert(NThreads * NPerThread == kNPerBlock); - static_assert(MWarps * MThreadPerWarp * MPerThread == kMPerBlock); return make_static_tile_distribution( - tile_distribution_encoding< - sequence<1>, - tuple, sequence>, - tuple, sequence<2, 1>>, - tuple, sequence<0, 1>>, - sequence<1, 2>, - sequence<2, 1>>{}); + tile_distribution_encoding, + tuple, + sequence>, + tuple, sequence<2, 1>>, + tuple, sequence<0, 1>>, + sequence<1, 2>, + sequence<2, 1>>{}); + } + + // similar to MakeOaccDramTileDistribution(), but duplicate same 1-warp encoding 4 times on M + // direction + template + CK_TILE_HOST_DEVICE static constexpr auto MakeOacc4DramTileDistribution() + { + constexpr index_t kMPerBlock = Problem::kM0; // real kMPerBlock we want is (4 * kM0) + constexpr index_t kNPerBlock = Problem::kN1; + static_assert(get_warp_size() <= kMPerBlock * kNPerBlock); + + constexpr index_t M1 = 1; // compose encoding base on 1 warp + constexpr index_t M2 = min(kMPerBlock / M1, get_warp_size()); + constexpr index_t N0 = get_warp_size() / M2; + constexpr index_t N1 = kNPerBlock / N0; + constexpr index_t M0 = kMPerBlock / (M2 * M1); + + return make_static_tile_distribution( + tile_distribution_encoding, + tuple, sequence>, + tuple, sequence<1, 2>>, + tuple, sequence<3, 0>>, + sequence<1, 2>, + sequence<1, 1>>{}); } template @@ -187,6 +265,7 @@ struct BlockFmhaFwdSplitKVCombinePipelineDefaultPolicy constexpr index_t kBlockSize = Problem::kBlockSize; constexpr index_t kMPerBlock = Problem::kM0; constexpr index_t kNPerBlock = Problem::kN1; + static_assert(kBlockSize <= kMPerBlock * kNPerBlock); constexpr index_t M1 = kBlockSize / get_warp_size(); constexpr index_t M2 = min(kMPerBlock / M1, get_warp_size()); diff --git a/include/ck_tile/ops/fmha/pipeline/block_fmha_fwd_splitkv_pipeline_nwarp_sshuffle_qr_ks_vs.hpp b/include/ck_tile/ops/fmha/pipeline/block_fmha_fwd_splitkv_pipeline_nwarp_sshuffle_qr_ks_vs.hpp new file mode 100644 index 0000000000..3726cd433c --- /dev/null +++ b/include/ck_tile/ops/fmha/pipeline/block_fmha_fwd_splitkv_pipeline_nwarp_sshuffle_qr_ks_vs.hpp @@ -0,0 +1,794 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved. + +#pragma once + +#include "ck_tile/core.hpp" +#include "ck_tile/ops/fmha/block/block_attention_bias_enum.hpp" +#include "ck_tile/ops/fmha/pipeline/block_fmha_fwd_splitkv_pipeline_nwarp_sshuffle_qr_ks_vs_default_policy.hpp" +#include "ck_tile/ops/reduce/block/block_reduce.hpp" + +namespace ck_tile { + +// This pipeline is qkv all located in LDS +template +struct BlockFmhaFwdSplitKVPipelineNWarpSShuffleQRKSVS +{ + using Problem = remove_cvref_t; + using Policy = remove_cvref_t; + using QDataType = remove_cvref_t; + using KDataType = remove_cvref_t; + using VDataType = remove_cvref_t; + using SaccDataType = remove_cvref_t; + using SMPLComputeDataType = remove_cvref_t; + using BiasDataType = remove_cvref_t; + using LSEDataType = remove_cvref_t; + using PDataType = remove_cvref_t; + using OaccDataType = remove_cvref_t; + using ODataType = remove_cvref_t; + using FmhaMask = remove_cvref_t; + + using BlockFmhaShape = remove_cvref_t; + using VLayout = remove_cvref_t; + static constexpr bool kQLoadOnce = true; // if q_tile load whole block length (hdim) at once + static_assert(kQLoadOnce == Policy::QLoadOnce); + + static constexpr index_t kBlockSize = Problem::kBlockSize; + + static constexpr index_t kM0 = BlockFmhaShape::kM0; + static constexpr index_t kN0 = BlockFmhaShape::kN0; + static constexpr index_t kK0 = BlockFmhaShape::kK0; + static constexpr index_t kN1 = BlockFmhaShape::kN1; + static constexpr index_t kK1 = BlockFmhaShape::kK1; + static constexpr index_t kQKHeaddim = BlockFmhaShape::kQKHeaddim; + static constexpr index_t kSubQKHeaddim = BlockFmhaShape::kSubQKHeaddim; + + static constexpr bool kIsGroupMode = Problem::kIsGroupMode; + static constexpr bool kPadSeqLenQ = Problem::kPadSeqLenQ; + static constexpr bool kPadSeqLenK = Problem::kPadSeqLenK; + static constexpr bool kPadHeadDimQ = Problem::kPadHeadDimQ; + static constexpr bool kPadHeadDimV = Problem::kPadHeadDimV; + static constexpr auto BiasEnum = Problem::BiasEnum; + static constexpr bool kStoreLSE = Problem::kStoreLSE; + static constexpr bool kIsPagedKV = Problem::kIsPagedKV; + static constexpr bool kHasUnevenSplits = Problem::kHasUnevenSplits; + + // last dimension vector length used to create tensor view(and decide buffer_load vector length) + // ... together with tensor distribution. tensor dist should able to overwrite this + static constexpr index_t kAlignmentQ = + kPadHeadDimQ ? 1 : Policy::template GetAlignmentQ(); + static constexpr index_t kAlignmentK = + kPadHeadDimQ ? 1 : Policy::template GetAlignmentK(); + static constexpr index_t kAlignmentV = []() { + if constexpr(std::is_same_v) + return kPadHeadDimV ? 1 : Policy::template GetAlignmentV(); + else + return kPadSeqLenK ? 1 : Policy::template GetAlignmentV(); + }(); + + static constexpr index_t kAlignmentOacc = + kPadHeadDimV ? 1 : Policy::template GetAlignmentOacc(); + + static constexpr index_t kAlignmentBias = + kPadSeqLenK ? 1 : Policy::template GetAlignmentBias(); + + static constexpr index_t kBlockPerCu = []() { + if constexpr(Problem::kBlockPerCu != -1) + return Problem::kBlockPerCu; + else + { + if constexpr(kQKHeaddim <= 32) + { + return 2; + } + else if constexpr(kQKHeaddim <= 64) + { + return 3; + } + else if constexpr(kQKHeaddim <= 128) + { + if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS) + return 1; + else + return 2; + } + else if constexpr(kQKHeaddim <= 256) + { + return 1; + } + } + }(); + + static constexpr const char* name = "qr_nwarp_sshuffle"; + + CK_TILE_HOST_DEVICE static constexpr ck_tile::index_t GetSmemSize() + { + return Policy::template GetSmemSize(); + } + + template + CK_TILE_HOST_DEVICE auto + operator()(const QDramBlockWindowTmp& q_dram_block_window_tmp, // M0*K0 tile + const QElementFunction& q_element_func, + const KDramBlockWindowLengths& k_dram_block_window_lengths, // N0*K0 tile + const KPageBlockNavigator& k_page_block_navigator, + const KElementFunction& k_element_func, + const VDramBlockWindowLengths& v_dram_block_window_lengths, // N1*K1 tile + const VPageBlockNavigator& v_page_block_navigator, + const VElementFunction& v_element_func, + const BiasDramBlockWindowTmp& bias_dram_block_window_tmp, // M0*N0 tile + const BiasElementFunction& bias_element_func, + LSEaccDramBlockWindowTmp& lse_acc_dram_window_tmp, // M0*1 tile + const LSEaccElementFunction& lse_acc_element_func, + const SAccElementFunction& s_acc_element_func, + const PComputeElementFunction& p_compute_element_func, + const OAccElementFunction& o_acc_element_func, + index_t num_splits, + index_t i_split, + FmhaMask mask, + PositionEncoding position_encoding, + float scale_s, + index_t kv_l2p_offset, // logical-to-physical offset of seqlen_k coordinate + void* smem_ptr) const + { + static_assert( + std::is_same_v> && + std::is_same_v> && + std::is_same_v>, + "wrong!"); + + static_assert(kM0 == QDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] && + kSubQKHeaddim == + QDramBlockWindowTmp{}.get_window_lengths()[number<1>{}] && + kN0 == KDramBlockWindowLengths{}[number<0>{}] && + kK0 == KDramBlockWindowLengths{}[number<1>{}] && + kN1 == VDramBlockWindowLengths{}[number<0>{}] && + kK1 == VDramBlockWindowLengths{}[number<1>{}] && + kM0 == BiasDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] && + kN0 == BiasDramBlockWindowTmp{}.get_window_lengths()[number<1>{}], + "wrong!"); + // Q tile in LDS + QDataType* q_lds_ptr = + static_cast(static_cast(static_cast(smem_ptr))); + auto q_lds = make_tensor_view( + q_lds_ptr, Policy::template MakeQLdsBlockDescriptor()); + + // K tile in LDS + KDataType* k_lds_ptr = + static_cast(static_cast(static_cast(smem_ptr))); + auto k_lds = make_tensor_view( + k_lds_ptr, Policy::template MakeKLdsBlockDescriptor()); + auto k_lds_window = + make_tile_window(k_lds, make_tuple(number{}, number{}), {0, 0}); + + // V tile in LDS + auto v_lds = make_tensor_view( + reinterpret_cast(static_cast(smem_ptr) + + max(Policy::template GetSmemSizeQ(), + Policy::template GetSmemSizeK())), + Policy::template MakeVLdsBlockDescriptor()); + auto v_lds_window = make_tile_window( + v_lds, Policy::template MakeVLdsBlockDescriptor().get_lengths(), {0, 0}); + + // S tile in LDS + auto s_lds = make_tensor_view( + reinterpret_cast(reinterpret_cast(smem_ptr) + + max(Policy::template GetSmemSizeQ(), + Policy::template GetSmemSizeK())), + Policy::template MakeSLdsBlockDescriptor()); + auto s_write_lds_window = make_tile_window( + s_lds, Policy::template MakeSLdsBlockDescriptor().get_lengths(), {0, 0}); + auto s_read_lds_window = + make_tile_window(s_lds, + Policy::template MakeSLdsBlockDescriptor().get_lengths(), + {0, 0}, + Policy::template MakeSRegTileDistribution()); + + // Block GEMM + constexpr auto gemm_0 = Policy::template GetQKBlockGemm(); + constexpr auto gemm_1 = Policy::template GetKVBlockGemm(); + + auto q_dram_window = + make_tile_window(q_dram_block_window_tmp.get_bottom_tensor_view(), + q_dram_block_window_tmp.get_window_lengths(), + q_dram_block_window_tmp.get_window_origin(), + Policy::template MakeQDramTileDistribution()); + + // load Q here, will store Q into LDS to maximize throughput + auto origin_q = load_tile(q_dram_window); + + using SaccBlockTileType = decltype(gemm_0.MakeCBlockTile()); + auto s_acc = SaccBlockTileType{}; + + // 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; }; + + using OaccBlockTileType = decltype(gemm_1.MakeCBlockTile()); + + auto o_acc = OaccBlockTileType{}; + + // infer Sacc, S, P, M, L, Oacc type + using SBlockTileType = decltype(cast_tile(o_acc)); + + using MLBlockTileType = decltype(block_tile_reduce( + SBlockTileType{}, sequence<1>{}, f_max, SMPLComputeDataType{0})); + + // init M, L + auto m = MLBlockTileType{}; + auto l = MLBlockTileType{}; + + clear_tile(o_acc); + set_tile(m, -numeric::infinity()); + clear_tile(l); + + const auto q_origin = q_dram_window.get_window_origin(); + const auto [logical_seqlen_k_start, logical_seqlen_k_end] = mask.GetTileRangeAlongX( + q_origin.at(number<0>{}), number{}, number{}, num_splits, i_split); + + // check early exit if no work to do + if constexpr(FmhaMask::IsMasking || kPadSeqLenK || kHasUnevenSplits) + { + const index_t logical_num_total_loop = + integer_divide_ceil(logical_seqlen_k_end - logical_seqlen_k_start, kN0); + if(logical_num_total_loop <= 0) + { + if constexpr(kStoreLSE) + { + auto lse_acc = + make_static_distributed_tensor(m.get_tile_distribution()); + + set_tile(lse_acc, -numeric::infinity()); + + if(get_thread_local_1d_id() < kM0) + { + store_tile(lse_acc_dram_window_tmp, + tile_elementwise_in(lse_acc_element_func, lse_acc)); + } + } + + // Note: here occ are all cleard, return it + // Note: q loaded but no fence, ignore it. + return o_acc; + } + } + + const index_t physical_seqlen_k_start = logical_seqlen_k_start + kv_l2p_offset; + const index_t physical_seqlen_k_end = logical_seqlen_k_end + kv_l2p_offset; + // make sure the first tile is completely located in page-block (page-block size should be + // divisible by kN0) + // relationship between each *_start variables: aligned_physical_seqlen_k_start <= + // physical_seqlen_k_start, logical_seqlen_k_start <= physical_seqlen_k_start + const index_t aligned_physical_seqlen_k_start = + [&, physical_seqlen_k_start_ = physical_seqlen_k_start] { + if constexpr(kIsPagedKV) + { + return kN0 * integer_divide_floor(physical_seqlen_k_start_, kN0); + } + else + { + return physical_seqlen_k_start_; + } + }(); + const index_t num_total_loop = + integer_divide_ceil(physical_seqlen_k_end - aligned_physical_seqlen_k_start, kN0); + + auto [i_page_block_k, k_dram_block_window] = k_page_block_navigator.make_tile_window( + k_dram_block_window_lengths, {aligned_physical_seqlen_k_start, 0}); + + const auto bias_origin = bias_dram_block_window_tmp.get_window_origin(); + auto bias_dram_window = + make_tile_window(bias_dram_block_window_tmp.get_bottom_tensor_view(), + bias_dram_block_window_tmp.get_window_lengths(), + {bias_origin.at(number<0>{}), + logical_seqlen_k_start - (physical_seqlen_k_start - + aligned_physical_seqlen_k_start)}, // M/N + Policy::template MakeBiasDramTileDistribution()); + + auto [i_page_block_v, v_dram_window] = v_page_block_navigator.make_tile_window( + v_dram_block_window_lengths, + {0, aligned_physical_seqlen_k_start}, // TODO: hdim split? + Policy::template MakeVDramTileDistribution()); + + // store Q into LDS + __builtin_amdgcn_sched_barrier(0); + auto q_lds_window_for_store = make_tile_window( + q_lds, Policy::template MakeQLdsBlockDescriptor().get_lengths(), {0, 0}); + + store_tile(q_lds_window_for_store, origin_q); + __builtin_amdgcn_sched_barrier(0); + + // load Q from LDS + __builtin_amdgcn_sched_barrier(0); + auto q_lds_window_for_load = make_tile_window( + q_lds, + Policy::template MakeQLdsBlockDescriptor().get_lengths(), + {0, 0}, + Policy::template MakeQRegTileDistribution()); + block_sync_lds(); + auto q = load_tile(q_lds_window_for_load); + __builtin_amdgcn_sched_barrier(0); + auto q_tile = tile_elementwise_in(q_element_func, q); + + // prefetch K tile + index_t i_total_loops = 0; + constexpr index_t k0_loops = kQKHeaddim / kK0; + constexpr index_t k1_loops = kN0 / kK1; + + static_assert(2 <= k0_loops); + static_assert(1 <= k1_loops); + + auto k_dram_window = make_tile_window( + k_dram_block_window, + Policy::template MakeKDramTileDistribution()); // K DRAM tile window for + + // load the first tile of the first iteration and store to LDS + auto k_block_tile = load_tile(k_dram_window); + // moving k_dram_window is an in-page-block operation, so there is + // no need to invoke k_page_block_navigator.move_tile_window() here. + move_tile_window(k_dram_window, {0, kK0}); + store_tile(k_lds_window, tile_elementwise_in(k_element_func, k_block_tile)); + + do + { + // STAGE 1, QK gemm + clear_tile(s_acc); // initialize C + + // load the second tile of the first iteration + k_block_tile = load_tile(k_dram_window); + + if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS) + { + __builtin_amdgcn_sched_barrier( + 0); // prevent from messing up the order of global loads + } + const auto bias_tile = load_tile(bias_dram_window); // load bias tile + if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS) + { + __builtin_amdgcn_sched_barrier( + 0); // prevent from messing up the order of global loads + } + + if constexpr(k0_loops > 2) + { + static_for<0, k0_loops - 2, 1>{}([&](auto i_k0) { + block_sync_lds(); + gemm_0(s_acc, + get_slice_tile(q_tile, + sequence<0, i_k0 * kK0>{}, + sequence{}), + k_lds_window); + block_sync_lds(); + move_tile_window(k_dram_window, {0, kK0}); + + store_tile( + k_lds_window, + tile_elementwise_in(k_element_func, k_block_tile)); // LDS write i + 1 + k_block_tile = load_tile(k_dram_window); // global read i + 2 + }); + } + + const auto v_prefetch = load_tile(v_dram_window); // prefetch load v tile + { // tail + block_sync_lds(); + gemm_0(s_acc, + get_slice_tile(q_tile, + sequence<0, (k0_loops - 2) * kK0>{}, + sequence{}), + k_lds_window); + block_sync_lds(); + + store_tile(k_lds_window, tile_elementwise_in(k_element_func, k_block_tile)); + block_sync_lds(); + + gemm_0(s_acc, + get_slice_tile(q_tile, + sequence<0, (k0_loops - 1) * kK0>{}, + sequence{}), + k_lds_window); + } + + // STAGE 2, scale_s, add bias, mask, softmax + if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS) + { + s_acc = tile_elementwise_in(s_acc_element_func, s_acc); + tile_elementwise_inout([&scale_s](auto& x) { x = x * scale_s; }, s_acc); + tile_elementwise_inout( + [&](auto& x, const auto& y) { +#if !CK_TILE_FMHA_FWD_FAST_EXP2 + x += type_convert(bias_element_func(y)); +#else + x += log2e_v * + type_convert(bias_element_func(y)); +#endif + }, + s_acc, + bias_tile); + } + else if constexpr(BiasEnum == BlockAttentionBiasEnum::ALIBI) + { + const auto k_origin = k_page_block_navigator.to_global_window_origin( + i_page_block_k, k_dram_block_window.get_window_origin()); + constexpr auto s_spans = decltype(s_acc)::get_distributed_spans(); + s_acc = tile_elementwise_in(s_acc_element_func, s_acc); + sweep_tile_span(s_spans[number<0>{}], [&](auto idx0) { + sweep_tile_span(s_spans[number<1>{}], [&](auto idx1) { + const auto tile_idx = get_x_indices_from_distributed_indices( + s_acc.get_tile_distribution(), make_tuple(idx0, idx1)); + + const auto row = q_origin.at(number<0>{}) + tile_idx.at(number<0>{}); + const auto col = k_origin.at(number<0>{}) + tile_idx.at(number<1>{}); + constexpr auto i_j_idx = make_tuple(idx0, idx1); + + s_acc(i_j_idx) *= scale_s; + // position_encoding accept only logical coordinates, do conversion here + position_encoding.update(s_acc(i_j_idx), row, col - kv_l2p_offset); + }); + }); + } + else + { + s_acc = tile_elementwise_in(s_acc_element_func, s_acc); +#if !CK_TILE_FMHA_FWD_FAST_EXP2 + tile_elementwise_inout([&scale_s](auto& x) { x = x * scale_s; }, s_acc); +#endif + } + move_tile_window(bias_dram_window, {0, kN0}); + + /// TODO: only check in first/last iteration without increasing code size + if constexpr(kHasUnevenSplits) + { + const auto k_origin = k_page_block_navigator.to_global_window_origin( + i_page_block_k, k_dram_block_window.get_window_origin()); + set_tile_if( + s_acc, + -numeric::infinity(), + [&, + physical_seqlen_k_start_ = physical_seqlen_k_start, + physical_seqlen_k_end_ = physical_seqlen_k_end](auto tile_idx) { + const auto col = k_origin.at(number<0>{}) + tile_idx.at(number<1>{}); + if constexpr(kIsPagedKV) + { + return col < physical_seqlen_k_start_ || physical_seqlen_k_end_ <= col; + } + else + { + return physical_seqlen_k_end_ <= col; + } + }); + } + + if constexpr(kPadSeqLenK || FmhaMask::IsMasking) + { + const auto k_origin = k_page_block_navigator.to_global_window_origin( + i_page_block_k, k_dram_block_window.get_window_origin()); + // mask accept only logical coordinates, do conversion here + bool need_perpixel_check = mask.IsEdgeTile(q_origin.at(number<0>{}), + k_origin.at(number<0>{}) - kv_l2p_offset, + number{}, + number{}); + if(need_perpixel_check) + { + set_tile_if( + s_acc, -numeric::infinity(), [&](auto tile_idx) { + const auto row = q_origin.at(number<0>{}) + tile_idx.at(number<0>{}); + const auto col = k_origin.at(number<0>{}) + tile_idx.at(number<1>{}); + return mask.IsOutOfBound(row, col - kv_l2p_offset); + }); + } + } + + __builtin_amdgcn_sched_barrier(0); + + // load the first tile for next iteration + if(i_total_loops < num_total_loop - 1) + { + // move K tile windows + i_page_block_k = k_page_block_navigator.move_tile_window( + i_page_block_k, k_dram_block_window, {kN0, 0}); + + k_dram_window = make_tile_window( + k_dram_block_window, + Policy::template MakeKDramTileDistribution()); // K DRAM tile window + + // laod the first tile of the first iteration and store to LDS + k_block_tile = load_tile(k_dram_window); + } + + __builtin_amdgcn_sched_barrier(0); + + const auto s = cast_tile(s_acc); // S{j} + + // shuffle through LDS so that the tile layout is consistent with required by Gemm1 + store_tile(s_write_lds_window, s); + block_sync_lds(); + auto s_new = load_tile(s_read_lds_window); + + auto m_local = block_tile_reduce( + s_new, + sequence<1>{}, + f_max, + -numeric::infinity()); // m_local = rowmax(S{j}) + block_tile_reduce_sync(m_local, f_max, bool_constant{}); + + const auto m_old = m; // m{j-1} + tile_elementwise_inout( + [](auto& e0, auto e1, auto e2) { e0 = max(e1, e2); }, m, m_old, m_local); // m{j} + + auto p_compute = make_static_distributed_tensor( + s_new.get_tile_distribution()); // Pcompute{j} + + static const auto get_validated_m = [](SMPLComputeDataType raw_m) { + /// NOTICE: bias might be materialized mask including -inf values, need + /// consideration + if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS || + FmhaMask::IsMasking) + { + return raw_m == -numeric::infinity() + ? type_convert(0.f) + : raw_m; + } + else + { + return raw_m; + } + }; + + constexpr auto p_spans = decltype(p_compute)::get_distributed_spans(); + sweep_tile_span(p_spans[number<0>{}], [&](auto idx0) { + constexpr auto i_idx = make_tuple(idx0); +#if CK_TILE_FMHA_FWD_FAST_EXP2 + auto row_max = scale_s * get_validated_m(m[i_idx]); +#endif + sweep_tile_span(p_spans[number<1>{}], [&](auto idx1) { + constexpr auto i_j_idx = make_tuple(idx0, idx1); +#if CK_TILE_FMHA_FWD_FAST_EXP2 + if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS || + BiasEnum == BlockAttentionBiasEnum::ALIBI) + { + p_compute(i_j_idx) = exp2(s_new[i_j_idx] - get_validated_m(m[i_idx])); + } + else + { + p_compute(i_j_idx) = exp2(scale_s * s_new[i_j_idx] - row_max); + } +#else + p_compute(i_j_idx) = exp(s_new[i_j_idx] - get_validated_m(m[i_idx])); +#endif + }); + }); + + auto rowsum_p = block_tile_reduce( + p_compute, sequence<1>{}, f_sum, SMPLComputeDataType{0}); // rowsum(Pcompute{j}) + + block_tile_reduce_sync(rowsum_p, f_sum, bool_constant{}); + + const auto p = + cast_tile(tile_elementwise_in(p_compute_element_func, p_compute)); + + // l{j}, Oacc{j} + constexpr auto o_spans = decltype(o_acc)::get_distributed_spans(); + sweep_tile_span(o_spans[number<0>{}], [&](auto idx0) { + constexpr auto i_idx = make_tuple(idx0); +#if CK_TILE_FMHA_FWD_FAST_EXP2 + const auto tmp = [&]() { + if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS || + BiasEnum == BlockAttentionBiasEnum::ALIBI) + { + return exp2(m_old[i_idx] - get_validated_m(m[i_idx])); + } + else + { + auto row_max = scale_s * get_validated_m(m[i_idx]); + return exp2(scale_s * m_old[i_idx] - row_max); + } + }(); +#else + const auto tmp = exp(m_old[i_idx] - get_validated_m(m[i_idx])); +#endif + l(i_idx) = tmp * l[i_idx] + rowsum_p[i_idx]; + sweep_tile_span(o_spans[number<1>{}], [&](auto idx1) { + constexpr auto i_j_idx = make_tuple(idx0, idx1); + // FIXME: this use different equation from FA v2 paper, + // but produce correc result. + // Is the equation wrong? + o_acc(i_j_idx) *= tmp; + }); + }); + + block_sync_lds(); + if constexpr(std::is_same_v) + { + auto v_shuffle_tmp = make_static_distributed_tensor( + Policy::template MakeShuffledVRegBlockDescriptor()); + shuffle_tile(v_shuffle_tmp, v_prefetch); + store_tile( + v_lds_window, + tile_elementwise_in(v_element_func, v_shuffle_tmp)); // store the prefetch + } + else + { + store_tile(v_lds_window, + tile_elementwise_in(v_element_func, v_prefetch)); // store the prefetch + } + i_page_block_v = + v_page_block_navigator.move_tile_window(i_page_block_v, v_dram_window, {0, kK1}); + + // STAGE 3, KV gemm + if constexpr(k1_loops > 1) + { + static_for<0, k1_loops - 1, 1>{}([&, + &i_page_block_v_ = i_page_block_v, + &v_dram_window_ = v_dram_window](auto i_k1) { + const auto v = load_tile(v_dram_window_); // load next v + block_sync_lds(); + + gemm_1(o_acc, + get_slice_tile( + p, sequence<0, i_k1 * kK1>{}, sequence{}), + v_lds_window); + block_sync_lds(); + + if constexpr(std::is_same_v) + { + auto v_shuffle_tmp = make_static_distributed_tensor( + Policy::template MakeShuffledVRegBlockDescriptor()); + shuffle_tile(v_shuffle_tmp, v); + store_tile(v_lds_window, + tile_elementwise_in(v_element_func, + v_shuffle_tmp)); // store the prefetch + } + else + { + store_tile(v_lds_window, + tile_elementwise_in(v_element_func, v)); // store next v + } + i_page_block_v_ = v_page_block_navigator.move_tile_window( + i_page_block_v_, v_dram_window_, {0, kK1}); + }); + } + + // tail + { + block_sync_lds(); + gemm_1(o_acc, + get_slice_tile( + p, sequence<0, (k1_loops - 1) * kK1>{}, sequence{}), + v_lds_window); + block_sync_lds(); + } + + __builtin_amdgcn_sched_barrier(0); + + // load the first tile for next iteration + if(i_total_loops < num_total_loop - 1) + { + // store the first tile for next iteration to LDS + // moving k_dram_window is an in-page-block operation, so there is + // no need to invoke k_page_block_navigator.move_tile_window() here. + move_tile_window(k_dram_window, {0, kK0}); + store_tile(k_lds_window, tile_elementwise_in(k_element_func, k_block_tile)); + } + } while(++i_total_loops < num_total_loop); + + if constexpr(kStoreLSE) + { + // store lse acc + auto lse_acc = make_static_distributed_tensor(m.get_tile_distribution()); + + constexpr auto lse_acc_spans = decltype(lse_acc)::get_distributed_spans(); + sweep_tile_span(lse_acc_spans[number<0>{}], [&, m_ = m, l_ = l](auto idx0) { + constexpr auto i_idx = make_tuple(idx0); +#if CK_TILE_FMHA_FWD_FAST_EXP2 + if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS || + BiasEnum == BlockAttentionBiasEnum::ALIBI) + { + lse_acc(i_idx) = m_[i_idx] / C_LOG2E + log(l_[i_idx]); + } + else + { + lse_acc(i_idx) = m_[i_idx] * scale_s / C_LOG2E + log(l_[i_idx]); + } +#else + lse_acc(i_idx) = m_[i_idx] + log(l_[i_idx]); +#endif + }); + + if(get_thread_local_1d_id() < kM0) + { + store_tile(lse_acc_dram_window_tmp, + tile_elementwise_in(lse_acc_element_func, lse_acc)); + } + } + + // finally, O + constexpr auto o_spans = decltype(o_acc)::get_distributed_spans(); + + sweep_tile_span(o_spans[number<0>{}], [&](auto idx0) { + constexpr auto i_idx = make_tuple(idx0); + const auto tmp = [&]() { + if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS || + FmhaMask::IsMasking) + { + return l[i_idx] == 0.f ? 0.f : 1 / l[i_idx]; + } + else + return 1 / l[i_idx]; + }(); + sweep_tile_span(o_spans[number<1>{}], [&](auto idx1) { + constexpr auto i_j_idx = make_tuple(idx0, idx1); + o_acc(i_j_idx) *= tmp; + }); + }); + + o_acc = tile_elementwise_in(o_acc_element_func, o_acc); + + return o_acc; + } + + template + CK_TILE_HOST_DEVICE auto + operator()(const QDramBlockWindowTmp& q_dram_block_window_tmp, // M0*K0 tile + const KDramBlockWindowLengths& k_dram_block_window_lengths, // N0*K0 tile + const KPageBlockNavigator& k_page_block_navigator, + const VDramBlockWindowLengths& v_dram_block_window_lengths, // N1*K1 tile + const VPageBlockNavigator& v_page_block_navigator, + const BiasDramBlockWindowTmp& bias_dram_block_window_tmp, // M0*N0 tile + LSEaccDramBlockWindowTmp& lse_acc_dram_block_window_tmp, // M0*1 tile + index_t num_splits, + index_t i_split, + FmhaMask mask, + PositionEncoding position_encoding, + float scale_s, + index_t kv_l2p_offset, // logical-to-physical offset of seqlen_k coordinate + void* smem_ptr) const + { + return operator()(q_dram_block_window_tmp, + identity{}, + k_dram_block_window_lengths, + k_page_block_navigator, + identity{}, + v_dram_block_window_lengths, + v_page_block_navigator, + identity{}, + bias_dram_block_window_tmp, + identity{}, + lse_acc_dram_block_window_tmp, + identity{}, + identity{}, + identity{}, + identity{}, + num_splits, + i_split, + mask, + position_encoding, + scale_s, + kv_l2p_offset, + smem_ptr); + } +}; + +} // namespace ck_tile diff --git a/include/ck_tile/ops/fmha/pipeline/block_fmha_fwd_splitkv_pipeline_nwarp_sshuffle_qr_ks_vs_default_policy.hpp b/include/ck_tile/ops/fmha/pipeline/block_fmha_fwd_splitkv_pipeline_nwarp_sshuffle_qr_ks_vs_default_policy.hpp new file mode 100644 index 0000000000..74d755ef39 --- /dev/null +++ b/include/ck_tile/ops/fmha/pipeline/block_fmha_fwd_splitkv_pipeline_nwarp_sshuffle_qr_ks_vs_default_policy.hpp @@ -0,0 +1,226 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved. + +#pragma once + +#include "ck_tile/core.hpp" +#include "ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qx_ks_vs_custom_policy.hpp" +#include "ck_tile/ops/gemm/block/block_gemm_asmem_bsmem_creg_v1_custom_policy.hpp" +#include "ck_tile/ops/gemm/block/block_gemm_asmem_bsmem_creg_v1.hpp" + +namespace ck_tile { + +// This pipeline is qkv all located in LDS +struct BlockFmhaFwdSplitKVPipelineNWarpSShuffleQRKSVSDefaultPolicy + : BlockFmhaPipelineQXKSVSCustomPolicy +{ + using BasePolicy = BlockFmhaPipelineQXKSVSCustomPolicy; + + template + CK_TILE_HOST_DEVICE static constexpr auto GetAlignmentQ() + { + constexpr index_t kBlockSize = Problem::kBlockSize; + constexpr index_t kMPerBlock = Problem::BlockFmhaShape::kM0; + constexpr index_t kKPerBlock = Problem::BlockFmhaShape::kSubQKHeaddim; + + constexpr index_t MaxVectorSize = 16 / sizeof(typename Problem::QDataType); + + // this should align with MakeQDramTileDistribution() + constexpr index_t ElemPerThread = (kMPerBlock * kKPerBlock) / kBlockSize; + static_assert(0 < ElemPerThread); + return min(ElemPerThread, MaxVectorSize); + } + + template + CK_TILE_HOST_DEVICE static constexpr auto GetAlignmentOacc() + { + using OaccDataType = remove_cvref_t; + + return static_cast(16 / sizeof(OaccDataType)); + } + + template + CK_TILE_HOST_DEVICE static constexpr auto MakeQDramTileDistribution() + { + constexpr index_t kBlockSize = Problem::kBlockSize; + constexpr index_t kMPerBlock = Problem::BlockFmhaShape::kM0; + constexpr index_t kKPerBlock = Problem::BlockFmhaShape::kSubQKHeaddim; + + constexpr index_t MaxVectorSize = 16 / sizeof(typename Problem::QDataType); + + constexpr index_t ElemPerThread = (kMPerBlock * kKPerBlock) / kBlockSize; + static_assert(0 < ElemPerThread); + constexpr index_t kMaxVecLoad = min(ElemPerThread, MaxVectorSize); + + constexpr index_t KPerThread = kMaxVecLoad; + constexpr index_t KThreads = kKPerBlock / KPerThread; + constexpr index_t MThreadPerWarp = get_warp_size() / KThreads; + constexpr index_t NumWarps = kBlockSize / get_warp_size(); + constexpr index_t MPerThread = kMPerBlock / (MThreadPerWarp * NumWarps); + + 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 + CK_TILE_HOST_DEVICE static constexpr auto MakeQRegTileDistribution() + { + return BasePolicy::template MakeQDramTileDistribution(); + } + + template + CK_TILE_HOST_DEVICE static constexpr auto GetSmemKPackQ() + { + // TODO: this is for 3d layout + using QDataType = remove_cvref_t; + return static_cast(16 / sizeof(QDataType)); + } + + template + CK_TILE_HOST_DEVICE static constexpr auto MakeQLdsBlockDescriptor() + { + constexpr index_t kBlockSize = Problem::kBlockSize; + constexpr index_t kMPerBlock = Problem::BlockFmhaShape::kM0; + constexpr index_t kKPerBlock = Problem::BlockFmhaShape::kSubQKHeaddim; + + constexpr index_t ElemPerThread = (kMPerBlock * kKPerBlock) / kBlockSize; + static_assert(0 < ElemPerThread); + constexpr index_t kKPack = min(ElemPerThread, GetSmemKPackQ()); + + constexpr auto q_lds_block_desc_0 = make_naive_tensor_descriptor( + make_tuple(number{}, number{}, number{}), + make_tuple(number<(kMPerBlock + 1) * kKPack>{}, number{}, number<1>{}), + number{}, + number<1>{}); + + constexpr auto q_lds_block_desc = transform_tensor_descriptor( + q_lds_block_desc_0, + make_tuple( + make_pass_through_transform(number{}), + make_merge_transform(make_tuple(number{}, number{}))), + make_tuple(sequence<1>{}, sequence<0, 2>{}), + make_tuple(sequence<0>{}, sequence<1>{})); + + return q_lds_block_desc; + } + + template + CK_TILE_HOST_DEVICE static constexpr auto GetSmemNPackS() + { + using SDataType = remove_cvref_t; + return static_cast(16 / sizeof(SDataType)); + } + + template + CK_TILE_HOST_DEVICE static constexpr auto MakeSLdsBlockDescriptor() + { + constexpr index_t kMPerBlock = Problem::BlockFmhaShape::kM0; + constexpr index_t kNPerBlock = Problem::BlockFmhaShape::kN0; + constexpr index_t kNPack = GetSmemNPackS(); + + constexpr auto s_lds_block_desc_0 = make_naive_tensor_descriptor( + make_tuple(number{}, number{}, number{}), + make_tuple(number<(kMPerBlock + 1) * kNPack>{}, number{}, number<1>{}), + number{}, + number<1>{}); + + constexpr auto s_lds_block_desc = transform_tensor_descriptor( + s_lds_block_desc_0, + make_tuple( + make_pass_through_transform(number{}), + make_merge_transform(make_tuple(number{}, number{}))), + make_tuple(sequence<1>{}, sequence<0, 2>{}), + make_tuple(sequence<0>{}, sequence<1>{})); + + return s_lds_block_desc; + } + + template + CK_TILE_HOST_DEVICE static constexpr auto MakeSRegTileDistribution() + { + using BlockGemm = remove_cvref_t())>; + + constexpr auto config = BlockGemm::Policy::template GetWarpGemmMWarpNWarp(); + using WG = remove_cvref_t())>; + constexpr index_t MWarp = config.template at<1>(); + constexpr index_t NWarp = config.template at<2>(); + + static_assert(MWarp == 1, "Check failed!"); + + constexpr index_t kMPerBlock = Problem::BlockFmhaShape::kM0; + constexpr index_t kKPerBlock = Problem::BlockFmhaShape::kK1; + constexpr index_t kTileK = Problem::BlockFmhaShape::kN0; + + // K2 is equal to Impl::kABKPerLane * kKIterPerWarpGemm + constexpr index_t K3 = WG::kK / WG::WarpGemmAttribute::Impl::kABKLane; + constexpr index_t K2 = WG::WarpGemmAttribute::Impl::kABKLane; + constexpr index_t K1 = kKPerBlock / (K2 * K3); + constexpr index_t K0 = kTileK / kKPerBlock; + constexpr index_t M2 = WG::WarpGemmAttribute::Impl::kAMLane; + constexpr index_t M1 = MWarp; + constexpr index_t M0 = kMPerBlock / (M2 * M1); + + constexpr auto s2_block_dstr_encoding = + tile_distribution_encoding, + tuple, sequence>, + tuple, sequence<2, 1>>, + tuple, sequence<2, 2>>, + sequence<1, 2, 2, 2>, + sequence<0, 0, 1, 3>>{}; + + constexpr auto s2_block_dstr = make_static_tile_distribution(s2_block_dstr_encoding); + + return s2_block_dstr; + } + + template + CK_TILE_HOST_DEVICE static constexpr ck_tile::index_t GetSmemSizeQ() + { + return MakeQLdsBlockDescriptor().get_element_space_size() * + sizeof(typename Problem::QDataType); + } + + template + CK_TILE_HOST_DEVICE static constexpr ck_tile::index_t GetSmemSizeK() + { + return MakeKLdsBlockDescriptor().get_element_space_size() * + sizeof(typename Problem::KDataType); + } + + template + CK_TILE_HOST_DEVICE static constexpr ck_tile::index_t GetSmemSizeV() + { + return MakeVLdsBlockDescriptor().get_element_space_size() * + sizeof(typename Problem::VDataType); + } + + template + CK_TILE_HOST_DEVICE static constexpr ck_tile::index_t GetSmemSizeS() + { + return MakeSLdsBlockDescriptor().get_element_space_size() * + sizeof(typename Problem::SaccDataType); + } + + template + CK_TILE_HOST_DEVICE static constexpr ck_tile::index_t GetSmemSize() + { + return max(GetSmemSizeQ(), GetSmemSizeK()) + + max(GetSmemSizeV(), GetSmemSizeS()); + } +}; + +} // namespace ck_tile diff --git a/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_problem.hpp b/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_problem.hpp index d9da2f088c..1fe19faaf9 100644 --- a/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_problem.hpp +++ b/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_problem.hpp @@ -106,28 +106,43 @@ struct BlockFmhaFwdSplitKVPipelineProblem static constexpr index_t kBlockPerCu = Traits::kBlockPerCu; }; +// extract tile size attributes to remove dependency on traits +template +struct BlockFmhaSplitKVCombinePipelineTileSizes +{ + static constexpr index_t MaxVectorSize = 16 / sizeof(OaccDataType_); + + static constexpr index_t kN1 = kN1_; + static constexpr index_t NThreads = kN1 / MaxVectorSize; + static constexpr index_t kM0 = get_warp_size() / NThreads; // MThreadPerWarp +}; + template struct BlockFmhaSplitKVCombinePipelineProblem + : BlockFmhaSplitKVCombinePipelineTileSizes { + using BaseType = BlockFmhaSplitKVCombinePipelineTileSizes; + using LSEDataType = remove_cvref_t; using OaccDataType = remove_cvref_t; using ODataType = remove_cvref_t; using Traits = remove_cvref_t; - static constexpr index_t kNumWarps = kM0_ / (get_warp_size() / 4); - static constexpr index_t kBlockSize = kNumWarps * get_warp_size(); - static constexpr bool kIsGroupMode = kIsGroupMode_; + static_assert(std::is_same_v); static constexpr index_t kHeadDimV = HeadDimV_; - static constexpr index_t kM0 = kM0_; - static constexpr index_t kN1 = kN1_; + static constexpr bool kIsGroupMode = kIsGroupMode_; + + using BaseType::kM0; + using BaseType::kN1; + + static_assert(kN1 <= kHeadDimV && kHeadDimV % kN1 == 0); // attributes from traits static constexpr bool kPadSeqLenQ = Traits::kPadSeqLenQ; @@ -136,6 +151,13 @@ struct BlockFmhaSplitKVCombinePipelineProblem static constexpr bool kDoFp8StaticQuant = Traits::kDoFp8StaticQuant; static constexpr index_t kBlockPerCu = Traits::kBlockPerCu; static constexpr index_t kMaxSplits = Traits::kMaxSplits; + static_assert(8 <= kMaxSplits); + + static constexpr index_t kNumWarps = 4; // always use 4 warps for each workgroup + static constexpr index_t kBlockSize = kNumWarps * get_warp_size(); + + static_assert(get_warp_size() <= (kM0 * kMaxSplits) && + (kM0 * kMaxSplits) % get_warp_size() == 0); }; template template CK_TILE_HOST_DEVICE static constexpr auto GetAlignmentQ() { + constexpr index_t MaxVectorSize = 16 / sizeof(typename Problem::QDataType); + using BlockGemm = remove_cvref_t())>; constexpr auto config = BlockGemm::Policy::template GetWarpGemmMWarpNWarp(); using WG = remove_cvref_t())>; - return WG::kK / WG::WarpGemmAttribute::Impl::kABKLane; + + return min(MaxVectorSize, WG::kK / WG::WarpGemmAttribute::Impl::kABKLane); } template CK_TILE_HOST_DEVICE static constexpr auto MakeQDramTileDistribution() { - constexpr auto config = BlockGemm::Policy::template GetWarpGemmMWarpNWarp(); - using WG = remove_cvref_t())>; - constexpr index_t MWarp = config.template at<1>(); - - constexpr index_t kMPerBlock = Problem::BlockFmhaShape::kM0; - constexpr index_t kKPerBlock = Problem::BlockFmhaShape::kSubQKHeaddim; - - constexpr index_t K2 = WG::kK / WG::WarpGemmAttribute::Impl::kABKLane; - 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); - - if constexpr(1 < Problem::kNumGemm0Warps) - { - return make_static_tile_distribution( - tile_distribution_encoding, - tuple, sequence>, - tuple, sequence<2, 1>>, - tuple, sequence<1, 2>>, - sequence<1, 2, 2>, - sequence<0, 0, 2>>{}); - } - else - { - static_assert(MWarp == 1); - - return make_static_tile_distribution( - tile_distribution_encoding, - tuple, sequence>, - tuple>, - tuple>, - sequence<1, 2, 2>, - sequence<0, 0, 2>>{}); - } + return BlockGemm::template MakeABlockTileDistribution< + Problem::BlockFmhaShape::kM0, + Problem::BlockFmhaShape::kSubQKHeaddim>(); } template @@ -105,7 +74,7 @@ struct BlockFmhaPipelineQXCustomPolicy constexpr auto warp_gemm = []() { constexpr index_t WarpGemmM = Problem::BlockFmhaShape::Gemm0WarpTile::at(number<0>{}); - static_assert(WarpGemmM == 16 || WarpGemmM == 32); + static_assert(WarpGemmM == 4 || WarpGemmM == 16 || WarpGemmM == 32); if constexpr(std::is_same_v && std::is_same_v && @@ -113,8 +82,10 @@ struct BlockFmhaPipelineQXCustomPolicy { if constexpr(WarpGemmM == 32) return WarpGemmMfmaF16F16F32M32N32K16SwizzleBTransposedCDistribution{}; - else // WarpGemmM == 16 + else if constexpr(WarpGemmM == 16) return WarpGemmMfmaF16F16F32M16N16K16TransposedCDistribution{}; + else // WarpGemmM == 4 + return WarpGemmMfmaF16F16F32M4N64K16{}; } else if constexpr(std::is_same_v && std::is_same_v && @@ -122,8 +93,10 @@ struct BlockFmhaPipelineQXCustomPolicy { if constexpr(WarpGemmM == 32) return WarpGemmMfmaBf16Bf16F32M32N32K16SwizzleBTransposedCDistribution{}; - else // WarpGemmM == 16 + else if constexpr(WarpGemmM == 16) return WarpGemmMfmaBf16Bf16F32M16N16K16TransposedCDistribution{}; + else // WarpGemmM == 4 + return WarpGemmMfmaBf16Bf16F32M4N64K16{}; } else if constexpr(std::is_same_v && std::is_same_v && diff --git a/include/ck_tile/ops/fmha/pipeline/tile_fmha_shape.hpp b/include/ck_tile/ops/fmha/pipeline/tile_fmha_shape.hpp index bb33b5f021..5ce80c2d1f 100644 --- a/include/ck_tile/ops/fmha/pipeline/tile_fmha_shape.hpp +++ b/include/ck_tile/ops/fmha/pipeline/tile_fmha_shape.hpp @@ -43,8 +43,6 @@ struct TileFmhaShape static constexpr index_t NumWarps = max(NumGemm0Warps, NumGemm1Warps); - static_assert(std::is_same_v); - static constexpr index_t kM0 = BlockTile::at(number<0>{}); // tile size along q seqlen static constexpr index_t kN0 = BlockTile::at(number<1>{}); // tile size along k seqlen static constexpr index_t kK0 = BlockTile::at(number<2>{}); // tile size along qk gemm unroll diff --git a/include/ck_tile/ops/gemm/block/block_gemm_areg_bsmem_creg_one_warp_v1.hpp b/include/ck_tile/ops/gemm/block/block_gemm_areg_bsmem_creg_one_warp_v1.hpp index ff23f63556..b99466b1ea 100644 --- a/include/ck_tile/ops/gemm/block/block_gemm_areg_bsmem_creg_one_warp_v1.hpp +++ b/include/ck_tile/ops/gemm/block/block_gemm_areg_bsmem_creg_one_warp_v1.hpp @@ -65,14 +65,6 @@ struct BlockGemmARegBSmemCRegOneWarpV1 const index_t iNWarp = 0; - constexpr auto a_block_outer_dstr_encoding = - tile_distribution_encoding, - tuple, sequence>, - tuple>, - tuple>, - sequence<1, 2>, - sequence<0, 0>>{}; - constexpr auto c_block_outer_dstr_encoding = tile_distribution_encoding, tuple, sequence>, @@ -81,19 +73,14 @@ struct BlockGemmARegBSmemCRegOneWarpV1 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); + auto a_block_tensor = make_static_distributed_tensor( + MakeABlockTileDistribution()); a_block_tensor.get_thread_buffer() = a_block_tensor_tmp.get_thread_buffer(); @@ -187,6 +174,33 @@ struct BlockGemmARegBSmemCRegOneWarpV1 }); } + template + CK_TILE_DEVICE static constexpr auto MakeABlockTileDistribution() + { + constexpr auto config = Policy::template GetWarpGemmMWarpNWarp(); + + using WG = remove_cvref_t())>; + + constexpr index_t MWarp = config.template at<1>(); + constexpr index_t NWarp = config.template at<2>(); + + constexpr index_t MIterPerWarp = MPerBlock / (MWarp * WG::kM); + constexpr index_t KIterPerWarp = KPerBlock / 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{}); + + return make_static_tile_distribution(a_block_dstr_encode); + } + CK_TILE_DEVICE static constexpr auto MakeCBlockTile() { constexpr index_t MPerBlock = BlockGemmShape::kM; diff --git a/include/ck_tile/ops/gemm/block/block_gemm_areg_bsmem_creg_v2.hpp b/include/ck_tile/ops/gemm/block/block_gemm_areg_bsmem_creg_v2.hpp index 173ef0a02e..0181c0eec8 100644 --- a/include/ck_tile/ops/gemm/block/block_gemm_areg_bsmem_creg_v2.hpp +++ b/include/ck_tile/ops/gemm/block/block_gemm_areg_bsmem_creg_v2.hpp @@ -59,14 +59,6 @@ struct BlockGemmARegBSmemCRegV2 const index_t iNWarp = get_warp_id() % NWarp; - constexpr auto a_block_outer_dstr_encoding = - tile_distribution_encoding, - tuple, sequence>, - tuple>, - tuple>, - sequence<1, 2>, - sequence<0, 0>>{}; - constexpr auto c_block_outer_dstr_encoding = tile_distribution_encoding< sequence<>, tuple, sequence>, @@ -75,19 +67,14 @@ struct BlockGemmARegBSmemCRegV2 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); + auto a_block_tensor = make_static_distributed_tensor( + MakeABlockTileDistribution()); a_block_tensor.get_thread_buffer() = a_block_tensor_tmp.get_thread_buffer(); @@ -182,6 +169,33 @@ struct BlockGemmARegBSmemCRegV2 }); } + template + CK_TILE_DEVICE static constexpr auto MakeABlockTileDistribution() + { + constexpr auto config = Policy::template GetWarpGemmMWarpNWarp(); + + using WG = remove_cvref_t())>; + + constexpr index_t MWarp = config.template at<1>(); + constexpr index_t NWarp = config.template at<2>(); + + constexpr index_t MIterPerWarp = MPerBlock / (MWarp * WG::kM); + constexpr index_t KIterPerWarp = KPerBlock / 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{}); + + return make_static_tile_distribution(a_block_dstr_encode); + } + CK_TILE_DEVICE static constexpr auto MakeCBlockTile() { constexpr index_t MPerBlock = BlockGemmShape::kM; diff --git a/include/ck_tile/ops/gemm/warp/warp_gemm.hpp b/include/ck_tile/ops/gemm/warp/warp_gemm.hpp index 89ea82c5bd..1fd12973f6 100644 --- a/include/ck_tile/ops/gemm/warp/warp_gemm.hpp +++ b/include/ck_tile/ops/gemm/warp/warp_gemm.hpp @@ -56,6 +56,14 @@ using WarpGemmMfmaF16F16F32M32N32K16SwizzleBTransposedCDistribution = WarpGemmAttributeMfmaImplF16F16F32M32N32K8, 2>>; +using WarpGemmMfmaF16F16F32M4N64K16 = WarpGemmImpl, + 4>>; + +using WarpGemmMfmaF16F16F32M64N4K16 = WarpGemmImpl, + 4>>; + // bf16 using WarpGemmMfmaBf16Bf16F32M32N32K8 = WarpGemmImpl< @@ -104,6 +112,14 @@ using WarpGemmMfmaBf16Bf16F32M32N32K16SwizzleBTransposedCDistribution = WarpGemmAttributeMfmaImplBf16Bf16F32M32N32K8, 2>>; +using WarpGemmMfmaBf16Bf16F32M4N64K16 = WarpGemmImpl, + 4>>; + +using WarpGemmMfmaBf16Bf16F32M64N4K16 = WarpGemmImpl, + 4>>; + // fp8 using WarpGemmMfma_f32_32x32x16_fp8_fp8 = WarpGemmImpl< diff --git a/include/ck_tile/ops/gemm/warp/warp_gemm_attribute_mfma.hpp b/include/ck_tile/ops/gemm/warp/warp_gemm_attribute_mfma.hpp index a9e466a796..e7d4c37966 100644 --- a/include/ck_tile/ops/gemm/warp/warp_gemm_attribute_mfma.hpp +++ b/include/ck_tile/ops/gemm/warp/warp_gemm_attribute_mfma.hpp @@ -28,6 +28,9 @@ struct WarpGemmAtrributeMfma CK_TILE_HOST_DEVICE static constexpr auto get_num_of_access() { return 1; } + static_assert(Impl::kAMBlock == 1 && Impl::kBNBlock == 1, + "Multi-block WarpGemmAttributeMfmaImpl is not supported"); + using AWarpDstrEncoding = tile_distribution_encoding< sequence<>, tuple, sequence>, @@ -94,30 +97,130 @@ struct WarpGemmAtrributeMfmaIterateK CK_TILE_HOST_DEVICE static constexpr auto get_num_of_access() { return kKIter; } - using AWarpDstrEncoding = tile_distribution_encoding< - sequence<>, - tuple, sequence>, - tuple>, - tuple>, - sequence<2>, - sequence<1>>; + static_assert(Impl::kAMBlock == 1 || Impl::kBNBlock == 1, + "Multi-block on both M & N directions is not supported"); - using BWarpDstrEncoding = tile_distribution_encoding< - sequence<>, - tuple, sequence>, - tuple>, - tuple>, - sequence<2>, - sequence<1>>; + CK_TILE_DEVICE static constexpr auto get_awarp_dstr_encoding() + { + if constexpr(Impl::kAMBlock == 1 && Impl::kBNBlock == 1) + { + return tile_distribution_encoding< + sequence<>, + tuple, + sequence>, + tuple>, + tuple>, + sequence<2>, + sequence<1>>{}; + } + else if constexpr(Impl::kAMBlock == 1 && 1 < Impl::kBNBlock) + { + // each M blocks share the same data + return tile_distribution_encoding< + sequence, + tuple, + sequence>, + tuple>, + tuple>, + sequence<2>, + sequence<1>>{}; + } + else if constexpr(1 < Impl::kAMBlock && Impl::kBNBlock == 1) + { + // single block to multi-block thread mapping + return tile_distribution_encoding< + sequence<>, + tuple, + sequence>, + tuple>, + tuple>, + sequence<2>, + sequence<1>>{}; + } + } - using CWarpDstrEncoding = tile_distribution_encoding< - sequence<>, - tuple, - sequence>, - tuple>, - tuple>, - sequence<1, 1>, - sequence<0, 2>>; + CK_TILE_DEVICE static constexpr auto get_bwarp_dstr_encoding() + { + if constexpr(Impl::kAMBlock == 1 && Impl::kBNBlock == 1) + { + return tile_distribution_encoding< + sequence<>, + tuple, + sequence>, + tuple>, + tuple>, + sequence<2>, + sequence<1>>{}; + } + else if constexpr(Impl::kAMBlock == 1 && 1 < Impl::kBNBlock) + { + // single block to multi-block thread mapping + return tile_distribution_encoding< + sequence<>, + tuple, + sequence>, + tuple>, + tuple>, + sequence<2>, + sequence<1>>{}; + } + else if constexpr(1 < Impl::kAMBlock && Impl::kBNBlock == 1) + { + // each N blocks share the same data + return tile_distribution_encoding< + sequence, + tuple, + sequence>, + tuple>, + tuple>, + sequence<2>, + sequence<1>>{}; + } + } + + CK_TILE_DEVICE static constexpr auto get_cwarp_dstr_encoding() + { + if constexpr(Impl::kAMBlock == 1 && Impl::kBNBlock == 1) + { + return tile_distribution_encoding< + sequence<>, + tuple, + sequence>, + tuple>, + tuple>, + sequence<1, 1>, + sequence<0, 2>>{}; + } + else if constexpr(Impl::kAMBlock == 1 && 1 < Impl::kBNBlock) + { + return tile_distribution_encoding< + sequence<>, + tuple, + sequence>, + tuple>, + tuple>, + sequence<1, 1>, + sequence<0, 2>>{}; + } + else if constexpr(1 < Impl::kAMBlock && Impl::kBNBlock == 1) + { + return tile_distribution_encoding< + sequence<>, + tuple< + sequence, + sequence>, + tuple>, + tuple>, + sequence<1, 1>, + sequence<0, 2>>{}; + } + } + + using AWarpDstrEncoding = decltype(get_awarp_dstr_encoding()); + + using BWarpDstrEncoding = decltype(get_bwarp_dstr_encoding()); + + using CWarpDstrEncoding = decltype(get_cwarp_dstr_encoding()); // c_vec += a_vec * b_vec template @@ -206,6 +309,9 @@ struct WarpGemmAtrributeMfmaTransposedCDistribution CK_TILE_HOST_DEVICE static constexpr auto get_num_of_access() { return 1; } + static_assert(Impl::kAMBlock == 1 && Impl::kBNBlock == 1, + "Multi-block WarpGemmAttributeMfmaImpl is not supported"); + using AWarpDstrEncoding = tile_distribution_encoding< sequence<>, tuple, sequence>, @@ -270,6 +376,9 @@ struct WarpGemmAtrributeMfmaTransposedCDistribution_SwizzleB CK_TILE_HOST_DEVICE static constexpr auto get_num_of_access() { return 1; } + static_assert(Impl::kAMBlock == 1 && Impl::kBNBlock == 1, + "Multi-block WarpGemmAttributeMfmaImpl is not supported"); + using AWarpDstrEncoding = tile_distribution_encoding< sequence<>, tuple, sequence>, @@ -341,30 +450,130 @@ struct WarpGemmAtrributeMfmaIterateKAndTransposedCDistribution CK_TILE_HOST_DEVICE static constexpr auto get_num_of_access() { return kKIter; } - using AWarpDstrEncoding = tile_distribution_encoding< - sequence<>, - tuple, sequence>, - tuple>, - tuple>, - sequence<2>, - sequence<1>>; + static_assert(Impl::kAMBlock == 1 || Impl::kBNBlock == 1, + "Multi-block on both M & N directions is not supported"); - using BWarpDstrEncoding = tile_distribution_encoding< - sequence<>, - tuple, sequence>, - tuple>, - tuple>, - sequence<2>, - sequence<1>>; + CK_TILE_DEVICE static constexpr auto get_awarp_dstr_encoding() + { + if constexpr(Impl::kAMBlock == 1 && Impl::kBNBlock == 1) + { + return tile_distribution_encoding< + sequence<>, + tuple, + sequence>, + tuple>, + tuple>, + sequence<2>, + sequence<1>>{}; + } + else if constexpr(Impl::kAMBlock == 1 && 1 < Impl::kBNBlock) + { + // single block to multi-block thread mapping + return tile_distribution_encoding< + sequence<>, + tuple, + sequence>, + tuple>, + tuple>, + sequence<2>, + sequence<1>>{}; + } + else if constexpr(1 < Impl::kAMBlock && Impl::kBNBlock == 1) + { + // each N blocks share the same data + return tile_distribution_encoding< + sequence, + tuple, + sequence>, + tuple>, + tuple>, + sequence<2>, + sequence<1>>{}; + } + } - using CWarpDstrEncoding = tile_distribution_encoding< - sequence<>, - tuple, - sequence>, - tuple>, - tuple>, - sequence<2, 2>, - sequence<0, 2>>; + CK_TILE_DEVICE static constexpr auto get_bwarp_dstr_encoding() + { + if constexpr(Impl::kAMBlock == 1 && Impl::kBNBlock == 1) + { + return tile_distribution_encoding< + sequence<>, + tuple, + sequence>, + tuple>, + tuple>, + sequence<2>, + sequence<1>>{}; + } + else if constexpr(Impl::kAMBlock == 1 && 1 < Impl::kBNBlock) + { + // each M blocks share the same data + return tile_distribution_encoding< + sequence, + tuple, + sequence>, + tuple>, + tuple>, + sequence<2>, + sequence<1>>{}; + } + else if constexpr(1 < Impl::kAMBlock && Impl::kBNBlock == 1) + { + // single block to multi-block thread mapping + return tile_distribution_encoding< + sequence<>, + tuple, + sequence>, + tuple>, + tuple>, + sequence<2>, + sequence<1>>{}; + } + } + + CK_TILE_DEVICE static constexpr auto get_cwarp_dstr_encoding() + { + if constexpr(Impl::kAMBlock == 1 && Impl::kBNBlock == 1) + { + return tile_distribution_encoding< + sequence<>, + tuple, + sequence>, + tuple>, + tuple>, + sequence<2, 2>, + sequence<0, 2>>{}; + } + else if constexpr(Impl::kAMBlock == 1 && 1 < Impl::kBNBlock) + { + return tile_distribution_encoding< + sequence<>, + tuple, + sequence>, + tuple>, + tuple>, + sequence<2, 2>, + sequence<0, 2>>{}; + } + else if constexpr(1 < Impl::kAMBlock && Impl::kBNBlock == 1) + { + return tile_distribution_encoding< + sequence<>, + tuple< + sequence, + sequence>, + tuple>, + tuple>, + sequence<2, 2>, + sequence<0, 2>>{}; + } + } + + using AWarpDstrEncoding = decltype(get_awarp_dstr_encoding()); + + using BWarpDstrEncoding = decltype(get_bwarp_dstr_encoding()); + + using CWarpDstrEncoding = decltype(get_cwarp_dstr_encoding()); template // c_vec += a_vec * b_vec @@ -457,6 +666,9 @@ struct WarpGemmAtrributeMfmaIterateKAndTransposedCDistribution_SwizzleB CK_TILE_HOST_DEVICE static constexpr auto get_num_of_access() { return kKIter; } + static_assert(Impl::kAMBlock == 1 && Impl::kBNBlock == 1, + "Multi-block WarpGemmAttributeMfmaImpl is not supported"); + using AWarpDstrEncoding = tile_distribution_encoding< sequence<>, tuple, sequence>, @@ -597,6 +809,9 @@ struct WarpGemmAtrributeMfmaIterateK_SwizzleA CK_TILE_HOST_DEVICE static constexpr auto get_num_of_access() { return kKIter; } + static_assert(Impl::kAMBlock == 1 && Impl::kBNBlock == 1, + "Multi-block WarpGemmAttributeMfmaImpl is not supported"); + using AWarpDstrEncoding = tile_distribution_encoding< sequence<>, tuple +struct WarpGemmAttributeMfmaImplF16F16F32M4N64K4 +{ + static constexpr WGAttrCtlEnum Ctrl = Ctrl_; + using ADataType = fp16_t; + using BDataType = fp16_t; + using CDataType = float; + + using AVecType = ext_vector_t; + using BVecType = ext_vector_t; + using CVecType = ext_vector_t; + + static constexpr index_t kM = 4; + static constexpr index_t kN = 64; + static constexpr index_t kK = 4; + + static constexpr index_t kAMBlock = 1; + static constexpr index_t kBNBlock = 16; + + // we only write down single block (4 threads) thread mapping here + static constexpr index_t kAMLane = 4; + static constexpr index_t kBNLane = 4; + static constexpr index_t kABKLane = 1; + static constexpr index_t kABKPerLane = 4; + + static constexpr index_t kCMLane = 1; + static constexpr index_t kCNLane = 4; + static constexpr index_t kCM0PerLane = 1; + static constexpr index_t kCM1PerLane = 4; + + // c_vec += a_vec * b_vec + template + CK_TILE_DEVICE void operator()(CVecType& c_vec, + const AVecType& a_vec, + const BVecType& b_vec, + bool_constant = {}) const + { + DISPATCH_MFMA_CTRL_("v_mfma_f32_4x4x4f16", Ctrl) + else + { +#if defined(__gfx9__) + c_vec = __builtin_amdgcn_mfma_f32_4x4x4f16(a_vec, b_vec, c_vec, 0, 0, 0); +#else + ignore = c_vec; + ignore = a_vec; + ignore = b_vec; +#endif + } + } + + // c_vec = a_vec * b_vec + CK_TILE_DEVICE CVecType operator()(const AVecType& a_vec, const BVecType& b_vec) const + { +#if defined(__gfx9__) + return bit_cast( + __builtin_amdgcn_mfma_f32_4x4x4f16(a_vec, b_vec, fp32x4_t{0.f}, 0, 0, 0)); +#else + ignore = a_vec; + ignore = b_vec; + return CVecType{0.f}; +#endif + } +}; + +template +struct WarpGemmAttributeMfmaImplF16F16F32M64N4K4 +{ + static constexpr WGAttrCtlEnum Ctrl = Ctrl_; + using ADataType = fp16_t; + using BDataType = fp16_t; + using CDataType = float; + + using AVecType = ext_vector_t; + using BVecType = ext_vector_t; + using CVecType = ext_vector_t; + + static constexpr index_t kM = 64; + static constexpr index_t kN = 4; + static constexpr index_t kK = 4; + + static constexpr index_t kAMBlock = 16; + static constexpr index_t kBNBlock = 1; + + // we only write down single block (4 threads) thread mapping here + static constexpr index_t kAMLane = 4; + static constexpr index_t kBNLane = 4; + static constexpr index_t kABKLane = 1; + static constexpr index_t kABKPerLane = 4; + + static constexpr index_t kCMLane = 1; + static constexpr index_t kCNLane = 4; + static constexpr index_t kCM0PerLane = 1; + static constexpr index_t kCM1PerLane = 4; + + // c_vec += a_vec * b_vec + template + CK_TILE_DEVICE void operator()(CVecType& c_vec, + const AVecType& a_vec, + const BVecType& b_vec, + bool_constant = {}) const + { + DISPATCH_MFMA_CTRL_("v_mfma_f32_4x4x4f16", Ctrl) + else + { +#if defined(__gfx9__) + c_vec = __builtin_amdgcn_mfma_f32_4x4x4f16(a_vec, b_vec, c_vec, 0, 0, 0); +#else + ignore = c_vec; + ignore = a_vec; + ignore = b_vec; +#endif + } + } + + // c_vec = a_vec * b_vec + CK_TILE_DEVICE CVecType operator()(const AVecType& a_vec, const BVecType& b_vec) const + { +#if defined(__gfx9__) + return bit_cast( + __builtin_amdgcn_mfma_f32_4x4x4f16(a_vec, b_vec, fp32x4_t{0.f}, 0, 0, 0)); +#else + ignore = a_vec; + ignore = b_vec; + return CVecType{0.f}; +#endif + } +}; + // Bf16 template struct WarpGemmAttributeMfmaImplBf16Bf16F32M32N32K8 @@ -199,6 +333,9 @@ struct WarpGemmAttributeMfmaImplBf16Bf16F32M32N32K8 static constexpr index_t kN = 32; static constexpr index_t kK = 8; + static constexpr index_t kAMBlock = 1; + static constexpr index_t kBNBlock = 1; + static constexpr index_t kAMLane = 32; static constexpr index_t kBNLane = 32; static constexpr index_t kABKLane = 2; @@ -285,6 +422,9 @@ struct WarpGemmAttributeMfmaImplBf16Bf16F32M16N16K16 static constexpr index_t kN = 16; static constexpr index_t kK = 16; + static constexpr index_t kAMBlock = 1; + static constexpr index_t kBNBlock = 1; + static constexpr index_t kAMLane = 16; static constexpr index_t kBNLane = 16; static constexpr index_t kABKLane = 4; @@ -354,6 +494,134 @@ struct WarpGemmAttributeMfmaImplBf16Bf16F32M16N16K16 } }; +template +struct WarpGemmAttributeMfmaImplBf16Bf16F32M4N64K4 +{ + static constexpr WGAttrCtlEnum Ctrl = Ctrl_; + using ADataType = bf16_t; + using BDataType = bf16_t; + using CDataType = float; + + using AVecType = ext_vector_t; + using BVecType = ext_vector_t; + using CVecType = ext_vector_t; + + static constexpr index_t kM = 4; + static constexpr index_t kN = 64; + static constexpr index_t kK = 4; + + static constexpr index_t kAMBlock = 1; + static constexpr index_t kBNBlock = 16; + + // we only write down single block (4 threads) thread mapping here + static constexpr index_t kAMLane = 4; + static constexpr index_t kBNLane = 4; + static constexpr index_t kABKLane = 1; + static constexpr index_t kABKPerLane = 4; + + static constexpr index_t kCMLane = 1; + static constexpr index_t kCNLane = 4; + static constexpr index_t kCM0PerLane = 1; + static constexpr index_t kCM1PerLane = 4; + + // c_vec += a_vec * b_vec + template + CK_TILE_DEVICE void operator()(CVecType& c_vec, + const AVecType& a_vec, + const BVecType& b_vec, + bool_constant = {}) const + { + DISPATCH_MFMA_CTRL_("v_mfma_f32_4x4x4bf16_1k", Ctrl) + else + { +#if defined(__gfx9__) + c_vec = __builtin_amdgcn_mfma_f32_4x4x4bf16_1k(a_vec, b_vec, c_vec, 0, 0, 0); +#else + ignore = c_vec; + ignore = a_vec; + ignore = b_vec; +#endif + } + } + + // c_vec = a_vec * b_vec + CK_TILE_DEVICE CVecType operator()(const AVecType& a_vec, const BVecType& b_vec) const + { +#if defined(__gfx9__) + return bit_cast( + __builtin_amdgcn_mfma_f32_4x4x4bf16_1k(a_vec, b_vec, fp32x4_t{0.f}, 0, 0, 0)); +#else + ignore = a_vec; + ignore = b_vec; + return CVecType{0.f}; +#endif + } +}; + +template +struct WarpGemmAttributeMfmaImplBf16Bf16F32M64N4K4 +{ + static constexpr WGAttrCtlEnum Ctrl = Ctrl_; + using ADataType = bf16_t; + using BDataType = bf16_t; + using CDataType = float; + + using AVecType = ext_vector_t; + using BVecType = ext_vector_t; + using CVecType = ext_vector_t; + + static constexpr index_t kM = 64; + static constexpr index_t kN = 4; + static constexpr index_t kK = 4; + + static constexpr index_t kAMBlock = 16; + static constexpr index_t kBNBlock = 1; + + // we only write down single block (4 threads) thread mapping here + static constexpr index_t kAMLane = 4; + static constexpr index_t kBNLane = 4; + static constexpr index_t kABKLane = 1; + static constexpr index_t kABKPerLane = 4; + + static constexpr index_t kCMLane = 1; + static constexpr index_t kCNLane = 4; + static constexpr index_t kCM0PerLane = 1; + static constexpr index_t kCM1PerLane = 4; + + // c_vec += a_vec * b_vec + template + CK_TILE_DEVICE void operator()(CVecType& c_vec, + const AVecType& a_vec, + const BVecType& b_vec, + bool_constant = {}) const + { + DISPATCH_MFMA_CTRL_("v_mfma_f32_4x4x4bf16_1k", Ctrl) + else + { +#if defined(__gfx9__) + c_vec = __builtin_amdgcn_mfma_f32_4x4x4bf16_1k(a_vec, b_vec, c_vec, 0, 0, 0); +#else + ignore = c_vec; + ignore = a_vec; + ignore = b_vec; +#endif + } + } + + // c_vec = a_vec * b_vec + CK_TILE_DEVICE CVecType operator()(const AVecType& a_vec, const BVecType& b_vec) const + { +#if defined(__gfx9__) + return bit_cast( + __builtin_amdgcn_mfma_f32_4x4x4bf16_1k(a_vec, b_vec, fp32x4_t{0.f}, 0, 0, 0)); +#else + ignore = a_vec; + ignore = b_vec; + return CVecType{0.f}; +#endif + } +}; + // FP8 template struct WarpGemmAttributeMfmaImpl_f32_32x32x16_f8_base @@ -371,6 +639,9 @@ struct WarpGemmAttributeMfmaImpl_f32_32x32x16_f8_base static constexpr index_t kN = 32; static constexpr index_t kK = 16; + static constexpr index_t kAMBlock = 1; + static constexpr index_t kBNBlock = 1; + static constexpr index_t kAMLane = 32; static constexpr index_t kBNLane = 32; static constexpr index_t kABKLane = 2; diff --git a/include/ck_tile/ops/gemm/warp/warp_gemm_dispatcher.hpp b/include/ck_tile/ops/gemm/warp/warp_gemm_dispatcher.hpp index 99cd5d787e..9c319b5e5f 100644 --- a/include/ck_tile/ops/gemm/warp/warp_gemm_dispatcher.hpp +++ b/include/ck_tile/ops/gemm/warp/warp_gemm_dispatcher.hpp @@ -29,6 +29,8 @@ template<> struct WarpGemmMfmaDispatcher struct WarpGemmMfmaDispatcher { using Type = WarpGemmMfmaF16F16F32M16N16K16TransposedCDistribution; }; template<> struct WarpGemmMfmaDispatcher { using Type = WarpGemmMfmaF16F16F32M16N16K32; }; template<> struct WarpGemmMfmaDispatcher { using Type = WarpGemmMfmaF16F16F32M16N16K32TransposedCDistribution; }; +template<> struct WarpGemmMfmaDispatcher { using Type = WarpGemmMfmaF16F16F32M4N64K16; }; +template<> struct WarpGemmMfmaDispatcher { using Type = WarpGemmMfmaF16F16F32M64N4K16; }; template<> struct WarpGemmMfmaDispatcher { using Type = WarpGemmMfmaF16F16F32M32N32K8SwizzleA; }; template<> struct WarpGemmMfmaDispatcher { using Type = WarpGemmMfmaF16F16F32M32N32K16SwizzleA; }; @@ -42,6 +44,8 @@ template<> struct WarpGemmMfmaDispatcher struct WarpGemmMfmaDispatcher { using Type = WarpGemmMfmaBf16Bf16F32M16N16K16TransposedCDistribution; }; template<> struct WarpGemmMfmaDispatcher { using Type = WarpGemmMfmaBf16Bf16F32M16N16K32; }; template<> struct WarpGemmMfmaDispatcher { using Type = WarpGemmMfmaBf16Bf16F32M16N16K32TransposedCDistribution; }; +template<> struct WarpGemmMfmaDispatcher { using Type = WarpGemmMfmaBf16Bf16F32M4N64K16; }; +template<> struct WarpGemmMfmaDispatcher { using Type = WarpGemmMfmaBf16Bf16F32M64N4K16; }; template<> struct WarpGemmMfmaDispatcher { using Type = WarpGemmMfmaBf16Bf16F32M32N32K8SwizzleA; }; template<> struct WarpGemmMfmaDispatcher { using Type = WarpGemmMfmaBf16Bf16F32M32N32K16SwizzleA; }; From 1c45ca35dd5c215e0c1db1f40f01556f467f52a8 Mon Sep 17 00:00:00 2001 From: carlushuang Date: Fri, 20 Dec 2024 16:40:45 +0800 Subject: [PATCH 22/25] hot-fix (#1768) --- .../ck_tile/ops/gemm/warp/warp_gemm_attribute_mfma_impl.hpp | 3 +++ 1 file changed, 3 insertions(+) diff --git a/include/ck_tile/ops/gemm/warp/warp_gemm_attribute_mfma_impl.hpp b/include/ck_tile/ops/gemm/warp/warp_gemm_attribute_mfma_impl.hpp index fa24711de2..21a865e792 100644 --- a/include/ck_tile/ops/gemm/warp/warp_gemm_attribute_mfma_impl.hpp +++ b/include/ck_tile/ops/gemm/warp/warp_gemm_attribute_mfma_impl.hpp @@ -839,6 +839,9 @@ struct WarpGemmAttributeMfmaImpl_i32_32x32x16_i8 static constexpr index_t kN = 32; static constexpr index_t kK = 16; + static constexpr index_t kAMBlock = 1; + static constexpr index_t kBNBlock = 1; + static constexpr index_t kAMLane = 32; static constexpr index_t kBNLane = 32; static constexpr index_t kABKLane = 2; From 07339c738396ebeae57374771ded4dcf11bddf1e Mon Sep 17 00:00:00 2001 From: Illia Silin <98187287+illsilin@users.noreply.github.com> Date: Fri, 20 Dec 2024 07:52:24 -0800 Subject: [PATCH 23/25] fix typo for CK_USE_OCP_FP8 (#1769) --- include/ck/config.h.in | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/include/ck/config.h.in b/include/ck/config.h.in index 55a498073f..2c37300e9b 100644 --- a/include/ck/config.h.in +++ b/include/ck/config.h.in @@ -115,8 +115,8 @@ #cmakedefine CK_USE_GFX94 @CK_USE_GFX94@ #endif -#ifndef DCK_USE_OCP_FP8 -#cmakedefine DCK_USE_OCP_FP8 @DCK_USE_OCP_FP8@ +#ifndef CK_USE_OCP_FP8 +#cmakedefine CK_USE_OCP_FP8 @CK_USE_OCP_FP8@ #endif #ifndef CK_USE_FNUZ_FP8 From 3d15f364b367b24ac709ea5687fa2d7d39f07cf9 Mon Sep 17 00:00:00 2001 From: carlushuang Date: Mon, 23 Dec 2024 10:59:02 +0800 Subject: [PATCH 24/25] [CK_TILE] optimize moe-sorting kernel (#1771) * opt moe sorting * remove commented code --- .../13_moe_sorting/moe_sorting_api.cpp | 53 ++-- .../13_moe_sorting/script/smoke_test.sh | 3 +- .../instances/fused_moesorting_api.cpp | 53 ++-- .../fused_moe/kernel/moe_sorting_kernel.hpp | 253 +++++++++++++++--- .../pipeline/moe_sorting_problem.hpp | 13 +- 5 files changed, 292 insertions(+), 83 deletions(-) diff --git a/example/ck_tile/13_moe_sorting/moe_sorting_api.cpp b/example/ck_tile/13_moe_sorting/moe_sorting_api.cpp index 25e99c5306..723fb3f69f 100644 --- a/example/ck_tile/13_moe_sorting/moe_sorting_api.cpp +++ b/example/ck_tile/13_moe_sorting/moe_sorting_api.cpp @@ -3,18 +3,42 @@ #include "moe_sorting_api.hpp" -#define MOE_SORTING_DISPATCH(unroll_num_) \ - constexpr ck_tile::index_t unroll_num = unroll_num_; \ - using ms_problem = ck_tile::MoeSortingProblem; \ - using kernel = ck_tile::MoeSortingKernel; \ - auto kargs = kernel::MakeKargs(a); \ - const dim3 grids = kernel::GridSize(a); \ - const dim3 blocks = kernel::BlockSize(a); \ - const auto lds_bytes = kernel::GetSmemSize(a); \ - float ave_time = ck_tile::launch_kernel( \ - s, ck_tile::make_kernel(kernel{}, grids, blocks, lds_bytes, kargs)); \ +#define MOE_SORTING_DISPATCH_ETILE(unroll_num_, expert_tile_) \ + constexpr ck_tile::index_t unroll_num = unroll_num_; \ + constexpr ck_tile::index_t expert_tile = expert_tile_; \ + using ms_problem = \ + ck_tile::MoeSortingProblem; \ + using kernel = ck_tile::MoeSortingKernel; \ + auto kargs = kernel::MakeKargs(a); \ + const dim3 grids = kernel::GridSize(a); \ + const dim3 blocks = kernel::BlockSize(a); \ + const auto lds_bytes = kernel::GetSmemSize(a); \ + float ave_time = ck_tile::launch_kernel( \ + s, ck_tile::make_kernel(kernel{}, grids, blocks, lds_bytes, kargs)); \ return ave_time; +#define MOE_SORTING_DISPATCH(unroll_num_) \ + if(a.num_experts <= 8) \ + { \ + MOE_SORTING_DISPATCH_ETILE(unroll_num_, 8) \ + } \ + else if(a.num_experts <= 16) \ + { \ + MOE_SORTING_DISPATCH_ETILE(unroll_num_, 16) \ + } \ + else if(a.num_experts <= 32) \ + { \ + MOE_SORTING_DISPATCH_ETILE(unroll_num_, 32) \ + } \ + else if(a.num_experts <= 64) \ + { \ + MOE_SORTING_DISPATCH_ETILE(unroll_num_, 64) \ + } \ + else \ + { \ + MOE_SORTING_DISPATCH_ETILE(unroll_num_, 0) \ + } + float moe_sorting(moe_sorting_trait t, moe_sorting_args a, ck_tile::stream_config s) { if(t.weight_type == "fp32" && t.index_type == "int32") @@ -49,21 +73,12 @@ float moe_sorting(moe_sorting_trait t, moe_sorting_args a, ck_tile::stream_confi case(6): { MOE_SORTING_DISPATCH(6); } - case(7): { - MOE_SORTING_DISPATCH(7); - } case(8): { MOE_SORTING_DISPATCH(8); } - case(9): { - MOE_SORTING_DISPATCH(9); - } case(10): { MOE_SORTING_DISPATCH(10); } - case(11): { - MOE_SORTING_DISPATCH(11); - } default: { MOE_SORTING_DISPATCH(4); } diff --git a/example/ck_tile/13_moe_sorting/script/smoke_test.sh b/example/ck_tile/13_moe_sorting/script/smoke_test.sh index 1fc5eafcb0..3ff8a7332d 100644 --- a/example/ck_tile/13_moe_sorting/script/smoke_test.sh +++ b/example/ck_tile/13_moe_sorting/script/smoke_test.sh @@ -16,4 +16,5 @@ $EXE -t=127 -e=99 -k=19 $EXE -t=71 -e=11 -k=11 $EXE -t=1 -e=1 -k=1 $EXE -t=99 -e=2 -k=1 -$EXE -t=333 -e=99 -k=13 \ No newline at end of file +$EXE -t=333 -e=99 -k=13 +$EXE -t=128 -e=32 -k=5 -moe_buf_size=262144 diff --git a/example/ck_tile/15_fused_moe/instances/fused_moesorting_api.cpp b/example/ck_tile/15_fused_moe/instances/fused_moesorting_api.cpp index 75aaf86b74..7ca24c5c9a 100644 --- a/example/ck_tile/15_fused_moe/instances/fused_moesorting_api.cpp +++ b/example/ck_tile/15_fused_moe/instances/fused_moesorting_api.cpp @@ -3,18 +3,42 @@ #include "fused_moesorting.hpp" -#define MOE_SORTING_DISPATCH(unroll_num_) \ - constexpr ck_tile::index_t unroll_num = unroll_num_; \ - using ms_problem = ck_tile::MoeSortingProblem; \ - using kernel = ck_tile::MoeSortingKernel; \ - auto kargs = kernel::MakeKargs(a); \ - const dim3 grids = kernel::GridSize(a); \ - const dim3 blocks = kernel::BlockSize(a); \ - const auto lds_bytes = kernel::GetSmemSize(a); \ - float ave_time = ck_tile::launch_kernel( \ - s, ck_tile::make_kernel(kernel{}, grids, blocks, lds_bytes, kargs)); \ +#define MOE_SORTING_DISPATCH_ETILE(unroll_num_, expert_tile_) \ + constexpr ck_tile::index_t unroll_num = unroll_num_; \ + constexpr ck_tile::index_t expert_tile = expert_tile_; \ + using ms_problem = \ + ck_tile::MoeSortingProblem; \ + using kernel = ck_tile::MoeSortingKernel; \ + auto kargs = kernel::MakeKargs(a); \ + const dim3 grids = kernel::GridSize(a); \ + const dim3 blocks = kernel::BlockSize(a); \ + const auto lds_bytes = kernel::GetSmemSize(a); \ + float ave_time = ck_tile::launch_kernel( \ + s, ck_tile::make_kernel(kernel{}, grids, blocks, lds_bytes, kargs)); \ return ave_time; +#define MOE_SORTING_DISPATCH(unroll_num_) \ + if(a.num_experts <= 8) \ + { \ + MOE_SORTING_DISPATCH_ETILE(unroll_num_, 8) \ + } \ + else if(a.num_experts <= 16) \ + { \ + MOE_SORTING_DISPATCH_ETILE(unroll_num_, 16) \ + } \ + else if(a.num_experts <= 32) \ + { \ + MOE_SORTING_DISPATCH_ETILE(unroll_num_, 32) \ + } \ + else if(a.num_experts <= 64) \ + { \ + MOE_SORTING_DISPATCH_ETILE(unroll_num_, 64) \ + } \ + else \ + { \ + MOE_SORTING_DISPATCH_ETILE(unroll_num_, 0) \ + } + float fused_moesorting(fused_moesorting_trait t, fused_moesorting_args a, ck_tile::stream_config s) { if(t.weight_type == "fp32" && t.index_type == "int32") @@ -49,21 +73,12 @@ float fused_moesorting(fused_moesorting_trait t, fused_moesorting_args a, ck_til case(6): { MOE_SORTING_DISPATCH(6); } - case(7): { - MOE_SORTING_DISPATCH(7); - } case(8): { MOE_SORTING_DISPATCH(8); } - case(9): { - MOE_SORTING_DISPATCH(9); - } case(10): { MOE_SORTING_DISPATCH(10); } - case(11): { - MOE_SORTING_DISPATCH(11); - } default: { MOE_SORTING_DISPATCH(4); } diff --git a/include/ck_tile/ops/fused_moe/kernel/moe_sorting_kernel.hpp b/include/ck_tile/ops/fused_moe/kernel/moe_sorting_kernel.hpp index d9e28ceb52..30e68996b6 100644 --- a/include/ck_tile/ops/fused_moe/kernel/moe_sorting_kernel.hpp +++ b/include/ck_tile/ops/fused_moe/kernel/moe_sorting_kernel.hpp @@ -130,7 +130,8 @@ struct MoeSortingKernel CK_TILE_HOST static constexpr auto GetSmemSize(const Hargs& h) { const auto blocks = BlockSize(h); - return ((blocks.x + 1) * h.num_experts + (h.num_experts + 1)) * sizeof(index_t); + // usually num_experts is power of 2, we pad 1 dword here for the row-size + return ((blocks.x + 1) * (h.num_experts + 1) + (h.num_experts + 1)) * sizeof(index_t); } CK_TILE_HOST static constexpr auto MakeKargs(const Hargs& h) @@ -154,6 +155,75 @@ struct MoeSortingKernel return k; } + // [a, b, c, d....] -> [a, a+b, a+b+c, a+b+c+d, ....] + template + __device__ inline void wave_cumsum(data_t& thread_data) const + { + // wave_size must be power of 2 + constexpr int row_mask = 0xf; + constexpr int bank_mask = 0xf; + constexpr bool bound_ctrl = true; // ! out-of-bound is zero ! + auto reduce_op = [&](auto x_, auto y_) { return x_ + y_; }; + + if constexpr(wave_size > 1) + { + thread_data = reduce_op( + thread_data, + __builtin_bit_cast(data_t, __builtin_amdgcn_mov_dpp(__builtin_bit_cast(int, thread_data), + 0x111, + row_mask, + bank_mask, + bound_ctrl))); // row_shr:1 + } + + if constexpr(wave_size > 2) + { + thread_data = reduce_op( + thread_data, + __builtin_bit_cast(data_t, __builtin_amdgcn_mov_dpp(__builtin_bit_cast(int, thread_data), + 0x112, + row_mask, + bank_mask, + bound_ctrl))); // row_shr:2 + } + if constexpr(wave_size > 4) + { + thread_data = + reduce_op(thread_data, + __builtin_bit_cast(data_t, __builtin_amdgcn_mov_dpp(__builtin_bit_cast(int, thread_data), + 0x114, + row_mask, + bank_mask, + bound_ctrl))); // row_shr:4 + } + if constexpr(wave_size > 8) + { + thread_data = + reduce_op(thread_data, + __builtin_bit_cast(data_t, __builtin_amdgcn_mov_dpp(__builtin_bit_cast(int, thread_data), + 0x118, + row_mask, + bank_mask, + bound_ctrl))); // row_shr:8 + } + + if constexpr(wave_size > 16) + { + // now row-0, row-0+row-1, row-1+row-2, row-2+row-3 + int v_remote_tmp = __builtin_amdgcn_ds_bpermute(((__lane_id() & 0x30) - 1) << 2, __builtin_bit_cast(int, thread_data)); + v_remote_tmp = __lane_id() >= 16 ? v_remote_tmp : 0; + thread_data = reduce_op(thread_data, __builtin_bit_cast(data_t, v_remote_tmp)); + } + + if constexpr(wave_size > 32) + { + // lane-id 48...63->31 + int v_remote_tmp = __builtin_amdgcn_ds_bpermute(((__lane_id() & 0x30) - 17) << 2, __builtin_bit_cast(int, thread_data)); + v_remote_tmp = __lane_id() >= 32 ? v_remote_tmp : 0; + thread_data = reduce_op(thread_data, __builtin_bit_cast(data_t, v_remote_tmp)); + } + } + CK_TILE_DEVICE index_t calc_index(index_t total_col, index_t row, index_t col) const { return row * total_col + col; @@ -187,48 +257,124 @@ struct MoeSortingKernel index_t* shared_mem = reinterpret_cast(smem); index_t* tokens_cnts = shared_mem; // 2d: (blockDim.x + 1, num_experts) - index_t* cumsum = shared_mem + (blockDim.x + 1) * num_experts; // 1: (num_experts + 1) + index_t* cumsum = shared_mem + (blockDim.x + 1) * (num_experts+1); // 1: (num_experts + 1) + for(int i = 0; i < num_experts; ++i) { - tokens_cnts[calc_index(num_experts, tid + 1, i)] = 0; + tokens_cnts[calc_index(num_experts+1, tid + 1, i)] = 0; } + #pragma unroll Problem_::InternalLoadUnroll for(int i = start_idx; i < numel && i < start_idx + tokens_per_thread; ++i) { - ++tokens_cnts[calc_index(num_experts, tid + 1, topk_id[i])]; + ++tokens_cnts[calc_index(num_experts+1, tid + 1, topk_id[i])]; } __syncthreads(); +#if 1 if(tid < num_experts) { - tokens_cnts[calc_index(num_experts, 0, tid)] = 0; - for(int i = 1; i <= static_cast(blockDim.x); ++i) + tokens_cnts[calc_index(num_experts+1, 0, tid)] = 0; + index_t local_c[8]; + index_t prev_c = 0; + // TODO: manually unroll. pragma unroll does not work well when we have dependency + for(int i = 1; i <= static_cast(blockDim.x); i+= 8) { - tokens_cnts[calc_index(num_experts, i, tid)] += - tokens_cnts[calc_index(num_experts, i - 1, tid)]; + local_c[0] = tokens_cnts[calc_index(num_experts+1, i + 0, tid)]; + local_c[1] = tokens_cnts[calc_index(num_experts+1, i + 1, tid)]; + local_c[2] = tokens_cnts[calc_index(num_experts+1, i + 2, tid)]; + local_c[3] = tokens_cnts[calc_index(num_experts+1, i + 3, tid)]; + local_c[4] = tokens_cnts[calc_index(num_experts+1, i + 4, tid)]; + local_c[5] = tokens_cnts[calc_index(num_experts+1, i + 5, tid)]; + local_c[6] = tokens_cnts[calc_index(num_experts+1, i + 6, tid)]; + local_c[7] = tokens_cnts[calc_index(num_experts+1, i + 7, tid)]; + + local_c[0] += prev_c; + local_c[1] += local_c[0]; + local_c[2] += local_c[1]; + local_c[3] += local_c[2]; + local_c[4] += local_c[3]; + local_c[5] += local_c[4]; + local_c[6] += local_c[5]; + local_c[7] += local_c[6]; + prev_c = local_c[7]; + + tokens_cnts[calc_index(num_experts+1, i + 0, tid)] = local_c[0]; + tokens_cnts[calc_index(num_experts+1, i + 1, tid)] = local_c[1]; + tokens_cnts[calc_index(num_experts+1, i + 2, tid)] = local_c[2]; + tokens_cnts[calc_index(num_experts+1, i + 3, tid)] = local_c[3]; + tokens_cnts[calc_index(num_experts+1, i + 4, tid)] = local_c[4]; + tokens_cnts[calc_index(num_experts+1, i + 5, tid)] = local_c[5]; + tokens_cnts[calc_index(num_experts+1, i + 6, tid)] = local_c[6]; + tokens_cnts[calc_index(num_experts+1, i + 7, tid)] = local_c[7]; + } + } +#else + // TODO: below code still working, but slow in expert=32/topk=5 case. Put here for future heuristic + { + if(tid < num_experts) + tokens_cnts[calc_index(num_experts+1, 0, tid)] = 0; + for(int i = 0; i < num_experts; i+=8) { + index_t local_c[8]; + #pragma unroll + for(int j = 0; j < 8; j++) { + local_c[j] = tokens_cnts[calc_index(num_experts+1, tid+1, i+j)]; + } + + #pragma unroll + for(int j = 0; j < 8; j++) { + wave_cumsum(local_c[j]); + } + + #pragma unroll + for(int j = 0; j < 8; j++) { + tokens_cnts[calc_index(num_experts+1, tid+1, i+j)] = local_c[j]; + } + } + } +#endif + + __syncthreads(); + if constexpr (Problem::ExpertTile == 0) { + if(tid == 0) + { + cumsum[0] = 0; + for(int i = 1; i <= num_experts; ++i) + { + auto current_units = [&]() { + index_t x_ = tokens_cnts[calc_index(num_experts+1, blockDim.x, i - 1)] + + unit_size_mdiv.divisor - 1; + index_t y_ = unit_size_mdiv.div(x_); + return max(y_, 1) * unit_size_mdiv.divisor; + }(); + cumsum[i] = cumsum[i - 1] + current_units; + } + *p_total_tokens_post_pad = cumsum[num_experts]; + } + } else { + // TODO: we have out-of-bound read here. But result is still OK (will ignore tid >= expert) + // for simplicity, not check experts here. + int local_cnt = tokens_cnts[calc_index(num_experts+1, blockDim.x, tid)]; + int blocks_pers_expert = unit_size_mdiv.div(local_cnt + unit_size_mdiv.divisor - 1); + int padded_tokens_per_expert = max(blocks_pers_expert, 1) * unit_size_mdiv.divisor; + int local_cumsum = padded_tokens_per_expert; + wave_cumsum(local_cumsum); + + if(tid == (num_experts - 1)) { + cumsum[0] = 0; + *p_total_tokens_post_pad = local_cumsum; + } + if(tid < num_experts) { + cumsum[tid + 1] = local_cumsum; } } - // __syncthreads(); - if(tid == 0) - { - cumsum[0] = 0; - for(int i = 1; i <= num_experts; ++i) - { - auto current_units = [&]() { - index_t x_ = tokens_cnts[calc_index(num_experts, blockDim.x, i - 1)] + - unit_size_mdiv.divisor - 1; - index_t y_ = unit_size_mdiv.div(x_); - return max(y_, 1) * unit_size_mdiv.divisor; - }(); - cumsum[i] = cumsum[i - 1] + current_units; - } - *p_total_tokens_post_pad = cumsum[num_experts]; - } __syncthreads(); if(tid < num_experts) { - for(int i = cumsum[tid]; i < cumsum[tid + 1]; i += unit_size_mdiv.divisor) + int e_start = cumsum[tid]; + int e_end = cumsum[tid + 1]; + for(int i = e_start; i < e_end; i += unit_size_mdiv.divisor) { p_sorted_expert_ids[unit_size_mdiv.div(i)] = tid; } @@ -238,8 +384,8 @@ struct MoeSortingKernel for(int i = start_idx; i < numel && i < start_idx + tokens_per_thread; ++i) { index_t expert_id = topk_id[i]; - index_t rank_post_pad = - tokens_cnts[calc_index(num_experts, tid, expert_id)] + cumsum[expert_id]; + index_t local_cnt = tokens_cnts[calc_index(num_experts+1, tid, expert_id)]; + index_t rank_post_pad = local_cnt + cumsum[expert_id]; #if CK_TILE_REFERENCE_MOE_SORTING_MOCK_ID uint32_t curr_token_id, curr_topk_id; topk_mdiv.divmod(i, curr_token_id, curr_topk_id); @@ -247,27 +393,54 @@ struct MoeSortingKernel #else p_sorted_token_ids[rank_post_pad] = topk_mdiv.div(i); #endif - p_sorted_weights[rank_post_pad] = weights[i]; - ++tokens_cnts[calc_index(num_experts, tid, expert_id)]; + p_sorted_weights[rank_post_pad] = weights[i]; + tokens_cnts[calc_index(num_experts+1, tid, expert_id)] = local_cnt+1; } - const index_t prefill_token = topk_mdiv.div(numel); - if(tid < num_experts) - { - index_t expert_offset = - cumsum[tid] + tokens_cnts[calc_index(num_experts, blockDim.x, tid)]; - while(expert_offset < cumsum[tid + 1]) + if constexpr (Problem::ExpertTile == 0) { + const index_t prefill_token = topk_mdiv.div(numel); + if(tid < num_experts) { + index_t expert_offset = + cumsum[tid] + tokens_cnts[calc_index(num_experts+1, blockDim.x, tid)]; + index_t expert_end = cumsum[tid + 1]; + while(expert_offset < expert_end) + { #if CK_TILE_REFERENCE_MOE_SORTING_MOCK_ID - p_sorted_token_ids[expert_offset] = - MOE_SORTING_MOCK_ID(prefill_token, topk_mdiv.divisor); + p_sorted_token_ids[expert_offset] = + MOE_SORTING_MOCK_ID(prefill_token, topk_mdiv.divisor); #else - p_sorted_token_ids[expert_offset] = prefill_token; + p_sorted_token_ids[expert_offset] = prefill_token; #endif - p_sorted_weights[expert_offset] = static_cast(0.0); - expert_offset++; + p_sorted_weights[expert_offset] = static_cast(0.0); + expert_offset++; + } } } + else { + const index_t prefill_token = topk_mdiv.div(numel); + // TODO: only support expert-tile like 8, 16, 32 + static constexpr index_t experts_per_wave = warpSize / Problem::ExpertTile; + { + index_t eid = tid / experts_per_wave; + index_t expert_offset = + cumsum[eid] + tokens_cnts[calc_index(num_experts+1, blockDim.x, eid)] + tid % experts_per_wave; + index_t expert_end = cumsum[eid + 1]; + if(eid < num_experts) { + while(expert_offset < expert_end) + { +#if CK_TILE_REFERENCE_MOE_SORTING_MOCK_ID + p_sorted_token_ids[expert_offset] = + MOE_SORTING_MOCK_ID(prefill_token, topk_mdiv.divisor); +#else + p_sorted_token_ids[expert_offset] = prefill_token; +#endif + p_sorted_weights[expert_offset] = static_cast(0.0); + expert_offset+=experts_per_wave; + } + } + } + } } CK_TILE_DEVICE void operator()(Kargs kargs) const diff --git a/include/ck_tile/ops/fused_moe/pipeline/moe_sorting_problem.hpp b/include/ck_tile/ops/fused_moe/pipeline/moe_sorting_problem.hpp index adde59e356..50005c4402 100644 --- a/include/ck_tile/ops/fused_moe/pipeline/moe_sorting_problem.hpp +++ b/include/ck_tile/ops/fused_moe/pipeline/moe_sorting_problem.hpp @@ -9,15 +9,20 @@ namespace ck_tile { -template +template struct MoeSortingProblem { // TODO: this kernel only support warp per row using WeightType = remove_cvref_t; using IndexType = remove_cvref_t; - static constexpr index_t WarpSize = get_warp_size(); - static constexpr index_t WarpsPerBlock = 1; - static constexpr index_t InternalLoadUnroll = InternalLoadUnroll_; + static constexpr index_t WarpSize = get_warp_size(); + static constexpr index_t WarpsPerBlock = 1; + static constexpr index_t InternalLoadUnroll = + InternalLoadUnroll_; // TODO: need better design(like tile size) + static constexpr index_t ExpertTile = ExpertTile_; // TODO: only used in store out }; } // namespace ck_tile From 4c2eff023a26821512a100171531dc8757ad0e8f Mon Sep 17 00:00:00 2001 From: Po Yen Chen Date: Wed, 25 Dec 2024 23:57:28 +0800 Subject: [PATCH 25/25] Correct the dtype checking logics (#1775) --- example/ck_tile/01_fmha/codegen/ops/fmha_fwd_splitkv.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/example/ck_tile/01_fmha/codegen/ops/fmha_fwd_splitkv.py b/example/ck_tile/01_fmha/codegen/ops/fmha_fwd_splitkv.py index df5b9cecc6..2f7edd5477 100644 --- a/example/ck_tile/01_fmha/codegen/ops/fmha_fwd_splitkv.py +++ b/example/ck_tile/01_fmha/codegen/ops/fmha_fwd_splitkv.py @@ -261,7 +261,7 @@ FMHA_FWD_SPLITKV_API_INNER_DISPATCH=""" {F_if}((t.is_group_mode == {F static_assert({F_bn1} % 32 == 0); if (t.has_lse) {{ - if constexpr (std::is_same_v<{F_dtype}, ck_tile::fp8_t>) {{ + if constexpr (std::is_same_v<{F_dtype}, FmhaFwdFp8>) {{ return -1; }} else {{ using traits2_ = fmha_fwd_splitkv_combine_traits_<{F_hdim}, {F_dtype}, {F_mode}, /*F_bn1=*/32, true, {F_squant}, {F_spad}, {F_dvpad}>; @@ -614,7 +614,7 @@ def get_fmha_fwd_splitkv_combine_tile_dict_from_dtype(dtype : str) -> Optional[d } elif dtype == 'fp8' or dtype == 'bf8': return { - '64' : FmhaFwdSplitKVCombineTileSize(32, -1), + '64' : FmhaFwdSplitKVCombineTileSize(32, -1), '128' : FmhaFwdSplitKVCombineTileSize(32, -1), '256' : FmhaFwdSplitKVCombineTileSize(32, -1), }