From 0606e5498e7aa085a91c083d9c49794d30d371dc Mon Sep 17 00:00:00 2001 From: Mateusz Ozga <110818320+mozga-amd@users.noreply.github.com> Date: Tue, 13 Aug 2024 16:15:47 +0200 Subject: [PATCH 01/20] Support large: 12d tensor size for reduction kenrel (#1465) --- example/12_reduce/reduce_blockwise.cpp | 29 ++++++++++++++++++- example/12_reduce/reduce_example_common.hpp | 5 ++-- .../gpu/device/impl/device_reduce_common.hpp | 6 ++-- .../device/impl/device_reduce_multiblock.hpp | 4 +-- .../device/impl/device_reduce_threadwise.hpp | 4 +-- .../impl/device_reduce_threadwise_multi_d.hpp | 2 +- 6 files changed, 39 insertions(+), 11 deletions(-) diff --git a/example/12_reduce/reduce_blockwise.cpp b/example/12_reduce/reduce_blockwise.cpp index 9a736d4cfa..309100cdde 100644 --- a/example/12_reduce/reduce_blockwise.cpp +++ b/example/12_reduce/reduce_blockwise.cpp @@ -1,5 +1,5 @@ // SPDX-License-Identifier: MIT -// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved. +// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved. #include #include @@ -255,34 +255,61 @@ int main(int argc, char* argv[]) else { // for testing half_t + pass = + pass && reduce_blockwise_test( + true, 2, true, {3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3}, {0, 1, 2}, 1.0f, 0.0f); pass = pass && reduce_blockwise_test( true, 2, true, {16, 64, 32, 960}, {0, 1, 2}, 1.0f, 0.0f); // for testing float + pass = + pass && reduce_blockwise_test( + true, 2, true, {3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3}, {0, 1, 2}, 1.0f, 0.0f); + pass = pass && reduce_blockwise_test( true, 2, true, {16, 64, 32, 960}, {0, 1, 2}, 1.0f, 0.0f); // for testing double + pass = + pass && reduce_blockwise_test( + true, 2, true, {3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3}, {0, 1, 2}, 1.0f, 0.0f); + pass = pass && reduce_blockwise_test( true, 2, true, {16, 64, 32, 960}, {0, 1, 2}, 1.0f, 0.0f); // for testing bhalf_t + pass = pass && + reduce_blockwise_test( + true, 2, true, {3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3}, {0, 1, 2}, 1.0f, 0.0f); + pass = pass && reduce_blockwise_test( true, 2, true, {16, 64, 32, 960}, {0, 1, 2}, 1.0f, 0.0f); // for testing int8_t + pass = + pass && reduce_blockwise_test( + true, 2, true, {3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3}, {0, 1, 2}, 1.0f, 0.0f); + pass = pass && reduce_blockwise_test( true, 2, true, {16, 64, 32, 960}, {0, 1, 2}, 1.0f, 0.0f); #ifdef CK_EXPERIMENTAL_BIT_INT_EXTENSION_INT4 // for testing int4_t using AVG operation + pass = + pass && reduce_blockwise_test( + true, 2, true, {3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3}, {0, 1, 2}, 1.0f, 0.0f); + pass = pass && reduce_blockwise_test( true, 2, true, {16, 64, 32, 960}, {0, 1, 2}, 1.0f, 0.0f); // for testing int4_t using MAX operation + pass = + pass && reduce_blockwise_test( + true, 2, true, {3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3}, {0, 1, 2}, 1.0f, 0.0f); + pass = pass && reduce_blockwise_test( true, 2, true, {16, 64, 32, 960}, {0, 1, 2}, 1.0f, 0.0f); #endif diff --git a/example/12_reduce/reduce_example_common.hpp b/example/12_reduce/reduce_example_common.hpp index 5f9a48804a..08cd6e7ff9 100644 --- a/example/12_reduce/reduce_example_common.hpp +++ b/example/12_reduce/reduce_example_common.hpp @@ -1,5 +1,5 @@ // SPDX-License-Identifier: MIT -// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved. +// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved. #pragma once @@ -38,7 +38,8 @@ struct ReduceShape static constexpr ck::index_t NumReduceDim_ = NumReduceDim; }; -using reduce_shape_instances = std::tuple, +using reduce_shape_instances = std::tuple, + ReduceShape<3, 1>, ReduceShape<3, 2>, ReduceShape<4, 1>, ReduceShape<4, 2>, diff --git a/include/ck/tensor_operation/gpu/device/impl/device_reduce_common.hpp b/include/ck/tensor_operation/gpu/device/impl/device_reduce_common.hpp index 2481c5c769..67956d9f3f 100644 --- a/include/ck/tensor_operation/gpu/device/impl/device_reduce_common.hpp +++ b/include/ck/tensor_operation/gpu/device/impl/device_reduce_common.hpp @@ -1,5 +1,5 @@ // SPDX-License-Identifier: MIT -// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved. +// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved. #pragma once @@ -19,7 +19,7 @@ namespace device { template std::pair get_2d_lengths(const std::vector& inLengths) { - static_assert(Rank <= 6, "bigger Rank size not supported!"); + static_assert(Rank <= 12, "bigger Rank size not supported!"); long_index_t invariant_total_length = 1; long_index_t reduce_total_length = 1; @@ -38,7 +38,7 @@ std::pair get_2d_lengths(const std::vector& template std::pair get_2d_lengths(const std::array& inLengths) { - static_assert(Rank <= 6, "bigger Rank size not supported!"); + static_assert(Rank <= 12, "bigger Rank size not supported!"); long_index_t invariant_total_length = 1; long_index_t reduce_total_length = 1; diff --git a/include/ck/tensor_operation/gpu/device/impl/device_reduce_multiblock.hpp b/include/ck/tensor_operation/gpu/device/impl/device_reduce_multiblock.hpp index bf3deeb57a..b4873e3403 100644 --- a/include/ck/tensor_operation/gpu/device/impl/device_reduce_multiblock.hpp +++ b/include/ck/tensor_operation/gpu/device/impl/device_reduce_multiblock.hpp @@ -1,5 +1,5 @@ // SPDX-License-Identifier: MIT -// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved. +// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved. #pragma once @@ -51,7 +51,7 @@ struct DeviceReduceMultiBlock : public DeviceReduce { - static_assert(Rank <= 6, "Bigger Rank size is not supported!"); + static_assert(Rank <= 12, "Bigger Rank size is not supported!"); static_assert(BlockSize == MThreadClusterSize * KThreadClusterSize, "Invalid thread cluster size assignments!"); diff --git a/include/ck/tensor_operation/gpu/device/impl/device_reduce_threadwise.hpp b/include/ck/tensor_operation/gpu/device/impl/device_reduce_threadwise.hpp index 609c4c2f5b..8291575fb8 100644 --- a/include/ck/tensor_operation/gpu/device/impl/device_reduce_threadwise.hpp +++ b/include/ck/tensor_operation/gpu/device/impl/device_reduce_threadwise.hpp @@ -1,5 +1,5 @@ // SPDX-License-Identifier: MIT -// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved. +// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved. #pragma once @@ -47,7 +47,7 @@ struct DeviceReduceThreadWise : public DeviceReduce { - static_assert(Rank <= 6, "Bigger Rank size is not supported!"); + static_assert(Rank <= 12, "Bigger Rank size is not supported!"); static_assert(((InSrcVectorDim == 0 && MThreadSliceSize % InSrcVectorSize == 0) || (InSrcVectorDim == 1 && KThreadSliceSize % InSrcVectorSize == 0)) && diff --git a/include/ck/tensor_operation/gpu/device/impl/device_reduce_threadwise_multi_d.hpp b/include/ck/tensor_operation/gpu/device/impl/device_reduce_threadwise_multi_d.hpp index 75abb4d2e4..764b9312f3 100644 --- a/include/ck/tensor_operation/gpu/device/impl/device_reduce_threadwise_multi_d.hpp +++ b/include/ck/tensor_operation/gpu/device/impl/device_reduce_threadwise_multi_d.hpp @@ -45,7 +45,7 @@ struct DeviceReduceThreadWiseMultiD : public DeviceReduceMultiD { - static_assert(Rank <= 6, "Bigger Rank size is not supported!"); + static_assert(Rank <= 12, "Bigger Rank size is not supported!"); static_assert(((InSrcVectorDim == 0 && MThreadSliceSize % InSrcVectorSize == 0) || (InSrcVectorDim == 1 && KThreadSliceSize % InSrcVectorSize == 0)) && From 50c423481b9eaeaa95ea8c99d77c1aca2257fdc1 Mon Sep 17 00:00:00 2001 From: AngryLoki Date: Wed, 14 Aug 2024 03:31:15 +0800 Subject: [PATCH 02/20] Fix compilation errors with libc++ (#1461) This fixes 2 issues when compiled with libc++. First issue is attempt to call std::numeric_limits>::min(). _Float16 is extension of libstdc++, it does not exist in C++ standard[2]. Luckily, there is NumericLimits class in composable_kernel, which does everything needed. Second issue with call to 'check_err' is ambiguous: there are 2 candidates. It happens because composable_kernel relies on idea that f8_t (defined as _BitInt(8)) does not pass is_integral trait. However, libc++ treats _BitInt(N) as integral (per standard "any implementation-defined extended integer types" can be integral). Closes: #1460 Signed-off-by: Sv. Lockal --- library/include/ck/library/utility/check_err.hpp | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/library/include/ck/library/utility/check_err.hpp b/library/include/ck/library/utility/check_err.hpp index a3df884eee..9f4212ebd4 100644 --- a/library/include/ck/library/utility/check_err.hpp +++ b/library/include/ck/library/utility/check_err.hpp @@ -146,7 +146,7 @@ check_err(const Range& out, bool res{true}; int err_count = 0; double err = 0; - double max_err = std::numeric_limits>::min(); + double max_err = NumericLimits>::Min(); for(std::size_t i = 0; i < ref.size(); ++i) { const double o = type_convert(*std::next(std::begin(out), i)); @@ -178,7 +178,9 @@ check_err(const Range& out, template std::enable_if_t<(std::is_same_v, ranges::range_value_t> && std::is_integral_v> && - !std::is_same_v, bhalf_t>) + !std::is_same_v, bhalf_t> && + !std::is_same_v, f8_t> && + !std::is_same_v, bf8_t>) #ifdef CK_EXPERIMENTAL_BIT_INT_EXTENSION_INT4 || std::is_same_v, int4_t> #endif From 3049b5467c7dc5ccab49697eed4f20375ead92f8 Mon Sep 17 00:00:00 2001 From: Haocong WANG Date: Wed, 14 Aug 2024 10:42:30 +0800 Subject: [PATCH 03/20] [GEMM] gemm_universal related optimization (#1453) * replace buffer_atomic with global_atomic * fixed global_atomic_add * added bf16 atomic_add * format * clang-format-12 * clean * clean * add guards * Update gtest.cmake * enabled splitk_gemm_multi_d * format * add ckProfiler * format * fixed naming * format * clean * clean * add guards * fix clang format * format * add kbatch printout * clean * Add rocm6.2 related gemm optimization * Limit bf16 atomic usage * remove redundant RCR gemm_universal instance * Add RRR fp8 gemm universal instance * Bug fix * Add GPU_TARGET guard to FP8/BF8 target * bug fix * update cmake * remove all fp8/bf8 example if arch not support * Enable fp8 RRR support in ckProfiler * limit greedy-reverse flag to gemm_universal in ckProfiler --------- Co-authored-by: Jing Zhang Co-authored-by: Jing Zhang Co-authored-by: zjing14 Co-authored-by: Illia Silin <98187287+illsilin@users.noreply.github.com> Co-authored-by: illsilin --- CMakeLists.txt | 13 +- .../07_grouped_convnd_fwd/CMakeLists.txt | 6 +- .../10_grouped_convnd_bwd_data/CMakeLists.txt | 6 +- .../11_grouped_conv_bwd_weight/CMakeLists.txt | 7 +- client_example/16_convnd_fwd/CMakeLists.txt | 2 +- client_example/20_splitk_gemm/CMakeLists.txt | 2 +- client_example/CMakeLists.txt | 13 +- example/20_grouped_conv_bwd_weight/common.hpp | 8 +- .../gemm_add_add_xdl_fp16.cpp | 1 + .../gemm_multiply_multiply_xdl_fp8.cpp | 18 +- example/CMakeLists.txt | 14 + include/ck/host_utility/device_prop.hpp | 6 + .../gpu/device/device_gemm_multiple_d.hpp | 43 ++ ...device_gemm_multiple_d_xdl_cshuffle_v3.hpp | 385 +++++++++++++----- .../impl/device_gemm_xdl_cshuffle_v3.hpp | 304 +++++++------- .../grid/gridwise_gemm_xdl_cshuffle_v3.hpp | 18 +- .../gridwise_gemm_xdl_cshuffle_v3_multi_d.hpp | 28 +- include/ck/utility/amd_buffer_addressing.hpp | 57 ++- include/ck/utility/dynamic_buffer.hpp | 6 +- .../gpu/gemm_multiply_multiply.hpp | 217 ++++------ .../gpu/gemm_universal.hpp | 211 ++++------ .../gpu/gemm/CMakeLists.txt | 22 +- .../gpu/gemm_ab_scale/CMakeLists.txt | 5 + .../gpu/gemm_multiply_multiply/CMakeLists.txt | 7 +- ...tiply_multiply_xdl_f8_f8_bf16_mk_nk_mn.hpp | 3 +- ...f8_bf16_mk_nk_mn_comp_default_instance.cpp | 22 +- ...8_bf16_mk_nk_mn_comp_kpadding_instance.cpp | 22 +- ...bf16_mk_nk_mn_comp_mnkpadding_instance.cpp | 32 -- ..._bf16_mk_nk_mn_comp_mnpadding_instance.cpp | 32 -- ..._bf16_mk_nk_mn_mem_v1_default_instance.cpp | 22 +- ...bf16_mk_nk_mn_mem_v1_kpadding_instance.cpp | 22 +- ...16_mk_nk_mn_mem_v1_mnkpadding_instance.cpp | 33 -- ..._bf16_mk_nk_mn_mem_v2_default_instance.cpp | 22 +- ...bf16_mk_nk_mn_mem_v2_kpadding_instance.cpp | 22 +- ...16_mk_nk_mn_mem_v2_mnkpadding_instance.cpp | 33 -- .../gpu/gemm_universal/CMakeLists.txt | 149 ++++--- ...bf16_mk_nk_mn_comp_mnkpadding_instance.cpp | 24 -- ..._bf16_mk_nk_mn_comp_mnpadding_instance.cpp | 24 -- ...16_mk_nk_mn_mem_v1_mnkpadding_instance.cpp | 25 -- ...16_mk_nk_mn_mem_v2_mnkpadding_instance.cpp | 25 -- ..._f16_mk_nk_mn_comp_mnkpadding_instance.cpp | 23 -- ...6_f16_mk_nk_mn_comp_mnpadding_instance.cpp | 23 -- ...16_mk_nk_mn_mem_v1_mnkpadding_instance.cpp | 24 -- ...16_mk_nk_mn_mem_v2_mnkpadding_instance.cpp | 24 -- ..._f16_mk_nk_mn_comp_mnkpadding_instance.cpp | 23 -- ...8_f16_mk_nk_mn_comp_mnpadding_instance.cpp | 26 -- ...16_mk_nk_mn_mem_v1_mnkpadding_instance.cpp | 24 -- ...16_mk_nk_mn_mem_v2_mnkpadding_instance.cpp | 24 -- ..._f16_mk_nk_mn_comp_mnkpadding_instance.cpp | 23 -- ...6_f16_mk_nk_mn_comp_mnpadding_instance.cpp | 26 -- ...16_mk_nk_mn_mem_v1_mnkpadding_instance.cpp | 24 -- ...16_mk_nk_mn_mem_v2_mnkpadding_instance.cpp | 24 -- ...gemm_xdl_universal_f8_f8_bf16_mk_kn_mn.hpp | 96 +++++ ...f8_bf16_mk_kn_mn_comp_default_instance.cpp | 23 ++ ...8_bf16_mk_kn_mn_comp_kpadding_instance.cpp | 23 ++ ...bf16_mk_kn_mn_comp_nkpadding_instance.cpp} | 8 +- ...bf16_mk_kn_mn_mem_v1_default_instance.cpp} | 8 +- ...f16_mk_kn_mn_mem_v1_kpadding_instance.cpp} | 9 +- ...16_mk_kn_mn_mem_v1_nkpadding_instance.cpp} | 8 +- ..._bf16_mk_kn_mn_mem_v2_default_instance.cpp | 24 ++ ...bf16_mk_kn_mn_mem_v2_kpadding_instance.cpp | 24 ++ ...f16_mk_kn_mn_mem_v2_nkpadding_instance.cpp | 24 ++ ...gemm_xdl_universal_f8_f8_bf16_mk_nk_mn.hpp | 11 +- .../grouped_conv3d_bwd_data/CMakeLists.txt | 2 +- .../grouped_conv3d_bwd_weight/CMakeLists.txt | 2 +- .../CMakeLists.txt | 2 +- .../CMakeLists.txt | 2 +- .../gpu/grouped_conv3d_fwd/CMakeLists.txt | 8 +- .../profile_gemm_multiply_multiply_impl.hpp | 193 +++++---- .../profiler/profile_gemm_universal_impl.hpp | 4 +- .../src/profile_gemm_multiply_multiply.cpp | 20 +- profiler/src/profile_gemm_universal.cpp | 4 + .../src/profile_grouped_gemm_fixed_nk.cpp | 8 +- .../test_gemm_universal_xdl.cpp | 9 +- 74 files changed, 1336 insertions(+), 1375 deletions(-) delete mode 100644 library/src/tensor_operation_instance/gpu/gemm_multiply_multiply/device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_comp_mnkpadding_instance.cpp delete mode 100644 library/src/tensor_operation_instance/gpu/gemm_multiply_multiply/device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_comp_mnpadding_instance.cpp delete mode 100644 library/src/tensor_operation_instance/gpu/gemm_multiply_multiply/device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_mem_v1_mnkpadding_instance.cpp delete mode 100644 library/src/tensor_operation_instance/gpu/gemm_multiply_multiply/device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_mem_v2_mnkpadding_instance.cpp delete mode 100644 library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_bf16_bf16_bf16/device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_comp_mnkpadding_instance.cpp delete mode 100644 library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_bf16_bf16_bf16/device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_comp_mnpadding_instance.cpp delete mode 100644 library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_bf16_bf16_bf16/device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_mem_v1_mnkpadding_instance.cpp delete mode 100644 library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_bf16_bf16_bf16/device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_mem_v2_mnkpadding_instance.cpp delete mode 100644 library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f16_f16_f16/device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_comp_mnkpadding_instance.cpp delete mode 100644 library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f16_f16_f16/device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_comp_mnpadding_instance.cpp delete mode 100644 library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f16_f16_f16/device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_mem_v1_mnkpadding_instance.cpp delete mode 100644 library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f16_f16_f16/device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_mem_v2_mnkpadding_instance.cpp delete mode 100644 library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f16_f8_f16/device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_comp_mnkpadding_instance.cpp delete mode 100644 library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f16_f8_f16/device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_comp_mnpadding_instance.cpp delete mode 100644 library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f16_f8_f16/device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_mem_v1_mnkpadding_instance.cpp delete mode 100644 library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f16_f8_f16/device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_mem_v2_mnkpadding_instance.cpp delete mode 100644 library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f16_f16/device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_comp_mnkpadding_instance.cpp delete mode 100644 library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f16_f16/device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_comp_mnpadding_instance.cpp delete mode 100644 library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f16_f16/device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_mem_v1_mnkpadding_instance.cpp delete mode 100644 library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f16_f16/device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_mem_v2_mnkpadding_instance.cpp create mode 100644 library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn.hpp create mode 100644 library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_comp_default_instance.cpp create mode 100644 library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_comp_kpadding_instance.cpp rename library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f8_bf16/{device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_comp_mnpadding_instance.cpp => device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_comp_nkpadding_instance.cpp} (59%) rename library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f8_bf16/{device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_mem_v1_mnkpadding_instance.cpp => device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_v1_default_instance.cpp} (58%) rename library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f8_bf16/{device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_comp_mnkpadding_instance.cpp => device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_v1_kpadding_instance.cpp} (58%) rename library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f8_bf16/{device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_mem_v2_mnkpadding_instance.cpp => device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_v1_nkpadding_instance.cpp} (58%) create mode 100644 library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_v2_default_instance.cpp create mode 100644 library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_v2_kpadding_instance.cpp create mode 100644 library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_v2_nkpadding_instance.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index 96a49b1c00..96c37ac943 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -62,8 +62,17 @@ if (DTYPES) endif() message("DTYPES macro set to ${DTYPES}") else() - add_definitions(-DCK_ENABLE_INT8 -DCK_ENABLE_FP8 -DCK_ENABLE_BF8 -DCK_ENABLE_FP16 -DCK_ENABLE_FP32 -DCK_ENABLE_FP64 -DCK_ENABLE_BF16) - set(CK_ENABLE_ALL_DTYPES "ON") + add_definitions(-DCK_ENABLE_INT8 -DCK_ENABLE_FP16 -DCK_ENABLE_FP32 -DCK_ENABLE_FP64 -DCK_ENABLE_BF16) + set(CK_ENABLE_INT8 "ON") + set(CK_ENABLE_FP16 "ON") + set(CK_ENABLE_FP32 "ON") + set(CK_ENABLE_FP64 "ON") + set(CK_ENABLE_BF16 "ON") + if (GPU_TARGETS MATCHES "gfx94") + add_definitions(-DCK_ENABLE_FP8 -DCK_ENABLE_BF8) + set(CK_ENABLE_FP8 "ON") + set(CK_ENABLE_BF8 "ON") + endif() endif() #for f8/bf8_t type diff --git a/client_example/07_grouped_convnd_fwd/CMakeLists.txt b/client_example/07_grouped_convnd_fwd/CMakeLists.txt index e8c046ff44..c953e21d02 100644 --- a/client_example/07_grouped_convnd_fwd/CMakeLists.txt +++ b/client_example/07_grouped_convnd_fwd/CMakeLists.txt @@ -5,17 +5,17 @@ if(GPU_TARGETS MATCHES "gfx9") add_executable(client_grouped_conv1d_fwd grouped_conv1d_fwd.cpp) target_link_libraries(client_grouped_conv1d_fwd PRIVATE composable_kernel::device_conv_operations) - if((DTYPES MATCHES "fp8") OR NOT DEFINED DTYPES) + if((DTYPES MATCHES "fp8") OR (NOT DEFINED DTYPES AND GPU_TARGETS MATCHES "gfx94")) add_executable(client_grouped_conv3d_fwd_fp8 grouped_conv3d_fwd_fp8.cpp) target_link_libraries(client_grouped_conv3d_fwd_fp8 PRIVATE composable_kernel::device_conv_operations) endif() - if((DTYPES MATCHES "bf8") OR NOT DEFINED DTYPES) + if((DTYPES MATCHES "bf8") OR (NOT DEFINED DTYPES AND GPU_TARGETS MATCHES "gfx94")) add_executable(client_grouped_conv3d_fwd_bf8 grouped_conv3d_fwd_bf8.cpp) target_link_libraries(client_grouped_conv3d_fwd_bf8 PRIVATE composable_kernel::device_conv_operations) endif() - if((DTYPES MATCHES "fp8" AND DTYPES MATCHES "bf8") OR NOT DEFINED DTYPES) + if((DTYPES MATCHES "fp8" AND DTYPES MATCHES "bf8") OR (NOT DEFINED DTYPES AND GPU_TARGETS MATCHES "gfx94")) add_executable(client_grouped_conv3d_fwd_fp8_bf8 grouped_conv3d_fwd_fp8_bf8.cpp) target_link_libraries(client_grouped_conv3d_fwd_fp8_bf8 PRIVATE composable_kernel::device_conv_operations) diff --git a/client_example/10_grouped_convnd_bwd_data/CMakeLists.txt b/client_example/10_grouped_convnd_bwd_data/CMakeLists.txt index 0cf308c6e1..d10c39ed80 100644 --- a/client_example/10_grouped_convnd_bwd_data/CMakeLists.txt +++ b/client_example/10_grouped_convnd_bwd_data/CMakeLists.txt @@ -4,5 +4,7 @@ target_link_libraries(client_grouped_conv2d_bwd_data PRIVATE composable_kernel:: add_executable(client_grouped_conv3d_bwd_data grouped_conv3d_bwd_data.cpp) target_link_libraries(client_grouped_conv3d_bwd_data PRIVATE composable_kernel::device_conv_operations) -add_executable(client_grouped_conv3d_bwd_data_input_fp16_comp_bf8f8 grouped_conv3d_bwd_data_input_fp16_comp_bf8f8.cpp) -target_link_libraries(client_grouped_conv3d_bwd_data_input_fp16_comp_bf8f8 PRIVATE composable_kernel::device_conv_operations) +if((DTYPES MATCHES "fp8" AND DTYPES MATCHES "bf8") OR (NOT DEFINED DTYPES AND GPU_TARGETS MATCHES "gfx94")) + add_executable(client_grouped_conv3d_bwd_data_input_fp16_comp_bf8f8 grouped_conv3d_bwd_data_input_fp16_comp_bf8f8.cpp) + target_link_libraries(client_grouped_conv3d_bwd_data_input_fp16_comp_bf8f8 PRIVATE composable_kernel::device_conv_operations) +endif() \ No newline at end of file diff --git a/client_example/11_grouped_conv_bwd_weight/CMakeLists.txt b/client_example/11_grouped_conv_bwd_weight/CMakeLists.txt index dddfabb787..60a6dc1021 100644 --- a/client_example/11_grouped_conv_bwd_weight/CMakeLists.txt +++ b/client_example/11_grouped_conv_bwd_weight/CMakeLists.txt @@ -2,10 +2,13 @@ add_executable(client_grouped_conv1d_bwd_weight_fp16 grouped_conv1d_bwd_weight_f add_executable(client_grouped_conv2d_bwd_weight_fp16 grouped_conv2d_bwd_weight_fp16.cpp) add_executable(client_grouped_conv3d_bwd_weight_fp16 grouped_conv3d_bwd_weight_fp16.cpp) add_executable(client_grouped_conv3d_bwd_weight_fp32 grouped_conv3d_bwd_weight_fp32.cpp) -add_executable(client_grouped_conv3d_bwd_weight_fp16_comp_bf8_fp8 grouped_conv3d_bwd_weight_fp16_comp_bf8_fp8.cpp) target_link_libraries(client_grouped_conv1d_bwd_weight_fp16 PRIVATE composable_kernel::device_conv_operations) target_link_libraries(client_grouped_conv2d_bwd_weight_fp16 PRIVATE composable_kernel::device_conv_operations) target_link_libraries(client_grouped_conv3d_bwd_weight_fp16 PRIVATE composable_kernel::device_conv_operations) target_link_libraries(client_grouped_conv3d_bwd_weight_fp32 PRIVATE composable_kernel::device_conv_operations) -target_link_libraries(client_grouped_conv3d_bwd_weight_fp16_comp_bf8_fp8 PRIVATE composable_kernel::device_conv_operations) + +if((DTYPES MATCHES "fp8" AND DTYPES MATCHES "bf8") OR (NOT DEFINED DTYPES AND GPU_TARGETS MATCHES "gfx94")) + add_executable(client_grouped_conv3d_bwd_weight_fp16_comp_bf8_fp8 grouped_conv3d_bwd_weight_fp16_comp_bf8_fp8.cpp) + target_link_libraries(client_grouped_conv3d_bwd_weight_fp16_comp_bf8_fp8 PRIVATE composable_kernel::device_conv_operations) +endif() \ No newline at end of file diff --git a/client_example/16_convnd_fwd/CMakeLists.txt b/client_example/16_convnd_fwd/CMakeLists.txt index 5279e3dfcf..8c1372e741 100644 --- a/client_example/16_convnd_fwd/CMakeLists.txt +++ b/client_example/16_convnd_fwd/CMakeLists.txt @@ -4,7 +4,7 @@ if((DTYPES MATCHES "fp16") OR NOT DEFINED DTYPES) endif() -if((DTYPES MATCHES "fp8") OR NOT DEFINED DTYPES) +if((DTYPES MATCHES "fp8") OR (NOT DEFINED DTYPES AND GPU_TARGETS MATCHES "gfx94")) add_executable(client_conv3d_fwd_fp16_comp_fp8 conv3d_fwd_fp16_comp_fp8.cpp) target_link_libraries(client_conv3d_fwd_fp16_comp_fp8 PRIVATE composable_kernel::device_conv_operations) endif() diff --git a/client_example/20_splitk_gemm/CMakeLists.txt b/client_example/20_splitk_gemm/CMakeLists.txt index 05fcaa8103..383c5d6630 100644 --- a/client_example/20_splitk_gemm/CMakeLists.txt +++ b/client_example/20_splitk_gemm/CMakeLists.txt @@ -1,4 +1,4 @@ -if(GPU_TARGETS MATCHES "gfx9" AND ((DTYPES MATCHES "fp8" AND DTYPES MATCHES "fp16") OR NOT DEFINED DTYPES)) +if((DTYPES MATCHES "fp8" AND DTYPES MATCHES "fp16") OR (NOT DEFINED DTYPES AND GPU_TARGETS MATCHES "gfx94")) add_executable(client_splitK_gemm splitK_gemm_fp16_f8.cpp) target_link_libraries(client_splitK_gemm PRIVATE composable_kernel::device_gemm_operations) endif() diff --git a/client_example/CMakeLists.txt b/client_example/CMakeLists.txt index d2222a840e..acb57d7bb0 100644 --- a/client_example/CMakeLists.txt +++ b/client_example/CMakeLists.txt @@ -34,8 +34,17 @@ if (DTYPES) endif() message("DTYPES macro set to ${DTYPES}") else() - add_definitions(-DCK_ENABLE_INT8 -DCK_ENABLE_FP8 -DCK_ENABLE_BF8 -DCK_ENABLE_FP16 -DCK_ENABLE_FP32 -DCK_ENABLE_FP64 -DCK_ENABLE_BF16) - set(CK_ENABLE_ALL_DTYPES "ON") + add_definitions(-DCK_ENABLE_INT8 -DCK_ENABLE_FP16 -DCK_ENABLE_FP32 -DCK_ENABLE_FP64 -DCK_ENABLE_BF16) + set(CK_ENABLE_INT8 "ON") + set(CK_ENABLE_FP16 "ON") + set(CK_ENABLE_FP32 "ON") + set(CK_ENABLE_FP64 "ON") + set(CK_ENABLE_BF16 "ON") + if (GPU_TARGETS MATCHES "gfx94") + add_definitions(-DCK_ENABLE_FP8 -DCK_ENABLE_BF8) + set(CK_ENABLE_FP8 "ON") + set(CK_ENABLE_BF8 "ON") + endif() endif() if (GPU_TARGETS) diff --git a/example/20_grouped_conv_bwd_weight/common.hpp b/example/20_grouped_conv_bwd_weight/common.hpp index 1d2206ff72..e0034bf7eb 100644 --- a/example/20_grouped_conv_bwd_weight/common.hpp +++ b/example/20_grouped_conv_bwd_weight/common.hpp @@ -23,12 +23,8 @@ using BF16 = ck::bhalf_t; using F16 = ck::half_t; using F32 = float; -#ifdef CK_ENABLE_FP8 -using F8 = ck::f8_t; -#endif -#ifdef CK_ENABLE_BF8 -using BF8 = ck::bf8_t; -#endif +using F8 = ck::f8_t; +using BF8 = ck::bf8_t; template using S = ck::Sequence; diff --git a/example/65_gemm_multiply_multiply/gemm_add_add_xdl_fp16.cpp b/example/65_gemm_multiply_multiply/gemm_add_add_xdl_fp16.cpp index 5fea43ffc3..580f38a79f 100644 --- a/example/65_gemm_multiply_multiply/gemm_add_add_xdl_fp16.cpp +++ b/example/65_gemm_multiply_multiply/gemm_add_add_xdl_fp16.cpp @@ -208,6 +208,7 @@ int main(int argc, char* argv[]) StrideB, std::array{StrideD, StrideD}, StrideE, + 1, a_element_op, b_element_op, cde_element_op); diff --git a/example/65_gemm_multiply_multiply/gemm_multiply_multiply_xdl_fp8.cpp b/example/65_gemm_multiply_multiply/gemm_multiply_multiply_xdl_fp8.cpp index b0b1aa73c1..cb4f60764e 100644 --- a/example/65_gemm_multiply_multiply/gemm_multiply_multiply_xdl_fp8.cpp +++ b/example/65_gemm_multiply_multiply/gemm_multiply_multiply_xdl_fp8.cpp @@ -69,7 +69,7 @@ using AElementOp = PassThrough; using BElementOp = PassThrough; using CDEElementOp = MultiplyMultiply; -static constexpr auto GemmSpec = ck::tensor_operation::device::GemmSpecialization::Default; +static constexpr auto GemmSpec = ck::tensor_operation::device::GemmSpecialization::MNPadding; using DeviceOpInstance = ck::tensor_operation::device::DeviceGemmMultiD_Xdl_CShuffle_V3 // clang-format off @@ -99,6 +99,8 @@ int main(int argc, char* argv[]) ck::index_t StrideD = 0; ck::index_t StrideE = N; + ck::index_t KBatch = 1; + if(argc == 1) { // use default case @@ -109,7 +111,7 @@ int main(int argc, char* argv[]) init_method = std::stoi(argv[2]); time_kernel = std::stoi(argv[3]); } - else if(argc == 11) + else if(argc == 12) { do_verification = std::stoi(argv[1]); init_method = std::stoi(argv[2]); @@ -123,13 +125,16 @@ int main(int argc, char* argv[]) StrideB = std::stoi(argv[8]); StrideD = std::stoi(argv[9]); StrideE = std::stoi(argv[10]); + + KBatch = std::stoi(argv[11]); } else { printf("arg1: verification (0=no, 1=yes)\n"); printf("arg2: initialization (0=no init, 1=integer value, 2=decimal value)\n"); printf("arg3: time kernel (0=no, 1=yes)\n"); - printf("arg4 to 9: M (256x), N(128x), K(32x), StrideA, StrideB, StrideD, StrideE\n"); + printf( + "arg4 to 9: M (256x), N(128x), K(32x), StrideA, StrideB, StrideD, StrideE, KBatch\n"); exit(0); } @@ -212,6 +217,7 @@ int main(int argc, char* argv[]) StrideB, std::array{I0, I0}, StrideE, + KBatch, a_element_op, b_element_op, cde_element_op); @@ -236,10 +242,12 @@ int main(int argc, char* argv[]) std::cout << "Perf: " << ave_time << " ms, " << tflops << " TFlops, " << gb_per_sec << " GB/s" << std::endl; - e_device_buf.FromDevice(e_m_n_device_result.mData.data()); - if(do_verification) { + invoker.Run(argument, StreamConfig{nullptr, false}); + + e_device_buf.FromDevice(e_m_n_device_result.mData.data()); + Tensor c_m_n({M, N}); using ReferenceGemmInstance = ck::tensor_operation::host::ReferenceGemm MakeInvokerPointer() = 0; }; +// GEMM: +// input : A[M, K], B[K, N], +// input : D0[M, N], D1[M, N], ... +// output : E[M, N] +// C = a_op(A) * b_op(B) +// E = cde_op(C, D0, D1, ...) +// Assume: +// D0, D1, ... and E have the same layout +template +struct DeviceGemmMultipleDSplitK : public BaseOperator +{ + static constexpr index_t NumDTensor = DsDataType::Size(); + + virtual std::unique_ptr + MakeArgumentPointer(const void* p_a, + const void* p_b, + std::array p_ds, + void* p_e, + ck::index_t M, + ck::index_t N, + ck::index_t K, + ck::index_t StrideA, + ck::index_t StrideB, + std::array StrideDs, + ck::index_t StrideE, + ck::index_t KBatch, + AElementwiseOperation a_element_op, + BElementwiseOperation b_element_op, + CDEElementwiseOperation cde_element_op) = 0; + + virtual std::unique_ptr MakeInvokerPointer() = 0; +}; + } // namespace device } // namespace tensor_operation } // namespace ck diff --git a/include/ck/tensor_operation/gpu/device/impl/device_gemm_multiple_d_xdl_cshuffle_v3.hpp b/include/ck/tensor_operation/gpu/device/impl/device_gemm_multiple_d_xdl_cshuffle_v3.hpp index 92aa47d53d..c0912726c2 100644 --- a/include/ck/tensor_operation/gpu/device/impl/device_gemm_multiple_d_xdl_cshuffle_v3.hpp +++ b/include/ck/tensor_operation/gpu/device/impl/device_gemm_multiple_d_xdl_cshuffle_v3.hpp @@ -69,17 +69,17 @@ template -struct DeviceGemmMultiD_Xdl_CShuffle_V3 : public DeviceGemmMultipleD +struct DeviceGemmMultiD_Xdl_CShuffle_V3 : public DeviceGemmMultipleDSplitK { static constexpr index_t NumDTensor = DsDataType::Size(); @@ -192,15 +192,11 @@ struct DeviceGemmMultiD_Xdl_CShuffle_V3 : public DeviceGemmMultipleD, bhalf_t>::value) - { - if(arg_.KBatch > 1) - hipGetErrorString( - hipMemsetAsync(arg_.p_c_grid, - 0, - arg_.M * arg_.N * sizeof(CDataType), - stream_config.stream_id_)); - } + if(arg_.KBatch > 1) + hipGetErrorString(hipMemsetAsync(arg_.p_c_grid, + 0, + arg_.M * arg_.N * sizeof(CDataType), + stream_config.stream_id_)); }; ave_time = ck::utility::launch_and_time_kernel_with_preprocess( @@ -234,38 +230,49 @@ struct DeviceGemmMultiD_Xdl_CShuffle_V3 : public DeviceGemmMultipleD 1) + { + const auto kernel = kernel_gemm_xdl_cshuffle_v3_multi_d< + GridwiseGemm, + true, + InMemoryDataOperationEnum::AtomicAdd, + minimum_occupancy>; + Run(kernel); + } + else { const auto kernel = - kernel_gemm_xdl_cshuffle_v3; + kernel_gemm_xdl_cshuffle_v3_multi_d; Run(kernel); } } // Tail number could be One to Seven else if constexpr(BlkGemmPipelineVer == BlockGemmPipelineVersion::v2) { + if(arg.KBatch > 1) { if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == TailNumber::One) { - const auto kernel = - kernel_gemm_xdl_cshuffle_v3; + const auto kernel = kernel_gemm_xdl_cshuffle_v3_multi_d< + GridwiseGemm, + true, + InMemoryDataOperationEnum::AtomicAdd, + minimum_occupancy, + TailNumber::One>; Run(kernel); } else if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == TailNumber::Full) { - const auto kernel = - kernel_gemm_xdl_cshuffle_v3; + const auto kernel = kernel_gemm_xdl_cshuffle_v3_multi_d< + GridwiseGemm, + true, + InMemoryDataOperationEnum::AtomicAdd, + minimum_occupancy, + TailNumber::Full>; Run(kernel); } @@ -273,12 +280,12 @@ struct DeviceGemmMultiD_Xdl_CShuffle_V3 : public DeviceGemmMultipleD; + const auto kernel = kernel_gemm_xdl_cshuffle_v3_multi_d< + GridwiseGemm, + true, + InMemoryDataOperationEnum::AtomicAdd, + minimum_occupancy, + TailNumber::Two>; Run(kernel); } } @@ -288,12 +295,12 @@ struct DeviceGemmMultiD_Xdl_CShuffle_V3 : public DeviceGemmMultipleD; + const auto kernel = kernel_gemm_xdl_cshuffle_v3_multi_d< + GridwiseGemm, + true, + InMemoryDataOperationEnum::AtomicAdd, + minimum_occupancy, + TailNumber::Three>; Run(kernel); } } @@ -303,12 +310,12 @@ struct DeviceGemmMultiD_Xdl_CShuffle_V3 : public DeviceGemmMultipleD; + const auto kernel = kernel_gemm_xdl_cshuffle_v3_multi_d< + GridwiseGemm, + true, + InMemoryDataOperationEnum::AtomicAdd, + minimum_occupancy, + TailNumber::Four>; Run(kernel); } } @@ -318,12 +325,12 @@ struct DeviceGemmMultiD_Xdl_CShuffle_V3 : public DeviceGemmMultipleD; + const auto kernel = kernel_gemm_xdl_cshuffle_v3_multi_d< + GridwiseGemm, + true, + InMemoryDataOperationEnum::AtomicAdd, + minimum_occupancy, + TailNumber::Five>; Run(kernel); } } @@ -332,12 +339,12 @@ struct DeviceGemmMultiD_Xdl_CShuffle_V3 : public DeviceGemmMultipleD; + const auto kernel = kernel_gemm_xdl_cshuffle_v3_multi_d< + GridwiseGemm, + true, + InMemoryDataOperationEnum::AtomicAdd, + minimum_occupancy, + TailNumber::Six>; Run(kernel); } } @@ -347,12 +354,124 @@ struct DeviceGemmMultiD_Xdl_CShuffle_V3 : public DeviceGemmMultipleD; + const auto kernel = kernel_gemm_xdl_cshuffle_v3_multi_d< + GridwiseGemm, + true, + InMemoryDataOperationEnum::AtomicAdd, + minimum_occupancy, + TailNumber::Seven>; + Run(kernel); + } + } + } + else + { + if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == TailNumber::One) + { + const auto kernel = + kernel_gemm_xdl_cshuffle_v3_multi_d; + Run(kernel); + } + else if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == + TailNumber::Full) + { + const auto kernel = + kernel_gemm_xdl_cshuffle_v3_multi_d; + Run(kernel); + } + + if constexpr(GridwiseGemm::BlockwiseGemmPipe::PrefetchStages > 2) + { + if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == TailNumber::Two) + { + const auto kernel = kernel_gemm_xdl_cshuffle_v3_multi_d< + GridwiseGemm, + true, + InMemoryDataOperationEnum::Set, + minimum_occupancy, + TailNumber::Two>; + Run(kernel); + } + } + + if constexpr(GridwiseGemm::BlockwiseGemmPipe::PrefetchStages > 3) + { + if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == + TailNumber::Three) + { + const auto kernel = kernel_gemm_xdl_cshuffle_v3_multi_d< + GridwiseGemm, + true, + InMemoryDataOperationEnum::Set, + minimum_occupancy, + TailNumber::Three>; + Run(kernel); + } + } + + if constexpr(GridwiseGemm::BlockwiseGemmPipe::PrefetchStages > 4) + { + if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == + TailNumber::Four) + { + const auto kernel = kernel_gemm_xdl_cshuffle_v3_multi_d< + GridwiseGemm, + true, + InMemoryDataOperationEnum::Set, + minimum_occupancy, + TailNumber::Four>; + Run(kernel); + } + } + + if constexpr(GridwiseGemm::BlockwiseGemmPipe::PrefetchStages > 5) + { + if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == + TailNumber::Five) + { + const auto kernel = kernel_gemm_xdl_cshuffle_v3_multi_d< + GridwiseGemm, + true, + InMemoryDataOperationEnum::Set, + minimum_occupancy, + TailNumber::Five>; + Run(kernel); + } + } + + if constexpr(GridwiseGemm::BlockwiseGemmPipe::PrefetchStages > 6) + { + if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == TailNumber::Six) + { + const auto kernel = kernel_gemm_xdl_cshuffle_v3_multi_d< + GridwiseGemm, + true, + InMemoryDataOperationEnum::Set, + minimum_occupancy, + TailNumber::Six>; + Run(kernel); + } + } + + if constexpr(GridwiseGemm::BlockwiseGemmPipe::PrefetchStages > 7) + { + if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == + TailNumber::Seven) + { + const auto kernel = kernel_gemm_xdl_cshuffle_v3_multi_d< + GridwiseGemm, + true, + InMemoryDataOperationEnum::Set, + minimum_occupancy, + TailNumber::Seven>; Run(kernel); } } @@ -361,51 +480,98 @@ struct DeviceGemmMultiD_Xdl_CShuffle_V3 : public DeviceGemmMultipleD 1) { if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == TailNumber::Odd) { - const auto kernel = - kernel_gemm_xdl_cshuffle_v3_2lds; + const auto kernel = kernel_gemm_xdl_cshuffle_v3_multi_d_2lds< + GridwiseGemm, + true, + InMemoryDataOperationEnum::AtomicAdd, + minimum_occupancy, + TailNumber::Odd>; Run(kernel); } else { - const auto kernel = - kernel_gemm_xdl_cshuffle_v3_2lds; + const auto kernel = kernel_gemm_xdl_cshuffle_v3_multi_d_2lds< + GridwiseGemm, + true, + InMemoryDataOperationEnum::AtomicAdd, + minimum_occupancy, + TailNumber::Even>; + Run(kernel); + } + } + else + { + if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == TailNumber::Odd) + { + const auto kernel = kernel_gemm_xdl_cshuffle_v3_multi_d_2lds< + GridwiseGemm, + true, + InMemoryDataOperationEnum::Set, + minimum_occupancy, + TailNumber::Odd>; + Run(kernel); + } + else + { + const auto kernel = kernel_gemm_xdl_cshuffle_v3_multi_d_2lds< + GridwiseGemm, + true, + InMemoryDataOperationEnum::Set, + minimum_occupancy, + TailNumber::Even>; Run(kernel); } } } else { + if(arg.KBatch > 1) + { + if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == TailNumber::Odd) + { + const auto kernel = kernel_gemm_xdl_cshuffle_v3_multi_d< + GridwiseGemm, + true, + InMemoryDataOperationEnum::AtomicAdd, + minimum_occupancy, + TailNumber::Odd>; + Run(kernel); + } + else + { + const auto kernel = kernel_gemm_xdl_cshuffle_v3_multi_d< + GridwiseGemm, + true, + InMemoryDataOperationEnum::AtomicAdd, + minimum_occupancy, + TailNumber::Even>; + Run(kernel); + } + } + else { if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == TailNumber::Odd) { const auto kernel = - kernel_gemm_xdl_cshuffle_v3; + kernel_gemm_xdl_cshuffle_v3_multi_d; Run(kernel); } else { const auto kernel = - kernel_gemm_xdl_cshuffle_v3; + kernel_gemm_xdl_cshuffle_v3_multi_d; Run(kernel); } } @@ -416,12 +582,22 @@ struct DeviceGemmMultiD_Xdl_CShuffle_V3 : public DeviceGemmMultipleD 1) + { + const auto kernel = kernel_gemm_xdl_cshuffle_v3_multi_d< + GridwiseGemm, + false, + InMemoryDataOperationEnum::AtomicAdd, + minimum_occupancy>; + Run(kernel); + } + else { const auto kernel = - kernel_gemm_xdl_cshuffle_v3; + kernel_gemm_xdl_cshuffle_v3_multi_d; Run(kernel); } } @@ -451,6 +627,11 @@ struct DeviceGemmMultiD_Xdl_CShuffle_V3 : public DeviceGemmMultipleD && arg.KBatch > 1) + { + return false; + } + if((arg.K % AK1 != 0 || arg.K % BK1 != 0) && !(GemmSpec == GemmSpecialization::MKPadding || GemmSpec == GemmSpecialization::NKPadding || GemmSpec == GemmSpecialization::MNKPadding || @@ -479,6 +660,7 @@ struct DeviceGemmMultiD_Xdl_CShuffle_V3 : public DeviceGemmMultipleD StrideDs, index_t StrideC, + index_t KBatch, AElementwiseOperation a_element_op, BElementwiseOperation b_element_op, CElementwiseOperation c_element_op) @@ -494,7 +676,7 @@ struct DeviceGemmMultiD_Xdl_CShuffle_V3 : public DeviceGemmMultipleD StrideDs, index_t StrideC, + index_t KBatch, AElementwiseOperation a_element_op, BElementwiseOperation b_element_op, CElementwiseOperation c_element_op) override @@ -529,7 +712,7 @@ struct DeviceGemmMultiD_Xdl_CShuffle_V3 : public DeviceGemmMultipleD, bhalf_t>::value) - { - if(arg_.KBatch > 1) - hipGetErrorString( - hipMemsetAsync(arg_.p_c_grid, - 0, - arg_.M * arg_.N * sizeof(CDataType), - stream_config.stream_id_)); - } + if(arg_.KBatch > 1) + hipGetErrorString(hipMemsetAsync(arg_.p_c_grid, + 0, + arg_.M * arg_.N * sizeof(CDataType), + stream_config.stream_id_)); }; ave_time = ck::utility::launch_and_time_kernel_with_preprocess( @@ -190,14 +186,11 @@ struct DeviceGemm_Xdl_CShuffleV3 : public DeviceGemmV2, bhalf_t>::value) - { - if(arg.KBatch > 1) - hipGetErrorString(hipMemsetAsync(arg.p_c_grid, - 0, - arg.M * arg.N * sizeof(CDataType), - stream_config.stream_id_)); - } + if(arg.KBatch > 1) + hipGetErrorString(hipMemsetAsync(arg.p_c_grid, + 0, + arg.M * arg.N * sizeof(CDataType), + stream_config.stream_id_)); ave_time = launch_and_time_kernel( stream_config, kernel, dim3(gdx, gdy, gdz), dim3(BlockSize), 0, arg); @@ -215,15 +208,12 @@ struct DeviceGemm_Xdl_CShuffleV3 : public DeviceGemmV2 1) { - if constexpr(!is_same, bhalf_t>::value) - { - const auto kernel = - kernel_gemm_xdl_cshuffle_v3; - Run(kernel); - } + const auto kernel = + kernel_gemm_xdl_cshuffle_v3; + Run(kernel); } else { @@ -240,118 +230,113 @@ struct DeviceGemm_Xdl_CShuffleV3 : public DeviceGemmV2 1) { - if constexpr(!is_same, bhalf_t>::value) + if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == TailNumber::One) { - if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == TailNumber::One) + const auto kernel = + kernel_gemm_xdl_cshuffle_v3; + Run(kernel); + } + else if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == + TailNumber::Full) + { + const auto kernel = + kernel_gemm_xdl_cshuffle_v3; + Run(kernel); + } + + if constexpr(GridwiseGemm::BlockwiseGemmPipe::PrefetchStages > 2) + { + if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == TailNumber::Two) { const auto kernel = kernel_gemm_xdl_cshuffle_v3< GridwiseGemm, true, InMemoryDataOperationEnum::AtomicAdd, minimum_occupancy, - TailNumber::One>; + TailNumber::Two>; Run(kernel); } - else if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == - TailNumber::Full) + } + + if constexpr(GridwiseGemm::BlockwiseGemmPipe::PrefetchStages > 3) + { + if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == + TailNumber::Three) { const auto kernel = kernel_gemm_xdl_cshuffle_v3< GridwiseGemm, true, InMemoryDataOperationEnum::AtomicAdd, minimum_occupancy, - TailNumber::Full>; + TailNumber::Three>; Run(kernel); } + } - if constexpr(GridwiseGemm::BlockwiseGemmPipe::PrefetchStages > 2) + if constexpr(GridwiseGemm::BlockwiseGemmPipe::PrefetchStages > 4) + { + if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == + TailNumber::Four) { - if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == - TailNumber::Two) - { - const auto kernel = kernel_gemm_xdl_cshuffle_v3< - GridwiseGemm, - true, - InMemoryDataOperationEnum::AtomicAdd, - minimum_occupancy, - TailNumber::Two>; - Run(kernel); - } + const auto kernel = kernel_gemm_xdl_cshuffle_v3< + GridwiseGemm, + true, + InMemoryDataOperationEnum::AtomicAdd, + minimum_occupancy, + TailNumber::Four>; + Run(kernel); } + } - if constexpr(GridwiseGemm::BlockwiseGemmPipe::PrefetchStages > 3) + if constexpr(GridwiseGemm::BlockwiseGemmPipe::PrefetchStages > 5) + { + if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == + TailNumber::Five) { - if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == - TailNumber::Three) - { - const auto kernel = kernel_gemm_xdl_cshuffle_v3< - GridwiseGemm, - true, - InMemoryDataOperationEnum::AtomicAdd, - minimum_occupancy, - TailNumber::Three>; - Run(kernel); - } + const auto kernel = kernel_gemm_xdl_cshuffle_v3< + GridwiseGemm, + true, + InMemoryDataOperationEnum::AtomicAdd, + minimum_occupancy, + TailNumber::Five>; + Run(kernel); } + } - if constexpr(GridwiseGemm::BlockwiseGemmPipe::PrefetchStages > 4) + if constexpr(GridwiseGemm::BlockwiseGemmPipe::PrefetchStages > 6) + { + if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == TailNumber::Six) { - if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == - TailNumber::Four) - { - const auto kernel = kernel_gemm_xdl_cshuffle_v3< - GridwiseGemm, - true, - InMemoryDataOperationEnum::AtomicAdd, - minimum_occupancy, - TailNumber::Four>; - Run(kernel); - } + const auto kernel = kernel_gemm_xdl_cshuffle_v3< + GridwiseGemm, + true, + InMemoryDataOperationEnum::AtomicAdd, + minimum_occupancy, + TailNumber::Six>; + Run(kernel); } + } - if constexpr(GridwiseGemm::BlockwiseGemmPipe::PrefetchStages > 5) + if constexpr(GridwiseGemm::BlockwiseGemmPipe::PrefetchStages > 7) + { + if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == + TailNumber::Seven) { - if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == - TailNumber::Five) - { - const auto kernel = kernel_gemm_xdl_cshuffle_v3< - GridwiseGemm, - true, - InMemoryDataOperationEnum::AtomicAdd, - minimum_occupancy, - TailNumber::Five>; - Run(kernel); - } - } - - if constexpr(GridwiseGemm::BlockwiseGemmPipe::PrefetchStages > 6) - { - if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == - TailNumber::Six) - { - const auto kernel = kernel_gemm_xdl_cshuffle_v3< - GridwiseGemm, - true, - InMemoryDataOperationEnum::AtomicAdd, - minimum_occupancy, - TailNumber::Six>; - Run(kernel); - } - } - - if constexpr(GridwiseGemm::BlockwiseGemmPipe::PrefetchStages > 7) - { - if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == - TailNumber::Seven) - { - const auto kernel = kernel_gemm_xdl_cshuffle_v3< - GridwiseGemm, - true, - InMemoryDataOperationEnum::AtomicAdd, - minimum_occupancy, - TailNumber::Seven>; - Run(kernel); - } + const auto kernel = kernel_gemm_xdl_cshuffle_v3< + GridwiseGemm, + true, + InMemoryDataOperationEnum::AtomicAdd, + minimum_occupancy, + TailNumber::Seven>; + Run(kernel); } } } @@ -473,28 +458,25 @@ struct DeviceGemm_Xdl_CShuffleV3 : public DeviceGemmV2 1) { - if constexpr(!is_same, bhalf_t>::value) + if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == TailNumber::Odd) { - if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == TailNumber::Odd) - { - const auto kernel = kernel_gemm_xdl_cshuffle_v3_2lds< - GridwiseGemm, - true, - InMemoryDataOperationEnum::AtomicAdd, - minimum_occupancy, - TailNumber::Odd>; - Run(kernel); - } - else - { - const auto kernel = kernel_gemm_xdl_cshuffle_v3_2lds< - GridwiseGemm, - true, - InMemoryDataOperationEnum::AtomicAdd, - minimum_occupancy, - TailNumber::Even>; - Run(kernel); - } + const auto kernel = kernel_gemm_xdl_cshuffle_v3_2lds< + GridwiseGemm, + true, + InMemoryDataOperationEnum::AtomicAdd, + minimum_occupancy, + TailNumber::Odd>; + Run(kernel); + } + else + { + const auto kernel = kernel_gemm_xdl_cshuffle_v3_2lds< + GridwiseGemm, + true, + InMemoryDataOperationEnum::AtomicAdd, + minimum_occupancy, + TailNumber::Even>; + Run(kernel); } } else @@ -525,28 +507,25 @@ struct DeviceGemm_Xdl_CShuffleV3 : public DeviceGemmV2 1) { - if constexpr(!is_same, bhalf_t>::value) + if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == TailNumber::Odd) { - if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == TailNumber::Odd) - { - const auto kernel = kernel_gemm_xdl_cshuffle_v3< - GridwiseGemm, - true, - InMemoryDataOperationEnum::AtomicAdd, - minimum_occupancy, - TailNumber::Odd>; - Run(kernel); - } - else - { - const auto kernel = kernel_gemm_xdl_cshuffle_v3< - GridwiseGemm, - true, - InMemoryDataOperationEnum::AtomicAdd, - minimum_occupancy, - TailNumber::Even>; - Run(kernel); - } + const auto kernel = + kernel_gemm_xdl_cshuffle_v3; + Run(kernel); + } + else + { + const auto kernel = + kernel_gemm_xdl_cshuffle_v3; + Run(kernel); } } else @@ -579,18 +558,14 @@ struct DeviceGemm_Xdl_CShuffleV3 : public DeviceGemmV2 1) { - if constexpr(!is_same, bhalf_t>::value) - { - const auto kernel = - kernel_gemm_xdl_cshuffle_v3; - Run(kernel); - } + const auto kernel = + kernel_gemm_xdl_cshuffle_v3; + Run(kernel); } else { @@ -628,6 +603,11 @@ struct DeviceGemm_Xdl_CShuffleV3 : public DeviceGemmV2 && arg.KBatch > 1) + { + return false; + } + if((arg.K % AK1 != 0 || arg.K % BK1 != 0) && !(GemmSpec == GemmSpecialization::MKPadding || GemmSpec == GemmSpecialization::NKPadding || GemmSpec == GemmSpecialization::MNKPadding || diff --git a/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_v3.hpp b/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_v3.hpp index 03750dbc36..2fea99b9a5 100644 --- a/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_v3.hpp +++ b/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_v3.hpp @@ -417,6 +417,13 @@ struct GridwiseGemm_xdl_cshuffle_v3 } }(); + // pad M and N + return transform_tensor_descriptor(c_grid_desc_mraw_nraw, + make_tuple(make_right_pad_transform(M, MPad - M), + make_right_pad_transform(N, NPad - N)), + make_tuple(Sequence<0>{}, Sequence<1>{}), + make_tuple(Sequence<0>{}, Sequence<1>{})); +#if 0 using GemmSpecialization = tensor_operation::device::GemmSpecialization; if constexpr(GemmSpec == GemmSpecialization::MNPadding || @@ -454,6 +461,7 @@ struct GridwiseGemm_xdl_cshuffle_v3 // not pad M or N return c_grid_desc_mraw_nraw; } +#endif } struct Problem @@ -953,7 +961,8 @@ struct GridwiseGemm_xdl_cshuffle_v3 if constexpr(!(GemmSpec == tensor_operation::device::GemmSpecialization::MPadding || GemmSpec == tensor_operation::device::GemmSpecialization::MNPadding || GemmSpec == tensor_operation::device::GemmSpecialization::MKPadding || - GemmSpec == tensor_operation::device::GemmSpecialization::MNKPadding)) + GemmSpec == tensor_operation::device::GemmSpecialization::MNKPadding) && + !(is_same::value)) { if(!(karg.M % MPerBlock == 0)) { @@ -970,7 +979,8 @@ struct GridwiseGemm_xdl_cshuffle_v3 if constexpr(!(GemmSpec == tensor_operation::device::GemmSpecialization::NPadding || GemmSpec == tensor_operation::device::GemmSpecialization::MNPadding || GemmSpec == tensor_operation::device::GemmSpecialization::NKPadding || - GemmSpec == tensor_operation::device::GemmSpecialization::MNKPadding)) + GemmSpec == tensor_operation::device::GemmSpecialization::MNKPadding) && + (is_same::value)) { if(!(karg.N % NPerBlock == 0)) { @@ -1105,7 +1115,9 @@ struct GridwiseGemm_xdl_cshuffle_v3 } if constexpr(!(is_same, half_t>::value || - is_same, float>::value)) + is_same, float>::value || + is_same, bhalf_t>::value || + is_same, int32_t>::value)) { if(!karg.IsReduceAdd()) { 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 3a1ac6c6de..64eaf4da04 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 @@ -36,10 +36,9 @@ __global__ void __launch_bounds__(CK_MAX_THREAD_PER_BLOCK, MinimumOccupancy) #endif // __attribute__((amdgpu_waves_per_eu(1, 1))) - kernel_gemm_xdl_cshuffle_v3(typename GridwiseGemm::Argument karg) + kernel_gemm_xdl_cshuffle_v3_multi_d(typename GridwiseGemm::Argument karg) { -#if(!defined(__HIP_DEVICE_COMPILE__) || defined(__gfx908__) || defined(__gfx90a__) || \ - defined(__gfx940__) || defined(__gfx941__) || defined(__gfx942__)) +#if(!defined(__HIP_DEVICE_COMPILE__) || defined(__gfx9__)) __shared__ char p_shared[GridwiseGemm::GetSharedMemoryNumberOfByte()]; auto splitk_batch_offset = typename GridwiseGemm::SplitKBatchOffset(karg); @@ -56,7 +55,7 @@ __global__ void karg.c_element_op); #else ignore = karg; -#endif // end of if (defined(__gfx908__) || defined(__gfx90a__)) +#endif // end of if (defined(__gfx9__)) } template {}, Sequence<1>{}), + make_tuple(Sequence<0>{}, Sequence<1>{})); +#if 0 using GemmSpecialization = tensor_operation::device::GemmSpecialization; if constexpr(GemmSpec == GemmSpecialization::MNPadding || @@ -491,6 +496,7 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3 // not pad M or N return c_grid_desc_mraw_nraw; } +#endif } __host__ __device__ static auto MakeDsGridDescriptor_M_N( @@ -1016,7 +1022,8 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3 if constexpr(!(GemmSpec == tensor_operation::device::GemmSpecialization::MPadding || GemmSpec == tensor_operation::device::GemmSpecialization::MNPadding || GemmSpec == tensor_operation::device::GemmSpecialization::MKPadding || - GemmSpec == tensor_operation::device::GemmSpecialization::MNKPadding)) + GemmSpec == tensor_operation::device::GemmSpecialization::MNKPadding) && + !(is_same::value)) { if(!(karg.M % MPerBlock == 0)) { @@ -1033,7 +1040,8 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3 if constexpr(!(GemmSpec == tensor_operation::device::GemmSpecialization::NPadding || GemmSpec == tensor_operation::device::GemmSpecialization::MNPadding || GemmSpec == tensor_operation::device::GemmSpecialization::NKPadding || - GemmSpec == tensor_operation::device::GemmSpecialization::MNKPadding)) + GemmSpec == tensor_operation::device::GemmSpecialization::MNKPadding) && + (is_same::value)) { if(!(karg.N % NPerBlock == 0)) { diff --git a/include/ck/utility/amd_buffer_addressing.hpp b/include/ck/utility/amd_buffer_addressing.hpp index ab22134fc6..d4ee5c886c 100644 --- a/include/ck/utility/amd_buffer_addressing.hpp +++ b/include/ck/utility/amd_buffer_addressing.hpp @@ -562,6 +562,34 @@ __device__ void amd_buffer_store_impl(const typename vector_type::type src dst_wave_addr_offset); } +template +__device__ void amd_global_atomic_add_impl(const typename vector_type::type src_thread_data, + T* addr) +{ + static_assert((is_same::value && (N == 2 || N == 4 || N == 8)) || + (is_same::value && (N == 2 || N == 4 || N == 8)), + "wrong! not implemented"); + + if constexpr(is_same::value) + { + vector_type tmp{src_thread_data}; + static_for<0, N / 2, 1>{}([&](auto i) { + __builtin_amdgcn_global_atomic_fadd_v2f16(bit_cast(addr) + i, + tmp.template AsType()[i]); + }); + } +#if defined(__gfx942__) + else if constexpr(is_same::value) + { + vector_type tmp{src_thread_data}; + static_for<0, N / 2, 1>{}([&](auto i) { + __builtin_amdgcn_global_atomic_fadd_v2bf16(bit_cast(addr) + i, + tmp.template AsType()[i]); + }); + } +#endif +} + template __device__ void amd_buffer_atomic_add_impl(const typename vector_type::type src_thread_data, int32x4_t dst_wave_buffer_resource, @@ -907,18 +935,29 @@ amd_buffer_atomic_add(const typename vector_type_maker::type::type src_thr using scalar_t = typename scalar_type::type; constexpr index_t vector_size = scalar_type::vector_size; -#if CK_EXPERIMENTAL_USE_BUFFER_ATOMIC_ADD_OOB_CHECK_OFFSET_TRICK - uint32_t dst_addr_shift = dst_thread_element_valid ? 0 : 0x80000000; - - amd_buffer_atomic_add_impl( - src_thread_data, dst_wave_buffer_resource, dst_addr_shift + dst_thread_addr_offset, 0); -#else - if(dst_thread_element_valid) + if constexpr(is_same::value) { - amd_buffer_atomic_add_impl( - src_thread_data, dst_wave_buffer_resource, dst_thread_addr_offset, 0); + if(dst_thread_element_valid) + { + amd_global_atomic_add_impl( + src_thread_data, p_dst_wave + dst_thread_element_offset); + } } + else + { +#if CK_EXPERIMENTAL_USE_BUFFER_ATOMIC_ADD_OOB_CHECK_OFFSET_TRICK + uint32_t dst_addr_shift = dst_thread_element_valid ? 0 : 0x80000000; + + amd_buffer_atomic_add_impl( + src_thread_data, dst_wave_buffer_resource, dst_addr_shift + dst_thread_addr_offset, 0); +#else + if(dst_thread_element_valid) + { + amd_buffer_atomic_add_impl( + src_thread_data, dst_wave_buffer_resource, dst_thread_addr_offset, 0); + } #endif + } } // buffer_atomic_max requires: diff --git a/include/ck/utility/dynamic_buffer.hpp b/include/ck/utility/dynamic_buffer.hpp index 76390e614e..0dcc514a2f 100644 --- a/include/ck/utility/dynamic_buffer.hpp +++ b/include/ck/utility/dynamic_buffer.hpp @@ -358,13 +358,15 @@ struct DynamicBuffer bool constexpr use_amd_buffer_addressing = is_same_v, int32_t> || is_same_v, float> || - (is_same_v, half_t> && scalar_per_x_vector % 2 == 0); + (is_same_v, half_t> && scalar_per_x_vector % 2 == 0) || + (is_same_v, bhalf_t> && scalar_per_x_vector % 2 == 0); #elif CK_USE_AMD_BUFFER_ATOMIC_ADD_INTEGER && (!CK_USE_AMD_BUFFER_ATOMIC_ADD_FLOAT) bool constexpr use_amd_buffer_addressing = is_same_v, int32_t>; #elif(!CK_USE_AMD_BUFFER_ATOMIC_ADD_INTEGER) && CK_USE_AMD_BUFFER_ATOMIC_ADD_FLOAT bool constexpr use_amd_buffer_addressing = is_same_v, float> || - (is_same_v, half_t> && scalar_per_x_vector % 2 == 0); + (is_same_v, half_t> && scalar_per_x_vector % 2 == 0) || + (is_same_v, bhalf_t> && scalar_per_x_vector % 2 == 0); #else bool constexpr use_amd_buffer_addressing = false; #endif diff --git a/library/include/ck/library/tensor_operation_instance/gpu/gemm_multiply_multiply.hpp b/library/include/ck/library/tensor_operation_instance/gpu/gemm_multiply_multiply.hpp index f8e8e8fdec..2077f904d3 100644 --- a/library/include/ck/library/tensor_operation_instance/gpu/gemm_multiply_multiply.hpp +++ b/library/include/ck/library/tensor_operation_instance/gpu/gemm_multiply_multiply.hpp @@ -18,134 +18,82 @@ namespace device { namespace instance { #if(defined(CK_ENABLE_BF16) || defined(CK_ENABLE_FP8)) void add_device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_comp_default_instances( - std::vector, - Row, - F8, - F8, - Tuple, - BF16, - PassThrough, - PassThrough, - MultiplyMultiply>>>& instances); + std::vector, + Row, + F8, + F8, + Tuple, + BF16, + PassThrough, + PassThrough, + MultiplyMultiply>>>& instances); void add_device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_comp_kpadding_instances( - std::vector, - Row, - F8, - F8, - Tuple, - BF16, - PassThrough, - PassThrough, - MultiplyMultiply>>>& instances); - -void add_device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_comp_mnpadding_instances( - std::vector, - Row, - F8, - F8, - Tuple, - BF16, - PassThrough, - PassThrough, - MultiplyMultiply>>>& instances); - -void add_device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_comp_mnkpadding_instances( - std::vector, - Row, - F8, - F8, - Tuple, - BF16, - PassThrough, - PassThrough, - MultiplyMultiply>>>& instances); + std::vector, + Row, + F8, + F8, + Tuple, + BF16, + PassThrough, + PassThrough, + MultiplyMultiply>>>& instances); void add_device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_mem_v1_default_instances( - std::vector, - Row, - F8, - F8, - Tuple, - BF16, - PassThrough, - PassThrough, - MultiplyMultiply>>>& instances); + std::vector, + Row, + F8, + F8, + Tuple, + BF16, + PassThrough, + PassThrough, + MultiplyMultiply>>>& instances); void add_device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_mem_v1_kpadding_instances( - std::vector, - Row, - F8, - F8, - Tuple, - BF16, - PassThrough, - PassThrough, - MultiplyMultiply>>>& instances); - -void add_device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_mem_v1_mnkpadding_instances( - std::vector, - Row, - F8, - F8, - Tuple, - BF16, - PassThrough, - PassThrough, - MultiplyMultiply>>>& instances); + std::vector, + Row, + F8, + F8, + Tuple, + BF16, + PassThrough, + PassThrough, + MultiplyMultiply>>>& instances); void add_device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_mem_v2_default_instances( - std::vector, - Row, - F8, - F8, - Tuple, - BF16, - PassThrough, - PassThrough, - MultiplyMultiply>>>& instances); + std::vector, + Row, + F8, + F8, + Tuple, + BF16, + PassThrough, + PassThrough, + MultiplyMultiply>>>& instances); void add_device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_mem_v2_kpadding_instances( - std::vector, - Row, - F8, - F8, - Tuple, - BF16, - PassThrough, - PassThrough, - MultiplyMultiply>>>& instances); - -void add_device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_mem_v2_mnkpadding_instances( - std::vector, - Row, - F8, - F8, - Tuple, - BF16, - PassThrough, - PassThrough, - MultiplyMultiply>>>& instances); + std::vector, + Row, + F8, + F8, + Tuple, + BF16, + PassThrough, + PassThrough, + MultiplyMultiply>>>& instances); #endif template -struct DeviceOperationInstanceFactory, @@ -167,17 +115,18 @@ struct DeviceOperationInstanceFactory> { - using DeviceOp = DeviceGemmMultipleD, - CLayout, - ADataType, - BDataType, - Tuple, - CDataType, - ck::tensor_operation::element_wise::PassThrough, - ck::tensor_operation::element_wise::PassThrough, - ck::tensor_operation::element_wise::MultiplyMultiply>; + using DeviceOp = + DeviceGemmMultipleDSplitK, + CLayout, + ADataType, + BDataType, + Tuple, + CDataType, + ck::tensor_operation::element_wise::PassThrough, + ck::tensor_operation::element_wise::PassThrough, + ck::tensor_operation::element_wise::MultiplyMultiply>; static auto GetInstances() { @@ -194,24 +143,16 @@ struct DeviceOperationInstanceFactory>>& instances); -void add_device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_comp_mnpadding_instances( - std::vector>>& - instances); - -void add_device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_comp_mnkpadding_instances( - std::vector>>& - instances); - void add_device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_mem_v1_default_instances( std::vector>>& @@ -97,11 +87,6 @@ void add_device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_mem_v1_kpadding_instance DeviceGemmV2>>& instances); -void add_device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_mem_v1_mnkpadding_instances( - std::vector>>& - instances); - void add_device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_mem_v2_default_instances( std::vector>>& @@ -111,13 +96,8 @@ void add_device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_mem_v2_kpadding_instance std::vector>>& instances); - -void add_device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_mem_v2_mnkpadding_instances( - std::vector>>& - instances); #endif -#if(defined(CK_ENABLE_FP16) || defined(CK_ENABLE_FP8)) +#if(defined(CK_ENABLE_FP16) && defined(CK_ENABLE_FP8)) void add_device_gemm_xdl_universal_f16_f8_f16_mk_kn_mn_comp_default_instances( std::vector>>& @@ -177,16 +157,6 @@ void add_device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_comp_kpadding_instances( DeviceGemmV2>>& instances); -void add_device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_comp_mnpadding_instances( - std::vector>>& - instances); - -void add_device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_comp_mnkpadding_instances( - std::vector>>& - instances); - void add_device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_mem_v1_default_instances( std::vector>>& @@ -196,12 +166,6 @@ void add_device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_mem_v1_kpadding_instances std::vector>>& instances); - -void add_device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_mem_v1_mnkpadding_instances( - std::vector>>& - instances); - void add_device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_mem_v2_default_instances( std::vector>>& @@ -212,10 +176,6 @@ void add_device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_mem_v2_kpadding_instances DeviceGemmV2>>& instances); -void add_device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_mem_v2_mnkpadding_instances( - std::vector>>& - instances); void add_device_gemm_xdl_universal_f8_f16_f16_mk_kn_mn_comp_default_instances( std::vector>>& @@ -275,16 +235,6 @@ void add_device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_comp_kpadding_instances( DeviceGemmV2>>& instances); -void add_device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_comp_mnpadding_instances( - std::vector>>& - instances); - -void add_device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_comp_mnkpadding_instances( - std::vector>>& - instances); - void add_device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_mem_v1_default_instances( std::vector>>& @@ -295,11 +245,6 @@ void add_device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_mem_v1_kpadding_instances DeviceGemmV2>>& instances); -void add_device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_mem_v1_mnkpadding_instances( - std::vector>>& - instances); - void add_device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_mem_v2_default_instances( std::vector>>& @@ -309,11 +254,6 @@ void add_device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_mem_v2_kpadding_instances std::vector>>& instances); - -void add_device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_mem_v2_mnkpadding_instances( - std::vector>>& - instances); #endif #ifdef CK_ENABLE_BF16 void add_device_gemm_xdl_universal_bf16_bf16_bf16_mk_kn_mn_comp_default_instances( @@ -376,16 +316,6 @@ void add_device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_comp_kpadding_instanc DeviceGemmV2>>& instances); -void add_device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_comp_mnpadding_instances( - std::vector>>& - instances); - -void add_device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_comp_mnkpadding_instances( - std::vector>>& - instances); - void add_device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_mem_v1_default_instances( std::vector>>& @@ -396,11 +326,6 @@ void add_device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_mem_v1_kpadding_insta DeviceGemmV2>>& instances); -void add_device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_mem_v1_mnkpadding_instances( - std::vector>>& - instances); - void add_device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_mem_v2_default_instances( std::vector>>& @@ -410,13 +335,53 @@ void add_device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_mem_v2_kpadding_insta std::vector>>& instances); - -void add_device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_mem_v2_mnkpadding_instances( - std::vector>>& - instances); #endif -#if(defined(CK_ENABLE_BF16) || defined(CK_ENABLE_FP8)) +#if(defined(CK_ENABLE_BF16) && defined(CK_ENABLE_FP8)) +void add_device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_comp_default_instances( + std::vector>>& + instances); + +void add_device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_comp_kpadding_instances( + std::vector>>& + instances); + +void add_device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_comp_nkpadding_instances( + std::vector>>& + instances); + +void add_device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_v1_default_instances( + std::vector>>& + instances); + +void add_device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_v1_kpadding_instances( + std::vector>>& + instances); + +void add_device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_v1_nkpadding_instances( + std::vector>>& + instances); + +void add_device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_v2_default_instances( + std::vector>>& + instances); + +void add_device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_v2_kpadding_instances( + std::vector>>& + instances); + +void add_device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_v2_nkpadding_instances( + std::vector>>& + instances); + void add_device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_comp_default_instances( std::vector>>& @@ -427,16 +392,6 @@ void add_device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_comp_kpadding_instances( DeviceGemmV2>>& instances); -void add_device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_comp_mnpadding_instances( - std::vector>>& - instances); - -void add_device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_comp_mnkpadding_instances( - std::vector>>& - instances); - void add_device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_mem_v1_default_instances( std::vector>>& @@ -447,11 +402,6 @@ void add_device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_mem_v1_kpadding_instances DeviceGemmV2>>& instances); -void add_device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_mem_v1_mnkpadding_instances( - std::vector>>& - instances); - void add_device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_mem_v2_default_instances( std::vector>>& @@ -461,11 +411,6 @@ void add_device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_mem_v2_kpadding_instances std::vector>>& instances); - -void add_device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_mem_v2_mnkpadding_instances( - std::vector>>& - instances); #endif template && is_same_v && is_same_v) { @@ -562,21 +499,14 @@ struct DeviceOperationInstanceFactory< { add_device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_comp_default_instances(op_ptrs); add_device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_comp_kpadding_instances(op_ptrs); - add_device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_comp_mnpadding_instances(op_ptrs); - add_device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_comp_mnkpadding_instances( - op_ptrs); add_device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_mem_v1_default_instances(op_ptrs); add_device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_mem_v1_kpadding_instances( op_ptrs); - add_device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_mem_v1_mnkpadding_instances( - op_ptrs); add_device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_mem_v2_default_instances(op_ptrs); add_device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_mem_v2_kpadding_instances( op_ptrs); - add_device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_mem_v2_mnkpadding_instances( - op_ptrs); } else if constexpr(is_same_v && is_same_v && is_same_v) @@ -608,21 +538,14 @@ struct DeviceOperationInstanceFactory< { add_device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_comp_default_instances(op_ptrs); add_device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_comp_kpadding_instances(op_ptrs); - add_device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_comp_mnpadding_instances(op_ptrs); - add_device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_comp_mnkpadding_instances( - op_ptrs); add_device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_mem_v1_default_instances(op_ptrs); add_device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_mem_v1_kpadding_instances( op_ptrs); - add_device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_mem_v1_mnkpadding_instances( - op_ptrs); add_device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_mem_v2_default_instances(op_ptrs); add_device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_mem_v2_kpadding_instances( op_ptrs); - add_device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_mem_v2_mnkpadding_instances( - op_ptrs); } else if constexpr(is_same_v && is_same_v && is_same_v) @@ -684,51 +607,55 @@ struct DeviceOperationInstanceFactory< op_ptrs); add_device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_comp_kpadding_instances( op_ptrs); - add_device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_comp_mnpadding_instances( - op_ptrs); - add_device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_comp_mnkpadding_instances( - op_ptrs); add_device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_mem_v1_default_instances( op_ptrs); add_device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_mem_v1_kpadding_instances( op_ptrs); - add_device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_mem_v1_mnkpadding_instances( - op_ptrs); add_device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_mem_v2_default_instances( op_ptrs); add_device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_mem_v2_kpadding_instances( op_ptrs); - add_device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_mem_v2_mnkpadding_instances( - op_ptrs); } } #endif -#if(defined(CK_ENABLE_BF16) || defined(CK_ENABLE_FP8)) +#if(defined(CK_ENABLE_BF16) && defined(CK_ENABLE_FP8)) if constexpr(is_same_v && is_same_v && is_same_v) { - if constexpr(is_same_v && is_same_v && + if constexpr(is_same_v && is_same_v && is_same_v) + { + add_device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_comp_default_instances(op_ptrs); + add_device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_comp_kpadding_instances(op_ptrs); + add_device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_comp_nkpadding_instances(op_ptrs); + + add_device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_v1_default_instances(op_ptrs); + add_device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_v1_kpadding_instances( + op_ptrs); + add_device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_v1_nkpadding_instances( + op_ptrs); + + add_device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_v2_default_instances(op_ptrs); + add_device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_v2_kpadding_instances( + op_ptrs); + add_device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_v2_nkpadding_instances( + op_ptrs); + } + else if constexpr(is_same_v && is_same_v && + is_same_v) { add_device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_comp_default_instances(op_ptrs); add_device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_comp_kpadding_instances(op_ptrs); - add_device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_comp_mnpadding_instances(op_ptrs); - add_device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_comp_mnkpadding_instances( - op_ptrs); add_device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_mem_v1_default_instances(op_ptrs); add_device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_mem_v1_kpadding_instances( op_ptrs); - add_device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_mem_v1_mnkpadding_instances( - op_ptrs); add_device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_mem_v2_default_instances(op_ptrs); add_device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_mem_v2_kpadding_instances( op_ptrs); - add_device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_mem_v2_mnkpadding_instances( - op_ptrs); } } #endif diff --git a/library/src/tensor_operation_instance/gpu/gemm/CMakeLists.txt b/library/src/tensor_operation_instance/gpu/gemm/CMakeLists.txt index e9cc1e854f..0cd54c7788 100644 --- a/library/src/tensor_operation_instance/gpu/gemm/CMakeLists.txt +++ b/library/src/tensor_operation_instance/gpu/gemm/CMakeLists.txt @@ -100,16 +100,18 @@ list(APPEND GEMM_INSTANCES device_gemm_xdl_c_shuffle_bf16_bf16_bf16_km_kn_mn_instance.cpp device_gemm_xdl_c_shuffle_bf16_bf16_bf16_km_nk_mn_instance.cpp) -list(APPEND GEMM_INSTANCES - device_gemm_xdl_c_shuffle_fp8_fp8_fp8_mk_kn_mn_v1_default_instance.cpp - device_gemm_xdl_c_shuffle_fp8_fp8_fp8_mk_kn_mn_v1_interwave_default_instance.cpp - device_gemm_xdl_c_shuffle_fp8_fp8_fp8_mk_kn_mn_v2_default_instance.cpp - device_gemm_xdl_c_shuffle_fp8_fp8_fp8_mk_kn_mn_v1_padded_instance.cpp - device_gemm_xdl_c_shuffle_fp8_fp8_fp8_mk_kn_mn_v1_interwave_padded_instance.cpp - device_gemm_xdl_c_shuffle_fp8_fp8_fp8_mk_kn_mn_v2_padded_instance.cpp - device_gemm_xdl_c_shuffle_fp8_fp8_fp8_mk_nk_mn_instance.cpp - device_gemm_xdl_c_shuffle_fp8_fp8_fp8_km_kn_mn_instance.cpp - device_gemm_xdl_c_shuffle_fp8_fp8_fp8_km_nk_mn_instance.cpp) +if(GPU_TARGETS MATCHES "gfx94") + list(APPEND GEMM_INSTANCES + device_gemm_xdl_c_shuffle_fp8_fp8_fp8_mk_kn_mn_v1_default_instance.cpp + device_gemm_xdl_c_shuffle_fp8_fp8_fp8_mk_kn_mn_v1_interwave_default_instance.cpp + device_gemm_xdl_c_shuffle_fp8_fp8_fp8_mk_kn_mn_v2_default_instance.cpp + device_gemm_xdl_c_shuffle_fp8_fp8_fp8_mk_kn_mn_v1_padded_instance.cpp + device_gemm_xdl_c_shuffle_fp8_fp8_fp8_mk_kn_mn_v1_interwave_padded_instance.cpp + device_gemm_xdl_c_shuffle_fp8_fp8_fp8_mk_kn_mn_v2_padded_instance.cpp + device_gemm_xdl_c_shuffle_fp8_fp8_fp8_mk_nk_mn_instance.cpp + device_gemm_xdl_c_shuffle_fp8_fp8_fp8_km_kn_mn_instance.cpp + device_gemm_xdl_c_shuffle_fp8_fp8_fp8_km_nk_mn_instance.cpp) +endif() list(APPEND GEMM_INSTANCES device_gemm_wmma_f16_f16_f16_mk_kn_mn_instance.cpp diff --git a/library/src/tensor_operation_instance/gpu/gemm_ab_scale/CMakeLists.txt b/library/src/tensor_operation_instance/gpu/gemm_ab_scale/CMakeLists.txt index 5621cf0eec..aab1c4e86e 100644 --- a/library/src/tensor_operation_instance/gpu/gemm_ab_scale/CMakeLists.txt +++ b/library/src/tensor_operation_instance/gpu/gemm_ab_scale/CMakeLists.txt @@ -11,4 +11,9 @@ list(APPEND GEMM_AB_SCALE_INSTANCES device_gemm_ab_scale_xdl_f8_f8_bf16/device_gemm_ab_scale_xdl_f8_f8_bf16_mk_nk_mn_128_128_128_mem_v1_mnkpadding_instance.cpp ) +set_source_files_properties(device_gemm_ab_scale_xdl_f8_f8_bf16/device_gemm_ab_scale_xdl_f8_f8_bf16_mk_nk_mn_128_128_128_comp_default_instance.cpp PROPERTIES COMPILE_OPTIONS ";-mllvm;-greedy-reverse-local-assignment=1") +set_source_files_properties(device_gemm_ab_scale_xdl_f8_f8_bf16/device_gemm_ab_scale_xdl_f8_f8_bf16_mk_nk_mn_128_128_128_comp_kpadding_instance.cpp PROPERTIES COMPILE_OPTIONS ";-mllvm;-greedy-reverse-local-assignment=1") +set_source_files_properties(device_gemm_ab_scale_xdl_f8_f8_bf16/device_gemm_ab_scale_xdl_f8_f8_bf16_mk_nk_mn_128_128_128_comp_mnpadding_instance.cpp PROPERTIES COMPILE_OPTIONS ";-mllvm;-greedy-reverse-local-assignment=1") +set_source_files_properties(device_gemm_ab_scale_xdl_f8_f8_bf16/device_gemm_ab_scale_xdl_f8_f8_bf16_mk_nk_mn_128_128_128_comp_mnkpadding_instance.cpp PROPERTIES COMPILE_OPTIONS ";-mllvm;-greedy-reverse-local-assignment=1") + add_instance_library(device_gemm_ab_scale_instance ${GEMM_AB_SCALE_INSTANCES}) diff --git a/library/src/tensor_operation_instance/gpu/gemm_multiply_multiply/CMakeLists.txt b/library/src/tensor_operation_instance/gpu/gemm_multiply_multiply/CMakeLists.txt index df092aaaf6..5e56aebcfd 100644 --- a/library/src/tensor_operation_instance/gpu/gemm_multiply_multiply/CMakeLists.txt +++ b/library/src/tensor_operation_instance/gpu/gemm_multiply_multiply/CMakeLists.txt @@ -4,14 +4,13 @@ set(GEMM_MULTIPLY_MULTIPLY_INSTANCES) list(APPEND GEMM_MULTIPLY_MULTIPLY_INSTANCES device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_comp_default_instance.cpp device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_comp_kpadding_instance.cpp - device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_comp_mnpadding_instance.cpp - device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_comp_mnkpadding_instance.cpp device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_mem_v1_default_instance.cpp device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_mem_v1_kpadding_instance.cpp - device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_mem_v1_mnkpadding_instance.cpp device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_mem_v2_default_instance.cpp device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_mem_v2_kpadding_instance.cpp - device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_mem_v2_mnkpadding_instance.cpp ) +set_source_files_properties(device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_comp_default_instance.cpp PROPERTIES COMPILE_OPTIONS ";-mllvm;-greedy-reverse-local-assignment=1") +set_source_files_properties(device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_comp_kpadding_instance.cpp PROPERTIES COMPILE_OPTIONS ";-mllvm;-greedy-reverse-local-assignment=1") + add_instance_library(device_gemm_multiply_multiply_instance ${GEMM_MULTIPLY_MULTIPLY_INSTANCES}) diff --git a/library/src/tensor_operation_instance/gpu/gemm_multiply_multiply/device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn.hpp b/library/src/tensor_operation_instance/gpu/gemm_multiply_multiply/device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn.hpp index 61d55cfa49..8a24af1b8b 100644 --- a/library/src/tensor_operation_instance/gpu/gemm_multiply_multiply/device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn.hpp +++ b/library/src/tensor_operation_instance/gpu/gemm_multiply_multiply/device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn.hpp @@ -46,8 +46,7 @@ using device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_comp_instances = std DeviceGemmMultiD_Xdl_CShuffle_V3< Row, Col, Tuple, Row, F8, F8, Tuple, BF16, F32, F32, PassThrough, PassThrough, MultiplyMultiply, 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, 8, 1>, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v4, F8>, DeviceGemmMultiD_Xdl_CShuffle_V3< Row, Col, Tuple, Row, F8, F8, Tuple, BF16, F32, F32, PassThrough, PassThrough, MultiplyMultiply, GemmSpec, 256, 128, 128, 128, 16, 16, 32, 32, 2, 2, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, 1, 1, S<1, 32, 1, 8>, S<8, 8, 1>, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v4, F8>, DeviceGemmMultiD_Xdl_CShuffle_V3< Row, Col, Tuple, Row, F8, F8, Tuple, BF16, F32, F32, PassThrough, PassThrough, MultiplyMultiply, GemmSpec, 256, 128, 128, 64, 16, 16, 32, 32, 2, 2, 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, 8, 1>, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v4, F8>, - DeviceGemmMultiD_Xdl_CShuffle_V3< Row, Col, Tuple, Row, F8, F8, Tuple, BF16, F32, F32, PassThrough, PassThrough, MultiplyMultiply, 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, 8, 1>, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v3, F8>, - DeviceGemmMultiD_Xdl_CShuffle_V3< Row, Col, Tuple, Row, F8, F8, Tuple, BF16, F32, F32, PassThrough, PassThrough, MultiplyMultiply, 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, 8, 1>, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v5, F8>, + DeviceGemmMultiD_Xdl_CShuffle_V3< Row, Col, Tuple, Row, F8, F8, Tuple, BF16, F32, F32, PassThrough, PassThrough, MultiplyMultiply, GemmSpec, 256, 256, 256, 128, 16, 16, 16, 16, 8, 8, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, 1, 2, S<1, 32, 1, 8>, S<8, 8, 1>, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v3, F8>, DeviceGemmMultiD_Xdl_CShuffle_V3< Row, Col, Tuple, Row, F8, F8, Tuple, BF16, F32, F32, PassThrough, PassThrough, MultiplyMultiply, GemmSpec, 256, 256, 256, 64, 16, 16, 16, 16, 8, 8, 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, 2, S<1, 32, 1, 8>, S<8, 8, 1>, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v3, F8>, DeviceGemmMultiD_Xdl_CShuffle_V3< Row, Col, Tuple, Row, F8, F8, Tuple, BF16, F32, F32, PassThrough, PassThrough, MultiplyMultiply, GemmSpec, 256, 224, 256, 128, 16, 16, 16, 16, 7, 8, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, 1, 2, S<1, 32, 1, 8>, S<8, 8, 1>, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v3, F8>, DeviceGemmMultiD_Xdl_CShuffle_V3< Row, Col, Tuple, Row, F8, F8, Tuple, BF16, F32, F32, PassThrough, PassThrough, MultiplyMultiply, GemmSpec, 256, 256, 224, 128, 16, 16, 16, 16, 8, 7, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, 2, 1, S<1, 64, 1, 4>, S<8, 8, 1>, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v3, F8>, diff --git a/library/src/tensor_operation_instance/gpu/gemm_multiply_multiply/device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_comp_default_instance.cpp b/library/src/tensor_operation_instance/gpu/gemm_multiply_multiply/device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_comp_default_instance.cpp index 81131b4de2..6527d93473 100644 --- a/library/src/tensor_operation_instance/gpu/gemm_multiply_multiply/device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_comp_default_instance.cpp +++ b/library/src/tensor_operation_instance/gpu/gemm_multiply_multiply/device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_comp_default_instance.cpp @@ -9,17 +9,17 @@ namespace device { namespace instance { void add_device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_comp_default_instances( - std::vector, - Row, - F8, - F8, - Tuple, - BF16, - PassThrough, - PassThrough, - MultiplyMultiply>>>& instances) + std::vector, + Row, + F8, + F8, + Tuple, + BF16, + PassThrough, + PassThrough, + MultiplyMultiply>>>& instances) { add_device_operation_instances( instances, diff --git a/library/src/tensor_operation_instance/gpu/gemm_multiply_multiply/device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_comp_kpadding_instance.cpp b/library/src/tensor_operation_instance/gpu/gemm_multiply_multiply/device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_comp_kpadding_instance.cpp index 149e4ad144..7f16a7a2c5 100644 --- a/library/src/tensor_operation_instance/gpu/gemm_multiply_multiply/device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_comp_kpadding_instance.cpp +++ b/library/src/tensor_operation_instance/gpu/gemm_multiply_multiply/device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_comp_kpadding_instance.cpp @@ -9,17 +9,17 @@ namespace device { namespace instance { void add_device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_comp_kpadding_instances( - std::vector, - Row, - F8, - F8, - Tuple, - BF16, - PassThrough, - PassThrough, - MultiplyMultiply>>>& instances) + std::vector, + Row, + F8, + F8, + Tuple, + BF16, + PassThrough, + PassThrough, + MultiplyMultiply>>>& instances) { add_device_operation_instances( instances, diff --git a/library/src/tensor_operation_instance/gpu/gemm_multiply_multiply/device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_comp_mnkpadding_instance.cpp b/library/src/tensor_operation_instance/gpu/gemm_multiply_multiply/device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_comp_mnkpadding_instance.cpp deleted file mode 100644 index ba71f924e0..0000000000 --- a/library/src/tensor_operation_instance/gpu/gemm_multiply_multiply/device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_comp_mnkpadding_instance.cpp +++ /dev/null @@ -1,32 +0,0 @@ -// SPDX-License-Identifier: MIT -// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved. - -#include "device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn.hpp" - -namespace ck { -namespace tensor_operation { -namespace device { -namespace instance { - -void add_device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_comp_mnkpadding_instances( - std::vector, - Row, - F8, - F8, - Tuple, - BF16, - PassThrough, - PassThrough, - MultiplyMultiply>>>& instances) -{ - add_device_operation_instances( - instances, - device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_comp_instances{}); -} - -} // namespace instance -} // namespace device -} // namespace tensor_operation -} // namespace ck diff --git a/library/src/tensor_operation_instance/gpu/gemm_multiply_multiply/device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_comp_mnpadding_instance.cpp b/library/src/tensor_operation_instance/gpu/gemm_multiply_multiply/device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_comp_mnpadding_instance.cpp deleted file mode 100644 index e76f4f82b3..0000000000 --- a/library/src/tensor_operation_instance/gpu/gemm_multiply_multiply/device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_comp_mnpadding_instance.cpp +++ /dev/null @@ -1,32 +0,0 @@ -// SPDX-License-Identifier: MIT -// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved. - -#include "device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn.hpp" - -namespace ck { -namespace tensor_operation { -namespace device { -namespace instance { - -void add_device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_comp_mnpadding_instances( - std::vector, - Row, - F8, - F8, - Tuple, - BF16, - PassThrough, - PassThrough, - MultiplyMultiply>>>& instances) -{ - add_device_operation_instances( - instances, - device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_comp_instances{}); -} - -} // namespace instance -} // namespace device -} // namespace tensor_operation -} // namespace ck diff --git a/library/src/tensor_operation_instance/gpu/gemm_multiply_multiply/device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_mem_v1_default_instance.cpp b/library/src/tensor_operation_instance/gpu/gemm_multiply_multiply/device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_mem_v1_default_instance.cpp index 03f360a457..d2e64be2f6 100644 --- a/library/src/tensor_operation_instance/gpu/gemm_multiply_multiply/device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_mem_v1_default_instance.cpp +++ b/library/src/tensor_operation_instance/gpu/gemm_multiply_multiply/device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_mem_v1_default_instance.cpp @@ -9,17 +9,17 @@ namespace device { namespace instance { void add_device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_mem_v1_default_instances( - std::vector, - Row, - F8, - F8, - Tuple, - BF16, - PassThrough, - PassThrough, - MultiplyMultiply>>>& instances) + std::vector, + Row, + F8, + F8, + Tuple, + BF16, + PassThrough, + PassThrough, + MultiplyMultiply>>>& instances) { add_device_operation_instances( instances, diff --git a/library/src/tensor_operation_instance/gpu/gemm_multiply_multiply/device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_mem_v1_kpadding_instance.cpp b/library/src/tensor_operation_instance/gpu/gemm_multiply_multiply/device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_mem_v1_kpadding_instance.cpp index 194615e0fa..3a57f860f0 100644 --- a/library/src/tensor_operation_instance/gpu/gemm_multiply_multiply/device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_mem_v1_kpadding_instance.cpp +++ b/library/src/tensor_operation_instance/gpu/gemm_multiply_multiply/device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_mem_v1_kpadding_instance.cpp @@ -9,17 +9,17 @@ namespace device { namespace instance { void add_device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_mem_v1_kpadding_instances( - std::vector, - Row, - F8, - F8, - Tuple, - BF16, - PassThrough, - PassThrough, - MultiplyMultiply>>>& instances) + std::vector, + Row, + F8, + F8, + Tuple, + BF16, + PassThrough, + PassThrough, + MultiplyMultiply>>>& instances) { add_device_operation_instances( instances, diff --git a/library/src/tensor_operation_instance/gpu/gemm_multiply_multiply/device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_mem_v1_mnkpadding_instance.cpp b/library/src/tensor_operation_instance/gpu/gemm_multiply_multiply/device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_mem_v1_mnkpadding_instance.cpp deleted file mode 100644 index ae82b5800e..0000000000 --- a/library/src/tensor_operation_instance/gpu/gemm_multiply_multiply/device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_mem_v1_mnkpadding_instance.cpp +++ /dev/null @@ -1,33 +0,0 @@ -// SPDX-License-Identifier: MIT -// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved. - -#include "device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn.hpp" - -namespace ck { -namespace tensor_operation { -namespace device { -namespace instance { - -void add_device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_mem_v1_mnkpadding_instances( - std::vector, - Row, - F8, - F8, - Tuple, - BF16, - PassThrough, - PassThrough, - MultiplyMultiply>>>& instances) -{ - add_device_operation_instances( - instances, - device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_mem_instances{}); -} - -} // namespace instance -} // namespace device -} // namespace tensor_operation -} // namespace ck diff --git a/library/src/tensor_operation_instance/gpu/gemm_multiply_multiply/device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_mem_v2_default_instance.cpp b/library/src/tensor_operation_instance/gpu/gemm_multiply_multiply/device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_mem_v2_default_instance.cpp index 47bf0df2c7..1515021f6d 100644 --- a/library/src/tensor_operation_instance/gpu/gemm_multiply_multiply/device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_mem_v2_default_instance.cpp +++ b/library/src/tensor_operation_instance/gpu/gemm_multiply_multiply/device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_mem_v2_default_instance.cpp @@ -9,17 +9,17 @@ namespace device { namespace instance { void add_device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_mem_v2_default_instances( - std::vector, - Row, - F8, - F8, - Tuple, - BF16, - PassThrough, - PassThrough, - MultiplyMultiply>>>& instances) + std::vector, + Row, + F8, + F8, + Tuple, + BF16, + PassThrough, + PassThrough, + MultiplyMultiply>>>& instances) { add_device_operation_instances( instances, diff --git a/library/src/tensor_operation_instance/gpu/gemm_multiply_multiply/device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_mem_v2_kpadding_instance.cpp b/library/src/tensor_operation_instance/gpu/gemm_multiply_multiply/device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_mem_v2_kpadding_instance.cpp index 88ee816202..1b80244f84 100644 --- a/library/src/tensor_operation_instance/gpu/gemm_multiply_multiply/device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_mem_v2_kpadding_instance.cpp +++ b/library/src/tensor_operation_instance/gpu/gemm_multiply_multiply/device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_mem_v2_kpadding_instance.cpp @@ -9,17 +9,17 @@ namespace device { namespace instance { void add_device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_mem_v2_kpadding_instances( - std::vector, - Row, - F8, - F8, - Tuple, - BF16, - PassThrough, - PassThrough, - MultiplyMultiply>>>& instances) + std::vector, + Row, + F8, + F8, + Tuple, + BF16, + PassThrough, + PassThrough, + MultiplyMultiply>>>& instances) { add_device_operation_instances( instances, diff --git a/library/src/tensor_operation_instance/gpu/gemm_multiply_multiply/device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_mem_v2_mnkpadding_instance.cpp b/library/src/tensor_operation_instance/gpu/gemm_multiply_multiply/device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_mem_v2_mnkpadding_instance.cpp deleted file mode 100644 index 2c8784bedb..0000000000 --- a/library/src/tensor_operation_instance/gpu/gemm_multiply_multiply/device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_mem_v2_mnkpadding_instance.cpp +++ /dev/null @@ -1,33 +0,0 @@ -// SPDX-License-Identifier: MIT -// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved. - -#include "device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn.hpp" - -namespace ck { -namespace tensor_operation { -namespace device { -namespace instance { - -void add_device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_mem_v2_mnkpadding_instances( - std::vector, - Row, - F8, - F8, - Tuple, - BF16, - PassThrough, - PassThrough, - MultiplyMultiply>>>& instances) -{ - add_device_operation_instances( - instances, - device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_mem_instances{}); -} - -} // namespace instance -} // namespace device -} // namespace tensor_operation -} // namespace ck diff --git a/library/src/tensor_operation_instance/gpu/gemm_universal/CMakeLists.txt b/library/src/tensor_operation_instance/gpu/gemm_universal/CMakeLists.txt index fedd480c3f..cc4ce76606 100644 --- a/library/src/tensor_operation_instance/gpu/gemm_universal/CMakeLists.txt +++ b/library/src/tensor_operation_instance/gpu/gemm_universal/CMakeLists.txt @@ -14,56 +14,10 @@ list(APPEND GEMM_UNIVERSAL_INSTANCES device_gemm_xdl_universal_f16_f16_f16/device_gemm_xdl_universal_f16_f16_f16_mk_kn_mn_mem_v2_mnkpadding_instance.cpp device_gemm_xdl_universal_f16_f16_f16/device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_comp_default_instance.cpp device_gemm_xdl_universal_f16_f16_f16/device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_comp_kpadding_instance.cpp - device_gemm_xdl_universal_f16_f16_f16/device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_comp_mnpadding_instance.cpp - device_gemm_xdl_universal_f16_f16_f16/device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_comp_mnkpadding_instance.cpp device_gemm_xdl_universal_f16_f16_f16/device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_mem_v1_default_instance.cpp device_gemm_xdl_universal_f16_f16_f16/device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_mem_v1_kpadding_instance.cpp - device_gemm_xdl_universal_f16_f16_f16/device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_mem_v1_mnkpadding_instance.cpp device_gemm_xdl_universal_f16_f16_f16/device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_mem_v2_default_instance.cpp device_gemm_xdl_universal_f16_f16_f16/device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_mem_v2_kpadding_instance.cpp - device_gemm_xdl_universal_f16_f16_f16/device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_mem_v2_mnkpadding_instance.cpp - - device_gemm_xdl_universal_f16_f8_f16/device_gemm_xdl_universal_f16_f8_f16_mk_kn_mn_comp_default_instance.cpp - device_gemm_xdl_universal_f16_f8_f16/device_gemm_xdl_universal_f16_f8_f16_mk_kn_mn_comp_kpadding_instance.cpp - device_gemm_xdl_universal_f16_f8_f16/device_gemm_xdl_universal_f16_f8_f16_mk_kn_mn_comp_mnpadding_instance.cpp - device_gemm_xdl_universal_f16_f8_f16/device_gemm_xdl_universal_f16_f8_f16_mk_kn_mn_comp_mnkpadding_instance.cpp - device_gemm_xdl_universal_f16_f8_f16/device_gemm_xdl_universal_f16_f8_f16_mk_kn_mn_mem_v1_default_instance.cpp - device_gemm_xdl_universal_f16_f8_f16/device_gemm_xdl_universal_f16_f8_f16_mk_kn_mn_mem_v1_kpadding_instance.cpp - device_gemm_xdl_universal_f16_f8_f16/device_gemm_xdl_universal_f16_f8_f16_mk_kn_mn_mem_v1_mnkpadding_instance.cpp - device_gemm_xdl_universal_f16_f8_f16/device_gemm_xdl_universal_f16_f8_f16_mk_kn_mn_mem_v2_default_instance.cpp - device_gemm_xdl_universal_f16_f8_f16/device_gemm_xdl_universal_f16_f8_f16_mk_kn_mn_mem_v2_kpadding_instance.cpp - device_gemm_xdl_universal_f16_f8_f16/device_gemm_xdl_universal_f16_f8_f16_mk_kn_mn_mem_v2_mnkpadding_instance.cpp - device_gemm_xdl_universal_f16_f8_f16/device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_comp_default_instance.cpp - device_gemm_xdl_universal_f16_f8_f16/device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_comp_kpadding_instance.cpp - device_gemm_xdl_universal_f16_f8_f16/device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_comp_mnpadding_instance.cpp - device_gemm_xdl_universal_f16_f8_f16/device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_comp_mnkpadding_instance.cpp - device_gemm_xdl_universal_f16_f8_f16/device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_mem_v1_default_instance.cpp - device_gemm_xdl_universal_f16_f8_f16/device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_mem_v1_kpadding_instance.cpp - device_gemm_xdl_universal_f16_f8_f16/device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_mem_v1_mnkpadding_instance.cpp - device_gemm_xdl_universal_f16_f8_f16/device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_mem_v2_default_instance.cpp - device_gemm_xdl_universal_f16_f8_f16/device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_mem_v2_kpadding_instance.cpp - device_gemm_xdl_universal_f16_f8_f16/device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_mem_v2_mnkpadding_instance.cpp - - device_gemm_xdl_universal_f8_f16_f16/device_gemm_xdl_universal_f8_f16_f16_mk_kn_mn_comp_default_instance.cpp - device_gemm_xdl_universal_f8_f16_f16/device_gemm_xdl_universal_f8_f16_f16_mk_kn_mn_comp_kpadding_instance.cpp - device_gemm_xdl_universal_f8_f16_f16/device_gemm_xdl_universal_f8_f16_f16_mk_kn_mn_comp_mnpadding_instance.cpp - device_gemm_xdl_universal_f8_f16_f16/device_gemm_xdl_universal_f8_f16_f16_mk_kn_mn_comp_mnkpadding_instance.cpp - device_gemm_xdl_universal_f8_f16_f16/device_gemm_xdl_universal_f8_f16_f16_mk_kn_mn_mem_v1_default_instance.cpp - device_gemm_xdl_universal_f8_f16_f16/device_gemm_xdl_universal_f8_f16_f16_mk_kn_mn_mem_v1_kpadding_instance.cpp - device_gemm_xdl_universal_f8_f16_f16/device_gemm_xdl_universal_f8_f16_f16_mk_kn_mn_mem_v1_mnkpadding_instance.cpp - device_gemm_xdl_universal_f8_f16_f16/device_gemm_xdl_universal_f8_f16_f16_mk_kn_mn_mem_v2_default_instance.cpp - device_gemm_xdl_universal_f8_f16_f16/device_gemm_xdl_universal_f8_f16_f16_mk_kn_mn_mem_v2_kpadding_instance.cpp - device_gemm_xdl_universal_f8_f16_f16/device_gemm_xdl_universal_f8_f16_f16_mk_kn_mn_mem_v2_mnkpadding_instance.cpp - device_gemm_xdl_universal_f8_f16_f16/device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_comp_default_instance.cpp - device_gemm_xdl_universal_f8_f16_f16/device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_comp_kpadding_instance.cpp - device_gemm_xdl_universal_f8_f16_f16/device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_comp_mnpadding_instance.cpp - device_gemm_xdl_universal_f8_f16_f16/device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_comp_mnkpadding_instance.cpp - device_gemm_xdl_universal_f8_f16_f16/device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_mem_v1_default_instance.cpp - device_gemm_xdl_universal_f8_f16_f16/device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_mem_v1_kpadding_instance.cpp - device_gemm_xdl_universal_f8_f16_f16/device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_mem_v1_mnkpadding_instance.cpp - device_gemm_xdl_universal_f8_f16_f16/device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_mem_v2_default_instance.cpp - device_gemm_xdl_universal_f8_f16_f16/device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_mem_v2_kpadding_instance.cpp - device_gemm_xdl_universal_f8_f16_f16/device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_mem_v2_mnkpadding_instance.cpp device_gemm_xdl_universal_bf16_bf16_bf16/device_gemm_xdl_universal_bf16_bf16_bf16_mk_kn_mn_comp_default_instance.cpp device_gemm_xdl_universal_bf16_bf16_bf16/device_gemm_xdl_universal_bf16_bf16_bf16_mk_kn_mn_comp_kpadding_instance.cpp @@ -77,25 +31,98 @@ list(APPEND GEMM_UNIVERSAL_INSTANCES device_gemm_xdl_universal_bf16_bf16_bf16/device_gemm_xdl_universal_bf16_bf16_bf16_mk_kn_mn_mem_v2_mnkpadding_instance.cpp device_gemm_xdl_universal_bf16_bf16_bf16/device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_comp_default_instance.cpp device_gemm_xdl_universal_bf16_bf16_bf16/device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_comp_kpadding_instance.cpp - device_gemm_xdl_universal_bf16_bf16_bf16/device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_comp_mnpadding_instance.cpp - device_gemm_xdl_universal_bf16_bf16_bf16/device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_comp_mnkpadding_instance.cpp device_gemm_xdl_universal_bf16_bf16_bf16/device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_mem_v1_default_instance.cpp device_gemm_xdl_universal_bf16_bf16_bf16/device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_mem_v1_kpadding_instance.cpp - device_gemm_xdl_universal_bf16_bf16_bf16/device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_mem_v1_mnkpadding_instance.cpp device_gemm_xdl_universal_bf16_bf16_bf16/device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_mem_v2_default_instance.cpp device_gemm_xdl_universal_bf16_bf16_bf16/device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_mem_v2_kpadding_instance.cpp - device_gemm_xdl_universal_bf16_bf16_bf16/device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_mem_v2_mnkpadding_instance.cpp - - device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_comp_default_instance.cpp - device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_comp_kpadding_instance.cpp - device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_comp_mnpadding_instance.cpp - device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_comp_mnkpadding_instance.cpp - device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_mem_v1_default_instance.cpp - device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_mem_v1_kpadding_instance.cpp - device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_mem_v1_mnkpadding_instance.cpp - device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_mem_v2_default_instance.cpp - device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_mem_v2_kpadding_instance.cpp - device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_mem_v2_mnkpadding_instance.cpp ) +set_source_files_properties(device_gemm_xdl_universal_f16_f16_f16/device_gemm_xdl_universal_f16_f16_f16_mk_kn_mn_comp_default_instance.cpp PROPERTIES COMPILE_OPTIONS ";-mllvm;-greedy-reverse-local-assignment=1") +set_source_files_properties(device_gemm_xdl_universal_f16_f16_f16/device_gemm_xdl_universal_f16_f16_f16_mk_kn_mn_comp_kpadding_instance.cpp PROPERTIES COMPILE_OPTIONS ";-mllvm;-greedy-reverse-local-assignment=1") +set_source_files_properties(device_gemm_xdl_universal_f16_f16_f16/device_gemm_xdl_universal_f16_f16_f16_mk_kn_mn_comp_mnpadding_instance.cpp PROPERTIES COMPILE_OPTIONS ";-mllvm;-greedy-reverse-local-assignment=1") +set_source_files_properties(device_gemm_xdl_universal_f16_f16_f16/device_gemm_xdl_universal_f16_f16_f16_mk_kn_mn_comp_mnkpadding_instance.cpp PROPERTIES COMPILE_OPTIONS ";-mllvm;-greedy-reverse-local-assignment=1") +set_source_files_properties(device_gemm_xdl_universal_f16_f16_f16/device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_comp_default_instance.cpp PROPERTIES COMPILE_OPTIONS ";-mllvm;-greedy-reverse-local-assignment=1") +set_source_files_properties(device_gemm_xdl_universal_f16_f16_f16/device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_comp_kpadding_instance.cpp PROPERTIES COMPILE_OPTIONS ";-mllvm;-greedy-reverse-local-assignment=1") + +set_source_files_properties(device_gemm_xdl_universal_bf16_bf16_bf16/device_gemm_xdl_universal_bf16_bf16_bf16_mk_kn_mn_comp_default_instance.cpp PROPERTIES COMPILE_OPTIONS ";-mllvm;-greedy-reverse-local-assignment=1") +set_source_files_properties(device_gemm_xdl_universal_bf16_bf16_bf16/device_gemm_xdl_universal_bf16_bf16_bf16_mk_kn_mn_comp_kpadding_instance.cpp PROPERTIES COMPILE_OPTIONS ";-mllvm;-greedy-reverse-local-assignment=1") +set_source_files_properties(device_gemm_xdl_universal_bf16_bf16_bf16/device_gemm_xdl_universal_bf16_bf16_bf16_mk_kn_mn_comp_mnpadding_instance.cpp PROPERTIES COMPILE_OPTIONS ";-mllvm;-greedy-reverse-local-assignment=1") +set_source_files_properties(device_gemm_xdl_universal_bf16_bf16_bf16/device_gemm_xdl_universal_bf16_bf16_bf16_mk_kn_mn_comp_mnkpadding_instance.cpp PROPERTIES COMPILE_OPTIONS ";-mllvm;-greedy-reverse-local-assignment=1") +set_source_files_properties(device_gemm_xdl_universal_bf16_bf16_bf16/device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_comp_default_instance.cpp PROPERTIES COMPILE_OPTIONS ";-mllvm;-greedy-reverse-local-assignment=1") +set_source_files_properties(device_gemm_xdl_universal_bf16_bf16_bf16/device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_comp_kpadding_instance.cpp PROPERTIES COMPILE_OPTIONS ";-mllvm;-greedy-reverse-local-assignment=1") + +if(GPU_TARGETS MATCHES "gfx94") + list(APPEND GEMM_UNIVERSAL_INSTANCES + device_gemm_xdl_universal_f16_f8_f16/device_gemm_xdl_universal_f16_f8_f16_mk_kn_mn_comp_default_instance.cpp + device_gemm_xdl_universal_f16_f8_f16/device_gemm_xdl_universal_f16_f8_f16_mk_kn_mn_comp_kpadding_instance.cpp + device_gemm_xdl_universal_f16_f8_f16/device_gemm_xdl_universal_f16_f8_f16_mk_kn_mn_comp_mnpadding_instance.cpp + device_gemm_xdl_universal_f16_f8_f16/device_gemm_xdl_universal_f16_f8_f16_mk_kn_mn_comp_mnkpadding_instance.cpp + device_gemm_xdl_universal_f16_f8_f16/device_gemm_xdl_universal_f16_f8_f16_mk_kn_mn_mem_v1_default_instance.cpp + device_gemm_xdl_universal_f16_f8_f16/device_gemm_xdl_universal_f16_f8_f16_mk_kn_mn_mem_v1_kpadding_instance.cpp + device_gemm_xdl_universal_f16_f8_f16/device_gemm_xdl_universal_f16_f8_f16_mk_kn_mn_mem_v1_mnkpadding_instance.cpp + device_gemm_xdl_universal_f16_f8_f16/device_gemm_xdl_universal_f16_f8_f16_mk_kn_mn_mem_v2_default_instance.cpp + device_gemm_xdl_universal_f16_f8_f16/device_gemm_xdl_universal_f16_f8_f16_mk_kn_mn_mem_v2_kpadding_instance.cpp + device_gemm_xdl_universal_f16_f8_f16/device_gemm_xdl_universal_f16_f8_f16_mk_kn_mn_mem_v2_mnkpadding_instance.cpp + device_gemm_xdl_universal_f16_f8_f16/device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_comp_default_instance.cpp + device_gemm_xdl_universal_f16_f8_f16/device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_comp_kpadding_instance.cpp + device_gemm_xdl_universal_f16_f8_f16/device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_mem_v1_default_instance.cpp + device_gemm_xdl_universal_f16_f8_f16/device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_mem_v1_kpadding_instance.cpp + device_gemm_xdl_universal_f16_f8_f16/device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_mem_v2_default_instance.cpp + device_gemm_xdl_universal_f16_f8_f16/device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_mem_v2_kpadding_instance.cpp + + device_gemm_xdl_universal_f8_f16_f16/device_gemm_xdl_universal_f8_f16_f16_mk_kn_mn_comp_default_instance.cpp + device_gemm_xdl_universal_f8_f16_f16/device_gemm_xdl_universal_f8_f16_f16_mk_kn_mn_comp_kpadding_instance.cpp + device_gemm_xdl_universal_f8_f16_f16/device_gemm_xdl_universal_f8_f16_f16_mk_kn_mn_comp_mnpadding_instance.cpp + device_gemm_xdl_universal_f8_f16_f16/device_gemm_xdl_universal_f8_f16_f16_mk_kn_mn_comp_mnkpadding_instance.cpp + device_gemm_xdl_universal_f8_f16_f16/device_gemm_xdl_universal_f8_f16_f16_mk_kn_mn_mem_v1_default_instance.cpp + device_gemm_xdl_universal_f8_f16_f16/device_gemm_xdl_universal_f8_f16_f16_mk_kn_mn_mem_v1_kpadding_instance.cpp + device_gemm_xdl_universal_f8_f16_f16/device_gemm_xdl_universal_f8_f16_f16_mk_kn_mn_mem_v1_mnkpadding_instance.cpp + device_gemm_xdl_universal_f8_f16_f16/device_gemm_xdl_universal_f8_f16_f16_mk_kn_mn_mem_v2_default_instance.cpp + device_gemm_xdl_universal_f8_f16_f16/device_gemm_xdl_universal_f8_f16_f16_mk_kn_mn_mem_v2_kpadding_instance.cpp + device_gemm_xdl_universal_f8_f16_f16/device_gemm_xdl_universal_f8_f16_f16_mk_kn_mn_mem_v2_mnkpadding_instance.cpp + device_gemm_xdl_universal_f8_f16_f16/device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_comp_default_instance.cpp + device_gemm_xdl_universal_f8_f16_f16/device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_comp_kpadding_instance.cpp + device_gemm_xdl_universal_f8_f16_f16/device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_mem_v1_default_instance.cpp + device_gemm_xdl_universal_f8_f16_f16/device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_mem_v1_kpadding_instance.cpp + device_gemm_xdl_universal_f8_f16_f16/device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_mem_v2_default_instance.cpp + device_gemm_xdl_universal_f8_f16_f16/device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_mem_v2_kpadding_instance.cpp + + device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_comp_default_instance.cpp + device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_comp_kpadding_instance.cpp + device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_comp_nkpadding_instance.cpp + device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_v1_default_instance.cpp + device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_v1_kpadding_instance.cpp + device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_v1_nkpadding_instance.cpp + device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_v2_default_instance.cpp + device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_v2_kpadding_instance.cpp + device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_v2_nkpadding_instance.cpp + device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_comp_default_instance.cpp + device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_comp_kpadding_instance.cpp + device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_mem_v1_default_instance.cpp + device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_mem_v1_kpadding_instance.cpp + device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_mem_v2_default_instance.cpp + device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_mem_v2_kpadding_instance.cpp + ) + + set_source_files_properties(device_gemm_xdl_universal_f16_f8_f16/device_gemm_xdl_universal_f16_f8_f16_mk_kn_mn_comp_default_instance.cpp PROPERTIES COMPILE_OPTIONS ";-mllvm;-greedy-reverse-local-assignment=1") + set_source_files_properties(device_gemm_xdl_universal_f16_f8_f16/device_gemm_xdl_universal_f16_f8_f16_mk_kn_mn_comp_kpadding_instance.cpp PROPERTIES COMPILE_OPTIONS ";-mllvm;-greedy-reverse-local-assignment=1") + set_source_files_properties(device_gemm_xdl_universal_f16_f8_f16/device_gemm_xdl_universal_f16_f8_f16_mk_kn_mn_comp_mnpadding_instance.cpp PROPERTIES COMPILE_OPTIONS ";-mllvm;-greedy-reverse-local-assignment=1") + set_source_files_properties(device_gemm_xdl_universal_f16_f8_f16/device_gemm_xdl_universal_f16_f8_f16_mk_kn_mn_comp_mnkpadding_instance.cpp PROPERTIES COMPILE_OPTIONS ";-mllvm;-greedy-reverse-local-assignment=1") + set_source_files_properties(device_gemm_xdl_universal_f16_f8_f16/device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_comp_default_instance.cpp PROPERTIES COMPILE_OPTIONS ";-mllvm;-greedy-reverse-local-assignment=1") + set_source_files_properties(device_gemm_xdl_universal_f16_f8_f16/device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_comp_kpadding_instance.cpp PROPERTIES COMPILE_OPTIONS ";-mllvm;-greedy-reverse-local-assignment=1") + + set_source_files_properties(device_gemm_xdl_universal_f8_f16_f16/device_gemm_xdl_universal_f8_f16_f16_mk_kn_mn_comp_default_instance.cpp PROPERTIES COMPILE_OPTIONS ";-mllvm;-greedy-reverse-local-assignment=1") + set_source_files_properties(device_gemm_xdl_universal_f8_f16_f16/device_gemm_xdl_universal_f8_f16_f16_mk_kn_mn_comp_kpadding_instance.cpp PROPERTIES COMPILE_OPTIONS ";-mllvm;-greedy-reverse-local-assignment=1") + set_source_files_properties(device_gemm_xdl_universal_f8_f16_f16/device_gemm_xdl_universal_f8_f16_f16_mk_kn_mn_comp_mnpadding_instance.cpp PROPERTIES COMPILE_OPTIONS ";-mllvm;-greedy-reverse-local-assignment=1") + set_source_files_properties(device_gemm_xdl_universal_f8_f16_f16/device_gemm_xdl_universal_f8_f16_f16_mk_kn_mn_comp_mnkpadding_instance.cpp PROPERTIES COMPILE_OPTIONS ";-mllvm;-greedy-reverse-local-assignment=1") + set_source_files_properties(device_gemm_xdl_universal_f8_f16_f16/device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_comp_default_instance.cpp PROPERTIES COMPILE_OPTIONS ";-mllvm;-greedy-reverse-local-assignment=1") + set_source_files_properties(device_gemm_xdl_universal_f8_f16_f16/device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_comp_kpadding_instance.cpp PROPERTIES COMPILE_OPTIONS ";-mllvm;-greedy-reverse-local-assignment=1") + + set_source_files_properties(device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_comp_default_instance.cpp PROPERTIES COMPILE_OPTIONS ";-mllvm;-greedy-reverse-local-assignment=1") + set_source_files_properties(device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_comp_kpadding_instance.cpp PROPERTIES COMPILE_OPTIONS ";-mllvm;-greedy-reverse-local-assignment=1") + set_source_files_properties(device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_comp_nkpadding_instance.cpp PROPERTIES COMPILE_OPTIONS ";-mllvm;-greedy-reverse-local-assignment=1") + set_source_files_properties(device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_comp_default_instance.cpp PROPERTIES COMPILE_OPTIONS ";-mllvm;-greedy-reverse-local-assignment=1") + set_source_files_properties(device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_comp_kpadding_instance.cpp PROPERTIES COMPILE_OPTIONS ";-mllvm;-greedy-reverse-local-assignment=1") +endif() + add_instance_library(device_gemm_universal_instance ${GEMM_UNIVERSAL_INSTANCES}) diff --git a/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_bf16_bf16_bf16/device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_comp_mnkpadding_instance.cpp b/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_bf16_bf16_bf16/device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_comp_mnkpadding_instance.cpp deleted file mode 100644 index b95ea76652..0000000000 --- a/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_bf16_bf16_bf16/device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_comp_mnkpadding_instance.cpp +++ /dev/null @@ -1,24 +0,0 @@ -// SPDX-License-Identifier: MIT -// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved. - -#include "device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn.hpp" - -namespace ck { -namespace tensor_operation { -namespace device { -namespace instance { - -void add_device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_comp_mnkpadding_instances( - std::vector>>& - instances) -{ - add_device_operation_instances( - instances, - device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_comp_instances{}); -} - -} // namespace instance -} // namespace device -} // namespace tensor_operation -} // namespace ck diff --git a/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_bf16_bf16_bf16/device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_comp_mnpadding_instance.cpp b/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_bf16_bf16_bf16/device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_comp_mnpadding_instance.cpp deleted file mode 100644 index 65af442696..0000000000 --- a/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_bf16_bf16_bf16/device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_comp_mnpadding_instance.cpp +++ /dev/null @@ -1,24 +0,0 @@ -// SPDX-License-Identifier: MIT -// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved. - -#include "device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn.hpp" - -namespace ck { -namespace tensor_operation { -namespace device { -namespace instance { - -void add_device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_comp_mnpadding_instances( - std::vector>>& - instances) -{ - add_device_operation_instances( - instances, - device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_comp_instances{}); -} - -} // namespace instance -} // namespace device -} // namespace tensor_operation -} // namespace ck diff --git a/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_bf16_bf16_bf16/device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_mem_v1_mnkpadding_instance.cpp b/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_bf16_bf16_bf16/device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_mem_v1_mnkpadding_instance.cpp deleted file mode 100644 index 4142c5e8b9..0000000000 --- a/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_bf16_bf16_bf16/device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_mem_v1_mnkpadding_instance.cpp +++ /dev/null @@ -1,25 +0,0 @@ -// SPDX-License-Identifier: MIT -// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved. - -#include "device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn.hpp" - -namespace ck { -namespace tensor_operation { -namespace device { -namespace instance { - -void add_device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_mem_v1_mnkpadding_instances( - std::vector>>& - instances) -{ - add_device_operation_instances( - instances, - device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_mem_instances{}); -} - -} // namespace instance -} // namespace device -} // namespace tensor_operation -} // namespace ck diff --git a/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_bf16_bf16_bf16/device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_mem_v2_mnkpadding_instance.cpp b/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_bf16_bf16_bf16/device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_mem_v2_mnkpadding_instance.cpp deleted file mode 100644 index 059c88ddd6..0000000000 --- a/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_bf16_bf16_bf16/device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_mem_v2_mnkpadding_instance.cpp +++ /dev/null @@ -1,25 +0,0 @@ -// SPDX-License-Identifier: MIT -// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved. - -#include "device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn.hpp" - -namespace ck { -namespace tensor_operation { -namespace device { -namespace instance { - -void add_device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_mem_v2_mnkpadding_instances( - std::vector>>& - instances) -{ - add_device_operation_instances( - instances, - device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_mem_instances{}); -} - -} // namespace instance -} // namespace device -} // namespace tensor_operation -} // namespace ck diff --git a/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f16_f16_f16/device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_comp_mnkpadding_instance.cpp b/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f16_f16_f16/device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_comp_mnkpadding_instance.cpp deleted file mode 100644 index 56fb3a129c..0000000000 --- a/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f16_f16_f16/device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_comp_mnkpadding_instance.cpp +++ /dev/null @@ -1,23 +0,0 @@ -// SPDX-License-Identifier: MIT -// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved. - -#include "device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn.hpp" - -namespace ck { -namespace tensor_operation { -namespace device { -namespace instance { - -void add_device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_comp_mnkpadding_instances( - std::vector>>& - instances) -{ - add_device_operation_instances( - instances, device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_comp_instances{}); -} - -} // namespace instance -} // namespace device -} // namespace tensor_operation -} // namespace ck diff --git a/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f16_f16_f16/device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_comp_mnpadding_instance.cpp b/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f16_f16_f16/device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_comp_mnpadding_instance.cpp deleted file mode 100644 index f63817ce53..0000000000 --- a/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f16_f16_f16/device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_comp_mnpadding_instance.cpp +++ /dev/null @@ -1,23 +0,0 @@ -// SPDX-License-Identifier: MIT -// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved. - -#include "device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn.hpp" - -namespace ck { -namespace tensor_operation { -namespace device { -namespace instance { - -void add_device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_comp_mnpadding_instances( - std::vector>>& - instances) -{ - add_device_operation_instances( - instances, device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_comp_instances{}); -} - -} // namespace instance -} // namespace device -} // namespace tensor_operation -} // namespace ck diff --git a/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f16_f16_f16/device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_mem_v1_mnkpadding_instance.cpp b/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f16_f16_f16/device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_mem_v1_mnkpadding_instance.cpp deleted file mode 100644 index 2aa313851f..0000000000 --- a/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f16_f16_f16/device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_mem_v1_mnkpadding_instance.cpp +++ /dev/null @@ -1,24 +0,0 @@ -// SPDX-License-Identifier: MIT -// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved. - -#include "device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn.hpp" - -namespace ck { -namespace tensor_operation { -namespace device { -namespace instance { - -void add_device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_mem_v1_mnkpadding_instances( - std::vector>>& - instances) -{ - add_device_operation_instances( - instances, - device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_mem_instances{}); -} - -} // namespace instance -} // namespace device -} // namespace tensor_operation -} // namespace ck diff --git a/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f16_f16_f16/device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_mem_v2_mnkpadding_instance.cpp b/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f16_f16_f16/device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_mem_v2_mnkpadding_instance.cpp deleted file mode 100644 index a1928ccc63..0000000000 --- a/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f16_f16_f16/device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_mem_v2_mnkpadding_instance.cpp +++ /dev/null @@ -1,24 +0,0 @@ -// SPDX-License-Identifier: MIT -// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved. - -#include "device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn.hpp" - -namespace ck { -namespace tensor_operation { -namespace device { -namespace instance { - -void add_device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_mem_v2_mnkpadding_instances( - std::vector>>& - instances) -{ - add_device_operation_instances( - instances, - device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_mem_instances{}); -} - -} // namespace instance -} // namespace device -} // namespace tensor_operation -} // namespace ck diff --git a/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f16_f8_f16/device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_comp_mnkpadding_instance.cpp b/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f16_f8_f16/device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_comp_mnkpadding_instance.cpp deleted file mode 100644 index 203ada9a37..0000000000 --- a/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f16_f8_f16/device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_comp_mnkpadding_instance.cpp +++ /dev/null @@ -1,23 +0,0 @@ -// SPDX-License-Identifier: MIT -// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved. - -#include "device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn.hpp" - -namespace ck { -namespace tensor_operation { -namespace device { -namespace instance { - -void add_device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_comp_mnkpadding_instances( - std::vector>>& - instances) -{ - add_device_operation_instances( - instances, device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_comp_instances{}); -} - -} // namespace instance -} // namespace device -} // namespace tensor_operation -} // namespace ck diff --git a/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f16_f8_f16/device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_comp_mnpadding_instance.cpp b/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f16_f8_f16/device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_comp_mnpadding_instance.cpp deleted file mode 100644 index da38705672..0000000000 --- a/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f16_f8_f16/device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_comp_mnpadding_instance.cpp +++ /dev/null @@ -1,26 +0,0 @@ -// SPDX-License-Identifier: MIT -// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved. - -#include "device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn.hpp" - -namespace ck { -namespace tensor_operation { -namespace device { -namespace instance { - -using F16 = ck::half_t; -using F32 = float; - -void add_device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_comp_mnpadding_instances( - std::vector>>& - instances) -{ - add_device_operation_instances( - instances, device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_comp_instances{}); -} - -} // namespace instance -} // namespace device -} // namespace tensor_operation -} // namespace ck diff --git a/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f16_f8_f16/device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_mem_v1_mnkpadding_instance.cpp b/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f16_f8_f16/device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_mem_v1_mnkpadding_instance.cpp deleted file mode 100644 index 13f34a1c58..0000000000 --- a/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f16_f8_f16/device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_mem_v1_mnkpadding_instance.cpp +++ /dev/null @@ -1,24 +0,0 @@ -// SPDX-License-Identifier: MIT -// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved. - -#include "device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn.hpp" - -namespace ck { -namespace tensor_operation { -namespace device { -namespace instance { - -void add_device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_mem_v1_mnkpadding_instances( - std::vector>>& - instances) -{ - add_device_operation_instances( - instances, - device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_mem_instances{}); -} - -} // namespace instance -} // namespace device -} // namespace tensor_operation -} // namespace ck diff --git a/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f16_f8_f16/device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_mem_v2_mnkpadding_instance.cpp b/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f16_f8_f16/device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_mem_v2_mnkpadding_instance.cpp deleted file mode 100644 index b601313ff1..0000000000 --- a/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f16_f8_f16/device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_mem_v2_mnkpadding_instance.cpp +++ /dev/null @@ -1,24 +0,0 @@ -// SPDX-License-Identifier: MIT -// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved. - -#include "device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn.hpp" - -namespace ck { -namespace tensor_operation { -namespace device { -namespace instance { - -void add_device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_mem_v2_mnkpadding_instances( - std::vector>>& - instances) -{ - add_device_operation_instances( - instances, - device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_mem_instances{}); -} - -} // namespace instance -} // namespace device -} // namespace tensor_operation -} // namespace ck diff --git a/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f16_f16/device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_comp_mnkpadding_instance.cpp b/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f16_f16/device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_comp_mnkpadding_instance.cpp deleted file mode 100644 index 74299cd552..0000000000 --- a/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f16_f16/device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_comp_mnkpadding_instance.cpp +++ /dev/null @@ -1,23 +0,0 @@ -// SPDX-License-Identifier: MIT -// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved. - -#include "device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn.hpp" - -namespace ck { -namespace tensor_operation { -namespace device { -namespace instance { - -void add_device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_comp_mnkpadding_instances( - std::vector>>& - instances) -{ - add_device_operation_instances( - instances, device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_comp_instances{}); -} - -} // namespace instance -} // namespace device -} // namespace tensor_operation -} // namespace ck diff --git a/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f16_f16/device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_comp_mnpadding_instance.cpp b/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f16_f16/device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_comp_mnpadding_instance.cpp deleted file mode 100644 index d0561c0af5..0000000000 --- a/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f16_f16/device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_comp_mnpadding_instance.cpp +++ /dev/null @@ -1,26 +0,0 @@ -// SPDX-License-Identifier: MIT -// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved. - -#include "device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn.hpp" - -namespace ck { -namespace tensor_operation { -namespace device { -namespace instance { - -using F16 = ck::half_t; -using F32 = float; - -void add_device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_comp_mnpadding_instances( - std::vector>>& - instances) -{ - add_device_operation_instances( - instances, device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_comp_instances{}); -} - -} // namespace instance -} // namespace device -} // namespace tensor_operation -} // namespace ck diff --git a/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f16_f16/device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_mem_v1_mnkpadding_instance.cpp b/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f16_f16/device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_mem_v1_mnkpadding_instance.cpp deleted file mode 100644 index 2814dee434..0000000000 --- a/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f16_f16/device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_mem_v1_mnkpadding_instance.cpp +++ /dev/null @@ -1,24 +0,0 @@ -// SPDX-License-Identifier: MIT -// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved. - -#include "device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn.hpp" - -namespace ck { -namespace tensor_operation { -namespace device { -namespace instance { - -void add_device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_mem_v1_mnkpadding_instances( - std::vector>>& - instances) -{ - add_device_operation_instances( - instances, - device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_mem_instances{}); -} - -} // namespace instance -} // namespace device -} // namespace tensor_operation -} // namespace ck diff --git a/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f16_f16/device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_mem_v2_mnkpadding_instance.cpp b/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f16_f16/device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_mem_v2_mnkpadding_instance.cpp deleted file mode 100644 index ae0816ffc4..0000000000 --- a/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f16_f16/device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_mem_v2_mnkpadding_instance.cpp +++ /dev/null @@ -1,24 +0,0 @@ -// SPDX-License-Identifier: MIT -// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved. - -#include "device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn.hpp" - -namespace ck { -namespace tensor_operation { -namespace device { -namespace instance { - -void add_device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_mem_v2_mnkpadding_instances( - std::vector>>& - instances) -{ - add_device_operation_instances( - instances, - device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_mem_instances{}); -} - -} // namespace instance -} // namespace device -} // namespace tensor_operation -} // namespace ck diff --git a/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn.hpp b/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn.hpp new file mode 100644 index 0000000000..12994aeecd --- /dev/null +++ b/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn.hpp @@ -0,0 +1,96 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved. + +#include "ck/ck.hpp" +#include "ck/tensor_operation/gpu/device/tensor_layout.hpp" +#include "ck/tensor_operation/gpu/device/gemm_specialization.hpp" +#include "ck/tensor_operation/gpu/device/impl/device_gemm_xdl_cshuffle_v3.hpp" + +#include "ck/library/tensor_operation_instance/add_device_operation_instance.hpp" + +namespace ck { +namespace tensor_operation { +namespace device { +namespace instance { + +using F8 = f8_t; +using BF16 = bhalf_t; +using F32 = float; + +using Row = tensor_layout::gemm::RowMajor; + +template +using S = Sequence; + +using PassThrough = element_wise::PassThrough; + +static constexpr auto GemmDefault = GemmSpecialization::Default; +static constexpr auto GemmKPadding = GemmSpecialization::KPadding; +static constexpr auto GemmMNPadding = GemmSpecialization::MNPadding; +static constexpr auto GemmNKPadding = GemmSpecialization::NKPadding; +static constexpr auto GemmMNKPadding = GemmSpecialization::MNKPadding; + +static constexpr auto Intrawave = BlockGemmPipelineScheduler::Intrawave; +static constexpr auto Interwave = BlockGemmPipelineScheduler::Interwave; + +template +using device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_comp_instances = std::tuple< + // clang-format off + //#########################| ALayout| BLayout| CLayout|AData| BData| CData| AccData| Cshuffle| A| B| C| GEMM| Block| MPer| NPer| KPer| AK1| BK1|MPer| NPer| MXdl| NXdl| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockLds| BBlockTransfer| BBlockTransfer| BBlockTransfer| BlockTransfer| BBlockTransfer| BBlockTransfer| BBlockLds| CShuffle| CShuffle| CBlockTransferClusterLengths| CBlockTransfer| Block-wiseGemm| Block-wiseGemm| + //#########################| | | | 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| + //#########################| | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | + + // Compute friendly + DeviceGemm_Xdl_CShuffleV3< Row, Row, Row, F8, F8, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 256, 256, 64, 16, 4, 32, 32, 4, 4, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, S<16, 16, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 16, 4, 0, 1, 1, S<1, 32, 1, 8>, 8, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v4, F8>, + DeviceGemm_Xdl_CShuffleV3< Row, Row, Row, F8, F8, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 128, 128, 128, 16, 4, 32, 32, 2, 2, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, S<32, 8, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 16, 4, 0, 1, 1, S<1, 32, 1, 8>, 8, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v4, F8>, + DeviceGemm_Xdl_CShuffleV3< Row, Row, Row, F8, F8, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 128, 128, 64, 16, 4, 32, 32, 2, 2, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 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, F8>, + DeviceGemm_Xdl_CShuffleV3< Row, Row, Row, F8, F8, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 256, 256, 128, 16, 8, 16, 16, 8, 8, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, S<16, 16, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 16, 8, 0, 1, 2, S<1, 32, 1, 8>, 8, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v3, F8>, + DeviceGemm_Xdl_CShuffleV3< Row, Row, Row, F8, F8, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 256, 256, 64, 16, 4, 16, 16, 8, 8, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, S<16, 16, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 16, 4, 0, 1, 2, S<1, 32, 1, 8>, 8, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v3, F8>, + DeviceGemm_Xdl_CShuffleV3< Row, Row, Row, F8, F8, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 224, 256, 128, 16, 8, 16, 16, 7, 8, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, S<16, 16, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 16, 8, 0, 1, 2, S<1, 32, 1, 8>, 8, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v3, F8>, + DeviceGemm_Xdl_CShuffleV3< Row, Row, Row, F8, F8, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 128, 128, 128, 16, 4, 32, 32, 2, 2, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, S<32, 8, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 16, 4, 0, 1, 1, S<1, 32, 1, 8>, 8, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v3, F8>, + DeviceGemm_Xdl_CShuffleV3< Row, Row, Row, F8, F8, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 128, 128, 128, 16, 4, 32, 32, 2, 2, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, S<32, 8, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 16, 4, 0, 1, 1, S<1, 32, 1, 8>, 8, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v5, F8>, + DeviceGemm_Xdl_CShuffleV3< Row, Row, Row, F8, F8, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 128, 256, 64, 16, 4, 32, 32, 2, 4, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, S<16, 16, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 16, 4, 0, 1, 1, S<1, 32, 1, 8>, 8, BlockGemmPipelineScheduler::Interwave, BlockGemmPipelineVersion::v1, F8>, + DeviceGemm_Xdl_CShuffleV3< Row, Row, Row, F8, F8, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 256, 128, 64, 16, 4, 32, 32, 4, 2, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 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, F8>, + DeviceGemm_Xdl_CShuffleV3< Row, Row, Row, F8, F8, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 128, 128, 128, 16, 4, 32, 32, 2, 2, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, S<32, 8, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 16, 4, 0, 1, 1, S<1, 32, 1, 8>, 8, BlockGemmPipelineScheduler::Interwave, BlockGemmPipelineVersion::v1, F8>, + DeviceGemm_Xdl_CShuffleV3< Row, Row, Row, F8, F8, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 128, 64, 128, 16, 4, 32, 32, 2, 1, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, S<32, 8, 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, F8>, + DeviceGemm_Xdl_CShuffleV3< Row, Row, Row, F8, F8, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 64, 64, 128, 16, 4, 32, 32, 1, 1, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, S<32, 8, 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, F8> + // clang-format on + >; + +template +using device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_instances = std::tuple< + // clang-format off + //#########################| ALayout| BLayout| CLayout|AData| BData| CData| AccData| Cshuffle| A| B| C| GEMM| Block| MPer| NPer| KPer| AK1| BK1|MPer| NPer| MXdl| NXdl| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockLds| BBlockTransfer| BBlockTransfer| BBlockTransfer| BlockTransfer| BBlockTransfer| BBlockTransfer| BBlockLds| CShuffle| CShuffle| CBlockTransferClusterLengths| CBlockTransfer| Block-wiseGemm| Block-wiseGemm| + //#########################| | | | 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| + //#########################| | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | + + // Latency friendly + DeviceGemm_Xdl_CShuffleV3< Row, Row, Row, F8, F8, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 128, 32, 16, 128, 16, 4, 16, 16, 1, 1, S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, S<32, 4, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 4, 4, 0, 1, 1, S<1, 16, 1, 8>, 2, BlkGemmPipeSched, BlockGemmPipelineVersion::v1, F8>, + DeviceGemm_Xdl_CShuffleV3< Row, Row, Row, F8, F8, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 64, 16, 16, 128, 16, 4, 16, 16, 1, 1, S<8, 8, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, S<32, 2, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 8, 4, 0, 1, 1, S<1, 16, 1, 4>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v1, F8>, + DeviceGemm_Xdl_CShuffleV3< Row, Row, Row, F8, F8, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 64, 16, 16, 256, 16, 4, 16, 16, 1, 1, S<8, 8, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, S<64, 1, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 16, 4, 0, 1, 1, S<1, 16, 1, 4>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v1, F8>, + DeviceGemm_Xdl_CShuffleV3< Row, Row, Row, F8, F8, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 64, 16, 16, 512, 16, 8, 16, 16, 1, 1, S<8, 8, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, S<64, 1, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 16, 8, 0, 1, 1, S<1, 16, 1, 4>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v1, F8>, + DeviceGemm_Xdl_CShuffleV3< Row, Row, Row, F8, F8, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 128, 16, 32, 128, 16, 4, 16, 16, 1, 1, S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, S<32, 4, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 8, 4, 0, 1, 1, S<1, 16, 1, 8>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v1, F8>, + DeviceGemm_Xdl_CShuffleV3< Row, Row, Row, F8, F8, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 128, 16, 32, 256, 16, 4, 16, 16, 1, 1, S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, S<64, 2, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 16, 4, 0, 1, 1, S<1, 16, 1, 8>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v1, F8>, + DeviceGemm_Xdl_CShuffleV3< Row, Row, Row, F8, F8, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 128, 16, 32, 512, 16, 8, 16, 16, 1, 1, S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, S<64, 2, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 16, 8, 0, 1, 1, S<1, 16, 1, 8>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v1, F8>, + // Memory friendly + DeviceGemm_Xdl_CShuffleV3< Row, Row, Row, F8, F8, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 256, 16, 128, 16, 4, 16, 16, 4, 1, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, S<32, 4, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 4, 4, 0, 1, 1, S<1, 32, 1, 8>, 2, BlkGemmPipeSched, BlockGemmPipelineVersion::v2, F8>, + DeviceGemm_Xdl_CShuffleV3< Row, Row, Row, F8, F8, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 128, 128, 16, 128, 16, 4, 16, 16, 4, 1, S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, S<32, 4, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 4, 4, 0, 1, 1, S<1, 16, 1, 8>, 2, BlkGemmPipeSched, BlockGemmPipelineVersion::v2, F8>, + DeviceGemm_Xdl_CShuffleV3< Row, Row, Row, F8, F8, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 128, 64, 16, 128, 16, 4, 16, 16, 2, 1, S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, S<32, 4, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 4, 4, 0, 1, 1, S<1, 16, 1, 8>, 2, BlkGemmPipeSched, BlockGemmPipelineVersion::v2, F8>, + DeviceGemm_Xdl_CShuffleV3< Row, Row, Row, F8, F8, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 128, 32, 16, 128, 16, 4, 16, 16, 1, 1, S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, S<32, 4, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 4, 4, 0, 1, 1, S<1, 16, 1, 8>, 2, BlkGemmPipeSched, BlockGemmPipelineVersion::v2, F8>, + DeviceGemm_Xdl_CShuffleV3< Row, Row, Row, F8, F8, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 64, 16, 16, 128, 16, 4, 16, 16, 1, 1, S<8, 8, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, S<32, 2, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 8, 4, 0, 1, 1, S<1, 16, 1, 4>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v2, F8>, + DeviceGemm_Xdl_CShuffleV3< Row, Row, Row, F8, F8, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 64, 16, 16, 256, 16, 4, 16, 16, 1, 1, S<8, 8, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, S<64, 1, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 16, 4, 0, 1, 1, S<1, 16, 1, 4>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v2, F8>, + DeviceGemm_Xdl_CShuffleV3< Row, Row, Row, F8, F8, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 64, 16, 16, 512, 16, 8, 16, 16, 1, 1, S<8, 8, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, S<64, 1, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 16, 8, 0, 1, 1, S<1, 16, 1, 4>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v2, F8>, + DeviceGemm_Xdl_CShuffleV3< Row, Row, Row, F8, F8, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 128, 16, 32, 128, 16, 4, 16, 16, 1, 1, S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, S<32, 4, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 8, 4, 0, 1, 1, S<1, 16, 1, 8>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v2, F8>, + DeviceGemm_Xdl_CShuffleV3< Row, Row, Row, F8, F8, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 128, 16, 32, 256, 16, 4, 16, 16, 1, 1, S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, S<64, 2, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 16, 4, 0, 1, 1, S<1, 16, 1, 8>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v2, F8>, + DeviceGemm_Xdl_CShuffleV3< Row, Row, Row, F8, F8, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 128, 16, 32, 512, 16, 8, 16, 16, 1, 1, S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, S<64, 2, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 16, 8, 0, 1, 1, S<1, 16, 1, 8>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v2, F8>, + DeviceGemm_Xdl_CShuffleV3< Row, Row, Row, F8, F8, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 128, 16, 64, 128, 16, 4, 16, 16, 1, 2, S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, S<32, 4, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 16, 4, 0, 1, 1, S<1, 16, 1, 8>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v2, F8>, + DeviceGemm_Xdl_CShuffleV3< Row, Row, Row, F8, F8, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 128, 16, 128, 128, 16, 8, 16, 16, 1, 4, S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, S<16, 8, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 16, 8, 0, 1, 1, S<1, 16, 1, 8>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v2, F8>, + DeviceGemm_Xdl_CShuffleV3< Row, Row, Row, F8, F8, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 16, 256, 128, 8, 8, 16, 16, 1, 4, S<16, 16, 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, 16, 8, 0, 1, 1, S<1, 16, 1, 16>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v2, F8> + // clang-format on + >; +} // namespace instance +} // namespace device +} // namespace tensor_operation +} // namespace ck diff --git a/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_comp_default_instance.cpp b/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_comp_default_instance.cpp new file mode 100644 index 0000000000..96f171e066 --- /dev/null +++ b/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_comp_default_instance.cpp @@ -0,0 +1,23 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved. + +#include "device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn.hpp" + +namespace ck { +namespace tensor_operation { +namespace device { +namespace instance { + +void add_device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_comp_default_instances( + std::vector>>& + instances) +{ + add_device_operation_instances( + instances, device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_comp_instances{}); +} + +} // namespace instance +} // namespace device +} // namespace tensor_operation +} // namespace ck diff --git a/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_comp_kpadding_instance.cpp b/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_comp_kpadding_instance.cpp new file mode 100644 index 0000000000..4965fe51c6 --- /dev/null +++ b/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_comp_kpadding_instance.cpp @@ -0,0 +1,23 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved. + +#include "device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn.hpp" + +namespace ck { +namespace tensor_operation { +namespace device { +namespace instance { + +void add_device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_comp_kpadding_instances( + std::vector>>& + instances) +{ + add_device_operation_instances( + instances, device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_comp_instances{}); +} + +} // namespace instance +} // namespace device +} // namespace tensor_operation +} // namespace ck diff --git a/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_comp_mnpadding_instance.cpp b/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_comp_nkpadding_instance.cpp similarity index 59% rename from library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_comp_mnpadding_instance.cpp rename to library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_comp_nkpadding_instance.cpp index 4aae579ba3..d325c47d8a 100644 --- a/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_comp_mnpadding_instance.cpp +++ b/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_comp_nkpadding_instance.cpp @@ -1,20 +1,20 @@ // SPDX-License-Identifier: MIT // Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved. -#include "device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn.hpp" +#include "device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn.hpp" namespace ck { namespace tensor_operation { namespace device { namespace instance { -void add_device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_comp_mnpadding_instances( +void add_device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_comp_nkpadding_instances( std::vector>>& + DeviceGemmV2>>& instances) { add_device_operation_instances( - instances, device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_comp_instances{}); + instances, device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_comp_instances{}); } } // namespace instance diff --git a/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_mem_v1_mnkpadding_instance.cpp b/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_v1_default_instance.cpp similarity index 58% rename from library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_mem_v1_mnkpadding_instance.cpp rename to library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_v1_default_instance.cpp index 12eba27bd8..ff459e626e 100644 --- a/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_mem_v1_mnkpadding_instance.cpp +++ b/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_v1_default_instance.cpp @@ -1,21 +1,21 @@ // SPDX-License-Identifier: MIT // Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved. -#include "device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn.hpp" +#include "device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn.hpp" namespace ck { namespace tensor_operation { namespace device { namespace instance { -void add_device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_mem_v1_mnkpadding_instances( +void add_device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_v1_default_instances( std::vector>>& + DeviceGemmV2>>& instances) { add_device_operation_instances( instances, - device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_mem_instances{}); + device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_instances{}); } } // namespace instance diff --git a/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_comp_mnkpadding_instance.cpp b/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_v1_kpadding_instance.cpp similarity index 58% rename from library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_comp_mnkpadding_instance.cpp rename to library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_v1_kpadding_instance.cpp index 6feeaf6112..8b09a1946e 100644 --- a/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_comp_mnkpadding_instance.cpp +++ b/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_v1_kpadding_instance.cpp @@ -1,20 +1,21 @@ // SPDX-License-Identifier: MIT // Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved. -#include "device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn.hpp" +#include "device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn.hpp" namespace ck { namespace tensor_operation { namespace device { namespace instance { -void add_device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_comp_mnkpadding_instances( +void add_device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_v1_kpadding_instances( std::vector>>& + DeviceGemmV2>>& instances) { add_device_operation_instances( - instances, device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_comp_instances{}); + instances, + device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_instances{}); } } // namespace instance diff --git a/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_mem_v2_mnkpadding_instance.cpp b/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_v1_nkpadding_instance.cpp similarity index 58% rename from library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_mem_v2_mnkpadding_instance.cpp rename to library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_v1_nkpadding_instance.cpp index a4362fed5e..ec9d1b70a3 100644 --- a/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_mem_v2_mnkpadding_instance.cpp +++ b/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_v1_nkpadding_instance.cpp @@ -1,21 +1,21 @@ // SPDX-License-Identifier: MIT // Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved. -#include "device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn.hpp" +#include "device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn.hpp" namespace ck { namespace tensor_operation { namespace device { namespace instance { -void add_device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_mem_v2_mnkpadding_instances( +void add_device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_v1_nkpadding_instances( std::vector>>& + DeviceGemmV2>>& instances) { add_device_operation_instances( instances, - device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_mem_instances{}); + device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_instances{}); } } // namespace instance diff --git a/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_v2_default_instance.cpp b/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_v2_default_instance.cpp new file mode 100644 index 0000000000..383db226d8 --- /dev/null +++ b/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_v2_default_instance.cpp @@ -0,0 +1,24 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved. + +#include "device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn.hpp" + +namespace ck { +namespace tensor_operation { +namespace device { +namespace instance { + +void add_device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_v2_default_instances( + std::vector>>& + instances) +{ + add_device_operation_instances( + instances, + device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_instances{}); +} + +} // namespace instance +} // namespace device +} // namespace tensor_operation +} // namespace ck diff --git a/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_v2_kpadding_instance.cpp b/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_v2_kpadding_instance.cpp new file mode 100644 index 0000000000..9647747359 --- /dev/null +++ b/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_v2_kpadding_instance.cpp @@ -0,0 +1,24 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved. + +#include "device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn.hpp" + +namespace ck { +namespace tensor_operation { +namespace device { +namespace instance { + +void add_device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_v2_kpadding_instances( + std::vector>>& + instances) +{ + add_device_operation_instances( + instances, + device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_instances{}); +} + +} // namespace instance +} // namespace device +} // namespace tensor_operation +} // namespace ck diff --git a/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_v2_nkpadding_instance.cpp b/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_v2_nkpadding_instance.cpp new file mode 100644 index 0000000000..b2c621665c --- /dev/null +++ b/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_v2_nkpadding_instance.cpp @@ -0,0 +1,24 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved. + +#include "device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn.hpp" + +namespace ck { +namespace tensor_operation { +namespace device { +namespace instance { + +void add_device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_v2_nkpadding_instances( + std::vector>>& + instances) +{ + add_device_operation_instances( + instances, + device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_instances{}); +} + +} // namespace instance +} // namespace device +} // namespace tensor_operation +} // namespace ck diff --git a/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn.hpp b/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn.hpp index 46027a5756..b621cad942 100644 --- a/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn.hpp +++ b/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn.hpp @@ -45,8 +45,7 @@ using device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_comp_instances = std::tuple< DeviceGemm_Xdl_CShuffleV3< Row, Col, Row, F8, F8, 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>, 8, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v4, F8>, DeviceGemm_Xdl_CShuffleV3< Row, Col, Row, F8, F8, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 128, 128, 128, 16, 16, 32, 32, 2, 2, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, 1, 1, S<1, 32, 1, 8>, 8, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v4, F8>, DeviceGemm_Xdl_CShuffleV3< Row, Col, Row, F8, F8, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 128, 128, 64, 16, 16, 32, 32, 2, 2, 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>, 8, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v4, F8>, - DeviceGemm_Xdl_CShuffleV3< Row, Col, Row, F8, F8, 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>, 8, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v3, F8>, - DeviceGemm_Xdl_CShuffleV3< Row, Col, Row, F8, F8, 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>, 8, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v5, F8>, + DeviceGemm_Xdl_CShuffleV3< Row, Col, Row, F8, F8, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 256, 256, 128, 16, 16, 16, 16, 8, 8, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, 1, 2, S<1, 32, 1, 8>, 8, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v3, F8>, DeviceGemm_Xdl_CShuffleV3< Row, Col, Row, F8, F8, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 256, 256, 64, 16, 16, 16, 16, 8, 8, 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, 2, S<1, 32, 1, 8>, 8, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v3, F8>, DeviceGemm_Xdl_CShuffleV3< Row, Col, Row, F8, F8, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 224, 256, 128, 16, 16, 16, 16, 7, 8, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, 1, 2, S<1, 32, 1, 8>, 8, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v3, F8>, DeviceGemm_Xdl_CShuffleV3< Row, Col, Row, F8, F8, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 256, 224, 128, 16, 16, 16, 16, 8, 7, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, 2, 1, S<1, 64, 1, 4>, 8, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v3, F8>, @@ -72,7 +71,11 @@ using device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_mem_instances = std::tuple< // Latency friendly DeviceGemm_Xdl_CShuffleV3< Row, Col, Row, F8, F8, 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>, 2, BlkGemmPipeSched, BlockGemmPipelineVersion::v1, F8>, DeviceGemm_Xdl_CShuffleV3< Row, Col, Row, F8, F8, 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>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v1, F8>, + DeviceGemm_Xdl_CShuffleV3< Row, Col, Row, F8, F8, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 64, 16, 16, 256, 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>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v1, F8>, + DeviceGemm_Xdl_CShuffleV3< Row, Col, Row, F8, F8, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 64, 16, 16, 512, 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>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v1, F8>, DeviceGemm_Xdl_CShuffleV3< Row, Col, Row, F8, F8, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 128, 16, 32, 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>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v1, F8>, + DeviceGemm_Xdl_CShuffleV3< Row, Col, Row, F8, F8, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 128, 16, 32, 256, 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>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v1, F8>, + DeviceGemm_Xdl_CShuffleV3< Row, Col, Row, F8, F8, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 128, 16, 32, 512, 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>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v1, F8>, // Memory friendly DeviceGemm_Xdl_CShuffleV3< Row, Col, Row, F8, F8, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 256, 32, 128, 16, 16, 32, 32, 2, 1, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, 1, 1, S<1, 32, 1, 8>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v2, F8>, DeviceGemm_Xdl_CShuffleV3< Row, Col, Row, F8, F8, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 256, 16, 128, 16, 16, 16, 16, 4, 1, S<8, 32, 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, 32, 1, 8>, 2, BlkGemmPipeSched, BlockGemmPipelineVersion::v2, F8>, @@ -83,7 +86,11 @@ using device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_mem_instances = std::tuple< DeviceGemm_Xdl_CShuffleV3< Row, Col, Row, F8, F8, 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>, 2, BlkGemmPipeSched, BlockGemmPipelineVersion::v2, F8>, DeviceGemm_Xdl_CShuffleV3< Row, Col, Row, F8, F8, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 64, 16, 16, 64, 16, 16, 16, 16, 1, 1, S<4, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, S<4, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, 1, 1, S<1, 16, 1, 4>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v2, F8>, DeviceGemm_Xdl_CShuffleV3< Row, Col, Row, F8, F8, 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>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v2, F8>, + DeviceGemm_Xdl_CShuffleV3< Row, Col, Row, F8, F8, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 64, 16, 16, 256, 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>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v2, F8>, + DeviceGemm_Xdl_CShuffleV3< Row, Col, Row, F8, F8, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 64, 16, 16, 512, 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>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v2, F8>, DeviceGemm_Xdl_CShuffleV3< Row, Col, Row, F8, F8, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 128, 16, 32, 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>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v2, F8>, + DeviceGemm_Xdl_CShuffleV3< Row, Col, Row, F8, F8, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 128, 16, 32, 256, 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>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v2, F8>, + DeviceGemm_Xdl_CShuffleV3< Row, Col, Row, F8, F8, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 128, 16, 32, 512, 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>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v2, F8>, DeviceGemm_Xdl_CShuffleV3< Row, Col, Row, F8, F8, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 128, 16, 64, 128, 16, 16, 16, 16, 1, 2, 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>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v2, F8>, DeviceGemm_Xdl_CShuffleV3< Row, Col, Row, F8, F8, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 128, 32, 64, 128, 16, 16, 32, 32, 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>, 8, BlkGemmPipeSched, BlockGemmPipelineVersion::v2, F8>, DeviceGemm_Xdl_CShuffleV3< Row, Col, Row, F8, F8, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 128, 16, 128, 128, 16, 16, 16, 16, 1, 4, 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>, 4, BlkGemmPipeSched, BlockGemmPipelineVersion::v2, F8>, diff --git a/library/src/tensor_operation_instance/gpu/grouped_conv3d_bwd_data/CMakeLists.txt b/library/src/tensor_operation_instance/gpu/grouped_conv3d_bwd_data/CMakeLists.txt index 29fa8fa3c5..7bb7e71c54 100644 --- a/library/src/tensor_operation_instance/gpu/grouped_conv3d_bwd_data/CMakeLists.txt +++ b/library/src/tensor_operation_instance/gpu/grouped_conv3d_bwd_data/CMakeLists.txt @@ -15,7 +15,7 @@ set(GROUPED_CONV3D_BWD_DATA wmma/device_grouped_conv3d_bwd_data_wmma_gndhwc_gkzyxc_gndhwk_i8_1x1s1p0_instance.cpp wmma/device_grouped_conv3d_bwd_data_wmma_ndhwgc_gkzyxc_ndhwgk_i8_1x1s1p0_instance.cpp) -if((DTYPES MATCHES "fp8" AND DTYPES MATCHES "bf8" AND DTYPES MATCHES "fp16") OR NOT DEFINED DTYPES) +if((DTYPES MATCHES "fp8" AND DTYPES MATCHES "bf8" AND DTYPES MATCHES "fp16") OR (NOT DEFINED DTYPES AND GPU_TARGETS MATCHES "gfx94")) list(APPEND GROUPED_CONV3D_BWD_DATA xdl/device_grouped_conv3d_bwd_data_xdl_ndhwgc_gkzyxc_ndhwgk_input_f16_comp_bf8_f8_instance.cpp) endif() diff --git a/library/src/tensor_operation_instance/gpu/grouped_conv3d_bwd_weight/CMakeLists.txt b/library/src/tensor_operation_instance/gpu/grouped_conv3d_bwd_weight/CMakeLists.txt index 8e939c15a9..1a9c455220 100644 --- a/library/src/tensor_operation_instance/gpu/grouped_conv3d_bwd_weight/CMakeLists.txt +++ b/library/src/tensor_operation_instance/gpu/grouped_conv3d_bwd_weight/CMakeLists.txt @@ -30,7 +30,7 @@ list(APPEND GROUPED_CONV3D_BWD_WEIGHT wmma/device_grouped_conv3d_bwd_weight_wmma_gndhwc_gkzyxc_gndhwk_i8_instance.cpp wmma/device_grouped_conv3d_bwd_weight_wmma_ndhwgc_gkzyxc_ndhwgk_i8_instance.cpp) -if((DTYPES MATCHES "fp8" AND DTYPES MATCHES "bf8" AND DTYPES MATCHES "fp16") OR NOT DEFINED DTYPES) +if((DTYPES MATCHES "fp8" AND DTYPES MATCHES "bf8" AND DTYPES MATCHES "fp16") OR (NOT DEFINED DTYPES AND GPU_TARGETS MATCHES "gfx94")) list(APPEND GROUPED_CONV3D_BWD_WEIGHT xdl/device_grouped_conv3d_bwd_weight_xdl_ndhwgc_gkzyxc_ndhwgk_f16_comp_bf8_fp8_instance.cpp) endif() diff --git a/library/src/tensor_operation_instance/gpu/grouped_conv3d_bwd_weight_bilinear/CMakeLists.txt b/library/src/tensor_operation_instance/gpu/grouped_conv3d_bwd_weight_bilinear/CMakeLists.txt index 329e8e4c7f..5781f07080 100644 --- a/library/src/tensor_operation_instance/gpu/grouped_conv3d_bwd_weight_bilinear/CMakeLists.txt +++ b/library/src/tensor_operation_instance/gpu/grouped_conv3d_bwd_weight_bilinear/CMakeLists.txt @@ -4,7 +4,7 @@ set(GROUPED_CONV3D_BWD_WEIGHT_BILINEAR xdl/device_grouped_conv3d_bwd_weight_xdl_bilinear_ndhwgc_gkzyxc_ndhwgk_f32_instance.cpp xdl/device_grouped_conv3d_bwd_weight_xdl_bilinear_ndhwgc_gkzyxc_ndhwgk_bf16_instance.cpp) -if((DTYPES MATCHES "fp8" AND DTYPES MATCHES "bf8" AND DTYPES MATCHES "fp16") OR NOT DEFINED DTYPES) +if((DTYPES MATCHES "fp8" AND DTYPES MATCHES "bf8" AND DTYPES MATCHES "fp16") OR (NOT DEFINED DTYPES AND GPU_TARGETS MATCHES "gfx94")) list(APPEND GROUPED_CONV3D_BWD_WEIGHT_BILINEAR xdl/device_grouped_conv3d_bwd_weight_xdl_bilinear_ndhwgc_gkzyxc_ndhwgk_f16_comp_bf8_fp8_instance.cpp) endif() diff --git a/library/src/tensor_operation_instance/gpu/grouped_conv3d_bwd_weight_scale/CMakeLists.txt b/library/src/tensor_operation_instance/gpu/grouped_conv3d_bwd_weight_scale/CMakeLists.txt index 9a42d1ec3a..be54eb4adf 100644 --- a/library/src/tensor_operation_instance/gpu/grouped_conv3d_bwd_weight_scale/CMakeLists.txt +++ b/library/src/tensor_operation_instance/gpu/grouped_conv3d_bwd_weight_scale/CMakeLists.txt @@ -4,7 +4,7 @@ set(GROUPED_CONV3D_BWD_WEIGHT_SCALE xdl/device_grouped_conv3d_bwd_weight_xdl_scale_ndhwgc_gkzyxc_ndhwgk_f32_instance.cpp xdl/device_grouped_conv3d_bwd_weight_xdl_scale_ndhwgc_gkzyxc_ndhwgk_bf16_instance.cpp) -if((DTYPES MATCHES "fp8" AND DTYPES MATCHES "bf8" AND DTYPES MATCHES "fp16") OR NOT DEFINED DTYPES) +if((DTYPES MATCHES "fp8" AND DTYPES MATCHES "bf8" AND DTYPES MATCHES "fp16") OR (NOT DEFINED DTYPES AND GPU_TARGETS MATCHES "gfx94")) list(APPEND GROUPED_CONV3D_BWD_WEIGHT_SCALE xdl/device_grouped_conv3d_bwd_weight_xdl_scale_ndhwgc_gkzyxc_ndhwgk_f16_comp_bf8_fp8_instance.cpp) endif() diff --git a/library/src/tensor_operation_instance/gpu/grouped_conv3d_fwd/CMakeLists.txt b/library/src/tensor_operation_instance/gpu/grouped_conv3d_fwd/CMakeLists.txt index 6e0e94ba64..9bb6d807e6 100644 --- a/library/src/tensor_operation_instance/gpu/grouped_conv3d_fwd/CMakeLists.txt +++ b/library/src/tensor_operation_instance/gpu/grouped_conv3d_fwd/CMakeLists.txt @@ -43,22 +43,22 @@ set(GROUPED_CONV3D_FWD wmma/device_grouped_conv3d_fwd_wmma_ndhwgc_gkzyxc_ndhwgk_i8_oddc_instance.cpp ) -if((DTYPES MATCHES "fp8" AND DTYPES MATCHES "fp16") OR NOT DEFINED DTYPES) +if((DTYPES MATCHES "fp8" AND DTYPES MATCHES "fp16") OR (NOT DEFINED DTYPES AND GPU_TARGETS MATCHES "gfx94")) list(APPEND GROUPED_CONV3D_FWD xdl/device_grouped_conv3d_fwd_xdl_ndhwgc_gkzyxc_ndhwgk_f16_comp_fp8_instance.cpp) endif() -if(DTYPES MATCHES "fp8" OR NOT DEFINED DTYPES) +if((DTYPES MATCHES "fp8") OR (NOT DEFINED DTYPES AND GPU_TARGETS MATCHES "gfx94")) list(APPEND GROUPED_CONV3D_FWD xdl/device_grouped_conv3d_fwd_xdl_ndhwgc_gkzyxc_ndhwgk_fp8_instance.cpp) endif() -if(DTYPES MATCHES "bf8" OR NOT DEFINED DTYPES) +if((DTYPES MATCHES "bf8") OR (NOT DEFINED DTYPES AND GPU_TARGETS MATCHES "gfx94")) list(APPEND GROUPED_CONV3D_FWD xdl/device_grouped_conv3d_fwd_xdl_ndhwgc_gkzyxc_ndhwgk_bf8_instance.cpp) endif() -if((DTYPES MATCHES "fp8" AND DTYPES MATCHES "bf8") OR NOT DEFINED DTYPES) +if((DTYPES MATCHES "fp8" AND DTYPES MATCHES "bf8") OR (NOT DEFINED DTYPES AND GPU_TARGETS MATCHES "gfx94")) list(APPEND GROUPED_CONV3D_FWD xdl/device_grouped_conv3d_fwd_xdl_ndhwgc_gkzyxc_ndhwgk_fp8_bf8_instance.cpp) list(APPEND GROUPED_CONV3D_FWD diff --git a/profiler/include/profiler/profile_gemm_multiply_multiply_impl.hpp b/profiler/include/profiler/profile_gemm_multiply_multiply_impl.hpp index 022399a9c0..7dd7b041ed 100644 --- a/profiler/include/profiler/profile_gemm_multiply_multiply_impl.hpp +++ b/profiler/include/profiler/profile_gemm_multiply_multiply_impl.hpp @@ -48,6 +48,7 @@ bool profile_gemm_multiply_multiply_impl(int do_verification, int StrideD0, int StrideD1, int StrideE, + int KBatch, int n_warmup, int n_iter, uint64_t rotating = 0) @@ -129,17 +130,17 @@ bool profile_gemm_multiply_multiply_impl(int do_verification, d1_device_buf.ToDevice(d1_m_n.mData.data()); using DeviceOp = - ck::tensor_operation::device::DeviceGemmMultipleD, - ELayout, - ADataType, - BDataType, - ck::Tuple, - EDataType, - AElementOp, - BElementOp, - CElementOp>; + ck::tensor_operation::device::DeviceGemmMultipleDSplitK, + ELayout, + ADataType, + BDataType, + ck::Tuple, + EDataType, + AElementOp, + BElementOp, + CElementOp>; // get device op instances const auto op_ptrs = ck::tensor_operation::device::instance::DeviceOperationInstanceFactory< @@ -182,103 +183,127 @@ bool profile_gemm_multiply_multiply_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 GEMM instances for(auto& op_ptr : op_ptrs) { - auto argument_ptr = - op_ptr->MakeArgumentPointer(static_cast(a_device_buf.GetDeviceBuffer()), - static_cast(b_device_buf.GetDeviceBuffer()), - std::array{d0_device_buf.GetDeviceBuffer(), - d1_device_buf.GetDeviceBuffer()}, - static_cast(c_device_buf.GetDeviceBuffer()), - M, - N, - K, - StrideA, - StrideB, - std::array{StrideD0, StrideD1}, - StrideE, - a_element_op, - b_element_op, - c_element_op); + std::vector kbatch_list = {1, 2, 4, 8, 16, 19, 32, 38}; - auto invoker_ptr = op_ptr->MakeInvokerPointer(); - - if(op_ptr->IsSupportedArgument(argument_ptr.get())) + if(KBatch > 0) { + kbatch_list = {KBatch}; + } - // re-init C to zero before profiling next kernel - c_device_buf.SetZero(); + for(std::size_t i = 0; i < kbatch_list.size(); i++) + { + auto kbatch_curr = kbatch_list[i]; - invoker_ptr->Run(argument_ptr.get(), StreamConfig{nullptr, false, 0, n_warmup, n_iter}); + auto argument_ptr = op_ptr->MakeArgumentPointer( + static_cast(a_device_buf.GetDeviceBuffer()), + static_cast(b_device_buf.GetDeviceBuffer()), + std::array{d0_device_buf.GetDeviceBuffer(), + d1_device_buf.GetDeviceBuffer()}, + static_cast(c_device_buf.GetDeviceBuffer()), + M, + N, + K, + StrideA, + StrideB, + std::array{StrideD0, StrideD1}, + StrideE, + kbatch_curr, + a_element_op, + b_element_op, + c_element_op); - if(do_verification) + auto invoker_ptr = op_ptr->MakeInvokerPointer(); + + if(op_ptr->IsSupportedArgument(argument_ptr.get())) { - c_device_buf.FromDevice(e_m_n_device_result.mData.data()); - pass = pass & ck::utils::check_err(e_m_n_device_result, e_m_n_host_result); + // re-init C to zero before profiling next kernel + c_device_buf.SetZero(); - if(do_log) + invoker_ptr->Run(argument_ptr.get(), + StreamConfig{nullptr, false, 0, n_warmup, n_iter}); + + if(do_verification) { - LogRangeAsType(std::cout << "a : ", a_m_k.mData, ",") << std::endl; - LogRangeAsType(std::cout << "b: ", b_k_n.mData, ",") << std::endl; - LogRangeAsType(std::cout << "c_host : ", e_m_n_host_result.mData, ",") - << std::endl; - LogRangeAsType(std::cout << "c_device: ", e_m_n_device_result.mData, ",") - << std::endl; + c_device_buf.FromDevice(e_m_n_device_result.mData.data()); + + pass = pass & ck::utils::check_err(e_m_n_device_result, e_m_n_host_result); + + if(do_log) + { + LogRangeAsType(std::cout << "a : ", a_m_k.mData, ",") << std::endl; + LogRangeAsType(std::cout << "b: ", b_k_n.mData, ",") << std::endl; + LogRangeAsType( + std::cout << "c_host : ", e_m_n_host_result.mData, ",") + << std::endl; + LogRangeAsType( + std::cout << "c_device: ", e_m_n_device_result.mData, ",") + << std::endl; + } } - } - std::string op_name = op_ptr->GetTypeString(); + std::string op_name = op_ptr->GetTypeString(); - float ave_time = invoker_ptr->Run( - argument_ptr.get(), - StreamConfig{ - nullptr, time_kernel, 0, n_warmup, n_iter, rotating_count > 1, rotating_count}); + float ave_time = invoker_ptr->Run(argument_ptr.get(), + StreamConfig{nullptr, + time_kernel, + 0, + n_warmup, + n_iter, + rotating_count > 1, + rotating_count}); - std::size_t flop = std::size_t(2) * M * N * K; + std::size_t flop = std::size_t(2) * M * N * K; - std::size_t num_btype = - sizeof(ADataType) * M * K + sizeof(BDataType) * K * N + sizeof(EDataType) * M * N; + std::size_t num_btype = sizeof(ADataType) * M * K + sizeof(BDataType) * K * N + + sizeof(EDataType) * M * N; - float tflops = static_cast(flop) / 1.E9 / ave_time; + float tflops = static_cast(flop) / 1.E9 / ave_time; - float gb_per_sec = num_btype / 1.E6 / ave_time; + float gb_per_sec = num_btype / 1.E6 / ave_time; - std::cout << "Perf: " << std::setw(10) << ave_time << " ms, " << tflops << " TFlops, " - << gb_per_sec << " GB/s, " << op_name << std::endl; + std::cout << "Perf: " << std::setw(10) << ave_time << " ms, " << tflops + << " TFlops, " << gb_per_sec << " GB/s, " << op_name << ", KBatch " + << kbatch_curr << std::endl; #if defined CK_ENABLE_FP8 - // set softer tolerances for fp8 - if constexpr(is_same_v || is_same_v || - is_same_v) - { - std::string msg = "Error: Incorrect results!"; - double rtol = 1e-1; - double atol = 1e-1; - pass = pass & ck::utils::check_err( - e_m_n_device_result, e_m_n_host_result, msg, rtol, atol); + // set softer tolerances for fp8 + if constexpr(is_same_v || is_same_v || + is_same_v) + { + std::string msg = "Error: Incorrect results!"; + double rtol = 1e-1; + double atol = 1e-1; + pass = pass & ck::utils::check_err( + e_m_n_device_result, e_m_n_host_result, msg, rtol, atol); + } + else + { +#endif + pass = pass & ck::utils::check_err(e_m_n_device_result, e_m_n_host_result); +#if defined CK_ENABLE_FP8 + } +#endif + + 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; + } } else { -#endif - pass = pass & ck::utils::check_err(e_m_n_device_result, e_m_n_host_result); -#if defined CK_ENABLE_FP8 + std::cout << op_ptr->GetTypeString() << " does not support this problem" + << std::endl; } -#endif - - if(tflops > best_tflops) - { - best_op_name = op_name; - best_tflops = tflops; - best_ave_time = ave_time; - best_gb_per_sec = gb_per_sec; - } - } - else - { - std::cout << op_ptr->GetTypeString() << " does not support this problem" << std::endl; } } @@ -318,9 +343,9 @@ bool profile_gemm_multiply_multiply_impl(int do_verification, } std::cout << " M = " << M << " N = " << N << " K = " << K << " StrideA = " << StrideA - << " StrideB = " << StrideB << " StrideE = " << StrideE << " : " << best_ave_time - << " ms, " << best_tflops << " TFlops, " << best_gb_per_sec << " GB/s, " - << best_op_name << std::endl; + << " StrideB = " << StrideB << " StrideE = " << StrideE << " 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/include/profiler/profile_gemm_universal_impl.hpp b/profiler/include/profiler/profile_gemm_universal_impl.hpp index b6dac96989..f6e1f12e2a 100644 --- a/profiler/include/profiler/profile_gemm_universal_impl.hpp +++ b/profiler/include/profiler/profile_gemm_universal_impl.hpp @@ -152,7 +152,7 @@ bool profile_gemm_universal_impl(int do_verification, // profile device GEMM instances for(auto& op_ptr : op_ptrs) { - std::vector kbatch_list = {1, 2, 4, 8, 12, 16, 19, 20, 32, 38}; + std::vector kbatch_list = {1, 2, 4, 8, 16, 19, 32, 38}; if(KBatch > 0) { @@ -249,7 +249,7 @@ bool profile_gemm_universal_impl(int do_verification, << " TFlops, " << gb_per_sec << " GB/s, " << op_name << ", KBatch " << kbatch_curr << std::endl; - if(tflops > best_tflops) + if(tflops > best_tflops && ave_time > 1e-10) { best_op_name = op_name; best_tflops = tflops; diff --git a/profiler/src/profile_gemm_multiply_multiply.cpp b/profiler/src/profile_gemm_multiply_multiply.cpp index 7fbc318fe5..b7e80ed798 100644 --- a/profiler/src/profile_gemm_multiply_multiply.cpp +++ b/profiler/src/profile_gemm_multiply_multiply.cpp @@ -34,7 +34,7 @@ enum struct GemmDataType int profile_gemm_multiply_multiply(int argc, char* argv[]) { - if(argc != 16 && argc != 19) + if(argc != 16 && argc != 20) { printf("arg1: tensor operation (" OP_NAME ": " OP_DESC ")\n"); printf("arg2: data type (0: fp32; 1: fp16; 2: bf16; 3: int8; 4: f8@f16; 5: f16@f8; 6: " @@ -50,9 +50,10 @@ int profile_gemm_multiply_multiply(int argc, char* argv[]) printf("arg7: time kernel (0=no, 1=yes)\n"); printf("arg8 to 15: M, N, K, StrideA, StrideB, StrideD0, StrideD1, StrideE\n"); printf("optional:\n"); - printf("arg16: number of warm-up cycles (default 1)\n"); - printf("arg17: number of iterations (default 10)\n"); - printf("arg18: memory for rotating buffer (default 0, size in MB)\n"); + printf("arg16: number of kbatch (default 1)\n"); + printf("arg17: number of warm-up cycles (default 1)\n"); + printf("arg18: number of iterations (default 10)\n"); + printf("arg19: memory for rotating buffer (default 0, size in MB)\n"); exit(1); } @@ -76,11 +77,13 @@ int profile_gemm_multiply_multiply(int argc, char* argv[]) int n_warmup = 1; int n_iter = 10; uint64_t rotating = 0; - if(argc == 19) + int KBatch = 1; + if(argc == 20) { - n_warmup = std::stoi(argv[16]); - n_iter = std::stoi(argv[17]); - rotating = std::stoull(argv[18]) * 1024 * 1024; + KBatch = std::stoi(argv[16]); + n_warmup = std::stoi(argv[17]); + n_iter = std::stoi(argv[18]); + rotating = std::stoull(argv[19]) * 1024 * 1024; } using F32 = float; @@ -146,6 +149,7 @@ int profile_gemm_multiply_multiply(int argc, char* argv[]) (StrideD0 < 0) ? DefaultStrideD0 : StrideD0, (StrideD1 < 0) ? DefaultStrideD1 : StrideD1, (StrideE < 0) ? DefaultStrideE : StrideE, + KBatch, n_warmup, n_iter, rotating); diff --git a/profiler/src/profile_gemm_universal.cpp b/profiler/src/profile_gemm_universal.cpp index ca220ddc47..cd61511926 100644 --- a/profiler/src/profile_gemm_universal.cpp +++ b/profiler/src/profile_gemm_universal.cpp @@ -171,6 +171,10 @@ int profile_gemm_universal(int argc, char* argv[]) { return profile(BF16{}, BF16{}, BF16{}, F32{}, BF16{}, Row{}, Col{}, Row{}); } + else if(data_type == GemmDataType::F8_F8_BF16 && layout == GemmMatrixLayout::MK_KN_MN) + { + return profile(F8{}, F8{}, F8{}, F32{}, BF16{}, Row{}, Row{}, Row{}); + } else if(data_type == GemmDataType::F8_F8_BF16 && layout == GemmMatrixLayout::MK_NK_MN) { return profile(F8{}, F8{}, F8{}, F32{}, BF16{}, Row{}, Col{}, Row{}); diff --git a/profiler/src/profile_grouped_gemm_fixed_nk.cpp b/profiler/src/profile_grouped_gemm_fixed_nk.cpp index 3d280c2f43..de90a33ef4 100644 --- a/profiler/src/profile_grouped_gemm_fixed_nk.cpp +++ b/profiler/src/profile_grouped_gemm_fixed_nk.cpp @@ -85,9 +85,11 @@ int profile_grouped_gemm_fixed_nk(int argc, char* argv[]) const auto StrideCs = argToIntArray(argv[13]); const int kbatch = argc == 15 ? std::stoi(argv[14]) : 1; - using F32 = float; - using F16 = ck::half_t; - using F8 = ck::f8_t; + using F32 = float; + using F16 = ck::half_t; +#if defined(CK_ENABLE_FP8) + using F8 = ck::f8_t; +#endif using BF16 = ck::bhalf_t; using I8 = int8_t; diff --git a/test/gemm_universal/test_gemm_universal_xdl.cpp b/test/gemm_universal/test_gemm_universal_xdl.cpp index 1644cd4f7b..63c1a48497 100644 --- a/test/gemm_universal/test_gemm_universal_xdl.cpp +++ b/test/gemm_universal/test_gemm_universal_xdl.cpp @@ -44,17 +44,22 @@ class TestGemmUniversal_MK_NK using KernelTypes_MK_KN = ::testing::Types< // ADataType, BDataType, ComputeDataType, CDataType std::tuple< F16, F16, F16, F16>, +#if (defined CK_ENABLE_FP8) std::tuple< F16, F8, F16, F16>, std::tuple< F8, F16, F16, F16>, + std::tuple< F8, F8, F8, BF16>, +#endif std::tuple< BF16, BF16, BF16, BF16> >; using KernelTypes_MK_NK = ::testing::Types< // ADataType, BDataType, ComputeDataType, CDataType std::tuple< F16, F16, F16, F16>, +#if (defined CK_ENABLE_FP8) std::tuple< F16, F8, F16, F16>, std::tuple< F8, F16, F16, F16>, - std::tuple< BF16, BF16, BF16, BF16>, - std::tuple< F8, F8, F8, BF16> + std::tuple< F8, F8, F8, BF16>, +#endif + std::tuple< BF16, BF16, BF16, BF16> >; // clang-format on From 49769ec8891720e0a6b7316325f3dd966663c41b Mon Sep 17 00:00:00 2001 From: trixirt Date: Wed, 14 Aug 2024 20:43:10 -0700 Subject: [PATCH 04/20] Check compiler flags before using (#1403) * Check compiler flags before using The user's compiler may not support these flags, so check. Resolves failures on Fedora. Signed-off-by: Tom Rix * fix syntax CMakeLists.txt Fix syntax in the check_cxx_compiler_flag. --------- Signed-off-by: Tom Rix Co-authored-by: Tom Rix Co-authored-by: Illia Silin <98187287+illsilin@users.noreply.github.com> --- CMakeLists.txt | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 96c37ac943..2039948a12 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -191,12 +191,18 @@ endif() configure_file(include/ck/config.h.in ${CMAKE_CURRENT_BINARY_DIR}/include/ck/config.h) if(NOT WIN32 AND ${hip_VERSION_FLAT} GREATER 500723302) - message("Adding the fno-offload-uniform-block compiler flag") - add_compile_options(-fno-offload-uniform-block) + check_cxx_compiler_flag("-fno-offload-uniform-block" HAS_NO_OFFLOAD_UNIFORM_BLOCK) + if(HAS_NO_OFFLOAD_UNIFORM_BLOCK) + message("Adding the fno-offload-uniform-block compiler flag") + add_compile_options(-fno-offload-uniform-block) + endif() endif() if(NOT WIN32 AND ${hip_VERSION_FLAT} GREATER 600140090) - message("Adding the enable-post-misched=0 compiler flag") - add_compile_options("SHELL: -mllvm -enable-post-misched=0") + check_cxx_compiler_flag("-mllvm -enable-post-misched=0" HAS_ENABLE_POST_MISCHED) + if(HAS_ENABLE_POST_MISCHED) + message("Adding the enable-post-misched=0 compiler flag") + add_compile_options("SHELL: -mllvm -enable-post-misched=0") + endif() endif() set(check-coerce) check_cxx_compiler_flag(" -mllvm -amdgpu-coerce-illegal-types=1" check-coerce) From 76bd0af6af17f4f56af4a8fe0896b1be1d36ab91 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Thu, 15 Aug 2024 13:59:40 -0700 Subject: [PATCH 05/20] Bump rocm-docs-core from 1.6.2 to 1.7.0 in /docs/sphinx (#1467) Bumps [rocm-docs-core](https://github.com/ROCm/rocm-docs-core) from 1.6.2 to 1.7.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.6.2...v1.7.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 49ff317876..8be30305a6 100644 --- a/docs/sphinx/requirements.in +++ b/docs/sphinx/requirements.in @@ -1,2 +1,2 @@ -rocm-docs-core==1.6.2 +rocm-docs-core==1.7.0 sphinxcontrib-bibtex==2.6.2 diff --git a/docs/sphinx/requirements.txt b/docs/sphinx/requirements.txt index bc7d0f689e..4cc4d30f79 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.6.2 +rocm-docs-core==1.7.0 # via -r requirements.in six==1.16.0 # via pybtex From 2581727d2a2256c37afcaedf1bc53a2e023f1f51 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Bart=C5=82omiej=20Kocot?= Date: Fri, 16 Aug 2024 16:48:30 +0200 Subject: [PATCH 06/20] Add performance and large tensor tests for grouped conv (#1456) * Add performance and large tensor tests for grouped conv * Resize tests * Resize tests * update the python script to parse the grouped_conv results * Remove int8 tests * change bwd wei layout --------- Co-authored-by: illsilin --- Jenkinsfile | 43 +++++- script/process_perf_data.py | 18 ++- script/process_qa_data.sh | 5 +- ...ta.sh => profile_grouped_conv_bwd_data.sh} | 0 ....sh => profile_grouped_conv_bwd_weight.sh} | 43 +++--- script/profile_grouped_conv_fwd.sh | 39 ++++++ script/run_full_performance_tests.sh | 34 +++-- script/run_performance_tests.sh | 15 +++ test/grouped_convnd_fwd/CMakeLists.txt | 6 + .../test_grouped_convnd_fwd.cpp | 65 +++------ ...est_grouped_convnd_fwd_large_cases_xdl.cpp | 127 ++++++++++++++++++ 11 files changed, 305 insertions(+), 90 deletions(-) rename script/{profile_conv_bwd_data.sh => profile_grouped_conv_bwd_data.sh} (100%) rename script/{profile_conv_fwd.sh => profile_grouped_conv_bwd_weight.sh} (90%) create mode 100755 script/profile_grouped_conv_fwd.sh create mode 100644 test/grouped_convnd_fwd/test_grouped_convnd_fwd_large_cases_xdl.cpp diff --git a/Jenkinsfile b/Jenkinsfile index 3fccb2881b..e9ea3d1c08 100644 --- a/Jenkinsfile +++ b/Jenkinsfile @@ -426,8 +426,9 @@ def runCKProfiler(Map conf=[:]){ archiveArtifacts "perf_resnet50_N4.log" archiveArtifacts "perf_batched_gemm.log" archiveArtifacts "perf_grouped_gemm.log" - archiveArtifacts "perf_conv_fwd.log" - archiveArtifacts "perf_conv_bwd_data.log" + archiveArtifacts "perf_grouped_conv_fwd.log" + archiveArtifacts "perf_grouped_conv_bwd_data.log" + archiveArtifacts "perf_grouped_conv_bwd_weight.log" archiveArtifacts "perf_gemm_bilinear.log" archiveArtifacts "perf_reduction.log" archiveArtifacts "perf_splitK_gemm.log" @@ -439,8 +440,9 @@ def runCKProfiler(Map conf=[:]){ stash name: "perf_resnet50_N4.log" stash name: "perf_batched_gemm.log" stash name: "perf_grouped_gemm.log" - stash name: "perf_conv_fwd.log" - stash name: "perf_conv_bwd_data.log" + stash name: "perf_grouped_conv_fwd.log" + stash name: "perf_grouped_conv_bwd_data.log" + stash name: "perf_grouped_conv_bwd_weight.log" stash name: "perf_gemm_bilinear.log" stash name: "perf_reduction.log" stash name: "perf_splitK_gemm.log" @@ -648,8 +650,9 @@ def process_results(Map conf=[:]){ unstash "perf_resnet50_N4.log" unstash "perf_batched_gemm.log" unstash "perf_grouped_gemm.log" - unstash "perf_conv_fwd.log" - unstash "perf_conv_bwd_data.log" + unstash "perf_grouped_conv_fwd.log" + unstash "perf_grouped_conv_bwd_data.log" + unstash "perf_grouped_conv_bwd_weight.log" unstash "perf_gemm_bilinear.log" unstash "perf_reduction.log" unstash "perf_splitK_gemm.log" @@ -746,6 +749,10 @@ pipeline { name: "RUN_PERFORMANCE_TESTS", defaultValue: true, description: "Run the performance tests (default: ON)") + booleanParam( + name: "RUN_GROUPED_CONV_LARGE_CASES_TESTS", + defaultValue: false, + description: "Run the grouped conv large cases tests (default: OFF)") booleanParam( name: "RUN_CK_TILE_TESTS", defaultValue: false, @@ -837,6 +844,30 @@ pipeline { } } } + stage("Run Grouped Conv Large Case Tests") + { + parallel + { + stage("Run Grouped Conv Large Case Tests on gfx90a") + { + when { + beforeAgent true + expression { params.RUN_GROUPED_CONV_LARGE_CASES_TESTS.toBoolean() } + } + agent{ label rocmnode("gfx90a")} + environment{ + setup_args = "NO_CK_BUILD" + execute_args = """ ../script/cmake-ck-dev.sh ../ gfx90a && \ + make -j64 test_grouped_convnd_fwd_large_cases_xdl && \ + ./bin/test_grouped_convnd_fwd_large_cases_xdl""" + } + steps{ + buildHipClangJobAndReboot(setup_args:setup_args, no_reboot:true, build_type: 'Release', execute_cmd: execute_args) + cleanWs() + } + } + } + } stage("Run CK_TILE Tests") { parallel diff --git a/script/process_perf_data.py b/script/process_perf_data.py index c6cb6e05c7..b82a7c2891 100644 --- a/script/process_perf_data.py +++ b/script/process_perf_data.py @@ -122,7 +122,7 @@ def parse_logfile(logfile): #sorted_kernels = [x for _,x in sorted(zip(tests,kernels))] test_list=list(range(1,len(tests)+1)) #parse conv_fwd and conv_bwd performance tests: - elif 'conv_fwd' in logfile or 'conv_bwd_data' in logfile: + elif 'conv_fwd' in logfile or 'conv_bwd' in logfile: for line in open(logfile): if 'tflops:' in line: lst=line.split() @@ -274,14 +274,26 @@ def main(): for i in range(1,len(results)+1): testlist.append("Test%i"%i) table_name="ck_grouped_gemm_tflops" - if 'conv_fwd' in filename: + if 'perf_conv_fwd' in filename: for i in range(1,len(results)+1): testlist.append("Test%i"%i) table_name="ck_conv_fwd_tflops" - if 'conv_bwd_data' in filename: + if 'perf_conv_bwd_data' in filename: for i in range(1,len(results)+1): testlist.append("Test%i"%i) table_name="ck_conv_bwd_data_tflops" + if 'grouped_conv_fwd' in filename: + for i in range(1,len(results)+1): + testlist.append("Test%i"%i) + table_name="ck_grouped_conv_fwd_tflops" + if 'grouped_conv_bwd_data' in filename: + for i in range(1,len(results)+1): + testlist.append("Test%i"%i) + table_name="ck_grouped_conv_bwd_data_tflops" + if 'grouped_conv_bwd_weight' in filename: + for i in range(1,len(results)+1): + testlist.append("Test%i"%i) + table_name="ck_grouped_conv_bwd_weight_tflops" if 'gemm_bilinear' in filename: for i in range(1,len(results)+1): testlist.append("Test%i"%i) diff --git a/script/process_qa_data.sh b/script/process_qa_data.sh index bf16f05cd0..d6083d2fc7 100755 --- a/script/process_qa_data.sh +++ b/script/process_qa_data.sh @@ -15,8 +15,9 @@ python3 process_perf_data.py perf_resnet50_N256.log python3 process_perf_data.py perf_resnet50_N4.log python3 process_perf_data.py perf_batched_gemm.log python3 process_perf_data.py perf_grouped_gemm.log -python3 process_perf_data.py perf_conv_fwd.log -python3 process_perf_data.py perf_conv_bwd_data.log +python3 process_perf_data.py perf_grouped_conv_fwd.log +python3 process_perf_data.py perf_grouped_conv_bwd_data.log +python3 process_perf_data.py perf_grouped_conv_bwd_weight.log python3 process_perf_data.py perf_gemm_bilinear.log python3 process_perf_data.py perf_reduction.log python3 process_perf_data.py perf_splitK_gemm.log diff --git a/script/profile_conv_bwd_data.sh b/script/profile_grouped_conv_bwd_data.sh similarity index 100% rename from script/profile_conv_bwd_data.sh rename to script/profile_grouped_conv_bwd_data.sh diff --git a/script/profile_conv_fwd.sh b/script/profile_grouped_conv_bwd_weight.sh similarity index 90% rename from script/profile_conv_fwd.sh rename to script/profile_grouped_conv_bwd_weight.sh index a1d2f450c9..e3652202d4 100755 --- a/script/profile_conv_fwd.sh +++ b/script/profile_grouped_conv_bwd_weight.sh @@ -12,27 +12,28 @@ INIT=$5 LOG=$6 TIME=$7 - N=$8 +N=$8 +SplitK=$9 # Resnet50 ######## op datatype layout verify init log time conv_dim G__ N__ K___ C___ Y X Hi__ Wi__ Strides Dilations LeftPads RightPads - $DRIVER $OP $DATATYPE $LAYOUT $VERIFY $INIT $LOG $TIME 2 1 $N 256 1024 1 1 14 14 1 1 1 1 0 0 0 0 - $DRIVER $OP $DATATYPE $LAYOUT $VERIFY $INIT $LOG $TIME 2 1 $N 512 1024 1 1 14 14 1 1 1 1 0 0 0 0 - $DRIVER $OP $DATATYPE $LAYOUT $VERIFY $INIT $LOG $TIME 2 1 $N 128 128 3 3 28 28 1 1 1 1 1 1 1 1 - $DRIVER $OP $DATATYPE $LAYOUT $VERIFY $INIT $LOG $TIME 2 1 $N 512 128 1 1 28 28 1 1 1 1 0 0 0 0 - $DRIVER $OP $DATATYPE $LAYOUT $VERIFY $INIT $LOG $TIME 2 1 $N 128 128 3 3 56 56 2 2 1 1 1 1 1 1 - $DRIVER $OP $DATATYPE $LAYOUT $VERIFY $INIT $LOG $TIME 2 1 $N 512 2048 1 1 7 7 1 1 1 1 0 0 0 0 - $DRIVER $OP $DATATYPE $LAYOUT $VERIFY $INIT $LOG $TIME 2 1 $N 1024 256 1 1 14 14 1 1 1 1 0 0 0 0 - $DRIVER $OP $DATATYPE $LAYOUT $VERIFY $INIT $LOG $TIME 2 1 $N 256 256 3 3 14 14 1 1 1 1 1 1 1 1 - $DRIVER $OP $DATATYPE $LAYOUT $VERIFY $INIT $LOG $TIME 2 1 $N 256 256 3 3 28 28 2 2 1 1 1 1 1 1 - $DRIVER $OP $DATATYPE $LAYOUT $VERIFY $INIT $LOG $TIME 2 1 $N 128 256 1 1 56 56 1 1 1 1 0 0 0 0 - $DRIVER $OP $DATATYPE $LAYOUT $VERIFY $INIT $LOG $TIME 2 1 $N 64 256 1 1 56 56 1 1 1 1 0 0 0 0 - $DRIVER $OP $DATATYPE $LAYOUT $VERIFY $INIT $LOG $TIME 2 1 $N 512 512 3 3 14 14 2 2 1 1 1 1 1 1 - $DRIVER $OP $DATATYPE $LAYOUT $VERIFY $INIT $LOG $TIME 2 1 $N 128 512 1 1 28 28 1 1 1 1 0 0 0 0 - $DRIVER $OP $DATATYPE $LAYOUT $VERIFY $INIT $LOG $TIME 2 1 $N 256 512 1 1 28 28 1 1 1 1 0 0 0 0 - $DRIVER $OP $DATATYPE $LAYOUT $VERIFY $INIT $LOG $TIME 2 1 $N 2048 512 1 1 7 7 1 1 1 1 0 0 0 0 - $DRIVER $OP $DATATYPE $LAYOUT $VERIFY $INIT $LOG $TIME 2 1 $N 512 512 3 3 7 7 1 1 1 1 1 1 1 1 - $DRIVER $OP $DATATYPE $LAYOUT $VERIFY $INIT $LOG $TIME 2 1 $N 256 64 1 1 56 56 1 1 1 1 0 0 0 0 - $DRIVER $OP $DATATYPE $LAYOUT $VERIFY $INIT $LOG $TIME 2 1 $N 64 64 1 1 56 56 1 1 1 1 0 0 0 0 - $DRIVER $OP $DATATYPE $LAYOUT $VERIFY $INIT $LOG $TIME 2 1 $N 64 64 3 3 56 56 1 1 1 1 1 1 1 1 - $DRIVER $OP $DATATYPE $LAYOUT $VERIFY $INIT $LOG $TIME 2 1 $N 64 3 7 7 224 224 2 2 1 1 3 3 3 3 + $DRIVER $OP $DATATYPE $LAYOUT $VERIFY $INIT $LOG $TIME 2 1 $N 256 1024 1 1 14 14 1 1 1 1 0 0 0 0 $SplitK + $DRIVER $OP $DATATYPE $LAYOUT $VERIFY $INIT $LOG $TIME 2 1 $N 512 1024 1 1 14 14 1 1 1 1 0 0 0 0 $SplitK + $DRIVER $OP $DATATYPE $LAYOUT $VERIFY $INIT $LOG $TIME 2 1 $N 128 128 3 3 28 28 1 1 1 1 1 1 1 1 $SplitK + $DRIVER $OP $DATATYPE $LAYOUT $VERIFY $INIT $LOG $TIME 2 1 $N 512 128 1 1 28 28 1 1 1 1 0 0 0 0 $SplitK + $DRIVER $OP $DATATYPE $LAYOUT $VERIFY $INIT $LOG $TIME 2 1 $N 128 128 3 3 56 56 2 2 1 1 1 1 1 1 $SplitK + $DRIVER $OP $DATATYPE $LAYOUT $VERIFY $INIT $LOG $TIME 2 1 $N 512 2048 1 1 7 7 1 1 1 1 0 0 0 0 $SplitK + $DRIVER $OP $DATATYPE $LAYOUT $VERIFY $INIT $LOG $TIME 2 1 $N 1024 256 1 1 14 14 1 1 1 1 0 0 0 0 $SplitK + $DRIVER $OP $DATATYPE $LAYOUT $VERIFY $INIT $LOG $TIME 2 1 $N 256 256 3 3 14 14 1 1 1 1 1 1 1 1 $SplitK + $DRIVER $OP $DATATYPE $LAYOUT $VERIFY $INIT $LOG $TIME 2 1 $N 256 256 3 3 28 28 2 2 1 1 1 1 1 1 $SplitK + $DRIVER $OP $DATATYPE $LAYOUT $VERIFY $INIT $LOG $TIME 2 1 $N 128 256 1 1 56 56 1 1 1 1 0 0 0 0 $SplitK + $DRIVER $OP $DATATYPE $LAYOUT $VERIFY $INIT $LOG $TIME 2 1 $N 64 256 1 1 56 56 1 1 1 1 0 0 0 0 $SplitK + $DRIVER $OP $DATATYPE $LAYOUT $VERIFY $INIT $LOG $TIME 2 1 $N 512 512 3 3 14 14 2 2 1 1 1 1 1 1 $SplitK + $DRIVER $OP $DATATYPE $LAYOUT $VERIFY $INIT $LOG $TIME 2 1 $N 128 512 1 1 28 28 1 1 1 1 0 0 0 0 $SplitK + $DRIVER $OP $DATATYPE $LAYOUT $VERIFY $INIT $LOG $TIME 2 1 $N 256 512 1 1 28 28 1 1 1 1 0 0 0 0 $SplitK + $DRIVER $OP $DATATYPE $LAYOUT $VERIFY $INIT $LOG $TIME 2 1 $N 2048 512 1 1 7 7 1 1 1 1 0 0 0 0 $SplitK + $DRIVER $OP $DATATYPE $LAYOUT $VERIFY $INIT $LOG $TIME 2 1 $N 512 512 3 3 7 7 1 1 1 1 1 1 1 1 $SplitK + $DRIVER $OP $DATATYPE $LAYOUT $VERIFY $INIT $LOG $TIME 2 1 $N 256 64 1 1 56 56 1 1 1 1 0 0 0 0 $SplitK + $DRIVER $OP $DATATYPE $LAYOUT $VERIFY $INIT $LOG $TIME 2 1 $N 64 64 1 1 56 56 1 1 1 1 0 0 0 0 $SplitK + $DRIVER $OP $DATATYPE $LAYOUT $VERIFY $INIT $LOG $TIME 2 1 $N 64 64 3 3 56 56 1 1 1 1 1 1 1 1 $SplitK + $DRIVER $OP $DATATYPE $LAYOUT $VERIFY $INIT $LOG $TIME 2 1 $N 64 3 7 7 224 224 2 2 1 1 3 3 3 3 $SplitK diff --git a/script/profile_grouped_conv_fwd.sh b/script/profile_grouped_conv_fwd.sh new file mode 100755 index 0000000000..9a974525ad --- /dev/null +++ b/script/profile_grouped_conv_fwd.sh @@ -0,0 +1,39 @@ +#!/bin/bash + +## GPU visibility +export HIP_VISIBLE_DEVICES=0 +DRIVER="../build/bin/ckProfiler" + +OP=$1 +DATATYPE=$2 +LAYOUT=$3 +INDEXTYPE=$4 +VERIFY=$5 +INIT=$6 +LOG=$7 +TIME=$8 + + N=$9 + +# Resnet50 +######## op datatype indextype layout verify init log time conv_dim G__ N__ K___ C___ Y X Hi__ Wi__ Strides Dilations LeftPads RightPads + $DRIVER $OP $DATATYPE $INDEXTYPE $LAYOUT $VERIFY $INIT $LOG $TIME 2 1 $N 256 1024 1 1 14 14 1 1 1 1 0 0 0 0 + $DRIVER $OP $DATATYPE $INDEXTYPE $LAYOUT $VERIFY $INIT $LOG $TIME 2 1 $N 512 1024 1 1 14 14 1 1 1 1 0 0 0 0 + $DRIVER $OP $DATATYPE $INDEXTYPE $LAYOUT $VERIFY $INIT $LOG $TIME 2 1 $N 128 128 3 3 28 28 1 1 1 1 1 1 1 1 + $DRIVER $OP $DATATYPE $INDEXTYPE $LAYOUT $VERIFY $INIT $LOG $TIME 2 1 $N 512 128 1 1 28 28 1 1 1 1 0 0 0 0 + $DRIVER $OP $DATATYPE $INDEXTYPE $LAYOUT $VERIFY $INIT $LOG $TIME 2 1 $N 128 128 3 3 56 56 2 2 1 1 1 1 1 1 + $DRIVER $OP $DATATYPE $INDEXTYPE $LAYOUT $VERIFY $INIT $LOG $TIME 2 1 $N 512 2048 1 1 7 7 1 1 1 1 0 0 0 0 + $DRIVER $OP $DATATYPE $INDEXTYPE $LAYOUT $VERIFY $INIT $LOG $TIME 2 1 $N 1024 256 1 1 14 14 1 1 1 1 0 0 0 0 + $DRIVER $OP $DATATYPE $INDEXTYPE $LAYOUT $VERIFY $INIT $LOG $TIME 2 1 $N 256 256 3 3 14 14 1 1 1 1 1 1 1 1 + $DRIVER $OP $DATATYPE $INDEXTYPE $LAYOUT $VERIFY $INIT $LOG $TIME 2 1 $N 256 256 3 3 28 28 2 2 1 1 1 1 1 1 + $DRIVER $OP $DATATYPE $INDEXTYPE $LAYOUT $VERIFY $INIT $LOG $TIME 2 1 $N 128 256 1 1 56 56 1 1 1 1 0 0 0 0 + $DRIVER $OP $DATATYPE $INDEXTYPE $LAYOUT $VERIFY $INIT $LOG $TIME 2 1 $N 64 256 1 1 56 56 1 1 1 1 0 0 0 0 + $DRIVER $OP $DATATYPE $INDEXTYPE $LAYOUT $VERIFY $INIT $LOG $TIME 2 1 $N 512 512 3 3 14 14 2 2 1 1 1 1 1 1 + $DRIVER $OP $DATATYPE $INDEXTYPE $LAYOUT $VERIFY $INIT $LOG $TIME 2 1 $N 128 512 1 1 28 28 1 1 1 1 0 0 0 0 + $DRIVER $OP $DATATYPE $INDEXTYPE $LAYOUT $VERIFY $INIT $LOG $TIME 2 1 $N 256 512 1 1 28 28 1 1 1 1 0 0 0 0 + $DRIVER $OP $DATATYPE $INDEXTYPE $LAYOUT $VERIFY $INIT $LOG $TIME 2 1 $N 2048 512 1 1 7 7 1 1 1 1 0 0 0 0 + $DRIVER $OP $DATATYPE $INDEXTYPE $LAYOUT $VERIFY $INIT $LOG $TIME 2 1 $N 512 512 3 3 7 7 1 1 1 1 1 1 1 1 + $DRIVER $OP $DATATYPE $INDEXTYPE $LAYOUT $VERIFY $INIT $LOG $TIME 2 1 $N 256 64 1 1 56 56 1 1 1 1 0 0 0 0 + $DRIVER $OP $DATATYPE $INDEXTYPE $LAYOUT $VERIFY $INIT $LOG $TIME 2 1 $N 64 64 1 1 56 56 1 1 1 1 0 0 0 0 + $DRIVER $OP $DATATYPE $INDEXTYPE $LAYOUT $VERIFY $INIT $LOG $TIME 2 1 $N 64 64 3 3 56 56 1 1 1 1 1 1 1 1 + $DRIVER $OP $DATATYPE $INDEXTYPE $LAYOUT $VERIFY $INIT $LOG $TIME 2 1 $N 64 3 7 7 224 224 2 2 1 1 3 3 3 3 diff --git a/script/run_full_performance_tests.sh b/script/run_full_performance_tests.sh index 01ac1b0a39..e167ce012b 100755 --- a/script/run_full_performance_tests.sh +++ b/script/run_full_performance_tests.sh @@ -90,21 +90,27 @@ print_log_header $gemm_bilinear_log $env_type $branch $host_name ./profile_gemm_bilinear.sh gemm_bilinear 1 2 $verify 1 0 1 2>&1 | tee -a $gemm_bilinear_log ./profile_gemm_bilinear.sh gemm_bilinear 1 3 $verify 1 0 1 2>&1 | tee -a $gemm_bilinear_log -#run conv_fwd tests -export conv_fwd_log="perf_conv_fwd.log" -print_log_header $conv_fwd_log $env_type $branch $host_name -./profile_conv_fwd.sh conv_fwd 0 1 $verify 1 0 1 256 2>&1 | tee -a $conv_fwd_log -./profile_conv_fwd.sh conv_fwd 1 1 $verify 1 0 1 256 2>&1 | tee -a $conv_fwd_log -./profile_conv_fwd.sh conv_fwd 2 1 $verify 1 0 1 256 2>&1 | tee -a $conv_fwd_log -./profile_conv_fwd.sh conv_fwd 3 1 $verify 1 0 1 256 2>&1 | tee -a $conv_fwd_log +#run grouped_fwd tests +export grouped_conv_fwd_log="perf_grouped_conv_fwd.log" +print_log_header $grouped_conv_fwd_log $env_type $branch $host_name +./profile_grouped_conv_fwd.sh grouped_conv_fwd 0 1 0 $verify 1 0 1 256 2>&1 | tee -a $grouped_conv_fwd_log +./profile_grouped_conv_fwd.sh grouped_conv_fwd 1 1 0 $verify 1 0 1 256 2>&1 | tee -a $grouped_conv_fwd_log +./profile_grouped_conv_fwd.sh grouped_conv_fwd 2 1 0 $verify 1 0 1 256 2>&1 | tee -a $grouped_conv_fwd_log -#run conv_bwd_data tests -export conv_bwd_data_log="perf_conv_bwd_data.log" -print_log_header $conv_bwd_data_log $env_type $branch $host_name -./profile_conv_bwd_data.sh conv_bwd_data 0 1 $verify 1 0 1 256 2>&1 | tee -a $conv_bwd_data_log -./profile_conv_bwd_data.sh conv_bwd_data 1 1 $verify 1 0 1 256 2>&1 | tee -a $conv_bwd_data_log -./profile_conv_bwd_data.sh conv_bwd_data 2 1 $verify 1 0 1 256 2>&1 | tee -a $conv_bwd_data_log -./profile_conv_bwd_data.sh conv_bwd_data 3 1 $verify 1 0 1 256 2>&1 | tee -a $conv_bwd_data_log +#run grouped_bwd_data tests +export grouped_conv_bwd_data_log="perf_grouped_conv_bwd_data.log" +print_log_header $grouped_conv_bwd_data_log $env_type $branch $host_name +./profile_grouped_conv_bwd_data.sh grouped_conv_bwd_data 0 1 $verify 1 0 1 256 2>&1 | tee -a $grouped_conv_bwd_data_log +./profile_grouped_conv_bwd_data.sh grouped_conv_bwd_data 1 1 $verify 1 0 1 256 2>&1 | tee -a $grouped_conv_bwd_data_log +./profile_grouped_conv_bwd_data.sh grouped_conv_bwd_data 2 1 $verify 1 0 1 256 2>&1 | tee -a $grouped_conv_bwd_data_log + +#run grouped_bwd_weight tests +export grouped_conv_bwd_weight_log="perf_grouped_conv_bwd_weight.log" +print_log_header $grouped_conv_bwd_weight_log $env_type $branch $host_name +./profile_grouped_conv_bwd_weight.sh grouped_conv_bwd_weight 0 2 $verify 1 0 1 256 1 2>&1 | tee -a $grouped_conv_bwd_weight_log +./profile_grouped_conv_bwd_weight.sh grouped_conv_bwd_weight 1 2 $verify 1 0 1 256 1 2>&1 | tee -a $grouped_conv_bwd_weight_log +./profile_grouped_conv_bwd_weight.sh grouped_conv_bwd_weight 2 2 $verify 1 0 1 256 1 2>&1 | tee -a $grouped_conv_bwd_weight_log +./profile_grouped_conv_bwd_weight.sh grouped_conv_bwd_weight 1 2 $verify 1 0 1 256 4 2>&1 | tee -a $grouped_conv_bwd_weight_log #run resnet50 tests export resnet256_log="perf_resnet50_N256.log" diff --git a/script/run_performance_tests.sh b/script/run_performance_tests.sh index 4e3a6fc8eb..317d270983 100755 --- a/script/run_performance_tests.sh +++ b/script/run_performance_tests.sh @@ -51,6 +51,21 @@ print_log_header $gemm_log $env_type $branch $host_name ./profile_gemm.sh gemm 2 3 $verify 1 0 1 | tee -a $gemm_log ./profile_gemm.sh gemm 3 3 $verify 1 0 1 | tee -a $gemm_log +#run grouped_fwd fp16 tests +export grouped_conv_fwd_log="perf_grouped_conv_fwd_fp16.log" +print_log_header $conv_fwd_log $env_type $branch $host_name +./profile_grouped_conv_fwd.sh grouped_conv_fwd 1 1 0 $verify 1 0 1 256 2>&1 | tee -a $grouped_conv_fwd_log + +#run grouped_bwd_data fp16 tests +export grouped_conv_bwd_data_log="perf_grouped_conv_bwd_data_fp16.log" +print_log_header $grouped_conv_bwd_data_log $env_type $branch $host_name +./profile_grouped_conv_bwd_data.sh grouped_conv_bwd_data 1 1 $verify 1 0 1 256 2>&1 | tee -a $grouped_conv_bwd_data_log + +#run grouped_bwd_weight fp16 tests +export grouped_conv_bwd_weight_log="perf_grouped_conv_bwd_weight_fp16.log" +print_log_header $grouped_conv_bwd_weight_log $env_type $branch $host_name +./profile_grouped_conv_bwd_weight.sh grouped_conv_bwd_weight 1 1 $verify 1 0 1 256 1 2>&1 | tee -a $grouped_conv_bwd_weight_log + #run resnet50 tests export resnet256_log="perf_resnet50_N256.log" print_log_header $resnet256_log $env_type $branch $host_name diff --git a/test/grouped_convnd_fwd/CMakeLists.txt b/test/grouped_convnd_fwd/CMakeLists.txt index f611e66243..4ceb4a2d99 100644 --- a/test/grouped_convnd_fwd/CMakeLists.txt +++ b/test/grouped_convnd_fwd/CMakeLists.txt @@ -7,6 +7,12 @@ if(GPU_TARGETS MATCHES "gfx9" OR GPU_TARGETS MATCHES "gfx11") endif() endif() +if(GPU_TARGETS MATCHES "gfx9") + add_executable(test_grouped_convnd_fwd_large_cases_xdl test_grouped_convnd_fwd_large_cases_xdl.cpp) + target_compile_options(test_grouped_convnd_fwd_large_cases_xdl PRIVATE -Wno-global-constructors -Wno-undef) + target_link_libraries(test_grouped_convnd_fwd_large_cases_xdl PRIVATE gtest_main getopt::getopt utility device_grouped_conv1d_fwd_instance device_grouped_conv2d_fwd_instance device_grouped_conv3d_fwd_instance) +endif() + add_gtest_executable(test_grouped_convnd_fwd_multi_ab_interface test_grouped_convnd_fwd_multi_ab_interface.cpp) if(result EQUAL 0) target_link_libraries(test_grouped_convnd_fwd_multi_ab_interface PRIVATE utility) diff --git a/test/grouped_convnd_fwd/test_grouped_convnd_fwd.cpp b/test/grouped_convnd_fwd/test_grouped_convnd_fwd.cpp index c86b18e77e..b960676574 100644 --- a/test/grouped_convnd_fwd/test_grouped_convnd_fwd.cpp +++ b/test/grouped_convnd_fwd/test_grouped_convnd_fwd.cpp @@ -17,7 +17,7 @@ class TestGroupedConvndFwd : public ::testing::Test using InLayout = std::tuple_element_t<1, Tuple>; using WeiLayout = std::tuple_element_t<2, Tuple>; using OutLayout = std::tuple_element_t<3, Tuple>; - using IndexType = std::tuple_element_t<4, Tuple>; + using IndexType = ck::index_t; std::vector conv_params; @@ -50,31 +50,28 @@ class TestGroupedConvndFwd : public ::testing::Test using namespace ck::tensor_layout::convolution; -using KernelTypes1d = ::testing::Types, - std::tuple, - std::tuple, - std::tuple>; +using KernelTypes1d = ::testing::Types, + std::tuple, + std::tuple, + std::tuple>; -using KernelTypes2d = ::testing::Types, - std::tuple, - std::tuple, - std::tuple, - std::tuple, - std::tuple, - std::tuple, - std::tuple>; +using KernelTypes2d = ::testing::Types, + std::tuple, + std::tuple, + std::tuple, + std::tuple, + std::tuple, + std::tuple, + std::tuple>; -using KernelTypes3d = ::testing::Types, - std::tuple, - std::tuple, - std::tuple, - std::tuple, - std::tuple, - std::tuple, - std::tuple>; - -using KernelTypes2dLargeCases = - ::testing::Types>; +using KernelTypes3d = ::testing::Types, + std::tuple, + std::tuple, + std::tuple, + std::tuple, + std::tuple, + std::tuple, + std::tuple>; template class TestGroupedConvndFwd1d : public TestGroupedConvndFwd @@ -91,15 +88,9 @@ class TestGroupedConvndFwd3d : public TestGroupedConvndFwd { }; -template -class TestGroupedConvndFwd2dLargeCases : public TestGroupedConvndFwd -{ -}; - TYPED_TEST_SUITE(TestGroupedConvndFwd1d, KernelTypes1d); TYPED_TEST_SUITE(TestGroupedConvndFwd2d, KernelTypes2d); TYPED_TEST_SUITE(TestGroupedConvndFwd3d, KernelTypes3d); -TYPED_TEST_SUITE(TestGroupedConvndFwd2dLargeCases, KernelTypes2dLargeCases); TYPED_TEST(TestGroupedConvndFwd1d, Test1D) { @@ -149,17 +140,3 @@ TYPED_TEST(TestGroupedConvndFwd3d, Test3D) {3, 96, 1, 1, 1, {3, 3, 3}, {4, 30, 160}, {1, 1, 1}, {1, 1, 1}, {1, 1, 1}, {1, 1, 1}}); this->template Run<3>(); } - -TYPED_TEST(TestGroupedConvndFwd2dLargeCases, Test2DLargeCases) -{ - // Case larger than 2GB - this->conv_params.push_back( - {2, 1, 64, 4, 192, {2, 2}, {224, 224}, {224, 224}, {1, 1}, {0, 0}, {0, 0}}); - // With supported NumGroupsToMerge > 1 - this->conv_params.push_back( - {2, 32, 64, 1, 1, {2, 2}, {672, 672}, {672, 672}, {1, 1}, {0, 0}, {0, 0}}); - // When image is larger than 2GB - this->conv_params.push_back( - {2, 1, 1, 256, 256, {3, 3}, {4096, 2048}, {1024, 1024}, {3, 3}, {1, 1}, {1, 1}}); - this->template Run<2>(); -} diff --git a/test/grouped_convnd_fwd/test_grouped_convnd_fwd_large_cases_xdl.cpp b/test/grouped_convnd_fwd/test_grouped_convnd_fwd_large_cases_xdl.cpp new file mode 100644 index 0000000000..3d734fa5e5 --- /dev/null +++ b/test/grouped_convnd_fwd/test_grouped_convnd_fwd_large_cases_xdl.cpp @@ -0,0 +1,127 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved. + +#include +#include +#include +#include +#include + +#include "profiler/profile_grouped_conv_fwd_impl.hpp" + +template +class TestGroupedConvndFwd : public ::testing::Test +{ + protected: + using DataType = std::tuple_element_t<0, Tuple>; + using InLayout = std::tuple_element_t<1, Tuple>; + using WeiLayout = std::tuple_element_t<2, Tuple>; + using OutLayout = std::tuple_element_t<3, Tuple>; + using IndexType = ck::long_index_t; + + std::vector conv_params; + + template + void Run() + { + EXPECT_FALSE(conv_params.empty()); + bool pass = true; + for(auto& param : conv_params) + { + pass = pass && ck::profiler::profile_grouped_conv_fwd_impl( + true, // do_verification + 1, // init_method: integer value + false, // do_log + false, // time_kernel + param); + } + EXPECT_TRUE(pass); + } +}; + +using namespace ck::tensor_layout::convolution; + +using KernelTypes2d = ::testing::Types, + std::tuple, + std::tuple>; + +using KernelTypes3d = ::testing::Types, + std::tuple, + std::tuple>; + +template +class TestGroupedConvndFwd2d : public TestGroupedConvndFwd +{ +}; + +template +class TestGroupedConvndFwd3d : public TestGroupedConvndFwd +{ +}; + +TYPED_TEST_SUITE(TestGroupedConvndFwd2d, KernelTypes2d); +TYPED_TEST_SUITE(TestGroupedConvndFwd3d, KernelTypes3d); + +TYPED_TEST(TestGroupedConvndFwd2d, Test2D) +{ + // Case larger than 2GB + this->conv_params.push_back( + {2, 1, 128, 4, 192, {2, 2}, {224, 224}, {224, 224}, {1, 1}, {0, 0}, {0, 0}}); + // With supported NumGroupsToMerge > 1 + this->conv_params.push_back( + {2, 32, 64, 1, 1, {2, 2}, {672, 672}, {672, 672}, {1, 1}, {0, 0}, {0, 0}}); + // When image is larger than 2GB + this->conv_params.push_back( + {2, 2, 2, 128, 128, {3, 3}, {4096, 2048}, {300, 300}, {3, 3}, {1, 1}, {1, 1}}); + this->template Run<2>(); +} + +TYPED_TEST(TestGroupedConvndFwd3d, Test3D) +{ + // Case larger than 2GB + this->conv_params.push_back({3, + 1, + 128, + 4, + 192, + {2, 2, 2}, + {2, 224, 224}, + {1, 224, 224}, + {1, 1, 1}, + {0, 0, 0}, + {0, 0, 0}}); + // With supported NumGroupsToMerge > 1 + this->conv_params.push_back({3, + 32, + 64, + 1, + 1, + {2, 2, 2}, + {360, 2, 672}, + {360, 2, 672}, + {1, 1, 1}, + {0, 0, 0}, + {0, 0, 0}}); + // When image is larger than 2GB + this->conv_params.push_back({3, + 1, + 2, + 128, + 128, + {3, 1, 3}, + {900, 2, 2048}, + {300, 1, 300}, + {3, 2, 3}, + {1, 1, 1}, + {1, 1, 1}}); + this->template Run<3>(); +} From 79a5d9c10c45a4290ba916695dce4625022c89df Mon Sep 17 00:00:00 2001 From: Dan Yao Date: Sat, 17 Aug 2024 04:40:10 +0800 Subject: [PATCH 07/20] [CK_TILE] FA bwd kernels optimization (#1397) * tmp save * fix batch deterministic bugs * fix group deterministic bugs * codegen update * reorder files * bias support * hd256 bias support * bwd smoke test update * simplify convert dq * fix hd256 dropout scratch * do{}while() -> while(){} * comments * remove FmhaBwdTilePartitioner * save clear_tile * refactor dropout * code cleanup * code cleanup * comments * fix epilogue problem * fix fwd dropout * group convert_dq opt * fix dq alignment * Do not store storerandval in bwd for flash attention integration * fix hd32 error and boost performance * revert * Remove duplicated WarpGemm definitions in the policy file * dropout patch for mrepeat 16*16 * code sync up * dq_acc stride * dq_acc stride stuff * codegen update * fwd dropout revert * fix hd128 scratches and boost performance * receipt 3 for simplified smoke test * more strides for fa integration * fix hd64 scratches and boost performance * non-iglp pipeline for headdim padding cases * dpad same as dvpad for flash attention integration * unpadded lse&d for group mode * Support unpad layout for group lse * Support unpad lse layout for splitkv * Fix stride for splitkv kernel * fix unpadded lse issue in fwd splitkv * comment * solve lds read&write conflicts * rename * bias rename * tile index revert --------- Co-authored-by: danyao12 Co-authored-by: rocking Co-authored-by: Qianfeng Zhang --- example/ck_tile/01_fmha/CMakeLists.txt | 7 +- .../ck_tile/01_fmha/codegen/cpp_symbol_map.py | 16 + .../ck_tile/01_fmha/codegen/ops/fmha_bwd.py | 547 ++-- example/ck_tile/01_fmha/fmha_bwd.cpp | 67 +- example/ck_tile/01_fmha/fmha_bwd.hpp | 106 +- example/ck_tile/01_fmha/fmha_fwd.cpp | 25 +- example/ck_tile/01_fmha/fmha_fwd.hpp | 4 - .../ck_tile/01_fmha/script/smoke_test_bwd.sh | 22 +- .../core/algorithm/coordinate_transform.hpp | 42 +- include/ck_tile/core/numeric/vector_type.hpp | 9 + .../ck_tile/core/tensor/tile_distribution.hpp | 5 +- include/ck_tile/core/utility/philox_rand.hpp | 33 + include/ck_tile/ops/fmha.hpp | 11 +- .../ck_tile/ops/fmha/block/block_dropout.hpp | 397 ++- .../ops/fmha/kernel/fmha_bwd_kernel.hpp | 885 +++--- .../fmha/kernel/fmha_bwd_tile_partitioner.hpp | 54 - .../ops/fmha/kernel/fmha_fwd_kernel.hpp | 6 +- .../fmha_fwd_splitkv_combine_kernel.hpp | 39 +- .../fmha/kernel/fmha_fwd_splitkv_kernel.hpp | 25 +- .../pipeline/block_fmha_bwd_convert_dq.hpp | 141 + .../fmha/pipeline/block_fmha_bwd_dot_do_o.hpp | 6 +- ...block_fmha_bwd_dot_do_o_default_policy.hpp | 20 - ...k_fmha_bwd_dq_dk_dv_pipeline_kr_ktr_vr.hpp | 782 +++++ ...a_bwd_dq_dk_dv_pipeline_kr_ktr_vr_iglp.hpp | 1037 +++++++ ...k_fmha_bwd_dq_dk_dv_pipeline_ks_kts_vr.hpp | 848 ------ ...k_dv_pipeline_ks_kts_vr_default_policy.hpp | 20 - ...block_fmha_bwd_dq_dk_dv_pipeline_ks_vr.hpp | 821 ----- ...dq_dk_dv_pipeline_ks_vr_default_policy.hpp | 20 - ...mha_bwd_dq_dk_dv_pipeline_qs_ks_vr_dos.hpp | 692 ----- ...v_pipeline_qs_ks_vr_dos_default_policy.hpp | 20 - ...block_fmha_bwd_pipeline_default_policy.hpp | 2631 ++++++++++------- .../pipeline/block_fmha_bwd_pipeline_enum.hpp | 5 +- .../block_fmha_bwd_pipeline_problem.hpp | 40 +- .../ops/fmha/pipeline/tile_fmha_traits.hpp | 10 + include/ck_tile/ops/gemm.hpp | 3 + .../block/block_gemm_areg_breg_creg_v1.hpp | 202 ++ ...k_gemm_areg_breg_creg_v1_custom_policy.hpp | 36 + ..._gemm_areg_breg_creg_v1_default_policy.hpp | 33 + .../block/block_gemm_areg_bsmem_creg_v1.hpp | 15 +- .../block/block_gemm_asmem_breg_creg_v1.hpp | 15 +- include/ck_tile/ops/gemm/warp/warp_gemm.hpp | 6 + .../gemm/warp/warp_gemm_attribute_mfma.hpp | 12 +- .../ops/gemm/warp/warp_gemm_dispatcher.hpp | 22 +- 43 files changed, 5515 insertions(+), 4222 deletions(-) delete mode 100644 include/ck_tile/ops/fmha/kernel/fmha_bwd_tile_partitioner.hpp create mode 100644 include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_convert_dq.hpp delete mode 100644 include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dot_do_o_default_policy.hpp create mode 100644 include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_kr_ktr_vr.hpp create mode 100644 include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_kr_ktr_vr_iglp.hpp delete mode 100644 include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_ks_kts_vr.hpp delete mode 100644 include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_ks_kts_vr_default_policy.hpp delete mode 100644 include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_ks_vr.hpp delete mode 100644 include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_ks_vr_default_policy.hpp delete mode 100644 include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_qs_ks_vr_dos.hpp delete mode 100644 include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_qs_ks_vr_dos_default_policy.hpp create mode 100644 include/ck_tile/ops/gemm/block/block_gemm_areg_breg_creg_v1.hpp create mode 100644 include/ck_tile/ops/gemm/block/block_gemm_areg_breg_creg_v1_custom_policy.hpp create mode 100644 include/ck_tile/ops/gemm/block/block_gemm_areg_breg_creg_v1_default_policy.hpp diff --git a/example/ck_tile/01_fmha/CMakeLists.txt b/example/ck_tile/01_fmha/CMakeLists.txt index e30e9e793c..5d59b246e1 100644 --- a/example/ck_tile/01_fmha/CMakeLists.txt +++ b/example/ck_tile/01_fmha/CMakeLists.txt @@ -6,7 +6,7 @@ execute_process( execute_process( COMMAND ${Python3_EXECUTABLE} ${CMAKE_CURRENT_LIST_DIR}/generate.py - --api bwd --list_blobs ${CMAKE_CURRENT_BINARY_DIR}/bwd_blob_list.txt + --api bwd --list_blobs ${CMAKE_CURRENT_BINARY_DIR}/bwd_blob_list.txt --receipt 3 ) # NOTE: for cmake, the FMHA_FWD_GEN_BLOBS/FMHA_BWD_GEN_BLOBS files must be in the same directory @@ -23,7 +23,7 @@ add_custom_command( add_custom_command( OUTPUT ${FMHA_BWD_GEN_BLOBS} COMMAND ${Python3_EXECUTABLE} ${CMAKE_CURRENT_LIST_DIR}/generate.py - --api bwd --output_dir ${CMAKE_CURRENT_BINARY_DIR} + --api bwd --output_dir ${CMAKE_CURRENT_BINARY_DIR} --receipt 3 ) set(EXAMPLE_FMHA_FWD "tile_example_fmha_fwd") @@ -55,11 +55,10 @@ set(EXAMPLE_FMHA_BWD_COMPILE_OPTIONS) # ... because they are auto-generated if(FMHA_FWD_FAST_EXP2) list(APPEND EXAMPLE_FMHA_FWD_COMPILE_OPTIONS -Wno-undefined-func-template -DCK_TILE_FMHA_FWD_FAST_EXP2=1 -fgpu-flush-denormals-to-zero) - list(APPEND EXAMPLE_FMHA_BWD_COMPILE_OPTIONS -Wno-undefined-func-template -DCK_TILE_FMHA_FWD_FAST_EXP2=1 -fgpu-flush-denormals-to-zero) else() list(APPEND EXAMPLE_FMHA_FWD_COMPILE_OPTIONS -Wno-undefined-func-template -DCK_TILE_FMHA_FWD_FAST_EXP2=0) - list(APPEND EXAMPLE_FMHA_BWD_COMPILE_OPTIONS -Wno-undefined-func-template -DCK_TILE_FMHA_FWD_FAST_EXP2=0) endif() +list(APPEND EXAMPLE_FMHA_BWD_COMPILE_OPTIONS -Wno-undefined-func-template -fgpu-flush-denormals-to-zero) # Allow comparing floating points directly in order to check sentinel values list(APPEND EXAMPLE_FMHA_FWD_COMPILE_OPTIONS -Wno-float-equal) 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 d3d215f7f5..a5862ad5d9 100644 --- a/example/ck_tile/01_fmha/codegen/cpp_symbol_map.py +++ b/example/ck_tile/01_fmha/codegen/cpp_symbol_map.py @@ -66,6 +66,22 @@ BIAS_CHECK_MAP = { "alibi" : "bias_enum::alibi" } +DROPOUT_MAP = { + "no" : "ck_tile::BlockDropoutBwd", + "dropout_wg32" : "ck_tile::BlockDropoutBwd", + "dropout_wg32_storerandval" : "ck_tile::BlockDropoutBwd", + "dropout_wg16" : "ck_tile::BlockDropoutBwd", + "dropout_wg16_storerandval" : "ck_tile::BlockDropoutBwd" +} + +DROPOUT_CHECK_MAP = { + "no" : "t.has_dropout == false", + "dropout_wg32" : "t.has_dropout == true && t.is_store_randval == false", + "dropout_wg32_storerandval" : "t.has_dropout == true && t.is_store_randval == true", + "dropout_wg16" : "t.has_dropout == true && t.is_store_randval == false", + "dropout_wg16_storerandval" : "t.has_dropout == true && t.is_store_randval == true", +} + MODE_MAP = { "batch" : "false", "group" : "true" diff --git a/example/ck_tile/01_fmha/codegen/ops/fmha_bwd.py b/example/ck_tile/01_fmha/codegen/ops/fmha_bwd.py index 0df115dc3d..096394c0c9 100644 --- a/example/ck_tile/01_fmha/codegen/ops/fmha_bwd.py +++ b/example/ck_tile/01_fmha/codegen/ops/fmha_bwd.py @@ -14,15 +14,13 @@ from codegen.cpp_symbol_map import * BWD_DQDKDV_PIPELINE_MAP = { - "ks_kts_vr" : "ck_tile::BlockFmhaBwdDQDKDVPipelineKSKTSVR", - "qs_ks_vr_dos" : "ck_tile::BlockFmhaBwdDQDKDVPipelineQSKSVROGradS", - "ks_vr" : "ck_tile::BlockFmhaBwdDQDKDVPipelineKSVR", + "kr_ktr_vr_iglp" : "ck_tile::BlockFmhaBwdDQDKDVPipelineKRKTRVRIGLP", + "kr_ktr_vr" : "ck_tile::BlockFmhaBwdDQDKDVPipelineKRKTRVR", } BWD_DQDKDV_PIPELINE_ENUM_MAP = { - "ks_kts_vr" : "ck_tile::BlockFmhaBwdPipelineEnum::KSKTSVR", - "qs_ks_vr_dos" : "ck_tile::BlockFmhaBwdPipelineEnum::QSKSVROGradS", - "ks_vr" : "ck_tile::BlockFmhaBwdPipelineEnum::KSVR", + "kr_ktr_vr_iglp" : "ck_tile::BlockFmhaBwdPipelineEnum::KRKTRVR_IGLP", + "kr_ktr_vr" : "ck_tile::BlockFmhaBwdPipelineEnum::KRKTRVR", } FMHA_BWD_KERNEL_HEADER = """// SPDX-License-Identifier: MIT @@ -34,39 +32,42 @@ FMHA_BWD_KERNEL_HEADER = """// SPDX-License-Identifier: MIT FMHA_BWD_DQ_DK_DV_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_bk1}, {F_bk2}, {F_bk3}, {F_bk4}, {F_bhdq}, {F_bhdv}>; +using fmha_block_tile_{F_idx} = ck_tile:: + sequence<{F_bm0}, {F_bn0}, {F_bk0}, {F_bk1}, {F_bk2}, {F_bk3}, {F_bk4}, {F_bhdq}, {F_bhdv}>; using fmha_block_warps0_{F_idx} = ck_tile::sequence<{F_rm0}, {F_rn0}, {F_rk0}>; using fmha_block_warps1_{F_idx} = ck_tile::sequence<{F_rm1}, {F_rn1}, {F_rk1}>; using fmha_block_warps2_{F_idx} = ck_tile::sequence<{F_rm2}, {F_rn2}, {F_rk2}>; -using fmha_warp_tile_{F_idx} = ck_tile::sequence<{F_wm}, {F_wn}, {F_wk}>; +using fmha_warp_tile0_{F_idx} = ck_tile::sequence<{F_wm0}, {F_wn0}, {F_wk0}>; +using fmha_warp_tile1_{F_idx} = ck_tile::sequence<{F_wm1}, {F_wn1}, {F_wk1}>; // TODO: simplify Gemm0~4BlockWarps in TileFmhaBwdShape // G0&G2 -> GSdP // G1&G3 -> GdKV // G4 -> GdQ using fmha_bwd_shape_{F_idx} = ck_tile::TileFmhaBwdShape; + fmha_block_warps0_{F_idx}, + fmha_warp_tile0_{F_idx}, + fmha_block_warps1_{F_idx}, + fmha_warp_tile1_{F_idx}, + fmha_block_warps0_{F_idx}, + fmha_warp_tile0_{F_idx}, + fmha_block_warps1_{F_idx}, + fmha_warp_tile1_{F_idx}, + fmha_block_warps2_{F_idx}, + fmha_warp_tile0_{F_idx}>; using fmha_bwd_trait_{F_idx} = ck_tile::TileFmhaTraits<{F_spad}, - {F_skpad}, - {F_dpad}, - {F_dvpad}, - {F_bias}, - {F_dbias}, - false, - {F_dropout}, - false, - {F_occupancy}>; -using fmha_mask_{F_idx} = {F_mask}; + {F_skpad}, + {F_dpad}, + {F_dvpad}, + {F_bias}, + {F_dbias}, + false, + false, + false, + {F_occupancy}>; +using fmha_mask_{F_idx} = {F_mask}; +using fmha_dropout_{F_idx} = {F_dropout}; using fmha_bwd_pipeline_problem_{F_idx} = ck_tile::BlockFmhaBwdPipelineProblem< typename FmhaBwdTypeConfig::QDataType, @@ -86,55 +87,72 @@ using fmha_bwd_pipeline_problem_{F_idx} = ck_tile::BlockFmhaBwdPipelineProblem< typename FmhaBwdTypeConfig::BiasGradDataType, fmha_bwd_shape_{F_idx}, {F_mode}, + {F_deterministic}, fmha_mask_{F_idx}, + fmha_dropout_{F_idx}, fmha_bwd_trait_{F_idx}>; -using fmha_bwd_pipeline_{F_idx} = {F_pipeline}< - fmha_bwd_pipeline_problem_{F_idx}>; +using fmha_bwd_pipeline_{F_idx} = {F_pipeline}; -using fmha_bwd_dk_epilogue_{F_idx} = - ck_tile::Default2DEpilogue::AccDataType, - typename FmhaBwdTypeConfig<{F_dtype}>::KGradDataType, - false, false>>; +using fmha_bwd_dk_epilogue_{F_idx} = ck_tile::Default2DEpilogue< + ck_tile::Default2DEpilogueProblem::AccDataType, + typename FmhaBwdTypeConfig<{F_dtype}>::KGradDataType, + {F_skpad}, + {F_dpad}>>; -using fmha_bwd_dv_epilogue_{F_idx} = - ck_tile::Default2DEpilogue::AccDataType, - typename FmhaBwdTypeConfig<{F_dtype}>::VGradDataType, - false, false>>; +using fmha_bwd_dv_epilogue_{F_idx} = ck_tile::Default2DEpilogue< + ck_tile::Default2DEpilogueProblem::AccDataType, + typename FmhaBwdTypeConfig<{F_dtype}>::VGradDataType, + {F_skpad}, + {F_dvpad}>>; using fmha_bwd_dq_dk_dv_kernel_{F_idx} = - ck_tile::FmhaBwdDQDKDVKernel, - fmha_bwd_pipeline_{F_idx}, - fmha_bwd_dk_epilogue_{F_idx}, - fmha_bwd_dv_epilogue_{F_idx}>; + ck_tile::FmhaBwdDQDKDVKernel; -using dq_dk_dv_trait_{F_idx} = fmha_bwd_dq_dk_dv_traits_<{F_hdim}, {F_dtype}, {F_mode}, {F_pipeline_enum}, fmha_mask_{F_idx}, {F_bias}, {F_dbias}, {F_dropout}, {F_spad}, {F_skpad}, {F_dpad}, {F_dvpad}>; +using dq_dk_dv_trait_{F_idx} = fmha_bwd_dq_dk_dv_traits_<{F_hdim}, + {F_dtype}, + {F_mode}, + {F_pipeline_enum}, + fmha_mask_{F_idx}, + fmha_dropout_{F_idx}, + {F_bias}, + {F_dbias}, + {F_spad}, + {F_skpad}, + {F_dpad}, + {F_dvpad}, + {F_deterministic}>; #include -template<> +template <> float fmha_bwd_dq_dk_dv_(const ck_tile::stream_config& s, fmha_bwd_args a) {{ using k_ = fmha_bwd_dq_dk_dv_kernel_{F_idx}; if(s.log_level_ > 0) std::cout << ", " << k_::GetName() << std::flush; - auto [kargs, grids] = fmha_bwd_dq_dk_dv_create_kargs_and_grids(a); - constexpr dim3 blocks = k_::BlockSize(); + auto [kargs, grids] = fmha_bwd_dq_dk_dv_create_kargs_and_grids(a); + constexpr dim3 blocks = k_::BlockSize(); constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; - return ck_tile::launch_kernel(s, ck_tile::make_kernel(k_{{}}, grids, blocks, 0, kargs)); + return ck_tile::launch_kernel( + s, ck_tile::make_kernel(k_{{}}, grids, blocks, 0, kargs)); }} -template<> -void fmha_bwd_dq_dk_dv_oneshot_(const ck_tile::stream_config& s, fmha_bwd_args a) +template <> +void fmha_bwd_dq_dk_dv_oneshot_(const ck_tile::stream_config& s, + fmha_bwd_args a) {{ - using k_ = fmha_bwd_dq_dk_dv_kernel_{F_idx}; - auto [kargs, grids] = fmha_bwd_dq_dk_dv_create_kargs_and_grids(a); - constexpr dim3 blocks = k_::BlockSize(); + using k_ = fmha_bwd_dq_dk_dv_kernel_{F_idx}; + auto [kargs, grids] = fmha_bwd_dq_dk_dv_create_kargs_and_grids(a); + constexpr dim3 blocks = k_::BlockSize(); constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; - ck_tile::make_kernel(k_{{}}, grids, blocks, 0, kargs)(ck_tile::stream_config{{s.stream_id_}}); + ck_tile::make_kernel(k_{{}}, grids, blocks, 0, kargs)( + ck_tile::stream_config{{s.stream_id_}}); }} -template<> +template <> std::string fmha_bwd_dq_dk_dv_get_name_() {{ using k_ = fmha_bwd_dq_dk_dv_kernel_{F_idx}; @@ -146,14 +164,15 @@ FMHA_BWD_API_FILENAME="fmha_bwd_api.cpp" FMHA_BWD_API=""" #include -template +template float fmha_bwd_(const ck_tile::stream_config& s, fmha_bwd_args a) {{ if(s.log_level_ > 0) - std::cout << ", " << fmha_bwd_dot_do_o_get_name_() << ", " << fmha_bwd_dq_dk_dv_get_name_() << std::flush; + std::cout << ", " << fmha_bwd_dot_do_o_get_name_() << ", " << fmha_bwd_dq_dk_dv_get_name_() << ", " << fmha_bwd_convert_dq_get_name_() << std::flush; return ck_tile::launch_kernel(s, - [=](const ck_tile::stream_config& s_){{ fmha_bwd_dot_do_o_oneshot_(s_, a); }}, - [=](const ck_tile::stream_config& s_){{ fmha_bwd_dq_dk_dv_oneshot_(s_, a); }} + [=](const ck_tile::stream_config& s_){{ fmha_bwd_dot_do_o_oneshot_(s_, a); }}, + [=](const ck_tile::stream_config& s_){{ fmha_bwd_dq_dk_dv_oneshot_(s_, a); }}, + [=](const ck_tile::stream_config& s_){{ fmha_bwd_convert_dq_oneshot_(s_, a); }} ); }} @@ -173,38 +192,36 @@ FMHA_BWD_API_PER_HDIM_CASE=""" {F_if} (t.hdim_q <= {F_hdim} && t.hdim_v < }} """ -FMHA_BWD_API_INNER_DISPATCH=""" {F_if}((t.is_group_mode == {F_mode}) && ({F_mask_check}) && (t.bias_type == {F_bias_check}) && (t.has_dbias == {F_dbias}) && (t.has_dropout == {F_dropout}) && - ({F_scheck}) && ({F_skcheck}) && ({F_dcheck}) && ({F_dvcheck})) {{ - using dq_dk_dv_trait_ = fmha_bwd_dq_dk_dv_traits_<{F_hdim}, {F_dtype}, {F_mode}, {F_pipeline_enum}, {F_mask}, {F_bias}, {F_dbias}, {F_dropout}, {F_spad0}, {F_skpad}, {F_dpad}, {F_dvpad}>; +FMHA_BWD_API_INNER_DISPATCH=""" {F_if}((t.is_group_mode == {F_mode}) && ({F_mask_check}) && (t.bias_type == {F_bias_check}) && (t.has_dbias == {F_dbias}) && ({F_dropout_check}) && + ({F_scheck}) && ({F_skcheck}) && ({F_dcheck}) && ({F_dvcheck}) && (t.is_deterministic == {F_deterministic})) {{ using dot_do_o_trait_ = fmha_bwd_dot_do_o_traits_<{F_hdim}, {F_dtype}, {F_mode}, {F_spad1}, {F_dvpad}>; - r = fmha_bwd_(s, a); + using dq_dk_dv_trait_ = fmha_bwd_dq_dk_dv_traits_<{F_hdim}, {F_dtype}, {F_mode}, {F_pipeline_enum}, {F_mask}, {F_dropout}, {F_bias}, {F_dbias}, {F_spad0}, {F_skpad}, {F_dpad}, {F_dvpad}, {F_deterministic}>; + using convert_dq_trait_ = fmha_bwd_convert_dq_traits_<{F_hdim}, {F_dtype}, {F_mode}, {F_spad1}, {F_dpad}, {F_deterministic}>; + r = fmha_bwd_(s, a); return r; }} """ @dataclass class FmhaBwdDQDKDVApiTrait: - pipeline : str + pipeline : str # sync with fmha_bwd_traits<>, to generate fallback calls - hdim : str - dtype : str # data type - mode : str # value from MODE_MAP - bm0 : int # tile size along q seqlen (block size) - bn0 : int # tile size along k seqlen - bhdq : int # q head_dim - bhdv : int # v head_dim - mask : str - bias : str - dbias : str - dropout : str - spad : str - skpad : str - dpad : str - dvpad : str - - @property - def name(self) -> str: - return f'{self.pipeline}-{self.hdim}-{self.dtype}-{self.mode}-{self.mask}-{self.bias}-{self.dbias}-{self.dropout}-{self.spad}-{self.skpad}-{self.dpad}-{self.dvpad}' + hdim : str + dtype : str # data type + mode : str # value from MODE_MAP + bm0 : int # tile size along q seqlen (block size) + bn0 : int # tile size along k seqlen + bhdq : int # q head_dim + bhdv : int # v head_dim + mask : str + bias : str + dbias : str + dropout : str + spad : str + skpad : str + dpad : str + dvpad : str + deterministic : str def scheck(self, spad1 : str) -> str: if self.mode == 'group': @@ -212,9 +229,9 @@ class FmhaBwdDQDKDVApiTrait: elif self.spad == 't' and spad1 == 't': return f'a.seqlen_q % {self.bm0} != 0' elif self.spad == 'f' and spad1 == 't': - return f'a.seqlen_q % {self.bm0} == 0 and a.seqlen_q % 256 != 0' # BlockSize + return f'a.seqlen_q % {self.bm0} == 0 and a.seqlen_q % 64 != 0' else: # self.skpad == 'f' and skpad1 == 'f' - return f'a.seqlen_q % 256 == 0' # BlockSize + return f'a.seqlen_q % 64 == 0' @property def skcheck(self) -> str: @@ -256,16 +273,19 @@ class FmhaBwdApiPool: per_hdim_case=str() for j, hdim in enumerate(self.dq_dk_dv_pool[dtype].keys()): traits=self.dq_dk_dv_pool[dtype][hdim] + hdim_int = int(hdim) inners=str() for k, trait in enumerate(traits): if_k = 'if' if k == 0 else 'else if' for spad1 in ["t", "f"]: - if ((spad1 == "f" and trait.spad == "t") or (trait.mode == "group" and spad1 == "f")): + if (spad1 == "f" and (trait.spad == "t" or trait.mode == "group")): continue - inners = inners + FMHA_BWD_API_INNER_DISPATCH.format(F_if=if_k, F_mode=MODE_MAP[trait.mode], F_mask=get_mask_map(self.mask_impl)[trait.mask], F_pipeline_enum=BWD_DQDKDV_PIPELINE_ENUM_MAP[trait.pipeline], - F_mask_check=get_mask_check_map(self.mask_impl)[trait.mask], F_bias_check=BIAS_CHECK_MAP[trait.bias], F_bias=BIAS_MAP[trait.bias], F_dbias=BOOL_MAP[trait.dbias], F_dropout=BOOL_MAP[trait.dropout], + inners = inners + FMHA_BWD_API_INNER_DISPATCH.format(F_if=if_k, F_mode=MODE_MAP[trait.mode], F_pipeline_enum=BWD_DQDKDV_PIPELINE_ENUM_MAP[trait.pipeline], + F_mask_check=get_mask_check_map(self.mask_impl)[trait.mask], F_mask=get_mask_map(self.mask_impl)[trait.mask], F_bias_check=BIAS_CHECK_MAP[trait.bias], + F_bias=BIAS_MAP[trait.bias], F_dbias=BOOL_MAP[trait.dbias], F_dropout_check=DROPOUT_CHECK_MAP[trait.dropout], F_dropout=DROPOUT_MAP[trait.dropout], F_scheck=trait.scheck(spad1=spad1), F_skcheck=trait.skcheck, F_dcheck=trait.dcheck, F_dvcheck=trait.dvcheck, F_hdim=hdim, F_dtype=DTYPE_MAP[dtype], - F_spad0=BOOL_MAP[trait.spad], F_spad1=BOOL_MAP[spad1], F_skpad=BOOL_MAP[trait.skpad], F_dpad=BOOL_MAP[trait.dpad], F_dvpad=BOOL_MAP[trait.dvpad]) + F_spad0=BOOL_MAP[trait.spad], F_spad1=BOOL_MAP[spad1], F_skpad=BOOL_MAP[trait.skpad], F_dpad=BOOL_MAP[trait.dpad], F_dvpad=BOOL_MAP[trait.dvpad], + F_deterministic=BOOL_MAP[trait.deterministic]) if_j = 'if' if j == 0 else 'else if' per_hdim_case = per_hdim_case + FMHA_BWD_API_PER_HDIM_CASE.format(F_if=if_j, F_hdim=hdim, F_inner_dispatch=inners) @@ -295,81 +315,89 @@ class FmhaBwdDQDKDVTileSize: F_bhdv : int # v head_dim F_rm0 : int # number of warps along q seqlen (block warps) in gemm0/gemm2 F_rn0 : int # number of warps along k seqlen (block warps) in gemm0/gemm2 - F_rk0 : int # number of warps along gemm-k (not used) in gemm0/gemm2 + F_rk0 : int # number of warps along headdim_qk/v (not used) in gemm0/gemm2 F_rm1 : int # number of warps along k seqlen (block warps) in gemm1/gemm3 - F_rn1 : int # number of warps along q seqlen (block warps) in gemm1/gemm3 - F_rk1 : int # number of warps along gemm-k (not used) in gemm1/gemm3 - F_rm2 : int # number of warps along k seqlen (block warps) in gemm4 - F_rn2 : int # number of warps along q seqlen (block warps) in gemm4 - F_rk2 : int # number of warps along gemm-k (not used) in gemm4 - F_wm : int # warp size along m (warp size) - F_wn : int # warp size along n - F_wk : int # warp size along k + F_rn1 : int # number of warps along headdim_qk/v (block warps) in gemm1/gemm3 + F_rk1 : int # number of warps along q seqlen (not used) in gemm1/gemm3 + F_rm2 : int # number of warps along q seqlen (block warps) in gemm4 + F_rn2 : int # number of warps along headdim_qk (block warps) in gemm4 + F_rk2 : int # number of warps along k seqlen (not used) in gemm4 + F_wm0 : int # warp size along m in gemm0/gemm2/gemm4 + F_wn0 : int # warp size along n in gemm0/gemm2/gemm4 + F_wk0 : int # warp size along k in gemm0/gemm2/gemm4 + F_wm1 : int # warp size along m in gemm1/gemm3 + F_wn1 : int # warp size along n in gemm1/gemm3 + F_wk1 : int # warp size along k in gemm1/gemm3 F_occupancy : int # occupancy @property def name(self) -> str: return f"b{self.F_bm0}x{self.F_bn0}x{self.F_bk0}x{self.F_bk1}x{self.F_bk2}x{self.F_bk3}x{self.F_bk4}x{self.F_bhdq}x{self.F_bhdv}" +\ 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}_r{self.F_rm2}x{self.F_rn2}x{self.F_rk2}" +\ - f"_w{self.F_wm}x{self.F_wn}x{self.F_wk}_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}_o{self.F_occupancy}" @dataclass class FmhaBwdDQDKDVKernel: - F_idx : int # this is not a tunable, but a counter to differentiate symbol - F_hdim : int # hdim - F_dtype : str # data type - F_tile : FmhaBwdDQDKDVTileSize - F_spad : str # true/false - F_skpad : str # - F_dpad : str # - F_dvpad : str # - F_bias : str # - F_dbias : str # - F_dropout : str # - F_mask : str # value from MASK_MAP - F_mode : str # value from MODE_MAP - F_pipeline : str - mask_impl : str + F_idx : int # this is not a tunable, but a counter to differentiate symbol + F_hdim : int # hdim + F_dtype : str # data type + F_tile : FmhaBwdDQDKDVTileSize + F_spad : str # true/false + F_skpad : str # + F_dpad : str # + F_dvpad : str # + F_bias : str # + F_dbias : str # + F_dropout : str # + F_mask : str # value from MASK_MAP + F_mode : str # value from MODE_MAP + F_deterministic : str # + F_pipeline : str # + mask_impl : str # @property def template(self) -> str: return FMHA_BWD_KERNEL_HEADER + \ FMHA_BWD_DQ_DK_DV_KERNEL_BODY.format( - F_idx = self.F_idx, - F_hdim = self.F_hdim, - F_dtype = DTYPE_MAP[self.F_dtype], - F_bm0 = self.F_tile.F_bm0, - F_bn0 = self.F_tile.F_bn0, - F_bk0 = self.F_tile.F_bk0, - F_bk1 = self.F_tile.F_bk1, - F_bk2 = self.F_tile.F_bk2, - F_bk3 = self.F_tile.F_bk3, - F_bk4 = self.F_tile.F_bk4, - F_bhdq = self.F_tile.F_bhdq, - F_bhdv = self.F_tile.F_bhdv, - F_rm0 = self.F_tile.F_rm0, - F_rn0 = self.F_tile.F_rn0, - F_rk0 = self.F_tile.F_rk0, - F_rm1 = self.F_tile.F_rm1, - F_rn1 = self.F_tile.F_rn1, - F_rk1 = self.F_tile.F_rk1, - F_rm2 = self.F_tile.F_rm2, - F_rn2 = self.F_tile.F_rn2, - F_rk2 = self.F_tile.F_rk2, - F_wm = self.F_tile.F_wm, - F_wn = self.F_tile.F_wn, - F_wk = self.F_tile.F_wk, - F_spad = BOOL_MAP[self.F_spad], - F_skpad = BOOL_MAP[self.F_skpad], - F_dpad = BOOL_MAP[self.F_dpad], - F_dvpad = BOOL_MAP[self.F_dvpad], - F_bias = BIAS_MAP[self.F_bias], - F_dbias = BOOL_MAP[self.F_dbias], - F_dropout = BOOL_MAP[self.F_dropout], - F_occupancy = self.F_tile.F_occupancy, - F_mask = get_mask_map(self.mask_impl)[self.F_mask], - F_mode = MODE_MAP[self.F_mode], + F_idx = self.F_idx, + F_hdim = self.F_hdim, + F_dtype = DTYPE_MAP[self.F_dtype], + F_bm0 = self.F_tile.F_bm0, + F_bn0 = self.F_tile.F_bn0, + F_bk0 = self.F_tile.F_bk0, + F_bk1 = self.F_tile.F_bk1, + F_bk2 = self.F_tile.F_bk2, + F_bk3 = self.F_tile.F_bk3, + F_bk4 = self.F_tile.F_bk4, + F_bhdq = self.F_tile.F_bhdq, + F_bhdv = self.F_tile.F_bhdv, + F_rm0 = self.F_tile.F_rm0, + F_rn0 = self.F_tile.F_rn0, + F_rk0 = self.F_tile.F_rk0, + F_rm1 = self.F_tile.F_rm1, + F_rn1 = self.F_tile.F_rn1, + F_rk1 = self.F_tile.F_rk1, + F_rm2 = self.F_tile.F_rm2, + F_rn2 = self.F_tile.F_rn2, + F_rk2 = self.F_tile.F_rk2, + 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_spad = BOOL_MAP[self.F_spad], + F_skpad = BOOL_MAP[self.F_skpad], + F_dpad = BOOL_MAP[self.F_dpad], + F_dvpad = BOOL_MAP[self.F_dvpad], + F_bias = BIAS_MAP[self.F_bias], + F_dbias = BOOL_MAP[self.F_dbias], + F_dropout = DROPOUT_MAP[self.F_dropout], + F_occupancy = self.F_tile.F_occupancy, + F_mask = get_mask_map(self.mask_impl)[self.F_mask], + F_mode = MODE_MAP[self.F_mode], + F_deterministic = BOOL_MAP[self.F_deterministic], F_pipeline_enum = BWD_DQDKDV_PIPELINE_ENUM_MAP[self.F_pipeline], - F_pipeline = BWD_DQDKDV_PIPELINE_MAP[self.F_pipeline]) + F_pipeline = BWD_DQDKDV_PIPELINE_MAP[self.F_pipeline]) @property def name(self) -> str: @@ -382,7 +410,7 @@ class FmhaBwdDQDKDVKernel: if n != '' : n = 'p' + n return n pn = pad_name() - n = f"fmha_bwd_d{self.F_hdim}_{self.F_dtype}_{self.F_mode}_" + self.F_tile.name + n = f"fmha_bwd_d{self.F_hdim}_{self.F_dtype}_{self.F_mode}_" + self.F_tile.name + f'_{self.F_pipeline}' if pn != '' : n += f'_{pn}' if self.F_bias != 'no' : n += f'_{self.F_bias}' if self.F_dbias == 't' : n += '_dbias' @@ -390,7 +418,8 @@ class FmhaBwdDQDKDVKernel: if self.F_mask == 's_mask': n += f'_mask' else: if self.F_mask != 'no' : n += f'_m{self.F_mask[0]}' - if self.F_dropout == 't' : n += '_dropout' + if self.F_dropout != 'no' : n += f'_{self.F_dropout}' + if self.F_deterministic == 't' : n += '_deterministic' return n @property @@ -413,19 +442,23 @@ class FmhaBwdDQDKDVKernel: spad=self.F_spad, skpad=self.F_skpad, dpad=self.F_dpad, - dvpad=self.F_dvpad) + dvpad=self.F_dvpad, + deterministic=self.F_deterministic + ) # TODO: design a more practical way to do it # this is current supported tile size & pipeline. def get_fmha_bwd_dq_dk_dv_tile_ppl_dict_from_dtype(dtype : str) -> Optional[dict]: if dtype == 'fp16' or dtype == 'bf16': return { - '32' : [FmhaBwdDQDKDVTileSize(128, 128, 32, 32, 32, 32, 32, 32, 32, 1, 4, 1, 4, 1, 1, 4, 1, 1, 32, 32, 16, 1), - "qs_ks_vr_dos"], - '64' : [FmhaBwdDQDKDVTileSize( 64, 128, 32, 32, 32, 32, 32, 64, 64, 1, 4, 1, 4, 1, 1, 2, 2, 1, 32, 32, 16, 1), - "qs_ks_vr_dos"], - '128' : [FmhaBwdDQDKDVTileSize( 64, 128, 32, 32, 32, 32, 32, 128, 128, 1, 4, 1, 4, 1, 1, 2, 2, 1, 32, 32, 16, 1), - "ks_vr"] + '32' : [FmhaBwdDQDKDVTileSize( 32, 128, 32, 32, 32, 32, 64, 32, 32, 1, 4, 1, 4, 1, 1, 2, 2, 1, 16, 16, 32, 16, 16, 16, 1), + "kr_ktr_vr_iglp", "kr_ktr_vr"], + '64' : [FmhaBwdDQDKDVTileSize( 32, 128, 64, 32, 64, 32, 32, 64, 64, 1, 4, 1, 4, 1, 1, 1, 4, 1, 16, 16, 32, 16, 16, 16, 1), + "kr_ktr_vr_iglp", "kr_ktr_vr"], + '128' : [FmhaBwdDQDKDVTileSize( 16, 128, 128, 16, 128, 16, 32, 128, 128, 1, 4, 1, 4, 1, 1, 1, 4, 1, 16, 16, 32, 16, 16, 16, 1), + "kr_ktr_vr_iglp", "kr_ktr_vr"], + '256' : [FmhaBwdDQDKDVTileSize( 16, 64, 256, 16, 256, 16, 32, 256, 256, 1, 4, 1, 4, 1, 1, 1, 4, 1, 16, 16, 32, 16, 16, 16, 1), + "kr_ktr_vr_iglp", "kr_ktr_vr"] } else: return None @@ -440,7 +473,7 @@ def get_bwd_dq_dk_dv_blobs(kernel_filter : Optional[str], receipt, mask_impl) -> d = get_fmha_bwd_dq_dk_dv_tile_ppl_dict_from_dtype(dtype) if d == None: continue - for hdim_str, mode, mask, bias, dbias, dropout, spad, skpad, dpad, dvpad in itertools.product(d.keys(), MODE_MAP.keys(), get_mask_map(mask_impl).keys(), BIAS_MAP.keys(), ["t", "f"], ["t", "f"], ["t", "f"], ["t", "f"], ["t", "f"], ["t", "f"]): + for hdim_str, mode, mask, bias, dbias, dropout, spad, skpad, dpad, dvpad, deterministic in itertools.product(d.keys(), MODE_MAP.keys(), get_mask_map(mask_impl).keys(), BIAS_MAP.keys(), ["t", "f"], DROPOUT_MAP.keys(), ["t", "f"], ["t", "f"], ["t", "f"], ["t", "f"], ["t", "f"]): tile = d[hdim_str][0] ppl = d[hdim_str][1] hdim = int(hdim_str) @@ -448,16 +481,29 @@ def get_bwd_dq_dk_dv_blobs(kernel_filter : Optional[str], receipt, mask_impl) -> continue if ((bias == "no" or bias == "alibi") and dbias == "t"): continue + if ("wg32" in dropout): + continue + if (dpad == "t" or dvpad == "t"): + ppl = d[hdim_str][2] k = FmhaBwdDQDKDVKernel(F_idx=0, F_hdim=hdim, F_dtype=dtype, F_tile=tile, F_spad=spad, F_skpad=skpad, F_dpad=dpad, F_dvpad=dvpad, F_bias=bias, F_dbias=dbias, F_dropout=dropout, F_mask=mask, F_mode=mode, - F_pipeline=ppl, mask_impl=mask_impl) + F_pipeline=ppl, mask_impl=mask_impl, F_deterministic=deterministic) if kernel_filter != None: if not fnmatch.fnmatch(k.name, kernel_filter): continue if receipt == 2: cond = dtype in ['fp16', 'bf16'] cond &= bias in ['no', 'alibi'] + cond &= dropout in ['no', 'dropout_wg32', 'dropout_wg16'] + cond &= dpad == dvpad + if not cond: + continue + if receipt == 3: + cond = dtype in ['fp16', 'bf16'] + cond &= bias in ['no', 'alibi'] + cond &= dpad == dvpad + cond &= deterministic == "f" if not cond: continue api_pool.register_dq_dk_dv_traits(k.api_trait()) @@ -468,53 +514,54 @@ def get_bwd_dq_dk_dv_blobs(kernel_filter : Optional[str], receipt, mask_impl) -> FMHA_BWD_DOT_DO_O_KERNEL_BODY=""" using fmha_dtype_{F_idx} = {F_dtype}; -using fmha_bwd_dot_do_o_trait_{F_idx} = ck_tile::TileFmhaBwdOGradDotOTraits<{F_spad}, - {F_dvpad}, - {F_occupancy}>; +using fmha_bwd_dot_do_o_trait_{F_idx} = + ck_tile::TileFmhaBwdOGradDotOTraits<{F_spad}, {F_dvpad}, {F_occupancy}>; using fmha_bwd_dot_do_o_pipeline_problem_{F_idx} = ck_tile::BlockFmhaBwdOGradDotOPipelineProblem< typename FmhaBwdTypeConfig::ODataType, typename FmhaBwdTypeConfig::OGradDataType, typename FmhaBwdTypeConfig::DDataType, - /* BlockSize = */ 256, + /* BlockSize = */ 64, {F_hdim}, {F_mode}, fmha_bwd_dot_do_o_trait_{F_idx}>; -using fmha_bwd_dot_do_o_{F_idx} = typename ck_tile::BlockFmhaBwdOGradDotO< - fmha_bwd_dot_do_o_pipeline_problem_{F_idx}>; +using fmha_bwd_dot_do_o_{F_idx} = + typename ck_tile::BlockFmhaBwdOGradDotO; using fmha_bwd_dot_do_o_kernel_{F_idx} = - ck_tile::FmhaBwdOGradDotOKernel, - fmha_bwd_dot_do_o_{F_idx}>; + ck_tile::FmhaBwdOGradDotOKernel; -using dot_do_o_trait_{F_idx} = fmha_bwd_dot_do_o_traits_<{F_hdim}, {F_dtype}, {F_mode}, {F_spad}, {F_dvpad}>; +using dot_do_o_trait_{F_idx} = + fmha_bwd_dot_do_o_traits_<{F_hdim}, {F_dtype}, {F_mode}, {F_spad}, {F_dvpad}>; #include -template<> +template <> float fmha_bwd_dot_do_o_(const ck_tile::stream_config& s, fmha_bwd_args a) {{ using k_ = fmha_bwd_dot_do_o_kernel_{F_idx}; if(s.log_level_ > 0) std::cout << ", " << k_::GetName() << std::flush; - auto [kargs, grids] = fmha_bwd_dot_do_o_create_kargs_and_grids(a); - constexpr dim3 blocks = k_::BlockSize(); + auto [kargs, grids] = fmha_bwd_dot_do_o_create_kargs_and_grids(a); + constexpr dim3 blocks = k_::BlockSize(); constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; - return ck_tile::launch_kernel(s, ck_tile::make_kernel(k_{{}}, grids, blocks, 0, kargs)); + return ck_tile::launch_kernel( + s, ck_tile::make_kernel(k_{{}}, grids, blocks, 0, kargs)); }} -template<> +template <> void fmha_bwd_dot_do_o_oneshot_(const ck_tile::stream_config& s, fmha_bwd_args a) {{ - using k_ = fmha_bwd_dot_do_o_kernel_{F_idx}; - auto [kargs, grids] = fmha_bwd_dot_do_o_create_kargs_and_grids(a); - constexpr dim3 blocks = k_::BlockSize(); + using k_ = fmha_bwd_dot_do_o_kernel_{F_idx}; + auto [kargs, grids] = fmha_bwd_dot_do_o_create_kargs_and_grids(a); + constexpr dim3 blocks = k_::BlockSize(); constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; - ck_tile::make_kernel(k_{{}}, grids, blocks, 0, kargs)(ck_tile::stream_config{{s.stream_id_}}); + ck_tile::make_kernel(k_{{}}, grids, blocks, 0, kargs)( + ck_tile::stream_config{{s.stream_id_}}); }} -template<> +template <> std::string fmha_bwd_dot_do_o_get_name_() {{ using k_ = fmha_bwd_dot_do_o_kernel_{F_idx}; @@ -584,12 +631,150 @@ def get_bwd_dot_do_o_blobs() -> List[FmhaBwdOGradDotOKernel]: return gen +FMHA_BWD_CONVERT_DQ_KERNEL_BODY=""" +using fmha_dtype_{F_idx} = {F_dtype}; + +using fmha_bwd_convert_dq_trait_{F_idx} = + ck_tile::TileFmhaBwdConvertQGradTraits<{F_spad}, {F_dpad}, {F_occupancy}>; + +using fmha_bwd_convert_dq_pipeline_problem_{F_idx} = + ck_tile::BlockFmhaBwdConvertQGradPipelineProblem< + typename FmhaBwdTypeConfig::AccDataType, + typename FmhaBwdTypeConfig::QGradDataType, + /* BlockSize = */ 256, + {F_bm0}, + {F_bn0}, + {F_hdim}, + {F_mode}, + {F_deterministic}, + fmha_bwd_convert_dq_trait_{F_idx}>; + +using fmha_bwd_convert_dq_{F_idx} = + typename ck_tile::BlockFmhaBwdConvertQGrad; + +using fmha_bwd_convert_dq_kernel_{F_idx} = + ck_tile::FmhaBwdConvertQGradKernel; + +using convert_dq_trait_{F_idx} = fmha_bwd_convert_dq_traits_<{F_hdim}, + {F_dtype}, + {F_mode}, + {F_spad}, + {F_dpad}, + {F_deterministic}>; + +#include + +template <> +float fmha_bwd_convert_dq_(const ck_tile::stream_config& s, fmha_bwd_args a) +{{ + using k_ = fmha_bwd_convert_dq_kernel_{F_idx}; + if(s.log_level_ > 0) + std::cout << ", " << k_::GetName() << std::flush; + auto [kargs, grids] = fmha_bwd_convert_dq_create_kargs_and_grids(a); + constexpr dim3 blocks = k_::BlockSize(); + constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; + return ck_tile::launch_kernel( + s, ck_tile::make_kernel(k_{{}}, grids, blocks, 0, kargs)); +}} + +template <> +void fmha_bwd_convert_dq_oneshot_(const ck_tile::stream_config& s, + fmha_bwd_args a) +{{ + using k_ = fmha_bwd_convert_dq_kernel_{F_idx}; + auto [kargs, grids] = fmha_bwd_convert_dq_create_kargs_and_grids(a); + constexpr dim3 blocks = k_::BlockSize(); + constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; + ck_tile::make_kernel(k_{{}}, grids, blocks, 0, kargs)( + ck_tile::stream_config{{s.stream_id_}}); +}} + +template <> +std::string fmha_bwd_convert_dq_get_name_() +{{ + using k_ = fmha_bwd_convert_dq_kernel_{F_idx}; + return k_::GetName(); +}} +""" + +@dataclass +class FmhaBwdConvertQGradKernel: + F_idx : int # this is not a tunable, but a counter to differentiate symbol + F_hdim : int # hdim + F_dtype : str # data type + F_bm0 : int # tile size along q seqlen (block size) + F_bn0 : int # tile size along k seqlen + F_spad : str # true/false + F_dpad : str # + F_mode : str # value from MODE_MAP + F_occupancy : int # + F_deterministic : str # + + @property + def template(self) -> str: + return FMHA_BWD_KERNEL_HEADER + \ + FMHA_BWD_CONVERT_DQ_KERNEL_BODY.format( + F_idx = self.F_idx, + F_hdim = self.F_hdim, + F_dtype = DTYPE_MAP[self.F_dtype], + F_bm0 = self.F_bm0, + F_bn0 = self.F_bn0, + F_spad = BOOL_MAP[self.F_spad], + F_dpad = BOOL_MAP[self.F_dpad], + F_mode = MODE_MAP[self.F_mode], + F_occupancy = self.F_occupancy, + F_deterministic = BOOL_MAP[self.F_deterministic]) + + @property + def name(self) -> str: + def pad_name() -> str: + n = '' + if self.F_spad == 't': n += 's' + if self.F_dpad == 't' : n += 'd' + if n != '' : n = 'p' + n + return n + pn = pad_name() + n = f"fmha_bwd_convert_dq_d{self.F_hdim}_{self.F_dtype}_b{self.F_bm0}x{self.F_bn0}_{self.F_mode}_o{self.F_occupancy}" + if pn != '' : n += f'_{pn}' + if self.F_deterministic == 't' : n += f'_deterministic' + return n + + @property + def filename(self) -> str: + return self.name + ".cpp" + +def get_bwd_convert_dq_blobs() -> List[FmhaBwdConvertQGradKernel]: + # TODO: we don't support tuning yet, so pick up one value for pad/occupancy + # support this in future + def get_occupancy(dtype, hdim): + return 2 + + gen = list() + + for dtype in DTYPE_MAP.keys(): + d = get_fmha_bwd_dq_dk_dv_tile_ppl_dict_from_dtype(dtype) + if d == None: + continue + for hdim_str, mode, spad, dpad, deterministic in itertools.product(d.keys(), MODE_MAP.keys(), ["t", "f"], ["t", "f"], ["t", "f"]): + hdim = int(hdim_str) + tile = d[hdim_str][0] + if (mode == "group" and spad == "f"): + continue + k = FmhaBwdConvertQGradKernel(F_idx=0, F_hdim=hdim, F_dtype=dtype, F_bm0=64, F_bn0=tile.F_bn0, + F_spad=spad, F_dpad=dpad, F_mode=mode, F_occupancy=get_occupancy(dtype, hdim), F_deterministic=deterministic) + gen.append(k) + + return gen + def write_single_bwd_dq_dk_dv_kernel(kernel: FmhaBwdDQDKDVKernel, autogen_dir: Path) -> None: (autogen_dir / kernel.filename).write_text(kernel.template) def write_single_bwd_dot_do_o_kernel(kernel: FmhaBwdOGradDotOKernel, autogen_dir: Path) -> None: (autogen_dir / kernel.filename).write_text(kernel.template) +def write_single_bwd_convert_dq_kernel(kernel: FmhaBwdConvertQGradKernel, autogen_dir: Path) -> None: + (autogen_dir / kernel.filename).write_text(kernel.template) + def write_bwd_api(api_pool : FmhaBwdApiPool, autogen_dir: Path) -> None: (autogen_dir / FMHA_BWD_API_FILENAME).write_text(api_pool.api) @@ -597,6 +782,9 @@ def write_blobs(output_dir : Path, kernel_filter : Optional[str], receipt, mask_ kernels = get_bwd_dot_do_o_blobs() for kernel in kernels: write_single_bwd_dot_do_o_kernel(kernel, output_dir) + kernels = get_bwd_convert_dq_blobs() + for kernel in kernels: + write_single_bwd_convert_dq_kernel(kernel, output_dir) api_pool, kernels = get_bwd_dq_dk_dv_blobs(kernel_filter, receipt, mask_impl) for kernel in kernels: write_single_bwd_dq_dk_dv_kernel(kernel, output_dir) @@ -605,6 +793,9 @@ def write_blobs(output_dir : Path, kernel_filter : Optional[str], receipt, mask_ def list_blobs(file_path : Path, kernel_filter : Optional[str], receipt, mask_impl) -> None: with file_path.open('a') as f: kernels = get_bwd_dot_do_o_blobs() + for kernel in kernels: + f.write(str(file_path.parent / GEN_DIR / kernel.filename) + "\n") + kernels = get_bwd_convert_dq_blobs() for kernel in kernels: f.write(str(file_path.parent / GEN_DIR / kernel.filename) + "\n") _, kernels = get_bwd_dq_dk_dv_blobs(kernel_filter, receipt, mask_impl) diff --git a/example/ck_tile/01_fmha/fmha_bwd.cpp b/example/ck_tile/01_fmha/fmha_bwd.cpp index b1249b5eda..efae4e284a 100644 --- a/example/ck_tile/01_fmha/fmha_bwd.cpp +++ b/example/ck_tile/01_fmha/fmha_bwd.cpp @@ -87,7 +87,11 @@ auto create_args(int argc, char* argv[]) .insert("drop_offset", "0", "offset for random number generator") .insert("timer", "gpu", "gpu:gpu timer, cpu:cpu timer") .insert("warmup", "5", "number of iterations before benchmark the kernel") - .insert("repeat", "20", "number of iterations to benchmark the kernel"); + .insert("repeat", "20", "number of iterations to benchmark the kernel") + .insert("deterministic", + "0", + "if set to 1 will use multi-buffer reduction strategy for dq, atomic opeartion " + "will not be used"); bool result = arg_parser.parse(argc, argv); return std::make_tuple(result, arg_parser); @@ -128,11 +132,6 @@ bool run(const ck_tile::ArgParser& arg_parser) ck_tile::index_t hdim_v = arg_parser.get_int("d_v"); if(hdim_v < 0) hdim_v = hdim_q; - if(hdim_q % 2 != 0 || hdim_v % 2 != 0) - { - std::cerr << "FMHA Bwd kernel currently only supports even headdim" << std::endl; - return false; - } bool i_perm = arg_parser.get_bool("iperm"); // if true, will be batch * nhead * seqlen * hdim bool o_perm = arg_parser.get_bool("operm"); // if false, will be batch * seqlen * nhead * hdim @@ -177,9 +176,10 @@ bool run(const ck_tile::ArgParser& arg_parser) seed.reset(); } - int stream_warmup = arg_parser.get_int("warmup"); - int stream_repeat = arg_parser.get_int("repeat"); - bool kname = arg_parser.get_bool("kname"); + int stream_warmup = arg_parser.get_int("warmup"); + int stream_repeat = arg_parser.get_int("repeat"); + bool kname = arg_parser.get_bool("kname"); + bool deterministic = arg_parser.get_bool("deterministic"); ck_tile::stream_config stream_config{nullptr, true, @@ -265,6 +265,9 @@ bool run(const ck_tile::ArgParser& arg_parser) (mode == mode_enum::batch ? seqlen_q : seqstart_q_host.back()); const ck_tile::index_t shape_seqlen_k = (mode == mode_enum::batch ? seqlen_k : seqstart_k_host.back()); + const ck_tile::index_t kN0 = (hdim_q <= 128) ? 128 : 64; + const ck_tile::index_t nsplits = + deterministic ? ck_tile::integer_divide_ceil(max_seqlen_k, kN0) : 1; ck_tile::HostTensor q_host( get_lengths(i_perm, shape_batch, nhead, shape_seqlen_q, hdim_q)); @@ -284,9 +287,9 @@ bool run(const ck_tile::ArgParser& arg_parser) ck_tile::HostTensor o_host( get_lengths(o_perm, shape_batch, nhead, shape_seqlen_q, hdim_v)); ck_tile::HostTensor lse_host( - std::array{batch, nhead, max_seqlen_q}); + std::array{shape_batch, nhead, shape_seqlen_q}); ck_tile::HostTensor d_host( - std::array{batch, nhead, max_seqlen_q}); + std::array{shape_batch, nhead, shape_seqlen_q}); ck_tile::HostTensor randval_host( p_drop > 0 ? get_lengths(true, shape_batch, nhead, shape_seqlen_q, max_seqlen_k) : std::array{1, 1, 1, 1}); @@ -302,6 +305,10 @@ bool run(const ck_tile::ArgParser& arg_parser) use_dbias ? get_lengths(i_perm, shape_batch, nhead, shape_seqlen_q, max_seqlen_k) : std::array{1, 1, 1, 1} /* dummy shape for simplifying code */); + ck_tile::HostTensor dq_acc_host( + i_perm + ? std::array{nsplits, shape_batch, nhead, shape_seqlen_q, hdim_q} + : std::array{nsplits, shape_batch, shape_seqlen_q, nhead, hdim_q}); if(init_method == 0) { @@ -362,6 +369,7 @@ bool run(const ck_tile::ArgParser& arg_parser) ck_tile::DeviceMem seqstart_q(seqstart_q_host.size() * sizeof(int32_t)); ck_tile::DeviceMem seqstart_k(seqstart_k_host.size() * sizeof(int32_t)); ck_tile::DeviceMem alibi_slope_buf(alibi_slope_host.get_element_space_size_in_bytes()); + ck_tile::DeviceMem dq_acc_buf(dq_acc_host.get_element_space_size_in_bytes()); q_buf.ToDevice(q_host.data()); k_buf.ToDevice(k_host.data()); @@ -387,8 +395,17 @@ bool run(const ck_tile::ArgParser& arg_parser) std::cout << "[" << prec << "|" << mode << "|" << io_layout(i_perm, o_perm) << "] b:" << batch << ", h:" << nhead << "/" << nhead_k << ", s:" << seqlen_q << "/" << seqlen_k << ", d:" << hdim_q << "/" << hdim_v << ", scale:" << scale << ", bias:" << bias - << ", dbias:" << use_dbias << ", p_drop:" << p_drop << ", mask:" << mask - << std::flush; + << ", dbias:" << use_dbias << ", p_drop:" << p_drop << ", s_randval:" << s_randval + << ", deterministic:" << deterministic << ", mask:" << mask << std::flush; + + std::size_t workspace_size = + dq_acc_host.get_element_space_size_in_bytes() * sizeof(AccDataType) / (1024 * 1024); + + if(deterministic == 1) + { + std::cout << "\nDeterministic mode ON: " << workspace_size + << " MByte memory workspace allocated" << std::endl; + } auto fmha_traits = fmha_bwd_traits{hdim_q, hdim_v, @@ -397,7 +414,9 @@ bool run(const ck_tile::ArgParser& arg_parser) mask.type, bias.type, use_dbias, - p_drop > 0.0f}; + p_drop > 0.0f, + s_randval, + deterministic}; auto fmha_args = [&]() { assert(nhead % nhead_k == 0); /// NOTE: we broadcast bias from [1, 1, seqlen_q, seqlen_k] to [batch, nhead, seqlen_q, @@ -422,7 +441,7 @@ bool run(const ck_tile::ArgParser& arg_parser) const ck_tile::index_t nhead_stride_o = (o_perm ? shape_seqlen_q * hdim_v : hdim_v); const ck_tile::index_t nhead_stride_randval = (shape_seqlen_q * max_seqlen_k); const ck_tile::index_t nhead_stride_do = (o_perm ? shape_seqlen_q * hdim_v : hdim_v); - const ck_tile::index_t nhead_stride_lsed = max_seqlen_q; + const ck_tile::index_t nhead_stride_lsed = shape_seqlen_q; const ck_tile::index_t nhead_stride_dbias = (i_perm ? shape_seqlen_q * max_seqlen_k : max_seqlen_k); // setup batch_stride_* arguments @@ -433,10 +452,12 @@ bool run(const ck_tile::ArgParser& arg_parser) const ck_tile::index_t batch_stride_o = (nhead * shape_seqlen_q * hdim_v); const ck_tile::index_t batch_stride_randval = (nhead * shape_seqlen_q * max_seqlen_k); const ck_tile::index_t batch_stride_do = (nhead * shape_seqlen_q * hdim_v); - const ck_tile::index_t batch_stride_lsed = (nhead * max_seqlen_q); + const ck_tile::index_t batch_stride_lsed = (nhead * shape_seqlen_q); const ck_tile::index_t batch_stride_dk = (nhead * shape_seqlen_k * hdim_q); const ck_tile::index_t batch_stride_dv = (nhead * shape_seqlen_k * hdim_v); const ck_tile::index_t batch_stride_dbias = (nhead * shape_seqlen_q * max_seqlen_k); + const ck_tile::index_t split_stride_dq_acc = + (shape_batch * nhead * shape_seqlen_q * hdim_q); return fmha_bwd_args{q_buf.GetDeviceBuffer(), k_buf.GetDeviceBuffer(), @@ -452,6 +473,7 @@ bool run(const ck_tile::ArgParser& arg_parser) dk_buf.GetDeviceBuffer(), dv_buf.GetDeviceBuffer(), dbias_buf.GetDeviceBuffer(), + dq_acc_buf.GetDeviceBuffer(), seqstart_q.GetDeviceBuffer(), seqstart_k.GetDeviceBuffer(), nullptr, @@ -473,6 +495,8 @@ bool run(const ck_tile::ArgParser& arg_parser) stride_o, stride_randval, stride_do, + stride_q, // stride_dq_acc + stride_q, // stride_dq stride_dk, stride_dv, stride_dbias, @@ -484,6 +508,10 @@ bool run(const ck_tile::ArgParser& arg_parser) nhead_stride_randval, nhead_stride_do, nhead_stride_lsed, + nhead_stride_q, // nhead_stride_dq_acc + nhead_stride_q, // nhead_stride_dq + nhead_stride_k, // nhead_stride_dk + nhead_stride_v, // nhead_stride_dv nhead_stride_dbias, batch_stride_q, batch_stride_k, @@ -493,15 +521,17 @@ bool run(const ck_tile::ArgParser& arg_parser) batch_stride_randval, batch_stride_do, batch_stride_lsed, + batch_stride_q, // batch_stride_dq_acc + batch_stride_q, // batch_stride_dq batch_stride_dk, batch_stride_dv, batch_stride_dbias, + split_stride_dq_acc, mask.left, mask.right, static_cast(mask.type), p_drop, p_undrop, - s_randval, {drop_seed, drop_offset}}; }(); @@ -719,7 +749,7 @@ bool run(const ck_tile::ArgParser& arg_parser) if(o_perm) o_host_ref.ForEach([&](auto& self, auto idx) { o_host(b, idx[0], idx[1] + query_offset, idx[2]) = self(idx); }); else o_host_ref.ForEach([&](auto& self, auto idx) { o_host(b, idx[1] + query_offset, idx[0], idx[2]) = self(idx); }); - lse_host_ref.ForEach([&](auto& self, auto idx) { lse_host(wb, idx[0], idx[1]) = self(idx); }); + lse_host_ref.ForEach([&](auto& self, auto idx) { lse_host(b, idx[0], idx[1] + query_offset) = self(idx); }); // clang-format on q_host_refs.push_back(q_host_ref); @@ -738,6 +768,7 @@ bool run(const ck_tile::ArgParser& arg_parser) lse_buf.ToDevice(lse_host.data()); dq_buf.SetZero(); dbias_buf.SetZero(); + dq_acc_buf.SetZero(); ck_tile::stream_config stream_config_v{ nullptr, true, 0, 0, 1, arg_parser.get_str("timer") == std::string("gpu")}; diff --git a/example/ck_tile/01_fmha/fmha_bwd.hpp b/example/ck_tile/01_fmha/fmha_bwd.hpp index 0c6b468951..aea42515dc 100644 --- a/example/ck_tile/01_fmha/fmha_bwd.hpp +++ b/example/ck_tile/01_fmha/fmha_bwd.hpp @@ -77,6 +77,7 @@ struct fmha_bwd_args void* dk_ptr; void* dv_ptr; void* dbias_ptr; + void* dq_acc_ptr; const void* seqstart_q_ptr; const void* seqstart_k_ptr; const void* seqlen_k_ptr; @@ -97,6 +98,8 @@ struct fmha_bwd_args ck_tile::index_t stride_o; ck_tile::index_t stride_randval; ck_tile::index_t stride_do; + ck_tile::index_t stride_dq_acc; + ck_tile::index_t stride_dq; ck_tile::index_t stride_dk; ck_tile::index_t stride_dv; ck_tile::index_t stride_dbias; @@ -108,6 +111,10 @@ struct fmha_bwd_args ck_tile::index_t nhead_stride_randval; ck_tile::index_t nhead_stride_do; ck_tile::index_t nhead_stride_lsed; + ck_tile::index_t nhead_stride_dq_acc; + ck_tile::index_t nhead_stride_dq; + ck_tile::index_t nhead_stride_dk; + ck_tile::index_t nhead_stride_dv; ck_tile::index_t nhead_stride_dbias; ck_tile::index_t batch_stride_q; ck_tile::index_t batch_stride_k; @@ -117,15 +124,17 @@ struct fmha_bwd_args ck_tile::index_t batch_stride_randval; ck_tile::index_t batch_stride_do; ck_tile::index_t batch_stride_lsed; + ck_tile::index_t batch_stride_dq_acc; + ck_tile::index_t batch_stride_dq; ck_tile::index_t batch_stride_dk; ck_tile::index_t batch_stride_dv; ck_tile::index_t batch_stride_dbias; + ck_tile::index_t split_stride_dq_acc; ck_tile::index_t window_size_left; ck_tile::index_t window_size_right; ck_tile::index_t mask_type; float p_drop; float p_undrop; - bool s_randval; std::tuple drop_seed_offset; }; @@ -145,10 +154,10 @@ auto fmha_bwd_dq_dk_dv_create_kargs_and_grids(fmha_bwd_args args) args.do_ptr, args.d_ptr, args.rand_val_ptr, - args.dq_ptr, args.dk_ptr, args.dv_ptr, args.dbias_ptr, + args.dq_acc_ptr, args.seqstart_q_ptr, args.seqstart_k_ptr, args.seqlen_k_ptr, @@ -163,6 +172,7 @@ auto fmha_bwd_dq_dk_dv_create_kargs_and_grids(fmha_bwd_args args) args.stride_bias, args.stride_randval, args.stride_do, + args.stride_dq_acc, args.stride_dk, args.stride_dv, args.stride_dbias, @@ -173,13 +183,15 @@ auto fmha_bwd_dq_dk_dv_create_kargs_and_grids(fmha_bwd_args args) args.nhead_stride_randval, args.nhead_stride_do, args.nhead_stride_lsed, + args.nhead_stride_dq_acc, + args.nhead_stride_dk, + args.nhead_stride_dv, args.nhead_stride_dbias, - args.batch_stride_lsed, + args.split_stride_dq_acc, args.window_size_left, args.window_size_right, args.mask_type, args.p_drop, - args.s_randval, args.drop_seed_offset); } else @@ -192,10 +204,10 @@ auto fmha_bwd_dq_dk_dv_create_kargs_and_grids(fmha_bwd_args args) args.do_ptr, args.d_ptr, args.rand_val_ptr, - args.dq_ptr, args.dk_ptr, args.dv_ptr, args.dbias_ptr, + args.dq_acc_ptr, args.seqlen_q, args.seqlen_k, args.hdim_q, @@ -209,6 +221,7 @@ auto fmha_bwd_dq_dk_dv_create_kargs_and_grids(fmha_bwd_args args) args.stride_bias, args.stride_randval, args.stride_do, + args.stride_dq_acc, args.stride_dk, args.stride_dv, args.stride_dbias, @@ -219,6 +232,9 @@ auto fmha_bwd_dq_dk_dv_create_kargs_and_grids(fmha_bwd_args args) args.nhead_stride_randval, args.nhead_stride_do, args.nhead_stride_lsed, + args.nhead_stride_dq_acc, + args.nhead_stride_dk, + args.nhead_stride_dv, args.nhead_stride_dbias, args.batch_stride_q, args.batch_stride_k, @@ -227,14 +243,15 @@ auto fmha_bwd_dq_dk_dv_create_kargs_and_grids(fmha_bwd_args args) args.batch_stride_randval, args.batch_stride_do, args.batch_stride_lsed, + args.batch_stride_dq_acc, args.batch_stride_dk, args.batch_stride_dv, args.batch_stride_dbias, + args.split_stride_dq_acc, args.window_size_left, args.window_size_right, args.mask_type, args.p_drop, - args.s_randval, args.drop_seed_offset); } }(); @@ -260,8 +277,7 @@ auto fmha_bwd_dot_do_o_create_kargs_and_grids(fmha_bwd_args args) args.stride_o, args.nhead_stride_do, args.nhead_stride_o, - args.nhead_stride_lsed, - args.batch_stride_lsed); + args.nhead_stride_lsed); } else { // create batch mode kernel arguments @@ -286,19 +302,59 @@ auto fmha_bwd_dot_do_o_create_kargs_and_grids(fmha_bwd_args args) return ck_tile::make_tuple(kargs, grids); } +template +auto fmha_bwd_convert_dq_create_kargs_and_grids(fmha_bwd_args args) +{ + auto kargs = [&] { + // create group mode kernel arguments + if constexpr(FmhaBwdConvertQGradKernel::kIsGroupMode) + { + return FmhaBwdConvertQGradKernel::MakeKargs(args.dq_acc_ptr, + args.dq_ptr, + args.seqstart_q_ptr, + args.seqstart_k_ptr, + args.hdim_q, + args.stride_dq, + args.stride_dq_acc, + args.nhead_stride_dq, + args.nhead_stride_dq_acc, + args.split_stride_dq_acc); + } + else + { // create batch mode kernel arguments + return FmhaBwdConvertQGradKernel::MakeKargs(args.dq_acc_ptr, + args.dq_ptr, + args.seqlen_q, + args.seqlen_k, + args.hdim_q, + args.stride_dq, + args.stride_dq_acc, + args.nhead_stride_dq, + args.nhead_stride_dq_acc, + args.batch_stride_dq, + args.batch_stride_dq_acc, + args.split_stride_dq_acc); + } + }(); + + dim3 grids = FmhaBwdConvertQGradKernel::GridSize(args.batch, args.nhead_q, args.max_seqlen_q); + return ck_tile::make_tuple(kargs, grids); +} + // this is used to pattern-match internl kernel implementation, not to instantiate kernel template + bool kPadDv_, + bool kIsDeterministic_> struct fmha_bwd_dq_dk_dv_traits_ { static constexpr ck_tile::index_t HDim = HDim_; @@ -306,13 +362,14 @@ struct fmha_bwd_dq_dk_dv_traits_ static constexpr bool kIsGroupMode = kIsGroupMode_; static constexpr auto FmhaBwdPipelineEnum = FmhaBwdPipelineEnum_; using FmhaMask = ck_tile::remove_cvref_t; + using FmhaDropout = ck_tile::remove_cvref_t; static constexpr auto BiasEnum = BiasEnum_; static constexpr bool kHasBiasGrad = kHasBiasGrad_; - static constexpr bool kHasDropout = kHasDropout_; static constexpr bool kPadS = kPadS_; static constexpr bool kPadSK = kPadSK_; static constexpr bool kPadD = kPadD_; static constexpr bool kPadDv = kPadDv_; + static constexpr bool kIsDeterministic = kIsDeterministic_; }; template @@ -343,6 +400,31 @@ void fmha_bwd_dot_do_o_oneshot_(const ck_tile::stream_config&, fmha_bwd_args); template std::string fmha_bwd_dot_do_o_get_name_(); +template +struct fmha_bwd_convert_dq_traits_ +{ + static constexpr ck_tile::index_t HDim = HDim_; + using DataType = ck_tile::remove_cvref_t; + static constexpr bool kIsGroupMode = kIsGroupMode_; + static constexpr bool kPadS = kPadS_; + static constexpr bool kPadD = kPadD_; + static constexpr bool kIsDeterministic = kIsDeterministic_; +}; + +template +float fmha_bwd_convert_dq_(const ck_tile::stream_config&, fmha_bwd_args); + +template +void fmha_bwd_convert_dq_oneshot_(const ck_tile::stream_config&, fmha_bwd_args); + +template +std::string fmha_bwd_convert_dq_get_name_(); + // This is the public API, will be generated by script struct fmha_bwd_traits { @@ -354,6 +436,8 @@ struct fmha_bwd_traits bias_enum bias_type; // 0:no bias, 1:elementwise bias, 2:alibi. sync with BlockAttentionBiasEnum bool has_dbias; bool has_dropout; + bool is_store_randval; + bool is_deterministic; // TODO: padding check is inside this api }; float fmha_bwd(fmha_bwd_traits, fmha_bwd_args, const ck_tile::stream_config&); diff --git a/example/ck_tile/01_fmha/fmha_fwd.cpp b/example/ck_tile/01_fmha/fmha_fwd.cpp index 28f7905734..4cc3e77c75 100644 --- a/example/ck_tile/01_fmha/fmha_fwd.cpp +++ b/example/ck_tile/01_fmha/fmha_fwd.cpp @@ -479,16 +479,18 @@ bool run(const ck_tile::ArgParser& arg_parser) : std::array{1, 1}); ck_tile::HostTensor lse_acc_host( - 1 < num_splits ? std::array{num_splits, batch, nhead, max_seqlen_q} - : std::array{1, 1, 1, 1}); + 1 < num_splits + ? std::array{num_splits, shape_batch, nhead, shape_seqlen_q} + : std::array{1, 1, 1, 1}); ck_tile::HostTensor o_acc_host( 1 < num_splits ? std::array{num_splits, batch, nhead, max_seqlen_q, hdim_v} : std::array{1, 1, 1, 1, 1}); - // self define lse data layout as [batch, nhead, max_seqlen_q] + // batch mode of lse data layout is [batch, nhead, seqlen_q] + // group mode of lse data layout is [nhead, total_seqlen_q] ck_tile::HostTensor lse_host( - lse ? std::array{batch, nhead, max_seqlen_q} + lse ? std::array{shape_batch, nhead, shape_seqlen_q} : std::array{1, 1, 1} /* dummy shape for simplifying code */); ck_tile::HostTensor o_host( @@ -669,8 +671,8 @@ bool run(const ck_tile::ArgParser& arg_parser) const ck_tile::index_t nhead_stride_bias = (i_perm ? 0 * shape_seqlen_q * shape_seqlen_k : 0 * shape_seqlen_k); const ck_tile::index_t nhead_stride_randval = (shape_seqlen_q * max_seqlen_k); - const ck_tile::index_t nhead_stride_lse = max_seqlen_q; - const ck_tile::index_t nhead_stride_lse_acc = max_seqlen_q; + const ck_tile::index_t nhead_stride_lse = shape_seqlen_q; + const ck_tile::index_t nhead_stride_lse_acc = shape_seqlen_q; const ck_tile::index_t nhead_stride_o_acc = (max_seqlen_q * hdim_v); const ck_tile::index_t nhead_stride_o = (o_perm ? shape_seqlen_q * hdim_v : hdim_v); // setup batch_stride_* arguments @@ -679,12 +681,12 @@ bool run(const ck_tile::ArgParser& arg_parser) const ck_tile::index_t batch_stride_v = (nhead_k * hdim_v * shape_seqlen_k); const ck_tile::index_t batch_stride_bias = (0 * nhead * shape_seqlen_q * shape_seqlen_k); const ck_tile::index_t batch_stride_randval = (nhead * shape_seqlen_q * max_seqlen_k); - const ck_tile::index_t batch_stride_lse = (nhead * max_seqlen_q); - const ck_tile::index_t batch_stride_lse_acc = (nhead * max_seqlen_q); + const ck_tile::index_t batch_stride_lse = (nhead * shape_seqlen_q); + const ck_tile::index_t batch_stride_lse_acc = (nhead * shape_seqlen_q); const ck_tile::index_t batch_stride_o_acc = (nhead * max_seqlen_q * hdim_v); const ck_tile::index_t batch_stride_o = (nhead * shape_seqlen_q * hdim_v); // setup split_stride_* arguments (only used in split-kv kernel) - const ck_tile::index_t split_stride_lse_acc = (batch * nhead * max_seqlen_q); + const ck_tile::index_t split_stride_lse_acc = (shape_batch * nhead * shape_seqlen_q); const ck_tile::index_t split_stride_o_acc = (batch * nhead * max_seqlen_q * hdim_v); return fmha_fwd_args{q_buf.GetDeviceBuffer(), @@ -996,8 +998,9 @@ bool run(const ck_tile::ArgParser& arg_parser) if(lse) { ck_tile::HostTensor lse_host_result({nhead, real_seqlen_q}); - lse_host_result.ForEach( - [&](auto& self, auto idx) { self(idx) = lse_host(wb, idx[0], idx[1]); }); + lse_host_result.ForEach([&](auto& self, auto idx) { + self(idx) = lse_host(b, idx[0], idx[1] + query_offset); + }); cur_pass = ck_tile::check_err(lse_host_result, lse_host_ref, diff --git a/example/ck_tile/01_fmha/fmha_fwd.hpp b/example/ck_tile/01_fmha/fmha_fwd.hpp index ee932ce5d9..c4c951c43a 100644 --- a/example/ck_tile/01_fmha/fmha_fwd.hpp +++ b/example/ck_tile/01_fmha/fmha_fwd.hpp @@ -185,7 +185,6 @@ auto fmha_fwd_create_kargs_and_grids(fmha_fwd_args args) args.nhead_stride_randval, args.nhead_stride_lse, args.nhead_stride_o, - args.batch_stride_lse, args.window_size_left, args.window_size_right, args.mask_type, @@ -284,7 +283,6 @@ auto fmha_fwd_splitkv_create_kargs_and_grids(fmha_fwd_args args) args.nhead_stride_randval, args.nhead_stride_lse_acc, args.nhead_stride_o_acc, - args.batch_stride_lse_acc, args.batch_stride_o_acc, args.split_stride_lse_acc, args.split_stride_o_acc, @@ -376,9 +374,7 @@ auto fmha_fwd_splitkv_combine_create_kargs_and_grids(fmha_fwd_args args) args.nhead_stride_o_acc, args.nhead_stride_lse, args.nhead_stride_o, - args.batch_stride_lse_acc, args.batch_stride_o_acc, - args.batch_stride_lse, args.split_stride_lse_acc, args.split_stride_o_acc); } diff --git a/example/ck_tile/01_fmha/script/smoke_test_bwd.sh b/example/ck_tile/01_fmha/script/smoke_test_bwd.sh index d6830aa2ec..dbb592820e 100755 --- a/example/ck_tile/01_fmha/script/smoke_test_bwd.sh +++ b/example/ck_tile/01_fmha/script/smoke_test_bwd.sh @@ -11,18 +11,19 @@ COMMON_ARGS='-v=1' set -x for prec in "fp16" "bf16" ; do for perm in 0 1 ; do -for hdim in 32 64 128 ; do +for hdim in 32 64 128 256 ; do for mode in 0 1 ; do -for bias in "n" "e" "a"; do -for dbias in 0 1 ; do -for p_drop in 0.0 0.2; do +for bias in "n" "a" ; do +for dbias in 0 ; do +for p_drop in 0.0 0.2 ; do +for deterministic in 0 ; do -$EXE -prec=$prec -b=1 -h=4 -h_k=2 -d=$hdim -s=259 -bias=$bias -dbias=$dbias -p_drop=$p_drop -iperm=$perm -operm=$perm -v=1 -mode=$mode -kname=$KNAME $COMMON_ARGS -$EXE -prec=$prec -b=2 -h=2 -d=$hdim -s=516 -s_k=253 -bias=$bias -dbias=$dbias -p_drop=$p_drop -iperm=$perm -operm=$perm -v=1 -mode=$mode -kname=$KNAME $COMMON_ARGS -$EXE -prec=$prec -b=1 -h=4 -h_k=1 -d=$hdim -s=500 -s_k=251 -bias=$bias -dbias=$dbias -p_drop=$p_drop -iperm=$perm -operm=$perm -mask=1 -v=1 -mode=$mode -kname=$KNAME $COMMON_ARGS -$EXE -prec=$prec -b=1 -h=2 -d=$hdim -s=900 -s_k=258 -bias=$bias -dbias=$dbias -p_drop=$p_drop -iperm=$perm -operm=$perm -mask=2 -v=1 -mode=$mode -kname=$KNAME $COMMON_ARGS -$EXE -prec=$prec -b=2 -h=1 -d=$hdim -s=987 -s_k=219 -bias=$bias -dbias=$dbias -p_drop=$p_drop -iperm=$perm -operm=$perm -mask=t:128,30 -v=1 -mode=$mode -kname=$KNAME $COMMON_ARGS -$EXE -prec=$prec -b=2 -h=3 -h_k=1 -d=$hdim -s=244 -s_k=499 -bias=$bias -dbias=$dbias -p_drop=$p_drop -iperm=$perm -operm=$perm -mask=b:4,35 -v=1 -mode=$mode -kname=$KNAME $COMMON_ARGS +$EXE -prec=$prec -b=1 -h=4 -h_k=2 -d=$hdim -s=259 -bias=$bias -dbias=$dbias -p_drop=$p_drop -iperm=$perm -operm=$perm -deterministic=$deterministic -v=1 -mode=$mode -kname=$KNAME $COMMON_ARGS +$EXE -prec=$prec -b=2 -h=2 -d=$hdim -s=516 -s_k=253 -bias=$bias -dbias=$dbias -p_drop=$p_drop -iperm=$perm -operm=$perm -deterministic=$deterministic -v=1 -mode=$mode -kname=$KNAME $COMMON_ARGS +$EXE -prec=$prec -b=1 -h=4 -h_k=1 -d=$hdim -s=500 -s_k=251 -bias=$bias -dbias=$dbias -p_drop=$p_drop -iperm=$perm -operm=$perm -mask=1 -deterministic=$deterministic -v=1 -mode=$mode -kname=$KNAME $COMMON_ARGS +$EXE -prec=$prec -b=1 -h=2 -d=$hdim -s=900 -s_k=258 -bias=$bias -dbias=$dbias -p_drop=$p_drop -iperm=$perm -operm=$perm -mask=2 -v=1 -deterministic=$deterministic -mode=$mode -kname=$KNAME $COMMON_ARGS +$EXE -prec=$prec -b=2 -h=1 -d=$hdim -s=987 -s_k=219 -bias=$bias -dbias=$dbias -p_drop=$p_drop -iperm=$perm -operm=$perm -mask=t:128,30 -deterministic=$deterministic -v=1 -mode=$mode -kname=$KNAME $COMMON_ARGS +$EXE -prec=$prec -b=2 -h=3 -h_k=1 -d=$hdim -s=244 -s_k=499 -bias=$bias -dbias=$dbias -p_drop=$p_drop -iperm=$perm -operm=$perm -mask=b:4,35 -deterministic=$deterministic -v=1 -mode=$mode -kname=$KNAME $COMMON_ARGS done done @@ -31,4 +32,5 @@ done done done done +done set +x diff --git a/include/ck_tile/core/algorithm/coordinate_transform.hpp b/include/ck_tile/core/algorithm/coordinate_transform.hpp index 71602e5d13..5c7e489804 100644 --- a/include/ck_tile/core/algorithm/coordinate_transform.hpp +++ b/include/ck_tile/core/algorithm/coordinate_transform.hpp @@ -1341,7 +1341,7 @@ struct modulo : public base_transform<1, 1> }; // 2D XOR, NOTE: "xor" is a keyword -template +template struct xor_t : public base_transform<2, 2> { static constexpr auto type_enum = coord_transform_enum::xor_t; @@ -1352,15 +1352,10 @@ struct xor_t : public base_transform<2, 2> using UpLengths = LowLengths; UpLengths up_lengths_; - RightShift right_shift_; - CK_TILE_HOST_DEVICE constexpr xor_t() : up_lengths_{}, right_shift_{} {} + CK_TILE_HOST_DEVICE constexpr xor_t() : up_lengths_{} {} - CK_TILE_HOST_DEVICE constexpr xor_t(const LowLengths& low_lengths, - const RightShift& right_shift) - : up_lengths_{low_lengths}, right_shift_{right_shift} - { - } + CK_TILE_HOST_DEVICE constexpr xor_t(const LowLengths& low_lengths) : up_lengths_{low_lengths} {} CK_TILE_HOST_DEVICE static constexpr auto get_type_enum() { @@ -1378,13 +1373,8 @@ struct xor_t : public base_transform<2, 2> idx_low(number<0>{}) = idx_up[number<0>{}]; - const auto idx_low_1_tmp = - (idx_up[number<1>{}] - idx_up[number<0>{}] * right_shift_) % up_lengths_[number<1>{}]; - - const auto idx_low_1 = - (idx_low_1_tmp >= 0) ? idx_low_1_tmp : up_lengths_[number<1>{}] + idx_low_1_tmp; - - idx_low(number<1>{}) = idx_low_1; + idx_low(number<1>{}) = + idx_up[number<1>{}] ^ (idx_up[number<0>{}] % up_lengths_[number<1>{}]); } template @@ -1419,8 +1409,7 @@ struct xor_t : public base_transform<2, 2> CK_TILE_HOST_DEVICE static constexpr bool is_known_at_compile_time() { - return ck_tile::is_known_at_compile_time::value && - ck_tile::is_known_at_compile_time::value; + return ck_tile::is_known_at_compile_time::value; } // MUST be static function @@ -1432,14 +1421,6 @@ struct xor_t : public base_transform<2, 2> array up_vector_lengths = low_vector_lengths; array up_vector_strides = low_vector_strides; - if constexpr(ck_tile::is_known_at_compile_time::value) - { - if(low_vector_lengths[1] != -1) - { - up_vector_lengths(1) = gcd(low_vector_lengths[1], abs(right_shift_)); - } - } - return make_tuple(up_vector_lengths, up_vector_strides); } @@ -1452,10 +1433,6 @@ struct xor_t : public base_transform<2, 2> print(up_lengths_); printf(", "); - // - printf("right_shift_: "); - print(right_shift_); - printf("}"); } }; @@ -1655,11 +1632,10 @@ CK_TILE_HOST_DEVICE constexpr auto make_modulo_transform(const Modulus& modulus, return modulo{modulus, up_length}; } -template -CK_TILE_HOST_DEVICE constexpr auto make_xor_transform(const LowLengths& low_lengths, - const RightShift& right_shift) +template +CK_TILE_HOST_DEVICE constexpr auto make_xor_transform(const LowLengths& low_lengths) { - return xor_t{low_lengths, right_shift}; + return xor_t{low_lengths}; } template diff --git a/include/ck_tile/core/numeric/vector_type.hpp b/include/ck_tile/core/numeric/vector_type.hpp index c23c12f295..3ef066a3eb 100644 --- a/include/ck_tile/core/numeric/vector_type.hpp +++ b/include/ck_tile/core/numeric/vector_type.hpp @@ -117,6 +117,15 @@ using int32x16_t = int32_t __attribute__((ext_vector_type(16))); using int32x32_t = int32_t __attribute__((ext_vector_type(32))); using int32x64_t = int32_t __attribute__((ext_vector_type(64))); +// u32 +// using uint32_t = ... +using uint32x2_t = uint32_t __attribute__((ext_vector_type(2))); +using uint32x4_t = uint32_t __attribute__((ext_vector_type(4))); +using uint32x8_t = uint32_t __attribute__((ext_vector_type(8))); +using uint32x16_t = uint32_t __attribute__((ext_vector_type(16))); +using uint32x32_t = uint32_t __attribute__((ext_vector_type(32))); +using uint32x64_t = uint32_t __attribute__((ext_vector_type(64))); + // i16 // using int16_t = ... using int16x2_t = int16_t __attribute__((ext_vector_type(2))); diff --git a/include/ck_tile/core/tensor/tile_distribution.hpp b/include/ck_tile/core/tensor/tile_distribution.hpp index 42a30232fb..24c932f0a6 100644 --- a/include/ck_tile/core/tensor/tile_distribution.hpp +++ b/include/ck_tile/core/tensor/tile_distribution.hpp @@ -746,8 +746,9 @@ CK_TILE_HOST_DEVICE constexpr auto slice_distribution_from_x( return make_tuple( make_static_tile_distribution( tile_distribution_encoding, // only need to + // change the + // h_lengths type typename Encoding::Ps2RHssMajor, typename Encoding::Ps2RHssMinor, typename Encoding::Ys2RHsMajor, diff --git a/include/ck_tile/core/utility/philox_rand.hpp b/include/ck_tile/core/utility/philox_rand.hpp index c49f44ae48..87abf5cc18 100644 --- a/include/ck_tile/core/utility/philox_rand.hpp +++ b/include/ck_tile/core/utility/philox_rand.hpp @@ -53,6 +53,39 @@ class philox out_tmp[3] = tmp_ph.w; } + CK_TILE_HOST_DEVICE void get_random_8x8(uint8_t* out, + const unsigned long long subsequence, + const index_t start_idx) const + { + uint4 tmp_ph; + tmp_ph = get_philox_4x32(subsequence); + + uint32x4_t tmp; + tmp[0] = tmp_ph.x; + tmp[1] = tmp_ph.y; + tmp[2] = tmp_ph.z; + tmp[3] = tmp_ph.w; + uint32_t* out_tmp = reinterpret_cast(&out[0]); + out_tmp[0] = tmp[start_idx]; + out_tmp[1] = tmp[start_idx + 2]; + } + + CK_TILE_HOST_DEVICE void get_random_4x8(uint8_t* out, + const unsigned long long subsequence, + const index_t start_idx) const + { + uint4 tmp_ph; + tmp_ph = get_philox_4x32(subsequence); + + uint32x4_t tmp; + tmp[0] = tmp_ph.x; + tmp[1] = tmp_ph.y; + tmp[2] = tmp_ph.z; + tmp[3] = tmp_ph.w; + uint32_t* out_tmp = reinterpret_cast(&out[0]); + out_tmp[0] = tmp[start_idx]; + } + private: struct ull2 { diff --git a/include/ck_tile/ops/fmha.hpp b/include/ck_tile/ops/fmha.hpp index 057d2b11ff..cad3009473 100644 --- a/include/ck_tile/ops/fmha.hpp +++ b/include/ck_tile/ops/fmha.hpp @@ -8,21 +8,16 @@ #include "ck_tile/ops/fmha/block/block_masking.hpp" #include "ck_tile/ops/fmha/block/block_position_encoding.hpp" #include "ck_tile/ops/fmha/kernel/fmha_bwd_kernel.hpp" -#include "ck_tile/ops/fmha/kernel/fmha_bwd_tile_partitioner.hpp" #include "ck_tile/ops/fmha/kernel/fmha_fwd_kernel.hpp" #include "ck_tile/ops/fmha/kernel/fmha_fwd_splitkv_combine_kernel.hpp" #include "ck_tile/ops/fmha/kernel/fmha_fwd_splitkv_combine_tile_partitioner.hpp" #include "ck_tile/ops/fmha/kernel/fmha_fwd_splitkv_kernel.hpp" #include "ck_tile/ops/fmha/kernel/fmha_fwd_splitkv_tile_partitioner.hpp" #include "ck_tile/ops/fmha/kernel/fmha_fwd_tile_partitioner.hpp" +#include "ck_tile/ops/fmha/pipeline/block_fmha_bwd_convert_dq.hpp" #include "ck_tile/ops/fmha/pipeline/block_fmha_bwd_dot_do_o.hpp" -#include "ck_tile/ops/fmha/pipeline/block_fmha_bwd_dot_do_o_default_policy.hpp" -#include "ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_ks_kts_vr.hpp" -#include "ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_ks_kts_vr_default_policy.hpp" -#include "ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_ks_vr.hpp" -#include "ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_ks_vr_default_policy.hpp" -#include "ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_qs_ks_vr_dos.hpp" -#include "ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_qs_ks_vr_dos_default_policy.hpp" +#include "ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_kr_ktr_vr.hpp" +#include "ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_kr_ktr_vr_iglp.hpp" #include "ck_tile/ops/fmha/pipeline/block_fmha_bwd_pipeline_default_policy.hpp" #include "ck_tile/ops/fmha/pipeline/block_fmha_bwd_pipeline_enum.hpp" #include "ck_tile/ops/fmha/pipeline/block_fmha_bwd_pipeline_problem.hpp" diff --git a/include/ck_tile/ops/fmha/block/block_dropout.hpp b/include/ck_tile/ops/fmha/block/block_dropout.hpp index 7ebb306cce..e036402e16 100644 --- a/include/ck_tile/ops/fmha/block/block_dropout.hpp +++ b/include/ck_tile/ops/fmha/block/block_dropout.hpp @@ -286,11 +286,226 @@ struct BlockDropout }); } + ck_tile::philox ph; + const float rp_undrop; + const uint8_t p_undrop_in_uint8_t; + const bool is_store_randval; +}; + +template +struct BlockDropoutBwd; + +template +struct BlockDropoutBwd +{ + static constexpr bool IsDropout = false; + static constexpr bool IsStoreRandval = IsStoreRandval_; + + template + __host__ __device__ static constexpr auto + MakeRandvalDramWindow(RandValDramBlockWindowTmp& randval_dram_block_window_tmp, + index_t seqlen_qk_start) + { + (void)randval_dram_block_window_tmp; + (void)seqlen_qk_start; + + return make_null_tile_window(make_tuple(number<0>{}, number<0>{})); + } +}; + +template +struct BlockDropoutBwd +{ + static constexpr bool IsDropout = true; + // true: 32*32 warp gemm + // false: 16*16 warp gemm + static constexpr bool IsWG32 = IsWG32_; + static constexpr bool IsStoreRandval = IsStoreRandval_; + + CK_TILE_HOST_DEVICE BlockDropoutBwd(index_t i_batch, + index_t i_head, + index_t nheads, + unsigned long long seed, + unsigned long long offset, + float rp_undrop_, + uint8_t p_undrop_in_uint8_t_) + : ph(seed, + offset + (i_batch * nheads + i_head) * get_warp_size() + + (IsWG32 ? get_lane_id() : ((get_lane_id() & 47) + ((get_warp_id() & 1) << 4)))), + rp_undrop(rp_undrop_), + p_undrop_in_uint8_t(p_undrop_in_uint8_t_) + { + } + + template + CK_TILE_HOST_DEVICE static constexpr auto + MakeRandvalDramWindow(RandValDramBlockWindowTmp& randval_dram_block_window_tmp, + index_t seqlen_qk_start) + { + constexpr auto config = + BlockGemm::Policy::template GetWarpGemmMWarpNWarp(); + using BlockGemmShape = remove_cvref_t; + using WG = remove_cvref_t())>; + constexpr index_t kMPerBlock = BlockGemmShape::kM; + constexpr index_t MWarp = config.template at<1>(); + constexpr index_t NWarp = config.template at<2>(); + constexpr bool MBwdWG16MultiIterCheck = (!IsFwd) && (!IsWG32) && (kMPerBlock > 16); + constexpr index_t kMPerStep = [&]() { + if constexpr(MBwdWG16MultiIterCheck) + { + return MWarp * WG::kM * 2; + } + else + { + return MWarp * WG::kM; + } + }(); + constexpr index_t kNPerStep = NWarp * WG::kN; + + const auto block_origin = randval_dram_block_window_tmp.get_window_origin(); + auto randval_dram_window = [&]() { + if constexpr(IsFwd) + { + return make_tile_window( + randval_dram_block_window_tmp.get_bottom_tensor_view(), + ck_tile::make_tuple(number{}, number{}), + {block_origin.at(number<0>{}), seqlen_qk_start}); // M/N + } + else + { + return make_tile_window( + randval_dram_block_window_tmp.get_bottom_tensor_view(), + ck_tile::make_tuple(number{}, number{}), + {seqlen_qk_start, block_origin.at(number<1>{})}); // M/N + } + }(); + + return randval_dram_window; + } + + template + CK_TILE_HOST_DEVICE static constexpr auto MakeRandValLdsBlockDescriptor() + { + constexpr auto config = + BlockGemm::Policy::template GetWarpGemmMWarpNWarp(); + using WG = remove_cvref_t())>; + constexpr index_t MWarp = config.template at<1>(); + constexpr index_t kMPerStep = MWarp * WG::kM; + constexpr index_t kNPerStep = WG::kN; + constexpr index_t kN1 = 8; + constexpr index_t kN0 = kNPerStep / kN1; + + constexpr auto randval_lds_block_desc_0 = make_naive_tensor_descriptor( + ck_tile::make_tuple(number{}, number{}, number{}), + ck_tile::make_tuple(number<(kMPerStep + 1) * kN1>{}, number{}, number<1>{}), + number{}, + number<1>{}); + + constexpr auto randval_lds_block_desc = transform_tensor_descriptor( + randval_lds_block_desc_0, + ck_tile::make_tuple( + make_pass_through_transform(number{}), + make_merge_transform(ck_tile::make_tuple(number{}, number{}))), + ck_tile::make_tuple(sequence<1>{}, sequence<0, 2>{}), + ck_tile::make_tuple(sequence<0>{}, sequence<1>{})); + + return randval_lds_block_desc; + } + + template + CK_TILE_HOST_DEVICE static constexpr auto MakeRandValTileDistribution() + { + constexpr auto config = + BlockGemm::Policy::template GetWarpGemmMWarpNWarp(); + using BlockGemmShape = remove_cvref_t; + constexpr index_t kMPerBlock = BlockGemmShape::kM; + constexpr index_t MWarp = config.template at<1>(); + constexpr index_t NWarp = config.template at<2>(); + constexpr bool MBwdWG16MultiIterCheck = (!IsFwd) && (!IsWG32) && (kMPerBlock > 16); + + constexpr index_t MIterPerWarp = [&]() { + if constexpr(MBwdWG16MultiIterCheck) + { + return 2; + } + else + { + return 1; + } + }(); + constexpr index_t NIterPerWarp = 1; + + constexpr auto randval_block_outer_part_dstr_encoding = tile_distribution_encoding< + sequence<>, + tuple, sequence>, + tuple>, + tuple>, + sequence<1, 2>, + sequence<0, 0>>{}; + + // Use Bwd WarpGemm to ensure that Fwd's random values ​​are consistent with Bwd. + // except headdim256. + constexpr auto randval_block_inner_part_dstr_encoding = []() { + if constexpr(std::is_same_v && + std::is_same_v && + std::is_same_v) + { + if constexpr(IsWG32) + return typename WarpGemmMfmaF16F16F32M32N32K16SwizzleA::CWarpDstrEncoding{}; + else + return typename WarpGemmMfmaF16F16F32M16N16K16::CWarpDstrEncoding{}; + } + else + { + if constexpr(IsWG32) + return typename WarpGemmMfmaBf16Bf16F32M32N32K16SwizzleA::CWarpDstrEncoding{}; + else + return typename WarpGemmMfmaBf16Bf16F32M16N16K16::CWarpDstrEncoding{}; + } + }(); + + constexpr auto randval_block_part_dstr_encode = + detail::make_embed_tile_distribution_encoding(randval_block_outer_part_dstr_encoding, + randval_block_inner_part_dstr_encoding); + + return make_static_tile_distribution(randval_block_part_dstr_encode); + } + + template + CK_TILE_HOST_DEVICE static constexpr auto MakeRandValLdsShuffleTileDistribution() + { + 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>(); + + constexpr index_t MIterPerWarp = 1; + constexpr index_t NIterPerWarp = 1; + + constexpr auto randval_block_outer_part_dstr_encoding = tile_distribution_encoding< + sequence<>, + tuple, sequence>, + tuple>, + tuple>, + sequence<1, 2>, + sequence<0, 0>>{}; + + constexpr auto randval_block_part_dstr_encode = + detail::make_embed_tile_distribution_encoding(randval_block_outer_part_dstr_encoding, + typename WG::CWarpDstrEncoding{}); + + return make_static_tile_distribution(randval_block_part_dstr_encode); + } + template - CK_TILE_HOST_DEVICE void Run(const index_t start_m0_idx, + CK_TILE_HOST_DEVICE void Run(void* randval_ptr, + const index_t start_m0_idx, + const index_t start_n0_idx, PComputeWindow& p_compute, RandValDramWindow& randval_dram_window) const { @@ -305,30 +520,177 @@ struct BlockDropout constexpr index_t kMPerStep = MWarp * WG::kM; constexpr index_t kNPerStep = NWarp * WG::kN; - // register distribute - auto randval = - make_static_distributed_tensor(MakeRandValTileDistribution()); - static_assert(randval.kThreadElementSpaceSize == 16); + // randval tile in LDS + auto randval_lds = make_tensor_view( + reinterpret_cast(randval_ptr), MakeRandValLdsBlockDescriptor()); - const int start_n0_idx = randval_dram_window.get_window_origin().at(number<1>{}); - static_for<0, kNPerBlock / kNPerStep, 1>{}([&](auto i_n0) { - static_for<0, kMPerBlock / kMPerStep, 1>{}([&](auto i_m0) { - int block_row_start = (start_m0_idx / WG::kM) + i_m0; - int block_col_start = (start_n0_idx / WG::kN) + (i_n0 * NWarp) + get_warp_id(); + auto randval_lds_window = make_tile_window( + randval_lds, MakeRandValLdsBlockDescriptor().get_lengths(), {0, 0}); + + // register distribute + auto randval_dist_generated = + make_static_distributed_tensor(MakeRandValTileDistribution()); + static_assert(randval_dist_generated.kThreadElementSpaceSize == 16); + + auto randval_lds_read_window = + make_tile_window(randval_lds_window.get_bottom_tensor_view(), + randval_lds_window.get_window_lengths(), + randval_lds_window.get_window_origin(), + MakeRandValLdsShuffleTileDistribution()); + + static_for<0, kMPerBlock / kMPerStep, 1>{}([&](auto i_m0) { + static_for<0, kNPerBlock / kNPerStep, 1>{}([&](auto i_n0) { + int block_row_start = (start_m0_idx / WG::kM) + (i_m0 * MWarp) + get_warp_id(); + int block_col_start = (start_n0_idx / WG::kN) + i_n0; uint2 rowcol = make_uint2(block_row_start, block_col_start); // generate random number uint8_t random_uint8_t[16]; ph.get_random_16x8(random_uint8_t, reinterpret_cast(rowcol)); + constexpr auto randval_dist_generated_spans = + decltype(randval_dist_generated)::get_distributed_spans(); + int i_random_idx = 0; + sweep_tile_span(randval_dist_generated_spans[number<0>{}], [&](auto idx0) { + sweep_tile_span(randval_dist_generated_spans[number<1>{}], [&](auto idx1) { + constexpr auto i_j_idx = ck_tile::make_tuple(idx0, idx1); + randval_dist_generated(i_j_idx) = random_uint8_t[i_random_idx++]; + }); + }); + // save to LDS + store_tile(randval_lds_window, randval_dist_generated); + block_sync_lds(); + // read from LDS to register + auto randval = load_tile(randval_lds_read_window); + constexpr auto randval_spans = decltype(randval)::get_distributed_spans(); + sweep_tile_span(randval_spans[number<0>{}], [&](auto idx0) { + sweep_tile_span(randval_spans[number<1>{}], [&](auto idx1) { + constexpr auto p_idx0 = tile_distributed_index{}; + constexpr auto p_idx1 = + tile_distributed_index{}; + constexpr auto p_idx = ck_tile::make_tuple(p_idx0, p_idx1); + constexpr auto r_idx = ck_tile::make_tuple(idx0, idx1); + p_compute(p_idx) = randval[r_idx] <= p_undrop_in_uint8_t + ? p_compute[p_idx] * rp_undrop + : PComputeDataType(0); + }); + }); + // save to Global + if constexpr(IsStoreRandval) + { + const auto randval_store = cast_tile(randval); + store_tile(randval_dram_window, randval_store); + move_tile_window(randval_dram_window, {0, kNPerStep}); + } + }); + if constexpr(IsStoreRandval) + { + move_tile_window(randval_dram_window, {kMPerStep, -kNPerBlock}); + } + }); + if constexpr(IsStoreRandval) + { + move_tile_window(randval_dram_window, {-kMPerBlock, kNPerBlock}); + } + } + + template + CK_TILE_HOST_DEVICE void Run(const index_t start_m0_idx, + const index_t start_n0_idx, + PComputeWindow& p_compute, + RandValDramWindow& randval_dram_window) const + { + 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>(); + using BlockGemmShape = remove_cvref_t; + constexpr index_t kMPerBlock = BlockGemmShape::kM; + constexpr index_t kNPerBlock = BlockGemmShape::kN; + constexpr bool MBwdWG16MultiIterCheck = (!IsWG32) && (kMPerBlock > 16); + constexpr bool MBwdWG16SingleIterCheck = (!IsWG32) && (kMPerBlock == 16); + constexpr index_t kMPerStep = [&]() { + if constexpr(MBwdWG16MultiIterCheck) + { + return MWarp * WG::kM * 2; + } + else + { + return MWarp * WG::kM; + } + }(); + constexpr index_t kNPerStep = NWarp * WG::kN; + + // register distribute + auto randval = make_static_distributed_tensor( + MakeRandValTileDistribution()); + if constexpr(IsWG32) + static_assert(randval.kThreadElementSpaceSize == 16); + else + static_assert(randval.kThreadElementSpaceSize == 4 || + randval.kThreadElementSpaceSize == 8); + + static_for<0, kNPerBlock / kNPerStep, 1>{}([&](auto i_n0) { + static_for<0, kMPerBlock / kMPerStep, 1>{}([&](auto i_m0) { + int block_row_start, block_col_start; + if constexpr(IsWG32) + { + block_row_start = (start_m0_idx / WG::kM) + i_m0; + block_col_start = (start_n0_idx / WG::kN) + (i_n0 * NWarp) + get_warp_id(); + } + else + { + block_row_start = start_m0_idx / 32 + i_m0; + block_col_start = (start_n0_idx / 32) + get_warp_id() / 2 + i_n0 * 2; + } + uint2 rowcol = make_uint2(block_row_start, block_col_start); + + // generate random number + uint8_t* random_uint8_t_; + if constexpr(MBwdWG16SingleIterCheck) + { + uint8_t random_uint8_t[4]; + // m0t0 ~m0t15/m0t32~m0t47: 0 + // m0t16~m0t31/m0t48~m0t63: 1 + // m1t0 ~m1t15/m1t32~m1t47: 2 + // m1t16~m1t31/m1t48~m1t63: 3 + const index_t start_idx = + ((get_lane_id() >> 4) & 1) + (((start_m0_idx >> 4) & 1) << 1); + ph.get_random_4x8( + random_uint8_t, reinterpret_cast(rowcol), start_idx); + random_uint8_t_ = random_uint8_t; + } + else if constexpr(MBwdWG16MultiIterCheck) + { + uint8_t random_uint8_t[8]; + // t0 ~t15/t32~t47: 0 + // t16~t31/t48~t63: 1 + const index_t start_idx = (get_lane_id() >> 4) & 1; + ph.get_random_8x8( + random_uint8_t, reinterpret_cast(rowcol), start_idx); + random_uint8_t_ = random_uint8_t; + } + else + { + uint8_t random_uint8_t[16]; + ph.get_random_16x8(random_uint8_t, + reinterpret_cast(rowcol)); + random_uint8_t_ = random_uint8_t; + } + constexpr auto randval_spans = decltype(randval)::get_distributed_spans(); int i_random_idx = 0; sweep_tile_span(randval_spans[number<0>{}], [&](auto idx0) { sweep_tile_span(randval_spans[number<1>{}], [&](auto idx1) { - constexpr auto r_idx = ck_tile::make_tuple(idx0, idx1); - randval(r_idx) = random_uint8_t[i_random_idx++]; - constexpr auto p_idx0 = - tile_distributed_index{}; + constexpr auto r_idx = ck_tile::make_tuple(idx0, idx1); + randval(r_idx) = random_uint8_t_[i_random_idx++]; + constexpr auto p_idx0 = tile_distributed_index{}; constexpr auto p_idx1 = tile_distributed_index{}; constexpr auto p_idx = ck_tile::make_tuple(p_idx0, p_idx1); p_compute(p_idx) = randval[r_idx] <= p_undrop_in_uint8_t @@ -337,19 +699,19 @@ struct BlockDropout }); }); // save to Global - if(is_store_randval) + if constexpr(IsStoreRandval) { const auto randval_store = cast_tile(randval); store_tile(randval_dram_window, randval_store); move_tile_window(randval_dram_window, {kMPerStep, 0}); } }); - if(is_store_randval) + if constexpr(IsStoreRandval) { move_tile_window(randval_dram_window, {-kMPerBlock, kNPerStep}); } }); - if(is_store_randval) + if constexpr(IsStoreRandval) { move_tile_window(randval_dram_window, {kMPerBlock, -kNPerBlock}); } @@ -358,7 +720,6 @@ struct BlockDropout ck_tile::philox ph; const float rp_undrop; const uint8_t p_undrop_in_uint8_t; - const bool is_store_randval; }; } // namespace ck_tile diff --git a/include/ck_tile/ops/fmha/kernel/fmha_bwd_kernel.hpp b/include/ck_tile/ops/fmha/kernel/fmha_bwd_kernel.hpp index e713cefbda..167494b193 100644 --- a/include/ck_tile/ops/fmha/kernel/fmha_bwd_kernel.hpp +++ b/include/ck_tile/ops/fmha/kernel/fmha_bwd_kernel.hpp @@ -23,13 +23,9 @@ namespace ck_tile { -template +template struct FmhaBwdDQDKDVKernel { - using TilePartitioner = ck_tile::remove_cvref_t; using FmhaPipeline = ck_tile::remove_cvref_t; using KGradEpiloguePipeline = ck_tile::remove_cvref_t; using VGradEpiloguePipeline = ck_tile::remove_cvref_t; @@ -59,9 +55,12 @@ struct FmhaBwdDQDKDVKernel static constexpr bool kPadHeadDimV = FmhaPipeline::kPadHeadDimV; static constexpr auto BiasEnum = FmhaPipeline::BiasEnum; static constexpr bool kHasBiasGrad = FmhaPipeline::kHasBiasGrad; - static constexpr bool kHasDropout = FmhaPipeline::kHasDropout; using FmhaMask = ck_tile::remove_cvref_t; - static constexpr bool kHasMask = FmhaMask::IsMasking; + using FmhaDropout = ck_tile::remove_cvref_t; + static constexpr bool kHasMask = FmhaMask::IsMasking; + static constexpr bool kHasDropout = FmhaDropout::IsDropout; + static constexpr bool kIsStoreRandval = FmhaDropout::IsStoreRandval; + static constexpr bool kIsDeterministic = FmhaPipeline::kIsDeterministic; // clang-format off template struct t2s; @@ -73,9 +72,12 @@ struct FmhaBwdDQDKDVKernel { // sync with generate.py // clang-format off - using bfs = typename FmhaPipeline::BlockFmhaShape; - using gbr = typename bfs::Gemm0BlockWarps; - using gwt = typename bfs::Gemm0WarpTile; + using bfs = typename FmhaPipeline::BlockFmhaShape; + using gbr0 = typename bfs::Gemm0BlockWarps; + using gbr1 = typename bfs::Gemm1BlockWarps; + using gbr4 = typename bfs::Gemm4BlockWarps; + using gwt0 = typename bfs::Gemm0WarpTile; + using gwt1 = typename bfs::Gemm1WarpTile; #define _SS_ std::string #define _TS_ std::to_string auto pn = [&] () { @@ -88,13 +90,17 @@ struct FmhaBwdDQDKDVKernel return _SS_("fmha_bwd_d") + _TS_(bfs::kQKHeaddim) + "_" + _SS_(t2s::name) + "_" + (kIsGroupMode ? "group" : "batch") + "_" + - "b" + _TS_(bfs::kM0) + "x" + _TS_(bfs::kN0) + "x" + _TS_(bfs::kK0) + "x" + - _TS_(bfs::kQKHeaddim) + "x" + _TS_(bfs::kVHeaddim) + "_" + - "r" + _TS_(gbr::at(ck_tile::number<0>{})) + "x" + _TS_(gbr::at(ck_tile::number<1>{})) + "x" + _TS_(gbr::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>{})) + "_" + + "b" + _TS_(bfs::kM0) + "x" + _TS_(bfs::kN0) + "x" + _TS_(bfs::kK0) + "x" + _TS_(bfs::kK1) + "x" + _TS_(bfs::kK2) + "x" + _TS_(bfs::kK3) + "x" + + _TS_(bfs::kK4) + "x" + _TS_(bfs::kQKHeaddim) + "x" + _TS_(bfs::kVHeaddim) + "_" + + "r" + _TS_(gbr0::at(ck_tile::number<0>{})) + "x" + _TS_(gbr0::at(ck_tile::number<1>{})) + "x" + _TS_(gbr0::at(ck_tile::number<2>{})) + "_" + + "r" + _TS_(gbr1::at(ck_tile::number<0>{})) + "x" + _TS_(gbr1::at(ck_tile::number<1>{})) + "x" + _TS_(gbr1::at(ck_tile::number<2>{})) + "_" + + "r" + _TS_(gbr4::at(ck_tile::number<0>{})) + "x" + _TS_(gbr4::at(ck_tile::number<1>{})) + "x" + _TS_(gbr4::at(ck_tile::number<2>{})) + "_" + + "w" + _TS_(gwt0::at(ck_tile::number<0>{})) + "x" + _TS_(gwt0::at(ck_tile::number<1>{})) + "x" + _TS_(gwt0::at(ck_tile::number<2>{})) + "_" + + "w" + _TS_(gwt1::at(ck_tile::number<0>{})) + "x" + _TS_(gwt1::at(ck_tile::number<1>{})) + "x" + _TS_(gwt1::at(ck_tile::number<2>{})) + "_" + ("o" + _TS_(kBlockPerCu) + "_") + _SS_(FmhaPipeline::name) + (pn.empty() ? "" : "_" + pn) + (BiasEnum == BlockAttentionBiasEnum::NO_BIAS ? _SS_("") : (_SS_("_") + BlockAttentionBiasEnumToStr::name)) + - (kHasBiasGrad ? "_dbias" : "") + (kHasMask ? "_" + _SS_(FmhaMask::name) : "") + (kHasDropout ? "_dropout" : "" ); + (kHasBiasGrad ? "_dbias" : "") + (kHasMask ? "_" + _SS_(FmhaMask::name) : "") + (kHasDropout ? "_dropout" : "" ) + + (kIsStoreRandval ? "_storerandval" : "" ) + (kIsDeterministic ? "_deterministic" : "" ); #undef _SS_ #undef _TS_ // clang-format on @@ -117,7 +123,7 @@ struct FmhaBwdDQDKDVKernel const void* lse_ptr; const void* do_ptr; const void* d_ptr; - void* dq_ptr; + void* dq_acc_ptr; void* dk_ptr; void* dv_ptr; @@ -131,14 +137,13 @@ struct FmhaBwdDQDKDVKernel ck_tile::index_t num_head_q; ck_tile::index_t nhead_ratio_qk; float raw_scale; -#if CK_TILE_FMHA_FWD_FAST_EXP2 float scale; -#endif ck_tile::index_t stride_q; ck_tile::index_t stride_k; ck_tile::index_t stride_v; ck_tile::index_t stride_do; + ck_tile::index_t stride_dq_acc; ck_tile::index_t stride_dk; ck_tile::index_t stride_dv; @@ -147,8 +152,9 @@ struct FmhaBwdDQDKDVKernel ck_tile::index_t nhead_stride_v; ck_tile::index_t nhead_stride_do; ck_tile::index_t nhead_stride_lsed; - - ck_tile::index_t batch_stride_lsed; + ck_tile::index_t nhead_stride_dq_acc; + ck_tile::index_t nhead_stride_dk; + ck_tile::index_t nhead_stride_dv; }; struct FmhaBwdCommonBiasKargs @@ -206,7 +212,6 @@ struct FmhaBwdDQDKDVKernel float rp_undrop = 1; float scale_rp_undrop = 1; uint8_t p_undrop_in_uint8_t = std::numeric_limits::max(); - bool is_store_randval = false; uint64_t drop_seed = 1; uint64_t drop_offset = 0; void* rand_val_ptr = nullptr; @@ -218,6 +223,10 @@ struct FmhaBwdDQDKDVKernel { ck_tile::index_t batch_stride_randval = 0; }; + struct FmhaBwdDeterministicKargs + { + ck_tile::index_t split_stride_dq_acc = 0; + }; struct FmhaBwdBatchModeKargs : FmhaBwdCommonKargs, @@ -228,12 +237,15 @@ struct FmhaBwdDQDKDVKernel FmhaBwdEmptyKargs<0>>>, std::conditional_t>, std::conditional_t>, - std::conditional_t> + std::conditional_t>, + std::conditional_t> { ck_tile::index_t batch_stride_q; ck_tile::index_t batch_stride_k; ck_tile::index_t batch_stride_v; ck_tile::index_t batch_stride_do; + ck_tile::index_t batch_stride_lsed; + ck_tile::index_t batch_stride_dq_acc; ck_tile::index_t batch_stride_dk; ck_tile::index_t batch_stride_dv; }; @@ -247,7 +259,8 @@ struct FmhaBwdDQDKDVKernel FmhaBwdEmptyKargs<0>>>, std::conditional_t>, std::conditional_t>, - std::conditional_t> + std::conditional_t>, + std::conditional_t> { const int32_t* seqstart_q_ptr; const int32_t* seqstart_k_ptr; @@ -266,10 +279,10 @@ struct FmhaBwdDQDKDVKernel const void* do_ptr, const void* d_ptr, void* rand_val_ptr, - void* dq_ptr, void* dk_ptr, void* dv_ptr, void* dbias_ptr, + void* dq_acc_ptr, ck_tile::index_t seqlen_q, ck_tile::index_t seqlen_k, ck_tile::index_t hdim_q, @@ -283,6 +296,7 @@ struct FmhaBwdDQDKDVKernel ck_tile::index_t stride_bias, ck_tile::index_t stride_randval, ck_tile::index_t stride_do, + ck_tile::index_t stride_dq_acc, ck_tile::index_t stride_dk, ck_tile::index_t stride_dv, ck_tile::index_t stride_dbias, @@ -293,6 +307,9 @@ struct FmhaBwdDQDKDVKernel ck_tile::index_t nhead_stride_randval, ck_tile::index_t nhead_stride_do, ck_tile::index_t nhead_stride_lsed, + ck_tile::index_t nhead_stride_dq_acc, + ck_tile::index_t nhead_stride_dk, + ck_tile::index_t nhead_stride_dv, ck_tile::index_t nhead_stride_dbias, ck_tile::index_t batch_stride_q, ck_tile::index_t batch_stride_k, @@ -301,14 +318,15 @@ struct FmhaBwdDQDKDVKernel ck_tile::index_t batch_stride_randval, ck_tile::index_t batch_stride_do, ck_tile::index_t batch_stride_lsed, + ck_tile::index_t batch_stride_dq_acc, ck_tile::index_t batch_stride_dk, ck_tile::index_t batch_stride_dv, ck_tile::index_t batch_stride_dbias, + ck_tile::index_t split_stride_dq_acc, ck_tile::index_t window_size_left, ck_tile::index_t window_size_right, ck_tile::index_t mask_type, float p_drop, - bool s_randval, const std::tuple& drop_seed_offset) { Kargs kargs{{q_ptr, @@ -317,7 +335,7 @@ struct FmhaBwdDQDKDVKernel lse_ptr, do_ptr, d_ptr, - dq_ptr, + dq_acc_ptr, dk_ptr, dv_ptr, seqlen_q, @@ -327,13 +345,12 @@ struct FmhaBwdDQDKDVKernel num_head_q, nhead_ratio_qk, scale, -#if CK_TILE_FMHA_FWD_FAST_EXP2 static_cast(scale * ck_tile::log2e_v<>), -#endif stride_q, stride_k, stride_v, stride_do, + stride_dq_acc, stride_dk, stride_dv, nhead_stride_q, @@ -341,15 +358,20 @@ struct FmhaBwdDQDKDVKernel nhead_stride_v, nhead_stride_do, nhead_stride_lsed, - batch_stride_lsed}, // args for common karg - {}, // placeholder for bias - {}, // placeholder for dbias - {}, // placeholder for mask - {}, // placeholder for dropout + nhead_stride_dq_acc, + nhead_stride_dk, + nhead_stride_dv}, // args for common karg + {}, // placeholder for bias + {}, // placeholder for dbias + {}, // placeholder for mask + {}, // placeholder for dropout + {}, // placeholder for deterministic batch_stride_q, batch_stride_k, batch_stride_v, batch_stride_do, + batch_stride_lsed, + batch_stride_dq_acc, batch_stride_dk, batch_stride_dv}; @@ -384,11 +406,18 @@ struct FmhaBwdDQDKDVKernel if constexpr(kHasDropout) { kargs.init_dropout(p_drop, drop_seed_offset, scale); - kargs.rand_val_ptr = rand_val_ptr; - kargs.stride_randval = stride_randval; - kargs.nhead_stride_randval = nhead_stride_randval; - kargs.batch_stride_randval = batch_stride_randval; - kargs.is_store_randval = s_randval; + if constexpr(kIsStoreRandval) + { + kargs.rand_val_ptr = rand_val_ptr; + kargs.stride_randval = stride_randval; + kargs.nhead_stride_randval = nhead_stride_randval; + kargs.batch_stride_randval = batch_stride_randval; + } + } + + if constexpr(kIsDeterministic) + { + kargs.split_stride_dq_acc = split_stride_dq_acc; } return kargs; @@ -404,10 +433,10 @@ struct FmhaBwdDQDKDVKernel const void* do_ptr, const void* d_ptr, void* rand_val_ptr, - void* dq_ptr, void* dk_ptr, void* dv_ptr, void* dbias_ptr, + void* dq_acc_ptr, const void* seqstart_q_ptr, const void* seqstart_k_ptr, const void* seqlen_k_ptr, @@ -422,6 +451,7 @@ struct FmhaBwdDQDKDVKernel ck_tile::index_t stride_bias, ck_tile::index_t stride_randval, ck_tile::index_t stride_do, + ck_tile::index_t stride_dq_acc, ck_tile::index_t stride_dk, ck_tile::index_t stride_dv, ck_tile::index_t stride_dbias, @@ -432,13 +462,15 @@ struct FmhaBwdDQDKDVKernel ck_tile::index_t nhead_stride_randval, ck_tile::index_t nhead_stride_do, ck_tile::index_t nhead_stride_lsed, + ck_tile::index_t nhead_stride_dq_acc, + ck_tile::index_t nhead_stride_dk, + ck_tile::index_t nhead_stride_dv, ck_tile::index_t nhead_stride_dbias, - ck_tile::index_t batch_stride_lsed, + ck_tile::index_t split_stride_dq_acc, ck_tile::index_t window_size_left, ck_tile::index_t window_size_right, ck_tile::index_t mask_type, float p_drop, - bool s_randval, const std::tuple& drop_seed_offset) { Kargs kargs{{q_ptr, @@ -447,7 +479,7 @@ struct FmhaBwdDQDKDVKernel lse_ptr, do_ptr, d_ptr, - dq_ptr, + dq_acc_ptr, dk_ptr, dv_ptr, -1, // seqlen will be updated by another pointer @@ -457,13 +489,12 @@ struct FmhaBwdDQDKDVKernel num_head_q, nhead_ratio_qk, scale, -#if CK_TILE_FMHA_FWD_FAST_EXP2 static_cast(scale * ck_tile::log2e_v<>), -#endif stride_q, stride_k, stride_v, stride_do, + stride_dq_acc, stride_dk, stride_dv, nhead_stride_q, @@ -471,11 +502,14 @@ struct FmhaBwdDQDKDVKernel nhead_stride_v, nhead_stride_do, nhead_stride_lsed, - batch_stride_lsed}, // args for common karg - {}, // placeholder for bias - {}, // placeholder for dbias - {}, // placeholder for mask - {}, // placeholder for dropout + nhead_stride_dq_acc, + nhead_stride_dk, + nhead_stride_dv}, // args for common karg + {}, // placeholder for bias + {}, // placeholder for dbias + {}, // placeholder for mask + {}, // placeholder for dropout + {}, // placeholder for deterministic reinterpret_cast(seqstart_q_ptr), reinterpret_cast(seqstart_k_ptr), reinterpret_cast(seqlen_k_ptr)}; @@ -506,10 +540,16 @@ struct FmhaBwdDQDKDVKernel if constexpr(kHasDropout) { kargs.init_dropout(p_drop, drop_seed_offset, scale); - kargs.rand_val_ptr = rand_val_ptr; - kargs.stride_randval = stride_randval; - kargs.nhead_stride_randval = nhead_stride_randval; - kargs.is_store_randval = s_randval; + if constexpr(kIsStoreRandval) + { + kargs.rand_val_ptr = rand_val_ptr; + kargs.stride_randval = stride_randval; + kargs.nhead_stride_randval = nhead_stride_randval; + } + } + if constexpr(kIsDeterministic) + { + kargs.split_stride_dq_acc = split_stride_dq_acc; } return kargs; @@ -518,7 +558,17 @@ struct FmhaBwdDQDKDVKernel CK_TILE_HOST static constexpr auto GridSize(ck_tile::index_t batch_size_, ck_tile::index_t nhead_, ck_tile::index_t seqlen_k_) { - return TilePartitioner::GridSize(batch_size_, nhead_, seqlen_k_); + return dim3( + ck_tile::integer_divide_ceil(seqlen_k_, FmhaPipeline::kN0), nhead_, batch_size_); + } + + CK_TILE_DEVICE static constexpr auto GetTileIndex() + { + const index_t i_block = blockIdx.x; + const index_t i_nhead = blockIdx.y; + const index_t i_batch = blockIdx.z; + + return ck_tile::make_tuple(i_block, i_nhead, i_batch); } CK_TILE_HOST static constexpr auto BlockSize() { return dim3(kBlockSize); } @@ -536,7 +586,7 @@ struct FmhaBwdDQDKDVKernel __shared__ char smem_ptr[GetSmemSize()]; // divide problem - const auto [i_tile_n, i_nhead, i_batch] = TilePartitioner{}(kargs.seqlen_k); + const auto [i_tile_n, i_nhead, i_batch] = GetTileIndex(); const index_t i_n0 = __builtin_amdgcn_readfirstlane(i_tile_n * FmhaPipeline::kN0); @@ -547,6 +597,7 @@ struct FmhaBwdDQDKDVKernel long_index_t batch_offset_randval = 0; long_index_t batch_offset_do = 0; long_index_t batch_offset_lsed = 0; + long_index_t batch_offset_dq_acc = 0; long_index_t batch_offset_dk = 0; long_index_t batch_offset_dv = 0; long_index_t batch_offset_dbias = 0; @@ -557,13 +608,14 @@ struct FmhaBwdDQDKDVKernel const long_index_t query_start = kargs.seqstart_q_ptr[i_batch]; const long_index_t key_start = kargs.seqstart_k_ptr[i_batch]; - batch_offset_q = query_start * kargs.stride_q; - batch_offset_k = key_start * kargs.stride_k; - batch_offset_v = key_start * kargs.stride_v; - batch_offset_do = query_start * kargs.stride_do; - batch_offset_lsed = static_cast(i_batch) * kargs.batch_stride_lsed; - batch_offset_dk = key_start * kargs.stride_dk; - batch_offset_dv = key_start * kargs.stride_dv; + batch_offset_q = query_start * kargs.stride_q; + batch_offset_k = key_start * kargs.stride_k; + batch_offset_v = key_start * kargs.stride_v; + batch_offset_do = query_start * kargs.stride_do; + batch_offset_lsed = query_start; + batch_offset_dq_acc = query_start * kargs.stride_dq_acc; + batch_offset_dk = key_start * kargs.stride_dk; + batch_offset_dv = key_start * kargs.stride_dv; if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS) { batch_offset_bias = query_start * kargs.stride_bias; @@ -576,7 +628,7 @@ struct FmhaBwdDQDKDVKernel { batch_offset_dbias = key_start; } - if constexpr(kHasDropout) + if constexpr(kIsStoreRandval) { batch_offset_randval = query_start * kargs.stride_randval; } @@ -603,13 +655,14 @@ struct FmhaBwdDQDKDVKernel } else { - batch_offset_q = static_cast(i_batch) * kargs.batch_stride_q; - batch_offset_k = static_cast(i_batch) * kargs.batch_stride_k; - batch_offset_v = static_cast(i_batch) * kargs.batch_stride_v; - batch_offset_do = static_cast(i_batch) * kargs.batch_stride_do; - batch_offset_lsed = static_cast(i_batch) * kargs.batch_stride_lsed; - batch_offset_dk = static_cast(i_batch) * kargs.batch_stride_dk; - batch_offset_dv = static_cast(i_batch) * kargs.batch_stride_dv; + batch_offset_q = static_cast(i_batch) * kargs.batch_stride_q; + batch_offset_k = static_cast(i_batch) * kargs.batch_stride_k; + batch_offset_v = static_cast(i_batch) * kargs.batch_stride_v; + batch_offset_do = static_cast(i_batch) * kargs.batch_stride_do; + batch_offset_lsed = static_cast(i_batch) * kargs.batch_stride_lsed; + batch_offset_dq_acc = static_cast(i_batch) * kargs.batch_stride_dq_acc; + batch_offset_dk = static_cast(i_batch) * kargs.batch_stride_dk; + batch_offset_dv = static_cast(i_batch) * kargs.batch_stride_dv; if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS) { batch_offset_bias = static_cast(i_batch) * kargs.batch_stride_bias; @@ -618,7 +671,7 @@ struct FmhaBwdDQDKDVKernel { batch_offset_dbias = static_cast(i_batch) * kargs.batch_stride_dbias; } - if constexpr(kHasDropout) + if constexpr(kIsStoreRandval) { batch_offset_randval = static_cast(i_batch) * kargs.batch_stride_randval; @@ -646,14 +699,11 @@ struct FmhaBwdDQDKDVKernel const OGradDataType* do_ptr = reinterpret_cast(kargs.do_ptr) + static_cast(i_nhead) * kargs.nhead_stride_do + batch_offset_do; - QGradDataType* dq_ptr = reinterpret_cast(kargs.dq_ptr) + - static_cast(i_nhead) * kargs.nhead_stride_q + - batch_offset_q; KGradDataType* dk_ptr = reinterpret_cast(kargs.dk_ptr) + - static_cast(i_nhead) * kargs.nhead_stride_k + + static_cast(i_nhead) * kargs.nhead_stride_dk + batch_offset_dk; VGradDataType* dv_ptr = reinterpret_cast(kargs.dv_ptr) + - static_cast(i_nhead) * kargs.nhead_stride_v + + static_cast(i_nhead) * kargs.nhead_stride_dv + batch_offset_dv; // Q/K/V/LSE/D/dO/dQ/dK/dV DRAM and DRAM window @@ -663,45 +713,10 @@ struct FmhaBwdDQDKDVKernel make_tuple(kargs.stride_q, 1), number{}, number<1>{}); - const auto q_dram = [&]() { - if constexpr(FmhaPipeline::kQLoadOnce) - { - return pad_tensor_view( - q_dram_naive, - make_tuple(number{}, number{}), - sequence{}); - } - else - { - return pad_tensor_view( - q_dram_naive, - make_tuple(number{}, number{}), - sequence{}); - } - }(); - - const auto qt_dram_naive = - transform_tensor_view(q_dram_naive, - make_tuple(make_pass_through_transform(kargs.hdim_q), - make_pass_through_transform(kargs.seqlen_q)), - make_tuple(sequence<1>{}, sequence<0>{}), - make_tuple(sequence<0>{}, sequence<1>{})); - const auto qt_dram = [&]() { - if constexpr(FmhaPipeline::kQTLoadOnce) - { - return pad_tensor_view( - qt_dram_naive, - make_tuple(number{}, number{}), - sequence{}); - } - else - { - return pad_tensor_view( - qt_dram_naive, - make_tuple(number{}, number{}), - sequence{}); - } - }(); + const auto q_dram = pad_tensor_view( + q_dram_naive, + make_tuple(number{}, number{}), + sequence{}); const auto k_dram_naive = make_naive_tensor_view( k_ptr, @@ -709,45 +724,10 @@ struct FmhaBwdDQDKDVKernel make_tuple(kargs.stride_k, 1), number{}, number<1>{}); - const auto k_dram = [&]() { - if constexpr(FmhaPipeline::kKLoadOnce) - { - return pad_tensor_view( - k_dram_naive, - make_tuple(number{}, number{}), - sequence{}); - } - else - { - return pad_tensor_view( - k_dram_naive, - make_tuple(number{}, number{}), - sequence{}); - } - }(); - - const auto kt_dram_naive = - transform_tensor_view(k_dram_naive, - make_tuple(make_pass_through_transform(kargs.hdim_q), - make_pass_through_transform(kargs.seqlen_k)), - make_tuple(sequence<1>{}, sequence<0>{}), - make_tuple(sequence<0>{}, sequence<1>{})); - const auto kt_dram = [&]() { - if constexpr(FmhaPipeline::kKTLoadOnce) - { - return pad_tensor_view( - kt_dram_naive, - make_tuple(number{}, number{}), - sequence{}); - } - else - { - return pad_tensor_view( - kt_dram_naive, - make_tuple(number{}, number{}), - sequence{}); - } - }(); + const auto k_dram = pad_tensor_view( + k_dram_naive, + make_tuple(number{}, number{}), + sequence{}); const auto v_dram = [&]() { const auto v_dram_naive = make_naive_tensor_view( @@ -756,20 +736,10 @@ struct FmhaBwdDQDKDVKernel make_tuple(kargs.stride_v, 1), number{}, number<1>{}); - if constexpr(FmhaPipeline::kVLoadOnce) - { - return pad_tensor_view( - v_dram_naive, - make_tuple(number{}, number{}), - sequence{}); - } - else - { - return pad_tensor_view( - v_dram_naive, - make_tuple(number{}, number{}), - sequence{}); - } + return pad_tensor_view( + v_dram_naive, + make_tuple(number{}, number{}), + sequence{}); }(); const auto lse_dram = [&]() { @@ -792,145 +762,89 @@ struct FmhaBwdDQDKDVKernel make_tuple(kargs.stride_do, 1), number{}, number<1>{}); - const auto do_dram = [&]() { - if constexpr(FmhaPipeline::kOGradLoadOnce) - { - return pad_tensor_view( - do_dram_naive, - make_tuple(number{}, number{}), - sequence{}); - } - else - { - return pad_tensor_view( - do_dram_naive, - make_tuple(number{}, number{}), - sequence{}); - } - }(); - - const auto dot_dram_naive = - transform_tensor_view(do_dram_naive, - make_tuple(make_pass_through_transform(kargs.hdim_v), - make_pass_through_transform(kargs.seqlen_q)), - make_tuple(sequence<1>{}, sequence<0>{}), - make_tuple(sequence<0>{}, sequence<1>{})); - const auto dot_dram = [&]() { - if constexpr(FmhaPipeline::kOGradTLoadOnce) - { - return pad_tensor_view( - dot_dram_naive, - make_tuple(number{}, number{}), - sequence{}); - } - else - { - return pad_tensor_view( - dot_dram_naive, - make_tuple(number{}, number{}), - sequence{}); - } - }(); - - auto dq_dram = [&]() { - const auto dq_dram_naive = make_naive_tensor_view( - dq_ptr, - make_tuple(kargs.seqlen_q, kargs.hdim_q), - make_tuple(kargs.stride_q, 1), - number{}, - number<1>{}); - - return pad_tensor_view( - dq_dram_naive, - make_tuple(number{}, number{}), - sequence{}); - }(); + const auto do_dram = pad_tensor_view( + do_dram_naive, + make_tuple(number{}, number{}), + sequence{}); auto q_dram_window = make_tile_window( q_dram, - [&]() { - if constexpr(FmhaPipeline::kQLoadOnce) - return make_tuple(number{}, - number{}); - else - return make_tuple(number{}, number{}); - }(), + make_tuple(number{}, number{}), {0, 0}); - auto qt_dram_window = - make_tile_window(qt_dram, - [&]() { - if constexpr(FmhaPipeline::kQTLoadOnce) - return make_tuple(number{}, - number{}); - else - return make_tuple(number{}, - number{}); - }(), - {0, 0}); - auto k_dram_window = make_tile_window( k_dram, - [&]() { - if constexpr(FmhaPipeline::kKLoadOnce) - return make_tuple(number{}, - number{}); - else - return make_tuple(number{}, number{}); - }(), + make_tuple(number{}, number{}), {i_n0, 0}); - auto kt_dram_window = - make_tile_window(kt_dram, - [&]() { - if constexpr(FmhaPipeline::kKTLoadOnce) - return make_tuple(number{}, - number{}); - else - return make_tuple(number{}, - number{}); - }(), - {0, i_n0}); - auto v_dram_window = make_tile_window( v_dram, - [&]() { - if constexpr(FmhaPipeline::kVLoadOnce) - return make_tuple(number{}, - number{}); - else - return make_tuple(number{}, number{}); - }(), + make_tuple(number{}, number{}), {i_n0, 0}); auto do_dram_window = make_tile_window( do_dram, - [&]() { - if constexpr(FmhaPipeline::kOGradLoadOnce) - return make_tuple(number{}, - number{}); - else - return make_tuple(number{}, number{}); - }(), + make_tuple(number{}, number{}), {0, 0}); - auto dot_dram_window = - make_tile_window(dot_dram, - [&]() { - if constexpr(FmhaPipeline::kOGradTLoadOnce) - return make_tuple(number{}, - number{}); - else - return make_tuple(number{}, - number{}); - }(), - {0, 0}); + auto dq_dram_window = [&, i_tile_n_ = i_tile_n, i_nhead_ = i_nhead]() { + if constexpr(kIsDeterministic) + { + AccDataType* dq_acc_ptr = + reinterpret_cast(kargs.dq_acc_ptr) + + static_cast(i_nhead_) * kargs.nhead_stride_dq_acc + + static_cast(i_tile_n_) * kargs.split_stride_dq_acc + + batch_offset_dq_acc; - auto dq_dram_window = make_tile_window( - dq_dram, - make_tuple(number{}, number{}), - {0, 0}); + auto dq_acc_dram = [&]() { + const auto dq_acc_dram_naive = + make_naive_tensor_view( + dq_acc_ptr, + make_tuple(kargs.seqlen_q, kargs.hdim_q), + make_tuple(kargs.stride_dq_acc, 1), + number{}, + number<1>{}); + + return pad_tensor_view( + dq_acc_dram_naive, + make_tuple(number{}, number{}), + sequence{}); + }(); + + return make_tile_window( + dq_acc_dram, + make_tuple(number{}, number{}), + {0, 0}); + } + else + { + AccDataType* dq_acc_ptr = + reinterpret_cast(kargs.dq_acc_ptr) + + static_cast(i_nhead_) * kargs.nhead_stride_dq_acc + + batch_offset_dq_acc; + + auto dq_acc_dram = [&]() { + const auto dq_acc_dram_naive = + make_naive_tensor_view( + dq_acc_ptr, + make_tuple(kargs.seqlen_q, kargs.hdim_q), + make_tuple(kargs.stride_dq_acc, 1), + number{}, + number<1>{}); + + return pad_tensor_view( + dq_acc_dram_naive, + make_tuple(number{}, number{}), + sequence{}); + }(); + + return make_tile_window( + dq_acc_dram, + make_tuple(number{}, number{}), + {0, 0}); + } + }(); auto lse_dram_window = make_tile_window(lse_dram, make_tuple(number{}), {0}); @@ -1008,9 +922,7 @@ struct FmhaBwdDQDKDVKernel // TODO: how to use s_read? AccDataType slope = *(reinterpret_cast(kargs.alibi_slope_ptr) + i_batch_ * kargs.alibi_slope_stride + i_nhead_); -#if CK_TILE_FMHA_FWD_FAST_EXP2 slope *= ck_tile::log2e_v<>; -#endif if constexpr(kHasMask) { return make_alibi_from_lr_mask(slope, @@ -1033,35 +945,34 @@ struct FmhaBwdDQDKDVKernel }(); // dropout - float rp_undrop = 1; - float scale_rp_undrop = 1; - uint8_t p_undrop_in_uint8_t = std::numeric_limits::max(); - uint64_t drop_seed = 0; - uint64_t drop_offset = 0; - bool is_store_randval = false; - + float rp_undrop = 1; + float scale_rp_undrop = 1; if constexpr(kHasDropout) { - rp_undrop = kargs.rp_undrop; - scale_rp_undrop = kargs.scale_rp_undrop; - p_undrop_in_uint8_t = kargs.p_undrop_in_uint8_t; - drop_seed = kargs.drop_seed; - drop_offset = kargs.drop_offset; - is_store_randval = kargs.is_store_randval; + rp_undrop = kargs.rp_undrop; + scale_rp_undrop = kargs.scale_rp_undrop; } - BlockDropout dropout(i_batch, - i_nhead, - kargs.num_head_q, - drop_seed, - drop_offset, - rp_undrop, - p_undrop_in_uint8_t, - is_store_randval); + auto dropout = [&, i_nhead_ = i_nhead, i_batch_ = i_batch]() { + if constexpr(kHasDropout) + { + return FmhaDropout{i_batch_, + i_nhead_, + kargs.num_head_q, + kargs.drop_seed, + kargs.drop_offset, + kargs.rp_undrop, + kargs.p_undrop_in_uint8_t}; + } + else + { + return FmhaDropout{}; + }; + }(); auto randval_dram_window = [&, i_nhead_ = i_nhead]() { constexpr auto randval_dram_window_lengths = make_tuple(number{}, number{}); - if constexpr(kHasDropout) + if constexpr(kIsStoreRandval) { RandValOutputDataType* rand_val_ptr = reinterpret_cast(kargs.rand_val_ptr) + @@ -1103,14 +1014,11 @@ struct FmhaBwdDQDKDVKernel }(); auto [dk_acc_tile, dv_acc_tile] = FmhaPipeline{}(q_dram_window, - qt_dram_window, k_dram_window, - kt_dram_window, v_dram_window, bias_dram_window, randval_dram_window, do_dram_window, - dot_dram_window, lse_dram_window, d_dram_window, dq_dram_window, @@ -1118,9 +1026,7 @@ struct FmhaBwdDQDKDVKernel mask, position_encoding, kargs.raw_scale, -#if CK_TILE_FMHA_FWD_FAST_EXP2 kargs.scale, -#endif rp_undrop, scale_rp_undrop, smem_ptr, @@ -1169,10 +1075,9 @@ struct FmhaBwdDQDKDVKernel } }; -template +template struct FmhaBwdOGradDotOKernel { - using TilePartitioner = ck_tile::remove_cvref_t; using FmhaBwdOGradDotO = ck_tile::remove_cvref_t; static constexpr ck_tile::index_t kBlockSize = FmhaBwdOGradDotO::kBlockSize; static constexpr ck_tile::index_t kBlockPerCu = FmhaBwdOGradDotO::kBlockPerCu; @@ -1234,13 +1139,13 @@ struct FmhaBwdOGradDotOKernel ck_tile::index_t nhead_stride_do; ck_tile::index_t nhead_stride_o; ck_tile::index_t nhead_stride_d; - ck_tile::index_t batch_stride_d; }; struct FmhaBwdOGradDotOBatchModeKargs : FmhaBwdOGradDotOCommonKargs { ck_tile::index_t batch_stride_do; ck_tile::index_t batch_stride_o; + ck_tile::index_t batch_stride_d; }; struct FmhaBwdOGradDotOGroupModeKargs : FmhaBwdOGradDotOCommonKargs @@ -1278,10 +1183,10 @@ struct FmhaBwdOGradDotOKernel stride_o, nhead_stride_do, nhead_stride_o, - nhead_stride_d, - batch_stride_d}, + nhead_stride_d}, batch_stride_do, - batch_stride_o}; + batch_stride_o, + batch_stride_d}; return kargs; } @@ -1298,8 +1203,7 @@ struct FmhaBwdOGradDotOKernel ck_tile::index_t stride_o, ck_tile::index_t nhead_stride_do, ck_tile::index_t nhead_stride_o, - ck_tile::index_t nhead_stride_d, - ck_tile::index_t batch_stride_d) + ck_tile::index_t nhead_stride_d) { Kargs kargs{{o_ptr, do_ptr, @@ -1311,8 +1215,7 @@ struct FmhaBwdOGradDotOKernel stride_o, nhead_stride_do, nhead_stride_o, - nhead_stride_d, - batch_stride_d}, + nhead_stride_d}, reinterpret_cast(seqstart_q_ptr)}; return kargs; @@ -1321,7 +1224,16 @@ struct FmhaBwdOGradDotOKernel CK_TILE_HOST static constexpr auto GridSize(ck_tile::index_t batch_size_, ck_tile::index_t nhead_, ck_tile::index_t seqlen_q_) { - return TilePartitioner::GridSize(batch_size_, nhead_, seqlen_q_); + return dim3(ck_tile::integer_divide_ceil(seqlen_q_, kM0), nhead_, batch_size_); + } + + CK_TILE_DEVICE static constexpr auto GetTileIndex() + { + const index_t i_block = blockIdx.x; + const index_t i_nhead = blockIdx.y; + const index_t i_batch = blockIdx.z; + + return ck_tile::make_tuple(i_block, i_nhead, i_batch); } CK_TILE_HOST static constexpr auto BlockSize() { return dim3(kBlockSize); } @@ -1331,7 +1243,7 @@ struct FmhaBwdOGradDotOKernel CK_TILE_DEVICE void operator()(Kargs kargs) const { // divide problem - const auto [i_tile_m, i_nhead, i_batch] = TilePartitioner{}(kargs.seqlen_q); + const auto [i_tile_m, i_nhead, i_batch] = GetTileIndex(); const index_t i_m0 = __builtin_amdgcn_readfirstlane(i_tile_m * kM0); @@ -1346,7 +1258,7 @@ struct FmhaBwdOGradDotOKernel batch_offset_o = query_start * kargs.stride_o; batch_offset_do = query_start * kargs.stride_do; - batch_offset_d = static_cast(i_batch) * kargs.batch_stride_d; + batch_offset_d = query_start; // get real # queries & # keys under group mode const auto adjusted_seqstart_q_ptr = kargs.seqstart_q_ptr + i_batch; @@ -1418,4 +1330,315 @@ struct FmhaBwdOGradDotOKernel } }; +template +struct FmhaBwdConvertQGradKernel +{ + using FmhaBwdConvertQGrad = ck_tile::remove_cvref_t; + static constexpr ck_tile::index_t kBlockSize = FmhaBwdConvertQGrad::kBlockSize; + static constexpr ck_tile::index_t kBlockPerCu = FmhaBwdConvertQGrad::kBlockPerCu; + static constexpr ck_tile::index_t kM0 = FmhaBwdConvertQGrad::kM0; + static constexpr ck_tile::index_t kN0 = FmhaBwdConvertQGrad::kN0; + static constexpr ck_tile::index_t kQKHeaddim = FmhaBwdConvertQGrad::kQKHeaddim; + + using AccDataType = ck_tile::remove_cvref_t; + using QGradDataType = ck_tile::remove_cvref_t; + + static constexpr bool kIsGroupMode = FmhaBwdConvertQGrad::kIsGroupMode; + static constexpr bool kPadSeqLenQ = FmhaBwdConvertQGrad::kPadSeqLenQ; + static constexpr bool kPadHeadDimQ = FmhaBwdConvertQGrad::kPadHeadDimQ; + static constexpr bool kIsDeterministic = FmhaBwdConvertQGrad::kIsDeterministic; + + // clang-format off + template struct t2s; + template <> struct t2s { static constexpr const char * name = "fp16"; }; + template <> struct t2s { static constexpr const char * name = "bf16"; }; + // clang-format on + + CK_TILE_HOST static std::string GetName() + { + // sync with generate.py + // clang-format off + + #define _SS_ std::string + #define _TS_ std::to_string + auto pn = [&] () { + std::string n; + if (kPadSeqLenQ) n += "s"; + if (kPadHeadDimQ) n += "d"; + return n.empty() ? n : std::string("p") + n; }(); + return + _SS_("fmha_bwd_convert_dq_d") + _TS_(kQKHeaddim) + "_" + _SS_(t2s::name) + + "_" + (kIsGroupMode ? "group" : "batch") + (kIsDeterministic ? "_deterministic" : "") + "_" + + ("o" + _TS_(kBlockPerCu)) + (pn.empty() ? "" : "_" + pn); + #undef _SS_ + #undef _TS_ + // clang-format on + } + + // to avoid duplicated base class prblem, introduce an template arg + template + struct FmhaBwdConvertQGradEmptyKargs + { + }; + + // kargs use aggregate initializer, so no constructor will provided + // use inheritance to minimize karg size + // user need to use MakeKargs() function to create kargs. + struct FmhaBwdConvertQGradCommonKargs + { + const void* dq_acc_ptr; + void* dq_ptr; + + ck_tile::index_t seqlen_q; + ck_tile::index_t seqlen_k; + ck_tile::index_t hdim_q; + + ck_tile::index_t stride_dq; + ck_tile::index_t stride_dq_acc; + ck_tile::index_t nhead_stride_dq; + ck_tile::index_t nhead_stride_dq_acc; + }; + + struct FmhaBwdConvertQGradDeterministicKargs + { + ck_tile::index_t split_stride_dq_acc = 0; + }; + + struct FmhaBwdConvertQGradBatchModeKargs + : FmhaBwdConvertQGradCommonKargs, + std::conditional_t> + { + ck_tile::index_t batch_stride_dq; + ck_tile::index_t batch_stride_dq_acc; + }; + + struct FmhaBwdConvertQGradGroupModeKargs + : FmhaBwdConvertQGradCommonKargs, + std::conditional_t> + { + const int32_t* seqstart_q_ptr; + const int32_t* seqstart_k_ptr; + }; + + using Kargs = std::conditional_t; + + template + CK_TILE_HOST static constexpr std::enable_if_t + MakeKargs(const void* dq_acc_ptr, + void* dq_ptr, + ck_tile::index_t seqlen_q, + ck_tile::index_t seqlen_k, + ck_tile::index_t hdim_q, + ck_tile::index_t stride_dq, + ck_tile::index_t stride_dq_acc, + ck_tile::index_t nhead_stride_dq, + ck_tile::index_t nhead_stride_dq_acc, + ck_tile::index_t batch_stride_dq, + ck_tile::index_t batch_stride_dq_acc, + ck_tile::index_t split_stride_dq_acc) + { + Kargs kargs{{dq_acc_ptr, + dq_ptr, + seqlen_q, + seqlen_k, + hdim_q, + stride_dq, + stride_dq_acc, + nhead_stride_dq, + nhead_stride_dq_acc}, + {}, + batch_stride_dq, + batch_stride_dq_acc}; + + if constexpr(kIsDeterministic) + { + kargs.split_stride_dq_acc = split_stride_dq_acc; + } + + return kargs; + } + + template + CK_TILE_HOST static constexpr std::enable_if_t + MakeKargs(const void* dq_acc_ptr, + void* dq_ptr, + const void* seqstart_q_ptr, + const void* seqstart_k_ptr, + ck_tile::index_t hdim_q, + ck_tile::index_t stride_dq, + ck_tile::index_t stride_dq_acc, + ck_tile::index_t nhead_stride_dq, + ck_tile::index_t nhead_stride_dq_acc, + ck_tile::index_t split_stride_dq_acc) + { + Kargs kargs{{dq_acc_ptr, + dq_ptr, + -1, // seqlen will be updated by another pointer + -1, // + hdim_q, + stride_dq, + stride_dq_acc, + nhead_stride_dq, + nhead_stride_dq_acc}, + {}, + reinterpret_cast(seqstart_q_ptr), + reinterpret_cast(seqstart_k_ptr)}; + + if constexpr(kIsDeterministic) + { + kargs.split_stride_dq_acc = split_stride_dq_acc; + } + + return kargs; + } + + CK_TILE_HOST static constexpr auto + GridSize(ck_tile::index_t batch_size_, ck_tile::index_t nhead_, ck_tile::index_t seqlen_q_) + { + return dim3(ck_tile::integer_divide_ceil(seqlen_q_, kM0), nhead_, batch_size_); + } + + CK_TILE_DEVICE static constexpr auto GetTileIndex() + { + const index_t i_block = blockIdx.x; + const index_t i_nhead = blockIdx.y; + const index_t i_batch = blockIdx.z; + + return ck_tile::make_tuple(i_block, i_nhead, i_batch); + } + + CK_TILE_HOST static constexpr auto BlockSize() { return dim3(kBlockSize); } + + CK_TILE_HOST_DEVICE static constexpr ck_tile::index_t GetSmemSize() { return 0; } + + CK_TILE_DEVICE void operator()(Kargs kargs) const + { + // divide problem + const auto [i_tile_m, i_nhead, i_batch] = GetTileIndex(); + + const index_t i_m0 = __builtin_amdgcn_readfirstlane(i_tile_m * kM0); + + long_index_t batch_offset_dq = 0; + long_index_t batch_offset_dq_acc = 0; + if constexpr(kIsGroupMode) + { + // get starting offset for each batch + const long_index_t query_start = kargs.seqstart_q_ptr[i_batch]; + batch_offset_dq = query_start * kargs.stride_dq; + batch_offset_dq_acc = query_start * kargs.stride_dq_acc; + + // get real # queries & # keys under group mode + const auto adjusted_seqstart_q_ptr = kargs.seqstart_q_ptr + i_batch; + kargs.seqlen_q = adjusted_seqstart_q_ptr[1] - adjusted_seqstart_q_ptr[0]; + if constexpr(kIsDeterministic) + { + const auto adjusted_seqstart_k_ptr = kargs.seqstart_k_ptr + i_batch; + kargs.seqlen_k = adjusted_seqstart_k_ptr[1] - adjusted_seqstart_k_ptr[0]; + } + // # of required blocks is different in each groups, terminate unnecessary blocks + // earlier + if(kargs.seqlen_q <= i_m0) + { + return; + } + } + else + { + batch_offset_dq = static_cast(i_batch) * kargs.batch_stride_dq; + batch_offset_dq_acc = static_cast(i_batch) * kargs.batch_stride_dq_acc; + } + + // for simplicity, batch stride we just modify the pointer + QGradDataType* dq_ptr = reinterpret_cast(kargs.dq_ptr) + + static_cast(i_nhead) * kargs.nhead_stride_dq + + batch_offset_dq; + + // dQAcc/dQ DRAM and DRAM window + const auto dq_acc_dram = [&, i_nhead_ = i_nhead]() { + if constexpr(kIsDeterministic) + { + const AccDataType* dq_acc_ptr = + reinterpret_cast(kargs.dq_acc_ptr) + + static_cast(i_nhead_) * (kargs.nhead_stride_dq_acc) + + batch_offset_dq_acc; + + const index_t nsplits = ck_tile::integer_divide_ceil(kargs.seqlen_k, kN0); + + auto dq_acc_dram_naive = make_naive_tensor_view( + dq_acc_ptr, + make_tuple(nsplits, kargs.seqlen_q, kargs.hdim_q), + make_tuple(kargs.split_stride_dq_acc, kargs.stride_dq_acc, 1), + number{}, + number<1>{}); + return pad_tensor_view(dq_acc_dram_naive, + make_tuple(number<1>{}, number{}, number{}), + sequence{}); + } + else + { + const AccDataType* dq_acc_ptr = + reinterpret_cast(kargs.dq_acc_ptr) + + static_cast(i_nhead_) * (kargs.nhead_stride_dq_acc) + + batch_offset_dq_acc; + + auto dq_acc_dram_naive = make_naive_tensor_view( + dq_acc_ptr, + make_tuple(kargs.seqlen_q, kargs.hdim_q), + make_tuple(kargs.stride_dq_acc, 1), + number{}, + number<1>{}); + return pad_tensor_view(dq_acc_dram_naive, + make_tuple(number{}, number{}), + sequence{}); + } + }(); + + auto dq_dram = [&]() { + auto dq_dram_naive = make_naive_tensor_view( + dq_ptr, + make_tuple(kargs.seqlen_q, kargs.hdim_q), + make_tuple(kargs.stride_dq, 1), + number{}, + number<1>{}); + return pad_tensor_view(dq_dram_naive, + make_tuple(number{}, number{}), + sequence{}); + }(); + + auto dq_acc_dram_window = [&]() { + if constexpr(kIsDeterministic) + { + return make_tile_window( + dq_acc_dram, + make_tuple(number<1>{}, number{}, number{}), + {0, i_m0, 0}); + } + else + { + return make_tile_window( + dq_acc_dram, make_tuple(number{}, number{}), {i_m0, 0}); + } + }(); + + auto dq_dram_window = + make_tile_window(dq_dram, make_tuple(number{}, number{}), {i_m0, 0}); + + if constexpr(kIsDeterministic) + { + const index_t nsplits = ck_tile::integer_divide_ceil(kargs.seqlen_k, kN0); + FmhaBwdConvertQGrad{}(dq_acc_dram_window, dq_dram_window, nsplits); + } + else + { + FmhaBwdConvertQGrad{}(dq_acc_dram_window, dq_dram_window); + } + } +}; + } // namespace ck_tile diff --git a/include/ck_tile/ops/fmha/kernel/fmha_bwd_tile_partitioner.hpp b/include/ck_tile/ops/fmha/kernel/fmha_bwd_tile_partitioner.hpp deleted file mode 100644 index bc875b8e5a..0000000000 --- a/include/ck_tile/ops/fmha/kernel/fmha_bwd_tile_partitioner.hpp +++ /dev/null @@ -1,54 +0,0 @@ -// SPDX-License-Identifier: MIT -// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved. - -#pragma once - -#include "ck_tile/core.hpp" - -namespace ck_tile { - -template -struct FmhaBwdTilePartitioner -{ - using BlockFmhaShape = ck_tile::remove_cvref_t; - - static constexpr ck_tile::index_t kN0 = BlockFmhaShape::kN0; - - CK_TILE_HOST static constexpr auto - GridSize(ck_tile::index_t batch_size_, ck_tile::index_t nhead_, ck_tile::index_t seqlen_k_) - { - // TODO: this may need tuning - return dim3(ck_tile::integer_divide_ceil(seqlen_k_, kN0), nhead_, batch_size_); - } - - CK_TILE_DEVICE auto operator()(ck_tile::index_t /*seqlen_k*/) - { - const index_t i_block = blockIdx.x; - const index_t i_nhead = blockIdx.y; - const index_t i_batch = blockIdx.z; - - return ck_tile::make_tuple(i_block, i_nhead, i_batch); - } -}; - -template -struct FmhaBwdOGradDotOTilePartitioner -{ - CK_TILE_HOST static constexpr auto - GridSize(ck_tile::index_t batch_size_, ck_tile::index_t nhead_, ck_tile::index_t seqlen_q_) - { - // TODO: this may need tuning - return dim3(ck_tile::integer_divide_ceil(seqlen_q_, kBlockSize), nhead_, batch_size_); - } - - CK_TILE_DEVICE auto operator()(ck_tile::index_t /*seqlen_q*/) - { - const index_t i_block = blockIdx.x; - const index_t i_nhead = blockIdx.y; - const index_t i_batch = blockIdx.z; - - return ck_tile::make_tuple(i_block, i_nhead, i_batch); - } -}; - -} // namespace ck_tile 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 5ecc3a4d80..49ef7bf6d9 100644 --- a/include/ck_tile/ops/fmha/kernel/fmha_fwd_kernel.hpp +++ b/include/ck_tile/ops/fmha/kernel/fmha_fwd_kernel.hpp @@ -86,7 +86,7 @@ struct FmhaFwdKernel "w" + _TS_(gwt::at(ck_tile::number<0>{})) + "x" + _TS_(gwt::at(ck_tile::number<1>{})) + "x" + _TS_(gwt::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)) + + (BiasEnum == BlockAttentionBiasEnum::NO_BIAS ? _SS_("") : (_SS_("_") + BlockAttentionBiasEnumToStr::name)) + (kHasMask ? "_" + _SS_(FmhaMask::name) : "") + (kStoreLSE ? "_lse" : "" ) + (kHasDropout ? "_dropout" : "" ) + (kDoFp8StaticQuant ? "_squant" : "" ); #undef _SS_ #undef _TS_ @@ -387,7 +387,6 @@ struct FmhaFwdKernel ck_tile::index_t nhead_stride_randval, ck_tile::index_t nhead_stride_lse, ck_tile::index_t nhead_stride_o, - ck_tile::index_t batch_stride_lse, ck_tile::index_t window_size_left, ck_tile::index_t window_size_right, ck_tile::index_t mask_type, @@ -448,7 +447,6 @@ struct FmhaFwdKernel { kargs.lse_ptr = lse_ptr; kargs.nhead_stride_lse = nhead_stride_lse; - kargs.batch_stride_lse = batch_stride_lse; } if constexpr(kDoFp8StaticQuant) { @@ -524,7 +522,7 @@ struct FmhaFwdKernel } if constexpr(kStoreLSE) { - batch_offset_lse = static_cast(i_batch) * kargs.batch_stride_lse; + batch_offset_lse = query_start; } if constexpr(kHasDropout) { 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 6f4313d5b6..e2c7db3e1b 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 @@ -55,7 +55,7 @@ struct FmhaFwdSplitKVCombineKernel (kBlockPerCuInput == -1 ? "" : ("o" + _TS_(kBlockPerCu) + "_")) + _SS_(FmhaPipeline::name) + (pn.empty() ? "" : "_" + pn) + - (kStoreLSE ? "_lse" : "" ) + + (kStoreLSE ? "_lse" : "" ) + (kDoFp8StaticQuant ? "_squant" : "" ); #undef _SS_ #undef _TS_ @@ -91,7 +91,6 @@ struct FmhaFwdSplitKVCombineKernel ck_tile::index_t nhead_stride_o_acc; ck_tile::index_t nhead_stride_o; - ck_tile::index_t batch_stride_lse_acc; ck_tile::index_t batch_stride_o_acc; ck_tile::index_t split_stride_lse_acc; @@ -116,6 +115,7 @@ struct FmhaFwdSplitKVCombineKernel std::conditional_t> { ck_tile::index_t batch_stride_o; + ck_tile::index_t batch_stride_lse_acc; }; struct GroupModeKargs @@ -166,13 +166,13 @@ struct FmhaFwdSplitKVCombineKernel nhead_stride_lse_acc, nhead_stride_o_acc, nhead_stride_o, - batch_stride_lse_acc, batch_stride_o_acc, split_stride_lse_acc, split_stride_o_acc}, // args for common karg {}, // placeholder for lse {}, // placeholder for fp8_static_quant args - batch_stride_o}; + batch_stride_o, + batch_stride_lse_acc}; if constexpr(kStoreLSE) { @@ -206,9 +206,7 @@ struct FmhaFwdSplitKVCombineKernel ck_tile::index_t nhead_stride_o_acc, ck_tile::index_t nhead_stride_lse, ck_tile::index_t nhead_stride_o, - ck_tile::index_t batch_stride_lse_acc, ck_tile::index_t batch_stride_o_acc, - ck_tile::index_t batch_stride_lse, ck_tile::index_t split_stride_lse_acc, ck_tile::index_t split_stride_o_acc) { @@ -225,7 +223,6 @@ struct FmhaFwdSplitKVCombineKernel nhead_stride_lse_acc, nhead_stride_o_acc, nhead_stride_o, - batch_stride_lse_acc, batch_stride_o_acc, split_stride_lse_acc, split_stride_o_acc}, // args for common karg @@ -237,7 +234,6 @@ struct FmhaFwdSplitKVCombineKernel { kargs.lse_ptr = lse_ptr; kargs.nhead_stride_lse = nhead_stride_lse; - kargs.batch_stride_lse = batch_stride_lse; } if constexpr(kDoFp8StaticQuant) { @@ -274,24 +270,25 @@ struct FmhaFwdSplitKVCombineKernel const index_t i_m0 = __builtin_amdgcn_readfirstlane(i_tile_m * FmhaPipeline::kM0); const index_t i_n1 = __builtin_amdgcn_readfirstlane(i_tile_n * FmhaPipeline::kN1); - const long_index_t batch_offset_lse_acc = - static_cast(i_batch) * kargs.batch_stride_lse_acc; const long_index_t batch_offset_o_acc = static_cast(i_batch) * kargs.batch_stride_o_acc; - long_index_t batch_offset_lse = 0; - long_index_t batch_offset_o = 0; - if constexpr(kStoreLSE) - { - batch_offset_lse = static_cast(i_batch) * kargs.batch_stride_lse; - } + long_index_t batch_offset_lse_acc = 0; + long_index_t batch_offset_lse = 0; + long_index_t batch_offset_o = 0; if constexpr(kIsGroupMode) { // get starting offset for each batch const long_index_t query_start = kargs.seqstart_q_ptr[i_batch]; - batch_offset_o = query_start * kargs.row_stride_o; + batch_offset_o = query_start * kargs.row_stride_o; + batch_offset_lse_acc = query_start; + + if constexpr(kStoreLSE) + { + batch_offset_lse = query_start; + } // get real # queries & # keys under group mode const auto adjusted_seqstart_q_ptr = kargs.seqstart_q_ptr + i_batch; @@ -306,7 +303,13 @@ struct FmhaFwdSplitKVCombineKernel } else { - batch_offset_o = static_cast(i_batch) * kargs.batch_stride_o; + batch_offset_o = static_cast(i_batch) * kargs.batch_stride_o; + batch_offset_lse_acc = static_cast(i_batch) * kargs.batch_stride_lse_acc; + + if constexpr(kStoreLSE) + { + batch_offset_lse = static_cast(i_batch) * kargs.batch_stride_lse; + } } // for simplicity, batch stride we just modify the pointer 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 45ed185ada..36c10db79c 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 @@ -85,7 +85,7 @@ struct FmhaFwdSplitKVKernel "w" + _TS_(gwt::at(ck_tile::number<0>{})) + "x" + _TS_(gwt::at(ck_tile::number<1>{})) + "x" + _TS_(gwt::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)) + + (BiasEnum == BlockAttentionBiasEnum::NO_BIAS ? _SS_("") : (_SS_("_") + BlockAttentionBiasEnumToStr::name)) + (kHasMask ? "_" + _SS_(FmhaMask::name) : "") + (kHasDropout ? "_dropout" : "" ) + (kDoFp8StaticQuant ? "_squant" : "" ); #undef _SS_ #undef _TS_ @@ -136,7 +136,6 @@ struct FmhaFwdSplitKVKernel ck_tile::index_t nhead_stride_lse_acc; ck_tile::index_t nhead_stride_o_acc; - ck_tile::index_t batch_stride_lse_acc; ck_tile::index_t batch_stride_o_acc; ck_tile::index_t split_stride_lse_acc; @@ -216,6 +215,7 @@ struct FmhaFwdSplitKVKernel ck_tile::index_t batch_stride_q; ck_tile::index_t batch_stride_k; ck_tile::index_t batch_stride_v; + ck_tile::index_t batch_stride_lse_acc; }; struct GroupModeKargs @@ -313,7 +313,6 @@ struct FmhaFwdSplitKVKernel nhead_stride_v, nhead_stride_lse_acc, nhead_stride_o_acc, - batch_stride_lse_acc, batch_stride_o_acc, split_stride_lse_acc, split_stride_o_acc}, // args for common karg @@ -323,7 +322,8 @@ struct FmhaFwdSplitKVKernel {}, // placeholder for dropout batch_stride_q, batch_stride_k, - batch_stride_v}; + batch_stride_v, + batch_stride_lse_acc}; if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS) { @@ -394,7 +394,6 @@ struct FmhaFwdSplitKVKernel ck_tile::index_t nhead_stride_randval, ck_tile::index_t nhead_stride_lse_acc, ck_tile::index_t nhead_stride_o_acc, - ck_tile::index_t batch_stride_lse_acc, ck_tile::index_t batch_stride_o_acc, ck_tile::index_t split_stride_lse_acc, ck_tile::index_t split_stride_o_acc, @@ -433,7 +432,6 @@ struct FmhaFwdSplitKVKernel nhead_stride_v, nhead_stride_lse_acc, nhead_stride_o_acc, - batch_stride_lse_acc, batch_stride_o_acc, split_stride_lse_acc, split_stride_o_acc}, // args for common karg @@ -511,8 +509,7 @@ struct FmhaFwdSplitKVKernel long_index_t batch_offset_v = 0; long_index_t batch_offset_bias = 0; long_index_t batch_offset_randval = 0; - const long_index_t batch_offset_lse_acc = - static_cast(i_batch) * kargs.batch_stride_lse_acc; + long_index_t batch_offset_lse_acc = 0; const long_index_t batch_offset_o_acc = static_cast(i_batch) * kargs.batch_stride_o_acc; @@ -522,8 +519,9 @@ struct FmhaFwdSplitKVKernel const long_index_t query_start = kargs.seqstart_q_ptr[i_batch]; const long_index_t key_start = kargs.seqstart_k_ptr[i_batch]; - batch_offset_q = query_start * kargs.stride_q; - batch_offset_k = key_start * kargs.stride_k; + batch_offset_q = query_start * kargs.stride_q; + batch_offset_k = key_start * kargs.stride_k; + batch_offset_lse_acc = query_start; if constexpr(std::is_same_v) { batch_offset_v = key_start * kargs.stride_v; @@ -564,9 +562,10 @@ struct FmhaFwdSplitKVKernel } else { - batch_offset_q = static_cast(i_batch) * kargs.batch_stride_q; - batch_offset_k = static_cast(i_batch) * kargs.batch_stride_k; - batch_offset_v = static_cast(i_batch) * kargs.batch_stride_v; + batch_offset_q = static_cast(i_batch) * kargs.batch_stride_q; + batch_offset_k = static_cast(i_batch) * kargs.batch_stride_k; + batch_offset_v = static_cast(i_batch) * kargs.batch_stride_v; + batch_offset_lse_acc = static_cast(i_batch) * kargs.batch_stride_lse_acc; if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS) { batch_offset_bias = static_cast(i_batch) * kargs.batch_stride_bias; diff --git a/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_convert_dq.hpp b/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_convert_dq.hpp new file mode 100644 index 0000000000..3da1104169 --- /dev/null +++ b/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_convert_dq.hpp @@ -0,0 +1,141 @@ +// 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_bwd_pipeline_default_policy.hpp" + +namespace ck_tile { + +template +struct BlockFmhaBwdConvertQGrad +{ + using AccDataType = remove_cvref_t; + using QGradDataType = remove_cvref_t; + + static constexpr index_t kM0 = Problem::kM0; + static constexpr index_t kN0 = Problem::kN0; + + static constexpr index_t kBlockPerCu = Problem::kBlockPerCu; + static constexpr index_t kBlockSize = Problem::kBlockSize; + static constexpr index_t kQKHeaddim = Problem::kQKHeaddim; + + static constexpr bool kIsGroupMode = Problem::kIsGroupMode; + static constexpr bool kPadSeqLenQ = Problem::kPadSeqLenQ; + static constexpr bool kPadHeadDimQ = Problem::kPadHeadDimQ; + static constexpr bool kIsDeterministic = Problem::kIsDeterministic; + + static constexpr index_t kAlignmentQGradAcc = + kPadHeadDimQ ? 1 : Policy::template GetAlignmentPostQGradAcc(); + static constexpr index_t kAlignmentQGrad = + kPadHeadDimQ ? 1 : Policy::template GetAlignmentPostQGrad(); + + CK_TILE_HOST_DEVICE static constexpr ck_tile::index_t GetSmemSize() { return 0; } + + // Convert only + template + CK_TILE_HOST_DEVICE void + operator()(const QGradAccDramBlockWindowTmp& dq_acc_dram_block_window_tmp, + QGradDramBlockWindowTmp& dq_dram_block_window_tmp) const + { + static_assert( + std::is_same_v> && + std::is_same_v>, + "wrong!"); + + static_assert(kM0 == QGradDramBlockWindowTmp{}.get_window_lengths()[number<0>{}], "wrong!"); + + auto dq_acc_dram_window = + make_tile_window(dq_acc_dram_block_window_tmp.get_bottom_tensor_view(), + dq_acc_dram_block_window_tmp.get_window_lengths(), + dq_acc_dram_block_window_tmp.get_window_origin(), + Policy::template MakePostQGradDramTileDistribution()); + + auto dq_acc = load_tile(dq_acc_dram_window); + const auto dq = cast_tile(dq_acc); + + store_tile(dq_dram_block_window_tmp, dq); + } + + // Reduce + Convert + template + CK_TILE_HOST_DEVICE void + operator()(const QGradAccDramBlockWindowTmp& dq_acc_dram_block_window_tmp, + QGradDramBlockWindowTmp& dq_dram_block_window_tmp, + index_t nsplits) const + { + static_assert( + std::is_same_v> && + std::is_same_v>, + "wrong!"); + + static_assert(kM0 == QGradDramBlockWindowTmp{}.get_window_lengths()[number<0>{}], "wrong!"); + + auto dq_acc_dram_window = + make_tile_window(dq_acc_dram_block_window_tmp.get_bottom_tensor_view(), + dq_acc_dram_block_window_tmp.get_window_lengths(), + dq_acc_dram_block_window_tmp.get_window_origin(), + Policy::template MakePostQGradAccDramTileDistribution()); + + auto dq_acc = decltype(load_tile(dq_acc_dram_window)){}; + clear_tile(dq_acc); + + constexpr auto dq_acc_spans = decltype(dq_acc)::get_distributed_spans(); + index_t i_total_loops = 0; + auto dq_acc_buf = load_tile(dq_acc_dram_window); + move_tile_window(dq_acc_dram_window, {1, 0, 0}); + + do + { + sweep_tile_span(dq_acc_spans[number<0>{}], [&](auto idx0) { + sweep_tile_span(dq_acc_spans[number<1>{}], [&](auto idx1) { + sweep_tile_span(dq_acc_spans[number<2>{}], [&](auto idx2) { + constexpr auto n_i_j_idx = make_tuple(idx0, idx1, idx2); + dq_acc(n_i_j_idx) += dq_acc_buf(n_i_j_idx); + }); + }); + }); + + dq_acc_buf = load_tile(dq_acc_dram_window); + move_tile_window(dq_acc_dram_window, {1, 0, 0}); + + i_total_loops += 1; + } while(i_total_loops < (nsplits - 1)); + + sweep_tile_span(dq_acc_spans[number<0>{}], [&](auto idx0) { + sweep_tile_span(dq_acc_spans[number<1>{}], [&](auto idx1) { + sweep_tile_span(dq_acc_spans[number<2>{}], [&](auto idx2) { + constexpr auto n_i_j_idx = make_tuple(idx0, idx1, idx2); + dq_acc(n_i_j_idx) += dq_acc_buf(n_i_j_idx); + }); + }); + }); + + // declare dq + constexpr auto dq_converted_dstr = + Policy::template MakePostQGradAccDramTileDistribution(); + auto dq_converted = make_static_distributed_tensor(dq_converted_dstr); + + sweep_tile_span(dq_acc_spans[number<0>{}], [&](auto idx0) { + sweep_tile_span(dq_acc_spans[number<1>{}], [&](auto idx1) { + sweep_tile_span(dq_acc_spans[number<2>{}], [&](auto idx2) { + constexpr auto n_i_j_idx = make_tuple(idx0, idx1, idx2); + dq_converted(n_i_j_idx) = type_convert(dq_acc[n_i_j_idx]); + }); + }); + }); + + constexpr auto dq_dstr = Policy::template MakePostQGradDramTileDistribution(); + auto dq = make_static_distributed_tensor(dq_dstr); + dq.get_thread_buffer() = dq_converted.get_thread_buffer(); + + store_tile(dq_dram_block_window_tmp, dq); + } +}; + +} // namespace ck_tile diff --git a/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dot_do_o.hpp b/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dot_do_o.hpp index f189937038..c38779d1d2 100644 --- a/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dot_do_o.hpp +++ b/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dot_do_o.hpp @@ -4,11 +4,11 @@ #pragma once #include "ck_tile/core.hpp" -#include "ck_tile/ops/fmha/pipeline/block_fmha_bwd_dot_do_o_default_policy.hpp" +#include "ck_tile/ops/fmha/pipeline/block_fmha_bwd_pipeline_default_policy.hpp" namespace ck_tile { -template +template struct BlockFmhaBwdOGradDotO { using ODataType = remove_cvref_t; @@ -26,7 +26,7 @@ struct BlockFmhaBwdOGradDotO static constexpr index_t kAlignmentO = kPadHeadDimV ? 1 : Policy::template GetAlignmentO(); static constexpr index_t kAlignmentOGrad = - kPadHeadDimV ? 1 : Policy::template GetAlignmentOGrad(); + kPadHeadDimV ? 1 : Policy::template GetAlignmentO(); CK_TILE_HOST_DEVICE static constexpr ck_tile::index_t GetSmemSize() { return 0; } diff --git a/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dot_do_o_default_policy.hpp b/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dot_do_o_default_policy.hpp deleted file mode 100644 index 7843ab33a1..0000000000 --- a/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dot_do_o_default_policy.hpp +++ /dev/null @@ -1,20 +0,0 @@ -// SPDX-License-Identifier: MIT -// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved. - -#pragma once - -#include "ck_tile/ops/fmha/pipeline/block_fmha_bwd_pipeline_default_policy.hpp" - -namespace ck_tile { - -// These templates are not used here. -using BlockFmhaBwdOGradDotODefaultPolicy = - BlockFmhaBwdPipelineDefaultPolicy; - -} // namespace ck_tile diff --git a/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_kr_ktr_vr.hpp b/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_kr_ktr_vr.hpp new file mode 100644 index 0000000000..131729992b --- /dev/null +++ b/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_kr_ktr_vr.hpp @@ -0,0 +1,782 @@ +// 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/block/block_dropout.hpp" +#include "ck_tile/ops/fmha/pipeline/block_fmha_bwd_pipeline_default_policy.hpp" +#include "ck_tile/ops/reduce/block/block_reduce.hpp" + +namespace ck_tile { + +template +struct BlockFmhaBwdDQDKDVPipelineKRKTRVR +{ + using QDataType = remove_cvref_t; + using KDataType = remove_cvref_t; + using VDataType = remove_cvref_t; + using GemmDataType = remove_cvref_t; + using BiasDataType = remove_cvref_t; + using LSEDataType = remove_cvref_t; + using AccDataType = remove_cvref_t; + using DDataType = remove_cvref_t; + using RandValOutputDataType = remove_cvref_t; + using ODataType = remove_cvref_t; + using OGradDataType = remove_cvref_t; + using QGradDataType = remove_cvref_t; + using KGradDataType = remove_cvref_t; + using VGradDataType = remove_cvref_t; + using BiasGradDataType = remove_cvref_t; + using FmhaMask = remove_cvref_t; + using FmhaDropout = remove_cvref_t; + using HotLoopScheduler = typename Policy::template HotLoopScheduler; + + using BlockFmhaShape = remove_cvref_t; + + static constexpr index_t kBlockPerCu = Problem::kBlockPerCu; + 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 kK1 = BlockFmhaShape::kK1; + static constexpr index_t kK2 = BlockFmhaShape::kK2; + static constexpr index_t kK3 = BlockFmhaShape::kK3; + static constexpr index_t kK4 = BlockFmhaShape::kK4; + static constexpr index_t kQKHeaddim = BlockFmhaShape::kQKHeaddim; + static constexpr index_t kVHeaddim = BlockFmhaShape::kVHeaddim; + + 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 kHasBiasGrad = Problem::kHasBiasGrad; + static constexpr bool kIsDeterministic = Problem::kIsDeterministic; + + // 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 = + kPadHeadDimV ? 1 : Policy::template GetAlignmentV(); + static constexpr index_t kAlignmentOGrad = + kPadHeadDimV ? 1 : Policy::template GetAlignmentOGrad(); + static constexpr index_t kAlignmentQGrad = 1; + static constexpr index_t kAlignmentKGrad = + kPadHeadDimQ ? 1 : Policy::template GetAlignmentKGrad(); + static constexpr index_t kAlignmentVGrad = + kPadHeadDimV ? 1 : Policy::template GetAlignmentVGrad(); + static constexpr index_t kAlignmentBias = + kPadSeqLenK ? 1 : Policy::template GetTransposedAlignmentBias(); + + static constexpr const char* name = "kr_ktr_vr"; + + 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, + const KDramBlockWindowTmp& k_dram_block_window_tmp, + const VDramBlockWindowTmp& v_dram_block_window_tmp, + const BiasDramBlockWindowTmp& bias_dram_block_window_tmp, + const RandValDramBlockWindowTmp& randval_dram_block_window_tmp, + const OGradDramBlockWindowTmp& do_dram_block_window_tmp, + const LSEDramBlockWindowTmp& lse_dram_block_window_tmp, + const DDramBlockWindowTmp& d_dram_block_window_tmp, + const QGradDramBlockWindowTmp& dq_dram_block_window_tmp, + const BiasGradDramBlockWindowTmp& dbias_dram_block_window_tmp, + FmhaMask mask, + PositionEncoding position_encoding, + float raw_scale, + float scale, + float rp_undrop, + float scale_rp_undrop, + void* smem_ptr, + FmhaDropout& dropout) const + { + static_assert( + std::is_same_v> && + std::is_same_v> && + std::is_same_v> && + std::is_same_v> && + std::is_same_v> && + std::is_same_v>, + "wrong!"); + + static_assert(kM0 == QDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] && + kN0 == KDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] && + kN0 == VDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] && + kM0 == BiasDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] && + kN0 == BiasDramBlockWindowTmp{}.get_window_lengths()[number<1>{}] && + kM0 == OGradDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] && + kM0 == LSEDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] && + kM0 == DDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] && + kM0 == QGradDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] && + kM0 == BiasGradDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] && + kN0 == BiasGradDramBlockWindowTmp{}.get_window_lengths()[number<1>{}], + "wrong!"); + + // Block GEMM + constexpr auto gemm_0 = Policy::template GetQKBlockGemm(); + constexpr auto gemm_1 = Policy::template GetPTOGradTBlockGemm(); + constexpr auto gemm_2 = Policy::template GetOGradVBlockGemm(); + constexpr auto gemm_3 = Policy::template GetSGradTQTBlockGemm(); + constexpr auto gemm_4 = Policy::template GetSGradKTBlockGemm(); + + // init VGrad & KGrad + auto dv_acc = decltype(gemm_1.MakeCBlockTile()){}; + auto dk_acc = decltype(gemm_3.MakeCBlockTile()){}; + + // K, HBM ->LDS ->Reg + auto k_dram_window = + make_tile_window(k_dram_block_window_tmp.get_bottom_tensor_view(), + k_dram_block_window_tmp.get_window_lengths(), + k_dram_block_window_tmp.get_window_origin(), + Policy::template MakeKDramTileDistribution()); + + const auto k_origin = k_dram_window.get_window_origin(); + // Early termination + const auto [seqlen_q_start, seqlen_q_end] = + mask.GetTileRangeAlongY(k_origin.at(number<0>{}), number{}, number{}); + + const auto num_total_loop = integer_divide_ceil(seqlen_q_end - seqlen_q_start, kM0); + + // check early exit if masked and no work to do. + if constexpr(FmhaMask::IsMasking) + { + if(num_total_loop <= 0) + { + // Note: here dk_acc&dv_acc are all cleard, return it + // Note: v loaded but no fence, ignore it. + return make_tuple(dk_acc, dv_acc); + } + } + KDataType* k_lds_ptr = + static_cast(static_cast(static_cast(smem_ptr))); + auto k_lds = make_tensor_view( + k_lds_ptr, Policy::template MakeKLdsWriteBlockDescriptor()); + + auto k_lds_write_window = + make_tile_window(k_lds, make_tuple(number{}, number{}), {0, 0}); + + auto k_lds_read_window = + make_tile_window(k_lds_write_window.get_bottom_tensor_view(), + make_tuple(number{}, number{}), + k_lds_write_window.get_window_origin(), + Policy::template MakeKRegSliceBlockDescriptor()); + + auto k_reg_tensor = make_static_distributed_tensor( + Policy::template MakeKRegBlockDescriptor()); + + //------------------------------------------------------------------ + // V, HBM ->LDS ->Reg + auto v_dram_window = + make_tile_window(v_dram_block_window_tmp.get_bottom_tensor_view(), + v_dram_block_window_tmp.get_window_lengths(), + v_dram_block_window_tmp.get_window_origin(), + Policy::template MakeVDramTileDistribution()); + + VDataType* v_lds_ptr = + static_cast(static_cast(static_cast(smem_ptr))); + + auto v_lds = make_tensor_view( + v_lds_ptr, Policy::template MakeVLdsWriteBlockDescriptor()); + + auto v_lds_write_window = + make_tile_window(v_lds, make_tuple(number{}, number{}), {0, 0}); + + auto v_lds_read_window = + make_tile_window(v_lds_write_window.get_bottom_tensor_view(), + make_tuple(number{}, number{}), + v_lds_write_window.get_window_origin(), + Policy::template MakeVRegSliceBlockDescriptor()); + + auto v_reg_tensor = make_static_distributed_tensor( + Policy::template MakeVRegBlockDescriptor()); + + //------------------------------------------------------------------ + // KT, Reg ->LDS ->Reg + auto shuffled_k_block_tile = make_static_distributed_tensor( + Policy::template MakeShuffledKRegWriteBlockDescriptor()); + + KDataType* kt_lds_ptr = static_cast(static_cast( + static_cast(smem_ptr) + Policy::template GetSmemSizeK())); + + auto shuffled_k_lds_write = make_tensor_view( + kt_lds_ptr, Policy::template MakeShuffledKLdsWriteBlockDescriptor()); + + auto shuffled_k_lds_write_window = make_tile_window( + shuffled_k_lds_write, make_tuple(number{}, number{}), {0, 0}); + + auto kt_lds_read = make_tensor_view( + kt_lds_ptr, Policy::template MakeKTLdsReadBlockDescriptor()); + + auto kt_lds_read_window = + make_tile_window(kt_lds_read, + make_tuple(number{}, number{}), + {0, 0}, + Policy::template MakeKTRegBlockDescriptor()); + + //------------------------------------------------------------------ + // Pre-Load KV into Registers + auto k_block_tile = load_tile(k_dram_window); + auto v_block_tile = load_tile(v_dram_window); + + store_tile(k_lds_write_window, k_block_tile); + shuffle_tile(shuffled_k_block_tile, k_block_tile); + store_tile(shuffled_k_lds_write_window, shuffled_k_block_tile); + + block_sync_lds(); + k_reg_tensor = load_tile(k_lds_read_window); + block_sync_lds(); + + auto kt_reg_tensor = load_tile(kt_lds_read_window); + + store_tile(v_lds_write_window, v_block_tile); + + block_sync_lds(); + + v_reg_tensor = load_tile(v_lds_read_window); + block_sync_lds(); + //---------------------------- Loop Load in ----------------------------// + // Q: HBM ->Reg ->LDS + auto q_dram_window = + make_tile_window(q_dram_block_window_tmp.get_bottom_tensor_view(), + q_dram_block_window_tmp.get_window_lengths(), + {seqlen_q_start, 0}, + Policy::template MakeQDramTileDistribution()); + + QDataType* q_lds_ptr = static_cast(static_cast( + static_cast(smem_ptr) + Policy::template GetSmemSizeQT() + + Policy::template GetSmemSizeOGrad() + + Policy::template GetSmemSizeOGradT())); + + auto q_lds = make_tensor_view( + q_lds_ptr, Policy::template MakeQLdsBlockDescriptor()); + + auto q_lds_window = + make_tile_window(q_lds, make_tuple(number{}, number{}), {0, 0}); + + auto q_lds_read_window = + make_tile_window(q_lds_window.get_bottom_tensor_view(), + make_tuple(number{}, number{}), + q_lds_window.get_window_origin(), + Policy::template MakeQRegSliceBlockDescriptor()); + + auto pt_reg_tensor = make_static_distributed_tensor( + Policy::template MakePTRegSliceBlockDescriptor()); + // QT: Reg -> Reg-> LDS + auto shuffled_q_block_tile = make_static_distributed_tensor( + Policy::template MakeShuffledQRegWriteBlockDescriptor()); + + QDataType* qt_lds_ptr = + static_cast(static_cast(static_cast(smem_ptr))); + + auto shuffled_q_lds_write = make_tensor_view( + qt_lds_ptr, Policy::template MakeShuffledQLdsWriteBlockDescriptor()); + + auto shuffled_q_lds_write_window = make_tile_window( + shuffled_q_lds_write, make_tuple(number{}, number{}), {0, 0}); + + auto qt_lds_read = make_tensor_view( + qt_lds_ptr, Policy::template MakeQTLdsReadBlockDescriptor()); + + auto qt_lds_read_window = + make_tile_window(qt_lds_read, + make_tuple(number{}, number{}), + {0, 0}, + Policy::template MakeQTRegSliceBlockDescriptor()); + + // dO: HBM ->Reg ->LDS + auto do_dram_window = + make_tile_window(do_dram_block_window_tmp.get_bottom_tensor_view(), + do_dram_block_window_tmp.get_window_lengths(), + {seqlen_q_start, 0}, + Policy::template MakeOGradDramTileDistribution()); + + OGradDataType* do_lds_ptr = static_cast(static_cast( + static_cast(smem_ptr) + Policy::template GetSmemSizeQT())); + + auto do_lds = make_tensor_view( + do_lds_ptr, Policy::template MakeOGradLdsBlockDescriptor()); + + auto do_lds_window = + make_tile_window(do_lds, make_tuple(number{}, number{}), {0, 0}); + + auto do_lds_read_window = + make_tile_window(do_lds_window.get_bottom_tensor_view(), + make_tuple(number{}, number{}), + do_lds_window.get_window_origin(), + Policy::template MakeOGradRegSliceBlockDescriptor()); + // dOT: Reg ->Reg ->LDS + auto shuffled_do_block_tile = make_static_distributed_tensor( + Policy::template MakeShuffledOGradRegWriteBlockDescriptor()); + + OGradDataType* dot_lds_ptr = static_cast(static_cast( + static_cast(smem_ptr) + Policy::template GetSmemSizeQT() + + Policy::template GetSmemSizeOGrad())); + + auto shuffled_do_lds_write = make_tensor_view( + dot_lds_ptr, Policy::template MakeShuffledOGradLdsWriteBlockDescriptor()); + + auto shuffled_do_lds_write_window = make_tile_window( + shuffled_do_lds_write, make_tuple(number{}, number{}), {0, 0}); + + auto dot_read_lds = make_tensor_view( + dot_lds_ptr, Policy::template MakeOGradTLdsReadBlockDescriptor()); + + auto dot_lds_read_window = + make_tile_window(dot_read_lds, + make_tuple(number{}, number{}), + {0, 0}, + Policy::template MakeOGradTRegSliceBlockDescriptor()); + + // dS: Reg -> Reg -> LDS + GemmDataType* ds_lds_ptr = static_cast(static_cast( + static_cast(smem_ptr) + Policy::template GetSmemSizeQT() + + Policy::template GetSmemSizeOGrad() + + Policy::template GetSmemSizeOGradT() + + Policy::template GetSmemSizeQ() + Policy::template GetSmemSizeLSE() + + Policy::template GetSmemSizeD())); + + auto ds_lds = make_tensor_view( + ds_lds_ptr, Policy::template MakeSGradLdsBlockDescriptor()); + + auto ds_lds_window = + make_tile_window(ds_lds, make_tuple(number{}, number{}), {0, 0}); + + auto ds_lds_read_window = + make_tile_window(ds_lds_window.get_bottom_tensor_view(), + make_tuple(number{}, number{}), + ds_lds_window.get_window_origin(), + Policy::template MakeSGradRegSliceBlockDescriptor()); + + auto dst_reg_tensor = make_static_distributed_tensor( + Policy::template MakeSGradTRegSliceBlockDescriptor()); + // Bias: HBM ->Reg ->Reg ->LDS + 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(), + {seqlen_q_start, bias_origin.at(number<1>{})}, + Policy::template MakeBiasTileDistribution()); + + BiasDataType* bias_lds_ptr = static_cast(static_cast( + static_cast(smem_ptr) + Policy::template GetSmemSizeQT() + + Policy::template GetSmemSizeOGrad() + + Policy::template GetSmemSizeOGradT() + + Policy::template GetSmemSizeQ() + Policy::template GetSmemSizeLSE() + + Policy::template GetSmemSizeD())); + + auto bias_lds = make_tensor_view( + bias_lds_ptr, Policy::template MakeBiasLdsBlockDescriptor()); + + auto bias_lds_write_window = + make_tile_window(bias_lds, make_tuple(number{}, number{}), {0, 0}); + + auto bias_s_lds_read_window = + make_tile_window(bias_lds_write_window.get_bottom_tensor_view(), + bias_lds_write_window.get_window_lengths(), + bias_lds_write_window.get_window_origin(), + Policy::template MakeBiasSTileDistribution()); + + static_assert(std::is_same_v, + "BiasDataType and BiasGradDataType should be the same!"); + + // LSE: HBM -> LDS ->Reg + auto lse_dram_window = make_tile_window( + lse_dram_block_window_tmp.get_bottom_tensor_view(), + lse_dram_block_window_tmp.get_window_lengths(), + {seqlen_q_start}, + Policy::template MakeLSEDDramTileDistribution()); + + LSEDataType* lse_lds_ptr = static_cast(static_cast( + static_cast(smem_ptr) + Policy::template GetSmemSizeQT() + + Policy::template GetSmemSizeOGrad() + + Policy::template GetSmemSizeOGradT() + + Policy::template GetSmemSizeQ())); + + auto lse_lds = make_tensor_view( + lse_lds_ptr, Policy::template MakeLSEDLdsWriteBlockDescriptor()); + + auto lse_lds_write_window = make_tile_window(lse_lds, make_tuple(number{}), {0}); + + auto lse_lds_read_window = make_tile_window( + lse_lds, + make_tuple(number{}), + {0}, + Policy::template MakeLSEDLdsReadBlockDescriptor()); + + // D: HBM ->Reg + auto d_dram_window = make_tile_window( + d_dram_block_window_tmp.get_bottom_tensor_view(), + d_dram_block_window_tmp.get_window_lengths(), + {seqlen_q_start}, + Policy::template MakeLSEDDramTileDistribution()); + + DDataType* d_lds_ptr = static_cast(static_cast( + static_cast(smem_ptr) + Policy::template GetSmemSizeQT() + + Policy::template GetSmemSizeOGrad() + + Policy::template GetSmemSizeOGradT() + + Policy::template GetSmemSizeQ() + Policy::template GetSmemSizeLSE())); + + auto d_lds = make_tensor_view( + d_lds_ptr, Policy::template MakeLSEDLdsWriteBlockDescriptor()); + + auto d_lds_write_window = make_tile_window(d_lds, make_tuple(number{}), {0}); + + auto d_lds_read_window = make_tile_window( + d_lds, + make_tuple(number{}), + {0}, + Policy::template MakeLSEDLdsReadBlockDescriptor()); + + // RandVal: HBM ->Reg + auto randval_dram_window = dropout.template MakeRandvalDramWindow( + randval_dram_block_window_tmp, seqlen_q_start); + + // BiasGrad + // Reg ->LDS ->Reg ->HBM + const auto dbias_origin = dbias_dram_block_window_tmp.get_window_origin(); + + auto dbias_dram_window = + make_tile_window(dbias_dram_block_window_tmp.get_bottom_tensor_view(), + dbias_dram_block_window_tmp.get_window_lengths(), + {seqlen_q_start, dbias_origin.at(number<1>{})}); // M/N + + auto dbias_lds_read_window = + make_tile_window(bias_lds, + make_tuple(number{}, number{}), + {0, 0}, + Policy::template MakeShuffledBiasTileDistribution()); + + // ----------------------------Loop write out------------------------------// + auto dq_dram_window = make_tile_window(dq_dram_block_window_tmp.get_bottom_tensor_view(), + dq_dram_block_window_tmp.get_window_lengths(), + {seqlen_q_start, 0}); + + using SPBlockTileType = decltype(gemm_0.MakeCBlockTile()); + using SPGradBlockTileType = decltype(gemm_2.MakeCBlockTile()); + using QGradBlockTileType = decltype(gemm_4.MakeCBlockTile()); + + index_t i_total_loops = 0; + index_t seqlen_q_step = seqlen_q_start; + static_assert(kQKHeaddim == kK0, "kQKHeaddim should equal to kK0"); + static_assert(kM0 == kK1, "kM0 should equal to kK1"); + static_assert(kVHeaddim == kK2, "kVHeaddim should equal to kK2"); + static_assert(kM0 == kK3, "kM0 should equal to kK3"); + constexpr index_t k4_loops = kN0 / kK4; + + clear_tile(dv_acc); + clear_tile(dk_acc); + + __builtin_amdgcn_sched_barrier(0); + // Hot loop + while(i_total_loops < num_total_loop) + { + auto q_block_tile = load_tile(q_dram_window); + move_tile_window(q_dram_window, {kM0, 0}); + + auto lse_block_tile = load_tile(lse_dram_window); + move_tile_window(lse_dram_window, {kM0}); + + store_tile(q_lds_window, q_block_tile); + shuffle_tile(shuffled_q_block_tile, q_block_tile); + store_tile(shuffled_q_lds_write_window, shuffled_q_block_tile); + + store_tile(lse_lds_write_window, lse_block_tile); + + block_sync_lds(); + + auto q_reg_tensor = load_tile(q_lds_read_window); + auto lse = load_tile(lse_lds_read_window); + + block_sync_lds(); + + // STAGE 1, Q@K Gemm0 + auto s_acc = SPBlockTileType{}; + + s_acc = gemm_0(q_reg_tensor, k_reg_tensor); + + // STAGE 2, Scale, Add bias, Mask, Softmax, Dropout + if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS) + { + const auto bias_tile = load_tile(bias_dram_window); + auto shuffled_bias_tile = make_static_distributed_tensor( + Policy::template MakeShuffledBiasTileDistribution()); + shuffle_tile(shuffled_bias_tile, bias_tile); + store_tile(bias_lds_write_window, shuffled_bias_tile); + block_sync_lds(); + auto bias_s_tile = load_tile(bias_s_lds_read_window); + tile_elementwise_inout( + [&](auto& x, const auto& y) { + x = scale * x + log2e_v * type_convert(y); + }, + s_acc, + bias_s_tile); + move_tile_window(bias_dram_window, {kM0, 0}); + __builtin_amdgcn_sched_barrier(0); + } + else if constexpr(BiasEnum == BlockAttentionBiasEnum::ALIBI) + { + constexpr auto s_spans = decltype(s_acc)::get_distributed_spans(); + 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 = seqlen_q_step + 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; + position_encoding.update(s_acc(i_j_idx), row, col); + }); + }); + } + + if constexpr(kPadSeqLenK || FmhaMask::IsMasking) + { + bool need_perpixel_check = mask.IsEdgeTile( + seqlen_q_step, k_origin.at(number<0>{}), number{}, number{}); + if(need_perpixel_check) + { + set_tile_if(s_acc, -numeric::infinity(), [&](auto tile_idx) { + const auto row = seqlen_q_step + tile_idx.at(number<0>{}); + const auto col = k_origin.at(number<0>{}) + tile_idx.at(number<1>{}); + return mask.IsOutOfBound(row, col); + }); + } + } + + static const auto get_validated_lse = [](LSEDataType raw_lse) { + if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS || + FmhaMask::IsMasking) + { + return raw_lse == -numeric::infinity() + ? type_convert(0.f) + : raw_lse; + } + else + { + return raw_lse; + } + }; + + auto p = SPBlockTileType{}; + constexpr auto p_spans = decltype(p)::get_distributed_spans(); + sweep_tile_span(p_spans[number<0>{}], [&](auto idx0) { + constexpr auto i_idx = make_tuple(idx0); + auto row_lse = log2e_v * get_validated_lse(lse[i_idx]); + + sweep_tile_span(p_spans[number<1>{}], [&](auto idx1) { + constexpr auto i_j_idx = make_tuple(idx0, idx1); + + if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS || + BiasEnum == BlockAttentionBiasEnum::ALIBI) + { + p(i_j_idx) = exp2(s_acc[i_j_idx] - row_lse); + } + else + { + p(i_j_idx) = exp2(scale * s_acc[i_j_idx] - row_lse); + } + }); + }); + + if constexpr(FmhaDropout::IsDropout) + { + dropout.template Run( + seqlen_q_step, k_origin.at(number<0>{}), p, randval_dram_window); + } + const auto p_gemm = [&]() { + if constexpr(FmhaDropout::IsDropout) + { + return tile_elementwise_in( + [](const auto& x) { return type_convert(x > 0.f ? x : 0.f); }, + p); + } + else + { + return cast_tile(p); + } + }(); + + // STAGE 3, P^T@OGrad^T Gemm1 + auto do_block_tile = load_tile(do_dram_window); + move_tile_window(do_dram_window, {kM0, 0}); + + auto d_block_tile = load_tile(d_dram_window); + move_tile_window(d_dram_window, {kM0}); + + store_tile(do_lds_window, do_block_tile); + shuffle_tile(shuffled_do_block_tile, do_block_tile); + store_tile(shuffled_do_lds_write_window, shuffled_do_block_tile); + + store_tile(d_lds_write_window, d_block_tile); + + block_sync_lds(); + + auto dot_reg_tensor = load_tile(dot_lds_read_window); + + block_sync_lds(); + + Policy::template PTFromGemm0CToGemm1A(pt_reg_tensor, p_gemm); + gemm_1(dv_acc, pt_reg_tensor, dot_reg_tensor); + + // STAGE 4, OGrad@V Gemm2 + auto do_reg_tensor = load_tile(do_lds_read_window); + auto d = load_tile(d_lds_read_window); + block_sync_lds(); + + auto dp_acc = SPGradBlockTileType{}; + + dp_acc = gemm_2(do_reg_tensor, v_reg_tensor); + + // STAGE 5, P^T(PGrad^T - D) + auto ds = SPGradBlockTileType{}; + constexpr auto ds_spans = decltype(ds)::get_distributed_spans(); + sweep_tile_span(ds_spans[number<0>{}], [&](auto idx0) { + constexpr auto i_idx = make_tuple(idx0); + sweep_tile_span(ds_spans[number<1>{}], [&](auto idx1) { + constexpr auto i_j_idx = make_tuple(idx0, idx1); + bool undrop_flag = p[i_j_idx] >= 0; + ds(i_j_idx) = p[i_j_idx] * (!FmhaDropout::IsDropout || undrop_flag + ? (dp_acc[i_j_idx] - d[i_idx]) + : d[i_idx]); + }); + }); + + if constexpr(kHasBiasGrad) + { + const auto dbias = [&]() { + if constexpr(FmhaDropout::IsDropout) + { + return tile_elementwise_in( + [&rp_undrop](const auto& x) { + return type_convert(x * rp_undrop); + }, + ds); + } + else + { + return cast_tile(ds); + } + }(); + store_tile(bias_lds_write_window, dbias); + block_sync_lds(); + auto shuffled_dbias_tile = load_tile(dbias_lds_read_window); + auto dbias_tile = make_static_distributed_tensor( + Policy::template MakeBiasTileDistribution()); + shuffle_tile(dbias_tile, shuffled_dbias_tile); + store_tile(dbias_dram_window, dbias_tile); + move_tile_window(dbias_dram_window, {kM0, 0}); + __builtin_amdgcn_sched_barrier(0); + } + + // STAGE 6, SGrad^T@Q^T Gemm3 + auto qt_reg_tensor = load_tile(qt_lds_read_window); + block_sync_lds(); + + const auto ds_gemm = cast_tile(ds); + + Policy::template SGradTFromGemm2CToGemm3A(dst_reg_tensor, ds_gemm); + + gemm_3(dk_acc, dst_reg_tensor, qt_reg_tensor); + + store_tile(ds_lds_window, ds_gemm); + + block_sync_lds(); + + auto ds_reg_tensor = load_tile(ds_lds_read_window); + auto ds_reg_tensor_next = decltype(ds_reg_tensor){}; + move_tile_window(ds_lds_read_window, {0, kK4}); + + // STAGE7 SGrad@K^T Gemm4 + auto dq_acc = QGradBlockTileType{}; + clear_tile(dq_acc); + + static_for<0, k4_loops, 1>{}([&](auto i_k4) { + if constexpr(i_k4 < k4_loops - 1) + { + ds_reg_tensor_next = load_tile(ds_lds_read_window); + move_tile_window(ds_lds_read_window, {0, kK4}); + } + auto kt_reg_tensor_slice = get_slice_tile(kt_reg_tensor, + sequence<0, i_k4 * kK4>{}, + sequence{}); + gemm_4(dq_acc, ds_reg_tensor, kt_reg_tensor_slice); + + if constexpr(i_k4 < k4_loops - 1) + { + ds_reg_tensor.get_thread_buffer() = ds_reg_tensor_next.get_thread_buffer(); + } + }); + move_tile_window(ds_lds_read_window, {0, -kN0}); + // QGrad Scale + if constexpr(FmhaDropout::IsDropout) + { + tile_elementwise_inout([&scale_rp_undrop](auto& x) { x = x * scale_rp_undrop; }, + dq_acc); + } + else + { + tile_elementwise_inout([&raw_scale](auto& x) { x = x * raw_scale; }, dq_acc); + } + if constexpr(kIsDeterministic) + { + store_tile(dq_dram_window, dq_acc); + } + else + { + update_tile(dq_dram_window, dq_acc); + } + move_tile_window(dq_dram_window, {kM0, 0}); + + i_total_loops += 1; + seqlen_q_step += kM0; + } + + // Results Scale + if constexpr(FmhaDropout::IsDropout) + { + tile_elementwise_inout([&scale_rp_undrop](auto& x) { x = x * scale_rp_undrop; }, + dk_acc); + tile_elementwise_inout([&rp_undrop](auto& x) { x = x * rp_undrop; }, dv_acc); + } + else + { + tile_elementwise_inout([&raw_scale](auto& x) { x = x * raw_scale; }, dk_acc); + } + + return make_tuple(dk_acc, dv_acc); + } +}; + +} // namespace ck_tile diff --git a/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_kr_ktr_vr_iglp.hpp b/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_kr_ktr_vr_iglp.hpp new file mode 100644 index 0000000000..9e6a2725c9 --- /dev/null +++ b/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_kr_ktr_vr_iglp.hpp @@ -0,0 +1,1037 @@ +// 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/block/block_dropout.hpp" +#include "ck_tile/ops/fmha/pipeline/block_fmha_bwd_pipeline_default_policy.hpp" +#include "ck_tile/ops/reduce/block/block_reduce.hpp" + +namespace ck_tile { + +template +struct BlockFmhaBwdDQDKDVPipelineKRKTRVRIGLP +{ + using QDataType = remove_cvref_t; + using KDataType = remove_cvref_t; + using VDataType = remove_cvref_t; + using GemmDataType = remove_cvref_t; + using BiasDataType = remove_cvref_t; + using LSEDataType = remove_cvref_t; + using AccDataType = remove_cvref_t; + using DDataType = remove_cvref_t; + using RandValOutputDataType = remove_cvref_t; + using ODataType = remove_cvref_t; + using OGradDataType = remove_cvref_t; + using QGradDataType = remove_cvref_t; + using KGradDataType = remove_cvref_t; + using VGradDataType = remove_cvref_t; + using BiasGradDataType = remove_cvref_t; + using FmhaMask = remove_cvref_t; + using FmhaDropout = remove_cvref_t; + using HotLoopScheduler = typename Policy::template HotLoopScheduler; + + using BlockFmhaShape = remove_cvref_t; + + static constexpr index_t kBlockPerCu = Problem::kBlockPerCu; + 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 kK1 = BlockFmhaShape::kK1; + static constexpr index_t kK2 = BlockFmhaShape::kK2; + static constexpr index_t kK3 = BlockFmhaShape::kK3; + static constexpr index_t kK4 = BlockFmhaShape::kK4; + static constexpr index_t kQKHeaddim = BlockFmhaShape::kQKHeaddim; + static constexpr index_t kVHeaddim = BlockFmhaShape::kVHeaddim; + + 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 kHasBiasGrad = Problem::kHasBiasGrad; + static constexpr bool kIsDeterministic = Problem::kIsDeterministic; + + // 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 = + kPadHeadDimV ? 1 : Policy::template GetAlignmentV(); + static constexpr index_t kAlignmentOGrad = + kPadHeadDimV ? 1 : Policy::template GetAlignmentOGrad(); + static constexpr index_t kAlignmentQGrad = 1; + static constexpr index_t kAlignmentKGrad = + kPadHeadDimQ ? 1 : Policy::template GetAlignmentKGrad(); + static constexpr index_t kAlignmentVGrad = + kPadHeadDimV ? 1 : Policy::template GetAlignmentVGrad(); + static constexpr index_t kAlignmentBias = + kPadSeqLenK ? 1 : Policy::template GetTransposedAlignmentBias(); + + static constexpr const char* name = "kr_ktr_vr_iglp"; + + 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, + const KDramBlockWindowTmp& k_dram_block_window_tmp, + const VDramBlockWindowTmp& v_dram_block_window_tmp, + const BiasDramBlockWindowTmp& bias_dram_block_window_tmp, + const RandValDramBlockWindowTmp& randval_dram_block_window_tmp, + const OGradDramBlockWindowTmp& do_dram_block_window_tmp, + const LSEDramBlockWindowTmp& lse_dram_block_window_tmp, + const DDramBlockWindowTmp& d_dram_block_window_tmp, + const QGradDramBlockWindowTmp& dq_dram_block_window_tmp, + const BiasGradDramBlockWindowTmp& dbias_dram_block_window_tmp, + FmhaMask mask, + PositionEncoding position_encoding, + float raw_scale, + float scale, + float rp_undrop, + float scale_rp_undrop, + void* smem_ptr, + FmhaDropout& dropout) const + { + static_assert( + std::is_same_v> && + std::is_same_v> && + std::is_same_v> && + std::is_same_v> && + std::is_same_v> && + std::is_same_v>, + "wrong!"); + + static_assert(kM0 == QDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] && + kN0 == KDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] && + kN0 == VDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] && + kM0 == BiasDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] && + kN0 == BiasDramBlockWindowTmp{}.get_window_lengths()[number<1>{}] && + kM0 == OGradDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] && + kM0 == LSEDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] && + kM0 == DDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] && + kM0 == QGradDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] && + kM0 == BiasGradDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] && + kN0 == BiasGradDramBlockWindowTmp{}.get_window_lengths()[number<1>{}], + "wrong!"); + + // Block GEMM + constexpr auto gemm_0 = Policy::template GetQKBlockGemm(); + constexpr auto gemm_1 = Policy::template GetPTOGradTBlockGemm(); + constexpr auto gemm_2 = Policy::template GetOGradVBlockGemm(); + constexpr auto gemm_3 = Policy::template GetSGradTQTBlockGemm(); + constexpr auto gemm_4 = Policy::template GetSGradKTBlockGemm(); + + // init VGrad & KGrad + auto dv_acc = decltype(gemm_1.MakeCBlockTile()){}; + auto dk_acc = decltype(gemm_3.MakeCBlockTile()){}; + + // K, HBM ->LDS ->Reg + auto k_dram_window = + make_tile_window(k_dram_block_window_tmp.get_bottom_tensor_view(), + k_dram_block_window_tmp.get_window_lengths(), + k_dram_block_window_tmp.get_window_origin(), + Policy::template MakeKDramTileDistribution()); + + const auto k_origin = k_dram_window.get_window_origin(); + // Early termination + const auto [seqlen_q_start, seqlen_q_end] = + mask.GetTileRangeAlongY(k_origin.at(number<0>{}), number{}, number{}); + + const auto num_total_loop = integer_divide_ceil(seqlen_q_end - seqlen_q_start, kM0); + + // check early exit if masked and no work to do. + if constexpr(FmhaMask::IsMasking) + { + if(num_total_loop <= 0) + { + // Note: here dk_acc&dv_acc are all cleard, return it + // Note: v loaded but no fence, ignore it. + return make_tuple(dk_acc, dv_acc); + } + } + KDataType* k_lds_ptr = + static_cast(static_cast(static_cast(smem_ptr))); + auto k_lds = make_tensor_view( + k_lds_ptr, Policy::template MakeKLdsWriteBlockDescriptor()); + + auto k_lds_write_window = + make_tile_window(k_lds, make_tuple(number{}, number{}), {0, 0}); + + auto k_lds_read_window = + make_tile_window(k_lds_write_window.get_bottom_tensor_view(), + make_tuple(number{}, number{}), + k_lds_write_window.get_window_origin(), + Policy::template MakeKRegSliceBlockDescriptor()); + + auto k_reg_tensor = make_static_distributed_tensor( + Policy::template MakeKRegBlockDescriptor()); + + //------------------------------------------------------------------ + // V, HBM ->LDS ->Reg + auto v_dram_window = + make_tile_window(v_dram_block_window_tmp.get_bottom_tensor_view(), + v_dram_block_window_tmp.get_window_lengths(), + v_dram_block_window_tmp.get_window_origin(), + Policy::template MakeVDramTileDistribution()); + + VDataType* v_lds_ptr = + static_cast(static_cast(static_cast(smem_ptr))); + + auto v_lds = make_tensor_view( + v_lds_ptr, Policy::template MakeVLdsWriteBlockDescriptor()); + + auto v_lds_write_window = + make_tile_window(v_lds, make_tuple(number{}, number{}), {0, 0}); + + auto v_lds_read_window = + make_tile_window(v_lds_write_window.get_bottom_tensor_view(), + make_tuple(number{}, number{}), + v_lds_write_window.get_window_origin(), + Policy::template MakeVRegSliceBlockDescriptor()); + + auto v_reg_tensor = make_static_distributed_tensor( + Policy::template MakeVRegBlockDescriptor()); + + //------------------------------------------------------------------ + // KT, Reg ->LDS ->Reg + auto shuffled_k_block_tile = make_static_distributed_tensor( + Policy::template MakeShuffledKRegWriteBlockDescriptor()); + + KDataType* kt_lds_ptr = static_cast(static_cast( + static_cast(smem_ptr) + Policy::template GetSmemSizeK())); + + auto shuffled_k_lds_write = make_tensor_view( + kt_lds_ptr, Policy::template MakeShuffledKLdsWriteBlockDescriptor()); + + auto shuffled_k_lds_write_window = make_tile_window( + shuffled_k_lds_write, make_tuple(number{}, number{}), {0, 0}); + + auto kt_lds_read = make_tensor_view( + kt_lds_ptr, Policy::template MakeKTLdsReadBlockDescriptor()); + + auto kt_lds_read_window = + make_tile_window(kt_lds_read, + make_tuple(number{}, number{}), + {0, 0}, + Policy::template MakeKTRegBlockDescriptor()); + + //------------------------------------------------------------------ + // Pre-Load KV into Registers + auto k_block_tile = load_tile(k_dram_window); + auto v_block_tile = load_tile(v_dram_window); + + store_tile(k_lds_write_window, k_block_tile); + shuffle_tile(shuffled_k_block_tile, k_block_tile); + store_tile(shuffled_k_lds_write_window, shuffled_k_block_tile); + + block_sync_lds(); + k_reg_tensor = load_tile(k_lds_read_window); + block_sync_lds(); + + auto kt_reg_tensor = load_tile(kt_lds_read_window); + + store_tile(v_lds_write_window, v_block_tile); + + block_sync_lds(); + + v_reg_tensor = load_tile(v_lds_read_window); + //---------------------------- Loop Load in ----------------------------// + // Q: HBM ->Reg ->LDS + auto q_dram_window = + make_tile_window(q_dram_block_window_tmp.get_bottom_tensor_view(), + q_dram_block_window_tmp.get_window_lengths(), + {seqlen_q_start, 0}, + Policy::template MakeQDramTileDistribution()); + + QDataType* q_lds_ptr = static_cast(static_cast( + static_cast(smem_ptr) + Policy::template GetSmemSizeQT() + + Policy::template GetSmemSizeOGrad() + + Policy::template GetSmemSizeOGradT())); + + auto q_lds = make_tensor_view( + q_lds_ptr, Policy::template MakeQLdsBlockDescriptor()); + + auto q_lds_window = + make_tile_window(q_lds, make_tuple(number{}, number{}), {0, 0}); + + auto q_lds_read_window = + make_tile_window(q_lds_window.get_bottom_tensor_view(), + make_tuple(number{}, number{}), + q_lds_window.get_window_origin(), + Policy::template MakeQRegSliceBlockDescriptor()); + + auto pt_reg_tensor = make_static_distributed_tensor( + Policy::template MakePTRegSliceBlockDescriptor()); + // QT: Reg -> Reg-> LDS + auto shuffled_q_block_tile = make_static_distributed_tensor( + Policy::template MakeShuffledQRegWriteBlockDescriptor()); + + QDataType* qt_lds_ptr = + static_cast(static_cast(static_cast(smem_ptr))); + + auto shuffled_q_lds_write = make_tensor_view( + qt_lds_ptr, Policy::template MakeShuffledQLdsWriteBlockDescriptor()); + + auto shuffled_q_lds_write_window = make_tile_window( + shuffled_q_lds_write, make_tuple(number{}, number{}), {0, 0}); + + auto qt_lds_read = make_tensor_view( + qt_lds_ptr, Policy::template MakeQTLdsReadBlockDescriptor()); + + auto qt_lds_read_window = + make_tile_window(qt_lds_read, + make_tuple(number{}, number{}), + {0, 0}, + Policy::template MakeQTRegSliceBlockDescriptor()); + + // dO: HBM ->Reg ->LDS + auto do_dram_window = + make_tile_window(do_dram_block_window_tmp.get_bottom_tensor_view(), + do_dram_block_window_tmp.get_window_lengths(), + {seqlen_q_start, 0}, + Policy::template MakeOGradDramTileDistribution()); + + OGradDataType* do_lds_ptr = static_cast(static_cast( + static_cast(smem_ptr) + Policy::template GetSmemSizeQT())); + + auto do_lds = make_tensor_view( + do_lds_ptr, Policy::template MakeOGradLdsBlockDescriptor()); + + auto do_lds_window = + make_tile_window(do_lds, make_tuple(number{}, number{}), {0, 0}); + + auto do_lds_read_window = + make_tile_window(do_lds_window.get_bottom_tensor_view(), + make_tuple(number{}, number{}), + do_lds_window.get_window_origin(), + Policy::template MakeOGradRegSliceBlockDescriptor()); + // dOT: Reg ->Reg ->LDS + auto shuffled_do_block_tile = make_static_distributed_tensor( + Policy::template MakeShuffledOGradRegWriteBlockDescriptor()); + + OGradDataType* dot_lds_ptr = static_cast(static_cast( + static_cast(smem_ptr) + Policy::template GetSmemSizeQT() + + Policy::template GetSmemSizeOGrad())); + + auto shuffled_do_lds_write = make_tensor_view( + dot_lds_ptr, Policy::template MakeShuffledOGradLdsWriteBlockDescriptor()); + + auto shuffled_do_lds_write_window = make_tile_window( + shuffled_do_lds_write, make_tuple(number{}, number{}), {0, 0}); + + auto dot_read_lds = make_tensor_view( + dot_lds_ptr, Policy::template MakeOGradTLdsReadBlockDescriptor()); + + auto dot_lds_read_window = + make_tile_window(dot_read_lds, + make_tuple(number{}, number{}), + {0, 0}, + Policy::template MakeOGradTRegSliceBlockDescriptor()); + + // dS: Reg -> Reg -> LDS + GemmDataType* ds_lds_ptr = static_cast(static_cast( + static_cast(smem_ptr) + Policy::template GetSmemSizeQT() + + Policy::template GetSmemSizeOGrad() + + Policy::template GetSmemSizeOGradT() + + Policy::template GetSmemSizeQ() + Policy::template GetSmemSizeLSE() + + Policy::template GetSmemSizeD())); + + auto ds_lds = make_tensor_view( + ds_lds_ptr, Policy::template MakeSGradLdsBlockDescriptor()); + + auto ds_lds_window = + make_tile_window(ds_lds, make_tuple(number{}, number{}), {0, 0}); + + auto ds_lds_read_window = + make_tile_window(ds_lds_window.get_bottom_tensor_view(), + make_tuple(number{}, number{}), + ds_lds_window.get_window_origin(), + Policy::template MakeSGradRegSliceBlockDescriptor()); + + auto dst_reg_tensor = make_static_distributed_tensor( + Policy::template MakeSGradTRegSliceBlockDescriptor()); + // Bias: HBM ->Reg ->Reg ->LDS + 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(), + {seqlen_q_start, bias_origin.at(number<1>{})}, + Policy::template MakeBiasTileDistribution()); + + BiasDataType* bias_lds_ptr = static_cast(static_cast( + static_cast(smem_ptr) + Policy::template GetSmemSizeQT() + + Policy::template GetSmemSizeOGrad() + + Policy::template GetSmemSizeOGradT() + + Policy::template GetSmemSizeQ() + Policy::template GetSmemSizeLSE() + + Policy::template GetSmemSizeD())); + + auto bias_lds = make_tensor_view( + bias_lds_ptr, Policy::template MakeBiasLdsBlockDescriptor()); + + auto bias_lds_write_window = + make_tile_window(bias_lds, make_tuple(number{}, number{}), {0, 0}); + + auto bias_s_lds_read_window = + make_tile_window(bias_lds_write_window.get_bottom_tensor_view(), + bias_lds_write_window.get_window_lengths(), + bias_lds_write_window.get_window_origin(), + Policy::template MakeBiasSTileDistribution()); + + static_assert(std::is_same_v, + "BiasDataType and BiasGradDataType should be the same!"); + + // LSE: HBM -> LDS ->Reg + auto lse_dram_window = make_tile_window( + lse_dram_block_window_tmp.get_bottom_tensor_view(), + lse_dram_block_window_tmp.get_window_lengths(), + {seqlen_q_start}, + Policy::template MakeLSEDDramTileDistribution()); + + LSEDataType* lse_lds_ptr = static_cast(static_cast( + static_cast(smem_ptr) + Policy::template GetSmemSizeQT() + + Policy::template GetSmemSizeOGrad() + + Policy::template GetSmemSizeOGradT() + + Policy::template GetSmemSizeQ())); + + auto lse_lds = make_tensor_view( + lse_lds_ptr, Policy::template MakeLSEDLdsWriteBlockDescriptor()); + + auto lse_lds_write_window = make_tile_window(lse_lds, make_tuple(number{}), {0}); + + auto lse_lds_read_window = make_tile_window( + lse_lds, + make_tuple(number{}), + {0}, + Policy::template MakeLSEDLdsReadBlockDescriptor()); + + // D: HBM ->Reg + auto d_dram_window = make_tile_window( + d_dram_block_window_tmp.get_bottom_tensor_view(), + d_dram_block_window_tmp.get_window_lengths(), + {seqlen_q_start}, + Policy::template MakeLSEDDramTileDistribution()); + + DDataType* d_lds_ptr = static_cast(static_cast( + static_cast(smem_ptr) + Policy::template GetSmemSizeQT() + + Policy::template GetSmemSizeOGrad() + + Policy::template GetSmemSizeOGradT() + + Policy::template GetSmemSizeQ() + Policy::template GetSmemSizeLSE())); + + auto d_lds = make_tensor_view( + d_lds_ptr, Policy::template MakeLSEDLdsWriteBlockDescriptor()); + + auto d_lds_write_window = make_tile_window(d_lds, make_tuple(number{}), {0}); + + auto d_lds_read_window = make_tile_window( + d_lds, + make_tuple(number{}), + {0}, + Policy::template MakeLSEDLdsReadBlockDescriptor()); + + // RandVal: HBM ->Reg + auto randval_dram_window = dropout.template MakeRandvalDramWindow( + randval_dram_block_window_tmp, seqlen_q_start); + + // BiasGrad + // Reg ->LDS ->Reg ->HBM + const auto dbias_origin = dbias_dram_block_window_tmp.get_window_origin(); + + auto dbias_dram_window = + make_tile_window(dbias_dram_block_window_tmp.get_bottom_tensor_view(), + dbias_dram_block_window_tmp.get_window_lengths(), + {seqlen_q_start, dbias_origin.at(number<1>{})}); // M/N + + auto dbias_lds_read_window = + make_tile_window(bias_lds, + make_tuple(number{}, number{}), + {0, 0}, + Policy::template MakeShuffledBiasTileDistribution()); + + // ----------------------------Loop write out------------------------------// + auto dq_dram_window = make_tile_window(dq_dram_block_window_tmp.get_bottom_tensor_view(), + dq_dram_block_window_tmp.get_window_lengths(), + {seqlen_q_start, 0}); + + using SPBlockTileType = decltype(gemm_0.MakeCBlockTile()); + using SPGradBlockTileType = decltype(gemm_2.MakeCBlockTile()); + using QGradBlockTileType = decltype(gemm_4.MakeCBlockTile()); + + index_t i_total_loops = 0; + index_t seqlen_q_step = seqlen_q_start; + static_assert(kQKHeaddim == kK0, "kQKHeaddim should equal to kK0"); + static_assert(kM0 == kK1, "kM0 should equal to kK1"); + static_assert(kVHeaddim == kK2, "kVHeaddim should equal to kK2"); + static_assert(kM0 == kK3, "kM0 should equal to kK3"); + constexpr index_t k4_loops = kN0 / kK4; + + /* + * Prefetch Q, LSE, dO, D + */ + auto q_block_tile = load_tile(q_dram_window); + move_tile_window(q_dram_window, {kM0, 0}); + auto lse_block_tile = load_tile(lse_dram_window); + move_tile_window(lse_dram_window, {kM0}); + + auto do_block_tile = load_tile(do_dram_window); + move_tile_window(do_dram_window, {kM0, 0}); + + auto d_block_tile = load_tile(d_dram_window); + move_tile_window(d_dram_window, {kM0}); + + /* + * Store prefetched data into LDS + */ + block_sync_lds(); + store_tile(q_lds_window, q_block_tile); + shuffle_tile(shuffled_q_block_tile, q_block_tile); + store_tile(shuffled_q_lds_write_window, shuffled_q_block_tile); + + store_tile(lse_lds_write_window, lse_block_tile); + + store_tile(do_lds_window, do_block_tile); + shuffle_tile(shuffled_do_block_tile, do_block_tile); + store_tile(shuffled_do_lds_write_window, shuffled_do_block_tile); + + store_tile(d_lds_write_window, d_block_tile); + block_sync_lds(); + + /* + * Prefetch LDS data into Reg to Asynchronous Data Movement and MFMA pipeline + */ + + auto q_reg_tensor = load_tile(q_lds_read_window); + auto lse = load_tile(lse_lds_read_window); + auto do_reg_tensor = load_tile(do_lds_read_window); + auto d = load_tile(d_lds_read_window); + + clear_tile(dv_acc); + clear_tile(dk_acc); + + __builtin_amdgcn_sched_barrier(0); + // Hot loop + while(i_total_loops < (num_total_loop - 1)) + { + // STAGE 1, Q@K Gemm0 + auto s_acc = SPBlockTileType{}; + + q_block_tile = load_tile(q_dram_window); + move_tile_window(q_dram_window, {kM0, 0}); + + lse_block_tile = load_tile(lse_dram_window); + move_tile_window(lse_dram_window, {kM0}); + + do_block_tile = load_tile(do_dram_window); + move_tile_window(do_dram_window, {kM0, 0}); + + d_block_tile = load_tile(d_dram_window); + move_tile_window(d_dram_window, {kM0}); + + s_acc = gemm_0(q_reg_tensor, k_reg_tensor); + + auto dot_reg_tensor = load_tile(dot_lds_read_window); + + HotLoopScheduler::template GemmStagedScheduler<0>(); + __builtin_amdgcn_sched_barrier(0); + // STAGE 2, Scale, Add bias, Mask, Softmax, Dropout + if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS) + { + const auto bias_tile = load_tile(bias_dram_window); + auto shuffled_bias_tile = make_static_distributed_tensor( + Policy::template MakeShuffledBiasTileDistribution()); + shuffle_tile(shuffled_bias_tile, bias_tile); + store_tile(bias_lds_write_window, shuffled_bias_tile); + block_sync_lds(); + auto bias_s_tile = load_tile(bias_s_lds_read_window); + tile_elementwise_inout( + [&](auto& x, const auto& y) { + x = scale * x + log2e_v * type_convert(y); + }, + s_acc, + bias_s_tile); + move_tile_window(bias_dram_window, {kM0, 0}); + __builtin_amdgcn_sched_barrier(0); + } + else if constexpr(BiasEnum == BlockAttentionBiasEnum::ALIBI) + { + constexpr auto s_spans = decltype(s_acc)::get_distributed_spans(); + 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 = seqlen_q_step + 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; + position_encoding.update(s_acc(i_j_idx), row, col); + }); + }); + } + + if constexpr(kPadSeqLenK || FmhaMask::IsMasking) + { + bool need_perpixel_check = mask.IsEdgeTile( + seqlen_q_step, k_origin.at(number<0>{}), number{}, number{}); + if(need_perpixel_check) + { + set_tile_if(s_acc, -numeric::infinity(), [&](auto tile_idx) { + const auto row = seqlen_q_step + tile_idx.at(number<0>{}); + const auto col = k_origin.at(number<0>{}) + tile_idx.at(number<1>{}); + return mask.IsOutOfBound(row, col); + }); + } + } + + static const auto get_validated_lse = [](LSEDataType raw_lse) { + if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS || + FmhaMask::IsMasking) + { + return raw_lse == -numeric::infinity() + ? type_convert(0.f) + : raw_lse; + } + else + { + return raw_lse; + } + }; + + auto p = SPBlockTileType{}; + constexpr auto p_spans = decltype(p)::get_distributed_spans(); + sweep_tile_span(p_spans[number<0>{}], [&](auto idx0) { + constexpr auto i_idx = make_tuple(idx0); + auto row_lse = log2e_v * get_validated_lse(lse[i_idx]); + + sweep_tile_span(p_spans[number<1>{}], [&](auto idx1) { + constexpr auto i_j_idx = make_tuple(idx0, idx1); + + if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS || + BiasEnum == BlockAttentionBiasEnum::ALIBI) + { + p(i_j_idx) = exp2(s_acc[i_j_idx] - row_lse); + } + else + { + p(i_j_idx) = exp2(scale * s_acc[i_j_idx] - row_lse); + } + }); + }); + + if constexpr(FmhaDropout::IsDropout) + { + dropout.template Run( + seqlen_q_step, k_origin.at(number<0>{}), p, randval_dram_window); + } + const auto p_gemm = [&]() { + if constexpr(FmhaDropout::IsDropout) + { + return tile_elementwise_in( + [](const auto& x) { return type_convert(x > 0.f ? x : 0.f); }, + p); + } + else + { + return cast_tile(p); + } + }(); + + // STAGE 3, P^T@OGrad^T Gemm1 + Policy::template PTFromGemm0CToGemm1A(pt_reg_tensor, p_gemm); + gemm_1(dv_acc, pt_reg_tensor, dot_reg_tensor); + + auto qt_reg_tensor = load_tile(qt_lds_read_window); + + HotLoopScheduler::template GemmStagedScheduler<1>(); + __builtin_amdgcn_sched_barrier(0); + // STAGE 4, OGrad@V Gemm2 + auto dp_acc = SPGradBlockTileType{}; + + dp_acc = gemm_2(do_reg_tensor, v_reg_tensor); + + block_sync_lds(); + + store_tile(q_lds_window, q_block_tile); + shuffle_tile(shuffled_q_block_tile, q_block_tile); + store_tile(shuffled_q_lds_write_window, shuffled_q_block_tile); + + store_tile(lse_lds_write_window, lse_block_tile); + + store_tile(do_lds_window, do_block_tile); + shuffle_tile(shuffled_do_block_tile, do_block_tile); + store_tile(shuffled_do_lds_write_window, shuffled_do_block_tile); + + store_tile(d_lds_write_window, d_block_tile); + + HotLoopScheduler::template GemmStagedScheduler<2>(); + __builtin_amdgcn_sched_barrier(0); + // STAGE 5, P^T(PGrad^T - D) + auto ds = SPGradBlockTileType{}; + constexpr auto ds_spans = decltype(ds)::get_distributed_spans(); + sweep_tile_span(ds_spans[number<0>{}], [&](auto idx0) { + constexpr auto i_idx = make_tuple(idx0); + sweep_tile_span(ds_spans[number<1>{}], [&](auto idx1) { + constexpr auto i_j_idx = make_tuple(idx0, idx1); + bool undrop_flag = p[i_j_idx] >= 0; + ds(i_j_idx) = p[i_j_idx] * (!FmhaDropout::IsDropout || undrop_flag + ? (dp_acc[i_j_idx] - d[i_idx]) + : d[i_idx]); + }); + }); + + if constexpr(kHasBiasGrad) + { + const auto dbias = [&]() { + if constexpr(FmhaDropout::IsDropout) + { + return tile_elementwise_in( + [&rp_undrop](const auto& x) { + return type_convert(x * rp_undrop); + }, + ds); + } + else + { + return cast_tile(ds); + } + }(); + store_tile(bias_lds_write_window, dbias); + block_sync_lds(); + auto shuffled_dbias_tile = load_tile(dbias_lds_read_window); + auto dbias_tile = make_static_distributed_tensor( + Policy::template MakeBiasTileDistribution()); + shuffle_tile(dbias_tile, shuffled_dbias_tile); + store_tile(dbias_dram_window, dbias_tile); + move_tile_window(dbias_dram_window, {kM0, 0}); + __builtin_amdgcn_sched_barrier(0); + } + + // STAGE 6, SGrad^T@Q^T Gemm3 + const auto ds_gemm = cast_tile(ds); + + Policy::template SGradTFromGemm2CToGemm3A(dst_reg_tensor, ds_gemm); + + gemm_3(dk_acc, dst_reg_tensor, qt_reg_tensor); + + store_tile(ds_lds_window, ds_gemm); + + block_sync_lds(); + + auto ds_reg_tensor = load_tile(ds_lds_read_window); + auto ds_reg_tensor_next = decltype(ds_reg_tensor){}; + move_tile_window(ds_lds_read_window, {0, kK4}); + q_reg_tensor = load_tile(q_lds_read_window); + lse = load_tile(lse_lds_read_window); + + HotLoopScheduler::template GemmStagedScheduler<3>(); + __builtin_amdgcn_sched_barrier(0); + // STAGE7 SGrad@K^T Gemm4 + auto dq_acc = QGradBlockTileType{}; + clear_tile(dq_acc); + + static_for<0, k4_loops, 1>{}([&](auto i_k4) { + if constexpr(i_k4 < k4_loops - 1) + { + ds_reg_tensor_next = load_tile(ds_lds_read_window); + move_tile_window(ds_lds_read_window, {0, kK4}); + } + auto kt_reg_tensor_slice = get_slice_tile(kt_reg_tensor, + sequence<0, i_k4 * kK4>{}, + sequence{}); + gemm_4(dq_acc, ds_reg_tensor, kt_reg_tensor_slice); + + if constexpr(i_k4 < k4_loops - 1) + { + ds_reg_tensor.get_thread_buffer() = ds_reg_tensor_next.get_thread_buffer(); + } + }); + move_tile_window(ds_lds_read_window, {0, -kN0}); + + do_reg_tensor = load_tile(do_lds_read_window); + d = load_tile(d_lds_read_window); + + HotLoopScheduler::template GemmStagedScheduler<4>(); + + // QGrad Scale + if constexpr(FmhaDropout::IsDropout) + { + tile_elementwise_inout([&scale_rp_undrop](auto& x) { x = x * scale_rp_undrop; }, + dq_acc); + } + else + { + tile_elementwise_inout([&raw_scale](auto& x) { x = x * raw_scale; }, dq_acc); + } + if constexpr(kIsDeterministic) + { + store_tile(dq_dram_window, dq_acc); + } + else + { + update_tile(dq_dram_window, dq_acc); + } + move_tile_window(dq_dram_window, {kM0, 0}); + + i_total_loops += 1; + seqlen_q_step += kM0; + } + __builtin_amdgcn_sched_barrier(0); + + // Tail + auto s_acc = SPBlockTileType{}; + + // STAGE 1, Q@K Gemm0 + s_acc = gemm_0(q_reg_tensor, k_reg_tensor); + + // STAGE 2, Scale, Add bias, Mask, Softmax, Dropout + if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS) + { + const auto bias_tile = load_tile(bias_dram_window); + auto shuffled_bias_tile = make_static_distributed_tensor( + Policy::template MakeShuffledBiasTileDistribution()); + shuffle_tile(shuffled_bias_tile, bias_tile); + store_tile(bias_lds_write_window, shuffled_bias_tile); + block_sync_lds(); + auto bias_s_tile = load_tile(bias_s_lds_read_window); + tile_elementwise_inout( + [&](auto& x, const auto& y) { + x = scale * x + log2e_v * type_convert(y); + }, + s_acc, + bias_s_tile); + } + else if constexpr(BiasEnum == BlockAttentionBiasEnum::ALIBI) + { + constexpr auto s_spans = decltype(s_acc)::get_distributed_spans(); + 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 = seqlen_q_step + 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; + position_encoding.update(s_acc(i_j_idx), row, col); + }); + }); + } + + if constexpr(kPadSeqLenK || FmhaMask::IsMasking) + { + bool need_perpixel_check = mask.IsEdgeTile( + seqlen_q_step, k_origin.at(number<0>{}), number{}, number{}); + if(need_perpixel_check) + { + set_tile_if(s_acc, -numeric::infinity(), [&](auto tile_idx) { + const auto row = seqlen_q_step + tile_idx.at(number<0>{}); + const auto col = k_origin.at(number<0>{}) + tile_idx.at(number<1>{}); + return mask.IsOutOfBound(row, col); + }); + } + } + + static const auto get_validated_lse = [](LSEDataType raw_lse) { + if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS || + FmhaMask::IsMasking) + { + return raw_lse == -numeric::infinity() ? type_convert(0.f) + : raw_lse; + } + else + { + return raw_lse; + } + }; + + auto p = SPBlockTileType{}; + constexpr auto p_spans = decltype(p)::get_distributed_spans(); + sweep_tile_span(p_spans[number<0>{}], [&](auto idx0) { + constexpr auto i_idx = make_tuple(idx0); + auto row_lse = log2e_v * get_validated_lse(lse[i_idx]); + + sweep_tile_span(p_spans[number<1>{}], [&](auto idx1) { + constexpr auto i_j_idx = make_tuple(idx0, idx1); + if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS || + BiasEnum == BlockAttentionBiasEnum::ALIBI) + { + p(i_j_idx) = exp2(s_acc[i_j_idx] - row_lse); + } + else + { + p(i_j_idx) = exp2(scale * s_acc[i_j_idx] - row_lse); + } + }); + }); + + if constexpr(FmhaDropout::IsDropout) + { + dropout.template Run( + seqlen_q_step, k_origin.at(number<0>{}), p, randval_dram_window); + } + + // STAGE 3, P^T@OGrad^T Gemm1 + const auto p_gemm = [&]() { + if constexpr(FmhaDropout::IsDropout) + { + return tile_elementwise_in( + [](const auto& x) { return type_convert(x > 0.f ? x : 0.f); }, p); + } + else + { + return cast_tile(p); + } + }(); + + Policy::template PTFromGemm0CToGemm1A( + pt_reg_tensor, p_gemm); + auto dot_reg_tensor = load_tile(dot_lds_read_window); + gemm_1(dv_acc, pt_reg_tensor, dot_reg_tensor); + + HotLoopScheduler::template GemmStagedScheduler<1>(); + + // STAGE 4, OGrad@V Gemm2 + auto dp_acc = SPGradBlockTileType{}; + + auto qt_reg_tensor = load_tile(qt_lds_read_window); + + dp_acc = gemm_2(do_reg_tensor, v_reg_tensor); + + HotLoopScheduler::template GemmStagedScheduler<2>(); + + // STAGE 5, P^T(PGrad^T - D) + auto ds = SPGradBlockTileType{}; + constexpr auto ds_spans = decltype(ds)::get_distributed_spans(); + sweep_tile_span(ds_spans[number<0>{}], [&](auto idx0) { + constexpr auto i_idx = make_tuple(idx0); + sweep_tile_span(ds_spans[number<1>{}], [&](auto idx1) { + constexpr auto i_j_idx = make_tuple(idx0, idx1); + bool undrop_flag = p[i_j_idx] >= 0; + ds(i_j_idx) = p[i_j_idx] * (!FmhaDropout::IsDropout || undrop_flag + ? (dp_acc[i_j_idx] - d[i_idx]) + : d[i_idx]); + }); + }); + + if constexpr(kHasBiasGrad) + { + const auto dbias = [&]() { + if constexpr(FmhaDropout::IsDropout) + { + return tile_elementwise_in( + [&rp_undrop](const auto& x) { + return type_convert(x * rp_undrop); + }, + ds); + } + else + { + return cast_tile(ds); + } + }(); + store_tile(bias_lds_write_window, dbias); + block_sync_lds(); + auto shuffled_dbias_tile = load_tile(dbias_lds_read_window); + auto dbias_tile = make_static_distributed_tensor( + Policy::template MakeBiasTileDistribution()); + shuffle_tile(dbias_tile, shuffled_dbias_tile); + store_tile(dbias_dram_window, dbias_tile); + } + + // STAGE 6, SGrad^T@Q^T Gemm3 + const auto ds_gemm = cast_tile(ds); + + Policy::template SGradTFromGemm2CToGemm3A(dst_reg_tensor, ds_gemm); + + gemm_3(dk_acc, dst_reg_tensor, qt_reg_tensor); + store_tile(ds_lds_window, ds_gemm); + + block_sync_lds(); + + auto ds_reg_tensor = load_tile(ds_lds_read_window); + auto ds_reg_tensor_next = decltype(ds_reg_tensor){}; + move_tile_window(ds_lds_read_window, {0, kK4}); + + HotLoopScheduler::template GemmStagedScheduler<3>(); + // STAGE 7, SGrad@K^T Gemm4 + auto dq_acc = QGradBlockTileType{}; + clear_tile(dq_acc); + + static_for<0, k4_loops, 1>{}([&](auto i_k4) { + if constexpr(i_k4 < k4_loops - 1) + { + ds_reg_tensor_next = load_tile(ds_lds_read_window); + move_tile_window(ds_lds_read_window, {0, kK4}); + } + auto kt_reg_tensor_slice = get_slice_tile( + kt_reg_tensor, sequence<0, i_k4 * kK4>{}, sequence{}); + + gemm_4(dq_acc, ds_reg_tensor, kt_reg_tensor_slice); + if constexpr(i_k4 < k4_loops - 1) + { + ds_reg_tensor.get_thread_buffer() = ds_reg_tensor_next.get_thread_buffer(); + } + }); + + HotLoopScheduler::template GemmStagedScheduler<4>(); + + // Results Scale + if constexpr(FmhaDropout::IsDropout) + { + tile_elementwise_inout([&scale_rp_undrop](auto& x) { x = x * scale_rp_undrop; }, + dq_acc); + tile_elementwise_inout([&scale_rp_undrop](auto& x) { x = x * scale_rp_undrop; }, + dk_acc); + tile_elementwise_inout([&rp_undrop](auto& x) { x = x * rp_undrop; }, dv_acc); + } + else + { + tile_elementwise_inout([&raw_scale](auto& x) { x = x * raw_scale; }, dq_acc); + tile_elementwise_inout([&raw_scale](auto& x) { x = x * raw_scale; }, dk_acc); + } + + if constexpr(kIsDeterministic) + { + store_tile(dq_dram_window, dq_acc); + } + else + { + update_tile(dq_dram_window, dq_acc); + } + + return make_tuple(dk_acc, dv_acc); + } +}; + +} // namespace ck_tile diff --git a/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_ks_kts_vr.hpp b/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_ks_kts_vr.hpp deleted file mode 100644 index 3444567508..0000000000 --- a/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_ks_kts_vr.hpp +++ /dev/null @@ -1,848 +0,0 @@ -// 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/block/block_dropout.hpp" -#include "ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_ks_kts_vr_default_policy.hpp" -#include "ck_tile/ops/reduce/block/block_reduce.hpp" - -namespace ck_tile { - -template -struct BlockFmhaBwdDQDKDVPipelineKSKTSVR -{ - using QDataType = remove_cvref_t; - using KDataType = remove_cvref_t; - using VDataType = remove_cvref_t; - using GemmDataType = remove_cvref_t; - using BiasDataType = remove_cvref_t; - using LSEDataType = remove_cvref_t; - using AccDataType = remove_cvref_t; - using DDataType = remove_cvref_t; - using RandValOutputDataType = remove_cvref_t; - using ODataType = remove_cvref_t; - using OGradDataType = remove_cvref_t; - using QGradDataType = remove_cvref_t; - using KGradDataType = remove_cvref_t; - using VGradDataType = remove_cvref_t; - using BiasGradDataType = remove_cvref_t; - using FmhaMask = remove_cvref_t; - - using BlockFmhaShape = remove_cvref_t; - - static constexpr index_t kBlockPerCu = Problem::kBlockPerCu; - 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 kK1 = BlockFmhaShape::kK1; - static constexpr index_t kK2 = BlockFmhaShape::kK2; - static constexpr index_t kK3 = BlockFmhaShape::kK3; - static constexpr index_t kK4 = BlockFmhaShape::kK4; - static constexpr index_t kQKHeaddim = BlockFmhaShape::kQKHeaddim; - static constexpr index_t kVHeaddim = BlockFmhaShape::kVHeaddim; - - static constexpr bool kQLoadOnce = false; - static constexpr bool kQTLoadOnce = false; - static constexpr bool kKLoadOnce = true; - static constexpr bool kKTLoadOnce = true; - static constexpr bool kVLoadOnce = true; - static constexpr bool kOGradLoadOnce = false; - static constexpr bool kOGradTLoadOnce = false; - - 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 kHasBiasGrad = Problem::kHasBiasGrad; - static constexpr bool kHasDropout = Problem::kHasDropout; - - // 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 = - kPadHeadDimV ? 1 : Policy::template GetAlignmentV(); - static constexpr index_t kAlignmentO = - kPadHeadDimV ? 1 : Policy::template GetAlignmentO(); - static constexpr index_t kAlignmentOGrad = - kPadHeadDimV ? 1 : Policy::template GetAlignmentOGrad(); - static constexpr index_t kAlignmentQGrad = - kPadHeadDimQ ? 2 : Policy::template GetAlignmentQGrad(); - static constexpr index_t kAlignmentKGrad = - kPadHeadDimQ ? 1 : Policy::template GetAlignmentKGrad(); - static constexpr index_t kAlignmentVGrad = - kPadHeadDimV ? 1 : Policy::template GetAlignmentVGrad(); - static constexpr index_t kAlignmentBias = - kPadSeqLenK ? 1 : Policy::template GetTransposedAlignmentBias(); - - static constexpr const char* name = "ks_kts_vr"; - - 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, - const QTDramBlockWindowTmp& qt_dram_block_window_tmp, - const KDramBlockWindowTmp& k_dram_block_window_tmp, - const KTDramBlockWindowTmp& kt_dram_block_window_tmp, - const VDramBlockWindowTmp& v_dram_block_window_tmp, - const BiasDramBlockWindowTmp& bias_dram_block_window_tmp, - const RandValDramBlockWindowTmp& randval_dram_block_window_tmp, - const OGradDramBlockWindowTmp& do_dram_block_window_tmp, - const OGradTDramBlockWindowTmp& dot_dram_block_window_tmp, - const LSEDramBlockWindowTmp& lse_dram_block_window_tmp, - const DDramBlockWindowTmp& d_dram_block_window_tmp, - const QGradDramBlockWindowTmp& dq_dram_block_window_tmp, - const BiasGradDramBlockWindowTmp& dbias_dram_block_window_tmp, - FmhaMask mask, - PositionEncoding position_encoding, - float raw_scale, -#if CK_TILE_FMHA_FWD_FAST_EXP2 - float scale, -#endif - float rp_undrop, - float scale_rp_undrop, - void* smem_ptr, - BlockDropout& dropout) const - { - static_assert( - std::is_same_v> && - std::is_same_v> && - std::is_same_v> && - std::is_same_v> && - std::is_same_v> && - std::is_same_v> && - std::is_same_v> && - std::is_same_v> && - std::is_same_v> && - std::is_same_v>, - "wrong!"); - - static_assert(kM0 == QDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] && - kQKHeaddim == QTDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] && - kN0 == KDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] && - kQKHeaddim == KTDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] && - kN0 == VDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] && - kM0 == BiasDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] && - kN0 == BiasDramBlockWindowTmp{}.get_window_lengths()[number<1>{}] && - kM0 == OGradDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] && - kVHeaddim == - OGradTDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] && - kM0 == LSEDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] && - kM0 == DDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] && - kM0 == QGradDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] && - kM0 == BiasGradDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] && - kN0 == BiasGradDramBlockWindowTmp{}.get_window_lengths()[number<1>{}], - "wrong!"); - - // Q tile in LDS - QDataType* q_lds_ptr = static_cast(static_cast( - static_cast(smem_ptr) + Policy::template GetSmemSizeK() + - Policy::template GetSmemSizeKT())); - auto q_lds = make_tensor_view( - q_lds_ptr, Policy::template MakeQLdsBlockDescriptor()); - auto q_lds_window = - make_tile_window(q_lds, make_tuple(number{}, number{}), {0, 0}); - - // QT tile in LDS - QDataType* qt_lds_ptr = static_cast(static_cast( - static_cast(smem_ptr) + Policy::template GetSmemSizeK() + - Policy::template GetSmemSizeKT())); - auto qt_lds = make_tensor_view( - qt_lds_ptr, Policy::template MakeQTLdsBlockDescriptor()); - auto qt_lds_window = - make_tile_window(qt_lds, make_tuple(number{}, number{}), {0, 0}); - - // K tile in LDS - auto k_lds = make_tensor_view( - reinterpret_cast(smem_ptr), - Policy::template MakeKLdsBlockDescriptor()); - auto k_lds_window = - make_tile_window(k_lds, make_tuple(number{}, number{}), {0, 0}); - - // KT tile in LDS - KDataType* kt_lds_ptr = static_cast(static_cast( - static_cast(smem_ptr) + Policy::template GetSmemSizeK())); - auto kt_lds = make_tensor_view( - kt_lds_ptr, Policy::template MakeKTLdsBlockDescriptor()); - auto kt_lds_window = - make_tile_window(kt_lds, make_tuple(number{}, number{}), {0, 0}); - - // OGrad tile in LDS - OGradDataType* do_lds_ptr = static_cast(static_cast( - static_cast(smem_ptr) + Policy::template GetSmemSizeK() + - Policy::template GetSmemSizeKT())); - auto do_lds = make_tensor_view( - do_lds_ptr, Policy::template MakeOGradLdsBlockDescriptor()); - auto do_lds_window = - make_tile_window(do_lds, make_tuple(number{}, number{}), {0, 0}); - - // OGradT tile in LDS - OGradDataType* dot_lds_ptr = static_cast(static_cast( - static_cast(smem_ptr) + Policy::template GetSmemSizeK() + - Policy::template GetSmemSizeKT())); - auto dot_lds = make_tensor_view( - dot_lds_ptr, Policy::template MakeOGradTLdsBlockDescriptor()); - auto dot_lds_window = - make_tile_window(dot_lds, make_tuple(number{}, number{}), {0, 0}); - - // SGrad tile in LDS - GemmDataType* ds_lds_ptr = static_cast(static_cast( - static_cast(smem_ptr) + Policy::template GetSmemSizeK() + - Policy::template GetSmemSizeKT())); - auto ds_lds = make_tensor_view( - ds_lds_ptr, Policy::template MakeSGradLdsBlockDescriptor()); - auto ds_lds_window = - make_tile_window(ds_lds, make_tuple(number{}, number{}), {0, 0}); - - // BiasT/BiasGradT tile in LDS, use the same size and layout - BiasDataType* biast_lds_ptr = static_cast(static_cast( - static_cast(smem_ptr) + Policy::template GetSmemSizeK() + - Policy::template GetSmemSizeKT())); - auto biast_lds = make_tensor_view( - biast_lds_ptr, Policy::template MakeBiasTLdsBlockDescriptor()); - auto biast_lds_shuffle_window = - make_tile_window(biast_lds, make_tuple(number{}, number{}), {0, 0}); - auto dbiast_lds_shuffle_window = - make_tile_window(biast_lds, - make_tuple(number{}, number{}), - {0, 0}, - Policy::template MakeShuffledBiasTileDistribution()); - - static_assert(std::is_same_v, - "BiasDataType and BiasGradDataType should be the same!"); - - // Block GEMM - constexpr auto gemm_0 = Policy::template GetQKBlockGemm(); - constexpr auto gemm_1 = Policy::template GetPTOGradTBlockGemm(); - constexpr auto gemm_2 = Policy::template GetOGradVBlockGemm(); - constexpr auto gemm_3 = Policy::template GetSGradTQTBlockGemm(); - constexpr auto gemm_4 = Policy::template GetSGradKTBlockGemm(); - - auto v_dram_window = make_tile_window( - v_dram_block_window_tmp.get_bottom_tensor_view(), - v_dram_block_window_tmp.get_window_lengths(), - v_dram_block_window_tmp.get_window_origin(), - Policy::template MakeVInRegDramTileDistribution()); - - auto v = load_tile(v_dram_window); // persistent V register tile - - using SPTBlockTileType = decltype(gemm_0.MakeCBlockTile()); - using SPGradTBlockTileType = decltype(gemm_2.MakeCBlockTile()); - using QGradBlockTileType = decltype(gemm_4.MakeCBlockTile()); - - // init VGrad & KGrad - auto dv_acc = decltype(gemm_1.MakeCBlockTile()){}; - auto dk_acc = decltype(gemm_3.MakeCBlockTile()){}; - - clear_tile(dv_acc); - clear_tile(dk_acc); - - auto k_dram_window = make_tile_window( - k_dram_block_window_tmp.get_bottom_tensor_view(), - k_dram_block_window_tmp.get_window_lengths(), - k_dram_block_window_tmp.get_window_origin(), - Policy::template MakeKDramTileDistribution()); // K DRAM tile window for - // load - - __builtin_amdgcn_sched_barrier(0); - const auto k_origin = k_dram_window.get_window_origin(); - const auto [seqlen_q_start, seqlen_q_end] = - mask.GetTileRangeAlongY(k_origin.at(number<0>{}), number{}, number{}); - - const auto num_total_loop = integer_divide_ceil(seqlen_q_end - seqlen_q_start, kM0); - - // check early exit if masked and no work to do. - if constexpr(FmhaMask::IsMasking) - { - if(num_total_loop <= 0) - { - // Note: here dk_acc&dv_acc are all cleard, return it - // Note: v loaded but no fence, ignore it. - return ck_tile::make_tuple(dk_acc, dv_acc); - } - } - - auto k_block_tile = load_tile(k_dram_window); - - store_tile(k_lds_window, k_block_tile); // // persistent K in LDS - - auto kt_dram_block_window = kt_dram_block_window_tmp; - - auto kt_dram_window = make_tile_window( - kt_dram_block_window.get_bottom_tensor_view(), - kt_dram_block_window.get_window_lengths(), - kt_dram_block_window.get_window_origin(), - Policy::template MakeKTDramTileDistribution()); // K^T DRAM tile window for - // load - - auto kt_block_tile = load_tile(kt_dram_window); - - auto kt_shuffle_tmp = make_static_distributed_tensor( - Policy::template MakeShuffledKTRegBlockDescriptor()); - shuffle_tile(kt_shuffle_tmp, kt_block_tile); - - store_tile(kt_lds_window, kt_shuffle_tmp); // persistent K^T in LDS - - auto q_dram_block_window = - make_tile_window(q_dram_block_window_tmp.get_bottom_tensor_view(), - q_dram_block_window_tmp.get_window_lengths(), - {seqlen_q_start, 0}); - - auto qt_dram_block_window = - make_tile_window(qt_dram_block_window_tmp.get_bottom_tensor_view(), - qt_dram_block_window_tmp.get_window_lengths(), - {0, seqlen_q_start}); - - auto do_dram_block_window = - make_tile_window(do_dram_block_window_tmp.get_bottom_tensor_view(), - do_dram_block_window_tmp.get_window_lengths(), - {seqlen_q_start, 0}); - - auto dot_dram_block_window = - make_tile_window(dot_dram_block_window_tmp.get_bottom_tensor_view(), - dot_dram_block_window_tmp.get_window_lengths(), - {0, seqlen_q_start}); - - auto dq_dram_block_window = - make_tile_window(dq_dram_block_window_tmp.get_bottom_tensor_view(), - dq_dram_block_window_tmp.get_window_lengths(), - {seqlen_q_start, 0}); - - auto lse_dram_block_window = - make_tile_window(lse_dram_block_window_tmp.get_bottom_tensor_view(), - lse_dram_block_window_tmp.get_window_lengths(), - {seqlen_q_start}); - - auto d_dram_block_window = - make_tile_window(d_dram_block_window_tmp.get_bottom_tensor_view(), - d_dram_block_window_tmp.get_window_lengths(), - {seqlen_q_start}); - - const auto bias_origin = bias_dram_block_window_tmp.get_window_origin(); - auto bias_dram_block_window = - make_tile_window(bias_dram_block_window_tmp.get_bottom_tensor_view(), - bias_dram_block_window_tmp.get_window_lengths(), - {seqlen_q_start, bias_origin.at(number<1>{})}); // M/N - - const auto dbias_origin = dbias_dram_block_window_tmp.get_window_origin(); - auto dbias_dram_block_window = - make_tile_window(dbias_dram_block_window_tmp.get_bottom_tensor_view(), - dbias_dram_block_window_tmp.get_window_lengths(), - {seqlen_q_start, dbias_origin.at(number<1>{})}); // M/N - - auto qt_dram_window = - make_tile_window(qt_dram_block_window.get_bottom_tensor_view(), - qt_dram_block_window.get_window_lengths(), - qt_dram_block_window.get_window_origin(), - Policy::template MakeQTDramTileDistribution()); - - auto dot_dram_window = - make_tile_window(dot_dram_block_window.get_bottom_tensor_view(), - dot_dram_block_window.get_window_lengths(), - dot_dram_block_window.get_window_origin(), - Policy::template MakeOGradTDramTileDistribution()); - - auto lse_dram_window = make_tile_window( - lse_dram_block_window.get_bottom_tensor_view(), - lse_dram_block_window.get_window_lengths(), - lse_dram_block_window.get_window_origin(), - Policy::template MakeLSEDDramTileDistribution()); - - auto d_dram_window = make_tile_window( - d_dram_block_window.get_bottom_tensor_view(), - d_dram_block_window.get_window_lengths(), - d_dram_block_window.get_window_origin(), - Policy::template MakeLSEDDramTileDistribution()); - - auto bias_dram_window = - make_tile_window(bias_dram_block_window.get_bottom_tensor_view(), - bias_dram_block_window.get_window_lengths(), - bias_dram_block_window.get_window_origin(), - Policy::template MakeBiasTileDistribution()); - - auto biast_lds_window = - make_tile_window(biast_lds_shuffle_window.get_bottom_tensor_view(), - biast_lds_shuffle_window.get_window_lengths(), - biast_lds_shuffle_window.get_window_origin(), - Policy::template MakeBiasTTileDistribution()); - - auto randval_dram_window = dropout.MakeRandvalDramWindow( - randval_dram_block_window_tmp, seqlen_q_start); - - index_t i_total_loops = 0; - constexpr index_t k0_loops = kQKHeaddim / kK0; - constexpr index_t k1_loops = kM0 / kK1; - constexpr index_t k2_loops = kVHeaddim / kK2; - constexpr index_t k3_loops = kM0 / kK3; - constexpr index_t k4_loops = kN0 / kK4; - do - { - auto q_dram_window = make_tile_window( - q_dram_block_window.get_bottom_tensor_view(), - q_dram_block_window.get_window_lengths(), - q_dram_block_window.get_window_origin(), - Policy::template MakeQDramTileDistribution()); // Q DRAM tile window for - // load - - auto do_dram_window = make_tile_window( - do_dram_block_window.get_bottom_tensor_view(), - do_dram_block_window.get_window_lengths(), - do_dram_block_window.get_window_origin(), - Policy::template MakeOGradDramTileDistribution()); // OGrad DRAM tile - // window for load - - // STAGE 1, Q@K Gemm0 - auto st_acc = SPTBlockTileType{}; - - auto q_block_tile = load_tile(q_dram_window); - { - move_tile_window(q_dram_window, {0, kK0}); - - clear_tile(st_acc); // Initialize S^T - - store_tile(q_lds_window, q_block_tile); // LDS write 0 - q_block_tile = load_tile(q_dram_window); // global read 1 - } - - 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(st_acc, - q_lds_window, - get_slice_tile(k_lds_window, - sequence<0, i_k0 * kK0>{}, - sequence{})); - block_sync_lds(); - move_tile_window(q_dram_window, {0, kK0}); - - store_tile(q_lds_window, - q_block_tile); // LDS write i + 1 - q_block_tile = load_tile(q_dram_window); // global read i + 2 - }); - } - - const auto dot_prefetch = load_tile(dot_dram_window); // prefetch load OGrad^T tile - { // tail - block_sync_lds(); - gemm_0(st_acc, - q_lds_window, - get_slice_tile(k_lds_window, - sequence<0, (k0_loops - 2) * kK0>{}, - sequence{})); - block_sync_lds(); - - store_tile(q_lds_window, q_block_tile); - block_sync_lds(); - - gemm_0(st_acc, - q_lds_window, - get_slice_tile(k_lds_window, - sequence<0, (k0_loops - 1) * kK0>{}, - sequence{})); - } - - // STAGE 2, Scale, Add bias, Mask, Softmax, Dropout - if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS) - { - block_sync_lds(); - auto bias_shuffle_tmp = make_static_distributed_tensor( - Policy::template MakeShuffledBiasTileDistribution()); - shuffle_tile(bias_shuffle_tmp, bias_tile); - store_tile(biast_lds_shuffle_window, bias_shuffle_tmp); - block_sync_lds(); - auto biast_tile = load_tile(biast_lds_window); - tile_elementwise_inout( - [&](auto& x, const auto& y) { -#if !CK_TILE_FMHA_FWD_FAST_EXP2 - x = raw_scale * x + type_convert(y); -#else - x = scale * x + log2e_v * type_convert(y); -#endif - }, - st_acc, - biast_tile); - move_tile_window(bias_dram_window, {kM0, 0}); - } - else if constexpr(BiasEnum == BlockAttentionBiasEnum::ALIBI) - { - const auto q_origin = q_dram_block_window.get_window_origin(); - constexpr auto st_spans = decltype(st_acc)::get_distributed_spans(); - sweep_tile_span(st_spans[number<0>{}], [&](auto idx0) { - sweep_tile_span(st_spans[number<1>{}], [&](auto idx1) { - const auto tile_idx = get_x_indices_from_distributed_indices( - st_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); - -#if !CK_TILE_FMHA_FWD_FAST_EXP2 - st_acc(i_j_idx) *= raw_scale; -#else - st_acc(i_j_idx) *= scale; -#endif - position_encoding.update(st_acc(i_j_idx), row, col); - }); - }); - } - else - { -#if !CK_TILE_FMHA_FWD_FAST_EXP2 - tile_elementwise_inout([&raw_scale](auto& x) { x = x * raw_scale; }, st_acc); -#endif - } - - if constexpr(kPadSeqLenK || FmhaMask::IsMasking) - { - const auto q_origin = q_dram_block_window.get_window_origin(); - bool need_perpixel_check = mask.IsEdgeTile(q_origin.at(number<0>{}), - k_origin.at(number<0>{}), - number{}, - number{}); - if(need_perpixel_check) - { - set_tile_if(st_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); - }); - } - } - - const auto lse = load_tile(lse_dram_window); - - static const auto get_validated_lse = [](LSEDataType raw_lse) { - if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS || - FmhaMask::IsMasking) - { - return raw_lse == -numeric::infinity() - ? type_convert(0.f) - : raw_lse; - } - else - { - return raw_lse; - } - }; - - auto pt = SPTBlockTileType{}; - constexpr auto pt_spans = decltype(pt)::get_distributed_spans(); - sweep_tile_span(pt_spans[number<0>{}], [&](auto idx0) { - constexpr auto i_idx = make_tuple(idx0); -#if CK_TILE_FMHA_FWD_FAST_EXP2 - auto row_lse = log2e_v * get_validated_lse(lse[i_idx]); -#endif - sweep_tile_span(pt_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) - { - pt(i_j_idx) = exp2(st_acc[i_j_idx] - row_lse); - } - else - { - pt(i_j_idx) = exp2(scale * st_acc[i_j_idx] - row_lse); - } -#else - pt(i_j_idx) = exp(st_acc[i_j_idx] - get_validated_lse(lse[i_idx])); -#endif - }); - }); - - auto dot_shuffle_tmp = make_static_distributed_tensor( - Policy::template MakeShuffledOGradTRegBlockDescriptor()); - block_sync_lds(); - { - shuffle_tile(dot_shuffle_tmp, dot_prefetch); - store_tile(dot_lds_window, - dot_shuffle_tmp); // store the prefetch - } - move_tile_window(dot_dram_window, {0, kK1}); - - if constexpr(kHasDropout) - { - dropout.Run( - seqlen_q_start + i_total_loops * kM0, pt, randval_dram_window); - } - - // STAGE 3, P^T@OGrad^T Gemm1 - const auto pt_gemm = [&]() { - if constexpr(kHasDropout) - { - return tile_elementwise_in( - [](const auto& x) { return type_convert(x > 0.f ? x : 0.f); }, - pt); - } - else - { - return cast_tile(pt); - } - }(); - - if constexpr(k1_loops > 1) - { - static_for<0, k1_loops - 1, 1>{}([&](auto i_k1) { - const auto dot = load_tile(dot_dram_window); // load next OGrad^T - block_sync_lds(); - gemm_1(dv_acc, - get_slice_tile(pt_gemm, - sequence{}, - sequence<(i_k1 + 1) * kK1, kN0>{}), - dot_lds_window); - block_sync_lds(); - shuffle_tile(dot_shuffle_tmp, dot); - store_tile(dot_lds_window, - dot_shuffle_tmp); // store the prefetch - - move_tile_window(dot_dram_window, {0, kK1}); - }); - } - auto do_block_tile = load_tile(do_dram_window); // prefetch load OGrad tile - // tail - { - block_sync_lds(); - gemm_1(dv_acc, - get_slice_tile( - pt_gemm, sequence<(k1_loops - 1) * kK1, 0>{}, sequence{}), - dot_lds_window); - block_sync_lds(); - } - - // STAGE 4, OGrad@V Gemm2 - auto dpt_acc = SPGradTBlockTileType{}; - - { - move_tile_window(do_dram_window, {0, kK2}); - - clear_tile(dpt_acc); // Initialize PGrad^T - - store_tile(do_lds_window, do_block_tile); // LDS write 0 - do_block_tile = load_tile(do_dram_window); // global read 1 - } - - if constexpr(k2_loops > 2) - { - static_for<0, k2_loops - 2, 1>{}([&](auto i_k2) { - block_sync_lds(); - gemm_2(dpt_acc, - do_lds_window, - get_slice_tile( - v, sequence<0, i_k2 * kK2>{}, sequence{})); - block_sync_lds(); - move_tile_window(do_dram_window, {0, kK2}); - - store_tile(do_lds_window, - do_block_tile); // LDS write i + 1 - do_block_tile = load_tile(do_dram_window); // global read i + 2 - }); - } - - const auto qt_prefetch = load_tile(qt_dram_window); // prefetch load Q^T tile - { // tail - block_sync_lds(); - gemm_2(dpt_acc, - do_lds_window, - get_slice_tile(v, - sequence<0, (k2_loops - 2) * kK2>{}, - sequence{})); - block_sync_lds(); - - store_tile(do_lds_window, do_block_tile); - block_sync_lds(); - - gemm_2(dpt_acc, - do_lds_window, - get_slice_tile(v, - sequence<0, (k2_loops - 1) * kK2>{}, - sequence{})); - } - - // STAGE 5, P^T(PGrad^T - D) - const auto d = load_tile(d_dram_window); - - auto dst = SPGradTBlockTileType{}; - constexpr auto dst_spans = decltype(dst)::get_distributed_spans(); - sweep_tile_span(dst_spans[number<0>{}], [&](auto idx0) { - constexpr auto i_idx = make_tuple(idx0); - sweep_tile_span(dst_spans[number<1>{}], [&](auto idx1) { - constexpr auto i_j_idx = make_tuple(idx0, idx1); - bool undrop_flag = pt[i_j_idx] >= 0; - dst(i_j_idx) = - pt[i_j_idx] * - (!kHasDropout || undrop_flag ? (dpt_acc[i_j_idx] - d[i_idx]) : d[i_idx]); - }); - }); - - if constexpr(kHasBiasGrad) - { - const auto dbiast = [&]() { - if constexpr(kHasDropout) - { - return tile_elementwise_in( - [&rp_undrop](const auto& x) { - return type_convert(x * rp_undrop); - }, - dst); - } - else - { - return cast_tile(dst); - } - }(); - store_tile(biast_lds_shuffle_window, dbiast); - block_sync_lds(); - auto dbiast_tile = load_tile(dbiast_lds_shuffle_window); - auto dbiast_shuffle_tmp = make_static_distributed_tensor( - Policy::template MakeBiasTileDistribution()); - shuffle_tile(dbiast_shuffle_tmp, dbiast_tile); - store_tile(dbias_dram_block_window, dbiast_shuffle_tmp); - move_tile_window(dbias_dram_block_window, {kM0, 0}); - } - - // STAGE 6, SGrad^T@Q^T Gemm3 - auto qt_shuffle_tmp = make_static_distributed_tensor( - Policy::template MakeShuffledQTRegBlockDescriptor()); - block_sync_lds(); - { - shuffle_tile(qt_shuffle_tmp, qt_prefetch); - store_tile(qt_lds_window, - qt_shuffle_tmp); // store the prefetch - } - move_tile_window(qt_dram_window, {0, kK3}); - - const auto dst_gemm = cast_tile(dst); - - if constexpr(k3_loops > 1) - { - static_for<0, k3_loops - 1, 1>{}([&](auto i_k3) { - const auto qt = load_tile(qt_dram_window); // load next Q^T - block_sync_lds(); - gemm_3(dk_acc, - get_slice_tile(dst_gemm, - sequence{}, - sequence<(i_k3 + 1) * kK3, kN0>{}), - qt_lds_window); - block_sync_lds(); - shuffle_tile(qt_shuffle_tmp, qt); - store_tile(qt_lds_window, - qt_shuffle_tmp); // store the prefetch - - move_tile_window(qt_dram_window, {0, kK3}); - }); - } - // tail - { - block_sync_lds(); - gemm_3(dk_acc, - get_slice_tile( - dst_gemm, sequence<(k3_loops - 1) * kK3, 0>{}, sequence{}), - qt_lds_window); - block_sync_lds(); - } - - // STAGE 7, SGrad@K^T Gemm4 - store_tile(ds_lds_window, dst_gemm); - - auto dq_acc = QGradBlockTileType{}; - clear_tile(dq_acc); // Initialize QGrad - - block_sync_lds(); - - static_for<0, k4_loops, 1>{}([&](auto i_k4) { - gemm_4(dq_acc, - get_slice_tile(ds_lds_window, - sequence<0, i_k4 * kK4>{}, - sequence{}), - get_slice_tile(kt_lds_window, - sequence<0, i_k4 * kK4>{}, - sequence{})); - }); - - // QGrad Scale - if constexpr(kHasDropout) - { - tile_elementwise_inout([&scale_rp_undrop](auto& x) { x = x * scale_rp_undrop; }, - dq_acc); - } - else - { - tile_elementwise_inout([&raw_scale](auto& x) { x = x * raw_scale; }, dq_acc); - } - const auto dq = cast_tile(dq_acc); - update_tile(dq_dram_block_window, dq); - - // move tile windows - move_tile_window(q_dram_block_window, {kM0, 0}); - move_tile_window(dq_dram_block_window, {kM0, 0}); - move_tile_window(do_dram_block_window, {kM0, 0}); - move_tile_window(lse_dram_window, {kM0}); - move_tile_window(d_dram_window, {kM0}); - } while(++i_total_loops < num_total_loop); - - // KGrad Scale - if constexpr(kHasDropout) - { - tile_elementwise_inout([&scale_rp_undrop](auto& x) { x = x * scale_rp_undrop; }, - dk_acc); - } - else - { - tile_elementwise_inout([&raw_scale](auto& x) { x = x * raw_scale; }, dk_acc); - } - // VGrad Scale - if constexpr(kHasDropout) - { - tile_elementwise_inout([&rp_undrop](auto& x) { x = x * rp_undrop; }, dv_acc); - } - - return ck_tile::make_tuple(dk_acc, dv_acc); - } -}; - -} // namespace ck_tile diff --git a/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_ks_kts_vr_default_policy.hpp b/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_ks_kts_vr_default_policy.hpp deleted file mode 100644 index a05fbf252f..0000000000 --- a/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_ks_kts_vr_default_policy.hpp +++ /dev/null @@ -1,20 +0,0 @@ -// SPDX-License-Identifier: MIT -// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved. - -#pragma once - -#include "ck_tile/ops/fmha/pipeline/block_fmha_bwd_pipeline_default_policy.hpp" - -namespace ck_tile { - -// This pipeline is v located in regs, k & k^t located in lds. -using BlockFmhaBwdDQDKDVPipelineKSKTSVRDefaultPolicy = - BlockFmhaBwdPipelineDefaultPolicy; - -} // namespace ck_tile diff --git a/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_ks_vr.hpp b/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_ks_vr.hpp deleted file mode 100644 index dec421c1e6..0000000000 --- a/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_ks_vr.hpp +++ /dev/null @@ -1,821 +0,0 @@ -// 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/block/block_dropout.hpp" -#include "ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_ks_vr_default_policy.hpp" -#include "ck_tile/ops/reduce/block/block_reduce.hpp" - -namespace ck_tile { - -template -struct BlockFmhaBwdDQDKDVPipelineKSVR -{ - using QDataType = remove_cvref_t; - using KDataType = remove_cvref_t; - using VDataType = remove_cvref_t; - using GemmDataType = remove_cvref_t; - using BiasDataType = remove_cvref_t; - using LSEDataType = remove_cvref_t; - using AccDataType = remove_cvref_t; - using DDataType = remove_cvref_t; - using RandValOutputDataType = remove_cvref_t; - using ODataType = remove_cvref_t; - using OGradDataType = remove_cvref_t; - using QGradDataType = remove_cvref_t; - using KGradDataType = remove_cvref_t; - using VGradDataType = remove_cvref_t; - using BiasGradDataType = remove_cvref_t; - using FmhaMask = remove_cvref_t; - - using BlockFmhaShape = remove_cvref_t; - - static constexpr index_t kBlockPerCu = Problem::kBlockPerCu; - 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 kK1 = BlockFmhaShape::kK1; - static constexpr index_t kK2 = BlockFmhaShape::kK2; - static constexpr index_t kK3 = BlockFmhaShape::kK3; - static constexpr index_t kK4 = BlockFmhaShape::kK4; - static constexpr index_t kQKHeaddim = BlockFmhaShape::kQKHeaddim; - static constexpr index_t kVHeaddim = BlockFmhaShape::kVHeaddim; - - static constexpr bool kQLoadOnce = false; - static constexpr bool kQTLoadOnce = false; - static constexpr bool kKLoadOnce = true; - static constexpr bool kKTLoadOnce = false; - static constexpr bool kVLoadOnce = true; - static constexpr bool kOGradLoadOnce = false; - static constexpr bool kOGradTLoadOnce = false; - - 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 kHasBiasGrad = Problem::kHasBiasGrad; - static constexpr bool kHasDropout = Problem::kHasDropout; - - // 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 = - kPadHeadDimV ? 1 : Policy::template GetAlignmentV(); - static constexpr index_t kAlignmentO = - kPadHeadDimV ? 1 : Policy::template GetAlignmentO(); - static constexpr index_t kAlignmentOGrad = - kPadHeadDimV ? 1 : Policy::template GetAlignmentOGrad(); - static constexpr index_t kAlignmentQGrad = - kPadHeadDimQ ? 2 : Policy::template GetAlignmentQGrad(); - static constexpr index_t kAlignmentKGrad = - kPadHeadDimQ ? 1 : Policy::template GetAlignmentKGrad(); - static constexpr index_t kAlignmentVGrad = - kPadHeadDimV ? 1 : Policy::template GetAlignmentVGrad(); - static constexpr index_t kAlignmentBias = - kPadSeqLenK ? 1 : Policy::template GetTransposedAlignmentBias(); - - static constexpr const char* name = "ks_vr"; - - 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, - const QTDramBlockWindowTmp& qt_dram_block_window_tmp, - const KDramBlockWindowTmp& k_dram_block_window_tmp, - const KTDramBlockWindowTmp& /*kt_dram_block_window_tmp*/, - const VDramBlockWindowTmp& v_dram_block_window_tmp, - const BiasDramBlockWindowTmp& bias_dram_block_window_tmp, - const RandValDramBlockWindowTmp& randval_dram_block_window_tmp, - const OGradDramBlockWindowTmp& do_dram_block_window_tmp, - const OGradTDramBlockWindowTmp& dot_dram_block_window_tmp, - const LSEDramBlockWindowTmp& lse_dram_block_window_tmp, - const DDramBlockWindowTmp& d_dram_block_window_tmp, - const QGradDramBlockWindowTmp& dq_dram_block_window_tmp, - const BiasGradDramBlockWindowTmp& dbias_dram_block_window_tmp, - FmhaMask mask, - PositionEncoding position_encoding, - float raw_scale, -#if CK_TILE_FMHA_FWD_FAST_EXP2 - float scale, -#endif - float rp_undrop, - float scale_rp_undrop, - void* smem_ptr, - BlockDropout& dropout) const - { - static_assert( - std::is_same_v> && - std::is_same_v> && - std::is_same_v> && - std::is_same_v> && - std::is_same_v> && - std::is_same_v> && - std::is_same_v> && - std::is_same_v> && - std::is_same_v>, - "wrong!"); - - static_assert(kM0 == QDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] && - kQKHeaddim == QTDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] && - kN0 == KDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] && - kN0 == VDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] && - kM0 == BiasDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] && - kN0 == BiasDramBlockWindowTmp{}.get_window_lengths()[number<1>{}] && - kM0 == OGradDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] && - kVHeaddim == - OGradTDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] && - kM0 == LSEDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] && - kM0 == DDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] && - kM0 == QGradDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] && - kM0 == BiasGradDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] && - kN0 == BiasGradDramBlockWindowTmp{}.get_window_lengths()[number<1>{}], - "wrong!"); - - // Q tile in LDS - QDataType* q_lds_ptr = static_cast(static_cast( - static_cast(smem_ptr) + Policy::template GetSmemSizeK())); - auto q_lds = make_tensor_view( - q_lds_ptr, Policy::template MakeQLdsBlockDescriptor()); - auto q_lds_window = - make_tile_window(q_lds, make_tuple(number{}, number{}), {0, 0}); - - // QT tile in LDS - QDataType* qt_lds_ptr = static_cast(static_cast( - static_cast(smem_ptr) + Policy::template GetSmemSizeK())); - auto qt_lds = make_tensor_view( - qt_lds_ptr, Policy::template MakeQTLdsBlockDescriptor()); - auto qt_lds_window = - make_tile_window(qt_lds, make_tuple(number{}, number{}), {0, 0}); - - // K tile in LDS - auto k_lds = make_tensor_view( - reinterpret_cast(smem_ptr), - Policy::template MakeKLdsBlockDescriptor()); - auto k_lds_window = - make_tile_window(k_lds, make_tuple(number{}, number{}), {0, 0}); - - // KT tile in LDS - auto kt_lds = make_tensor_view( - reinterpret_cast(smem_ptr), - Policy::template MakeKLdsBlockDescriptorAsKT()); - auto kt_lds_window = - make_tile_window(kt_lds, make_tuple(number{}, number{}), {0, 0}); - - // OGrad tile in LDS - OGradDataType* do_lds_ptr = static_cast(static_cast( - static_cast(smem_ptr) + Policy::template GetSmemSizeK())); - auto do_lds = make_tensor_view( - do_lds_ptr, Policy::template MakeOGradLdsBlockDescriptor()); - auto do_lds_window = - make_tile_window(do_lds, make_tuple(number{}, number{}), {0, 0}); - - // OGradT tile in LDS - OGradDataType* dot_lds_ptr = static_cast(static_cast( - static_cast(smem_ptr) + Policy::template GetSmemSizeK())); - auto dot_lds = make_tensor_view( - dot_lds_ptr, Policy::template MakeOGradTLdsBlockDescriptor()); - auto dot_lds_window = - make_tile_window(dot_lds, make_tuple(number{}, number{}), {0, 0}); - - // SGrad tile in LDS - GemmDataType* ds_lds_ptr = static_cast(static_cast( - static_cast(smem_ptr) + Policy::template GetSmemSizeK())); - auto ds_lds = make_tensor_view( - ds_lds_ptr, Policy::template MakeSGradLdsBlockDescriptor()); - auto ds_lds_window = - make_tile_window(ds_lds, make_tuple(number{}, number{}), {0, 0}); - - // BiasT/BiasGradT tile in LDS, use the same size and layout - BiasDataType* biast_lds_ptr = static_cast(static_cast( - static_cast(smem_ptr) + Policy::template GetSmemSizeK())); - auto biast_lds = make_tensor_view( - biast_lds_ptr, Policy::template MakeBiasTLdsBlockDescriptor()); - auto biast_lds_shuffle_window = - make_tile_window(biast_lds, make_tuple(number{}, number{}), {0, 0}); - auto dbiast_lds_shuffle_window = - make_tile_window(biast_lds, - make_tuple(number{}, number{}), - {0, 0}, - Policy::template MakeShuffledBiasTileDistribution()); - - static_assert(std::is_same_v, - "BiasDataType and BiasGradDataType should be the same!"); - - // Block GEMM - constexpr auto gemm_0 = Policy::template GetQKBlockGemm(); - constexpr auto gemm_1 = Policy::template GetPTOGradTBlockGemm(); - constexpr auto gemm_2 = Policy::template GetOGradVBlockGemm(); - constexpr auto gemm_3 = Policy::template GetSGradTQTBlockGemm(); - constexpr auto gemm_4 = Policy::template GetSGradKTBlockGemm(); - - auto v_dram_window = make_tile_window( - v_dram_block_window_tmp.get_bottom_tensor_view(), - v_dram_block_window_tmp.get_window_lengths(), - v_dram_block_window_tmp.get_window_origin(), - Policy::template MakeVInRegDramTileDistribution()); - - auto v = load_tile(v_dram_window); // persistent V register tile - - using SPTBlockTileType = decltype(gemm_0.MakeCBlockTile()); - using SPGradTBlockTileType = decltype(gemm_2.MakeCBlockTile()); - using QGradBlockTileType = decltype(gemm_4.MakeCBlockTile()); - - // init VGrad & KGrad - auto dv_acc = decltype(gemm_1.MakeCBlockTile()){}; - auto dk_acc = decltype(gemm_3.MakeCBlockTile()){}; - - clear_tile(dv_acc); - clear_tile(dk_acc); - - auto k_dram_window = make_tile_window( - k_dram_block_window_tmp.get_bottom_tensor_view(), - k_dram_block_window_tmp.get_window_lengths(), - k_dram_block_window_tmp.get_window_origin(), - Policy::template MakeKDramTileDistribution()); // K DRAM tile window for - // load - - __builtin_amdgcn_sched_barrier(0); - const auto k_origin = k_dram_window.get_window_origin(); - const auto [seqlen_q_start, seqlen_q_end] = - mask.GetTileRangeAlongY(k_origin.at(number<0>{}), number{}, number{}); - - const auto num_total_loop = integer_divide_ceil(seqlen_q_end - seqlen_q_start, kM0); - - // check early exit if masked and no work to do. - if constexpr(FmhaMask::IsMasking) - { - if(num_total_loop <= 0) - { - // Note: here dk_acc&dv_acc are all cleard, return it - // Note: v loaded but no fence, ignore it. - return ck_tile::make_tuple(dk_acc, dv_acc); - } - } - - auto k_block_tile = load_tile(k_dram_window); - - store_tile(k_lds_window, k_block_tile); // // persistent K in LDS - - auto q_dram_block_window = - make_tile_window(q_dram_block_window_tmp.get_bottom_tensor_view(), - q_dram_block_window_tmp.get_window_lengths(), - {seqlen_q_start, 0}); - - auto qt_dram_block_window = - make_tile_window(qt_dram_block_window_tmp.get_bottom_tensor_view(), - qt_dram_block_window_tmp.get_window_lengths(), - {0, seqlen_q_start}); - - auto do_dram_block_window = - make_tile_window(do_dram_block_window_tmp.get_bottom_tensor_view(), - do_dram_block_window_tmp.get_window_lengths(), - {seqlen_q_start, 0}); - - auto dot_dram_block_window = - make_tile_window(dot_dram_block_window_tmp.get_bottom_tensor_view(), - dot_dram_block_window_tmp.get_window_lengths(), - {0, seqlen_q_start}); - - auto dq_dram_block_window = - make_tile_window(dq_dram_block_window_tmp.get_bottom_tensor_view(), - dq_dram_block_window_tmp.get_window_lengths(), - {seqlen_q_start, 0}); - - auto lse_dram_block_window = - make_tile_window(lse_dram_block_window_tmp.get_bottom_tensor_view(), - lse_dram_block_window_tmp.get_window_lengths(), - {seqlen_q_start}); - - auto d_dram_block_window = - make_tile_window(d_dram_block_window_tmp.get_bottom_tensor_view(), - d_dram_block_window_tmp.get_window_lengths(), - {seqlen_q_start}); - - const auto bias_origin = bias_dram_block_window_tmp.get_window_origin(); - auto bias_dram_block_window = - make_tile_window(bias_dram_block_window_tmp.get_bottom_tensor_view(), - bias_dram_block_window_tmp.get_window_lengths(), - {seqlen_q_start, bias_origin.at(number<1>{})}); // M/N - - const auto dbias_origin = dbias_dram_block_window_tmp.get_window_origin(); - auto dbias_dram_block_window = - make_tile_window(dbias_dram_block_window_tmp.get_bottom_tensor_view(), - dbias_dram_block_window_tmp.get_window_lengths(), - {seqlen_q_start, dbias_origin.at(number<1>{})}); // M/N - - auto qt_dram_window = - make_tile_window(qt_dram_block_window.get_bottom_tensor_view(), - qt_dram_block_window.get_window_lengths(), - qt_dram_block_window.get_window_origin(), - Policy::template MakeQTDramTileDistribution()); - - auto dot_dram_window = - make_tile_window(dot_dram_block_window.get_bottom_tensor_view(), - dot_dram_block_window.get_window_lengths(), - dot_dram_block_window.get_window_origin(), - Policy::template MakeOGradTDramTileDistribution()); - - auto lse_dram_window = make_tile_window( - lse_dram_block_window.get_bottom_tensor_view(), - lse_dram_block_window.get_window_lengths(), - lse_dram_block_window.get_window_origin(), - Policy::template MakeLSEDDramTileDistribution()); - - auto d_dram_window = make_tile_window( - d_dram_block_window.get_bottom_tensor_view(), - d_dram_block_window.get_window_lengths(), - d_dram_block_window.get_window_origin(), - Policy::template MakeLSEDDramTileDistribution()); - - auto bias_dram_window = - make_tile_window(bias_dram_block_window.get_bottom_tensor_view(), - bias_dram_block_window.get_window_lengths(), - bias_dram_block_window.get_window_origin(), - Policy::template MakeBiasTileDistribution()); - - auto biast_lds_window = - make_tile_window(biast_lds_shuffle_window.get_bottom_tensor_view(), - biast_lds_shuffle_window.get_window_lengths(), - biast_lds_shuffle_window.get_window_origin(), - Policy::template MakeBiasTTileDistribution()); - - auto randval_dram_window = dropout.MakeRandvalDramWindow( - randval_dram_block_window_tmp, seqlen_q_start); - - index_t i_total_loops = 0; - constexpr index_t k0_loops = kQKHeaddim / kK0; - constexpr index_t k1_loops = kM0 / kK1; - constexpr index_t k2_loops = kVHeaddim / kK2; - constexpr index_t k3_loops = kM0 / kK3; - constexpr index_t k4_loops = kN0 / kK4; - do - { - auto q_dram_window = make_tile_window( - q_dram_block_window.get_bottom_tensor_view(), - q_dram_block_window.get_window_lengths(), - q_dram_block_window.get_window_origin(), - Policy::template MakeQDramTileDistribution()); // Q DRAM tile window for - // load - - auto do_dram_window = make_tile_window( - do_dram_block_window.get_bottom_tensor_view(), - do_dram_block_window.get_window_lengths(), - do_dram_block_window.get_window_origin(), - Policy::template MakeOGradDramTileDistribution()); // OGrad DRAM tile - // window for load - - // STAGE 1, Q@K Gemm0 - auto st_acc = SPTBlockTileType{}; - - auto q_block_tile = load_tile(q_dram_window); - { - move_tile_window(q_dram_window, {0, kK0}); - - clear_tile(st_acc); // Initialize S^T - - store_tile(q_lds_window, q_block_tile); // LDS write 0 - q_block_tile = load_tile(q_dram_window); // global read 1 - } - - 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(st_acc, - q_lds_window, - get_slice_tile(k_lds_window, - sequence<0, i_k0 * kK0>{}, - sequence{})); - block_sync_lds(); - move_tile_window(q_dram_window, {0, kK0}); - - store_tile(q_lds_window, - q_block_tile); // LDS write i + 1 - q_block_tile = load_tile(q_dram_window); // global read i + 2 - }); - } - - const auto dot_prefetch = load_tile(dot_dram_window); // prefetch load OGrad^T tile - { // tail - block_sync_lds(); - gemm_0(st_acc, - q_lds_window, - get_slice_tile(k_lds_window, - sequence<0, (k0_loops - 2) * kK0>{}, - sequence{})); - block_sync_lds(); - - store_tile(q_lds_window, q_block_tile); - block_sync_lds(); - - gemm_0(st_acc, - q_lds_window, - get_slice_tile(k_lds_window, - sequence<0, (k0_loops - 1) * kK0>{}, - sequence{})); - } - - // STAGE 2, Scale, Add bias, Mask, Softmax, Dropout - if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS) - { - block_sync_lds(); - auto bias_shuffle_tmp = make_static_distributed_tensor( - Policy::template MakeShuffledBiasTileDistribution()); - shuffle_tile(bias_shuffle_tmp, bias_tile); - store_tile(biast_lds_shuffle_window, bias_shuffle_tmp); - block_sync_lds(); - auto biast_tile = load_tile(biast_lds_window); - tile_elementwise_inout( - [&](auto& x, const auto& y) { -#if !CK_TILE_FMHA_FWD_FAST_EXP2 - x = raw_scale * x + type_convert(y); -#else - x = scale * x + log2e_v * type_convert(y); -#endif - }, - st_acc, - biast_tile); - move_tile_window(bias_dram_window, {kM0, 0}); - } - else if constexpr(BiasEnum == BlockAttentionBiasEnum::ALIBI) - { - const auto q_origin = q_dram_block_window.get_window_origin(); - constexpr auto st_spans = decltype(st_acc)::get_distributed_spans(); - sweep_tile_span(st_spans[number<0>{}], [&](auto idx0) { - sweep_tile_span(st_spans[number<1>{}], [&](auto idx1) { - const auto tile_idx = get_x_indices_from_distributed_indices( - st_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); - -#if !CK_TILE_FMHA_FWD_FAST_EXP2 - st_acc(i_j_idx) *= raw_scale; -#else - st_acc(i_j_idx) *= scale; -#endif - position_encoding.update(st_acc(i_j_idx), row, col); - }); - }); - } - else - { -#if !CK_TILE_FMHA_FWD_FAST_EXP2 - tile_elementwise_inout([&raw_scale](auto& x) { x = x * raw_scale; }, st_acc); -#endif - } - - if constexpr(kPadSeqLenK || FmhaMask::IsMasking) - { - const auto q_origin = q_dram_block_window.get_window_origin(); - bool need_perpixel_check = mask.IsEdgeTile(q_origin.at(number<0>{}), - k_origin.at(number<0>{}), - number{}, - number{}); - if(need_perpixel_check) - { - set_tile_if(st_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); - }); - } - } - - const auto lse = load_tile(lse_dram_window); - - static const auto get_validated_lse = [](LSEDataType raw_lse) { - if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS || - FmhaMask::IsMasking) - { - return raw_lse == -numeric::infinity() - ? type_convert(0.f) - : raw_lse; - } - else - { - return raw_lse; - } - }; - - auto pt = SPTBlockTileType{}; - constexpr auto pt_spans = decltype(pt)::get_distributed_spans(); - sweep_tile_span(pt_spans[number<0>{}], [&](auto idx0) { - constexpr auto i_idx = make_tuple(idx0); -#if CK_TILE_FMHA_FWD_FAST_EXP2 - auto row_lse = log2e_v * get_validated_lse(lse[i_idx]); -#endif - sweep_tile_span(pt_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) - { - pt(i_j_idx) = exp2(st_acc[i_j_idx] - row_lse); - } - else - { - pt(i_j_idx) = exp2(scale * st_acc[i_j_idx] - row_lse); - } -#else - pt(i_j_idx) = exp(st_acc[i_j_idx] - get_validated_lse(lse[i_idx])); -#endif - }); - }); - - auto dot_shuffle_tmp = make_static_distributed_tensor( - Policy::template MakeShuffledOGradTRegBlockDescriptor()); - block_sync_lds(); - { - shuffle_tile(dot_shuffle_tmp, dot_prefetch); - store_tile(dot_lds_window, - dot_shuffle_tmp); // store the prefetch - } - move_tile_window(dot_dram_window, {0, kK1}); - - if constexpr(kHasDropout) - { - dropout.Run( - seqlen_q_start + i_total_loops * kM0, pt, randval_dram_window); - } - - // STAGE 3, P^T@OGrad^T Gemm1 - const auto pt_gemm = [&]() { - if constexpr(kHasDropout) - { - return tile_elementwise_in( - [](const auto& x) { return type_convert(x > 0.f ? x : 0.f); }, - pt); - } - else - { - return cast_tile(pt); - } - }(); - - if constexpr(k1_loops > 1) - { - static_for<0, k1_loops - 1, 1>{}([&](auto i_k1) { - const auto dot = load_tile(dot_dram_window); // load next OGrad^T - block_sync_lds(); - gemm_1(dv_acc, - get_slice_tile(pt_gemm, - sequence{}, - sequence<(i_k1 + 1) * kK1, kN0>{}), - dot_lds_window); - block_sync_lds(); - shuffle_tile(dot_shuffle_tmp, dot); - store_tile(dot_lds_window, - dot_shuffle_tmp); // store the prefetch - - move_tile_window(dot_dram_window, {0, kK1}); - }); - } - auto do_block_tile = load_tile(do_dram_window); // prefetch load OGrad tile - // tail - { - block_sync_lds(); - gemm_1(dv_acc, - get_slice_tile( - pt_gemm, sequence<(k1_loops - 1) * kK1, 0>{}, sequence{}), - dot_lds_window); - block_sync_lds(); - } - - // STAGE 4, OGrad@V Gemm2 - auto dpt_acc = SPGradTBlockTileType{}; - - { - move_tile_window(do_dram_window, {0, kK2}); - - clear_tile(dpt_acc); // Initialize PGrad^T - - store_tile(do_lds_window, do_block_tile); // LDS write 0 - do_block_tile = load_tile(do_dram_window); // global read 1 - } - - if constexpr(k2_loops > 2) - { - static_for<0, k2_loops - 2, 1>{}([&](auto i_k2) { - block_sync_lds(); - gemm_2(dpt_acc, - do_lds_window, - get_slice_tile( - v, sequence<0, i_k2 * kK2>{}, sequence{})); - block_sync_lds(); - move_tile_window(do_dram_window, {0, kK2}); - - store_tile(do_lds_window, - do_block_tile); // LDS write i + 1 - do_block_tile = load_tile(do_dram_window); // global read i + 2 - }); - } - - const auto qt_prefetch = load_tile(qt_dram_window); // prefetch load Q^T tile - { // tail - block_sync_lds(); - gemm_2(dpt_acc, - do_lds_window, - get_slice_tile(v, - sequence<0, (k2_loops - 2) * kK2>{}, - sequence{})); - block_sync_lds(); - - store_tile(do_lds_window, do_block_tile); - block_sync_lds(); - - gemm_2(dpt_acc, - do_lds_window, - get_slice_tile(v, - sequence<0, (k2_loops - 1) * kK2>{}, - sequence{})); - } - - // STAGE 5, P^T(PGrad^T - D) - const auto d = load_tile(d_dram_window); - - auto dst = SPGradTBlockTileType{}; - constexpr auto dst_spans = decltype(dst)::get_distributed_spans(); - sweep_tile_span(dst_spans[number<0>{}], [&](auto idx0) { - constexpr auto i_idx = make_tuple(idx0); - sweep_tile_span(dst_spans[number<1>{}], [&](auto idx1) { - constexpr auto i_j_idx = make_tuple(idx0, idx1); - bool undrop_flag = pt[i_j_idx] >= 0; - dst(i_j_idx) = - pt[i_j_idx] * - (!kHasDropout || undrop_flag ? (dpt_acc[i_j_idx] - d[i_idx]) : d[i_idx]); - }); - }); - - if constexpr(kHasBiasGrad) - { - const auto dbiast = [&]() { - if constexpr(kHasDropout) - { - return tile_elementwise_in( - [&rp_undrop](const auto& x) { - return type_convert(x * rp_undrop); - }, - dst); - } - else - { - return cast_tile(dst); - } - }(); - store_tile(biast_lds_shuffle_window, dbiast); - block_sync_lds(); - auto dbiast_tile = load_tile(dbiast_lds_shuffle_window); - auto dbiast_shuffle_tmp = make_static_distributed_tensor( - Policy::template MakeBiasTileDistribution()); - shuffle_tile(dbiast_shuffle_tmp, dbiast_tile); - store_tile(dbias_dram_block_window, dbiast_shuffle_tmp); - move_tile_window(dbias_dram_block_window, {kM0, 0}); - } - - // STAGE 6, SGrad^T@Q^T Gemm3 - auto qt_shuffle_tmp = make_static_distributed_tensor( - Policy::template MakeShuffledQTRegBlockDescriptor()); - block_sync_lds(); - { - shuffle_tile(qt_shuffle_tmp, qt_prefetch); - store_tile(qt_lds_window, - qt_shuffle_tmp); // store the prefetch - } - move_tile_window(qt_dram_window, {0, kK3}); - - const auto dst_gemm = cast_tile(dst); - - if constexpr(k3_loops > 1) - { - static_for<0, k3_loops - 1, 1>{}([&](auto i_k3) { - const auto qt = load_tile(qt_dram_window); // load next Q^T - block_sync_lds(); - gemm_3(dk_acc, - get_slice_tile(dst_gemm, - sequence{}, - sequence<(i_k3 + 1) * kK3, kN0>{}), - qt_lds_window); - block_sync_lds(); - shuffle_tile(qt_shuffle_tmp, qt); - store_tile(qt_lds_window, - qt_shuffle_tmp); // store the prefetch - - move_tile_window(qt_dram_window, {0, kK3}); - }); - } - // tail - { - block_sync_lds(); - gemm_3(dk_acc, - get_slice_tile( - dst_gemm, sequence<(k3_loops - 1) * kK3, 0>{}, sequence{}), - qt_lds_window); - block_sync_lds(); - } - - // STAGE 7, SGrad@K^T Gemm4 - store_tile(ds_lds_window, dst_gemm); - - auto dq_acc = QGradBlockTileType{}; - clear_tile(dq_acc); // Initialize QGrad - - block_sync_lds(); - - static_for<0, k4_loops, 1>{}([&](auto i_k4) { - gemm_4(dq_acc, - get_slice_tile(ds_lds_window, - sequence<0, i_k4 * kK4>{}, - sequence{}), - get_slice_tile(kt_lds_window, - sequence<0, i_k4 * kK4>{}, - sequence{})); - }); - - // QGrad Scale - if constexpr(kHasDropout) - { - tile_elementwise_inout([&scale_rp_undrop](auto& x) { x = x * scale_rp_undrop; }, - dq_acc); - } - else - { - tile_elementwise_inout([&raw_scale](auto& x) { x = x * raw_scale; }, dq_acc); - } - const auto dq = cast_tile(dq_acc); - update_tile(dq_dram_block_window, dq); - - // move tile windows - move_tile_window(q_dram_block_window, {kM0, 0}); - move_tile_window(dq_dram_block_window, {kM0, 0}); - move_tile_window(do_dram_block_window, {kM0, 0}); - move_tile_window(lse_dram_window, {kM0}); - move_tile_window(d_dram_window, {kM0}); - } while(++i_total_loops < num_total_loop); - - // KGrad Scale - if constexpr(kHasDropout) - { - tile_elementwise_inout([&scale_rp_undrop](auto& x) { x = x * scale_rp_undrop; }, - dk_acc); - } - else - { - tile_elementwise_inout([&raw_scale](auto& x) { x = x * raw_scale; }, dk_acc); - } - // VGrad Scale - if constexpr(kHasDropout) - { - tile_elementwise_inout([&rp_undrop](auto& x) { x = x * rp_undrop; }, dv_acc); - } - - return ck_tile::make_tuple(dk_acc, dv_acc); - } -}; - -} // namespace ck_tile diff --git a/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_ks_vr_default_policy.hpp b/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_ks_vr_default_policy.hpp deleted file mode 100644 index cc4e6304d0..0000000000 --- a/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_ks_vr_default_policy.hpp +++ /dev/null @@ -1,20 +0,0 @@ -// SPDX-License-Identifier: MIT -// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved. - -#pragma once - -#include "ck_tile/ops/fmha/pipeline/block_fmha_bwd_pipeline_default_policy.hpp" - -namespace ck_tile { - -// This pipeline is v located in regs, k located in lds. -using BlockFmhaBwdDQDKDVPipelineKSVRDefaultPolicy = - BlockFmhaBwdPipelineDefaultPolicy; - -} // namespace ck_tile diff --git a/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_qs_ks_vr_dos.hpp b/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_qs_ks_vr_dos.hpp deleted file mode 100644 index 97487bb71e..0000000000 --- a/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_qs_ks_vr_dos.hpp +++ /dev/null @@ -1,692 +0,0 @@ -// 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/block/block_dropout.hpp" -#include "ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_qs_ks_vr_dos_default_policy.hpp" -#include "ck_tile/ops/reduce/block/block_reduce.hpp" - -namespace ck_tile { - -template -struct BlockFmhaBwdDQDKDVPipelineQSKSVROGradS -{ - using QDataType = remove_cvref_t; - using KDataType = remove_cvref_t; - using VDataType = remove_cvref_t; - using GemmDataType = remove_cvref_t; - using BiasDataType = remove_cvref_t; - using LSEDataType = remove_cvref_t; - using AccDataType = remove_cvref_t; - using DDataType = remove_cvref_t; - using RandValOutputDataType = remove_cvref_t; - using ODataType = remove_cvref_t; - using OGradDataType = remove_cvref_t; - using QGradDataType = remove_cvref_t; - using KGradDataType = remove_cvref_t; - using VGradDataType = remove_cvref_t; - using BiasGradDataType = remove_cvref_t; - using FmhaMask = remove_cvref_t; - - using BlockFmhaShape = remove_cvref_t; - - static constexpr index_t kBlockPerCu = Problem::kBlockPerCu; - 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 kK1 = BlockFmhaShape::kK1; - static constexpr index_t kK2 = BlockFmhaShape::kK2; - static constexpr index_t kK3 = BlockFmhaShape::kK3; - static constexpr index_t kK4 = BlockFmhaShape::kK4; - static constexpr index_t kQKHeaddim = BlockFmhaShape::kQKHeaddim; - static constexpr index_t kVHeaddim = BlockFmhaShape::kVHeaddim; - - static constexpr bool kQLoadOnce = true; - static constexpr bool kQTLoadOnce = false; - static constexpr bool kKLoadOnce = true; - static constexpr bool kKTLoadOnce = false; - static constexpr bool kVLoadOnce = true; - static constexpr bool kOGradLoadOnce = true; - static constexpr bool kOGradTLoadOnce = false; - - 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 kHasBiasGrad = Problem::kHasBiasGrad; - static constexpr bool kHasDropout = Problem::kHasDropout; - - // 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 = - kPadHeadDimV ? 1 : Policy::template GetAlignmentV(); - static constexpr index_t kAlignmentO = - kPadHeadDimV ? 1 : Policy::template GetAlignmentO(); - static constexpr index_t kAlignmentOGrad = - kPadHeadDimV ? 1 : Policy::template GetAlignmentOGrad(); - static constexpr index_t kAlignmentQGrad = - kPadHeadDimQ ? 2 : Policy::template GetAlignmentQGrad(); - static constexpr index_t kAlignmentKGrad = - kPadHeadDimQ ? 1 : Policy::template GetAlignmentKGrad(); - static constexpr index_t kAlignmentVGrad = - kPadHeadDimV ? 1 : Policy::template GetAlignmentVGrad(); - static constexpr index_t kAlignmentBias = - kPadSeqLenK ? 1 : Policy::template GetTransposedAlignmentBias(); - - static constexpr const char* name = "qs_ks_vr_dos"; - - 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, - const QTDramBlockWindowTmp& /*qt_dram_block_window_tmp*/, - const KDramBlockWindowTmp& k_dram_block_window_tmp, - const KTDramBlockWindowTmp& /*kt_dram_block_window_tmp*/, - const VDramBlockWindowTmp& v_dram_block_window_tmp, - const BiasDramBlockWindowTmp& bias_dram_block_window_tmp, - const RandValDramBlockWindowTmp& randval_dram_block_window_tmp, - const OGradDramBlockWindowTmp& do_dram_block_window_tmp, - const OGradTDramBlockWindowTmp& /*dot_dram_block_window_tmp*/, - const LSEDramBlockWindowTmp& lse_dram_block_window_tmp, - const DDramBlockWindowTmp& d_dram_block_window_tmp, - const QGradDramBlockWindowTmp& dq_dram_block_window_tmp, - const BiasGradDramBlockWindowTmp& dbias_dram_block_window_tmp, - FmhaMask mask, - PositionEncoding position_encoding, - float raw_scale, -#if CK_TILE_FMHA_FWD_FAST_EXP2 - float scale, -#endif - float rp_undrop, - float scale_rp_undrop, - void* smem_ptr, - BlockDropout& dropout) const - { - static_assert( - std::is_same_v> && - std::is_same_v> && - std::is_same_v> && - std::is_same_v> && - std::is_same_v> && - std::is_same_v> && - std::is_same_v>, - "wrong!"); - - static_assert(kM0 == QDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] && - kN0 == KDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] && - kN0 == VDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] && - kM0 == BiasDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] && - kN0 == BiasDramBlockWindowTmp{}.get_window_lengths()[number<1>{}] && - kM0 == OGradDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] && - kM0 == LSEDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] && - kM0 == DDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] && - kM0 == QGradDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] && - kM0 == BiasGradDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] && - kN0 == BiasGradDramBlockWindowTmp{}.get_window_lengths()[number<1>{}], - "wrong!"); - - // Q tile in LDS - QDataType* q_lds_ptr = static_cast(static_cast( - static_cast(smem_ptr) + Policy::template GetSmemSizeK())); - auto q_lds = make_tensor_view( - q_lds_ptr, Policy::template MakeQLdsBlockDescriptor()); - auto q_lds_window = - make_tile_window(q_lds, make_tuple(number{}, number{}), {0, 0}); - - // QT tile in LDS - auto qt_lds = make_tensor_view( - q_lds_ptr, Policy::template MakeQLdsBlockDescriptorAsQT()); - auto qt_lds_window = - make_tile_window(qt_lds, make_tuple(number{}, number{}), {0, 0}); - - // K tile in LDS - auto k_lds = make_tensor_view( - reinterpret_cast(smem_ptr), - Policy::template MakeKLdsBlockDescriptor()); - auto k_lds_window = - make_tile_window(k_lds, make_tuple(number{}, number{}), {0, 0}); - - // KT tile in LDS - auto kt_lds = make_tensor_view( - reinterpret_cast(smem_ptr), - Policy::template MakeKLdsBlockDescriptorAsKT()); - auto kt_lds_window = - make_tile_window(kt_lds, make_tuple(number{}, number{}), {0, 0}); - - // OGrad tile in LDS - OGradDataType* do_lds_ptr = static_cast(static_cast( - static_cast(smem_ptr) + Policy::template GetSmemSizeK() + - Policy::template GetSmemSizeQ())); - auto do_lds = make_tensor_view( - do_lds_ptr, Policy::template MakeOGradLdsBlockDescriptor()); - auto do_lds_window = - make_tile_window(do_lds, make_tuple(number{}, number{}), {0, 0}); - - // OGradT tile in LDS - auto dot_lds = make_tensor_view( - do_lds_ptr, Policy::template MakeOGradLdsBlockDescriptorAsOGradT()); - auto dot_lds_window = - make_tile_window(dot_lds, make_tuple(number{}, number{}), {0, 0}); - - // SGrad tile in LDS - GemmDataType* ds_lds_ptr = static_cast(static_cast( - static_cast(smem_ptr) + Policy::template GetSmemSizeK() + - Policy::template GetSmemSizeQ() + - Policy::template GetSmemSizeOGrad())); - auto ds_lds = make_tensor_view( - ds_lds_ptr, Policy::template MakeSGradLdsBlockDescriptor()); - auto ds_lds_window = - make_tile_window(ds_lds, make_tuple(number{}, number{}), {0, 0}); - - // BiasT/BiasGradT tile in LDS, use the same size and layout - BiasDataType* biast_lds_ptr = static_cast(static_cast( - static_cast(smem_ptr) + Policy::template GetSmemSizeK() + - Policy::template GetSmemSizeQ() + - Policy::template GetSmemSizeOGrad())); - auto biast_lds = make_tensor_view( - biast_lds_ptr, Policy::template MakeBiasTLdsBlockDescriptor()); - auto biast_lds_shuffle_window = - make_tile_window(biast_lds, make_tuple(number{}, number{}), {0, 0}); - auto dbiast_lds_shuffle_window = - make_tile_window(biast_lds, - make_tuple(number{}, number{}), - {0, 0}, - Policy::template MakeShuffledBiasTileDistribution()); - - static_assert(std::is_same_v, - "BiasDataType and BiasGradDataType should be the same!"); - - // Block GEMM - constexpr auto gemm_0 = Policy::template GetQKBlockGemm(); - constexpr auto gemm_1 = Policy::template GetPTOGradTBlockGemm(); - constexpr auto gemm_2 = Policy::template GetOGradVBlockGemm(); - constexpr auto gemm_3 = Policy::template GetSGradTQTBlockGemm(); - constexpr auto gemm_4 = Policy::template GetSGradKTBlockGemm(); - - auto v_dram_window = make_tile_window( - v_dram_block_window_tmp.get_bottom_tensor_view(), - v_dram_block_window_tmp.get_window_lengths(), - v_dram_block_window_tmp.get_window_origin(), - Policy::template MakeVInRegDramTileDistribution()); - - auto v = load_tile(v_dram_window); // persistent V register tile - - using SPTBlockTileType = decltype(gemm_0.MakeCBlockTile()); - using SPGradTBlockTileType = decltype(gemm_2.MakeCBlockTile()); - using QGradBlockTileType = decltype(gemm_4.MakeCBlockTile()); - - // init VGrad & KGrad - auto dv_acc = decltype(gemm_1.MakeCBlockTile()){}; - auto dk_acc = decltype(gemm_3.MakeCBlockTile()){}; - - clear_tile(dv_acc); - clear_tile(dk_acc); - - auto k_dram_window = make_tile_window( - k_dram_block_window_tmp.get_bottom_tensor_view(), - k_dram_block_window_tmp.get_window_lengths(), - k_dram_block_window_tmp.get_window_origin(), - Policy::template MakeKDramTileDistribution()); // K DRAM tile window for - // load - - __builtin_amdgcn_sched_barrier(0); - const auto k_origin = k_dram_window.get_window_origin(); - const auto [seqlen_q_start, seqlen_q_end] = - mask.GetTileRangeAlongY(k_origin.at(number<0>{}), number{}, number{}); - - const auto num_total_loop = integer_divide_ceil(seqlen_q_end - seqlen_q_start, kM0); - - // check early exit if masked and no work to do. - if constexpr(FmhaMask::IsMasking) - { - if(num_total_loop <= 0) - { - // Note: here dk_acc&dv_acc are all cleard, return it - // Note: v loaded but no fence, ignore it. - return ck_tile::make_tuple(dk_acc, dv_acc); - } - } - - auto k_block_tile = load_tile(k_dram_window); - - store_tile(k_lds_window, k_block_tile); // // persistent K in LDS - - auto q_dram_block_window = - make_tile_window(q_dram_block_window_tmp.get_bottom_tensor_view(), - q_dram_block_window_tmp.get_window_lengths(), - {seqlen_q_start, 0}); - - auto do_dram_block_window = - make_tile_window(do_dram_block_window_tmp.get_bottom_tensor_view(), - do_dram_block_window_tmp.get_window_lengths(), - {seqlen_q_start, 0}); - - auto dq_dram_block_window = - make_tile_window(dq_dram_block_window_tmp.get_bottom_tensor_view(), - dq_dram_block_window_tmp.get_window_lengths(), - {seqlen_q_start, 0}); - - auto lse_dram_block_window = - make_tile_window(lse_dram_block_window_tmp.get_bottom_tensor_view(), - lse_dram_block_window_tmp.get_window_lengths(), - {seqlen_q_start}); - - auto d_dram_block_window = - make_tile_window(d_dram_block_window_tmp.get_bottom_tensor_view(), - d_dram_block_window_tmp.get_window_lengths(), - {seqlen_q_start}); - - const auto bias_origin = bias_dram_block_window_tmp.get_window_origin(); - auto bias_dram_block_window = - make_tile_window(bias_dram_block_window_tmp.get_bottom_tensor_view(), - bias_dram_block_window_tmp.get_window_lengths(), - {seqlen_q_start, bias_origin.at(number<1>{})}); // M/N - - const auto dbias_origin = dbias_dram_block_window_tmp.get_window_origin(); - auto dbias_dram_block_window = - make_tile_window(dbias_dram_block_window_tmp.get_bottom_tensor_view(), - dbias_dram_block_window_tmp.get_window_lengths(), - {seqlen_q_start, dbias_origin.at(number<1>{})}); // M/N - - auto lse_dram_window = make_tile_window( - lse_dram_block_window.get_bottom_tensor_view(), - lse_dram_block_window.get_window_lengths(), - lse_dram_block_window.get_window_origin(), - Policy::template MakeLSEDDramTileDistribution()); - - auto d_dram_window = make_tile_window( - d_dram_block_window.get_bottom_tensor_view(), - d_dram_block_window.get_window_lengths(), - d_dram_block_window.get_window_origin(), - Policy::template MakeLSEDDramTileDistribution()); - - auto bias_dram_window = - make_tile_window(bias_dram_block_window.get_bottom_tensor_view(), - bias_dram_block_window.get_window_lengths(), - bias_dram_block_window.get_window_origin(), - Policy::template MakeBiasTileDistribution()); - - auto biast_lds_window = - make_tile_window(biast_lds_shuffle_window.get_bottom_tensor_view(), - biast_lds_shuffle_window.get_window_lengths(), - biast_lds_shuffle_window.get_window_origin(), - Policy::template MakeBiasTTileDistribution()); - - auto randval_dram_window = dropout.MakeRandvalDramWindow( - randval_dram_block_window_tmp, seqlen_q_start); - - index_t i_total_loops = 0; - constexpr index_t k0_loops = kQKHeaddim / kK0; - constexpr index_t k1_loops = kM0 / kK1; - constexpr index_t k2_loops = kVHeaddim / kK2; - constexpr index_t k3_loops = kM0 / kK3; - constexpr index_t k4_loops = kN0 / kK4; - do - { - auto q_dram_window = make_tile_window( - q_dram_block_window.get_bottom_tensor_view(), - q_dram_block_window.get_window_lengths(), - q_dram_block_window.get_window_origin(), - Policy::template MakeQDramTileDistribution()); // Q DRAM tile window for - // load - - auto do_dram_window = make_tile_window( - do_dram_block_window.get_bottom_tensor_view(), - do_dram_block_window.get_window_lengths(), - do_dram_block_window.get_window_origin(), - Policy::template MakeOGradDramTileDistribution()); // OGrad DRAM tile - // window for load - - // STAGE 1, Q@K Gemm0 - auto st_acc = SPTBlockTileType{}; - - auto q_block_tile = load_tile(q_dram_window); - clear_tile(st_acc); // Initialize S^T - store_tile(q_lds_window, q_block_tile); // LDS write - - 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 > 1) - { - static_for<0, k0_loops - 1, 1>{}([&](auto i_k0) { - block_sync_lds(); - gemm_0(st_acc, - get_slice_tile(q_lds_window, - sequence<0, i_k0 * kK0>{}, - sequence{}), - get_slice_tile(k_lds_window, - sequence<0, i_k0 * kK0>{}, - sequence{})); - block_sync_lds(); - }); - } - - auto do_block_tile = load_tile(do_dram_window); // prefetch load OGrad tile - { // tail - block_sync_lds(); - gemm_0(st_acc, - get_slice_tile(q_lds_window, - sequence<0, (k0_loops - 1) * kK0>{}, - sequence{}), - get_slice_tile(k_lds_window, - sequence<0, (k0_loops - 1) * kK0>{}, - sequence{})); - block_sync_lds(); - } - - // STAGE 2, Scale, Add bias, Mask, Softmax, Dropout - if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS) - { - block_sync_lds(); - auto bias_shuffle_tmp = make_static_distributed_tensor( - Policy::template MakeShuffledBiasTileDistribution()); - shuffle_tile(bias_shuffle_tmp, bias_tile); - store_tile(biast_lds_shuffle_window, bias_shuffle_tmp); - block_sync_lds(); - auto biast_tile = load_tile(biast_lds_window); - tile_elementwise_inout( - [&](auto& x, const auto& y) { -#if !CK_TILE_FMHA_FWD_FAST_EXP2 - x = raw_scale * x + type_convert(y); -#else - x = scale * x + log2e_v * type_convert(y); -#endif - }, - st_acc, - biast_tile); - move_tile_window(bias_dram_window, {kM0, 0}); - } - else if constexpr(BiasEnum == BlockAttentionBiasEnum::ALIBI) - { - const auto q_origin = q_dram_block_window.get_window_origin(); - constexpr auto st_spans = decltype(st_acc)::get_distributed_spans(); - sweep_tile_span(st_spans[number<0>{}], [&](auto idx0) { - sweep_tile_span(st_spans[number<1>{}], [&](auto idx1) { - const auto tile_idx = get_x_indices_from_distributed_indices( - st_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); - -#if !CK_TILE_FMHA_FWD_FAST_EXP2 - st_acc(i_j_idx) *= raw_scale; -#else - st_acc(i_j_idx) *= scale; -#endif - position_encoding.update(st_acc(i_j_idx), row, col); - }); - }); - } - else - { -#if !CK_TILE_FMHA_FWD_FAST_EXP2 - tile_elementwise_inout([&raw_scale](auto& x) { x = x * raw_scale; }, st_acc); -#endif - } - - if constexpr(kPadSeqLenK || FmhaMask::IsMasking) - { - const auto q_origin = q_dram_block_window.get_window_origin(); - bool need_perpixel_check = mask.IsEdgeTile(q_origin.at(number<0>{}), - k_origin.at(number<0>{}), - number{}, - number{}); - if(need_perpixel_check) - { - set_tile_if(st_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); - }); - } - } - - const auto lse = load_tile(lse_dram_window); - - static const auto get_validated_lse = [](LSEDataType raw_lse) { - if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS || - FmhaMask::IsMasking) - { - return raw_lse == -numeric::infinity() - ? type_convert(0.f) - : raw_lse; - } - else - { - return raw_lse; - } - }; - - auto pt = SPTBlockTileType{}; - constexpr auto pt_spans = decltype(pt)::get_distributed_spans(); - sweep_tile_span(pt_spans[number<0>{}], [&](auto idx0) { - constexpr auto i_idx = make_tuple(idx0); -#if CK_TILE_FMHA_FWD_FAST_EXP2 - auto row_lse = log2e_v * get_validated_lse(lse[i_idx]); -#endif - sweep_tile_span(pt_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) - { - pt(i_j_idx) = exp2(st_acc[i_j_idx] - row_lse); - } - else - { - pt(i_j_idx) = exp2(scale * st_acc[i_j_idx] - row_lse); - } -#else - pt(i_j_idx) = exp(st_acc[i_j_idx] - get_validated_lse(lse[i_idx])); -#endif - }); - }); - - if constexpr(kHasDropout) - { - dropout.Run( - seqlen_q_start + i_total_loops * kM0, pt, randval_dram_window); - } - - // STAGE 3, P^T@OGrad^T Gemm1 - block_sync_lds(); - store_tile(do_lds_window, do_block_tile); // store the prefetch - - const auto pt_gemm = [&]() { - if constexpr(kHasDropout) - { - return tile_elementwise_in( - [](const auto& x) { return type_convert(x > 0.f ? x : 0.f); }, - pt); - } - else - { - return cast_tile(pt); - } - }(); - - static_for<0, k1_loops, 1>{}([&](auto i_k1) { - block_sync_lds(); - gemm_1(dv_acc, - get_slice_tile( - pt_gemm, sequence{}, sequence<(i_k1 + 1) * kK1, kN0>{}), - get_slice_tile(dot_lds_window, - sequence<0, i_k1 * kK1>{}, - sequence{})); - block_sync_lds(); - }); - - // STAGE 4, OGrad@V Gemm2 - auto dpt_acc = SPGradTBlockTileType{}; - clear_tile(dpt_acc); // Initialize PGrad^T - - static_for<0, k2_loops, 1>{}([&](auto i_k2) { - block_sync_lds(); - gemm_2(dpt_acc, - get_slice_tile(do_lds_window, - sequence<0, i_k2 * kK2>{}, - sequence{}), - get_slice_tile( - v, sequence<0, i_k2 * kK2>{}, sequence{})); - block_sync_lds(); - }); - - // STAGE 5, P^T(PGrad^T - D) - const auto d = load_tile(d_dram_window); - - auto dst = SPGradTBlockTileType{}; - constexpr auto dst_spans = decltype(dst)::get_distributed_spans(); - sweep_tile_span(dst_spans[number<0>{}], [&](auto idx0) { - constexpr auto i_idx = make_tuple(idx0); - sweep_tile_span(dst_spans[number<1>{}], [&](auto idx1) { - constexpr auto i_j_idx = make_tuple(idx0, idx1); - bool undrop_flag = pt[i_j_idx] >= 0; - dst(i_j_idx) = - pt[i_j_idx] * - (!kHasDropout || undrop_flag ? (dpt_acc[i_j_idx] - d[i_idx]) : d[i_idx]); - }); - }); - - if constexpr(kHasBiasGrad) - { - const auto dbiast = [&]() { - if constexpr(kHasDropout) - { - return tile_elementwise_in( - [&rp_undrop](const auto& x) { - return type_convert(x * rp_undrop); - }, - dst); - } - else - { - return cast_tile(dst); - } - }(); - store_tile(biast_lds_shuffle_window, dbiast); - block_sync_lds(); - auto dbiast_tile = load_tile(dbiast_lds_shuffle_window); - auto dbiast_shuffle_tmp = make_static_distributed_tensor( - Policy::template MakeBiasTileDistribution()); - shuffle_tile(dbiast_shuffle_tmp, dbiast_tile); - store_tile(dbias_dram_block_window, dbiast_shuffle_tmp); - move_tile_window(dbias_dram_block_window, {kM0, 0}); - } - - // STAGE 6, SGrad^T@Q^T Gemm3 - block_sync_lds(); - const auto dst_gemm = cast_tile(dst); - - static_for<0, k3_loops, 1>{}([&](auto i_k3) { - block_sync_lds(); - gemm_3(dk_acc, - get_slice_tile( - dst_gemm, sequence{}, sequence<(i_k3 + 1) * kK3, kN0>{}), - get_slice_tile(qt_lds_window, - sequence<0, i_k3 * kK3>{}, - sequence{})); - block_sync_lds(); - }); - - // STAGE 7, SGrad@K^T Gemm4 - store_tile(ds_lds_window, dst_gemm); - - auto dq_acc = QGradBlockTileType{}; - clear_tile(dq_acc); // Initialize QGrad - - block_sync_lds(); - - static_for<0, k4_loops, 1>{}([&](auto i_k4) { - gemm_4(dq_acc, - get_slice_tile(ds_lds_window, - sequence<0, i_k4 * kK4>{}, - sequence{}), - get_slice_tile(kt_lds_window, - sequence<0, i_k4 * kK4>{}, - sequence{})); - }); - - // QGrad Scale - if constexpr(kHasDropout) - { - tile_elementwise_inout([&scale_rp_undrop](auto& x) { x = x * scale_rp_undrop; }, - dq_acc); - } - else - { - tile_elementwise_inout([&raw_scale](auto& x) { x = x * raw_scale; }, dq_acc); - } - const auto dq = cast_tile(dq_acc); - update_tile(dq_dram_block_window, dq); - - // move tile windows - move_tile_window(q_dram_block_window, {kM0, 0}); - move_tile_window(dq_dram_block_window, {kM0, 0}); - move_tile_window(do_dram_block_window, {kM0, 0}); - move_tile_window(lse_dram_window, {kM0}); - move_tile_window(d_dram_window, {kM0}); - } while(++i_total_loops < num_total_loop); - - // KGrad Scale - if constexpr(kHasDropout) - { - tile_elementwise_inout([&scale_rp_undrop](auto& x) { x = x * scale_rp_undrop; }, - dk_acc); - } - else - { - tile_elementwise_inout([&raw_scale](auto& x) { x = x * raw_scale; }, dk_acc); - } - // VGrad Scale - if constexpr(kHasDropout) - { - tile_elementwise_inout([&rp_undrop](auto& x) { x = x * rp_undrop; }, dv_acc); - } - - return ck_tile::make_tuple(dk_acc, dv_acc); - } -}; - -} // namespace ck_tile diff --git a/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_qs_ks_vr_dos_default_policy.hpp b/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_qs_ks_vr_dos_default_policy.hpp deleted file mode 100644 index ac81990e07..0000000000 --- a/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_qs_ks_vr_dos_default_policy.hpp +++ /dev/null @@ -1,20 +0,0 @@ -// SPDX-License-Identifier: MIT -// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved. - -#pragma once - -#include "ck_tile/ops/fmha/pipeline/block_fmha_bwd_pipeline_default_policy.hpp" - -namespace ck_tile { - -// This pipeline is v located in regs, q & k & do located in lds. -using BlockFmhaBwdDQDKDVPipelineQSKSVROGradSDefaultPolicy = - BlockFmhaBwdPipelineDefaultPolicy; - -} // namespace ck_tile diff --git a/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_pipeline_default_policy.hpp b/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_pipeline_default_policy.hpp index d867772a1f..4143c34ff8 100644 --- a/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_pipeline_default_policy.hpp +++ b/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_pipeline_default_policy.hpp @@ -11,6 +11,8 @@ #include "ck_tile/ops/gemm/warp/warp_gemm_dispatcher.hpp" #include "ck_tile/ops/gemm/block/block_gemm_areg_bsmem_creg_v1_custom_policy.hpp" #include "ck_tile/ops/gemm/block/block_gemm_areg_bsmem_creg_v1.hpp" +#include "ck_tile/ops/gemm/block/block_gemm_areg_breg_creg_v1_custom_policy.hpp" +#include "ck_tile/ops/gemm/block/block_gemm_areg_breg_creg_v1.hpp" #include "ck_tile/ops/gemm/block/block_gemm_asmem_breg_creg_v1_custom_policy.hpp" #include "ck_tile/ops/gemm/block/block_gemm_asmem_breg_creg_v1.hpp" #include "ck_tile/ops/gemm/block/block_gemm_asmem_bsmem_creg_v1_custom_policy.hpp" @@ -18,60 +20,215 @@ namespace ck_tile { -template struct BlockFmhaBwdPipelineDefaultPolicy { - static constexpr bool QLoadOnce = - QLoadOnce_; // if q load whole block length (qkhdim) to LDS at once - static constexpr bool QTLoadOnce = - QTLoadOnce_; // if q^t load whole block length (qkhdim) to LDS at once - static constexpr bool KLoadOnce = - KLoadOnce_; // if k load whole block length (qkhdim) to LDS at once - static constexpr bool KTLoadOnce = - KTLoadOnce_; // if k^t load whole block length (qkhdim) to LDS at once - static constexpr bool VLoadOnce = - VLoadOnce_; // if v load whole block length (vhdim) to Vgprs at once - static constexpr bool OGradLoadOnce = - OGradLoadOnce_; // if do load whole block length (vhdim) to LDS at once - static constexpr bool OGradTLoadOnce = - OGradTLoadOnce_; // if do^t load whole block length (vhdim) to LDS at once + template + CK_TILE_HOST_DEVICE static constexpr auto GetQKBlockGemm() + { + using BlockGemmProblem = + BlockGemmPipelineProblem>; + + using WarpGemm = WarpGemmMfmaDispatcher< + typename Problem::QDataType, + typename Problem::KDataType, + typename Problem::AccDataType, + Problem::BlockFmhaShape::Gemm0WarpTile::at(number<0>{}), + Problem::BlockFmhaShape::Gemm0WarpTile::at(number<1>{}), + Problem::BlockFmhaShape::Gemm0WarpTile::at(number<2>{}), + false, + Problem::BlockFmhaShape::Gemm0WarpTile::at(number<0>{}) == 16 ? false : true>; + + using BlockGemmPolicy = + BlockGemmARegBRegCRegV1CustomPolicy; + + return BlockGemmARegBRegCRegV1{}; + } + + template + CK_TILE_HOST_DEVICE static constexpr auto GetPTOGradTBlockGemm() + { + using BlockGemmProblem = + BlockGemmPipelineProblem>; + + using WarpGemm = + WarpGemmMfmaDispatcher{}), + Problem::BlockFmhaShape::Gemm1WarpTile::at(number<1>{}), + Problem::BlockFmhaShape::Gemm1WarpTile::at(number<2>{}), + true>; + + using BlockGemmPolicy = + BlockGemmARegBRegCRegV1CustomPolicy; + + return BlockGemmARegBRegCRegV1{}; + } + + template + CK_TILE_HOST_DEVICE static constexpr auto GetOGradVBlockGemm() + { + using BlockGemmProblem = + BlockGemmPipelineProblem>; + + using WarpGemm = WarpGemmMfmaDispatcher< + typename Problem::OGradDataType, + typename Problem::VDataType, + typename Problem::AccDataType, + Problem::BlockFmhaShape::Gemm2WarpTile::at(number<0>{}), + Problem::BlockFmhaShape::Gemm2WarpTile::at(number<1>{}), + Problem::BlockFmhaShape::Gemm2WarpTile::at(number<2>{}), + false, + Problem::BlockFmhaShape::Gemm0WarpTile::at(number<0>{}) == 16 ? false : true>; + + using BlockGemmPolicy = + BlockGemmARegBRegCRegV1CustomPolicy; + + return BlockGemmARegBRegCRegV1{}; + } + + template + CK_TILE_HOST_DEVICE static constexpr auto GetSGradTQTBlockGemm() + { + using BlockGemmProblem = + BlockGemmPipelineProblem>; + + using WarpGemm = + WarpGemmMfmaDispatcher{}), + Problem::BlockFmhaShape::Gemm3WarpTile::at(number<1>{}), + Problem::BlockFmhaShape::Gemm3WarpTile::at(number<2>{}), + true>; + + using BlockGemmPolicy = + BlockGemmARegBRegCRegV1CustomPolicy; + + return BlockGemmARegBRegCRegV1{}; + } + + template + CK_TILE_HOST_DEVICE static constexpr auto GetSGradKTBlockGemm() + { + using BlockGemmProblem = + BlockGemmPipelineProblem>; + + using WarpGemm = + WarpGemmMfmaDispatcher{}), + Problem::BlockFmhaShape::Gemm4WarpTile::at(number<1>{}), + Problem::BlockFmhaShape::Gemm4WarpTile::at(number<2>{}), + false>; + + using BlockGemmPolicy = + BlockGemmARegBRegCRegV1CustomPolicy; + + return BlockGemmARegBRegCRegV1{}; + } // these are for global load template CK_TILE_HOST_DEVICE static constexpr auto GetAlignmentQ() { - using QDataType = remove_cvref_t; - return 16 / sizeof(QDataType); + using QDataType = remove_cvref_t; + constexpr index_t kBlockSize = Problem::kBlockSize; + constexpr index_t kMNPerBlock = Problem::BlockFmhaShape::kM0; + constexpr index_t kKPerBlock = Problem::BlockFmhaShape::kK0; + constexpr index_t kMaxVecLoad = 16 / sizeof(QDataType); + constexpr index_t kMinVecLoad = 4 / sizeof(QDataType); + + constexpr index_t total_pixels = kMNPerBlock * kKPerBlock / kBlockSize; + + constexpr index_t kVecLoad = ((total_pixels / kMaxVecLoad) >= kMinVecLoad) + ? kMaxVecLoad + : (total_pixels / kMinVecLoad); + + return kVecLoad; } template CK_TILE_HOST_DEVICE static constexpr auto GetAlignmentK() { - using KDataType = remove_cvref_t; - return 16 / sizeof(KDataType); + using KDataType = remove_cvref_t; + constexpr index_t kBlockSize = Problem::kBlockSize; + constexpr index_t kMNPerBlock = Problem::BlockFmhaShape::kN0; + constexpr index_t kKPerBlock = Problem::BlockFmhaShape::kK0; + constexpr index_t kMaxVecLoad = 16 / sizeof(KDataType); + constexpr index_t kMinVecLoad = 4 / sizeof(KDataType); + + constexpr index_t total_pixels = kMNPerBlock * kKPerBlock / kBlockSize; + + constexpr index_t kVecLoad = ((total_pixels / kMaxVecLoad) >= kMinVecLoad) + ? kMaxVecLoad + : (total_pixels / kMinVecLoad); + + return kVecLoad; } template CK_TILE_HOST_DEVICE static constexpr auto GetAlignmentV() { - if constexpr(VLoadOnce) - { - using BlockGemm = remove_cvref_t())>; - constexpr auto config = BlockGemm::Policy::template GetWarpGemmMWarpNWarp(); - using WG = remove_cvref_t())>; - return WG::kK / WG::WarpGemmAttribute::Impl::kABKLane; - } - else - { - using VDataType = remove_cvref_t; - return 16 / sizeof(VDataType); - } + using VDataType = remove_cvref_t; + constexpr index_t kBlockSize = Problem::kBlockSize; + constexpr index_t kMNPerBlock = Problem::BlockFmhaShape::kN0; + constexpr index_t kKPerBlock = Problem::BlockFmhaShape::kK2; + constexpr index_t kMaxVecLoad = 16 / sizeof(VDataType); + constexpr index_t total_pixels = kMNPerBlock * kKPerBlock / kBlockSize; + + return total_pixels > kMaxVecLoad ? kMaxVecLoad : total_pixels; } template @@ -84,20 +241,39 @@ struct BlockFmhaBwdPipelineDefaultPolicy template CK_TILE_HOST_DEVICE static constexpr auto GetAlignmentOGrad() { - using OGradDataType = remove_cvref_t; - return 16 / sizeof(OGradDataType); + using OGradDataType = remove_cvref_t; + constexpr index_t kBlockSize = Problem::kBlockSize; + constexpr index_t kMNPerBlock = Problem::BlockFmhaShape::kM0; + constexpr index_t kKPerBlock = Problem::BlockFmhaShape::kK2; + constexpr index_t kMaxVecLoad = 16 / sizeof(OGradDataType); + constexpr index_t kMinVecLoad = 4 / sizeof(OGradDataType); + + constexpr index_t total_pixels = kMNPerBlock * kKPerBlock / kBlockSize; + + constexpr index_t kVecLoad = ((total_pixels / kMaxVecLoad) >= kMinVecLoad) + ? kMaxVecLoad + : (total_pixels / kMinVecLoad); + + return kVecLoad; } template - CK_TILE_HOST_DEVICE static constexpr auto GetAlignmentQGrad() + CK_TILE_HOST_DEVICE static constexpr auto GetAlignmentBias() { - using BlockGemm = remove_cvref_t())>; - constexpr auto config = BlockGemm::Policy::template GetWarpGemmMWarpNWarp(); - using WG = remove_cvref_t())>; - using CWarpDstr = typename WG::CWarpDstr; - constexpr auto vec = - CWarpDstr{}.get_ys_to_d_descriptor().get_lengths().at(number{}); - return vec; + using BiasDataType = remove_cvref_t; + constexpr index_t kBlockSize = Problem::kBlockSize; + constexpr index_t kMPerBlock = Problem::BlockFmhaShape::kM0; + constexpr index_t kNPerBlock = Problem::BlockFmhaShape::kN0; + constexpr index_t kMaxVecLoad = 16 / sizeof(BiasDataType); + constexpr index_t kMinVecLoad = 4 / sizeof(BiasDataType); + + constexpr index_t total_pixels = kMPerBlock * kNPerBlock / kBlockSize; + + constexpr index_t kVecLoad = ((total_pixels / kMaxVecLoad) >= kMinVecLoad) + ? kMaxVecLoad + : (total_pixels / kMinVecLoad); + + return kVecLoad; } template @@ -128,60 +304,35 @@ struct BlockFmhaBwdPipelineDefaultPolicy CK_TILE_HOST_DEVICE static constexpr auto GetTransposedAlignmentQ() { constexpr index_t kBlockSize = Problem::kBlockSize; - constexpr index_t kNPerBlock = Problem::BlockFmhaShape::kQKHeaddim; - constexpr index_t kKPerBlock = [&]() { - if constexpr(QTLoadOnce) - return Problem::BlockFmhaShape::kM0; - else - return Problem::BlockFmhaShape::kK3; - }(); + constexpr index_t kNPerBlock = Problem::BlockFmhaShape::kM0; + constexpr index_t kKPerBlock = Problem::BlockFmhaShape::kK0; + constexpr index_t total_pixels = kNPerBlock * kKPerBlock / kBlockSize; - // TODO: not correct! - if constexpr(total_pixels > 4) - return 4; - else - return 2; + return total_pixels / GetAlignmentQ(); } template CK_TILE_HOST_DEVICE static constexpr auto GetTransposedAlignmentK() { - constexpr index_t kBlockSize = Problem::kBlockSize; - constexpr index_t kNPerBlock = Problem::BlockFmhaShape::kQKHeaddim; - constexpr index_t kKPerBlock = [&]() { - if constexpr(KTLoadOnce) - return Problem::BlockFmhaShape::kN0; - else - return Problem::BlockFmhaShape::kK4; - }(); + constexpr index_t kBlockSize = Problem::kBlockSize; + constexpr index_t kNPerBlock = Problem::BlockFmhaShape::kN0; + constexpr index_t kKPerBlock = Problem::BlockFmhaShape::kK0; constexpr index_t total_pixels = kNPerBlock * kKPerBlock / kBlockSize; - // TODO: not correct! - if constexpr(total_pixels > 4) - return 4; - else - return 2; + return total_pixels / GetAlignmentK(); } template CK_TILE_HOST_DEVICE static constexpr auto GetTransposedAlignmentOGrad() { constexpr index_t kBlockSize = Problem::kBlockSize; - constexpr index_t kNPerBlock = Problem::BlockFmhaShape::kVHeaddim; - constexpr index_t kKPerBlock = [&]() { - if constexpr(OGradTLoadOnce) - return Problem::BlockFmhaShape::kM0; - else - return Problem::BlockFmhaShape::kK1; - }(); + constexpr index_t kNPerBlock = Problem::BlockFmhaShape::kM0; + constexpr index_t kKPerBlock = Problem::BlockFmhaShape::kK2; + constexpr index_t total_pixels = kNPerBlock * kKPerBlock / kBlockSize; - // TODO: not correct! - if constexpr(total_pixels > 4) - return 4; - else - return 2; + return total_pixels / GetAlignmentOGrad(); } template @@ -193,554 +344,56 @@ struct BlockFmhaBwdPipelineDefaultPolicy constexpr index_t total_pixels = kMPerBlock * kNPerBlock / kBlockSize; - // TODO: not correct! - if constexpr(total_pixels > 32) - return 8; - else - return 4; - } - - // these are for lds - template - CK_TILE_HOST_DEVICE static constexpr auto GetSmemKPackQ() - { - // TODO: this is for 3d layout - using QDataType = remove_cvref_t; - return 16 / sizeof(QDataType); + return total_pixels / GetAlignmentBias(); } template - CK_TILE_HOST_DEVICE static constexpr auto GetSmemKPackK() + CK_TILE_HOST_DEVICE static constexpr auto GetAlignmentPostQGradAcc() { - // TODO: this is for 3d layout - using KDataType = remove_cvref_t; - return 16 / sizeof(KDataType); + using AccDataType = remove_cvref_t; + return 16 / sizeof(AccDataType); } template - CK_TILE_HOST_DEVICE static constexpr auto GetSmemKPackV() + CK_TILE_HOST_DEVICE static constexpr auto GetAlignmentPostQGrad() { - // TODO: this is for 3d layout - using VDataType = remove_cvref_t; - return 16 / sizeof(VDataType); + return GetAlignmentPostQGradAcc(); } template - CK_TILE_HOST_DEVICE static constexpr auto GetSmemKPackBias() + CK_TILE_HOST_DEVICE static constexpr auto MakeKDramTileDistribution() { - // TODO: this is for 3d layout - using BiasDataType = remove_cvref_t; - return 16 / sizeof(BiasDataType); - } + constexpr index_t kBlockSize = Problem::kBlockSize; - template - CK_TILE_HOST_DEVICE static constexpr auto GetSmemKPackOGrad() - { - // TODO: this is for 3d layout - using OGradDataType = remove_cvref_t; - return 16 / sizeof(OGradDataType); - } - - template - CK_TILE_HOST_DEVICE static constexpr auto GetSmemKPackSGrad() - { - // TODO: this is for 3d layout - using GemmDataType = remove_cvref_t; - return 16 / sizeof(GemmDataType); - } - - template - CK_TILE_HOST_DEVICE static constexpr auto MakeVInRegDramTileDistribution() - { constexpr index_t kNPerBlock = Problem::BlockFmhaShape::kN0; - constexpr index_t kKPerBlock = Problem::BlockFmhaShape::kVHeaddim; + constexpr index_t kKPerBlock = Problem::BlockFmhaShape::kK0; - 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>(); - - constexpr index_t NIterPerWarp = kNPerBlock / (NWarp * WG::kN); - constexpr index_t KIterPerWarp = kKPerBlock / WG::kK; - - constexpr auto v_block_outer_dstr_encoding = - tile_distribution_encoding, - tuple, sequence>, - tuple>, - tuple>, - sequence<1, 2>, - sequence<0, 0>>{}; - - constexpr auto v_block_dstr_encode = detail::make_embed_tile_distribution_encoding( - v_block_outer_dstr_encoding, typename WG::BWarpDstrEncoding{}); - - constexpr auto v_block_dstr = make_static_tile_distribution(v_block_dstr_encode); - - return v_block_dstr; - } - - // 3d + padding - template - CK_TILE_HOST_DEVICE static constexpr auto MakeXLdsBlockDescriptor() - { - constexpr auto x_lds_block_desc_0 = make_naive_tensor_descriptor( - make_tuple(number{}, number{}, number{}), - make_tuple(number<(MNPerBlock + 1) * KPack>{}, number{}, number<1>{}), - number<8>{}, - number<1>{}); - - constexpr auto x_lds_block_desc = transform_tensor_descriptor( - x_lds_block_desc_0, - make_tuple(make_pass_through_transform(MNPerBlock), - make_merge_transform(make_tuple(KPerBlock / KPack, KPack))), - make_tuple(sequence<1>{}, sequence<0, 2>{}), - make_tuple(sequence<0>{}, sequence<1>{})); - - return x_lds_block_desc; - } - - // 3d + padding - template - CK_TILE_HOST_DEVICE static constexpr auto MakeXLdsBlockDescriptorAsXT() - { - constexpr auto x_lds_block_desc_0 = make_naive_tensor_descriptor( - make_tuple(number{}, number{}, number{}), - make_tuple(number<(MNPerBlock + 1) * KPack>{}, number{}, number<1>{}), - number<8>{}, - number<1>{}); - - constexpr auto xt_lds_block_desc = transform_tensor_descriptor( - x_lds_block_desc_0, - make_tuple(make_pass_through_transform(MNPerBlock), - make_merge_transform(make_tuple(KPerBlock / KPack, KPack))), - make_tuple(sequence<1>{}, sequence<0, 2>{}), - make_tuple(sequence<1>{}, sequence<0>{})); - - return xt_lds_block_desc; - } - - template - CK_TILE_HOST_DEVICE static constexpr auto MakeXTLdsBlockDescriptor() - { - static_assert(PixelsPerRow % KPack == 0); - constexpr index_t NPerRow = PixelsPerRow / KPack; - static_assert(MNPerBlock % NPerRow == 0); - static_assert(KPerBlock % KPack == 0); - - constexpr auto xt_lds_block_desc_0 = make_naive_tensor_descriptor( - make_tuple(number{}, - number{}, - number{}, - number{}), - make_tuple(number<(MNPerBlock / NPerRow) * (PixelsPerRow + KPack)>{}, - number{}, - number{}, - number<1>{}), - number{}, - number<1>{}); - - constexpr auto xt_lds_block_desc = transform_tensor_descriptor( - xt_lds_block_desc_0, - make_tuple( - make_merge_transform(make_tuple(number{}, number{})), - make_merge_transform(make_tuple(number{}, number{}))), - make_tuple(sequence<1, 2>{}, sequence<0, 3>{}), - make_tuple(sequence<0>{}, sequence<1>{})); - - return xt_lds_block_desc; - } - - template - CK_TILE_HOST_DEVICE static constexpr auto MakeQLdsBlockDescriptor() - { - constexpr index_t kMPerBlock = Problem::BlockFmhaShape::kM0; - constexpr index_t kKPerBlock = [&]() { - if constexpr(QLoadOnce) - return Problem::BlockFmhaShape::kQKHeaddim; - else - return Problem::BlockFmhaShape::kK0; - }(); - constexpr index_t kKPack = GetSmemKPackQ(); - - return MakeXLdsBlockDescriptor(); - } - - template - CK_TILE_HOST_DEVICE static constexpr auto MakeQLdsBlockDescriptorAsQT() - { - constexpr index_t kMPerBlock = Problem::BlockFmhaShape::kM0; - constexpr index_t kKPerBlock = [&]() { - if constexpr(QLoadOnce) - return Problem::BlockFmhaShape::kQKHeaddim; - else - return Problem::BlockFmhaShape::kK0; - }(); - constexpr index_t kKPack = GetSmemKPackQ(); - - return MakeXLdsBlockDescriptorAsXT(); - } - - template - CK_TILE_HOST_DEVICE static constexpr auto MakeKLdsBlockDescriptor() - { - constexpr index_t kNPerBlock = Problem::BlockFmhaShape::kN0; - constexpr index_t kKPerBlock = [&]() { - if constexpr(KLoadOnce) - return Problem::BlockFmhaShape::kQKHeaddim; - else - return Problem::BlockFmhaShape::kK0; - }(); - constexpr index_t kKPack = GetSmemKPackK(); - - return MakeXLdsBlockDescriptor(); - } - - template - CK_TILE_HOST_DEVICE static constexpr auto MakeKLdsBlockDescriptorAsKT() - { - constexpr index_t kNPerBlock = Problem::BlockFmhaShape::kN0; - constexpr index_t kKPerBlock = [&]() { - if constexpr(KLoadOnce) - return Problem::BlockFmhaShape::kQKHeaddim; - else - return Problem::BlockFmhaShape::kK0; - }(); - constexpr index_t kKPack = GetSmemKPackK(); - - return MakeXLdsBlockDescriptorAsXT(); - } - - template - CK_TILE_HOST_DEVICE static constexpr auto MakeVLdsBlockDescriptor() - { - constexpr index_t kNPerBlock = Problem::BlockFmhaShape::kN0; - constexpr index_t kKPerBlock = Problem::BlockFmhaShape::kK2; - constexpr index_t kKPack = GetSmemKPackV(); - - return MakeXLdsBlockDescriptor(); - } - - template - CK_TILE_HOST_DEVICE static constexpr auto MakeOGradLdsBlockDescriptor() - { - constexpr index_t kMPerBlock = Problem::BlockFmhaShape::kM0; - constexpr index_t kKPerBlock = [&]() { - if constexpr(OGradLoadOnce) - return Problem::BlockFmhaShape::kVHeaddim; - else - return Problem::BlockFmhaShape::kK2; - }(); - constexpr index_t kKPack = GetSmemKPackOGrad(); - - return MakeXLdsBlockDescriptor(); - } - - template - CK_TILE_HOST_DEVICE static constexpr auto MakeOGradLdsBlockDescriptorAsOGradT() - { - constexpr index_t kMPerBlock = Problem::BlockFmhaShape::kM0; - constexpr index_t kKPerBlock = [&]() { - if constexpr(OGradLoadOnce) - return Problem::BlockFmhaShape::kVHeaddim; - else - return Problem::BlockFmhaShape::kK2; - }(); - constexpr index_t kKPack = GetSmemKPackOGrad(); - - return MakeXLdsBlockDescriptorAsXT(); - } - - template - CK_TILE_HOST_DEVICE static constexpr auto MakeSGradLdsBlockDescriptor() - { - constexpr index_t kMPerBlock = Problem::BlockFmhaShape::kM0; - constexpr index_t kKPerBlock = Problem::BlockFmhaShape::kN0; - constexpr index_t kKPack = GetSmemKPackSGrad(); - - return MakeXLdsBlockDescriptor(); - } - - template - CK_TILE_HOST_DEVICE static constexpr auto MakeQTLdsBlockDescriptor() - { - using QDataType = remove_cvref_t; - constexpr index_t Banks = 32; // TODO: need change based on arch - constexpr index_t PixelsPerRow = Banks * 4 / sizeof(QDataType); - constexpr index_t kKPack = GetSmemKPackQ(); - constexpr index_t kNPerBlock = Problem::BlockFmhaShape::kQKHeaddim; - constexpr index_t kKPerBlock = [&]() { - if constexpr(QTLoadOnce) - return Problem::BlockFmhaShape::kM0; - else - return Problem::BlockFmhaShape::kK3; - }(); - - return MakeXTLdsBlockDescriptor(); - } - - template - CK_TILE_HOST_DEVICE static constexpr auto MakeKTLdsBlockDescriptor() - { - using KDataType = remove_cvref_t; - constexpr index_t Banks = 32; // TODO: need change based on arch - constexpr index_t PixelsPerRow = Banks * 4 / sizeof(KDataType); - constexpr index_t kKPack = GetSmemKPackK(); - constexpr index_t kNPerBlock = Problem::BlockFmhaShape::kQKHeaddim; - constexpr index_t kKPerBlock = [&]() { - if constexpr(KTLoadOnce) - return Problem::BlockFmhaShape::kN0; - else - return Problem::BlockFmhaShape::kK4; - }(); - - return MakeXTLdsBlockDescriptor(); - } - - template - CK_TILE_HOST_DEVICE static constexpr auto MakeOGradTLdsBlockDescriptor() - { - using OGradDataType = remove_cvref_t; - constexpr index_t Banks = 32; // TODO: need change based on arch - constexpr index_t PixelsPerRow = Banks * 4 / sizeof(OGradDataType); - constexpr index_t kKPack = GetSmemKPackOGrad(); - constexpr index_t kNPerBlock = Problem::BlockFmhaShape::kVHeaddim; - constexpr index_t kKPerBlock = [&]() { - if constexpr(OGradTLoadOnce) - return Problem::BlockFmhaShape::kM0; - else - return Problem::BlockFmhaShape::kK1; - }(); - - return MakeXTLdsBlockDescriptor(); - } - - template - CK_TILE_HOST_DEVICE static constexpr auto MakeBiasTLdsBlockDescriptor() - { - using BiasDataType = remove_cvref_t; - constexpr index_t Banks = 32; // TODO: need change based on arch - constexpr index_t PixelsPerRow = Banks * 4 / sizeof(BiasDataType); - constexpr index_t kKPack = GetSmemKPackBias(); - constexpr index_t kMPerBlock = Problem::BlockFmhaShape::kM0; - constexpr index_t kNPerBlock = Problem::BlockFmhaShape::kN0; - - static_assert(PixelsPerRow % kKPack == 0); - constexpr index_t NPerRow = PixelsPerRow / kKPack; - static_assert(kNPerBlock % NPerRow == 0); - static_assert(kMPerBlock % kKPack == 0); - - constexpr auto biast_lds_block_desc_0 = make_naive_tensor_descriptor( - make_tuple(number{}, - number{}, - number{}, - number{}), - make_tuple(number<(kNPerBlock / NPerRow) * (PixelsPerRow + kKPack)>{}, - number{}, - number{}, - number<1>{}), - number{}, - number<1>{}); - - constexpr auto biast_lds_block_desc = transform_tensor_descriptor( - biast_lds_block_desc_0, - make_tuple( - make_merge_transform(make_tuple(number{}, number{})), - make_merge_transform(make_tuple(number{}, number{}))), - make_tuple(sequence<1, 2>{}, sequence<0, 3>{}), - make_tuple(sequence<1>{}, sequence<0>{})); - - return biast_lds_block_desc; - } - - template - CK_TILE_HOST_DEVICE static constexpr ck_tile::index_t GetSmemSizeQ() - { - constexpr index_t smem_size_q = sizeof(typename Problem::QDataType) * - MakeQLdsBlockDescriptor().get_element_space_size(); - return smem_size_q; - } - - template - CK_TILE_HOST_DEVICE static constexpr ck_tile::index_t GetSmemSizeQT() - { - constexpr index_t smem_size_qt = [&]() { - if constexpr(QLoadOnce && !QTLoadOnce) - return 0; - else - return sizeof(typename Problem::QDataType) * - MakeQTLdsBlockDescriptor().get_element_space_size(); - }(); - return smem_size_qt; - } - - template - CK_TILE_HOST_DEVICE static constexpr ck_tile::index_t GetSmemSizeK() - { - constexpr index_t smem_size_k = sizeof(typename Problem::KDataType) * - MakeKLdsBlockDescriptor().get_element_space_size(); - return smem_size_k; - } - - template - CK_TILE_HOST_DEVICE static constexpr ck_tile::index_t GetSmemSizeKT() - { - constexpr index_t smem_size_kt = [&]() { - if constexpr(KLoadOnce && !KTLoadOnce) - return 0; - else - return sizeof(typename Problem::KDataType) * - MakeKTLdsBlockDescriptor().get_element_space_size(); - }(); - return smem_size_kt; - } - - template - CK_TILE_HOST_DEVICE static constexpr ck_tile::index_t GetSmemSizeV() - { - constexpr index_t smem_size_v = [&]() { - if constexpr(VLoadOnce) - return 0; - else - return sizeof(typename Problem::VDataType) * - MakeVLdsBlockDescriptor().get_element_space_size(); - }(); - return smem_size_v; - } - - template - CK_TILE_HOST_DEVICE static constexpr ck_tile::index_t GetSmemSizeOGrad() - { - constexpr index_t smem_size_do = - sizeof(typename Problem::OGradDataType) * - MakeOGradLdsBlockDescriptor().get_element_space_size(); - return smem_size_do; - } - - template - CK_TILE_HOST_DEVICE static constexpr ck_tile::index_t GetSmemSizeOGradT() - { - constexpr index_t smem_size_dot = [&]() { - if constexpr(OGradLoadOnce && !OGradTLoadOnce) - return 0; - else - return sizeof(typename Problem::OGradDataType) * - MakeOGradTLdsBlockDescriptor().get_element_space_size(); - }(); - return smem_size_dot; - } - - template - CK_TILE_HOST_DEVICE static constexpr ck_tile::index_t GetSmemSizeSGrad() - { - constexpr index_t smem_size_ds = - sizeof(typename Problem::GemmDataType) * - MakeSGradLdsBlockDescriptor().get_element_space_size(); - return smem_size_ds; - } - - template - CK_TILE_HOST_DEVICE static constexpr ck_tile::index_t GetSmemSizeBias() - { - constexpr index_t smem_size_bias = [&]() { - if constexpr(Problem::BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS) - return sizeof(typename Problem::BiasDataType) * - MakeBiasTLdsBlockDescriptor().get_element_space_size(); - else - return 0; - }(); - return smem_size_bias; - } - - template - CK_TILE_HOST_DEVICE static constexpr ck_tile::index_t GetSmemSize() - { - constexpr index_t smem_size_q = GetSmemSizeQ(); - constexpr index_t smem_size_qt = GetSmemSizeQT(); - constexpr index_t smem_size_k = GetSmemSizeK(); - constexpr index_t smem_size_kt = GetSmemSizeKT(); - constexpr index_t smem_size_v = GetSmemSizeV(); - constexpr index_t smem_size_do = GetSmemSizeOGrad(); - constexpr index_t smem_size_dot = GetSmemSizeOGradT(); - constexpr index_t smem_size_ds = GetSmemSizeSGrad(); - constexpr index_t smem_size_bias = GetSmemSizeBias(); - constexpr index_t smem_size_transpose = max(smem_size_ds, smem_size_bias); - - index_t smem_size = 0; - - if constexpr(QLoadOnce && OGradLoadOnce) - smem_size += smem_size_q + smem_size_qt + smem_size_do + smem_size_dot + - smem_size_transpose; // 1~4 & 10 - else if(QLoadOnce && !OGradLoadOnce && !OGradTLoadOnce) - smem_size += smem_size_q + smem_size_qt + - max(smem_size_do, - smem_size_dot, - smem_size_transpose); // 5/7/11 TODO: Multiple buffers strategy - else if(!QLoadOnce && !QTLoadOnce && OGradLoadOnce) - smem_size += smem_size_do + smem_size_dot + - max(smem_size_q, - smem_size_qt, - smem_size_transpose); // 6/8/12 TODO: Multiple buffers strategy - else if(!QLoadOnce && !QTLoadOnce && !OGradLoadOnce && !OGradTLoadOnce) - smem_size += max(smem_size_q, - smem_size_qt, - smem_size_do, - smem_size_dot, - smem_size_transpose); // 9/13 TODO: Multiple buffers strategy - - // 14/15 needs to be adjusted - if constexpr(KLoadOnce) - smem_size += (smem_size_k + smem_size_kt); // 1~13 - else - smem_size = - max(smem_size_k, smem_size_kt, smem_size); // 14/15 TODO: Multiple buffers strategy - - return max(smem_size, smem_size_v); // 15 - } - - template - CK_TILE_HOST_DEVICE static constexpr auto MakeLSEDDramTileDistribution() - { - 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>(); - - constexpr index_t kMPerBlock = Problem::BlockFmhaShape::kM0; - - constexpr index_t N1 = WG::WarpGemmAttribute::Impl::kCNLane; - constexpr index_t N0 = NWarp; - - constexpr index_t M4 = WG::WarpGemmAttribute::Impl::kCM1PerLane * 2; - constexpr index_t M3 = WG::WarpGemmAttribute::Impl::kCMLane; - constexpr index_t M2 = WG::WarpGemmAttribute::Impl::kCM0PerLane / 2; - constexpr index_t M1 = MWarp; - constexpr index_t M0 = kMPerBlock / (M1 * WG::WarpGemmAttribute::Impl::kM); + constexpr index_t K1 = GetAlignmentK(); + constexpr index_t K0 = kKPerBlock / K1; + constexpr index_t N1 = get_warp_size() / K0; + constexpr index_t N0 = kBlockSize / get_warp_size(); + constexpr index_t N2 = kNPerBlock / (N1 * N0); return make_static_tile_distribution( - tile_distribution_encoding, - tuple>, - tuple, sequence<1, 0>>, - tuple, sequence<3, 1>>, - sequence<1, 1, 1>, - sequence<0, 2, 4>>{}); + tile_distribution_encoding, + tuple, sequence>, + tuple, sequence<1, 2>>, + tuple, sequence<1, 0>>, + sequence<1, 2>, + sequence<2, 1>>{}); } template CK_TILE_HOST_DEVICE static constexpr auto MakeVDramTileDistribution() { - using VDataType = remove_cvref_t; - constexpr index_t kBlockSize = Problem::kBlockSize; constexpr index_t kNPerBlock = Problem::BlockFmhaShape::kN0; constexpr index_t kKPerBlock = Problem::BlockFmhaShape::kK2; - constexpr index_t K1 = 16 / sizeof(VDataType); + constexpr index_t K1 = GetAlignmentV(); constexpr index_t K0 = kKPerBlock / K1; constexpr index_t N2 = get_warp_size() / K0; - // coalesce reading for each blocks constexpr index_t N1 = kBlockSize / get_warp_size(); constexpr index_t N0 = kNPerBlock / (N2 * N1); @@ -759,56 +412,21 @@ struct BlockFmhaBwdPipelineDefaultPolicy constexpr index_t kBlockSize = Problem::kBlockSize; constexpr index_t kMPerBlock = Problem::BlockFmhaShape::kM0; - constexpr index_t kKPerBlock = [&]() { - if constexpr(QLoadOnce) - return Problem::BlockFmhaShape::kQKHeaddim; - else - return Problem::BlockFmhaShape::kK0; - }(); + constexpr index_t kKPerBlock = Problem::BlockFmhaShape::kK0; constexpr index_t K1 = GetAlignmentQ(); constexpr index_t K0 = kKPerBlock / K1; - constexpr index_t M2 = get_warp_size() / K0; - // coalesce reading for each blocks - constexpr index_t M1 = kBlockSize / get_warp_size(); - constexpr index_t M0 = kMPerBlock / (M2 * M1); + constexpr index_t M1 = get_warp_size() / K0; + constexpr index_t M0 = kBlockSize / get_warp_size(); + constexpr index_t M2 = kMPerBlock / (M1 * M0); return make_static_tile_distribution( tile_distribution_encoding, tuple, sequence>, tuple, sequence<1, 2>>, - tuple, sequence<2, 0>>, + tuple, sequence<1, 0>>, sequence<1, 2>, - sequence<0, 1>>{}); - } - - template - CK_TILE_HOST_DEVICE static constexpr auto MakeKDramTileDistribution() - { - constexpr index_t kBlockSize = Problem::kBlockSize; - - constexpr index_t kNPerBlock = Problem::BlockFmhaShape::kN0; - constexpr index_t kKPerBlock = [&]() { - if constexpr(KLoadOnce) - return Problem::BlockFmhaShape::kQKHeaddim; - else - return Problem::BlockFmhaShape::kK0; - }(); - - constexpr index_t K1 = GetAlignmentK(); - constexpr index_t K0 = kKPerBlock / K1; - constexpr index_t N2 = get_warp_size() / K0; - // coalesce reading for each blocks - constexpr index_t N1 = kBlockSize / get_warp_size(); - constexpr index_t N0 = kNPerBlock / (N2 * N1); - - return make_static_tile_distribution( - tile_distribution_encoding, - tuple, sequence>, - tuple, sequence<1, 2>>, - tuple, sequence<2, 0>>, - sequence<1, 2>, - sequence<0, 1>>{}); + sequence<2, 1>>{}); } template @@ -817,27 +435,72 @@ struct BlockFmhaBwdPipelineDefaultPolicy constexpr index_t kBlockSize = Problem::kBlockSize; constexpr index_t kMPerBlock = Problem::BlockFmhaShape::kM0; - constexpr index_t kKPerBlock = [&]() { - if constexpr(OGradLoadOnce) - return Problem::BlockFmhaShape::kVHeaddim; - else - return Problem::BlockFmhaShape::kK2; - }(); + constexpr index_t kKPerBlock = Problem::BlockFmhaShape::kK2; constexpr index_t K1 = GetAlignmentOGrad(); constexpr index_t K0 = kKPerBlock / K1; - constexpr index_t M2 = get_warp_size() / K0; - // coalesce reading for each blocks - constexpr index_t M1 = kBlockSize / get_warp_size(); - constexpr index_t M0 = kMPerBlock / (M2 * M1); + constexpr index_t M1 = get_warp_size() / K0; + constexpr index_t M0 = kBlockSize / get_warp_size(); + constexpr index_t M2 = kMPerBlock / (M1 * M0); return make_static_tile_distribution( tile_distribution_encoding, tuple, sequence>, tuple, sequence<1, 2>>, - tuple, sequence<2, 0>>, + tuple, sequence<1, 0>>, sequence<1, 2>, - sequence<0, 1>>{}); + sequence<2, 1>>{}); + } + + template + CK_TILE_HOST_DEVICE static constexpr auto MakeLSEDDramTileDistribution() + { + constexpr auto config = BlockGemm::Policy::template GetWarpGemmMWarpNWarp(); + constexpr index_t MWarp = config.template at<1>(); + constexpr index_t NWarp = config.template at<2>(); + + constexpr index_t kMPerBlock = Problem::BlockFmhaShape::kM0; + + // Duplicate dimension + constexpr index_t N0 = NWarp; + constexpr index_t N1 = + (get_warp_size() / kMPerBlock) > 1 ? (get_warp_size() / kMPerBlock) : 1; + + constexpr index_t M0 = MWarp; + constexpr index_t M1 = (get_warp_size() / kMPerBlock) > 1 ? kMPerBlock : get_warp_size(); + constexpr index_t M2 = + (get_warp_size() / kMPerBlock) > 1 ? 1 : (kMPerBlock / get_warp_size()); + + return make_static_tile_distribution( + tile_distribution_encoding, + tuple>, + tuple, sequence<0, 1>>, + tuple, sequence<1, 1>>, + sequence<1>, + sequence<2>>{}); + } + + template + CK_TILE_HOST_DEVICE static constexpr auto MakeBiasTileDistribution() + { + constexpr index_t kBlockSize = Problem::kBlockSize; + + constexpr index_t kMPerBlock = Problem::BlockFmhaShape::kM0; + constexpr index_t kNPerBlock = Problem::BlockFmhaShape::kN0; + + constexpr index_t N1 = GetAlignmentBias(); + constexpr index_t N0 = kNPerBlock / N1; + constexpr index_t M1 = get_warp_size() / N0; + constexpr index_t M0 = kBlockSize / get_warp_size(); + constexpr index_t M2 = kMPerBlock / (M1 * M0); + + return make_static_tile_distribution( + tile_distribution_encoding, + tuple, sequence>, + tuple, sequence<1, 2>>, + tuple, sequence<1, 0>>, + sequence<1, 2>, + sequence<2, 1>>{}); } template @@ -881,463 +544,1377 @@ struct BlockFmhaBwdPipelineDefaultPolicy } template - CK_TILE_DEVICE static constexpr auto MakeQTDramTileDistribution() + CK_TILE_HOST_DEVICE static constexpr auto MakePostQGradAccDramTileDistribution() { + using AccDataType = remove_cvref_t; + constexpr index_t kBlockSize = Problem::kBlockSize; - constexpr index_t kNPerBlock = Problem::BlockFmhaShape::kQKHeaddim; - constexpr index_t kKPerBlock = [&]() { - if constexpr(QTLoadOnce) - return Problem::BlockFmhaShape::kM0; - else - return Problem::BlockFmhaShape::kK3; - }(); + constexpr index_t kMPerBlock = Problem::kM0; + constexpr index_t kKPerBlock = Problem::kQKHeaddim; - constexpr index_t N1 = GetTransposedAlignmentQ(); - constexpr index_t N0 = kNPerBlock / N1; // P + constexpr index_t K1 = 16 / sizeof(AccDataType); + constexpr index_t K0 = kKPerBlock / K1; - constexpr index_t total_pixels = kNPerBlock * kKPerBlock / kBlockSize; - static_assert(total_pixels % N1 == 0); // TODO: this is not always true? - constexpr index_t K3 = total_pixels / N1; - constexpr index_t kKPack = GetSmemKPackQ(); - static_assert(kKPack % K3 == 0); - constexpr index_t K2 = kKPack / K3; // TODO: this dimention could be outside single wave - constexpr index_t K1 = get_warp_size() / (K2 * N0); - constexpr index_t K0 = kBlockSize / get_warp_size(); - static_assert(kKPerBlock == K0 * K1 * K2 * K3); + constexpr index_t M2 = get_warp_size() / K0; + constexpr index_t M1 = kBlockSize / get_warp_size(); + constexpr index_t M0 = kMPerBlock / (M1 * M2); return make_static_tile_distribution( tile_distribution_encoding, - tuple, sequence>, - tuple, sequence<2, 1, 2>>, - tuple, sequence<1, 0, 2>>, - sequence<2, 1>, - sequence<3, 1>>{}); + tuple, sequence, sequence>, + tuple, sequence<2, 3>>, + tuple, sequence<2, 0>>, + sequence<1, 2, 3>, + sequence<0, 0, 1>>{}); } template - CK_TILE_HOST_DEVICE static constexpr auto MakeShuffledQTRegBlockDescriptor() + CK_TILE_HOST_DEVICE static constexpr auto MakePostQGradDramTileDistribution() { - constexpr index_t kBlockSize = Problem::kBlockSize; - constexpr index_t kNPerBlock = Problem::BlockFmhaShape::kQKHeaddim; - constexpr index_t kKPerBlock = [&]() { - if constexpr(QTLoadOnce) - return Problem::BlockFmhaShape::kM0; - else - return Problem::BlockFmhaShape::kK3; - }(); + using AccDataType = remove_cvref_t; - constexpr index_t N1 = GetTransposedAlignmentQ(); - constexpr index_t N0 = kNPerBlock / N1; - constexpr index_t total_pixels = kNPerBlock * kKPerBlock / kBlockSize; - static_assert(total_pixels % N1 == 0); // TODO: this is not always true? - constexpr index_t K3 = total_pixels / N1; - constexpr index_t kKPack = GetSmemKPackQ(); - static_assert(kKPack % K3 == 0); - constexpr index_t K2 = kKPack / K3; // TODO: this dimention could be outside single wave - constexpr index_t K1 = get_warp_size() / (K2 * N0); - constexpr index_t K0 = kBlockSize / get_warp_size(); + constexpr index_t kBlockSize = Problem::kBlockSize; + constexpr index_t kMPerBlock = Problem::kM0; + constexpr index_t kKPerBlock = Problem::kQKHeaddim; + + constexpr index_t K1 = 16 / sizeof(AccDataType); + constexpr index_t K0 = kKPerBlock / K1; + + constexpr index_t M2 = get_warp_size() / K0; + constexpr index_t M1 = kBlockSize / get_warp_size(); + constexpr index_t M0 = kMPerBlock / (M1 * M2); return make_static_tile_distribution( tile_distribution_encoding, - tuple, sequence>, - tuple, sequence<2, 1, 2>>, - tuple, sequence<1, 0, 2>>, + tuple, sequence>, + tuple, sequence<1, 2>>, + tuple, sequence<2, 0>>, sequence<1, 2>, - sequence<1, 3>>{}); + sequence<0, 1>>{}); + } + + // these are for lds + template + CK_TILE_HOST_DEVICE static constexpr auto GetSmemKPackQ() + { + return GetAlignmentQ(); } template - CK_TILE_DEVICE static constexpr auto MakeKTDramTileDistribution() + CK_TILE_HOST_DEVICE static constexpr auto GetSmemKPackQT() { - constexpr index_t kBlockSize = Problem::kBlockSize; - constexpr index_t kNPerBlock = Problem::BlockFmhaShape::kQKHeaddim; - constexpr index_t kKPerBlock = [&]() { - if constexpr(KTLoadOnce) - return Problem::BlockFmhaShape::kN0; - else - return Problem::BlockFmhaShape::kK4; - }(); - - constexpr index_t N1 = GetTransposedAlignmentK(); - constexpr index_t N0 = kNPerBlock / N1; // P - - constexpr index_t total_pixels = kNPerBlock * kKPerBlock / kBlockSize; - static_assert(total_pixels % N1 == 0); // TODO: this is not always true? - constexpr index_t K3 = total_pixels / N1; - constexpr index_t kKPack = GetSmemKPackK(); - static_assert(kKPack % K3 == 0); - constexpr index_t K2 = kKPack / K3; // TODO: this dimention could be outside single wave - constexpr index_t K1 = get_warp_size() / (K2 * N0); - constexpr index_t K0 = kBlockSize / get_warp_size(); - static_assert(kKPerBlock == K0 * K1 * K2 * K3); - - return make_static_tile_distribution( - tile_distribution_encoding, - tuple, sequence>, - tuple, sequence<2, 1, 2>>, - tuple, sequence<1, 0, 2>>, - sequence<2, 1>, - sequence<3, 1>>{}); + return GetTransposedAlignmentQ(); } template - CK_TILE_HOST_DEVICE static constexpr auto MakeShuffledKTRegBlockDescriptor() + CK_TILE_HOST_DEVICE static constexpr auto GetSmemKPackK() { - constexpr index_t kBlockSize = Problem::kBlockSize; - constexpr index_t kNPerBlock = Problem::BlockFmhaShape::kQKHeaddim; - constexpr index_t kKPerBlock = [&]() { - if constexpr(KTLoadOnce) - return Problem::BlockFmhaShape::kN0; - else - return Problem::BlockFmhaShape::kK4; - }(); - - constexpr index_t N1 = GetTransposedAlignmentK(); - constexpr index_t N0 = kNPerBlock / N1; - constexpr index_t total_pixels = kNPerBlock * kKPerBlock / kBlockSize; - static_assert(total_pixels % N1 == 0); // TODO: this is not always true? - constexpr index_t K3 = total_pixels / N1; - constexpr index_t kKPack = GetSmemKPackK(); - static_assert(kKPack % K3 == 0); - constexpr index_t K2 = kKPack / K3; // TODO: this dimention could be outside single wave - constexpr index_t K1 = get_warp_size() / (K2 * N0); - constexpr index_t K0 = kBlockSize / get_warp_size(); - - return make_static_tile_distribution( - tile_distribution_encoding, - tuple, sequence>, - tuple, sequence<2, 1, 2>>, - tuple, sequence<1, 0, 2>>, - sequence<1, 2>, - sequence<1, 3>>{}); + return GetAlignmentK(); } template - CK_TILE_DEVICE static constexpr auto MakeOGradTDramTileDistribution() + CK_TILE_HOST_DEVICE static constexpr auto GetSmemKPackKT() { - constexpr index_t kBlockSize = Problem::kBlockSize; - constexpr index_t kNPerBlock = Problem::BlockFmhaShape::kVHeaddim; - constexpr index_t kKPerBlock = [&]() { - if constexpr(OGradTLoadOnce) - return Problem::BlockFmhaShape::kM0; - else - return Problem::BlockFmhaShape::kK1; - }(); - - constexpr index_t N1 = GetTransposedAlignmentOGrad(); - constexpr index_t N0 = kNPerBlock / N1; // P - - constexpr index_t total_pixels = kNPerBlock * kKPerBlock / kBlockSize; - static_assert(total_pixels % N1 == 0); // TODO: this is not always true? - constexpr index_t K3 = total_pixels / N1; - constexpr index_t kKPack = GetSmemKPackOGrad(); - static_assert(kKPack % K3 == 0); - constexpr index_t K2 = kKPack / K3; // TODO: this dimention could be outside single wave - constexpr index_t K1 = get_warp_size() / (K2 * N0); - constexpr index_t K0 = kBlockSize / get_warp_size(); - static_assert(kKPerBlock == K0 * K1 * K2 * K3); - - return make_static_tile_distribution( - tile_distribution_encoding, - tuple, sequence>, - tuple, sequence<2, 1, 2>>, - tuple, sequence<1, 0, 2>>, - sequence<2, 1>, - sequence<3, 1>>{}); + return GetTransposedAlignmentK(); } template - CK_TILE_HOST_DEVICE static constexpr auto MakeShuffledOGradTRegBlockDescriptor() + CK_TILE_HOST_DEVICE static constexpr auto GetSmemKPackV() { - constexpr index_t kBlockSize = Problem::kBlockSize; - constexpr index_t kNPerBlock = Problem::BlockFmhaShape::kVHeaddim; - constexpr index_t kKPerBlock = [&]() { - if constexpr(OGradTLoadOnce) - return Problem::BlockFmhaShape::kM0; - else - return Problem::BlockFmhaShape::kK1; - }(); - - constexpr index_t N1 = GetTransposedAlignmentOGrad(); - constexpr index_t N0 = kNPerBlock / N1; - constexpr index_t total_pixels = kNPerBlock * kKPerBlock / kBlockSize; - static_assert(total_pixels % N1 == 0); // TODO: this is not always true? - constexpr index_t K3 = total_pixels / N1; - constexpr index_t kKPack = GetSmemKPackOGrad(); - static_assert(kKPack % K3 == 0); - constexpr index_t K2 = kKPack / K3; // TODO: this dimention could be outside single wave - constexpr index_t K1 = get_warp_size() / (K2 * N0); - constexpr index_t K0 = kBlockSize / get_warp_size(); - - return make_static_tile_distribution( - tile_distribution_encoding, - tuple, sequence>, - tuple, sequence<2, 1, 2>>, - tuple, sequence<1, 0, 2>>, - sequence<1, 2>, - sequence<1, 3>>{}); + return GetAlignmentV(); } template - CK_TILE_DEVICE static constexpr auto MakeBiasTileDistribution() + CK_TILE_HOST_DEVICE static constexpr auto GetSmemKPackBias() + { + return GetAlignmentBias(); + } + + template + CK_TILE_HOST_DEVICE static constexpr auto GetSmemKPackBiasT() + { + return GetTransposedAlignmentBias(); + } + + template + CK_TILE_HOST_DEVICE static constexpr auto GetSmemKPackOGrad() + { + return GetAlignmentOGrad(); + } + + template + CK_TILE_HOST_DEVICE static constexpr auto GetSmemKPackOGradT() + { + return GetTransposedAlignmentOGrad(); + } + + template + CK_TILE_HOST_DEVICE static constexpr auto GetSmemKPackSGrad() + { + // TODO: this is for 3d layout + using GemmDataType = remove_cvref_t; + return 16 / sizeof(GemmDataType); + } + + template + CK_TILE_HOST_DEVICE static constexpr auto MakeXLdsBlockDescriptor() + { + constexpr auto DataTypeSize = 2; // sizeof(F16/BF16) + constexpr auto MNLdsLayer = + (32 * 4 / KPerBlock / DataTypeSize) < 1 ? 1 : (32 * 4 / KPerBlock / DataTypeSize); + + constexpr auto x_lds_block_desc_0 = make_naive_tensor_descriptor( + make_tuple(number{}, + number{}, + number{}), + make_tuple(number{}, number{}, number<1>{}), + number{}, + number<1>{}); + + constexpr auto x_lds_block_desc_permuted = transform_tensor_descriptor( + x_lds_block_desc_0, + make_tuple(make_xor_transform(make_tuple(number{}, + number{})), + make_pass_through_transform(number{})), + make_tuple(sequence<1, 0>{}, sequence<2>{}), + make_tuple(sequence<1, 0>{}, sequence<2>{})); + + constexpr auto x_lds_block_desc_xk0_mnldslayer_mn_xk1 = transform_tensor_descriptor( + x_lds_block_desc_permuted, + make_tuple(make_unmerge_transform( + make_tuple(number{}, number{})), + make_pass_through_transform(number{}), + make_pass_through_transform(number{})), + make_tuple(sequence<0>{}, sequence<1>{}, sequence<2>{}), + make_tuple(sequence<0, 2>{}, sequence<1>{}, sequence<3>{})); + + constexpr auto x_lds_block_desc = transform_tensor_descriptor( + x_lds_block_desc_xk0_mnldslayer_mn_xk1, + make_tuple(make_merge_transform_v3_division_mod( + make_tuple(number{}, number{})), + make_merge_transform_v3_division_mod( + make_tuple(number{}, number{}))), + make_tuple(sequence<1, 2>{}, sequence<0, 3>{}), + make_tuple(sequence<0>{}, sequence<1>{})); + + return x_lds_block_desc; + } + + template + CK_TILE_HOST_DEVICE static constexpr auto MakeXTLdsBlockDescriptor() + { + // kfold and mpair dimension is not always required. + // more dimension in merge_transform increase the difficulty of generating immarg offset + // for compiler. + constexpr auto MNPerXDL = Problem::BlockFmhaShape::Gemm0WarpTile::at(number<0>{}); + constexpr auto kBlockSize = Problem::kBlockSize; + + constexpr auto MN0 = MNPerBlock / KPack; + constexpr auto MN1 = KPack; + + constexpr auto KThreadWrite = kBlockSize / MN0; + constexpr auto K0Number = KPerBlock / KPackT; + constexpr auto K0PerThreadWrite = K0Number / KThreadWrite; + constexpr auto KThreadRead = get_warp_size() / MNPerXDL; // assume 32x32x8 mfma + constexpr auto K0PerThreadRead = K0Number / KThreadRead; + + constexpr auto kfold = (KPackT * MN0 * 2 > 128) ? 1 : 128 / (KPackT * MN0 * 2); + constexpr auto KThreadReadPerm = + (kfold * K0PerThreadWrite / K0PerThreadRead) > 1 + ? KThreadRead / (kfold * K0PerThreadWrite / K0PerThreadRead) + : KThreadRead; + + // 1<=mnpair<=n0 + constexpr auto mnpair = + (KPackT * MNPerXDL * 2 > 128) + ? 1 + : ((128 / (KPackT * MNPerXDL * 2)) > MN0 ? MN0 : 128 / (KPackT * MNPerXDL * 2)); + + constexpr auto xt_lds_block_desc_raw = make_naive_tensor_descriptor( + make_tuple(number{}, + number{}, + number{}, + number{}, + number{}, + KPackT), + make_tuple(number{}, + number{}, + number{}, + number{}, + number{}, + number<1>{}), + number{}, + number<1>{}); + + constexpr auto xt_lds_block_desc_permuted = transform_tensor_descriptor( + xt_lds_block_desc_raw, + make_tuple( + make_pass_through_transform(number{}), + make_pass_through_transform(number{}), + make_xor_transform( + make_tuple(number{}, number{})), + make_pass_through_transform(number{}), + make_pass_through_transform(KPackT)), + make_tuple( + sequence<0>{}, sequence<1>{}, sequence<2, 3>{}, sequence<4>{}, sequence<5>{}), + make_tuple( + sequence<0>{}, sequence<1>{}, sequence<2, 3>{}, sequence<4>{}, sequence<5>{})); + + constexpr auto xt_lds_block_desc_unmerged = transform_tensor_descriptor( + xt_lds_block_desc_permuted, + make_tuple( + make_pass_through_transform(number{}), + make_pass_through_transform(number{}), + make_unmerge_transform(make_tuple(number{}, number{})), + make_unmerge_transform(make_tuple(number{}, number{})), + make_pass_through_transform(number{}), + make_pass_through_transform(KPackT)), + make_tuple(sequence<0>{}, + sequence<1>{}, + sequence<2>{}, + sequence<3>{}, + sequence<4>{}, + sequence<5>{}), + make_tuple(sequence<1>{}, + sequence<2>{}, + sequence<0, 3>{}, + sequence<4, 5>{}, + sequence<6>{}, + sequence<7>{})); + + constexpr auto xt_lds_block_desc = transform_tensor_descriptor( + xt_lds_block_desc_unmerged, + make_tuple(make_merge_transform_v3_division_mod( + make_tuple(number{}, + number{}, + number{}, + number{}, + number{})), + make_merge_transform_v3_division_mod( + make_tuple(number{}, number{}, number{}))), + make_tuple(sequence<0, 1, 4, 2, 7>{}, sequence<5, 6, 3>{}), + make_tuple(sequence<0>{}, sequence<1>{})); + + return xt_lds_block_desc; + } + + template + CK_TILE_HOST_DEVICE static constexpr auto MakeKLdsWriteBlockDescriptor() { - constexpr index_t kBlockSize = Problem::kBlockSize; - constexpr index_t kMPerBlock = Problem::BlockFmhaShape::kM0; constexpr index_t kNPerBlock = Problem::BlockFmhaShape::kN0; + constexpr index_t kKPerBlock = Problem::BlockFmhaShape::kK0; + constexpr index_t kKPack = GetSmemKPackK(); - constexpr index_t N1 = GetTransposedAlignmentBias(); - constexpr index_t N0 = kNPerBlock / N1; // P + return MakeXLdsBlockDescriptor(); + } - constexpr index_t total_pixels = kMPerBlock * kNPerBlock / kBlockSize; - static_assert(total_pixels % N1 == 0); // TODO: this is not always true? - constexpr index_t M3 = total_pixels / N1; - constexpr index_t kKPack = GetSmemKPackBias(); - static_assert(kKPack % M3 == 0); - constexpr index_t M2 = kKPack / M3; // TODO: this dimention could be outside single wave - constexpr index_t M1 = get_warp_size() / (M2 * N0); - constexpr index_t M0 = kBlockSize / get_warp_size(); - static_assert(kMPerBlock == M0 * M1 * M2 * M3); + template + CK_TILE_HOST_DEVICE static constexpr auto MakeKRegSliceBlockDescriptor() + { + using BlockGemm = remove_cvref_t())>; + constexpr auto config = BlockGemm::Policy::template GetWarpGemmMWarpNWarp(); + using WarpGemm = remove_cvref_t())>; + + constexpr index_t MWarp = Problem::BlockFmhaShape::Gemm0BlockWarps::at(number<0>{}); + constexpr index_t NWarp = Problem::BlockFmhaShape::Gemm0BlockWarps::at(number<1>{}); + + constexpr index_t kNPerBlock = Problem::BlockFmhaShape::kN0; + constexpr index_t kKPerBlock = Problem::BlockFmhaShape::kK0; + + constexpr index_t NIterPerWarp = kNPerBlock / (NWarp * WarpGemm::kN); + constexpr index_t KIterPerWarp = kKPerBlock / WarpGemm::kK; + + constexpr auto k_block_outer_dstr_encoding = + tile_distribution_encoding, + tuple, sequence>, + tuple>, + tuple>, + sequence<1, 2>, + sequence<0, 0>>{}; + + constexpr auto k_block_dstr_encode = detail::make_embed_tile_distribution_encoding( + k_block_outer_dstr_encoding, typename WarpGemm::BWarpDstrEncoding{}); + + constexpr auto k_block_dstr = make_static_tile_distribution(k_block_dstr_encode); + + return k_block_dstr; + } + + template + CK_TILE_HOST_DEVICE static constexpr auto MakeKRegBlockDescriptor() + { + using BlockGemm = remove_cvref_t())>; + constexpr auto config = BlockGemm::Policy::template GetWarpGemmMWarpNWarp(); + using WarpGemm = remove_cvref_t())>; + + constexpr index_t MWarp = Problem::BlockFmhaShape::Gemm0BlockWarps::at(number<0>{}); + constexpr index_t NWarp = Problem::BlockFmhaShape::Gemm0BlockWarps::at(number<1>{}); + + constexpr index_t kNPerBlock = Problem::BlockFmhaShape::kN0; + constexpr index_t kKPerBlock = Problem::BlockFmhaShape::kQKHeaddim; + + constexpr index_t NIterPerWarp = kNPerBlock / (NWarp * WarpGemm::kN); + constexpr index_t KIterPerWarp = kKPerBlock / WarpGemm::kK; + + constexpr auto k_block_outer_dstr_encoding = + tile_distribution_encoding, + tuple, sequence>, + tuple>, + tuple>, + sequence<1, 2>, + sequence<0, 0>>{}; + + constexpr auto k_block_dstr_encode = detail::make_embed_tile_distribution_encoding( + k_block_outer_dstr_encoding, typename WarpGemm::BWarpDstrEncoding{}); + + constexpr auto k_block_dstr = make_static_tile_distribution(k_block_dstr_encode); + + return k_block_dstr; + } + + template + CK_TILE_HOST_DEVICE static constexpr auto MakeVLdsWriteBlockDescriptor() + { + constexpr index_t kNPerBlock = Problem::BlockFmhaShape::kN0; + constexpr index_t kKPerBlock = Problem::BlockFmhaShape::kK2; + + constexpr index_t kVPack = GetSmemKPackV(); + + return MakeXLdsBlockDescriptor(); + } + + template + CK_TILE_HOST_DEVICE static constexpr auto MakeVRegSliceBlockDescriptor() + { + using BlockGemm = remove_cvref_t())>; + constexpr auto config = BlockGemm::Policy::template GetWarpGemmMWarpNWarp(); + using WarpGemm = remove_cvref_t())>; + + constexpr index_t MWarp = Problem::BlockFmhaShape::Gemm2BlockWarps::at(number<0>{}); + constexpr index_t NWarp = Problem::BlockFmhaShape::Gemm2BlockWarps::at(number<1>{}); + + constexpr index_t kNPerBlock = Problem::BlockFmhaShape::kN0; + constexpr index_t kKPerBlock = Problem::BlockFmhaShape::kK2; + + constexpr index_t NIterPerWarp = kNPerBlock / (NWarp * WarpGemm::kN); + constexpr index_t KIterPerWarp = kKPerBlock / WarpGemm::kK; + + constexpr auto v_block_outer_dstr_encoding = + tile_distribution_encoding, + tuple, sequence>, + tuple>, + tuple>, + sequence<1, 2>, + sequence<0, 0>>{}; + + constexpr auto v_block_dstr_encode = detail::make_embed_tile_distribution_encoding( + v_block_outer_dstr_encoding, typename WarpGemm::BWarpDstrEncoding{}); + + constexpr auto v_block_dstr = make_static_tile_distribution(v_block_dstr_encode); + + return v_block_dstr; + } + + template + CK_TILE_HOST_DEVICE static constexpr auto MakeVRegBlockDescriptor() + { + using BlockGemm = remove_cvref_t())>; + constexpr auto config = BlockGemm::Policy::template GetWarpGemmMWarpNWarp(); + using WarpGemm = remove_cvref_t())>; + + constexpr index_t MWarp = Problem::BlockFmhaShape::Gemm2BlockWarps::at(number<0>{}); + constexpr index_t NWarp = Problem::BlockFmhaShape::Gemm2BlockWarps::at(number<1>{}); + + constexpr index_t kNPerBlock = Problem::BlockFmhaShape::kN0; + constexpr index_t kKPerBlock = Problem::BlockFmhaShape::kVHeaddim; + + constexpr index_t NIterPerWarp = kNPerBlock / (NWarp * WarpGemm::kN); + constexpr index_t KIterPerWarp = kKPerBlock / WarpGemm::kK; + + constexpr auto v_block_outer_dstr_encoding = + tile_distribution_encoding, + tuple, sequence>, + tuple>, + tuple>, + sequence<1, 2>, + sequence<0, 0>>{}; + + constexpr auto v_block_dstr_encode = detail::make_embed_tile_distribution_encoding( + v_block_outer_dstr_encoding, typename WarpGemm::BWarpDstrEncoding{}); + + constexpr auto v_block_dstr = make_static_tile_distribution(v_block_dstr_encode); + + return v_block_dstr; + } + + template + CK_TILE_HOST_DEVICE static constexpr auto MakeShuffledKRegWriteBlockDescriptor() + { + constexpr index_t kBlockSize = Problem::kBlockSize; + + constexpr index_t kKPerBlock = Problem::BlockFmhaShape::kK0; + + constexpr index_t K1 = GetAlignmentK(); + constexpr index_t K0 = kKPerBlock / K1; + constexpr index_t N2 = GetTransposedAlignmentK(); + constexpr index_t N1 = get_warp_size() / K0; + constexpr index_t N0 = kBlockSize / get_warp_size(); return make_static_tile_distribution( tile_distribution_encoding, - tuple, sequence>, - tuple, sequence<1, 2, 1>>, - tuple, sequence<1, 0, 2>>, + tuple, sequence>, + tuple, sequence<1, 2>>, + tuple, sequence<1, 0>>, + sequence<2, 1>, + sequence<1, 2>>{}); + } + + template + CK_TILE_HOST_DEVICE static constexpr auto MakeShuffledKLdsWriteBlockDescriptor() + { + // Hold all data + constexpr index_t kNPerBlock = Problem::BlockFmhaShape::kQKHeaddim; + constexpr index_t kKPerBlock = Problem::BlockFmhaShape::kN0; + + constexpr index_t kKPack = GetSmemKPackK(); + constexpr index_t kKPackT = GetSmemKPackKT(); + + return MakeXTLdsBlockDescriptor(); + } + + template + CK_TILE_HOST_DEVICE static constexpr auto MakeKTLdsReadBlockDescriptor() + { + constexpr index_t kNPerBlock = Problem::BlockFmhaShape::kQKHeaddim; + constexpr index_t kKPerBlock = Problem::BlockFmhaShape::kN0; + + auto shuffled_k_lds_block_desc = MakeShuffledKLdsWriteBlockDescriptor(); + + return transform_tensor_descriptor( + shuffled_k_lds_block_desc, + make_tuple(make_pass_through_transform(number{}), + make_pass_through_transform(number{})), + make_tuple(sequence<1>{}, sequence<0>{}), + make_tuple(sequence<0>{}, sequence<1>{})); + } + + template + CK_TILE_HOST_DEVICE static constexpr auto MakeKTRegBlockDescriptor() + { + using BlockGemm = remove_cvref_t())>; + constexpr auto config = BlockGemm::Policy::template GetWarpGemmMWarpNWarp(); + using WarpGemm = remove_cvref_t())>; + + constexpr index_t MWarp = Problem::BlockFmhaShape::Gemm4BlockWarps::at(number<0>{}); + constexpr index_t NWarp = Problem::BlockFmhaShape::Gemm4BlockWarps::at(number<1>{}); + + constexpr index_t kNPerBlock = Problem::BlockFmhaShape::kQKHeaddim; + constexpr index_t kKPerBlock = Problem::BlockFmhaShape::kN0; + + constexpr index_t NIterPerWarp = kNPerBlock / (NWarp * WarpGemm::kN); + constexpr index_t KIterPerWarp = kKPerBlock / WarpGemm::kK; + + constexpr auto kt_block_outer_dstr_encoding = + tile_distribution_encoding, + tuple, sequence>, + tuple>, + tuple>, sequence<1, 2>, - sequence<3, 1>>{}); + sequence<0, 0>>{}; + + constexpr auto kt_block_dstr_encode = detail::make_embed_tile_distribution_encoding( + kt_block_outer_dstr_encoding, typename WarpGemm::BWarpDstrEncoding{}); + + constexpr auto kt_block_dstr = make_static_tile_distribution(kt_block_dstr_encode); + + return kt_block_dstr; + } + + template + CK_TILE_HOST_DEVICE static constexpr auto MakeQLdsBlockDescriptor() + { + constexpr index_t kMPerBlock = Problem::BlockFmhaShape::kM0; + constexpr index_t kKPerBlock = Problem::BlockFmhaShape::kK0; + + constexpr index_t kKPack = GetSmemKPackQ(); + + return MakeXLdsBlockDescriptor(); + } + + template + CK_TILE_HOST_DEVICE static constexpr auto MakeQRegSliceBlockDescriptor() + { + using BlockGemm = remove_cvref_t())>; + constexpr auto config = BlockGemm::Policy::template GetWarpGemmMWarpNWarp(); + using WarpGemm = remove_cvref_t())>; + + constexpr index_t MWarp = Problem::BlockFmhaShape::Gemm0BlockWarps::at(number<0>{}); + constexpr index_t NWarp = Problem::BlockFmhaShape::Gemm0BlockWarps::at(number<1>{}); + + constexpr index_t kMPerBlock = Problem::BlockFmhaShape::kM0; + constexpr index_t kKPerBlock = Problem::BlockFmhaShape::kK0; + + constexpr index_t MIterPerWarp = kMPerBlock / (MWarp * WarpGemm::kM); + constexpr index_t KIterPerWarp = kKPerBlock / WarpGemm::kK; + + constexpr auto q_block_outer_dstr_encoding = + tile_distribution_encoding, + tuple, sequence>, + tuple>, + tuple>, + sequence<1, 2>, + sequence<0, 0>>{}; + + constexpr auto q_block_dstr_encode = detail::make_embed_tile_distribution_encoding( + q_block_outer_dstr_encoding, typename WarpGemm::AWarpDstrEncoding{}); + + constexpr auto q_block_dstr = make_static_tile_distribution(q_block_dstr_encode); + + return q_block_dstr; + } + + template + CK_TILE_HOST_DEVICE static constexpr auto MakeShuffledQRegWriteBlockDescriptor() + { + constexpr index_t kBlockSize = Problem::kBlockSize; + + constexpr index_t kKPerBlock = Problem::BlockFmhaShape::kK0; + + constexpr index_t K1 = GetAlignmentQ(); + constexpr index_t K0 = kKPerBlock / K1; + constexpr index_t N2 = GetTransposedAlignmentQ(); + constexpr index_t N1 = get_warp_size() / K0; + constexpr index_t N0 = kBlockSize / get_warp_size(); + + return make_static_tile_distribution( + tile_distribution_encoding, + tuple, sequence>, + tuple, sequence<1, 2>>, + tuple, sequence<1, 0>>, + sequence<2, 1>, + sequence<1, 2>>{}); + } + + template + CK_TILE_HOST_DEVICE static constexpr auto MakeShuffledQLdsWriteBlockDescriptor() + { + // Hold full block data + constexpr index_t kNPerBlock = Problem::BlockFmhaShape::kQKHeaddim; + constexpr index_t kKPerBlock = Problem::BlockFmhaShape::kM0; + + constexpr index_t kKPack = GetSmemKPackQ(); + constexpr index_t kKPackT = GetSmemKPackQT(); + + return MakeXTLdsBlockDescriptor(); + } + + template + CK_TILE_HOST_DEVICE static constexpr auto MakeQTLdsReadBlockDescriptor() + { + // Hold full block data + constexpr index_t kNPerBlock = Problem::BlockFmhaShape::kQKHeaddim; + constexpr index_t kKPerBlock = Problem::BlockFmhaShape::kM0; + + auto shuffled_q_lds_block_desc = MakeShuffledQLdsWriteBlockDescriptor(); + + return transform_tensor_descriptor( + shuffled_q_lds_block_desc, + make_tuple(make_pass_through_transform(number{}), + make_pass_through_transform(number{})), + make_tuple(sequence<1>{}, sequence<0>{}), + make_tuple(sequence<0>{}, sequence<1>{})); + } + + template + CK_TILE_HOST_DEVICE static constexpr auto MakeQTRegSliceBlockDescriptor() + { + using BlockGemm = remove_cvref_t())>; + constexpr auto config = BlockGemm::Policy::template GetWarpGemmMWarpNWarp(); + using WarpGemm = remove_cvref_t())>; + + constexpr index_t MWarp = Problem::BlockFmhaShape::Gemm3BlockWarps::at(number<0>{}); + constexpr index_t NWarp = Problem::BlockFmhaShape::Gemm3BlockWarps::at(number<1>{}); + + constexpr index_t kNPerBlock = Problem::BlockFmhaShape::kQKHeaddim; + constexpr index_t kKPerBlock = Problem::BlockFmhaShape::kK3; + + constexpr index_t NIterPerWarp = kNPerBlock / (NWarp * WarpGemm::kN); + constexpr index_t KIterPerWarp = kKPerBlock / WarpGemm::kK; + + constexpr auto qt_block_outer_dstr_encoding = + tile_distribution_encoding, + tuple, sequence>, + tuple>, + tuple>, + sequence<1, 2>, + sequence<0, 0>>{}; + + constexpr auto qt_block_dstr_encode = detail::make_embed_tile_distribution_encoding( + qt_block_outer_dstr_encoding, typename WarpGemm::BWarpDstrEncoding{}); + + constexpr auto qt_block_dstr = make_static_tile_distribution(qt_block_dstr_encode); + + return qt_block_dstr; + } + + template + CK_TILE_HOST_DEVICE static constexpr auto MakeSGradTRegSliceBlockDescriptor() + { + using BlockGemm = remove_cvref_t())>; + constexpr auto config = BlockGemm::Policy::template GetWarpGemmMWarpNWarp(); + using WarpGemm = remove_cvref_t())>; + + constexpr index_t MWarp = Problem::BlockFmhaShape::Gemm3BlockWarps::at(number<0>{}); + constexpr index_t NWarp = Problem::BlockFmhaShape::Gemm3BlockWarps::at(number<1>{}); + + constexpr index_t kMPerBlock = Problem::BlockFmhaShape::kN0; + constexpr index_t kKPerBlock = Problem::BlockFmhaShape::kK3; + + constexpr index_t MIterPerWarp = kMPerBlock / (MWarp * WarpGemm::kM); + constexpr index_t KIterPerWarp = kKPerBlock / WarpGemm::kK; + + constexpr auto dst_block_outer_dstr_encoding = + tile_distribution_encoding, + tuple, sequence>, + tuple>, + tuple>, + sequence<1, 2>, + sequence<0, 0>>{}; + + constexpr auto dst_block_dstr_encode = detail::make_embed_tile_distribution_encoding( + dst_block_outer_dstr_encoding, typename WarpGemm::AWarpDstrEncoding{}); + + constexpr auto dst_block_dstr = make_static_tile_distribution(dst_block_dstr_encode); + + return dst_block_dstr; + } + + template + CK_TILE_HOST_DEVICE static constexpr auto MakeLSEDLdsWriteBlockDescriptor() + { + constexpr index_t kMPerBlock = Problem::BlockFmhaShape::kM0; + using LSEDType = remove_cvref_t; + constexpr index_t kMPack = 16 / sizeof(LSEDType); + + constexpr auto lsed_lds_block_desc = + make_naive_tensor_descriptor(make_tuple(number{}), + make_tuple(number<1>{}), + number{}, + number<1>{}); + + return lsed_lds_block_desc; + } + + template + CK_TILE_HOST_DEVICE static constexpr auto MakeLSEDLdsReadBlockDescriptor() + { + 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>(); + + constexpr index_t kMPerBlock = Problem::BlockFmhaShape::kM0; + + constexpr index_t N1 = WG::WarpGemmAttribute::Impl::kCNLane; + constexpr index_t N0 = NWarp; + + // M4 *2 and M2 /2 when swizzle mode enabled + constexpr index_t SwizzleConfig = WG::kM == 16 ? 1 : 2; + // constexpr index_t SwizzleConfig = 1; + constexpr index_t M4 = WG::WarpGemmAttribute::Impl::kCM1PerLane * SwizzleConfig; + constexpr index_t M3 = WG::WarpGemmAttribute::Impl::kCMLane; + constexpr index_t M2 = WG::WarpGemmAttribute::Impl::kCM0PerLane / SwizzleConfig; + constexpr index_t M1 = MWarp; + constexpr index_t M0 = kMPerBlock / (M1 * WG::WarpGemmAttribute::Impl::kM); + + return make_static_tile_distribution( + tile_distribution_encoding, + tuple>, + tuple, sequence<1, 0>>, + tuple, sequence<3, 1>>, + sequence<1, 1, 1>, + sequence<0, 2, 4>>{}); + } + + template + CK_TILE_HOST_DEVICE static constexpr auto MakeOGradLdsBlockDescriptor() + { + // Hold full block data + constexpr index_t kMPerBlock = Problem::BlockFmhaShape::kM0; + constexpr index_t kKPerBlock = Problem::BlockFmhaShape::kK2; + + constexpr index_t kKPack = GetSmemKPackOGrad(); + + return MakeXLdsBlockDescriptor(); + } + + template + CK_TILE_HOST_DEVICE static constexpr auto MakeOGradRegSliceBlockDescriptor() + { + using BlockGemm = remove_cvref_t())>; + constexpr auto config = BlockGemm::Policy::template GetWarpGemmMWarpNWarp(); + using WarpGemm = remove_cvref_t())>; + + constexpr index_t MWarp = Problem::BlockFmhaShape::Gemm2BlockWarps::at(number<0>{}); + constexpr index_t NWarp = Problem::BlockFmhaShape::Gemm2BlockWarps::at(number<1>{}); + + constexpr index_t kMPerBlock = Problem::BlockFmhaShape::kM0; + constexpr index_t kKPerBlock = Problem::BlockFmhaShape::kK2; + + constexpr index_t MIterPerWarp = kMPerBlock / (MWarp * WarpGemm::kM); + constexpr index_t KIterPerWarp = kKPerBlock / WarpGemm::kK; + + constexpr auto do_block_outer_dstr_encoding = + tile_distribution_encoding, + tuple, sequence>, + tuple>, + tuple>, + sequence<1, 2>, + sequence<0, 0>>{}; + + constexpr auto do_block_dstr_encode = detail::make_embed_tile_distribution_encoding( + do_block_outer_dstr_encoding, typename WarpGemm::AWarpDstrEncoding{}); + + constexpr auto do_block_dstr = make_static_tile_distribution(do_block_dstr_encode); + + return do_block_dstr; + } + + template + CK_TILE_HOST_DEVICE static constexpr auto MakeShuffledOGradRegWriteBlockDescriptor() + { + constexpr index_t kBlockSize = Problem::kBlockSize; + + constexpr index_t kKPerBlock = Problem::BlockFmhaShape::kK2; + + constexpr index_t K1 = GetAlignmentOGrad(); + constexpr index_t K0 = kKPerBlock / K1; + constexpr index_t N2 = GetTransposedAlignmentOGrad(); + constexpr index_t N1 = get_warp_size() / K0; + constexpr index_t N0 = kBlockSize / get_warp_size(); + + return make_static_tile_distribution( + tile_distribution_encoding, + tuple, sequence>, + tuple, sequence<1, 2>>, + tuple, sequence<1, 0>>, + sequence<2, 1>, + sequence<1, 2>>{}); + } + + template + CK_TILE_HOST_DEVICE static constexpr auto MakeShuffledOGradLdsWriteBlockDescriptor() + { + // Hold all data + constexpr index_t kNPerBlock = Problem::BlockFmhaShape::kVHeaddim; + constexpr index_t kKPerBlock = Problem::BlockFmhaShape::kM0; + + constexpr index_t kKPack = GetSmemKPackOGrad(); + constexpr index_t kKPackT = GetSmemKPackOGradT(); + + return MakeXTLdsBlockDescriptor(); + } + + template + CK_TILE_HOST_DEVICE static constexpr auto MakeOGradTLdsReadBlockDescriptor() + { + // Hold all data + constexpr index_t kNPerBlock = Problem::BlockFmhaShape::kVHeaddim; + constexpr index_t kKPerBlock = Problem::BlockFmhaShape::kM0; + auto shuffled_do_lds_block_desc = MakeShuffledOGradLdsWriteBlockDescriptor(); + + return transform_tensor_descriptor( + shuffled_do_lds_block_desc, + make_tuple(make_pass_through_transform(number{}), + make_pass_through_transform(number{})), + make_tuple(sequence<1>{}, sequence<0>{}), + make_tuple(sequence<0>{}, sequence<1>{})); + } + + template + CK_TILE_HOST_DEVICE static constexpr auto MakeOGradTRegSliceBlockDescriptor() + { + using BlockGemm = remove_cvref_t())>; + constexpr auto config = BlockGemm::Policy::template GetWarpGemmMWarpNWarp(); + using WarpGemm = remove_cvref_t())>; + + constexpr index_t MWarp = Problem::BlockFmhaShape::Gemm1BlockWarps::at(number<0>{}); + constexpr index_t NWarp = Problem::BlockFmhaShape::Gemm1BlockWarps::at(number<1>{}); + + constexpr index_t kNPerBlock = Problem::BlockFmhaShape::kVHeaddim; + // constexpr index_t kNPerBlock = 32; + constexpr index_t kKPerBlock = Problem::BlockFmhaShape::kK1; + + constexpr index_t NIterPerWarp = kNPerBlock / (NWarp * WarpGemm::kN); + constexpr index_t KIterPerWarp = kKPerBlock / WarpGemm::kK; + + constexpr auto dot_block_outer_dstr_encoding = + tile_distribution_encoding, + tuple, sequence>, + tuple>, + tuple>, + sequence<1, 2>, + sequence<0, 0>>{}; + + constexpr auto dot_block_dstr_encode = detail::make_embed_tile_distribution_encoding( + dot_block_outer_dstr_encoding, typename WarpGemm::BWarpDstrEncoding{}); + + constexpr auto dot_block_dstr = make_static_tile_distribution(dot_block_dstr_encode); + + return dot_block_dstr; + } + + template + CK_TILE_HOST_DEVICE static constexpr auto MakePTRegSliceBlockDescriptor() + { + using BlockGemm = remove_cvref_t())>; + constexpr auto config = BlockGemm::Policy::template GetWarpGemmMWarpNWarp(); + using WarpGemm = remove_cvref_t())>; + + constexpr index_t MWarp = Problem::BlockFmhaShape::Gemm1BlockWarps::at(number<0>{}); + constexpr index_t NWarp = Problem::BlockFmhaShape::Gemm1BlockWarps::at(number<1>{}); + + constexpr index_t kMPerBlock = Problem::BlockFmhaShape::kN0; + constexpr index_t kKPerBlock = Problem::BlockFmhaShape::kK1; + + constexpr index_t MIterPerWarp = kMPerBlock / (MWarp * WarpGemm::kM); + constexpr index_t KIterPerWarp = kKPerBlock / WarpGemm::kK; + + constexpr auto pt_block_outer_dstr_encoding = + tile_distribution_encoding, + tuple, sequence>, + tuple>, + tuple>, + sequence<1, 2>, + sequence<0, 0>>{}; + + constexpr auto pt_block_dstr_encode = detail::make_embed_tile_distribution_encoding( + pt_block_outer_dstr_encoding, typename WarpGemm::AWarpDstrEncoding{}); + + constexpr auto pt_block_dstr = make_static_tile_distribution(pt_block_dstr_encode); + + return pt_block_dstr; + } + + template + CK_TILE_HOST_DEVICE static constexpr auto MakeSGradLdsBlockDescriptor() + { + constexpr index_t kMPerBlock = Problem::BlockFmhaShape::kM0; + constexpr index_t kKPerBlock = Problem::BlockFmhaShape::kN0; + constexpr index_t kKPack = GetSmemKPackSGrad(); + + return MakeXLdsBlockDescriptor(); + } + + template + CK_TILE_HOST_DEVICE static constexpr auto MakeSGradRegSliceBlockDescriptor() + { + using BlockGemm = remove_cvref_t())>; + constexpr auto config = BlockGemm::Policy::template GetWarpGemmMWarpNWarp(); + using WarpGemm = remove_cvref_t())>; + + constexpr index_t MWarp = Problem::BlockFmhaShape::Gemm4BlockWarps::at(number<0>{}); + constexpr index_t NWarp = Problem::BlockFmhaShape::Gemm4BlockWarps::at(number<1>{}); + + constexpr index_t kMPerBlock = Problem::BlockFmhaShape::kM0; + constexpr index_t kKPerBlock = Problem::BlockFmhaShape::kK4; + + constexpr index_t MIterPerWarp = kMPerBlock / (MWarp * WarpGemm::kM); + constexpr index_t KIterPerWarp = kKPerBlock / WarpGemm::kK; + + constexpr auto ds_block_outer_dstr_encoding = + tile_distribution_encoding, + tuple, sequence>, + tuple>, + tuple>, + sequence<1, 2>, + sequence<0, 0>>{}; + + constexpr auto ds_block_dstr_encode = detail::make_embed_tile_distribution_encoding( + ds_block_outer_dstr_encoding, typename WarpGemm::AWarpDstrEncoding{}); + + constexpr auto ds_block_dstr = make_static_tile_distribution(ds_block_dstr_encode); + + return ds_block_dstr; + } + + template + CK_TILE_DEVICE static constexpr void PTFromGemm0CToGemm1A(PTOutTensor& pt_out, + const PInTensor& p_in) + { + if constexpr(Problem::BlockFmhaShape::Gemm1WarpTile::at(number<0>{}) == 16) + { + using BlockGemm = remove_cvref_t())>; + constexpr auto config = BlockGemm::Policy::template GetWarpGemmMWarpNWarp(); + using WarpGemm = remove_cvref_t())>; + + constexpr index_t MWarp = Problem::BlockFmhaShape::Gemm1BlockWarps::at(number<0>{}); + + constexpr index_t kMPerBlock = Problem::BlockFmhaShape::kN0; + constexpr index_t kKPerBlock = Problem::BlockFmhaShape::kK1; + + constexpr index_t MIterPerWarp = kMPerBlock / (MWarp * WarpGemm::kM); + constexpr index_t KIterPerWarp = kKPerBlock / WarpGemm::kK; + + using AWarpDstr = typename WarpGemm::AWarpDstr; + using CWarpDstr = typename WarpGemm::CWarpDstr; + auto pt_warp_tensor = + make_static_distributed_tensor(CWarpDstr{}); + + constexpr auto a_warp_y_lengths = + to_sequence(AWarpDstr{}.get_ys_to_d_descriptor().get_lengths()); + constexpr auto c_warp_y_lengths = + to_sequence(CWarpDstr{}.get_ys_to_d_descriptor().get_lengths()); + + constexpr auto a_warp_y_index_zeros = uniform_sequence_gen_t{}; + constexpr auto c_warp_y_index_zeros = uniform_sequence_gen_t{}; + + static_for<0, KIterPerWarp, 1>{}([&](auto kIter) { + static_for<0, MIterPerWarp, 1>{}([&](auto mIter) { + pt_warp_tensor.get_thread_buffer() = p_in.get_y_sliced_thread_data( + merge_sequences(sequence{}, c_warp_y_index_zeros), + merge_sequences(sequence<1, 1>{}, c_warp_y_lengths)); + + pt_out.set_y_sliced_thread_data( + merge_sequences(sequence{}, a_warp_y_index_zeros), + merge_sequences(sequence<1, 1>{}, a_warp_y_lengths), + pt_warp_tensor.get_thread_buffer()); + }); + }); + } + else + { + pt_out.get_thread_buffer() = p_in.get_thread_buffer(); + } + } + + template + CK_TILE_DEVICE static constexpr void SGradTFromGemm2CToGemm3A(SGradTOutTensor& dst_out, + const SGradInTensor& ds_in) + { + if constexpr(Problem::BlockFmhaShape::Gemm3WarpTile::at(number<0>{}) == 16) + { + using BlockGemm = remove_cvref_t())>; + constexpr auto config = BlockGemm::Policy::template GetWarpGemmMWarpNWarp(); + using WarpGemm = remove_cvref_t())>; + + constexpr index_t MWarp = Problem::BlockFmhaShape::Gemm3BlockWarps::at(number<0>{}); + + constexpr index_t kMPerBlock = Problem::BlockFmhaShape::kN0; + constexpr index_t kKPerBlock = Problem::BlockFmhaShape::kK3; + + constexpr index_t MIterPerWarp = kMPerBlock / (MWarp * WarpGemm::kM); + constexpr index_t KIterPerWarp = kKPerBlock / WarpGemm::kK; + + using AWarpDstr = typename WarpGemm::AWarpDstr; + using CWarpDstr = typename WarpGemm::CWarpDstr; + auto dst_warp_tensor = + make_static_distributed_tensor(CWarpDstr{}); + + constexpr auto a_warp_y_lengths = + to_sequence(AWarpDstr{}.get_ys_to_d_descriptor().get_lengths()); + constexpr auto c_warp_y_lengths = + to_sequence(CWarpDstr{}.get_ys_to_d_descriptor().get_lengths()); + + constexpr auto a_warp_y_index_zeros = uniform_sequence_gen_t{}; + constexpr auto c_warp_y_index_zeros = uniform_sequence_gen_t{}; + + static_for<0, KIterPerWarp, 1>{}([&](auto kIter) { + static_for<0, MIterPerWarp, 1>{}([&](auto mIter) { + dst_warp_tensor.get_thread_buffer() = ds_in.get_y_sliced_thread_data( + merge_sequences(sequence{}, c_warp_y_index_zeros), + merge_sequences(sequence<1, 1>{}, c_warp_y_lengths)); + + dst_out.set_y_sliced_thread_data( + merge_sequences(sequence{}, a_warp_y_index_zeros), + merge_sequences(sequence<1, 1>{}, a_warp_y_lengths), + dst_warp_tensor.get_thread_buffer()); + }); + }); + } + else + { + dst_out.get_thread_buffer() = ds_in.get_thread_buffer(); + } } template CK_TILE_HOST_DEVICE static constexpr auto MakeShuffledBiasTileDistribution() { constexpr index_t kBlockSize = Problem::kBlockSize; - constexpr index_t kMPerBlock = Problem::BlockFmhaShape::kM0; + constexpr index_t kNPerBlock = Problem::BlockFmhaShape::kN0; - constexpr index_t N1 = GetTransposedAlignmentBias(); - constexpr index_t N0 = kNPerBlock / N1; - constexpr index_t total_pixels = kMPerBlock * kNPerBlock / kBlockSize; - static_assert(total_pixels % N1 == 0); // TODO: this is not always true? - constexpr index_t M3 = total_pixels / N1; - constexpr index_t kKPack = GetSmemKPackBias(); - static_assert(kKPack % M3 == 0); - constexpr index_t M2 = kKPack / M3; // TODO: this dimention could be outside single wave - constexpr index_t M1 = get_warp_size() / (M2 * N0); + constexpr index_t N1 = GetAlignmentBias(); + constexpr index_t N0 = kNPerBlock / N1; + constexpr index_t M2 = GetTransposedAlignmentBias(); + constexpr index_t M1 = get_warp_size() / N0; constexpr index_t M0 = kBlockSize / get_warp_size(); return make_static_tile_distribution( tile_distribution_encoding, - tuple, sequence>, - tuple, sequence<1, 2, 1>>, - tuple, sequence<1, 0, 2>>, + tuple, sequence>, + tuple, sequence<1, 2>>, + tuple, sequence<1, 0>>, sequence<2, 1>, - sequence<1, 3>>{}); + sequence<1, 2>>{}); + } + + template + CK_TILE_HOST_DEVICE static constexpr auto MakeBiasLdsBlockDescriptor() + { + // Hold full block data + constexpr index_t kNPerBlock = Problem::BlockFmhaShape::kN0; + constexpr index_t kMPerBlock = Problem::BlockFmhaShape::kM0; + + constexpr index_t kKPack = GetSmemKPackBias(); + constexpr index_t kKPackT = GetSmemKPackBiasT(); + + return MakeXTLdsBlockDescriptor(); } template - CK_TILE_HOST_DEVICE static constexpr auto MakeBiasTTileDistribution() + CK_TILE_HOST_DEVICE static constexpr auto MakeBiasSTileDistribution() { using c_block_tensor_type = decltype(BlockGemm{}.MakeCBlockTile()); return c_block_tensor_type::get_tile_distribution(); } template - CK_TILE_HOST_DEVICE static constexpr auto GetQKBlockGemm() + CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSizeQ() { - using BlockGemmProblem = - BlockGemmPipelineProblem>; + constexpr index_t smem_size_q = sizeof(typename Problem::QDataType) * + MakeQLdsBlockDescriptor().get_element_space_size(); + return smem_size_q; + } - constexpr auto warp_gemm = []() { - if constexpr(std::is_same_v && - std::is_same_v && - std::is_same_v) - { - return WarpGemmMfmaF16F16F32M32N32K16SwizzleA{}; - } - else if constexpr(std::is_same_v && - std::is_same_v && - std::is_same_v) - { - return WarpGemmMfmaBf16Bf16F32M32N32K16SwizzleA{}; - } + template + CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSizeQT() + { + constexpr index_t smem_size_qt = + sizeof(typename Problem::QDataType) * + MakeShuffledQLdsWriteBlockDescriptor().get_element_space_size(); + + return smem_size_qt; + } + + template + CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSizeK() + { + constexpr index_t smem_size_k = + sizeof(typename Problem::KDataType) * + MakeKLdsWriteBlockDescriptor().get_element_space_size(); + return smem_size_k; + } + + template + CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSizeKT() + { + constexpr index_t smem_size_kt = + sizeof(typename Problem::KDataType) * + MakeKTLdsReadBlockDescriptor().get_element_space_size(); + return smem_size_kt; + } + + template + CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSizeLSE() + { + constexpr index_t smem_size_lse = + sizeof(typename Problem::LSEDataType) * + MakeLSEDLdsWriteBlockDescriptor().get_element_space_size(); + return smem_size_lse; + } + + template + CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSizeD() + { + constexpr index_t smem_size_d = + sizeof(typename Problem::DDataType) * + MakeLSEDLdsWriteBlockDescriptor().get_element_space_size(); + return smem_size_d; + } + + template + CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSizeV() + { + constexpr index_t smem_size_v = + sizeof(typename Problem::VDataType) * + MakeVLdsWriteBlockDescriptor().get_element_space_size(); + return smem_size_v; + } + + template + CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSizeOGrad() + { + constexpr index_t smem_size_do = + sizeof(typename Problem::OGradDataType) * + MakeOGradLdsBlockDescriptor().get_element_space_size(); + return smem_size_do; + } + + template + CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSizeOGradT() + { + constexpr index_t smem_size_dot = + sizeof(typename Problem::OGradDataType) * + MakeShuffledOGradLdsWriteBlockDescriptor().get_element_space_size(); + return smem_size_dot; + } + + template + CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSizeSGrad() + { + constexpr index_t smem_size_ds = + sizeof(typename Problem::GemmDataType) * + MakeSGradLdsBlockDescriptor().get_element_space_size(); + return smem_size_ds; + } + + template + CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSizeBias() + { + constexpr index_t smem_size_bias = [&]() { + if constexpr(Problem::BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS) + return sizeof(typename Problem::BiasDataType) * + MakeBiasLdsBlockDescriptor().get_element_space_size(); + else + return 0; }(); - - using BlockGemmPolicy = - BlockGemmASmemBSmemCRegV1CustomPolicy; - - return BlockGemmASmemBSmemCRegV1{}; + return smem_size_bias; } template - CK_TILE_HOST_DEVICE static constexpr auto GetPTOGradTBlockGemm() + CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSize() { - using BlockGemmProblem = - BlockGemmPipelineProblem>; + constexpr index_t smem_size_q = GetSmemSizeQ(); + constexpr index_t smem_size_qt = GetSmemSizeQT(); + constexpr index_t smem_size_lse = GetSmemSizeLSE(); + constexpr index_t smem_size_k = GetSmemSizeK(); + constexpr index_t smem_size_kt = GetSmemSizeKT(); + constexpr index_t smem_size_v = GetSmemSizeV(); + constexpr index_t smem_size_do = GetSmemSizeOGrad(); + constexpr index_t smem_size_dot = GetSmemSizeOGradT(); + constexpr index_t smem_size_d = GetSmemSizeD(); + constexpr index_t smem_size_ds = GetSmemSizeSGrad(); + constexpr index_t smem_size_bias = GetSmemSizeBias(); - using WarpGemm = - WarpGemmMfmaDispatcher{}), - Problem::BlockFmhaShape::Gemm1WarpTile::at(number<1>{}), - Problem::BlockFmhaShape::Gemm1WarpTile::at(number<2>{}), - true>; - using BlockGemmPolicy = - BlockGemmARegBSmemCRegV1CustomPolicy; - return BlockGemmARegBSmemCRegV1{}; + constexpr index_t smem_size_stage0_0 = smem_size_k + smem_size_kt; + constexpr index_t smem_size_stage0_1 = smem_size_v; + constexpr index_t smem_size_stage1 = smem_size_qt + smem_size_q + +smem_size_dot + + smem_size_do + smem_size_lse + smem_size_d + + max(smem_size_bias, smem_size_ds); + + return max(smem_size_stage0_0, smem_size_stage0_1, smem_size_stage1); } - template - CK_TILE_HOST_DEVICE static constexpr auto GetOGradVBlockGemm() + template + struct HotLoopScheduler { - using BlockGemmProblem = - BlockGemmPipelineProblem>; + using Problem = Problem_; - constexpr auto warp_gemm = []() { - if constexpr(std::is_same_v && - std::is_same_v && - std::is_same_v) - { - return WarpGemmMfmaF16F16F32M32N32K16SwizzleA{}; - } - else if constexpr(std::is_same_v && - std::is_same_v && - std::is_same_v) - { - return WarpGemmMfmaBf16Bf16F32M32N32K16SwizzleA{}; - } - }(); + template + CK_TILE_DEVICE static constexpr void GemmStagedScheduler() + { + } - using BlockGemmPolicy = - BlockGemmASmemBRegCRegV1CustomPolicy; + template <> + CK_TILE_DEVICE static constexpr void GemmStagedScheduler<0>() + { + // Mem: Q, LSE, OGrad, D global load, OGrad^T LDS load + // Comp: Q x K + constexpr index_t VMEM_READ_INST = + Q_VMEM_READ + OGrad_VMEM_READ + LSE_VMEM_READ + D_VMEM_READ; + constexpr index_t LDS_READ_INST = OGradT_LDS_READ; + constexpr index_t MFMA_INST = Gemm0MFMA; - return BlockGemmASmemBRegCRegV1{}; - } + // Evenly distributed to relieve SQ->TA FIFO pressure + constexpr index_t MFMA_PER_VMEM_READ = MFMA_INST / VMEM_READ_INST; + constexpr index_t MFMA_Remainder = MFMA_INST - MFMA_PER_VMEM_READ * VMEM_READ_INST; + // To hide instruction issue latency + constexpr index_t LDS_READ_PER_MFMA = LDS_READ_INST / MFMA_INST; - // template - // CK_TILE_HOST_DEVICE static constexpr auto GetOGradVBlockGemm() - // { - // using BlockGemmProblem = - // BlockGemmPipelineProblem>; - // constexpr auto warp_gemm = []() { - // if constexpr(std::is_same_v && - // std::is_same_v && - // std::is_same_v) - // { - // return WarpGemmMfmaF16F16F32M32N32K16SwizzleA{}; - // } - // else if constexpr(std::is_same_v && - // std::is_same_v && - // std::is_same_v) - // { - // return WarpGemmMfmaBf16Bf16F32M32N32K16SwizzleA{}; - // } - // }(); + static_for<0, VMEM_READ_INST, 1>{}([&](auto i) { + ignore = i; + __builtin_amdgcn_sched_group_barrier(0x020, 1, 0); // VMEM read + static_for<0, MFMA_PER_VMEM_READ, 1>{}([&](auto j) { + ignore = j; + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA + __builtin_amdgcn_sched_group_barrier(0x100, LDS_READ_PER_MFMA, 0); // DS read + }); + }); + static_for<0, MFMA_Remainder, 1>{}([&](auto i) { + ignore = i; + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA + __builtin_amdgcn_sched_group_barrier(0x100, LDS_READ_PER_MFMA, 0); // DS read + }); + } - // using BlockGemmPolicy = - // BlockGemmASmemBSmemCRegV1CustomPolicy; + template <> + CK_TILE_DEVICE static constexpr void GemmStagedScheduler<1>() + { + // Mem: Q^T LDS load + // Comp: OGrad x V + constexpr index_t LDS_READ_INST = QT_LDS_READ; + constexpr index_t MFMA_INST = Gemm1MFMA; - // return BlockGemmASmemBSmemCRegV1{}; - // } + // To hide instruction issue latency + constexpr index_t LDS_READ_PER_MFMA = LDS_READ_INST / MFMA_INST; - template - CK_TILE_HOST_DEVICE static constexpr auto GetSGradTQTBlockGemm() - { - using BlockGemmProblem = - BlockGemmPipelineProblem>; + static_for<0, MFMA_INST, 1>{}([&](auto i) { + ignore = i; + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA + __builtin_amdgcn_sched_group_barrier(0x100, LDS_READ_PER_MFMA, 0); // DS read + }); + } - using WarpGemm = - WarpGemmMfmaDispatcher{}), - Problem::BlockFmhaShape::Gemm3WarpTile::at(number<1>{}), - Problem::BlockFmhaShape::Gemm3WarpTile::at(number<2>{}), - true>; - using BlockGemmPolicy = - BlockGemmARegBSmemCRegV1CustomPolicy; - return BlockGemmARegBSmemCRegV1{}; - } + template <> + CK_TILE_DEVICE static constexpr void GemmStagedScheduler<2>() + { + // Mem: Q, QT, LSE, OGrad, OGradT, D, LDS store + // Comp: PT x OGrad + constexpr index_t LDS_WRITE_INST = Q_LDS_WRITE + QT_LDS_WRITE + OGrad_LDS_WRITE + + OGradT_LDS_WRITE + LSE_LDS_WRITE + D_LDS_WRITE; + constexpr index_t MFMA_INST = Gemm2MFMA; - template - CK_TILE_HOST_DEVICE static constexpr auto GetSGradKTBlockGemm() - { - using BlockGemmProblem = - BlockGemmPipelineProblem>; + // To hide instruction issue latency + constexpr index_t LDS_WRITE_PER_MFMA = LDS_WRITE_INST / MFMA_INST; - using WarpGemm = - WarpGemmMfmaDispatcher{}), - Problem::BlockFmhaShape::Gemm4WarpTile::at(number<1>{}), - Problem::BlockFmhaShape::Gemm4WarpTile::at(number<2>{}), - true>; - using BlockGemmPolicy = - BlockGemmASmemBSmemCRegV1CustomPolicy; - return BlockGemmASmemBSmemCRegV1{}; - } + static_for<0, MFMA_INST, 1>{}([&](auto i) { + ignore = i; + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA + __builtin_amdgcn_sched_group_barrier(0x200, LDS_WRITE_PER_MFMA, 0); // DS write + }); + } + + template <> + CK_TILE_DEVICE static constexpr void GemmStagedScheduler<3>() + { + // Mem: SGradT LDS store, SGrad, Q, LSE LDS load. + // Comp: SGradT x QT + constexpr index_t LDS_WRITE_INST = SGradT_LDS_WRITE; + constexpr index_t LDS_READ_INST = SGradT_LDS_READ_P1 + Q_LDS_READ + LSE_LDS_READ; + constexpr index_t MFMA_INST = Gemm3MFMA; + + // To hide instruction issue latency + constexpr index_t LDS_WRITE_PER_MFMA = + LDS_WRITE_INST / MFMA_INST >= 1 ? LDS_WRITE_INST / MFMA_INST : 1; + constexpr index_t MFMA_INST_LDS_WRITE = LDS_WRITE_INST / LDS_WRITE_PER_MFMA; + + constexpr index_t LDS_READ_PER_MFMA = + (MFMA_INST - MFMA_INST_LDS_WRITE) > 0 + ? LDS_READ_INST / (MFMA_INST - MFMA_INST_LDS_WRITE) > 0 + ? LDS_READ_INST / (MFMA_INST - MFMA_INST_LDS_WRITE) + : 1 + : 0; + + static_for<0, MFMA_INST_LDS_WRITE, 1>{}([&](auto i) { + ignore = i; + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA + __builtin_amdgcn_sched_group_barrier(0x200, LDS_WRITE_PER_MFMA, 0); // DS Write + }); + + static_for<0, MFMA_INST - MFMA_INST_LDS_WRITE, 1>{}([&](auto i) { + ignore = i; + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA + __builtin_amdgcn_sched_group_barrier(0x100, LDS_READ_PER_MFMA, 0); // DS Read + }); + } + + template <> + CK_TILE_DEVICE static constexpr void GemmStagedScheduler<4>() + { + // Mem: SGrad, OGrad, D LDS load. + // Comp: SGrad x KT + constexpr index_t LDS_READ_INST = SGradT_LDS_READ_P2 + OGrad_LDS_READ + D_LDS_READ; + constexpr index_t MFMA_INST = Gemm4MFMA; + + // To hide instruction issue latency + constexpr index_t LDS_READ_PER_MFMA = + LDS_READ_INST / MFMA_INST > 0 ? LDS_READ_INST / MFMA_INST : 1; + + static_for<0, MFMA_INST, 1>{}([&](auto i) { + ignore = i; + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA + __builtin_amdgcn_sched_group_barrier(0x100, LDS_READ_PER_MFMA, 0); // DS Read + }); + } + + private: + static constexpr index_t kBlockSize = Problem::kBlockSize; + static constexpr index_t kM0 = Problem::BlockFmhaShape::kM0; + static constexpr index_t kN0 = Problem::BlockFmhaShape::kN0; + static constexpr index_t kQKHeaddim = Problem::BlockFmhaShape::kQKHeaddim; + static constexpr index_t kVHeaddim = Problem::BlockFmhaShape::kVHeaddim; + static constexpr index_t kK4 = Problem::BlockFmhaShape::kK4; + + static constexpr index_t WarpGemmM = + Problem::BlockFmhaShape::Gemm0WarpTile::at(number<0>{}); + static constexpr index_t WarpGemmN = + Problem::BlockFmhaShape::Gemm0WarpTile::at(number<1>{}); + static constexpr index_t WarpGemmK = WarpGemmM == 16 ? 16 : 8; + static constexpr index_t Gemm4MWarp = + Problem::BlockFmhaShape::Gemm4BlockWarps::at(number<0>{}); + static constexpr index_t Gemm4NWarp = + Problem::BlockFmhaShape::Gemm4BlockWarps::at(number<1>{}); + + // Compute + static constexpr index_t Gemm0MFMA = + kM0 * kN0 * kQKHeaddim / + (kBlockSize / get_warp_size() * WarpGemmM * WarpGemmN * WarpGemmK); + static constexpr index_t Gemm1MFMA = + kM0 * kN0 * kVHeaddim / + (kBlockSize / get_warp_size() * WarpGemmM * WarpGemmN * WarpGemmK); + static constexpr index_t Gemm2MFMA = + kN0 * kVHeaddim * kM0 / + (kBlockSize / get_warp_size() * WarpGemmM * WarpGemmN * WarpGemmK); + static constexpr index_t Gemm3MFMA = + kN0 * kQKHeaddim * kM0 / + (kBlockSize / get_warp_size() * WarpGemmM * WarpGemmN * WarpGemmK); + static constexpr index_t Gemm4MFMA = + kM0 * kQKHeaddim * kN0 / + (kBlockSize / get_warp_size() * WarpGemmM * WarpGemmN * WarpGemmK); + + // VMEM + static constexpr index_t Q_VMEM_READ = + kM0 * kQKHeaddim / kBlockSize / GetAlignmentQ(); + static constexpr index_t OGrad_VMEM_READ = + kM0 * kVHeaddim / kBlockSize / GetAlignmentOGrad(); + static constexpr index_t LSE_VMEM_READ = 1; + static constexpr index_t D_VMEM_READ = 1; + + // LDS Read + static constexpr index_t OGradT_LDS_READ = + kM0 * kVHeaddim / get_warp_size() / GetTransposedAlignmentOGrad(); + static constexpr index_t QT_LDS_READ = + kM0 * kQKHeaddim / get_warp_size() / GetTransposedAlignmentQ(); + static constexpr index_t SGradT_LDS_READ_P1 = + kM0 * kK4 / (get_warp_size() * Gemm4MWarp) / GetSmemKPackSGrad(); + static constexpr index_t Q_LDS_READ = + kM0 * kQKHeaddim / kBlockSize / GetAlignmentQ(); + static constexpr index_t LSE_LDS_READ = WarpGemmM == 16 ? kM0 / (4 * 4) : kM0 / (2 * 4); + static constexpr index_t SGradT_LDS_READ_P2 = + kM0 * (kN0 - kK4) / (get_warp_size() * Gemm4MWarp) / GetSmemKPackSGrad(); + static constexpr index_t OGrad_LDS_READ = + kM0 * kVHeaddim / kBlockSize / GetAlignmentOGrad(); + static constexpr index_t D_LDS_READ = WarpGemmM == 16 ? kM0 / (4 * 4) : kM0 / (2 * 4); + + // LDS Write + static constexpr index_t Q_LDS_WRITE = + kM0 * kQKHeaddim / Problem::kBlockSize / GetAlignmentQ(); + static constexpr index_t QT_LDS_WRITE = + kM0 * kQKHeaddim / kBlockSize / GetTransposedAlignmentQ(); + static constexpr index_t OGrad_LDS_WRITE = + kM0 * kVHeaddim / kBlockSize / GetAlignmentOGrad(); + static constexpr index_t OGradT_LDS_WRITE = + kM0 * kVHeaddim / kBlockSize / GetTransposedAlignmentOGrad(); + static constexpr index_t LSE_LDS_WRITE = 1; + static constexpr index_t D_LDS_WRITE = 1; + static constexpr index_t SGradT_LDS_WRITE = kM0 * kN0 / kBlockSize; + }; }; } // namespace ck_tile diff --git a/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_pipeline_enum.hpp b/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_pipeline_enum.hpp index a54a9fcb32..27f58ef2f8 100644 --- a/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_pipeline_enum.hpp +++ b/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_pipeline_enum.hpp @@ -8,9 +8,8 @@ namespace ck_tile { // This class is used for codegen pattern matching enum class BlockFmhaBwdPipelineEnum { - KSKTSVR = 0, - QSKSVROGradS, - KSVR, + KRKTRVR_IGLP = 0, + KRKTRVR, }; } // namespace ck_tile diff --git a/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_pipeline_problem.hpp b/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_pipeline_problem.hpp index 7b787e9f36..c4c4a745a7 100644 --- a/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_pipeline_problem.hpp +++ b/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_pipeline_problem.hpp @@ -24,7 +24,9 @@ template struct BlockFmhaBwdPipelineProblem { @@ -45,10 +47,12 @@ struct BlockFmhaBwdPipelineProblem using BiasGradDataType = remove_cvref_t; using BlockFmhaShape = remove_cvref_t; using FmhaMask = remove_cvref_t; + using FmhaDropout = remove_cvref_t; using Traits = remove_cvref_t; - static constexpr index_t kBlockSize = BlockFmhaShape::NumWarps * get_warp_size(); - static constexpr bool kIsGroupMode = kIsGroupMode_; + static constexpr index_t kBlockSize = BlockFmhaShape::NumWarps * get_warp_size(); + static constexpr bool kIsGroupMode = kIsGroupMode_; + static constexpr bool kIsDeterministic = kIsDeterministic_; // attributes from traits static constexpr bool kPadSeqLenQ = Traits::kPadSeqLenQ; @@ -57,7 +61,6 @@ struct BlockFmhaBwdPipelineProblem static constexpr bool kPadHeadDimV = Traits::kPadHeadDimV; static constexpr auto BiasEnum = Traits::BiasEnum; static constexpr bool kHasBiasGrad = Traits::kHasBiasGrad; - static constexpr bool kHasDropout = Traits::kHasDropout; static constexpr index_t kBlockPerCu = Traits::kBlockPerCu; }; @@ -88,4 +91,35 @@ struct BlockFmhaBwdOGradDotOPipelineProblem static constexpr index_t kBlockPerCu = Traits::kBlockPerCu; }; +template +struct BlockFmhaBwdConvertQGradPipelineProblem +{ + using AccDataType = remove_cvref_t; + using QGradDataType = remove_cvref_t; + using Traits = remove_cvref_t; + + static_assert(0 < kBlockSize_ && kBlockSize_ % get_warp_size() == 0, + "kBlockSize should be divisible by get_warp_size()"); + + static constexpr index_t kBlockSize = kBlockSize_; + static constexpr index_t kM0 = kM0_; + static constexpr index_t kN0 = kN0_; + static constexpr index_t kQKHeaddim = kQKHeaddim_; + static constexpr bool kIsGroupMode = kIsGroupMode_; + static constexpr bool kIsDeterministic = kIsDeterministic_; + + // attributes from traits + static constexpr bool kPadSeqLenQ = Traits::kPadSeqLenQ; + static constexpr bool kPadHeadDimQ = Traits::kPadHeadDimQ; + static constexpr index_t kBlockPerCu = Traits::kBlockPerCu; +}; + } // namespace ck_tile diff --git a/include/ck_tile/ops/fmha/pipeline/tile_fmha_traits.hpp b/include/ck_tile/ops/fmha/pipeline/tile_fmha_traits.hpp index a59431e39d..be4fdfd711 100644 --- a/include/ck_tile/ops/fmha/pipeline/tile_fmha_traits.hpp +++ b/include/ck_tile/ops/fmha/pipeline/tile_fmha_traits.hpp @@ -86,4 +86,14 @@ struct TileFmhaBwdOGradDotOTraits static constexpr index_t kBlockPerCu = kBlockPerCu_; }; +template +struct TileFmhaBwdConvertQGradTraits +{ + static constexpr bool kPadSeqLenQ = kPadSeqLenQ_; + static constexpr bool kPadHeadDimQ = kPadHeadDimQ_; + static constexpr index_t kBlockPerCu = kBlockPerCu_; +}; + } // namespace ck_tile diff --git a/include/ck_tile/ops/gemm.hpp b/include/ck_tile/ops/gemm.hpp index a89536e6eb..dd313c5480 100644 --- a/include/ck_tile/ops/gemm.hpp +++ b/include/ck_tile/ops/gemm.hpp @@ -5,6 +5,9 @@ #include "ck_tile/ops/gemm/block/block_gemm_areg_bgmem_creg_v1.hpp" #include "ck_tile/ops/gemm/block/block_gemm_areg_bgmem_creg_v1_default_policy.hpp" +#include "ck_tile/ops/gemm/block/block_gemm_areg_breg_creg_v1.hpp" +#include "ck_tile/ops/gemm/block/block_gemm_areg_breg_creg_v1_custom_policy.hpp" +#include "ck_tile/ops/gemm/block/block_gemm_areg_breg_creg_v1_default_policy.hpp" #include "ck_tile/ops/gemm/block/block_gemm_areg_bsmem_creg_v1.hpp" #include "ck_tile/ops/gemm/block/block_gemm_areg_bsmem_creg_v1_custom_policy.hpp" #include "ck_tile/ops/gemm/block/block_gemm_areg_bsmem_creg_v1_default_policy.hpp" diff --git a/include/ck_tile/ops/gemm/block/block_gemm_areg_breg_creg_v1.hpp b/include/ck_tile/ops/gemm/block/block_gemm_areg_breg_creg_v1.hpp new file mode 100644 index 0000000000..9a5c2aae5c --- /dev/null +++ b/include/ck_tile/ops/gemm/block/block_gemm_areg_breg_creg_v1.hpp @@ -0,0 +1,202 @@ +// 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/block/block_gemm_areg_breg_creg_v1_default_policy.hpp" + +namespace ck_tile { + +// A is block distributed tensor +// B is block distributed tensor +// C is block distributed tensor +template +struct BlockGemmARegBRegCRegV1 +{ + using Problem = remove_cvref_t; + using Policy = remove_cvref_t; + using ADataType = remove_cvref_t; + using BDataType = remove_cvref_t; + using CDataType = remove_cvref_t; + using BlockGemmShape = remove_cvref_t; + + static constexpr index_t kBlockSize = Problem::kBlockSize; + + // C += A * B + template + CK_TILE_DEVICE void operator()(CBlockTensor& c_block_tensor, + const ABlockTensor& a_block_tensor, + const BBlockTensor& b_block_tensor) const + { + static_assert(std::is_same_v> && + std::is_same_v> && + std::is_same_v>, + "wrong!"); + + constexpr index_t MPerBlock = BlockGemmShape::kM; + constexpr index_t NPerBlock = BlockGemmShape::kN; + constexpr index_t KPerBlock = BlockGemmShape::kK; + + 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 NIterPerWarp = NPerBlock / (NWarp * WG::kN); + constexpr index_t KIterPerWarp = KPerBlock / WG::kK; + + // M->N Warp + constexpr auto a_block_outer_dstr_encoding = + tile_distribution_encoding, + tuple, sequence>, + tuple>, + tuple>, + sequence<1, 2>, + sequence<0, 0>>{}; + + constexpr auto b_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>, + tuple>, + tuple>, + sequence<1, 2>, + sequence<0, 0>>{}; + + constexpr auto a_block_dstr_encode = detail::make_embed_tile_distribution_encoding( + a_block_outer_dstr_encoding, typename WG::AWarpDstrEncoding{}); + + constexpr auto b_block_dstr_encode = detail::make_embed_tile_distribution_encoding( + b_block_outer_dstr_encoding, typename WG::BWarpDstrEncoding{}); + + constexpr auto c_block_dstr_encode = detail::make_embed_tile_distribution_encoding( + c_block_outer_dstr_encoding, typename WG::CWarpDstrEncoding{}); + + // check ABC-block-distribution + static_assert( + std::is_same_v, + remove_cvref_t>, + "A distribution is wrong!"); + static_assert( + std::is_same_v, + remove_cvref_t>, + "B distribution is wrong!"); + static_assert( + std::is_same_v, + remove_cvref_t>, + "C distribution is wrong!"); + + using AWarpDstr = typename WG::AWarpDstr; + using BWarpDstr = typename WG::BWarpDstr; + using CWarpDstr = typename WG::CWarpDstr; + + using AWarpTensor = typename WG::AWarpTensor; + using BWarpTensor = typename WG::BWarpTensor; + using CWarpTensor = typename WG::CWarpTensor; + + constexpr auto a_warp_y_lengths = + to_sequence(AWarpDstr{}.get_ys_to_d_descriptor().get_lengths()); + constexpr auto b_warp_y_lengths = + to_sequence(BWarpDstr{}.get_ys_to_d_descriptor().get_lengths()); + constexpr auto c_warp_y_lengths = + to_sequence(CWarpDstr{}.get_ys_to_d_descriptor().get_lengths()); + + constexpr auto a_warp_y_index_zeros = uniform_sequence_gen_t{}; + constexpr auto b_warp_y_index_zeros = uniform_sequence_gen_t{}; + constexpr auto c_warp_y_index_zeros = uniform_sequence_gen_t{}; + + // hot loop: + static_for<0, KIterPerWarp, 1>{}([&](auto kIter) { + static_for<0, MIterPerWarp, 1>{}([&](auto mIter) { + // read A warp tensor from A Block window + AWarpTensor a_warp_tensor; + + a_warp_tensor.get_thread_buffer() = a_block_tensor.get_y_sliced_thread_data( + merge_sequences(sequence{}, a_warp_y_index_zeros), + merge_sequences(sequence<1, 1>{}, a_warp_y_lengths)); + + static_for<0, NIterPerWarp, 1>{}([&](auto nIter) { + // read B warp tensor from B block tensor + BWarpTensor b_warp_tensor; + + b_warp_tensor.get_thread_buffer() = b_block_tensor.get_y_sliced_thread_data( + merge_sequences(sequence{}, b_warp_y_index_zeros), + merge_sequences(sequence<1, 1>{}, b_warp_y_lengths)); + + // read C warp tensor from C block tensor + CWarpTensor c_warp_tensor; + + c_warp_tensor.get_thread_buffer() = c_block_tensor.get_y_sliced_thread_data( + merge_sequences(sequence{}, c_warp_y_index_zeros), + merge_sequences(sequence<1, 1>{}, c_warp_y_lengths)); + + // warp GEMM + WG{}(c_warp_tensor, a_warp_tensor, b_warp_tensor); + + // write C warp tensor into C block tensor + c_block_tensor.set_y_sliced_thread_data( + merge_sequences(sequence{}, c_warp_y_index_zeros), + merge_sequences(sequence<1, 1>{}, c_warp_y_lengths), + c_warp_tensor.get_thread_buffer()); + }); + }); + }); + } + + CK_TILE_DEVICE constexpr auto MakeCBlockTile() const + { + constexpr index_t MPerBlock = BlockGemmShape::kM; + constexpr index_t NPerBlock = BlockGemmShape::kN; + + 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 NIterPerWarp = NPerBlock / (NWarp * WG::kN); + // constexpr index_t KIterPerWarp = KPerBlock / WG::kK; + + constexpr auto c_block_outer_dstr_encoding = tile_distribution_encoding< + sequence<>, + tuple, sequence>, + tuple>, + tuple>, + sequence<1, 2>, + sequence<0, 0>>{}; + + constexpr auto c_block_dstr_encode = detail::make_embed_tile_distribution_encoding( + c_block_outer_dstr_encoding, typename WG::CWarpDstrEncoding{}); + constexpr auto c_block_dstr = make_static_tile_distribution(c_block_dstr_encode); + auto c_block_tensor = make_static_distributed_tensor(c_block_dstr); + return c_block_tensor; + } + + // C = A * B + template + CK_TILE_DEVICE auto operator()(const ABlockTensor& a_block_tensor, + const BBlockTensor& b_block_tensor) const + { + auto c_block_tensor = MakeCBlockTile(); + operator()(c_block_tensor, a_block_tensor, b_block_tensor); + return c_block_tensor; + } +}; + +} // namespace ck_tile diff --git a/include/ck_tile/ops/gemm/block/block_gemm_areg_breg_creg_v1_custom_policy.hpp b/include/ck_tile/ops/gemm/block/block_gemm_areg_breg_creg_v1_custom_policy.hpp new file mode 100644 index 0000000000..9d494c2831 --- /dev/null +++ b/include/ck_tile/ops/gemm/block/block_gemm_areg_breg_creg_v1_custom_policy.hpp @@ -0,0 +1,36 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved. + +#pragma once + +#include "ck_tile/core.hpp" + +namespace ck_tile { + +template +struct BlockGemmARegBRegCRegV1CustomPolicy +{ + using AType = remove_cvref_t; + using BType = remove_cvref_t; + using CType = remove_cvref_t; + + using BlockWarps = remove_cvref_t; + + static constexpr index_t kMWarps = BlockWarps::at(number<0>{}); + static constexpr index_t kNWarps = BlockWarps::at(number<1>{}); + static constexpr index_t kKWarps = BlockWarps::at(number<2>{}); + + using WarpGemm = remove_cvref_t; + + template + CK_TILE_HOST_DEVICE static constexpr auto GetWarpGemmMWarpNWarp() + { + return make_tuple(WarpGemm{}, kMWarps, kNWarps); + } +}; + +} // namespace ck_tile diff --git a/include/ck_tile/ops/gemm/block/block_gemm_areg_breg_creg_v1_default_policy.hpp b/include/ck_tile/ops/gemm/block/block_gemm_areg_breg_creg_v1_default_policy.hpp new file mode 100644 index 0000000000..b849c48daf --- /dev/null +++ b/include/ck_tile/ops/gemm/block/block_gemm_areg_breg_creg_v1_default_policy.hpp @@ -0,0 +1,33 @@ +// 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" + +namespace ck_tile { + +// Default policy for BlockGemmARegBRegCRegV1 +// Default policy class should not be templated, put template on member functions instead +struct BlockGemmARegBRegCRegV1DefaultPolicy +{ + template + CK_TILE_HOST_DEVICE static constexpr auto GetWarpGemmMWarpNWarp() + { + if constexpr(std::is_same_v && + std::is_same_v && + std::is_same_v) + { + return make_tuple(WarpGemmMfmaF16F16F32M32N32K8TransposedCDistribution{}, 4, 1); + } + else if constexpr(std::is_same_v && + std::is_same_v && + std::is_same_v) + { + return make_tuple(WarpGemmMfmaBf16Bf16F32M32N32K8TransposedCDistribution{}, 4, 1); + } + } +}; + +} // namespace ck_tile diff --git a/include/ck_tile/ops/gemm/block/block_gemm_areg_bsmem_creg_v1.hpp b/include/ck_tile/ops/gemm/block/block_gemm_areg_bsmem_creg_v1.hpp index 84883d6ed8..beab457b90 100644 --- a/include/ck_tile/ops/gemm/block/block_gemm_areg_bsmem_creg_v1.hpp +++ b/include/ck_tile/ops/gemm/block/block_gemm_areg_bsmem_creg_v1.hpp @@ -35,16 +35,13 @@ struct BlockGemmARegBSmemCRegV1 std::is_same_v>, "wrong!"); - // constexpr index_t MPerBlock = ABlockTensorTmp{}.get_lengths()[number<0>{}]; - // constexpr index_t NPerBlock = BBlockWindowTmp{}.get_window_lengths()[number<0>{}]; - // constexpr index_t KPerBlock = ABlockTensorTmp{}.get_lengths()[number<1>{}]; - constexpr index_t MPerBlock = BlockGemmShape::kM; - constexpr index_t NPerBlock = BlockGemmShape::kN; - constexpr index_t KPerBlock = BlockGemmShape::kK; + constexpr index_t MPerBlock = ABlockTensorTmp{}.get_lengths()[number<0>{}]; + constexpr index_t NPerBlock = BBlockWindowTmp{}.get_window_lengths()[number<0>{}]; + constexpr index_t KPerBlock = ABlockTensorTmp{}.get_lengths()[number<1>{}]; - // static_assert(MPerBlock == BlockGemmShape::kM && NPerBlock == BlockGemmShape::kN && - // KPerBlock == BlockGemmShape::kK, - // "wrong!"); + static_assert(MPerBlock == BlockGemmShape::kM && NPerBlock == BlockGemmShape::kN && + KPerBlock == BlockGemmShape::kK, + "wrong!"); constexpr auto config = Policy::template GetWarpGemmMWarpNWarp(); diff --git a/include/ck_tile/ops/gemm/block/block_gemm_asmem_breg_creg_v1.hpp b/include/ck_tile/ops/gemm/block/block_gemm_asmem_breg_creg_v1.hpp index 65ce1a9b8f..3d142df4d4 100644 --- a/include/ck_tile/ops/gemm/block/block_gemm_asmem_breg_creg_v1.hpp +++ b/include/ck_tile/ops/gemm/block/block_gemm_asmem_breg_creg_v1.hpp @@ -35,16 +35,13 @@ struct BlockGemmASmemBRegCRegV1 std::is_same_v>, "wrong!"); - // constexpr index_t MPerBlock = ABlockWindowTmp{}.get_window_lengths()[number<0>{}]; - // constexpr index_t NPerBlock = BBlockTensorTmp{}.get_lengths()[number<0>{}]; - // constexpr index_t KPerBlock = ABlockWindowTmp{}.get_window_lengths()[number<1>{}]; - constexpr index_t MPerBlock = BlockGemmShape::kM; - constexpr index_t NPerBlock = BlockGemmShape::kN; - constexpr index_t KPerBlock = BlockGemmShape::kK; + constexpr index_t MPerBlock = ABlockWindowTmp{}.get_window_lengths()[number<0>{}]; + constexpr index_t NPerBlock = BBlockTensorTmp{}.get_lengths()[number<0>{}]; + constexpr index_t KPerBlock = ABlockWindowTmp{}.get_window_lengths()[number<1>{}]; - // static_assert(MPerBlock == BlockGemmShape::kM && NPerBlock == BlockGemmShape::kN && - // KPerBlock == BlockGemmShape::kK, - // "wrong!"); + static_assert(MPerBlock == BlockGemmShape::kM && NPerBlock == BlockGemmShape::kN && + KPerBlock == BlockGemmShape::kK, + "wrong!"); constexpr auto config = Policy::template GetWarpGemmMWarpNWarp(); diff --git a/include/ck_tile/ops/gemm/warp/warp_gemm.hpp b/include/ck_tile/ops/gemm/warp/warp_gemm.hpp index 5b4419b79f..7ca4a697a7 100644 --- a/include/ck_tile/ops/gemm/warp/warp_gemm.hpp +++ b/include/ck_tile/ops/gemm/warp/warp_gemm.hpp @@ -22,6 +22,9 @@ using WarpGemmMfmaF16F16F32M32N32K16 = using WarpGemmMfmaF16F16F32M16N16K32 = WarpGemmImpl>; +using WarpGemmMfmaF16F16F32M32N32K8SwizzleA = WarpGemmImpl< + WarpGemmAtrributeMfmaIterateK_SwizzleA>; + using WarpGemmMfmaF16F16F32M32N32K16SwizzleA = WarpGemmImpl< WarpGemmAtrributeMfmaIterateK_SwizzleA>; @@ -59,6 +62,9 @@ using WarpGemmMfmaBf16Bf16F32M32N32K16 = using WarpGemmMfmaBf16Bf16F32M16N16K32 = WarpGemmImpl>; +using WarpGemmMfmaBf16Bf16F32M32N32K8SwizzleA = WarpGemmImpl< + WarpGemmAtrributeMfmaIterateK_SwizzleA>; + using WarpGemmMfmaBf16Bf16F32M32N32K16SwizzleA = WarpGemmImpl< WarpGemmAtrributeMfmaIterateK_SwizzleA>; 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 fd5b004d36..d80e5198e6 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 @@ -119,9 +119,9 @@ struct WarpGemmAtrributeMfmaIterateK static_for<0, kKIter, 1>{}([&](auto iKIter) { Impl{}(c_vec, - reinterpret_cast(a_vec) + reinterpret_cast(a_vec) .template get_as()[iKIter], - reinterpret_cast(b_vec) + reinterpret_cast(b_vec) .template get_as()[iKIter]); }); } @@ -135,15 +135,15 @@ struct WarpGemmAtrributeMfmaIterateK // c = a * b auto c_vec = Impl{}( - reinterpret_cast(a_vec).template get_as()[I0], - reinterpret_cast(b_vec).template get_as()[I0]); + reinterpret_cast(a_vec).template get_as()[I0], + reinterpret_cast(b_vec).template get_as()[I0]); // c += a * b static_for<1, kKIter, 1>{}([&](auto iKIter) { Impl{}(c_vec, - reinterpret_cast(a_vec) + reinterpret_cast(a_vec) .template get_as()[iKIter], - reinterpret_cast(b_vec) + reinterpret_cast(b_vec) .template get_as()[iKIter]); }); 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 8d12130308..99cd5d787e 100644 --- a/include/ck_tile/ops/gemm/warp/warp_gemm_dispatcher.hpp +++ b/include/ck_tile/ops/gemm/warp/warp_gemm_dispatcher.hpp @@ -15,7 +15,8 @@ template + bool TransposeC, + bool SwizzleA = false> struct WarpGemmMfmaDispatcher; // clang-format off @@ -29,6 +30,9 @@ template<> struct WarpGemmMfmaDispatcher struct WarpGemmMfmaDispatcher { using Type = WarpGemmMfmaF16F16F32M16N16K32; }; template<> struct WarpGemmMfmaDispatcher { using Type = WarpGemmMfmaF16F16F32M16N16K32TransposedCDistribution; }; +template<> struct WarpGemmMfmaDispatcher { using Type = WarpGemmMfmaF16F16F32M32N32K8SwizzleA; }; +template<> struct WarpGemmMfmaDispatcher { using Type = WarpGemmMfmaF16F16F32M32N32K16SwizzleA; }; + // bf16 template<> struct WarpGemmMfmaDispatcher { using Type = WarpGemmMfmaBf16Bf16F32M32N32K8; }; template<> struct WarpGemmMfmaDispatcher { using Type = WarpGemmMfmaBf16Bf16F32M32N32K8TransposedCDistribution; }; @@ -39,6 +43,9 @@ template<> struct WarpGemmMfmaDispatcher struct WarpGemmMfmaDispatcher { using Type = WarpGemmMfmaBf16Bf16F32M16N16K32; }; template<> struct WarpGemmMfmaDispatcher { using Type = WarpGemmMfmaBf16Bf16F32M16N16K32TransposedCDistribution; }; +template<> struct WarpGemmMfmaDispatcher { using Type = WarpGemmMfmaBf16Bf16F32M32N32K8SwizzleA; }; +template<> struct WarpGemmMfmaDispatcher { using Type = WarpGemmMfmaBf16Bf16F32M32N32K16SwizzleA; }; + // fp8 template<> struct WarpGemmMfmaDispatcher { using Type = WarpGemmMfma_f32_32x32x16_fp8_fp8; }; template<> struct WarpGemmMfmaDispatcher { using Type = WarpGemmMfma_f32_32x32x16_fp8_fp8_CTransposed; }; @@ -58,8 +65,15 @@ template -using WarpGemmMfmaDispatcher = typename impl:: - WarpGemmMfmaDispatcher::Type; + bool TransposeC, + bool SwizzleA = false> +using WarpGemmMfmaDispatcher = typename impl::WarpGemmMfmaDispatcher::Type; } // namespace ck_tile From c8b6b64240e840a7decf76dfaa13c37da5294c4a Mon Sep 17 00:00:00 2001 From: Illia Silin <98187287+illsilin@users.noreply.github.com> Date: Fri, 16 Aug 2024 15:07:52 -0700 Subject: [PATCH 08/20] Re-enable fp8 types for all architectures. (#1470) * re-enable fp8 and bf8 for all targets * restore the fp8 gemm instances * re-enable conv_3d fp8 on all architectures * diasble several fp8 gemm instances on all architectures except gfx94 * clang format fix --- CMakeLists.txt | 9 +++----- .../gpu/gemm/CMakeLists.txt | 23 +++++++++---------- .../gpu/gemm_universal/CMakeLists.txt | 5 ++-- ...gemm_xdl_universal_f8_f8_bf16_mk_kn_mn.hpp | 6 +++-- .../grouped_conv3d_bwd_data/CMakeLists.txt | 2 +- .../grouped_conv3d_bwd_weight/CMakeLists.txt | 2 +- .../CMakeLists.txt | 2 +- .../CMakeLists.txt | 2 +- .../gpu/grouped_conv3d_fwd/CMakeLists.txt | 8 +++---- 9 files changed, 28 insertions(+), 31 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 2039948a12..8a08ddd19c 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -62,17 +62,14 @@ if (DTYPES) endif() message("DTYPES macro set to ${DTYPES}") else() - add_definitions(-DCK_ENABLE_INT8 -DCK_ENABLE_FP16 -DCK_ENABLE_FP32 -DCK_ENABLE_FP64 -DCK_ENABLE_BF16) + add_definitions(-DCK_ENABLE_INT8 -DCK_ENABLE_FP16 -DCK_ENABLE_FP32 -DCK_ENABLE_FP64 -DCK_ENABLE_BF16 -DCK_ENABLE_FP8 -DCK_ENABLE_BF8) set(CK_ENABLE_INT8 "ON") set(CK_ENABLE_FP16 "ON") set(CK_ENABLE_FP32 "ON") set(CK_ENABLE_FP64 "ON") set(CK_ENABLE_BF16 "ON") - if (GPU_TARGETS MATCHES "gfx94") - add_definitions(-DCK_ENABLE_FP8 -DCK_ENABLE_BF8) - set(CK_ENABLE_FP8 "ON") - set(CK_ENABLE_BF8 "ON") - endif() + set(CK_ENABLE_FP8 "ON") + set(CK_ENABLE_BF8 "ON") endif() #for f8/bf8_t type diff --git a/library/src/tensor_operation_instance/gpu/gemm/CMakeLists.txt b/library/src/tensor_operation_instance/gpu/gemm/CMakeLists.txt index 0cd54c7788..e4efae6173 100644 --- a/library/src/tensor_operation_instance/gpu/gemm/CMakeLists.txt +++ b/library/src/tensor_operation_instance/gpu/gemm/CMakeLists.txt @@ -100,18 +100,17 @@ list(APPEND GEMM_INSTANCES device_gemm_xdl_c_shuffle_bf16_bf16_bf16_km_kn_mn_instance.cpp device_gemm_xdl_c_shuffle_bf16_bf16_bf16_km_nk_mn_instance.cpp) -if(GPU_TARGETS MATCHES "gfx94") - list(APPEND GEMM_INSTANCES - device_gemm_xdl_c_shuffle_fp8_fp8_fp8_mk_kn_mn_v1_default_instance.cpp - device_gemm_xdl_c_shuffle_fp8_fp8_fp8_mk_kn_mn_v1_interwave_default_instance.cpp - device_gemm_xdl_c_shuffle_fp8_fp8_fp8_mk_kn_mn_v2_default_instance.cpp - device_gemm_xdl_c_shuffle_fp8_fp8_fp8_mk_kn_mn_v1_padded_instance.cpp - device_gemm_xdl_c_shuffle_fp8_fp8_fp8_mk_kn_mn_v1_interwave_padded_instance.cpp - device_gemm_xdl_c_shuffle_fp8_fp8_fp8_mk_kn_mn_v2_padded_instance.cpp - device_gemm_xdl_c_shuffle_fp8_fp8_fp8_mk_nk_mn_instance.cpp - device_gemm_xdl_c_shuffle_fp8_fp8_fp8_km_kn_mn_instance.cpp - device_gemm_xdl_c_shuffle_fp8_fp8_fp8_km_nk_mn_instance.cpp) -endif() +list(APPEND GEMM_INSTANCES + device_gemm_xdl_c_shuffle_fp8_fp8_fp8_mk_kn_mn_v1_default_instance.cpp + device_gemm_xdl_c_shuffle_fp8_fp8_fp8_mk_kn_mn_v1_interwave_default_instance.cpp + device_gemm_xdl_c_shuffle_fp8_fp8_fp8_mk_kn_mn_v2_default_instance.cpp + device_gemm_xdl_c_shuffle_fp8_fp8_fp8_mk_kn_mn_v1_padded_instance.cpp + device_gemm_xdl_c_shuffle_fp8_fp8_fp8_mk_kn_mn_v1_interwave_padded_instance.cpp + device_gemm_xdl_c_shuffle_fp8_fp8_fp8_mk_kn_mn_v2_padded_instance.cpp + device_gemm_xdl_c_shuffle_fp8_fp8_fp8_mk_nk_mn_instance.cpp + device_gemm_xdl_c_shuffle_fp8_fp8_fp8_km_kn_mn_instance.cpp + device_gemm_xdl_c_shuffle_fp8_fp8_fp8_km_nk_mn_instance.cpp) + list(APPEND GEMM_INSTANCES device_gemm_wmma_f16_f16_f16_mk_kn_mn_instance.cpp diff --git a/library/src/tensor_operation_instance/gpu/gemm_universal/CMakeLists.txt b/library/src/tensor_operation_instance/gpu/gemm_universal/CMakeLists.txt index cc4ce76606..5b6e985e59 100644 --- a/library/src/tensor_operation_instance/gpu/gemm_universal/CMakeLists.txt +++ b/library/src/tensor_operation_instance/gpu/gemm_universal/CMakeLists.txt @@ -51,8 +51,7 @@ set_source_files_properties(device_gemm_xdl_universal_bf16_bf16_bf16/device_gemm set_source_files_properties(device_gemm_xdl_universal_bf16_bf16_bf16/device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_comp_default_instance.cpp PROPERTIES COMPILE_OPTIONS ";-mllvm;-greedy-reverse-local-assignment=1") set_source_files_properties(device_gemm_xdl_universal_bf16_bf16_bf16/device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_comp_kpadding_instance.cpp PROPERTIES COMPILE_OPTIONS ";-mllvm;-greedy-reverse-local-assignment=1") -if(GPU_TARGETS MATCHES "gfx94") - list(APPEND GEMM_UNIVERSAL_INSTANCES +list(APPEND GEMM_UNIVERSAL_INSTANCES device_gemm_xdl_universal_f16_f8_f16/device_gemm_xdl_universal_f16_f8_f16_mk_kn_mn_comp_default_instance.cpp device_gemm_xdl_universal_f16_f8_f16/device_gemm_xdl_universal_f16_f8_f16_mk_kn_mn_comp_kpadding_instance.cpp device_gemm_xdl_universal_f16_f8_f16/device_gemm_xdl_universal_f16_f8_f16_mk_kn_mn_comp_mnpadding_instance.cpp @@ -123,6 +122,6 @@ if(GPU_TARGETS MATCHES "gfx94") set_source_files_properties(device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_comp_nkpadding_instance.cpp PROPERTIES COMPILE_OPTIONS ";-mllvm;-greedy-reverse-local-assignment=1") set_source_files_properties(device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_comp_default_instance.cpp PROPERTIES COMPILE_OPTIONS ";-mllvm;-greedy-reverse-local-assignment=1") set_source_files_properties(device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_comp_kpadding_instance.cpp PROPERTIES COMPILE_OPTIONS ";-mllvm;-greedy-reverse-local-assignment=1") -endif() + add_instance_library(device_gemm_universal_instance ${GEMM_UNIVERSAL_INSTANCES}) diff --git a/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn.hpp b/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn.hpp index 12994aeecd..3b930e9894 100644 --- a/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn.hpp +++ b/library/src/tensor_operation_instance/gpu/gemm_universal/device_gemm_xdl_universal_f8_f8_bf16/device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn.hpp @@ -35,12 +35,13 @@ static constexpr auto Interwave = BlockGemmPipelineScheduler::Interwave; template using device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_comp_instances = std::tuple< - // clang-format off +// clang-format off //#########################| ALayout| BLayout| CLayout|AData| BData| CData| AccData| Cshuffle| A| B| C| GEMM| Block| MPer| NPer| KPer| AK1| BK1|MPer| NPer| MXdl| NXdl| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockLds| BBlockTransfer| BBlockTransfer| BBlockTransfer| BlockTransfer| BBlockTransfer| BBlockTransfer| BBlockLds| CShuffle| CShuffle| CBlockTransferClusterLengths| CBlockTransfer| Block-wiseGemm| Block-wiseGemm| //#########################| | | | 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__ + //Only enable these instances on gfx94x // Compute friendly DeviceGemm_Xdl_CShuffleV3< Row, Row, Row, F8, F8, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 256, 256, 64, 16, 4, 32, 32, 4, 4, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, S<16, 16, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 16, 4, 0, 1, 1, S<1, 32, 1, 8>, 8, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v4, F8>, DeviceGemm_Xdl_CShuffleV3< Row, Row, Row, F8, F8, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 128, 128, 128, 16, 4, 32, 32, 2, 2, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, S<32, 8, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 16, 4, 0, 1, 1, S<1, 32, 1, 8>, 8, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v4, F8>, @@ -55,6 +56,7 @@ using device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_comp_instances = std::tuple< DeviceGemm_Xdl_CShuffleV3< Row, Row, Row, F8, F8, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 128, 128, 128, 16, 4, 32, 32, 2, 2, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, S<32, 8, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 16, 4, 0, 1, 1, S<1, 32, 1, 8>, 8, BlockGemmPipelineScheduler::Interwave, BlockGemmPipelineVersion::v1, F8>, DeviceGemm_Xdl_CShuffleV3< Row, Row, Row, F8, F8, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 128, 64, 128, 16, 4, 32, 32, 2, 1, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, S<32, 8, 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, F8>, DeviceGemm_Xdl_CShuffleV3< Row, Row, Row, F8, F8, BF16, F32, BF16, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 64, 64, 128, 16, 4, 32, 32, 1, 1, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, S<32, 8, 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, F8> +#endif // clang-format on >; diff --git a/library/src/tensor_operation_instance/gpu/grouped_conv3d_bwd_data/CMakeLists.txt b/library/src/tensor_operation_instance/gpu/grouped_conv3d_bwd_data/CMakeLists.txt index 7bb7e71c54..29fa8fa3c5 100644 --- a/library/src/tensor_operation_instance/gpu/grouped_conv3d_bwd_data/CMakeLists.txt +++ b/library/src/tensor_operation_instance/gpu/grouped_conv3d_bwd_data/CMakeLists.txt @@ -15,7 +15,7 @@ set(GROUPED_CONV3D_BWD_DATA wmma/device_grouped_conv3d_bwd_data_wmma_gndhwc_gkzyxc_gndhwk_i8_1x1s1p0_instance.cpp wmma/device_grouped_conv3d_bwd_data_wmma_ndhwgc_gkzyxc_ndhwgk_i8_1x1s1p0_instance.cpp) -if((DTYPES MATCHES "fp8" AND DTYPES MATCHES "bf8" AND DTYPES MATCHES "fp16") OR (NOT DEFINED DTYPES AND GPU_TARGETS MATCHES "gfx94")) +if((DTYPES MATCHES "fp8" AND DTYPES MATCHES "bf8" AND DTYPES MATCHES "fp16") OR NOT DEFINED DTYPES) list(APPEND GROUPED_CONV3D_BWD_DATA xdl/device_grouped_conv3d_bwd_data_xdl_ndhwgc_gkzyxc_ndhwgk_input_f16_comp_bf8_f8_instance.cpp) endif() diff --git a/library/src/tensor_operation_instance/gpu/grouped_conv3d_bwd_weight/CMakeLists.txt b/library/src/tensor_operation_instance/gpu/grouped_conv3d_bwd_weight/CMakeLists.txt index 1a9c455220..8e939c15a9 100644 --- a/library/src/tensor_operation_instance/gpu/grouped_conv3d_bwd_weight/CMakeLists.txt +++ b/library/src/tensor_operation_instance/gpu/grouped_conv3d_bwd_weight/CMakeLists.txt @@ -30,7 +30,7 @@ list(APPEND GROUPED_CONV3D_BWD_WEIGHT wmma/device_grouped_conv3d_bwd_weight_wmma_gndhwc_gkzyxc_gndhwk_i8_instance.cpp wmma/device_grouped_conv3d_bwd_weight_wmma_ndhwgc_gkzyxc_ndhwgk_i8_instance.cpp) -if((DTYPES MATCHES "fp8" AND DTYPES MATCHES "bf8" AND DTYPES MATCHES "fp16") OR (NOT DEFINED DTYPES AND GPU_TARGETS MATCHES "gfx94")) +if((DTYPES MATCHES "fp8" AND DTYPES MATCHES "bf8" AND DTYPES MATCHES "fp16") OR NOT DEFINED DTYPES) list(APPEND GROUPED_CONV3D_BWD_WEIGHT xdl/device_grouped_conv3d_bwd_weight_xdl_ndhwgc_gkzyxc_ndhwgk_f16_comp_bf8_fp8_instance.cpp) endif() diff --git a/library/src/tensor_operation_instance/gpu/grouped_conv3d_bwd_weight_bilinear/CMakeLists.txt b/library/src/tensor_operation_instance/gpu/grouped_conv3d_bwd_weight_bilinear/CMakeLists.txt index 5781f07080..329e8e4c7f 100644 --- a/library/src/tensor_operation_instance/gpu/grouped_conv3d_bwd_weight_bilinear/CMakeLists.txt +++ b/library/src/tensor_operation_instance/gpu/grouped_conv3d_bwd_weight_bilinear/CMakeLists.txt @@ -4,7 +4,7 @@ set(GROUPED_CONV3D_BWD_WEIGHT_BILINEAR xdl/device_grouped_conv3d_bwd_weight_xdl_bilinear_ndhwgc_gkzyxc_ndhwgk_f32_instance.cpp xdl/device_grouped_conv3d_bwd_weight_xdl_bilinear_ndhwgc_gkzyxc_ndhwgk_bf16_instance.cpp) -if((DTYPES MATCHES "fp8" AND DTYPES MATCHES "bf8" AND DTYPES MATCHES "fp16") OR (NOT DEFINED DTYPES AND GPU_TARGETS MATCHES "gfx94")) +if((DTYPES MATCHES "fp8" AND DTYPES MATCHES "bf8" AND DTYPES MATCHES "fp16") OR NOT DEFINED DTYPES) list(APPEND GROUPED_CONV3D_BWD_WEIGHT_BILINEAR xdl/device_grouped_conv3d_bwd_weight_xdl_bilinear_ndhwgc_gkzyxc_ndhwgk_f16_comp_bf8_fp8_instance.cpp) endif() diff --git a/library/src/tensor_operation_instance/gpu/grouped_conv3d_bwd_weight_scale/CMakeLists.txt b/library/src/tensor_operation_instance/gpu/grouped_conv3d_bwd_weight_scale/CMakeLists.txt index be54eb4adf..9a42d1ec3a 100644 --- a/library/src/tensor_operation_instance/gpu/grouped_conv3d_bwd_weight_scale/CMakeLists.txt +++ b/library/src/tensor_operation_instance/gpu/grouped_conv3d_bwd_weight_scale/CMakeLists.txt @@ -4,7 +4,7 @@ set(GROUPED_CONV3D_BWD_WEIGHT_SCALE xdl/device_grouped_conv3d_bwd_weight_xdl_scale_ndhwgc_gkzyxc_ndhwgk_f32_instance.cpp xdl/device_grouped_conv3d_bwd_weight_xdl_scale_ndhwgc_gkzyxc_ndhwgk_bf16_instance.cpp) -if((DTYPES MATCHES "fp8" AND DTYPES MATCHES "bf8" AND DTYPES MATCHES "fp16") OR (NOT DEFINED DTYPES AND GPU_TARGETS MATCHES "gfx94")) +if((DTYPES MATCHES "fp8" AND DTYPES MATCHES "bf8" AND DTYPES MATCHES "fp16") OR NOT DEFINED DTYPES) list(APPEND GROUPED_CONV3D_BWD_WEIGHT_SCALE xdl/device_grouped_conv3d_bwd_weight_xdl_scale_ndhwgc_gkzyxc_ndhwgk_f16_comp_bf8_fp8_instance.cpp) endif() diff --git a/library/src/tensor_operation_instance/gpu/grouped_conv3d_fwd/CMakeLists.txt b/library/src/tensor_operation_instance/gpu/grouped_conv3d_fwd/CMakeLists.txt index 9bb6d807e6..5a346f1e94 100644 --- a/library/src/tensor_operation_instance/gpu/grouped_conv3d_fwd/CMakeLists.txt +++ b/library/src/tensor_operation_instance/gpu/grouped_conv3d_fwd/CMakeLists.txt @@ -43,22 +43,22 @@ set(GROUPED_CONV3D_FWD wmma/device_grouped_conv3d_fwd_wmma_ndhwgc_gkzyxc_ndhwgk_i8_oddc_instance.cpp ) -if((DTYPES MATCHES "fp8" AND DTYPES MATCHES "fp16") OR (NOT DEFINED DTYPES AND GPU_TARGETS MATCHES "gfx94")) +if((DTYPES MATCHES "fp8" AND DTYPES MATCHES "fp16") OR NOT DEFINED DTYPES) list(APPEND GROUPED_CONV3D_FWD xdl/device_grouped_conv3d_fwd_xdl_ndhwgc_gkzyxc_ndhwgk_f16_comp_fp8_instance.cpp) endif() -if((DTYPES MATCHES "fp8") OR (NOT DEFINED DTYPES AND GPU_TARGETS MATCHES "gfx94")) +if((DTYPES MATCHES "fp8") OR NOT DEFINED DTYPES) list(APPEND GROUPED_CONV3D_FWD xdl/device_grouped_conv3d_fwd_xdl_ndhwgc_gkzyxc_ndhwgk_fp8_instance.cpp) endif() -if((DTYPES MATCHES "bf8") OR (NOT DEFINED DTYPES AND GPU_TARGETS MATCHES "gfx94")) +if((DTYPES MATCHES "bf8") OR NOT DEFINED DTYPES) list(APPEND GROUPED_CONV3D_FWD xdl/device_grouped_conv3d_fwd_xdl_ndhwgc_gkzyxc_ndhwgk_bf8_instance.cpp) endif() -if((DTYPES MATCHES "fp8" AND DTYPES MATCHES "bf8") OR (NOT DEFINED DTYPES AND GPU_TARGETS MATCHES "gfx94")) +if((DTYPES MATCHES "fp8" AND DTYPES MATCHES "bf8") OR NOT DEFINED DTYPES) list(APPEND GROUPED_CONV3D_FWD xdl/device_grouped_conv3d_fwd_xdl_ndhwgc_gkzyxc_ndhwgk_fp8_bf8_instance.cpp) list(APPEND GROUPED_CONV3D_FWD From a6a796650587d845e28c2c2e99535aa97065d36b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Bart=C5=82omiej=20Kocot?= Date: Mon, 19 Aug 2024 17:24:56 +0200 Subject: [PATCH 09/20] Add script to convert MIOpen driver to ckProfiler (#1472) * Add script to convert MIOpen driver to ckProfiler * Fix --- .../profile_grouped_conv_bwd_weight_impl.hpp | 169 ++++---- .../src/profile_grouped_conv_bwd_weight.cpp | 3 +- script/convert_miopen_driver_to_profiler.py | 386 ++++++++++++++++++ 3 files changed, 480 insertions(+), 78 deletions(-) create mode 100644 script/convert_miopen_driver_to_profiler.py diff --git a/profiler/include/profiler/profile_grouped_conv_bwd_weight_impl.hpp b/profiler/include/profiler/profile_grouped_conv_bwd_weight_impl.hpp index 356aec7a08..5318de5e8b 100644 --- a/profiler/include/profiler/profile_grouped_conv_bwd_weight_impl.hpp +++ b/profiler/include/profiler/profile_grouped_conv_bwd_weight_impl.hpp @@ -136,9 +136,10 @@ bool profile_grouped_conv_bwd_weight_impl(int do_verification, std::cout << "found " << op_ptrs.size() << " instances" << std::endl; std::string best_op_name; - float best_avg_time = 0; - float best_tflops = 0; - float best_gb_per_sec = 0; + float best_avg_time = 0; + float best_tflops = 0; + float best_gb_per_sec = 0; + ck::index_t best_split_k = 1; // profile device Conv instances bool all_pass = true; @@ -167,99 +168,115 @@ bool profile_grouped_conv_bwd_weight_impl(int do_verification, range_copy(conv_param.input_left_pads_, begin(input_left_pads)); range_copy(conv_param.input_right_pads_, begin(input_right_pads)); + std::vector split_k_list = {1, 2, 4, 8, 16, 32, 64, 128}; + + if(split_k > 0) + { + split_k_list = {split_k}; + } + for(auto& op_ptr : op_ptrs) { - auto argument_ptr = - op_ptr->MakeArgumentPointer(static_cast(in_device_buf.GetDeviceBuffer()), - static_cast(wei_device_buf.GetDeviceBuffer()), - static_cast(out_device_buf.GetDeviceBuffer()), - input_lengths, - input_strides, - filter_lengths, - weights_strides, - output_lengths, - output_strides, - conv_filter_strides, - conv_filter_dilations, - input_left_pads, - input_right_pads, - in_element_op, - wei_element_op, - out_element_op, - split_k); - - const std::size_t workspace_sz = op_ptr->GetWorkSpaceSize(argument_ptr.get()); - DeviceMem workspace_dev(workspace_sz); - op_ptr->SetWorkSpacePointer(argument_ptr.get(), workspace_dev.GetDeviceBuffer()); - - if(op_ptr->IsSupportedArgument(argument_ptr.get())) + for(std::size_t split_k_id = 0; split_k_id < split_k_list.size(); split_k_id++) { - // using atomic add, so need to reset input - wei_device_buf.SetZero(); + auto argument_ptr = op_ptr->MakeArgumentPointer( + static_cast(in_device_buf.GetDeviceBuffer()), + static_cast(wei_device_buf.GetDeviceBuffer()), + static_cast(out_device_buf.GetDeviceBuffer()), + input_lengths, + input_strides, + filter_lengths, + weights_strides, + output_lengths, + output_strides, + conv_filter_strides, + conv_filter_dilations, + input_left_pads, + input_right_pads, + in_element_op, + wei_element_op, + out_element_op, + split_k_list[split_k_id]); - std::string op_name = op_ptr->GetTypeString(); + const std::size_t workspace_sz = op_ptr->GetWorkSpaceSize(argument_ptr.get()); + DeviceMem workspace_dev(workspace_sz); + op_ptr->SetWorkSpacePointer(argument_ptr.get(), workspace_dev.GetDeviceBuffer()); - auto invoker_ptr = op_ptr->MakeInvokerPointer(); - - float avg_time = - invoker_ptr->Run(argument_ptr.get(), StreamConfig{nullptr, time_kernel}); - - std::size_t flop = conv_param.GetFlops(); - std::size_t num_btype = conv_param.GetByte(); - - float tflops = static_cast(flop) / 1.E9 / avg_time; - float gb_per_sec = num_btype / 1.E6 / avg_time; - - std::cout << "Perf: " << std::setw(10) << avg_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_avg_time = avg_time; - best_gb_per_sec = gb_per_sec; - } + // using atomic add, so need to reset input + wei_device_buf.SetZero(); - if(do_verification) - { - wei_device_buf.FromDevice(weight_device_result.mData.data()); + std::string op_name = op_ptr->GetTypeString(); - bool pass = ck::utils::check_err(weight_device_result, weight_host_result); + auto invoker_ptr = op_ptr->MakeInvokerPointer(); - if(!pass) + float avg_time = + invoker_ptr->Run(argument_ptr.get(), StreamConfig{nullptr, time_kernel}); + + std::size_t flop = conv_param.GetFlops(); + std::size_t num_btype = conv_param.GetByte(); + + float tflops = static_cast(flop) / 1.E9 / avg_time; + float gb_per_sec = num_btype / 1.E6 / avg_time; + + std::cout << "Perf: " << std::setw(10) << avg_time << " ms, " << tflops + << " TFlops, " << gb_per_sec << " GB/s, " << op_name << ", SplitK " + << split_k_list[split_k_id] << std::endl; + + if(tflops > best_tflops) { - std::cout << "Fail info: " << op_ptr->GetTypeString() << std::endl; + best_op_name = op_name; + best_tflops = tflops; + best_avg_time = avg_time; + best_gb_per_sec = gb_per_sec; + best_split_k = split_k_list[split_k_id]; } - all_pass &= pass; - - if(do_log) + if(do_verification) { - LogRangeAsType(std::cout << "output : ", output.mData, ",") << std::endl; - ; - LogRangeAsType( - std::cout << "weight (device): ", weight_device_result.mData, ",") - << std::endl; - ; - LogRangeAsType( - std::cout << "weight (host): ", weight_host_result.mData, ",") - << std::endl; - ; - LogRangeAsType(std::cout << "input: ", input.mData, ",") << std::endl; - ; + wei_device_buf.FromDevice(weight_device_result.mData.data()); + + bool pass = ck::utils::check_err(weight_device_result, weight_host_result); + + if(!pass) + { + std::cout << "Fail info: " << op_ptr->GetTypeString() << std::endl; + } + + all_pass &= pass; + + if(do_log) + { + LogRangeAsType(std::cout << "output : ", output.mData, ",") + << std::endl; + ; + LogRangeAsType( + std::cout << "weight (device): ", weight_device_result.mData, ",") + << std::endl; + ; + LogRangeAsType( + std::cout << "weight (host): ", weight_host_result.mData, ",") + << std::endl; + ; + LogRangeAsType(std::cout << "input: ", input.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; + } } } std::cout << "Best configuration parameters:" << "\nname: " << best_op_name << "\navg_time: " << best_avg_time - << "\ntflops: " << best_tflops << "\nGB/s: " << best_gb_per_sec << std::endl; + << "\ntflops: " << best_tflops << "\nGB/s: " << best_gb_per_sec << ", SplitK " + << best_split_k << std::endl; return all_pass; } diff --git a/profiler/src/profile_grouped_conv_bwd_weight.cpp b/profiler/src/profile_grouped_conv_bwd_weight.cpp index 6ed7cf5e48..7dd75a5e0a 100644 --- a/profiler/src/profile_grouped_conv_bwd_weight.cpp +++ b/profiler/src/profile_grouped_conv_bwd_weight.cpp @@ -1,5 +1,5 @@ // SPDX-License-Identifier: MIT -// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved. +// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved. #include #include @@ -81,7 +81,6 @@ int profile_grouped_conv_bwd_weight(int argc, char* argv[]) const auto params = ck::utils::conv::parse_conv_param(num_dim_spatial, 9, argv); ck::index_t split_k = std::stoi(argv[8 + 1 + 4 + 6 * num_dim_spatial]); - split_k = std::max(1, split_k); using F32 = float; using F16 = ck::half_t; diff --git a/script/convert_miopen_driver_to_profiler.py b/script/convert_miopen_driver_to_profiler.py new file mode 100644 index 0000000000..47135f3401 --- /dev/null +++ b/script/convert_miopen_driver_to_profiler.py @@ -0,0 +1,386 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved. +# Convert miopen driver command to ck Profiler +# Example: python3 ../script/convert_miopen_driver_to_profiler.py +# /opt/rocm/bin/MIOpenDriver conv -n 32 -c 64 -H 28 -W 28 -k 64 -y 3 -x 3 +# -p 1 -q 1 -u 2 -v 2 -l 1 -j 1 -m conv -g 32 -F 1 -t 1 + +import argparse +import subprocess + + +def init_const_args(args): + args.ck_profiler_cmd = '../build/bin/ckProfiler' + # use decimal values + args.init_method = 2 + # don't print tensor values + args.log_value = 0 + + +def run_ck_profiler_cmd(cmd): + print("ckProfiler command:") + print(cmd) + subprocess.run(cmd) + + +def parse_data_type(args): + if args.data_type == "fp32": + if args.ck_profier_op == "grouped_conv_bwd_weight" or \ + args.ck_profier_op == "grouped_conv_bwd_weight" or \ + args.ck_profier_op == "grouped_conv_fwd": + args.data_type = 0 + if args.data_type == "fp16": + if args.ck_profier_op == "grouped_conv_bwd_weight" or \ + args.ck_profier_op == "grouped_conv_bwd_data" or \ + args.ck_profier_op == "grouped_conv_fwd": + args.data_type = 1 + if args.data_type == "int8": + if args.ck_profier_op == "grouped_conv_bwd_weight": + args.data_type = 4 + if args.ck_profier_op == "grouped_conv_bwd_data": + print('Not supported data type for grouped_conv_bwd_data') + exit(1) + if args.ck_profier_op == "grouped_conv_fwd": + args.data_type = 3 + if args.data_type == "bfp16": + if args.ck_profier_op == "grouped_conv_bwd_weight" or \ + args.ck_profier_op == "grouped_conv_bwd_data" or \ + args.ck_profier_op == "grouped_conv_fwd": + args.data_type = 2 + + +def add_conv_params_to_cmd(args, cmd): + if args.spatial_dim == 1: + cmd += [str(args.fil_w), str(args.in_w)] + cmd += [str(args.conv_stride_w), str(args.dilation_w)] + cmd += [str(args.pad_w), str(args.pad_w)] + elif args.spatial_dim == 2: + cmd += [str(args.fil_h), str(args.fil_w)] + cmd += [str(args.in_h), str(args.in_w)] + cmd += [str(args.conv_stride_h), str(args.conv_stride_w)] + cmd += [str(args.dilation_h), str(args.dilation_w)] + cmd += [str(args.pad_h), str(args.pad_w)] + cmd += [str(args.pad_h), str(args.pad_w)] + elif args.spatial_dim == 3: + cmd += [str(args.fil_d), str(args.fil_h), str(args.fil_w)] + cmd += [str(args.in_d), str(args.in_h), str(args.in_w)] + cmd += [str(args.conv_stride_d), str(args.conv_stride_h)] + cmd += [str(args.conv_stride_w)] + cmd += [str(args.dilation_d), + str(args.dilation_h), + str(args.dilation_w)] + cmd += [str(args.pad_d), str(args.pad_h), str(args.pad_w)] + cmd += [str(args.pad_d), str(args.pad_h), str(args.pad_w)] + else: + print('Not supported spatial dim (supported: 1, 2, 3)') + exit(1) + + +def run_ck_grouped_conv_fwd(args): + args.ck_profier_op = "grouped_conv_fwd" + parse_data_type(args) + # default for MIOpen NHWGC + args.layout = 1 + # use int32 by default + args.index_type = 0 + + cmd = [str(args.ck_profiler_cmd), str(args.ck_profier_op)] + cmd += [str(args.data_type), str(args.layout), str(args.index_type)] + cmd += [str(args.verify), str(args.init_method)] + cmd += [str(args.log_value), str(args.time)] + cmd += [str(args.spatial_dim), str(args.group_count)] + cmd += [str(args.batchsize), str(args.out_channels)] + cmd += [str(args.in_channels)] + add_conv_params_to_cmd(args, cmd) + + run_ck_profiler_cmd(cmd) + + +def run_ck_grouped_conv_bwd_data(args): + args.ck_profier_op = "grouped_conv_bwd_data" + parse_data_type(args) + # default for MIOpen NHWGC + args.layout = 1 + + cmd = [str(args.ck_profiler_cmd), str(args.ck_profier_op)] + cmd += [str(args.data_type), str(args.layout)] + cmd += [str(args.verify), str(args.init_method)] + cmd += [str(args.log_value), str(args.time)] + cmd += [str(args.spatial_dim), str(args.group_count)] + cmd += [str(args.batchsize), str(args.out_channels)] + cmd += [str(args.in_channels)] + add_conv_params_to_cmd(args, cmd) + + run_ck_profiler_cmd(cmd) + + +def run_ck_grouped_conv_bwd_weight(args): + args.ck_profier_op = "grouped_conv_bwd_weight" + parse_data_type(args) + # default for MIOpen NHWGC + args.layout = 2 + # Test all split K value from the list {1, 2, 4, 8, 32, 64, 128} + args.split_k_value = -1 + + cmd = [str(args.ck_profiler_cmd), str(args.ck_profier_op)] + cmd += [str(args.data_type), str(args.layout)] + cmd += [str(args.verify), str(args.init_method)] + cmd += [str(args.log_value), str(args.time)] + cmd += [str(args.spatial_dim), str(args.group_count)] + cmd += [str(args.batchsize), str(args.out_channels)] + cmd += [str(args.in_channels)] + add_conv_params_to_cmd(args, cmd) + + cmd += [str(args.split_k_value)] + run_ck_profiler_cmd(cmd) + +# Get name of miopen driver, remove it from unknown +def process_miopen_driver_name(args, unknown): + if "convint8" in unknown: + args.data_type = 'int8' + unknown.remove("convint8") + elif "convbfp16" in unknown: + args.data_type = 'bfp16' + unknown.remove("convbfp16") + elif "convfp16" in unknown: + args.data_type = 'fp16' + unknown.remove("convfp16") + elif "conv" in unknown: + args.data_type = 'fp32' + unknown.remove("conv") + else: + print('Not supported driver (supported: conv, convfp16, convint8,' + ' convbfp16).') + exit(1) + + +def run_ck_profiler(args): + # MIOpen get number of channel per all groups, CK profiler get number of + # channel per group + args.in_channels = int(args.in_channels / args.group_count) + args.out_channels = int(args.out_channels / args.group_count) + + if args.forw == 0 or args.forw == 1 or args.forw == 3 or args.forw == 5: + run_ck_grouped_conv_fwd(args) + if args.forw == 0 or args.forw == 2 or args.forw == 3 or args.forw == 6: + run_ck_grouped_conv_bwd_data(args) + if args.forw == 0 or args.forw == 4 or args.forw == 5 or args.forw == 6: + run_ck_grouped_conv_bwd_weight(args) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + prog="converter", + description="Convert miopen driver command to ck Profiler" + "\nExample: python3 " + "../script/convert_miopen_driver_to_profiler.py " + "/opt/rocm/bin/MIOpenDriver conv -n 32 -c 64 -H 28 -W 28 " + "-k 64 -y 3 -x 3 -p 1 -q 1 -u 1 -v 1 -l 1 -j 1 -m conv -g " + "32 -F 1 -t 1", + ) + parser.add_argument( + "-in_layout", + "-I", + default=-1, + type=int, + required=False, + help="Input Layout (Default=NCHW for 2d conv, NCDHW for 3d conv)" + ) + parser.add_argument( + "-forw", + "-F", + default=0, + type=int, + required=False, + help="Flag enables fwd, bwd, wrw convolutions" + "\n0 fwd+bwd+wrw (default)" + "\n1 fwd only" + "\n2 bwd only" + "\n4 wrw only" + "\n3 fwd+bwd" + "\n5 fwd+wrw" + "\n6 bwd+wrw" + ) + parser.add_argument( + "-spatial_dim", + "-_", + default=2, + type=int, + required=False, + help="convolution spatial dimension (Default-2)" + ) + parser.add_argument( + "-batchsize", + "-n", + default=100, + type=int, + required=False, + help="Mini-batch size (Default=100)" + ) + parser.add_argument( + "-in_channels", + "-c", + default=3, + type=int, + required=False, + help="Number of Input Channels (Default=3)" + ) + parser.add_argument( + "-in_d", + "-!", + default=32, + type=int, + required=False, + help="Input Depth (Default=32)" + ) + parser.add_argument( + "-in_h", + "-H", + default=32, + type=int, + required=False, + help="Input Height (Default=32)" + ) + parser.add_argument( + "-in_w", + "-W", + default=32, + type=int, + required=False, + help="Input Width (Default=32)" + ) + parser.add_argument( + "-out_channels", + "-k", + default=32, + type=int, + required=False, + help="Number of Output Channels (Default=32)" + ) + parser.add_argument( + "-fil_d", + "-@", + default=3, + type=int, + required=False, + help="Filter Depth (Default=3)" + ) + parser.add_argument( + "-fil_h", + "-y", + default=3, + type=int, + required=False, + help="Filter Height (Default=3)" + ) + parser.add_argument( + "-fil_w", + "-x", + default=3, + type=int, + required=False, + help="Filter Width (Default=3)" + ) + parser.add_argument( + "-conv_stride_d", + "-#", + default=1, + type=int, + required=False, + help="Convolution Stride for Depth (Default=1)" + ) + parser.add_argument( + "-conv_stride_h", + "-u", + default=1, + type=int, + required=False, + help="Convolution Stride for Height (Default=1)" + ) + parser.add_argument( + "-conv_stride_w", + "-v", + default=1, + type=int, + required=False, + help="Convolution Stride for Width (Default=1)" + ) + parser.add_argument( + "-pad_d", + "-$", + default=1, + type=int, + required=False, + help="Zero Padding for Depth (Default=0)" + ) + parser.add_argument( + "-pad_h", + "-p", + default=1, + type=int, + required=False, + help="Zero Padding for Height (Default=0)" + ) + parser.add_argument( + "-pad_w", + "-q", + default=1, + type=int, + required=False, + help="Zero Padding for Width (Default=0)" + ) + parser.add_argument( + "-verify", + "-V", + default=1, + type=int, + required=False, + help="Verify Each Layer (Default=1)" + ) + parser.add_argument( + "-time", + "-t", + default=0, + type=int, + required=False, + help="Time Each Layer (Default=0)" + ) + parser.add_argument( + "-dilation_d", + "-^", + default=1, + type=int, + required=False, + help="Dilation of Filter Depth (Default=1)" + ) + parser.add_argument( + "-dilation_h", + "-l", + default=1, + type=int, + required=False, + help="Dilation of Filter Height (Default=1)" + ) + parser.add_argument( + "-dilation_w", + "-j", + default=1, + type=int, + required=False, + help="Dilation of Filter Width (Default=1)" + ) + parser.add_argument( + "-group_count", + "-g", + type=int, + default=1, + required=False, + help="Number of Groups (Default=1)" + ) + + args, unknown = parser.parse_known_args() + init_const_args(args) + process_miopen_driver_name(args, unknown) + print("Ignored args:") + print(unknown) + run_ck_profiler(args) From f48529b51161bc502ce1b7b3e9a3d1d5217f453b Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Mon, 19 Aug 2024 23:02:07 -0700 Subject: [PATCH 10/20] Bump rocm-docs-core from 1.7.0 to 1.7.1 in /docs/sphinx (#1475) Bumps [rocm-docs-core](https://github.com/ROCm/rocm-docs-core) from 1.7.0 to 1.7.1. - [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.7.0...v1.7.1) --- updated-dependencies: - dependency-name: rocm-docs-core dependency-type: direct:production update-type: version-update:semver-patch ... 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 8be30305a6..6bc6e07eea 100644 --- a/docs/sphinx/requirements.in +++ b/docs/sphinx/requirements.in @@ -1,2 +1,2 @@ -rocm-docs-core==1.7.0 +rocm-docs-core==1.7.1 sphinxcontrib-bibtex==2.6.2 diff --git a/docs/sphinx/requirements.txt b/docs/sphinx/requirements.txt index 4cc4d30f79..0e02dbb727 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.7.0 +rocm-docs-core==1.7.1 # via -r requirements.in six==1.16.0 # via pybtex From a94113a9418bb2333c72ff28b7a1fc44283fddba Mon Sep 17 00:00:00 2001 From: Andriy Roshchenko <107577548+andriy-ca@users.noreply.github.com> Date: Tue, 20 Aug 2024 09:30:56 -0600 Subject: [PATCH 11/20] Adding Instances and Examples for FP8-based Scaled Convolution with ReLU Activation and AMAX Reduction. (#1469) * Enable CMakePresets build * Verify Convolution, Scaling and ReLU algorithms. * Add tensor element-wise scale and type cast operation. * Reduction implemented but does not work. * Exploration of Reduction functionality. * Completed example for Convolution scaled with ReLu activation and AMAX reduction. * WIP: Add required instances for convolution. * WIP: Create client example. Implement convolution stage. * Add elementwise instances. * Add elementwise scale + convert example. * Add reduction instances. * WIP: Client example for AMAX reduction. * WIP: Add instances for multistage reduction. * WIP: Implementation of multistage reduction. * Refactoring. * Clean up. * Guard off FP8 instances when the data type is not available. * Improve output readability. * Addressing reviewer's comments. --- .../24_grouped_conv_activation/CMakeLists.txt | 8 + .../common.hpp | 835 ++++++++++++++++++ .../conv3d_fwd_convscale_relu_amax_fp8.cpp | 58 ++ example/62_convnd_activ/CMakeLists.txt | 1 + .../convscale_reduce/CMakeLists.txt | 11 + .../convnd_fwd_convscale_reduce_common.hpp | 502 +++++++++++ ...convnd_fwd_xdl_convscale_relu_amax_fp8.cpp | 82 ++ .../run_convnd_fwd_example.inc | 98 ++ ...ped_conv_fwd_xdl_outelementop_instance.hpp | 37 + ...ped_convolution_forward_convscale_relu.hpp | 84 ++ .../gpu/permute_scale.hpp | 13 + .../device_permute_scale_instances.hpp | 46 + ...ce_instance_blockwise_f32_f32_f32_amax.hpp | 27 +- ..._relu_ndhwgc_gkzyxc_ndhwgk_f8_instance.cpp | 49 + .../gpu/permute_scale/CMakeLists.txt | 5 +- ...ce_permute_scale_6d_fp32_fp8_instances.cpp | 28 + ...ce_instance_blockwise_f32_f32_f32_amax.cpp | 27 +- 17 files changed, 1891 insertions(+), 20 deletions(-) create mode 100644 client_example/24_grouped_conv_activation/grouped_convnd_fwd_convscale_reduce/common.hpp create mode 100644 client_example/24_grouped_conv_activation/grouped_convnd_fwd_convscale_reduce/conv3d_fwd_convscale_relu_amax_fp8.cpp create mode 100644 example/62_convnd_activ/convscale_reduce/CMakeLists.txt create mode 100644 example/62_convnd_activ/convscale_reduce/convnd_fwd_convscale_reduce_common.hpp create mode 100644 example/62_convnd_activ/convscale_reduce/convnd_fwd_xdl_convscale_relu_amax_fp8.cpp create mode 100644 example/62_convnd_activ/convscale_reduce/run_convnd_fwd_example.inc create mode 100644 library/src/tensor_operation_instance/gpu/permute_scale/device_permute_scale_6d_fp32_fp8_instances.cpp diff --git a/client_example/24_grouped_conv_activation/CMakeLists.txt b/client_example/24_grouped_conv_activation/CMakeLists.txt index ae9b33b94e..60f4ee41f7 100644 --- a/client_example/24_grouped_conv_activation/CMakeLists.txt +++ b/client_example/24_grouped_conv_activation/CMakeLists.txt @@ -47,6 +47,14 @@ target_link_libraries(client_conv3d_fwd_convscale_add_fp8 PRIVATE composable_ker add_executable(client_conv3d_fwd_convscale_relu_fp8 grouped_convnd_fwd_convscale_relu/conv3d_fwd_convscale_relu_fp8.cpp) target_link_libraries(client_conv3d_fwd_convscale_relu_fp8 PRIVATE composable_kernel::device_conv_operations) +# Fwd convscale + ReLU + AMAX +add_executable(client_conv3d_fwd_convscale_relu_amax_fp8 + grouped_convnd_fwd_convscale_reduce/conv3d_fwd_convscale_relu_amax_fp8.cpp) +target_link_libraries(client_conv3d_fwd_convscale_relu_amax_fp8 + PRIVATE composable_kernel::device_conv_operations + composable_kernel::device_other_operations + composable_kernel::device_reduction_operations + utility) # Fwd convscale add_executable(client_conv3d_fwd_convscale_fp8 grouped_convnd_fwd_convscale/conv3d_fwd_convscale_fp8.cpp) diff --git a/client_example/24_grouped_conv_activation/grouped_convnd_fwd_convscale_reduce/common.hpp b/client_example/24_grouped_conv_activation/grouped_convnd_fwd_convscale_reduce/common.hpp new file mode 100644 index 0000000000..b76c5191ec --- /dev/null +++ b/client_example/24_grouped_conv_activation/grouped_convnd_fwd_convscale_reduce/common.hpp @@ -0,0 +1,835 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved. + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "ck/ck.hpp" +#include "ck/library/utility/algorithm.hpp" +#include "ck/tensor_operation/gpu/device/device_elementwise.hpp" +#include "ck/tensor_operation/gpu/device/device_reduce.hpp" +#include "ck/tensor_operation/gpu/device/reduction_operator_mapping.hpp" +#include "ck/tensor_operation/gpu/element/combined_element_wise_operation.hpp" +#include "ck/tensor_operation/gpu/element/unary_element_wise_operation.hpp" +#include "ck/utility/tuple.hpp" +#include "ck/utility/type.hpp" +#include "ck/library/tensor_operation_instance/gpu/grouped_convolution_forward_convscale_relu.hpp" +#include "ck/utility/reduction_enums.hpp" +#include "ck/library/tensor_operation_instance/gpu/permute_scale.hpp" +#include "ck/library/tensor_operation_instance/gpu/reduce/reduce.hpp" +#include "ck/library/utility/host_tensor.hpp" + +namespace ew = ck::tensor_operation::element_wise; + +using PassThrough = ew::PassThrough; +using ConvScaleRelu = ew::UnaryCombinedOp; +using ConvScale = ew::UnaryCombinedOp; + +struct SimpleDeviceMem +{ + SimpleDeviceMem() = delete; + + SimpleDeviceMem(std::size_t mem_size) : p_mem_{} + { + (void)hipMalloc(static_cast(&p_mem_), mem_size); + } + + void* GetDeviceBuffer() { return p_mem_; } + + ~SimpleDeviceMem() { (void)hipFree(p_mem_); } + + void* p_mem_; +}; + +template +std::size_t +GetFlops(const std::array& output_lengths, + const std::array& weights_lengths, + const std::size_t& ds_size) +{ + // 2 * G * N * K * C * * + + // + ds_size * => + // => * ( 2 * C * + ds_size) => + // => G * N * K * * (2 * C * + + // ds_size) + ck::index_t G = weights_lengths[0]; + ck::index_t N = output_lengths[1]; + ck::index_t K = weights_lengths[1]; + ck::index_t C = weights_lengths[2]; + + return G * N * K * + std::accumulate(std::next(std::begin(output_lengths), NumNonSpatialDim), + std::end(output_lengths), + static_cast(1), + std::multiplies<>()) * + (ds_size + static_cast(2) * C * + std::accumulate(std::next(std::begin(weights_lengths), NumNonSpatialDim), + std::end(weights_lengths), + static_cast(1), + std::multiplies<>())); +} + +template +std::size_t GetTensorSize(const std::array& lengths) +{ + + return std::accumulate(std::begin(lengths), + std::end(lengths), + static_cast(1), + std::multiplies()); +} + +template +std::size_t +GetInputByte(const std::array& input_lengths) +{ + // sizeof(InDataType) * (G * N * C * ) + + return sizeof(InDataType) * GetTensorSize(input_lengths); +} + +template +std::size_t +GetWeightByte(const std::array& weights_lengths) +{ + // sizeof(WeiDataType) * (G * K * C * ) + + return sizeof(WeiDataType) * GetTensorSize(weights_lengths); +} + +template +std::size_t +GetOutputByte(const std::array& output_lengths) +{ + // sizeof(OutDataType) * (G * N * K * ); + return sizeof(OutDataType) * GetTensorSize(output_lengths); +} + +template +bool ConvolutionScale(SimpleDeviceMem& in, + SimpleDeviceMem& wei, + SimpleDeviceMem& out, + ConvElementOp elementwise_op, + const std::array& in_lengths, + const std::array& in_strides, + const std::array& wei_lengths, + const std::array& wei_strides, + const std::array& out_lengths, + const std::array& out_strides); + +template +bool TensorScaleConvert(SimpleDeviceMem& in, + SimpleDeviceMem& out, + float scale_out, + const std::array& lengths, + const std::array& strides); + +template +bool TensorFullReduction(SimpleDeviceMem& tensor, + SimpleDeviceMem& out_amax, + const std::array& lengths, + const std::array& strides); + +template +bool run_grouped_conv_fwd_convscale_reduce( + std::array in_lengths, + std::array wei_lengths, + std::array out_lengths) +{ + + namespace ctc = ck::tensor_layout::convolution; + static_assert(NumDimSpatial == 3 && ck::is_same_v && + ck::is_same_v && + ck::is_same_v, + "Unsupported configuration"); + + const ck::index_t G = in_lengths[4]; + const ck::index_t N = in_lengths[0]; + const ck::index_t K = wei_lengths[1]; + const ck::index_t C = in_lengths[5]; + const ck::index_t Z = wei_lengths[2]; + const ck::index_t Y = wei_lengths[3]; + const ck::index_t X = wei_lengths[4]; + const ck::index_t Di = in_lengths[1]; + const ck::index_t Hi = in_lengths[2]; + const ck::index_t Wi = in_lengths[3]; + const ck::index_t Do = out_lengths[1]; + const ck::index_t Ho = out_lengths[2]; + const ck::index_t Wo = out_lengths[3]; + + const std::size_t in_mem_size = sizeof(InDataType) * N * Di * Hi * Wi * G * C; + const std::size_t wei_mem_size = sizeof(WeiDataType) * G * K * Z * Y * X * C; + const std::size_t conv_out_mem_size = sizeof(ConvOutDataType) * N * Do * Ho * Wo * G * K; + const std::size_t out_mem_size = sizeof(OutDataType) * N * Do * Ho * Wo * G * K; + + SimpleDeviceMem in(in_mem_size); + SimpleDeviceMem wei(wei_mem_size); + SimpleDeviceMem conv_out(conv_out_mem_size); + SimpleDeviceMem out(out_mem_size); + + float scale_in = float(std::rand()) / float(RAND_MAX); + float scale_wei = float(std::rand()) / float(RAND_MAX); + float scale_out = float(std::rand()) / float(RAND_MAX); + + // We have NDHWGC/GKZYXC/NDHWGK (x, weight, y) in memory space. + // However, CK's API only accepts lengths and strides with order of GNCDHW/GKCZYX/GNKDHW. + // Hence, we need to adjust the order of strides. + const std::array input_lengths{G, N, C, Di, Hi, Wi}; + const std::array input_strides{ + C, Di * Hi * Wi * G * C, 1, Hi * Wi * G * C, Wi * G * C, G * C}; + const std::array weights_lengths{G, K, C, Z, Y, X}; + const std::array weights_strides{ + K * Z * Y * X * C, Z * Y * X * C, 1, Y * X * C, X * C, C}; + const std::array output_lengths{G, N, K, Do, Ho, Wo}; + const std::array output_strides{ + K, Do * Ho * Wo * G * K, 1, Ho * Wo * G * K, Wo * G * K, G * K}; + + /* + * FP8 Convolution with Scaling + */ + std::cout << "\n\nConvolution with scale Benchmarking:" << std::endl; + auto elementwise_op = ConvElementOp{ew::Scale{scale_in}, ew::Scale{scale_wei}, {}}; + auto conv_ok = ConvolutionScale(in, + wei, + conv_out, + elementwise_op, + input_lengths, + input_strides, + weights_lengths, + weights_strides, + output_lengths, + output_strides); + + if(!conv_ok) + return false; + + /* + * Scale with output weight and convert to FP8 + */ + std::cout << "\n\nElement-wise scale + convert Benchmarking:" << std::endl; + auto elem_wise_ok = TensorScaleConvert( + conv_out, out, scale_out, output_lengths, output_strides); + + if(!elem_wise_ok) + return false; + + /* + * Compute AMAX + */ + std::cout << "\n\nAMAX Benchmarking:" << std::endl; + SimpleDeviceMem amax_device(sizeof(ConvOutDataType)); + auto reduction_ok = + TensorFullReduction(conv_out, amax_device, output_lengths, output_strides); + + if(!reduction_ok) + return false; + + return true; +} + +template +bool ConvolutionScale(SimpleDeviceMem& in, + SimpleDeviceMem& wei, + SimpleDeviceMem& out, + ConvElementOp elementwise_op, + const std::array& in_lengths, + const std::array& in_strides, + const std::array& wei_lengths, + const std::array& wei_strides, + const std::array& out_lengths, + const std::array& out_strides) +{ + + const std::array conv_filter_strides{1, 1, 1}; + const std::array conv_filter_dilations{1, 1, 1}; + const std::array input_left_pads{1, 1, 1}; + const std::array input_right_pads{1, 1, 1}; + + const auto in_mem_size = GetInputByte(in_lengths); + const auto wei_mem_size = GetWeightByte(wei_lengths); + const auto out_mem_size = GetOutputByte(out_lengths); + + std::size_t ds_size = 2; // 2 element-wise scale multipliers + if constexpr(ck::is_same_v) + { + ds_size += 1; // +1 element-wise relu + } + std::size_t flop = GetFlops(out_lengths, wei_lengths, ds_size); + std::size_t num_bytes = + in_mem_size + wei_mem_size + sizeof(float) + sizeof(float) + out_mem_size; + + using ConvDeviceOp = + ck::tensor_operation::device::DeviceGroupedConvFwdMultipleABD, + OutLayout, + InDataType, + WeiDataType, + ck::Tuple<>, + OutDataType, + PassThrough, + PassThrough, + ConvElementOp, + AComputeType, + BComputeType>; + // get device op instances + const auto conv_ptrs = ck::tensor_operation::device::instance::DeviceOperationInstanceFactory< + ConvDeviceOp>::GetInstances(); + + std::cout << "found " << conv_ptrs.size() << " instances" << std::endl; + + std::string conv_best_op_name; + int conv_best_op_id = -1; + float conv_best_avg_time = std::numeric_limits::max(); + float conv_best_gb_per_sec = 0; + float conv_best_tflops = 0; + + // profile device operation instances + std::cout << "Run all convolution instances and do timing" << std::endl; + + for(int i = 0; i < conv_ptrs.size(); ++i) + { + auto& op_ptr = conv_ptrs[i]; + auto argument_ptr = op_ptr->MakeArgumentPointer( + in.GetDeviceBuffer(), + wei.GetDeviceBuffer(), + std::array{}, + out.GetDeviceBuffer(), + in_lengths, + in_strides, + wei_lengths, + wei_strides, + std::array, 0>{}, + std::array, 0>{}, + out_lengths, + out_strides, + conv_filter_strides, + conv_filter_dilations, + input_left_pads, + input_right_pads, + PassThrough{}, + PassThrough{}, + elementwise_op); + + auto invoker_ptr = op_ptr->MakeInvokerPointer(); + std::string op_name = op_ptr->GetTypeString(); + + if(op_ptr->IsSupportedArgument(argument_ptr.get())) + { + float avg_time = invoker_ptr->Run(argument_ptr.get(), StreamConfig{nullptr, true}); + + float tflops = static_cast(flop) / 1.E9 / avg_time; + float gb_per_sec = num_bytes / 1.E6 / avg_time; + + std::cout << "Perf: " << std::setw(10) << avg_time << " ms, " << tflops << " TFlops, " + << gb_per_sec << " GB/s, " << op_name << std::endl; + + if(tflops > conv_best_tflops) + { + conv_best_op_id = i; + conv_best_op_name = op_name; + conv_best_avg_time = avg_time; + conv_best_gb_per_sec = gb_per_sec; + conv_best_tflops = tflops; + } + } + else + { + std::cerr << op_name << " does not support this problem" << std::endl; + } + } + + if(conv_best_op_id < 0) + { + std::cerr << "no suitable instance" << std::endl; + return false; + } + + std::cout << "Best Perf: " << std::setw(10) << conv_best_avg_time << " ms, " << conv_best_tflops + << " TFlops, " << conv_best_gb_per_sec << " GB/s, " << conv_best_op_name << std::endl; + + // run the best instance + { + auto& op_ptr = conv_ptrs[conv_best_op_id]; + std::cout << "Run the best instance without timing: " << op_ptr->GetTypeString() + << std::endl; + auto argument_ptr = op_ptr->MakeArgumentPointer( + in.GetDeviceBuffer(), + wei.GetDeviceBuffer(), + std::array{}, + out.GetDeviceBuffer(), + in_lengths, + in_strides, + wei_lengths, + wei_strides, + std::array, 0>{}, + std::array, 0>{}, + out_lengths, + out_strides, + conv_filter_strides, + conv_filter_dilations, + input_left_pads, + input_right_pads, + PassThrough{}, + PassThrough{}, + elementwise_op); + + auto invoker_ptr = op_ptr->MakeInvokerPointer(); + + if(op_ptr->IsSupportedArgument(argument_ptr.get())) + { + invoker_ptr->Run(argument_ptr.get(), StreamConfig{nullptr, false}); + } + + std::cout << "Done" << std::endl; + } + + return true; +} + +template +bool TensorScaleConvert(SimpleDeviceMem& in, + SimpleDeviceMem& out, + float scale_out, + const std::array& lengths, + const std::array& strides) +{ + + const auto tensor_size = GetTensorSize(lengths); + + const std::size_t in_mem_size = sizeof(InDataType) * tensor_size; + const std::size_t out_mem_size = sizeof(OutDataType) * tensor_size; + + std::size_t flop = 2 * tensor_size; // element-wise scale + convert + + std::size_t bytes = + in_mem_size + sizeof(float) + out_mem_size; // read from in, scale, write to out + + using DeviceScaleConvert = + ck::tensor_operation::device::DeviceElementwise, + ck::Tuple, + ew::Scale, + NumDimSpatial + NumNonSpatialDim>; + + // get device op instances + const auto op_ptrs = ck::tensor_operation::device::instance::DeviceOperationInstanceFactory< + DeviceScaleConvert>::GetInstances(); + + std::cout << "found " << op_ptrs.size() << " instances" << std::endl; + + std::string best_op_name; + int best_op_id = -1; + float best_avg_time = std::numeric_limits::max(); + float best_gb_per_sec = 0; + float best_tflops = 0; + + // profile device operation instances + std::cout << "Run all DeviceScaleConvert instances and do timing" << std::endl; + + auto scale_convert = ew::Scale{scale_out}; + + for(int i = 0; i < op_ptrs.size(); ++i) + { + auto& op_ptr = op_ptrs[i]; + auto argument_ptr = op_ptr->MakeArgumentPointer(lengths, + {strides}, + {strides}, + {in.GetDeviceBuffer()}, + {out.GetDeviceBuffer()}, + scale_convert); + + auto invoker_ptr = op_ptr->MakeInvokerPointer(); + std::string op_name = op_ptr->GetTypeString(); + + if(op_ptr->IsSupportedArgument(argument_ptr.get())) + { + float avg_time = invoker_ptr->Run(argument_ptr.get(), StreamConfig{nullptr, true}); + + float tflops = static_cast(flop) / 1.E9 / avg_time; + float gb_per_sec = bytes / 1.E6 / avg_time; + + std::cout << "Perf: " << std::setw(10) << avg_time << " ms, " << tflops << " TFlops, " + << gb_per_sec << " GB/s, " << op_name << std::endl; + + if(tflops > best_tflops) + { + best_op_id = i; + best_op_name = op_name; + best_avg_time = avg_time; + best_gb_per_sec = gb_per_sec; + best_tflops = tflops; + } + } + else + { + std::cerr << op_name << " does not support this problem" << std::endl; + } + } + + if(best_op_id < 0) + { + std::cerr << "no suitable instance found." << std::endl; + return false; + } + else + { + std::cout << "Best Perf: " << std::setw(10) << best_avg_time << " ms, " << best_tflops + << " TFlops, " << best_gb_per_sec << " GB/s, " << best_op_name << std::endl; + + // run the best intance + auto& op_ptr = op_ptrs[best_op_id]; + std::cout << "Run the best instance without timing: " << op_ptr->GetTypeString() + << std::endl; + auto argument_ptr = op_ptr->MakeArgumentPointer(lengths, + {strides}, + {strides}, + {in.GetDeviceBuffer()}, + {out.GetDeviceBuffer()}, + scale_convert); + + auto invoker_ptr = op_ptr->MakeInvokerPointer(); + + if(op_ptr->IsSupportedArgument(argument_ptr.get())) + { + invoker_ptr->Run(argument_ptr.get(), StreamConfig{nullptr, false}); + } + + std::cout << "Done" << std::endl; + } + + return true; +} + +template +bool TensorFullReduction(SimpleDeviceMem& tensor, + SimpleDeviceMem& out_amax, + const std::array& lengths, + const std::array& strides) +{ + const auto spatial_dim_size = std::accumulate(std::next(std::begin(lengths), NumNonSpatialDim), + std::end(lengths), + static_cast(1), + std::multiplies<>()); + const auto tensor_size = GetTensorSize(lengths); + + auto copy = [](const auto& x, auto& y) { ck::ranges::copy(x, y.begin()); }; + + // Get the reduction operation + using ReduceOperation = typename ck::reduce_binary_operator::opType; + using InElementwiseOperation = + typename ck::reduce_unary_operator::InElementwiseOperation; + using AccElementwiseOperation = + typename ck::reduce_unary_operator::AccElementwiseOperation; + + InElementwiseOperation in_elementwise_op; + AccElementwiseOperation acc_elementwise_op; + std::tie(in_elementwise_op, acc_elementwise_op) = + ck::reduce_unary_operator::GetElementwiseOperator( + static_cast(tensor_size)); + + std::array reduce_out_lengths{1}; + std::array reduce_out_strides{1}; + + SimpleDeviceMem partial_reduce_tensor(sizeof(OutDataType) * spatial_dim_size); + std::array reduce_part_lengths; + std::copy(std::next(std::begin(lengths), NumNonSpatialDim), + std::end(lengths), + std::begin(reduce_part_lengths)); + std::array reduce_part_strides; + copy(HostTensorDescriptor(reduce_part_lengths).GetStrides(), reduce_part_strides); + + { + std::cout << "\nReduction of nonspatial dimensions:" << std::endl; + using DeviceOp = + ck::tensor_operation::device::DeviceReduce; // OutputIndex + const auto op_ptrs = ck::tensor_operation::device::instance::DeviceOperationInstanceFactory< + DeviceOp>::GetInstances(); + + std::cout << "found " << op_ptrs.size() << " instances" << std::endl; + + std::string best_op_name; + int best_op_id = -1; + float best_ave_time = std::numeric_limits::max(); + float best_gb_per_sec = 0; + + std::array reduce_dims; + std::iota(reduce_dims.begin(), reduce_dims.end(), 0); // 0,..., NumNonSpatialDim-1 + + ck::index_t num_in_elements = tensor_size; + ck::index_t num_out_elements = spatial_dim_size; + + // profile device operation instances + std::cout << "Run partial reduction and do timing" << std::endl; + + for(int i = 0; i < op_ptrs.size(); ++i) + { + auto& op_ptr = op_ptrs[i]; + + auto argument_ptr = op_ptr->MakeArgumentPointer(lengths, + strides, + reduce_part_lengths, + reduce_part_strides, + reduce_dims, + 1.0, + 0.0, + tensor.GetDeviceBuffer(), + nullptr, + partial_reduce_tensor.GetDeviceBuffer(), + nullptr, + in_elementwise_op, + PassThrough{}); + auto invoker_ptr = op_ptr->MakeInvokerPointer(); + std::string op_name = op_ptr->GetTypeString(); + + if(op_ptr->IsSupportedArgument(argument_ptr.get())) + { + float ave_time = invoker_ptr->Run(argument_ptr.get(), StreamConfig{nullptr, true}); + std::size_t num_bytes = + num_in_elements * sizeof(InDataType) + num_out_elements * sizeof(OutDataType); + + float gb_per_sec = num_bytes / 1.E6 / ave_time; + + std::cout << "Perf: " << std::setw(10) << ave_time << " ms, " << gb_per_sec + << " GB/s, " << op_name << std::endl; + + if(ave_time < best_ave_time) + { + best_op_id = i; + best_op_name = op_name; + best_ave_time = ave_time; + best_gb_per_sec = gb_per_sec; + } + } + else + { + std::cout << op_name << " does not support this problem" << std::endl; + } + } + + if(best_op_id < 0) + { + std::cerr << "no suitable instance found." << std::endl; + return false; + } + else + { + std::cout << "Best Perf: " << best_ave_time << " ms, " << best_gb_per_sec << " GB/s, " + << best_op_name << std::endl; + + // run the best instance + auto& op_ptr = op_ptrs[best_op_id]; + std::cout << "Run the best instance without timing: " << op_ptr->GetTypeString() + << std::endl; + auto argument_ptr = op_ptr->MakeArgumentPointer(lengths, + strides, + reduce_part_lengths, + reduce_part_strides, + reduce_dims, + 1.0, + 0.0, + tensor.GetDeviceBuffer(), + nullptr, + partial_reduce_tensor.GetDeviceBuffer(), + nullptr, + in_elementwise_op, + PassThrough{}); + + auto invoker_ptr = op_ptr->MakeInvokerPointer(); + + if(op_ptr->IsSupportedArgument(argument_ptr.get())) + { + invoker_ptr->Run(argument_ptr.get(), StreamConfig{nullptr, false}); + } + + std::cout << "Done" << std::endl; + } + } + + { + std::cout << "\nReduction of spatial dimensions:" << std::endl; + using DeviceOp = ck::tensor_operation::device::DeviceReduce; // OutputIndex + const auto op_ptrs = ck::tensor_operation::device::instance::DeviceOperationInstanceFactory< + DeviceOp>::GetInstances(); + + std::cout << "found " << op_ptrs.size() << " instances" << std::endl; + + std::string best_op_name; + int best_op_id = -1; + float best_ave_time = std::numeric_limits::max(); + float best_gb_per_sec = 0; + + std::array reduce_dims; + std::iota(reduce_dims.begin(), reduce_dims.end(), 0); // 0,..., NumDimSpatial-1 + + ck::index_t num_in_elements = spatial_dim_size; + ck::index_t num_out_elements = 1; + + // profile device operation instances + std::cout << "Run final reduction and do timing" << std::endl; + + for(int i = 0; i < op_ptrs.size(); ++i) + { + auto& op_ptr = op_ptrs[i]; + + auto argument_ptr = op_ptr->MakeArgumentPointer(reduce_part_lengths, + reduce_part_strides, + reduce_out_lengths, + reduce_out_strides, + reduce_dims, + 1.0, + 0.0, + partial_reduce_tensor.GetDeviceBuffer(), + nullptr, + out_amax.GetDeviceBuffer(), + nullptr, + PassThrough{}, + acc_elementwise_op); + auto invoker_ptr = op_ptr->MakeInvokerPointer(); + std::string op_name = op_ptr->GetTypeString(); + + if(op_ptr->IsSupportedArgument(argument_ptr.get())) + { + float ave_time = invoker_ptr->Run(argument_ptr.get(), StreamConfig{nullptr, true}); + + std::size_t num_bytes = + num_in_elements * sizeof(OutDataType) + num_out_elements * sizeof(OutDataType); + + float gb_per_sec = num_bytes / 1.E6 / ave_time; + + std::cout << "Perf: " << std::setw(10) << ave_time << " ms, " << gb_per_sec + << " GB/s, " << op_name << std::endl; + + if(ave_time < best_ave_time) + { + best_op_id = i; + best_op_name = op_name; + best_ave_time = ave_time; + best_gb_per_sec = gb_per_sec; + } + } + else + { + std::cout << op_name << " does not support this problem" << std::endl; + } + } + + if(best_op_id < 0) + { + std::cerr << "no suitable instance found." << std::endl; + return false; + } + else + { + std::cout << "Best Perf: " << best_ave_time << " ms, " << best_gb_per_sec << " GB/s, " + << best_op_name << std::endl; + + // run the best instance + auto& op_ptr = op_ptrs[best_op_id]; + std::cout << "Run the best instance without timing: " << op_ptr->GetTypeString() + << std::endl; + auto argument_ptr = op_ptr->MakeArgumentPointer(reduce_part_lengths, + reduce_part_strides, + reduce_out_lengths, + reduce_out_strides, + reduce_dims, + 1.0, + 0.0, + partial_reduce_tensor.GetDeviceBuffer(), + nullptr, + out_amax.GetDeviceBuffer(), + nullptr, + PassThrough{}, + acc_elementwise_op); + + auto invoker_ptr = op_ptr->MakeInvokerPointer(); + + if(op_ptr->IsSupportedArgument(argument_ptr.get())) + { + invoker_ptr->Run(argument_ptr.get(), StreamConfig{nullptr, false}); + } + + std::cout << "Done" << std::endl; + } + } + + return true; +} diff --git a/client_example/24_grouped_conv_activation/grouped_convnd_fwd_convscale_reduce/conv3d_fwd_convscale_relu_amax_fp8.cpp b/client_example/24_grouped_conv_activation/grouped_convnd_fwd_convscale_reduce/conv3d_fwd_convscale_relu_amax_fp8.cpp new file mode 100644 index 0000000000..182642c030 --- /dev/null +++ b/client_example/24_grouped_conv_activation/grouped_convnd_fwd_convscale_reduce/conv3d_fwd_convscale_relu_amax_fp8.cpp @@ -0,0 +1,58 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved. + +#include "common.hpp" + +#include "ck/ck.hpp" +#include "ck/tensor_operation/gpu/device/tensor_layout.hpp" + +using InDataType = ck::f8_t; +using WeiDataType = ck::f8_t; +using CShuffleDataType = float; +using ConvOutDataType = float; // data type of convolution result +using OutDataType = ck::f8_t; // data type of final result +using AComputeDataType = ck::f8_t; +using BComputeDataType = ck::f8_t; + +using ConvElementOp = ConvScaleRelu; + +using InLayout = ck::tensor_layout::convolution::NDHWGC; +using WeiLayout = ck::tensor_layout::convolution::GKZYXC; +using OutLayout = ck::tensor_layout::convolution::NDHWGK; + +constexpr auto ReduceOpId = ck::ReduceTensorOp::AMAX; + +static constexpr ck::index_t NumDimSpatial = 3; +static constexpr ck::index_t G = 1; +static constexpr ck::index_t N = 64; +static constexpr ck::index_t K = 128; +static constexpr ck::index_t C = 64; +static constexpr ck::index_t Z = 3; +static constexpr ck::index_t Y = 3; +static constexpr ck::index_t X = 3; +static constexpr ck::index_t Di = 28; +static constexpr ck::index_t Hi = 28; +static constexpr ck::index_t Wi = 3; +static constexpr ck::index_t Do = 28; +static constexpr ck::index_t Ho = 28; +static constexpr ck::index_t Wo = 3; + +int main() +{ + return run_grouped_conv_fwd_convscale_reduce( + {N, Di, Hi, Wi, G, C}, {G, K, Z, Y, X, C}, {N, Do, Ho, Wo, G, K}) + ? EXIT_SUCCESS + : EXIT_FAILURE; +} diff --git a/example/62_convnd_activ/CMakeLists.txt b/example/62_convnd_activ/CMakeLists.txt index 7114b4ad6d..ab136d99ba 100644 --- a/example/62_convnd_activ/CMakeLists.txt +++ b/example/62_convnd_activ/CMakeLists.txt @@ -3,6 +3,7 @@ add_subdirectory(convinvscale) add_subdirectory(convscale) add_subdirectory(convscale_relu) add_subdirectory(convscale_add) +add_subdirectory(convscale_reduce) add_subdirectory(multi_AB) add_subdirectory(unary) diff --git a/example/62_convnd_activ/convscale_reduce/CMakeLists.txt b/example/62_convnd_activ/convscale_reduce/CMakeLists.txt new file mode 100644 index 0000000000..b3c6621509 --- /dev/null +++ b/example/62_convnd_activ/convscale_reduce/CMakeLists.txt @@ -0,0 +1,11 @@ +list(APPEND gpu_list gfx908 gfx90a gfx940 gfx941 gfx942) +set(target 0) +foreach(gpu IN LISTS GPU_TARGETS) + if(gpu IN_LIST gpu_list AND target EQUAL 0) + add_custom_target(example_convnd_activ_xdl_convscale_reduce) + add_example_executable(example_convnd_fwd_xdl_convscale_relu_amax_fp8 convnd_fwd_xdl_convscale_relu_amax_fp8.cpp) + add_example_dependencies(example_convnd_activ_xdl_convscale_reduce example_convnd_fwd_xdl_convscale_relu_amax_fp8 ) + + set(target 1) + endif() +endforeach() diff --git a/example/62_convnd_activ/convscale_reduce/convnd_fwd_convscale_reduce_common.hpp b/example/62_convnd_activ/convscale_reduce/convnd_fwd_convscale_reduce_common.hpp new file mode 100644 index 0000000000..6940c20695 --- /dev/null +++ b/example/62_convnd_activ/convscale_reduce/convnd_fwd_convscale_reduce_common.hpp @@ -0,0 +1,502 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved. + +#include +#include + +#include "ck/ck.hpp" + +#include "ck/library/utility/algorithm.hpp" +#include "ck/library/utility/check_err.hpp" +#include "ck/library/utility/device_memory.hpp" +#include "ck/library/utility/host_tensor.hpp" +#include "ck/library/utility/host_tensor_generator.hpp" +#include "ck/library/utility/convolution_parameter.hpp" +#include "ck/library/reference_tensor_operation/cpu/reference_conv_fwd.hpp" +#include "ck/library/reference_tensor_operation/cpu/reference_reduce.hpp" +#include "ck/tensor_operation/gpu/element/combined_element_wise_operation.hpp" +#include "ck/tensor_operation/gpu/element/unary_element_wise_operation.hpp" +#include "ck/tensor_operation/gpu/device/impl/device_elementwise_dynamic_vector_dims_impl.hpp" +#include "ck/tensor_operation/gpu/device/impl/device_reduce_multiblock.hpp" +#include "ck/utility/reduction_operator.hpp" +#include "ck/utility/reduction_enums.hpp" +#include "ck/tensor_operation/gpu/device/reduction_operator_mapping.hpp" +#include "ck/utility/type.hpp" + +namespace ew = ck::tensor_operation::element_wise; + +using PassThrough = ew::PassThrough; +using ConvScaleRelu = ew::UnaryCombinedOp; +using ConvScale = ew::UnaryCombinedOp; + +using UnaryScaleConvert = ew::Scale; + +void print_helper_msg() +{ + std::cout << "arg1: verification (0=no, 1=yes)\n" + << "arg2: initialization (0=no init, 1=integer value, 2=decimal value)\n" + << "arg3: time kernel (0=no, 1=yes)\n" + << ck::utils::conv::get_conv_param_parser_helper_msg() << std::endl; +} + +template +inline __host__ __device__ constexpr double get_rtol() +{ + if constexpr(std::is_same_v) + { + return 1e-3; + } + else if constexpr(std::is_same_v) + { + return 1e-6; + } + else if constexpr(std::is_same_v) + { + return 1e-3; + } + else if constexpr(std::is_same_v) + { + return 5e-2; + } + else if constexpr(std::is_same_v) + { + return 1e-1; + } + else if constexpr(std::is_same_v) + { + return 1e-1; + } + else if constexpr(std::is_same_v) + { + return 1e-1; // 240 and 224 are acceptable + } + else if constexpr(std::is_same_v) + { + return 1.5e-1; // 57344 and 49152 are acceptable + } + else + { + return 1e-3; + } +} + +template +inline __host__ __device__ constexpr double get_atol() +{ + if constexpr(std::is_same_v) + { + return 1e-3; + } + else if constexpr(std::is_same_v) + { + return 1e-6; + } + else if constexpr(std::is_same_v) + { + return 1e-3; + } + else if constexpr(std::is_same_v) + { + return 5e-2; + } + else if constexpr(std::is_same_v) + { + return 1e-1; + } + else if constexpr(std::is_same_v) + { + return 1e-1; + } + else if constexpr(std::is_same_v) + { + return 16.1; // 240 and 224 are acceptable + } + else if constexpr(std::is_same_v) + { + return 8192.1; // 57344 and 49152 are acceptable + } + else + { + return 1e-3; + } +} + +template +bool run_grouped_conv_fwd(bool do_verification, + int init_method, + bool time_kernel, + const ck::utils::conv::ConvParam& conv_param, + const HostTensorDescriptor& in_g_n_c_wis_desc, + const HostTensorDescriptor& wei_g_k_c_xs_desc, + const HostTensorDescriptor& out_g_n_k_wos_desc, + const InElementOp& in_element_op, + const WeiElementOp& wei_element_op) +{ + Tensor in(in_g_n_c_wis_desc); + Tensor wei(wei_g_k_c_xs_desc); + Tensor host_conv(out_g_n_k_wos_desc); + Tensor device_conv(out_g_n_k_wos_desc); + Tensor out_host(out_g_n_k_wos_desc); + Tensor out_device(out_g_n_k_wos_desc); + + std::cout << "in: " << in.mDesc << std::endl; + std::cout << "wei: " << wei.mDesc << std::endl; + std::cout << "out: " << out_host.mDesc << std::endl; + + switch(init_method) + { + case 0: break; + case 1: + in.GenerateTensorValue(GeneratorTensor_2{-5, 5}); + wei.GenerateTensorValue(GeneratorTensor_2{-5, 5}); + break; + case 11: // used for debugging + in.GenerateTensorValue(GeneratorTensor_1{1}); + wei.GenerateTensorValue(GeneratorTensor_1{1}); + break; + default: + in.GenerateTensorValue(GeneratorTensor_3{-1.0, 1.0}); + wei.GenerateTensorValue(GeneratorTensor_3{-0.5, 0.5}); + } + + DeviceMem in_device_buf(sizeof(InDataType) * in.mDesc.GetElementSpaceSize()); + DeviceMem wei_device_buf(sizeof(WeiDataType) * wei.mDesc.GetElementSpaceSize()); + DeviceMem conv_device_buf(conv_param.GetOutputByte()); + DeviceMem out_device_buf(conv_param.GetOutputByte()); + + in_device_buf.ToDevice(in.mData.data()); + wei_device_buf.ToDevice(wei.mData.data()); + + std::array a_g_n_c_wis_lengths{}; + std::array a_g_n_c_wis_strides{}; + std::array b_g_k_c_xs_lengths{}; + std::array b_g_k_c_xs_strides{}; + std::array e_g_n_k_wos_lengths{}; + std::array e_g_n_k_wos_strides{}; + std::array conv_filter_strides{}; + std::array conv_filter_dilations{}; + std::array input_left_pads{}; + std::array input_right_pads{}; + + auto copy = [](const auto& x, auto& y) { ck::ranges::copy(x, y.begin()); }; + + copy(in_g_n_c_wis_desc.GetLengths(), a_g_n_c_wis_lengths); + copy(in_g_n_c_wis_desc.GetStrides(), a_g_n_c_wis_strides); + copy(wei_g_k_c_xs_desc.GetLengths(), b_g_k_c_xs_lengths); + copy(wei_g_k_c_xs_desc.GetStrides(), b_g_k_c_xs_strides); + copy(out_g_n_k_wos_desc.GetLengths(), e_g_n_k_wos_lengths); + copy(out_g_n_k_wos_desc.GetStrides(), e_g_n_k_wos_strides); + copy(conv_param.conv_filter_strides_, conv_filter_strides); + copy(conv_param.conv_filter_dilations_, conv_filter_dilations); + copy(conv_param.input_left_pads_, input_left_pads); + copy(conv_param.input_right_pads_, input_right_pads); + + // random scale values + float scale_in = float(std::rand()) / float(RAND_MAX); + float scale_wei = float(std::rand()) / float(RAND_MAX); + float scale_out = float(std::rand()) / float(RAND_MAX); + + std::cout << std::endl; + std::cout << "scale_in: " << scale_in << std::endl; + std::cout << "scale_wei: " << scale_wei << std::endl; + std::cout << "scale_out: " << scale_out << std::endl; + + // convolution elementwise operation + auto conv_element_op = ConvElementOp{ew::Scale{scale_in}, ew::Scale{scale_wei}, {}}; + auto scale_convert = UnaryScaleConvert{scale_out}; // elementwise scale and type cast + + // do Conv + auto conv = DeviceConvNDFwdInstance{}; + auto conv_invoker = conv.MakeInvoker(); + auto conv_argument = + conv.MakeArgument(in_device_buf.GetDeviceBuffer(), + wei_device_buf.GetDeviceBuffer(), + std::array{}, + conv_device_buf.GetDeviceBuffer(), + a_g_n_c_wis_lengths, + a_g_n_c_wis_strides, + b_g_k_c_xs_lengths, + b_g_k_c_xs_strides, + std::array, 0>{}, + std::array, 0>{}, + e_g_n_k_wos_lengths, + e_g_n_k_wos_strides, + conv_filter_strides, + conv_filter_dilations, + input_left_pads, + input_right_pads, + in_element_op, + wei_element_op, + conv_element_op); + + if(!conv.IsSupportedArgument(conv_argument)) + { + throw std::runtime_error( + "wrong! device_conv with the specified compilation parameters does " + "not support this Conv problem"); + } + + std::string kernels = conv.GetTypeString(); + + float avg_time = conv_invoker.Run(conv_argument, StreamConfig{nullptr, time_kernel}); + + using DeviceElementwiseScale = ck::tensor_operation::device::DeviceElementwiseImpl< + ck::Tuple, // InDataTypeTuple + ck::Tuple, // OutDataTypeTuple + UnaryScaleConvert, // UnaryScaleConvert + NDimSpatial + 3, // NumDim + 256, // BlockSize + 128, // M0PerBlock + 128, // M1PerBlock + 8, // M0PerThread + 8, // M1PerThread + ck::Sequence<1, 0>, // ThreadClusterArrangeOrder + ck::Sequence<8>, // InScalarPerVectorSeq + ck::Sequence<8>>; // OutScalarPerVectorSeq + + auto device_ew_scale = DeviceElementwiseScale{}; + auto scale_invoker = device_ew_scale.MakeInvoker(); + auto scale_argument = device_ew_scale.MakeArgument(e_g_n_k_wos_lengths, + {e_g_n_k_wos_strides}, + {e_g_n_k_wos_strides}, + {conv_device_buf.GetDeviceBuffer()}, + {out_device_buf.GetDeviceBuffer()}, + scale_convert); + + if(!device_ew_scale.IsSupportedArgument(scale_argument)) + { + throw std::runtime_error( + "wrong! DeviceElementwiseScale with the specified compilation parameters does " + "not support this problem"); + } + + kernels += std::string("\n\t\t ") + device_ew_scale.GetTypeString(); + + avg_time += scale_invoker.Run(scale_argument, StreamConfig{nullptr, time_kernel}); + + constexpr auto ReduceOpId = ck::ReduceTensorOp::AMAX; + using ReduceOperation = typename ck::reduce_binary_operator::opType; + using InElementwiseOperation = + typename ck::reduce_unary_operator::InElementwiseOperation; + using AccElementwiseOperation = + typename ck::reduce_unary_operator::AccElementwiseOperation; + using DeviceReduceInstance = + ck::tensor_operation::device::DeviceReduceMultiBlock; // OutDstVectorSize + + std::vector outLengths = {1}; + Tensor amax_host(outLengths); + Tensor amax_from_device(outLengths); + auto amax_host_strides = amax_host.mDesc.GetStrides(); + + std::array reduce_dims; + std::iota(reduce_dims.begin(), reduce_dims.end(), 0); // 0,..., NDimSpatial+3-1 + + std::array reduce_out_lengths{1}; + std::array reduce_out_strides{static_cast(amax_host_strides[0])}; + + DeviceMem amax_device(sizeof(ConvOutDataType) * amax_host.mDesc.GetElementSpaceSize()); + DeviceMem index_device; + + InElementwiseOperation in_elementwise_op; + AccElementwiseOperation acc_elementwise_op; + std::tie(in_elementwise_op, acc_elementwise_op) = + ck::reduce_unary_operator::GetElementwiseOperator( + static_cast(host_conv.mDesc.GetElementSize())); + + // Hack convolution output strides for reduction as kernel expects stride 1 for the last + // dimension. It only works because the reduction is done on the whole tensor and result is + // independent of the order of elements. + std::array reduction_strides{}; + copy(HostTensorDescriptor(e_g_n_k_wos_lengths).GetStrides(), reduction_strides); + + auto device_reduce = DeviceReduceInstance{}; + auto reduce_invoker = device_reduce.MakeInvokerPointer(); + auto reduce_argument = device_reduce.MakeArgumentPointer(e_g_n_k_wos_lengths, + reduction_strides, + reduce_out_lengths, + reduce_out_strides, + reduce_dims, + 1.0, + 0.0, + conv_device_buf.GetDeviceBuffer(), + nullptr, + amax_device.GetDeviceBuffer(), + nullptr, + in_elementwise_op, + acc_elementwise_op); + + if(!device_reduce.IsSupportedArgument(reduce_argument.get())) + { + throw std::runtime_error( + "wrong! DeviceReduceInstance with the specified compilation parameters does " + "not support this runtime parameters!"); + }; + + kernels += std::string("\n\t\t ") + device_reduce.GetTypeString(); + + float reduce_time = + reduce_invoker->Run(reduce_argument.get(), StreamConfig{nullptr, time_kernel}); + + if(time_kernel) + std::cout << "\nReduce time: " << reduce_time << " ms" << std::endl; + + avg_time += reduce_time; + + std::size_t flop = conv_param.GetFlops(); // convolution FLOPs + auto conv_out_elems = host_conv.GetElementSize(); // number of elements in conv result tensor + + // 3 element-wise scale multipliers + 1 AMAX + std::size_t elementwise_ops = 3 + 1; + if constexpr(ck::is_same_v) + { + elementwise_ops += 1; // +1 element-wise relu + } + + flop += elementwise_ops * conv_out_elems; + + // convolution + elementwise scaling (in + wei + output byte count) + std::size_t num_btype = conv_param.GetByte(); + num_btype += sizeof(float) + sizeof(float); // + 2 scales + + // elementwise scaling + F8 conversion + num_btype += conv_param.GetOutputByte() + sizeof(float) + + conv_param.GetOutputByte(); + + // AMAX + num_btype += conv_param.GetOutputByte() + sizeof(float); + + if(time_kernel) + { + float tflops = static_cast(flop) / 1.E9 / avg_time; + float gb_per_sec = num_btype / 1.E6 / avg_time; + std::cout << "Perf: " << avg_time << " ms, " << tflops << " TFlops, " << gb_per_sec + << " GB/s, " << std::endl; + } + + std::cout << "\nKernels: " << kernels << std::endl; + + if(do_verification) + { + auto ref_conv = ck::tensor_operation::host::ReferenceConvFwd(); + + auto ref_invoker = ref_conv.MakeInvoker(); + auto ref_argument = ref_conv.MakeArgument(in, + wei, + host_conv, + conv_param.conv_filter_strides_, + conv_param.conv_filter_dilations_, + conv_param.input_left_pads_, + conv_param.input_right_pads_, + in_element_op, + wei_element_op, + conv_element_op); + + ref_invoker.Run(ref_argument); + + conv_device_buf.FromDevice(device_conv.mData.data()); + + out_device_buf.FromDevice(out_device.mData.data()); + + out_host.ForEach([&](auto&, auto idx) { scale_convert(out_host(idx), host_conv(idx)); }); + + std::cout << "\nComparing output to reference: " << std::endl; + auto tight_tol_check = ck::utils::check_err(out_device, out_host, "Error: "); + if(!tight_tol_check) + { + std::cout << "\n\tRecompare applying tolerances...\n"; + std::cout << "\t\trtol = " << get_rtol() << std::endl; + std::cout << "\t\tatol = " << get_atol() << std::endl; + auto loose_tol_check = ck::utils::check_err(out_device, + out_host, + "Error: incorrect convolution results!", + get_rtol(), + get_atol()); + if(!loose_tol_check) + { + return false; + } + } + std::cout << "Success!" << std::endl; + + /// Verify AMAX + + using RefReduceInstance = + ck::tensor_operation::host::ReferenceReduce; + + auto ref_reduce = RefReduceInstance{}; + auto ref_reduce_invoker = ref_reduce.MakeInvokerPointer(); + auto ref_reduce_argument = ref_reduce.MakeArgumentPointer(e_g_n_k_wos_lengths, + e_g_n_k_wos_strides, + reduce_out_lengths, + reduce_out_strides, + reduce_dims, + 1.0, + 0.0, + host_conv.mData.data(), + nullptr, + amax_host.mData.data(), + nullptr, + in_elementwise_op, + acc_elementwise_op); + + if(!ref_reduce.IsSupportedArgument(ref_reduce_argument.get())) + { + throw std::runtime_error( + "wrong! RefReduceInstance with the specified compilation parameters does " + "not support this runtime parameters!"); + }; + + ref_reduce_invoker->Run(ref_reduce_argument.get()); + + amax_device.FromDevice(amax_from_device.mData.data()); + + std::cout << "\namax: " << amax_from_device.mData[0] << std::endl; + std::cout << "amax_ref: " << amax_host.mData[0] << std::endl; + + return ck::utils::check_err(amax_from_device, amax_host, "Error: incorrect AMAX results!"); + } + + return true; +} diff --git a/example/62_convnd_activ/convscale_reduce/convnd_fwd_xdl_convscale_relu_amax_fp8.cpp b/example/62_convnd_activ/convscale_reduce/convnd_fwd_xdl_convscale_relu_amax_fp8.cpp new file mode 100644 index 0000000000..df6bf7bd5c --- /dev/null +++ b/example/62_convnd_activ/convscale_reduce/convnd_fwd_xdl_convscale_relu_amax_fp8.cpp @@ -0,0 +1,82 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved. + +#include "convnd_fwd_convscale_reduce_common.hpp" + +#include "ck/tensor_operation/gpu/device/impl/device_grouped_conv_fwd_multiple_abd_xdl_cshuffle.hpp" + +using InDataType = ck::f8_t; +using WeiDataType = ck::f8_t; +using AccDataType = float; +using CShuffleDataType = float; +using ConvOutDataType = float; // data type of convolution result +using OutDataType = ck::f8_t; // data type of final result +using AComputeDataType = ck::f8_t; +using BComputeDataType = ck::f8_t; + +template +using S = ck::Sequence; + +using InElementOp = PassThrough; +using WeiElementOp = PassThrough; +using OutElementOp = ConvScaleRelu; + +static constexpr auto ConvSpec = + ck::tensor_operation::device::ConvolutionForwardSpecialization::Default; + +static constexpr auto GemmSpec = ck::tensor_operation::device::GemmSpecialization::MNKPadding; + +template +using DeviceGroupedConvNDFwdInstance = + ck::tensor_operation::device::DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle< + NDimSpatial, + InLayout, + WeiLayout, + ck::Tuple<>, + OutLayout, + InDataType, + WeiDataType, + AccDataType, + CShuffleDataType, + ck::Tuple<>, + ConvOutDataType, + InElementOp, + WeiElementOp, + OutElementOp, + ConvSpec, // ConvForwardSpecialization + GemmSpec, // GemmSpecialization + 1, // + 256, // BlockSize + 128, // MPerBlock + 256, // NPerBlock + 32, // KPerBlock + 8, // AK1 + 8, // BK1 + 32, // MPerXdl + 32, // NPerXdl + 2, // MXdlPerWave + 4, // NXdlPerWave + S<4, 64, 1>, // ABlockTransferThreadClusterLengths_AK0_M_AK1 + S<1, 0, 2>, // ABlockTransferThreadClusterArrangeOrder + S<1, 0, 2>, // ABlockTransferSrcAccessOrder + 2, // ABlockTransferSrcVectorDim + 8, // ABlockTransferSrcScalarPerVector + 8, // ABlockTransferDstScalarPerVector_AK1 + 1, // 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 + 1, + 1, + S<1, 32, 1, 8>, + 8, + AComputeDataType, + BComputeDataType>; + +#include "run_convnd_fwd_example.inc" + +int main(int argc, char* argv[]) { return run_convnd_fwd_example(argc, argv) ? 0 : 1; } diff --git a/example/62_convnd_activ/convscale_reduce/run_convnd_fwd_example.inc b/example/62_convnd_activ/convscale_reduce/run_convnd_fwd_example.inc new file mode 100644 index 0000000000..24775f21b5 --- /dev/null +++ b/example/62_convnd_activ/convscale_reduce/run_convnd_fwd_example.inc @@ -0,0 +1,98 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved. + +#pragma once + +bool run_convnd_fwd_example(int argc, char* argv[]) +{ + print_helper_msg(); + + bool do_verification = true; + int init_method = 1; + bool time_kernel = false; + + ck::utils::conv::ConvParam conv_param{ + 2, 1, 128, 256, 192, {3, 3}, {71, 71}, {2, 2}, {1, 1}, {1, 1}, {1, 1}}; + + if(argc == 1) + { + // use default + } + else if(argc == 4) + { + do_verification = std::stoi(argv[1]); + init_method = std::stoi(argv[2]); + time_kernel = std::stoi(argv[3]); + } + else + { + do_verification = std::stoi(argv[1]); + init_method = std::stoi(argv[2]); + time_kernel = std::stoi(argv[3]); + const ck::index_t num_dim_spatial = std::stoi(argv[4]); + + conv_param = ck::utils::conv::parse_conv_param(num_dim_spatial, 5, argv); + } + + // instantiate in and wei element ops, will + // instantiate out_element_op below for every iteration + const auto in_element_op = InElementOp{}; + const auto wei_element_op = WeiElementOp{}; + + const auto run = [&](auto ndim_spatial, auto in_layout, auto wei_layout, auto out_layout) { + constexpr ck::index_t ndim_spatial_value = ndim_spatial.value; + + using InLayout = decltype(in_layout); + using WeiLayout = decltype(wei_layout); + using OutLayout = decltype(out_layout); + + const auto in_g_n_c_wis_desc = + ck::utils::conv::make_input_host_tensor_descriptor_g_n_c_wis_packed( + conv_param); + + const auto wei_g_k_c_xs_desc = + ck::utils::conv::make_weight_host_tensor_descriptor_g_k_c_xs_packed( + conv_param); + + const auto out_g_n_k_wos_desc = + ck::utils::conv::make_output_host_tensor_descriptor_g_n_k_wos_packed( + conv_param); + + return run_grouped_conv_fwd< + ndim_spatial_value, + InDataType, + WeiDataType, + ConvOutDataType, + OutDataType, + InElementOp, + WeiElementOp, + OutElementOp, + DeviceGroupedConvNDFwdInstance>( + do_verification, + init_method, + time_kernel, + conv_param, + in_g_n_c_wis_desc, + wei_g_k_c_xs_desc, + out_g_n_k_wos_desc, + in_element_op, + wei_element_op); + }; + + namespace ctc = ck::tensor_layout::convolution; + + if(conv_param.num_dim_spatial_ == 1) + { + return run(ck::Number<1>{}, ctc::GNWC{}, ctc::GKXC{}, ctc::GNWK{}); + } + else if(conv_param.num_dim_spatial_ == 2) + { + return run(ck::Number<2>{}, ctc::GNHWC{}, ctc::GKYXC{}, ctc::GNHWK{}); + } + else if(conv_param.num_dim_spatial_ == 3) + { + return run(ck::Number<3>{}, ctc::GNDHWC{}, ctc::GKZYXC{}, ctc::GNDHWK{}); + } + + return true; +} diff --git a/library/include/ck/library/tensor_operation_instance/gpu/grouped_conv_fwd/device_grouped_conv_fwd_xdl_outelementop_instance.hpp b/library/include/ck/library/tensor_operation_instance/gpu/grouped_conv_fwd/device_grouped_conv_fwd_xdl_outelementop_instance.hpp index e3bec17514..89e9b2e763 100644 --- a/library/include/ck/library/tensor_operation_instance/gpu/grouped_conv_fwd/device_grouped_conv_fwd_xdl_outelementop_instance.hpp +++ b/library/include/ck/library/tensor_operation_instance/gpu/grouped_conv_fwd/device_grouped_conv_fwd_xdl_outelementop_instance.hpp @@ -184,6 +184,43 @@ using device_grouped_conv_fwd_xdl_outelementop_bf8_f8_instances = std::tuple< // clang-format on >; +template +using device_grouped_conv_fwd_xdl_outelementop_f8_f8_f32_instances = std::tuple< +// clang-format off + //########################################| NumDim| A| B| Ds| E| AData| BData| AccData| CShuffle| Ds| EData| A| B| CDE| ConvForward| GEMM| NumGemmK| Block| MPer| NPer| KPer| AK1| BK1| MPer| NPer| MXdl| NXdl| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockLds| BBlockTransfer| BBlockTransfer| BBlockTransfer| BlockTransfer| BBlockTransfer| BBlockTransfer| BBlockLds| CShuffle| CShuffle| CBlockTransferClusterLengths| CBlockTransfer| Compute| Compute| + //########################################| Spatial| Layout| Layout| Layout| Layout| Type| Type| Type| DataType| DataType| Type| Elementwise| Elementwise| Elementwise| Specialization| Specialization| Prefetch| 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_MWaveMPerXdl| ScalarPerVector| TypeA| TypeB| + //########################################| | | | | | | | | | | | Operation| Operation| Operation| | | Stage| | | | | | | | | Wave| Wave| Lengths_K0_M_K1| ArrangeOrder| | | PerVector| PerVector_K1| | Lengths_K0_N_K1| ArrangeOrder| | | PerVector| PerVector_K1| | PerShuffle| PerShuffle| _NBlock_NWaveNPerXdl| _NWaveNPerXdl| | | + //########################################| | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | +#ifdef CK_ENABLE_FP8 + // generic instance + DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle, F32, PassThrough, PassThrough, OutElementOp, ConvSpec, GemmMNKPadding, 1, 64, 64, 64, 32, 8, 8, 32, 32, 2, 2, S<4, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 1, 8, 1, S<4, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 1, 8, 1, 1, 1, S<1, 16, 1, 4>, 1, F8, F8>, + // instances for small conv.K and conv.C + DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle, F32, PassThrough, PassThrough, OutElementOp, ConvSpec, GemmMNKPadding, 1, 64, 64, 32, 32, 8, 8, 32, 32, 2, 1, S<4, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S<4, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, 1, 1, S<1, 16, 1, 4>, 1, F8, F8>, + DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle, F32, PassThrough, PassThrough, OutElementOp, ConvSpec, GemmMNKPadding, 1, 256, 128, 128, 32, 8, 8, 32, 32, 2, 2, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 1, 8, 1, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 1, 8, 1, 1, 1, S<1, 32, 1, 8>, 8, F8, F8>, + + DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle, F32, PassThrough, PassThrough, OutElementOp, ConvSpec, GemmMNKPadding, 1, 256, 256, 128, 32, 8, 8, 32, 32, 4, 2, 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, 1, S<1, 32, 1, 8>, 8, F8, F8>, + DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle, F32, PassThrough, PassThrough, OutElementOp, ConvSpec, GemmMNKPadding, 1, 256, 128, 256, 32, 8, 8, 32, 32, 2, 4, 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, 1, S<1, 32, 1, 8>, 8, F8, F8>, + DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle, F32, PassThrough, PassThrough, OutElementOp, ConvSpec, GemmMNKPadding, 1, 128, 128, 128, 32, 8, 8, 32, 32, 4, 2, S<4, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S<4, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, 1, 1, S<1, 16, 1, 8>, 8, F8, F8>, + DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle, F32, PassThrough, PassThrough, OutElementOp, ConvSpec, GemmMNKPadding, 1, 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, 1, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, 1, 1, S<1, 32, 1, 8>, 8, F8, F8>, + DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle, F32, PassThrough, PassThrough, OutElementOp, ConvSpec, GemmMNKPadding, 1, 128, 128, 64, 32, 8, 8, 32, 32, 2, 2, S<4, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S<4, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, 1, 1, S<1, 32, 1, 4>, 8, F8, F8>, + DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle, F32, PassThrough, PassThrough, OutElementOp, ConvSpec, GemmMNKPadding, 1, 128, 64, 128, 32, 8, 8, 32, 32, 2, 2, S<4, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S<4, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, 1, 1, S<1, 16, 1, 8>, 8, F8, F8>, + DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle, F32, PassThrough, PassThrough, OutElementOp, ConvSpec, GemmMNKPadding, 1, 64, 64, 64, 32, 8, 8, 32, 32, 2, 2, S<4, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S<4, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, 1, 1, S<1, 16, 1, 4>, 8, F8, F8>, + DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle, F32, PassThrough, PassThrough, OutElementOp, ConvSpec, GemmMNKPadding, 1, 256, 128, 64, 32, 8, 8, 32, 32, 2, 1, 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, 1, S<1, 32, 1, 8>, 8, F8, F8>, + DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle, F32, PassThrough, PassThrough, OutElementOp, ConvSpec, GemmMNKPadding, 1, 256, 64, 128, 32, 8, 8, 32, 32, 1, 2, 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, 1, S<1, 32, 1, 8>, 8, F8, F8>, + DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle, F32, PassThrough, PassThrough, OutElementOp, ConvSpec, GemmMNKPadding, 1, 128, 128, 32, 32, 8, 8, 32, 32, 2, 1, S<4, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S<4, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, 1, 1, S<1, 32, 1, 4>, 8, F8, F8>, + DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle, F32, PassThrough, PassThrough, OutElementOp, ConvSpec, GemmMNKPadding, 1, 128, 32, 128, 32, 8, 8, 32, 32, 1, 2, S<4, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S<4, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, 1, 1, S<1, 16, 1, 8>, 8, F8, F8>, + DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle, F32, PassThrough, PassThrough, OutElementOp, ConvSpec, GemmMNKPadding, 1, 64, 64, 32, 32, 8, 8, 32, 32, 2, 1, S<4, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S<4, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, 1, 1, S<1, 16, 1, 4>, 8, F8, F8>, + DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle, F32, PassThrough, PassThrough, OutElementOp, ConvSpec, GemmMNKPadding, 1, 64, 32, 64, 32, 8, 8, 32, 32, 1, 2, S<4, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S<4, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, 1, 1, S<1, 16, 1, 4>, 8, F8, F8> +#endif + // clang-format on + >; + } // namespace instance } // namespace device } // namespace tensor_operation diff --git a/library/include/ck/library/tensor_operation_instance/gpu/grouped_convolution_forward_convscale_relu.hpp b/library/include/ck/library/tensor_operation_instance/gpu/grouped_convolution_forward_convscale_relu.hpp index ad86d066f7..419f5a609a 100644 --- a/library/include/ck/library/tensor_operation_instance/gpu/grouped_convolution_forward_convscale_relu.hpp +++ b/library/include/ck/library/tensor_operation_instance/gpu/grouped_convolution_forward_convscale_relu.hpp @@ -8,6 +8,7 @@ #include "ck/ck.hpp" #include "ck/tensor_operation/gpu/device/device_grouped_conv_fwd_multiple_abd.hpp" +#include "ck/tensor_operation/gpu/element/combined_element_wise_operation.hpp" #include "ck/tensor_operation/gpu/element/unary_element_wise_operation.hpp" #include "ck/library/tensor_operation_instance/device_operation_instance_factory.hpp" @@ -99,6 +100,89 @@ struct DeviceOperationInstanceFactory< } }; +namespace ew = ck::tensor_operation::element_wise; +using CombConvScaleRelu = ew::UnaryCombinedOp; + +#ifdef CK_ENABLE_FP8 +void add_device_grouped_conv3d_fwd_xdl_combconvscale_relu_ndhwgc_gkzyxc_ndhwgk_f8_f8_f32_instances( + std::vector, + NDHWGK, + F8, + F8, + ck::Tuple<>, + F32, + PassThrough, + PassThrough, + CombConvScaleRelu, + F8, + F8>>>& instances); +#endif + +template +struct DeviceOperationInstanceFactory< + ck::tensor_operation::device::DeviceGroupedConvFwdMultipleABD> +{ + using DeviceOp = DeviceGroupedConvFwdMultipleABD; + + static auto GetInstances() + { + std::vector> op_ptrs; + if constexpr(NumDimSpatial == 3 && is_same_v && + is_same_v && is_same_v) + { +#ifdef CK_ENABLE_FP8 + if constexpr(is_same_v && is_same_v && + is_same_v && is_same_v && + is_same_v) + { + add_device_grouped_conv3d_fwd_xdl_combconvscale_relu_ndhwgc_gkzyxc_ndhwgk_f8_f8_f32_instances( + op_ptrs); + } +#endif + } + return op_ptrs; + } +}; + } // namespace instance } // namespace device } // namespace tensor_operation diff --git a/library/include/ck/library/tensor_operation_instance/gpu/permute_scale.hpp b/library/include/ck/library/tensor_operation_instance/gpu/permute_scale.hpp index 4f5d022f9c..eb71f9d8e5 100644 --- a/library/include/ck/library/tensor_operation_instance/gpu/permute_scale.hpp +++ b/library/include/ck/library/tensor_operation_instance/gpu/permute_scale.hpp @@ -70,6 +70,12 @@ void add_device_permute_scale_6d_f32_instances( DeviceElementwise, ck::Tuple, element_wise::Scale, 6>>>&); #endif +#ifdef CK_ENABLE_FP8 +void add_device_permute_scale_6d_f32_f8_instances( + std::vector, ck::Tuple, element_wise::Scale, 6>>>&); +#endif + template > && + is_same_v>) + { + add_device_permute_scale_6d_f32_f8_instances(op_ptrs); + } #endif } return op_ptrs; diff --git a/library/include/ck/library/tensor_operation_instance/gpu/permute_scale/device_permute_scale_instances.hpp b/library/include/ck/library/tensor_operation_instance/gpu/permute_scale/device_permute_scale_instances.hpp index 8a22005413..204c9a310d 100644 --- a/library/include/ck/library/tensor_operation_instance/gpu/permute_scale/device_permute_scale_instances.hpp +++ b/library/include/ck/library/tensor_operation_instance/gpu/permute_scale/device_permute_scale_instances.hpp @@ -10,6 +10,7 @@ namespace tensor_operation { namespace device { namespace instance { +using F8 = ck::f8_t; using F16 = ck::half_t; using F32 = float; @@ -183,6 +184,51 @@ using device_permute_scale_f32_instances = std::tuple< DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 32, 32, 16, 4, 4, ck::Sequence<1, 0>, ck::Sequence<1>, ck::Sequence<1>>, DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 32, 16, 32, 4, 4, ck::Sequence<1, 0>, ck::Sequence<1>, ck::Sequence<1>> >; + +#ifdef CK_ENABLE_FP8 +template +using device_permute_scale_f32_f8_instances = std::tuple< + DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 256, 64, 64, 4, 4, ck::Sequence<1, 0>, ck::Sequence<4>, ck::Sequence<4>>, + DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 256, 128, 32, 4, 4, ck::Sequence<1, 0>, ck::Sequence<4>, ck::Sequence<4>>, + DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 256, 32, 128, 4, 4, ck::Sequence<1, 0>, ck::Sequence<4>, ck::Sequence<4>>, + DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 128, 64, 32, 4, 4, ck::Sequence<1, 0>, ck::Sequence<4>, ck::Sequence<4>>, + DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 128, 32, 64, 4, 4, ck::Sequence<1, 0>, ck::Sequence<4>, ck::Sequence<4>>, + DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 128, 16, 128, 4, 4, ck::Sequence<1, 0>, ck::Sequence<4>, ck::Sequence<4>>, + DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 128, 128, 16, 4, 4, ck::Sequence<1, 0>, ck::Sequence<4>, ck::Sequence<4>>, + DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 64, 32, 32, 4, 4, ck::Sequence<1, 0>, ck::Sequence<4>, ck::Sequence<4>>, + DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 64, 16, 64, 4, 4, ck::Sequence<1, 0>, ck::Sequence<4>, ck::Sequence<4>>, + DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 64, 64, 16, 4, 4, ck::Sequence<1, 0>, ck::Sequence<4>, ck::Sequence<4>>, + DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 32, 32, 16, 4, 4, ck::Sequence<1, 0>, ck::Sequence<4>, ck::Sequence<4>>, + DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 32, 16, 32, 4, 4, ck::Sequence<1, 0>, ck::Sequence<4>, ck::Sequence<4>>, + + DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 256, 128, 128, 8, 8, ck::Sequence<1, 0>, ck::Sequence<8>, ck::Sequence<8>>, + DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 256, 256, 64, 8, 8, ck::Sequence<1, 0>, ck::Sequence<8>, ck::Sequence<8>>, + DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 256, 64, 256, 8, 8, ck::Sequence<1, 0>, ck::Sequence<8>, ck::Sequence<8>>, + DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 128, 128, 64, 8, 8, ck::Sequence<1, 0>, ck::Sequence<8>, ck::Sequence<8>>, + DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 128, 64, 128, 8, 8, ck::Sequence<1, 0>, ck::Sequence<8>, ck::Sequence<8>>, + DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 128, 32, 256, 8, 8, ck::Sequence<1, 0>, ck::Sequence<8>, ck::Sequence<8>>, + DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 128, 256, 32, 8, 8, ck::Sequence<1, 0>, ck::Sequence<8>, ck::Sequence<8>>, + DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 64, 64, 64, 8, 8, ck::Sequence<1, 0>, ck::Sequence<8>, ck::Sequence<8>>, + DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 64, 32, 128, 8, 8, ck::Sequence<1, 0>, ck::Sequence<8>, ck::Sequence<8>>, + DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 64, 128, 32, 8, 8, ck::Sequence<1, 0>, ck::Sequence<8>, ck::Sequence<8>>, + DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 32, 64, 32, 8, 8, ck::Sequence<1, 0>, ck::Sequence<8>, ck::Sequence<8>>, + DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 32, 32, 64, 8, 8, ck::Sequence<1, 0>, ck::Sequence<8>, ck::Sequence<8>>, + + DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 256, 64, 64, 4, 4, ck::Sequence<1, 0>, ck::Sequence<1>, ck::Sequence<1>>, + DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 256, 128, 32, 4, 4, ck::Sequence<1, 0>, ck::Sequence<1>, ck::Sequence<1>>, + DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 256, 32, 128, 4, 4, ck::Sequence<1, 0>, ck::Sequence<1>, ck::Sequence<1>>, + DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 128, 64, 32, 4, 4, ck::Sequence<1, 0>, ck::Sequence<1>, ck::Sequence<1>>, + DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 128, 32, 64, 4, 4, ck::Sequence<1, 0>, ck::Sequence<1>, ck::Sequence<1>>, + DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 128, 16, 128, 4, 4, ck::Sequence<1, 0>, ck::Sequence<1>, ck::Sequence<1>>, + DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 128, 128, 16, 4, 4, ck::Sequence<1, 0>, ck::Sequence<1>, ck::Sequence<1>>, + DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 64, 32, 32, 4, 4, ck::Sequence<1, 0>, ck::Sequence<1>, ck::Sequence<1>>, + DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 64, 16, 64, 4, 4, ck::Sequence<1, 0>, ck::Sequence<1>, ck::Sequence<1>>, + DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 64, 64, 16, 4, 4, ck::Sequence<1, 0>, ck::Sequence<1>, ck::Sequence<1>>, + DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 32, 32, 16, 4, 4, ck::Sequence<1, 0>, ck::Sequence<1>, ck::Sequence<1>>, + DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 32, 16, 32, 4, 4, ck::Sequence<1, 0>, ck::Sequence<1>, ck::Sequence<1>> + >; +#endif // clang-format on } // namespace instance diff --git a/library/include/ck/library/tensor_operation_instance/gpu/reduce/device_reduce_instance_blockwise_f32_f32_f32_amax.hpp b/library/include/ck/library/tensor_operation_instance/gpu/reduce/device_reduce_instance_blockwise_f32_f32_f32_amax.hpp index ec3bc852e8..142d3f4227 100644 --- a/library/include/ck/library/tensor_operation_instance/gpu/reduce/device_reduce_instance_blockwise_f32_f32_f32_amax.hpp +++ b/library/include/ck/library/tensor_operation_instance/gpu/reduce/device_reduce_instance_blockwise_f32_f32_f32_amax.hpp @@ -14,15 +14,24 @@ namespace device { namespace instance { // clang-format off -// InDataType | AccDataType | OutDataType | Rank | NumReduceDim | ReduceOperation | InElementwiseOp | AccElementwiseOp | PropagateNan | UseIndex -extern template void add_device_reduce_instance_blockwise(std::vector>&); -extern template void add_device_reduce_instance_blockwise(std::vector>&); -extern template void add_device_reduce_instance_blockwise(std::vector>&); -extern template void add_device_reduce_instance_blockwise(std::vector>&); -extern template void add_device_reduce_instance_blockwise(std::vector>&); -extern template void add_device_reduce_instance_blockwise(std::vector>&); -extern template void add_device_reduce_instance_blockwise(std::vector>&); -extern template void add_device_reduce_instance_blockwise(std::vector>&); +// InDataType | AccDataType | OutDataType | Rank | NumReduceDim | ReduceOperation | InElementwiseOp | AccElementwiseOp | PropagateNan | UseIndex +extern template void add_device_reduce_instance_blockwise< F32, F32, F32, 4, 3, ReduceAMax, UnaryAbs, PassThrough, false, false>(std::vector>&); +extern template void add_device_reduce_instance_blockwise< F32, F32, F32, 4, 4, ReduceAMax, UnaryAbs, PassThrough, false, false>(std::vector>&); +extern template void add_device_reduce_instance_blockwise< F32, F32, F32, 4, 1, ReduceAMax, UnaryAbs, PassThrough, false, false>(std::vector>&); +extern template void add_device_reduce_instance_blockwise< F32, F32, F32, 2, 1, ReduceAMax, UnaryAbs, PassThrough, false, false>(std::vector>&); +extern template void add_device_reduce_instance_blockwise< F32, F32, F32, 4, 3, ReduceAMax, UnaryAbs, PassThrough, false, true>(std::vector>&); +extern template void add_device_reduce_instance_blockwise< F32, F32, F32, 4, 4, ReduceAMax, UnaryAbs, PassThrough, false, true>(std::vector>&); +extern template void add_device_reduce_instance_blockwise< F32, F32, F32, 4, 1, ReduceAMax, UnaryAbs, PassThrough, false, true>(std::vector>&); +extern template void add_device_reduce_instance_blockwise< F32, F32, F32, 2, 1, ReduceAMax, UnaryAbs, PassThrough, false, true>(std::vector>&); +extern template void add_device_reduce_instance_blockwise< F32, F32, F32, 6, 6, ReduceAMax, UnaryAbs, PassThrough, true, false>(std::vector>&); +extern template void add_device_reduce_instance_blockwise< F32, F32, F32, 5, 5, ReduceAMax, UnaryAbs, PassThrough, true, false>(std::vector>&); +extern template void add_device_reduce_instance_blockwise< F32, F32, F32, 4, 4, ReduceAMax, UnaryAbs, PassThrough, true, false>(std::vector>&); +extern template void add_device_reduce_instance_blockwise< F32, F32, F32, 6, 3, ReduceAMax, UnaryAbs, PassThrough, true, false>(std::vector>&); +extern template void add_device_reduce_instance_blockwise< F32, F32, F32, 5, 3, ReduceAMax, UnaryAbs, PassThrough, true, false>(std::vector>&); +extern template void add_device_reduce_instance_blockwise< F32, F32, F32, 4, 3, ReduceAMax, UnaryAbs, PassThrough, true, false>(std::vector>&); +extern template void add_device_reduce_instance_blockwise< F32, F32, F32, 3, 3, ReduceAMax, PassThrough, PassThrough, true, false>(std::vector>&); +extern template void add_device_reduce_instance_blockwise< F32, F32, F32, 2, 2, ReduceAMax, PassThrough, PassThrough, true, false>(std::vector>&); +extern template void add_device_reduce_instance_blockwise< F32, F32, F32, 1, 1, ReduceAMax, PassThrough, PassThrough, true, false>(std::vector>&); // clang-format on } // namespace instance diff --git a/library/src/tensor_operation_instance/gpu/grouped_conv3d_fwd_convscale_relu/xdl/device_grouped_conv3d_fwd_xdl_convscale_relu_ndhwgc_gkzyxc_ndhwgk_f8_instance.cpp b/library/src/tensor_operation_instance/gpu/grouped_conv3d_fwd_convscale_relu/xdl/device_grouped_conv3d_fwd_xdl_convscale_relu_ndhwgc_gkzyxc_ndhwgk_f8_instance.cpp index 472da0da78..1fda1f4ee6 100644 --- a/library/src/tensor_operation_instance/gpu/grouped_conv3d_fwd_convscale_relu/xdl/device_grouped_conv3d_fwd_xdl_convscale_relu_ndhwgc_gkzyxc_ndhwgk_f8_instance.cpp +++ b/library/src/tensor_operation_instance/gpu/grouped_conv3d_fwd_convscale_relu/xdl/device_grouped_conv3d_fwd_xdl_convscale_relu_ndhwgc_gkzyxc_ndhwgk_f8_instance.cpp @@ -3,6 +3,7 @@ #include "ck/library/tensor_operation_instance/gpu/grouped_conv_fwd/device_grouped_conv_fwd_xdl_outelementop_instance.hpp" #include "ck/library/tensor_operation_instance/add_device_operation_instance.hpp" +#include "ck/tensor_operation/gpu/element/combined_element_wise_operation.hpp" #include "ck/tensor_operation/gpu/element/unary_element_wise_operation.hpp" namespace ck { @@ -57,6 +58,54 @@ void add_device_grouped_conv3d_fwd_xdl_convscale_relu_ndhwgc_gkzyxc_ndhwgk_f8_in ConvScaleRelu>{}); } +namespace ew = ck::tensor_operation::element_wise; +using CombConvScaleRelu = ew::UnaryCombinedOp; + +void add_device_grouped_conv3d_fwd_xdl_combconvscale_relu_ndhwgc_gkzyxc_ndhwgk_f8_f8_f32_instances( + std::vector, + NDHWGK, + F8, + F8, + ck::Tuple<>, + F32, + PassThrough, + PassThrough, + CombConvScaleRelu, + F8, + F8>>>& instances) +{ + add_device_operation_instances( + instances, + device_grouped_conv_fwd_xdl_outelementop_f8_f8_f32_instances<3, + NDHWGC, + GKZYXC, + ck::Tuple<>, + NDHWGK, + ConvFwdDefault, + CombConvScaleRelu>{}); + add_device_operation_instances( + instances, + device_grouped_conv_fwd_xdl_outelementop_f8_f8_f32_instances<3, + NDHWGC, + GKZYXC, + ck::Tuple<>, + NDHWGK, + ConvFwd1x1P0, + CombConvScaleRelu>{}); + add_device_operation_instances( + instances, + device_grouped_conv_fwd_xdl_outelementop_f8_f8_f32_instances<3, + NDHWGC, + GKZYXC, + ck::Tuple<>, + NDHWGK, + ConvFwd1x1S1P0, + CombConvScaleRelu>{}); +} + } // namespace instance } // namespace device } // namespace tensor_operation diff --git a/library/src/tensor_operation_instance/gpu/permute_scale/CMakeLists.txt b/library/src/tensor_operation_instance/gpu/permute_scale/CMakeLists.txt index fc0da56a96..427bf54ca1 100644 --- a/library/src/tensor_operation_instance/gpu/permute_scale/CMakeLists.txt +++ b/library/src/tensor_operation_instance/gpu/permute_scale/CMakeLists.txt @@ -1,4 +1,4 @@ -add_instance_library(device_permute_scale_instance +add_instance_library(device_permute_scale_instance device_permute_scale_1d_fp16_instances.cpp device_permute_scale_2d_fp16_instances.cpp device_permute_scale_3d_fp16_instances.cpp @@ -10,4 +10,5 @@ add_instance_library(device_permute_scale_instance device_permute_scale_3d_fp32_instances.cpp device_permute_scale_4d_fp32_instances.cpp device_permute_scale_5d_fp32_instances.cpp - device_permute_scale_6d_fp32_instances.cpp) + device_permute_scale_6d_fp32_instances.cpp + device_permute_scale_6d_fp32_fp8_instances.cpp) diff --git a/library/src/tensor_operation_instance/gpu/permute_scale/device_permute_scale_6d_fp32_fp8_instances.cpp b/library/src/tensor_operation_instance/gpu/permute_scale/device_permute_scale_6d_fp32_fp8_instances.cpp new file mode 100644 index 0000000000..95d83a5439 --- /dev/null +++ b/library/src/tensor_operation_instance/gpu/permute_scale/device_permute_scale_6d_fp32_fp8_instances.cpp @@ -0,0 +1,28 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved. + +#include "ck/library/tensor_operation_instance/gpu/permute_scale/device_permute_scale_instances.hpp" +#include "ck/library/tensor_operation_instance/add_device_operation_instance.hpp" + +namespace ck { +namespace tensor_operation { +namespace device { +namespace instance { + +using Scale = element_wise::Scale; + +void add_device_permute_scale_6d_f32_f8_instances( + std::vector, ck::Tuple, Scale, 6>>>& + instances) +{ +#ifdef CK_ENABLE_FP8 + add_device_operation_instances(instances, device_permute_scale_f32_f8_instances<6, Scale>{}); +#else + ignore = instances; +#endif +} + +} // namespace instance +} // namespace device +} // namespace tensor_operation +} // namespace ck diff --git a/library/src/tensor_operation_instance/gpu/reduce/device_reduce_instance_blockwise_f32_f32_f32_amax.cpp b/library/src/tensor_operation_instance/gpu/reduce/device_reduce_instance_blockwise_f32_f32_f32_amax.cpp index 17f45c3327..0c071e92f5 100644 --- a/library/src/tensor_operation_instance/gpu/reduce/device_reduce_instance_blockwise_f32_f32_f32_amax.cpp +++ b/library/src/tensor_operation_instance/gpu/reduce/device_reduce_instance_blockwise_f32_f32_f32_amax.cpp @@ -10,15 +10,24 @@ namespace device { namespace instance { // clang-format off -// InDataType | AccDataType | OutDataType | Rank | NumReduceDim | ReduceOperation | InElementwiseOp | AccElementwiseOp | PropagateNan | UseIndex -template void add_device_reduce_instance_blockwise(std::vector>&); -template void add_device_reduce_instance_blockwise(std::vector>&); -template void add_device_reduce_instance_blockwise(std::vector>&); -template void add_device_reduce_instance_blockwise(std::vector>&); -template void add_device_reduce_instance_blockwise(std::vector>&); -template void add_device_reduce_instance_blockwise(std::vector>&); -template void add_device_reduce_instance_blockwise(std::vector>&); -template void add_device_reduce_instance_blockwise(std::vector>&); +// InDataType | AccDataType | OutDataType | Rank | NumReduceDim | ReduceOperation | InElementwiseOp | AccElementwiseOp | PropagateNan | UseIndex +template void add_device_reduce_instance_blockwise< F32, F32, F32, 4, 3, ReduceAMax, UnaryAbs, PassThrough, false, false>(std::vector>&); +template void add_device_reduce_instance_blockwise< F32, F32, F32, 4, 4, ReduceAMax, UnaryAbs, PassThrough, false, false>(std::vector>&); +template void add_device_reduce_instance_blockwise< F32, F32, F32, 4, 1, ReduceAMax, UnaryAbs, PassThrough, false, false>(std::vector>&); +template void add_device_reduce_instance_blockwise< F32, F32, F32, 2, 1, ReduceAMax, UnaryAbs, PassThrough, false, false>(std::vector>&); +template void add_device_reduce_instance_blockwise< F32, F32, F32, 4, 3, ReduceAMax, UnaryAbs, PassThrough, false, true>(std::vector>&); +template void add_device_reduce_instance_blockwise< F32, F32, F32, 4, 4, ReduceAMax, UnaryAbs, PassThrough, false, true>(std::vector>&); +template void add_device_reduce_instance_blockwise< F32, F32, F32, 4, 1, ReduceAMax, UnaryAbs, PassThrough, false, true>(std::vector>&); +template void add_device_reduce_instance_blockwise< F32, F32, F32, 2, 1, ReduceAMax, UnaryAbs, PassThrough, false, true>(std::vector>&); +template void add_device_reduce_instance_blockwise< F32, F32, F32, 6, 6, ReduceAMax, UnaryAbs, PassThrough, true, false>(std::vector>&); +template void add_device_reduce_instance_blockwise< F32, F32, F32, 5, 5, ReduceAMax, UnaryAbs, PassThrough, true, false>(std::vector>&); +template void add_device_reduce_instance_blockwise< F32, F32, F32, 4, 4, ReduceAMax, UnaryAbs, PassThrough, true, false>(std::vector>&); +template void add_device_reduce_instance_blockwise< F32, F32, F32, 6, 3, ReduceAMax, UnaryAbs, PassThrough, true, false>(std::vector>&); +template void add_device_reduce_instance_blockwise< F32, F32, F32, 5, 3, ReduceAMax, UnaryAbs, PassThrough, true, false>(std::vector>&); +template void add_device_reduce_instance_blockwise< F32, F32, F32, 4, 3, ReduceAMax, UnaryAbs, PassThrough, true, false>(std::vector>&); +template void add_device_reduce_instance_blockwise< F32, F32, F32, 3, 3, ReduceAMax, PassThrough, PassThrough, true, false>(std::vector>&); +template void add_device_reduce_instance_blockwise< F32, F32, F32, 2, 2, ReduceAMax, PassThrough, PassThrough, true, false>(std::vector>&); +template void add_device_reduce_instance_blockwise< F32, F32, F32, 1, 1, ReduceAMax, PassThrough, PassThrough, true, false>(std::vector>&); // clang-format on } // namespace instance From dc82daa86e3b882dd2774f915749b17683bbb946 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Bart=C5=82omiej=20Kocot?= Date: Tue, 20 Aug 2024 19:04:14 +0200 Subject: [PATCH 12/20] Convert MIOpen driver to ckProfiler script typos fix (#1476) --- .../include/profiler/profile_grouped_conv_bwd_weight_impl.hpp | 4 ---- script/convert_miopen_driver_to_profiler.py | 2 +- 2 files changed, 1 insertion(+), 5 deletions(-) diff --git a/profiler/include/profiler/profile_grouped_conv_bwd_weight_impl.hpp b/profiler/include/profiler/profile_grouped_conv_bwd_weight_impl.hpp index 5318de5e8b..3758af2477 100644 --- a/profiler/include/profiler/profile_grouped_conv_bwd_weight_impl.hpp +++ b/profiler/include/profiler/profile_grouped_conv_bwd_weight_impl.hpp @@ -250,18 +250,14 @@ bool profile_grouped_conv_bwd_weight_impl(int do_verification, { LogRangeAsType(std::cout << "output : ", output.mData, ",") << std::endl; - ; LogRangeAsType( std::cout << "weight (device): ", weight_device_result.mData, ",") << std::endl; - ; LogRangeAsType( std::cout << "weight (host): ", weight_host_result.mData, ",") << std::endl; - ; LogRangeAsType(std::cout << "input: ", input.mData, ",") << std::endl; - ; } } } diff --git a/script/convert_miopen_driver_to_profiler.py b/script/convert_miopen_driver_to_profiler.py index 47135f3401..06c656c8e5 100644 --- a/script/convert_miopen_driver_to_profiler.py +++ b/script/convert_miopen_driver_to_profiler.py @@ -26,7 +26,7 @@ def run_ck_profiler_cmd(cmd): def parse_data_type(args): if args.data_type == "fp32": if args.ck_profier_op == "grouped_conv_bwd_weight" or \ - args.ck_profier_op == "grouped_conv_bwd_weight" or \ + args.ck_profier_op == "grouped_conv_bwd_data" or \ args.ck_profier_op == "grouped_conv_fwd": args.data_type = 0 if args.data_type == "fp16": From e20f20efbf97db1d360c218ec3c3b355e3c6ac2a Mon Sep 17 00:00:00 2001 From: Rostyslav Geyyer <46627076+geyyer@users.noreply.github.com> Date: Wed, 21 Aug 2024 11:09:48 -0500 Subject: [PATCH 13/20] Set RNE fp8 conversion as a default (#1458) * Set RNE fp8 conversion as a default * Update f8 tests * Disable failing test on gfx11 * Update bf8 tests * Add a flag * Fix the flag * Raise flag for gfx10 as well * Temp commit for tolerance testing * Update tolerances --- example/01_gemm/gemm_xdl_fp8.cpp | 4 +- example/01_gemm/run_gemm_example.inc | 10 ++--- include/ck/ck.hpp | 6 +-- .../include/ck/library/utility/check_err.hpp | 5 ++- test/data_type/CMakeLists.txt | 10 +++++ test/data_type/test_bf8.cpp | 42 ++++++++++--------- test/data_type/test_fp8.cpp | 42 ++++++++++--------- 7 files changed, 69 insertions(+), 50 deletions(-) diff --git a/example/01_gemm/gemm_xdl_fp8.cpp b/example/01_gemm/gemm_xdl_fp8.cpp index 7d8538681b..fe41602301 100644 --- a/example/01_gemm/gemm_xdl_fp8.cpp +++ b/example/01_gemm/gemm_xdl_fp8.cpp @@ -1,5 +1,5 @@ // SPDX-License-Identifier: MIT -// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved. +// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved. #include "common.hpp" @@ -7,7 +7,7 @@ using ADataType = ck::f8_t; using BDataType = ck::f8_t; -using CDataType = ck::half_t; +using CDataType = ck::f8_t; using AccDataType = float; using CShuffleDataType = float; diff --git a/example/01_gemm/run_gemm_example.inc b/example/01_gemm/run_gemm_example.inc index cb15186c3b..0e42c91d41 100644 --- a/example/01_gemm/run_gemm_example.inc +++ b/example/01_gemm/run_gemm_example.inc @@ -1,5 +1,5 @@ // SPDX-License-Identifier: MIT -// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved. +// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved. #pragma once @@ -34,11 +34,11 @@ inline __host__ __device__ constexpr double get_rtol() } else if constexpr(std::is_same_v) { - return 1e-1; // 240 and 224 are acceptable + return 2e-1; } else if constexpr(std::is_same_v) { - return 1.5e-1; // 57344 and 49152 are acceptable + return 2e-1; } else { @@ -75,11 +75,11 @@ inline __host__ __device__ constexpr double get_atol() } else if constexpr(std::is_same_v) { - return 16.1; // 240 and 224 are acceptable + return 2e-1; } else if constexpr(std::is_same_v) { - return 8192.1; // 57344 and 49152 are acceptable + return 2e-1; } else { diff --git a/include/ck/ck.hpp b/include/ck/ck.hpp index 9528a30b4b..5f74d51a65 100644 --- a/include/ck/ck.hpp +++ b/include/ck/ck.hpp @@ -1,5 +1,5 @@ // SPDX-License-Identifier: MIT -// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved. +// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved. #pragma once @@ -153,8 +153,8 @@ CK_DECLARE_ENV_VAR_BOOL(CK_LOGGING) // LDS direct loads using inline assembly #define CK_USE_AMD_LDS_DIRECT_LOAD_INLINE_ASM 0 -// set stochastic rounding as default for f8 conversions -#define CK_USE_SR_F8_CONVERSION 1 +// set rounding to nearest even as default for f8 conversions +#define CK_USE_SR_F8_CONVERSION 0 // block synchronization only s_wait lgkmcnt(0), not vmcnt(0) #define CK_EXPERIMENTAL_BLOCK_SYNC_LDS_WITHOUT_SYNC_VMEM 1 diff --git a/library/include/ck/library/utility/check_err.hpp b/library/include/ck/library/utility/check_err.hpp index 9f4212ebd4..58479f2127 100644 --- a/library/include/ck/library/utility/check_err.hpp +++ b/library/include/ck/library/utility/check_err.hpp @@ -1,5 +1,5 @@ // SPDX-License-Identifier: MIT -// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved. +// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved. #pragma once @@ -272,7 +272,8 @@ check_err(const Range& out, } if(!res) { - std::cerr << std::setw(12) << std::setprecision(7) << "max err: " << max_err << std::endl; + std::cerr << std::setw(12) << std::setprecision(7) << "max err: " << max_err + << " number of errors: " << err_count << std::endl; } return res; } diff --git a/test/data_type/CMakeLists.txt b/test/data_type/CMakeLists.txt index 0ebfc931ac..75e098ce48 100644 --- a/test/data_type/CMakeLists.txt +++ b/test/data_type/CMakeLists.txt @@ -1,3 +1,13 @@ +if (GPU_TARGETS) + if (GPU_TARGETS MATCHES "gfx10" OR GPU_TARGETS MATCHES "gfx11") + add_definitions(-DCK_SKIP_FLAKY_F8_TEST) + set(CK_SKIP_FLAKY_F8_TEST "ON") + endif() +else() + add_definitions(-DCK_SKIP_FLAKY_F8_TEST) + set(CK_SKIP_FLAKY_F8_TEST "ON") +endif() + if (USE_BITINT_EXTENSION_INT4) add_gtest_executable(test_int4 test_int4.cpp) if(result EQUAL 0) diff --git a/test/data_type/test_bf8.cpp b/test/data_type/test_bf8.cpp index 6a5fa281e8..6f50db68c7 100644 --- a/test/data_type/test_bf8.cpp +++ b/test/data_type/test_bf8.cpp @@ -1,11 +1,12 @@ // SPDX-License-Identifier: MIT -// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved. +// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved. #include "gtest/gtest.h" #include "ck/utility/data_type.hpp" #include "ck/utility/type_convert.hpp" using ck::bf8_t; +using ck::f8_convert_rne; using ck::f8_convert_sr; using ck::half_t; using ck::type_convert; @@ -24,33 +25,36 @@ TEST(BF8, ConvertFP32Nearest) // fix the tolerance value float abs_tol = 1e-6; // convert 0 float to bf8 and back, check if holds - ASSERT_NEAR(0.0f, type_convert(type_convert(0.0f)), abs_tol); + ASSERT_NEAR(0.0f, type_convert(f8_convert_rne(0.0f)), abs_tol); + // don't run the next test on gfx11 devices +#ifndef CK_SKIP_FLAKY_F8_TEST // convert minimal float to bf8 and back, check if holds ASSERT_NEAR(std::numeric_limits::min(), - type_convert(type_convert(std::numeric_limits::min())), + type_convert(f8_convert_rne(std::numeric_limits::min())), abs_tol); +#endif // convert maximal bf8_t to float and check if equal to 57344.0 - ASSERT_NEAR(57344.0f, type_convert(type_convert(57344.0f)), abs_tol); + ASSERT_NEAR(57344.0f, type_convert(f8_convert_rne(57344.0f)), abs_tol); // convert maximal float to bf8 and back, check if clipped to 57344.0 ASSERT_NEAR(57344.0f, - type_convert(type_convert(std::numeric_limits::max())), + type_convert(f8_convert_rne(std::numeric_limits::max())), abs_tol); // convert inf float to bf8_t and check if it is qNan ASSERT_NEAR(type_convert(0x80), - type_convert(std::numeric_limits::infinity()), + f8_convert_rne(std::numeric_limits::infinity()), abs_tol); // positive norm float value to bf8 and back, check if holds float pos_float = 0.0000762939f; - ASSERT_NEAR(pos_float, type_convert(type_convert(pos_float)), abs_tol); + ASSERT_NEAR(pos_float, type_convert(f8_convert_rne(pos_float)), abs_tol); // negative norm float value to bf8 and back, check if holds float neg_float = -0.0000610351f; - ASSERT_NEAR(neg_float, type_convert(type_convert(neg_float)), abs_tol); + ASSERT_NEAR(neg_float, type_convert(f8_convert_rne(neg_float)), abs_tol); // positive subnorm float value to bf8 and back, check if holds pos_float = 0.0000305175f; - ASSERT_NEAR(pos_float, type_convert(type_convert(pos_float)), abs_tol); + ASSERT_NEAR(pos_float, type_convert(f8_convert_rne(pos_float)), abs_tol); // negative subnorm float value to bf8 and back, check if holds neg_float = -0.0000152587f; - ASSERT_NEAR(neg_float, type_convert(type_convert(neg_float)), abs_tol); + ASSERT_NEAR(neg_float, type_convert(f8_convert_rne(neg_float)), abs_tol); } TEST(BF8, ConvertFP32Stochastic) @@ -92,34 +96,34 @@ TEST(BF8, ConvertFP16Nearest) // fix the tolerance value float abs_tol = 1e-3; // convert 0 fp16 to bf8 and back, check if holds - ASSERT_NEAR(half_t{0.0}, type_convert(type_convert(half_t{0.0})), abs_tol); + ASSERT_NEAR(half_t{0.0}, type_convert(f8_convert_rne(half_t{0.0})), abs_tol); // convert minimal fp16 to bf8 and back, check if holds ASSERT_NEAR(ck::NumericLimits::Min(), - type_convert(type_convert(ck::NumericLimits::Min())), + type_convert(f8_convert_rne(ck::NumericLimits::Min())), abs_tol); // convert maximal bf8_t to fp16 and check if equal to 57344.0 ASSERT_NEAR( - half_t{57344.0}, type_convert(type_convert(half_t{57344.0})), abs_tol); + half_t{57344.0}, type_convert(f8_convert_rne(half_t{57344.0})), abs_tol); // convert maximal fp16 to bf8 and back, check if clipped to 57344.0 ASSERT_NEAR(half_t{57344.0}, - type_convert(type_convert(ck::NumericLimits::Max())), + type_convert(f8_convert_rne(ck::NumericLimits::Max())), abs_tol); // convert QuietNaN fp16 to bf8_t and check if it is QuietNaN ASSERT_NEAR(type_convert(0x80), - type_convert(ck::NumericLimits::QuietNaN()), + f8_convert_rne(ck::NumericLimits::QuietNaN()), abs_tol); // positive norm fp16 value to bf8 and back, check if holds half_t pos_half = half_t{0.0000762939}; - ASSERT_NEAR(pos_half, type_convert(type_convert(pos_half)), abs_tol); + ASSERT_NEAR(pos_half, type_convert(f8_convert_rne(pos_half)), abs_tol); // negative norm fp16 value to bf8 and back, check if holds half_t neg_half = half_t{-0.0000610351}; - ASSERT_NEAR(neg_half, type_convert(type_convert(neg_half)), abs_tol); + ASSERT_NEAR(neg_half, type_convert(f8_convert_rne(neg_half)), abs_tol); // positive subnorm fp16 value to bf8 and back, check if holds pos_half = half_t{0.0000305175}; - ASSERT_NEAR(pos_half, type_convert(type_convert(pos_half)), abs_tol); + ASSERT_NEAR(pos_half, type_convert(f8_convert_rne(pos_half)), abs_tol); // negative subnorm fp16 value to bf8 and back, check if holds neg_half = half_t{-0.0000152587}; - ASSERT_NEAR(neg_half, type_convert(type_convert(neg_half)), abs_tol); + ASSERT_NEAR(neg_half, type_convert(f8_convert_rne(neg_half)), abs_tol); } TEST(BF8, ConvertFP16Stochastic) diff --git a/test/data_type/test_fp8.cpp b/test/data_type/test_fp8.cpp index 0612a1cf44..25d9d9d2fb 100644 --- a/test/data_type/test_fp8.cpp +++ b/test/data_type/test_fp8.cpp @@ -1,10 +1,11 @@ // SPDX-License-Identifier: MIT -// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved. +// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved. #include "gtest/gtest.h" #include "ck/utility/data_type.hpp" #include "ck/utility/type_convert.hpp" +using ck::f8_convert_rne; using ck::f8_convert_sr; using ck::f8_t; using ck::half_t; @@ -24,33 +25,36 @@ TEST(FP8, ConvertFP32Nearest) // fix the tolerance value float abs_tol = 1e-6; // convert 0 float to fp8 and back, check if holds - ASSERT_NEAR(0.0f, type_convert(type_convert(0.0f)), abs_tol); + ASSERT_NEAR(0.0f, type_convert(f8_convert_rne(0.0f)), abs_tol); + // don't run the next test on gfx11 devices +#ifndef CK_SKIP_FLAKY_F8_TEST // convert minimal float to fp8 and back, check if holds ASSERT_NEAR(std::numeric_limits::min(), - type_convert(type_convert(std::numeric_limits::min())), + type_convert(f8_convert_rne(std::numeric_limits::min())), abs_tol); +#endif // convert maximal f8_t to float and check if equal to 240.0 - ASSERT_NEAR(240.0f, type_convert(type_convert(240.0f)), abs_tol); + ASSERT_NEAR(240.0f, type_convert(f8_convert_rne(240.0f)), abs_tol); // convert maximal float to fp8 and back, check if clipped to 240.0 ASSERT_NEAR(240.0f, - type_convert(type_convert(std::numeric_limits::max())), + type_convert(f8_convert_rne(std::numeric_limits::max())), abs_tol); // convert inf float to f8_t and check if it is qNan ASSERT_NEAR(type_convert(0x80), - type_convert(std::numeric_limits::infinity()), + f8_convert_rne(std::numeric_limits::infinity()), abs_tol); // positive norm float value to fp8 and back, check if holds float pos_float = 0.017578125f; - ASSERT_NEAR(pos_float, type_convert(type_convert(pos_float)), abs_tol); + ASSERT_NEAR(pos_float, type_convert(f8_convert_rne(pos_float)), abs_tol); // negative norm float value to fp8 and back, check if holds float neg_float = -0.015625f; - ASSERT_NEAR(neg_float, type_convert(type_convert(neg_float)), abs_tol); + ASSERT_NEAR(neg_float, type_convert(f8_convert_rne(neg_float)), abs_tol); // positive subnorm float value to fp8 and back, check if holds pos_float = 0.00390625f; - ASSERT_NEAR(pos_float, type_convert(type_convert(pos_float)), abs_tol); + ASSERT_NEAR(pos_float, type_convert(f8_convert_rne(pos_float)), abs_tol); // negative subnorm float value to fp8 and back, check if holds neg_float = -0.001953125f; - ASSERT_NEAR(neg_float, type_convert(type_convert(neg_float)), abs_tol); + ASSERT_NEAR(neg_float, type_convert(f8_convert_rne(neg_float)), abs_tol); } TEST(FP8, ConvertFP32Stochastic) @@ -92,33 +96,33 @@ TEST(FP8, ConvertFP16Nearest) // fix the tolerance value float abs_tol = 1e-3; // convert 0 fp16 to fp8 and back, check if holds - ASSERT_NEAR(half_t{0.0}, type_convert(type_convert(half_t{0.0})), abs_tol); + ASSERT_NEAR(half_t{0.0}, type_convert(f8_convert_rne(half_t{0.0})), abs_tol); // convert minimal fp16 to fp8 and back, check if holds ASSERT_NEAR(ck::NumericLimits::Min(), - type_convert(type_convert(ck::NumericLimits::Min())), + type_convert(f8_convert_rne(ck::NumericLimits::Min())), abs_tol); // convert maximal f8_t to fp16 and check if equal to 240.0 - ASSERT_NEAR(half_t{240.0}, type_convert(type_convert(half_t{240.0})), abs_tol); + ASSERT_NEAR(half_t{240.0}, type_convert(f8_convert_rne(half_t{240.0})), abs_tol); // convert maximal fp16 to fp8 and back, check if clipped to 240.0 ASSERT_NEAR(half_t{240.0}, - type_convert(type_convert(ck::NumericLimits::Max())), + type_convert(f8_convert_rne(ck::NumericLimits::Max())), abs_tol); // convert QuietNaN fp16 to f8_t and check if it is QuietNaN ASSERT_NEAR(type_convert(0x80), - type_convert(ck::NumericLimits::QuietNaN()), + f8_convert_rne(ck::NumericLimits::QuietNaN()), abs_tol); // positive norm fp16 value to fp8 and back, check if holds half_t pos_half = half_t{0.017578125}; - ASSERT_NEAR(pos_half, type_convert(type_convert(pos_half)), abs_tol); + ASSERT_NEAR(pos_half, type_convert(f8_convert_rne(pos_half)), abs_tol); // negative norm fp16 value to fp8 and back, check if holds half_t neg_half = half_t{-0.015625}; - ASSERT_NEAR(neg_half, type_convert(type_convert(neg_half)), abs_tol); + ASSERT_NEAR(neg_half, type_convert(f8_convert_rne(neg_half)), abs_tol); // positive subnorm fp16 value to fp8 and back, check if holds pos_half = half_t{0.00390625}; - ASSERT_NEAR(pos_half, type_convert(type_convert(pos_half)), abs_tol); + ASSERT_NEAR(pos_half, type_convert(f8_convert_rne(pos_half)), abs_tol); // negative subnorm fp16 value to fp8 and back, check if holds neg_half = half_t{-0.001953125}; - ASSERT_NEAR(neg_half, type_convert(type_convert(neg_half)), abs_tol); + ASSERT_NEAR(neg_half, type_convert(f8_convert_rne(neg_half)), abs_tol); } TEST(FP8, ConvertFP16Stochastic) From c3515f277c878b972c24479a21eaf3d84acd5be7 Mon Sep 17 00:00:00 2001 From: Andriy Roshchenko <107577548+andriy-ca@users.noreply.github.com> Date: Wed, 21 Aug 2024 16:22:41 -0600 Subject: [PATCH 14/20] Adding Instances and Examples for FP8-based Scaled Convolution and AMAX Reduction. (#1473) * Enable CMakePresets build * Verify Convolution, Scaling and ReLU algorithms. * Add tensor element-wise scale and type cast operation. * Reduction implemented but does not work. * Exploration of Reduction functionality. * Completed example for Convolution scaled with ReLu activation and AMAX reduction. * WIP: Add required instances for convolution. * WIP: Create client example. Implement convolution stage. * Add elementwise instances. * Add elementwise scale + convert example. * Add reduction instances. * WIP: Client example for AMAX reduction. * WIP: Add instances for multistage reduction. * WIP: Implementation of multistage reduction. * Refactoring. * Clean up. * Add CMakePresets.json * Guard off FP8 instances when the data type is not available. * Add example for Scaled FP8 Convolution with AMAX reduction. * Refactor CombConvScaleRelu instances. * Add CombConvScale instances. * Add client example for Scaled FP8 Convolution with AMAX reduction. * Cleanup. --- .../24_grouped_conv_activation/CMakeLists.txt | 24 ++++-- .../common.hpp | 19 ++-- .../conv3d_fwd_convscale_amax_fp8.cpp | 58 +++++++++++++ .../convscale_reduce/CMakeLists.txt | 5 +- .../convnd_fwd_xdl_convscale_amax_fp8.cpp | 82 ++++++++++++++++++ .../combined_element_wise_operation.hpp | 4 +- .../grouped_convolution_forward_convscale.hpp | 86 ++++++++++++++++++- ...ped_convolution_forward_convscale_relu.hpp | 4 +- .../device_permute_scale_instances.hpp | 12 +-- .../CMakeLists.txt | 3 +- ...dhwgc_gkzyxc_ndhwgk_f8_f8_f32_instance.cpp | 61 +++++++++++++ .../CMakeLists.txt | 3 +- ...dhwgc_gkzyxc_ndhwgk_f8_f8_f32_instance.cpp | 61 +++++++++++++ ..._relu_ndhwgc_gkzyxc_ndhwgk_f8_instance.cpp | 54 +----------- 14 files changed, 389 insertions(+), 87 deletions(-) create mode 100644 client_example/24_grouped_conv_activation/grouped_convnd_fwd_convscale_reduce/conv3d_fwd_convscale_amax_fp8.cpp create mode 100644 example/62_convnd_activ/convscale_reduce/convnd_fwd_xdl_convscale_amax_fp8.cpp create mode 100644 library/src/tensor_operation_instance/gpu/grouped_conv3d_fwd_convscale/xdl/device_grouped_conv3d_fwd_xdl_combconvscale_ndhwgc_gkzyxc_ndhwgk_f8_f8_f32_instance.cpp create mode 100644 library/src/tensor_operation_instance/gpu/grouped_conv3d_fwd_convscale_relu/xdl/device_grouped_conv3d_fwd_xdl_combconvscale_relu_ndhwgc_gkzyxc_ndhwgk_f8_f8_f32_instance.cpp diff --git a/client_example/24_grouped_conv_activation/CMakeLists.txt b/client_example/24_grouped_conv_activation/CMakeLists.txt index 60f4ee41f7..dc55250bfe 100644 --- a/client_example/24_grouped_conv_activation/CMakeLists.txt +++ b/client_example/24_grouped_conv_activation/CMakeLists.txt @@ -1,6 +1,6 @@ if(GPU_TARGETS MATCHES "gfx9") # Fwd scaleadd scaleadd relu -add_executable(client_grouped_convnd_fwd_scaleadd_scaleadd_relu_fp32 +add_executable(client_grouped_convnd_fwd_scaleadd_scaleadd_relu_fp32 grouped_convnd_fwd_scaleadd_scaleadd_relu/grouped_conv_fwd_scaleadd_scaleadd_relu_fp32.cpp) target_link_libraries(client_grouped_convnd_fwd_scaleadd_scaleadd_relu_fp32 PRIVATE composable_kernel::device_conv_operations) @@ -36,7 +36,7 @@ add_executable(client_grouped_convnd_fwd_bilinear_residual_fp16 grouped_convnd_fwd_bilinear/grouped_conv_fwd_bilinear_residual_fp16.cpp) target_link_libraries(client_grouped_convnd_fwd_bilinear_residual_fp16 PRIVATE composable_kernel::device_conv_operations) # Fwd convinvscale -add_executable(client_conv3d_fwd_convinvscale_fp8 +add_executable(client_conv3d_fwd_convinvscale_fp8 grouped_convnd_fwd_convinvscale/conv3d_fwd_convinvscale_fp8.cpp) target_link_libraries(client_conv3d_fwd_convinvscale_fp8 PRIVATE composable_kernel::device_conv_operations) # Fwd convscale + Bias @@ -50,10 +50,18 @@ target_link_libraries(client_conv3d_fwd_convscale_relu_fp8 PRIVATE composable_ke # Fwd convscale + ReLU + AMAX add_executable(client_conv3d_fwd_convscale_relu_amax_fp8 grouped_convnd_fwd_convscale_reduce/conv3d_fwd_convscale_relu_amax_fp8.cpp) -target_link_libraries(client_conv3d_fwd_convscale_relu_amax_fp8 - PRIVATE composable_kernel::device_conv_operations - composable_kernel::device_other_operations - composable_kernel::device_reduction_operations +target_link_libraries(client_conv3d_fwd_convscale_relu_amax_fp8 + PRIVATE composable_kernel::device_conv_operations + composable_kernel::device_other_operations + composable_kernel::device_reduction_operations + utility) +# Fwd convscale + AMAX +add_executable(client_conv3d_fwd_convscale_amax_fp8 + grouped_convnd_fwd_convscale_reduce/conv3d_fwd_convscale_amax_fp8.cpp) +target_link_libraries(client_conv3d_fwd_convscale_amax_fp8 + PRIVATE composable_kernel::device_conv_operations + composable_kernel::device_other_operations + composable_kernel::device_reduction_operations utility) # Fwd convscale add_executable(client_conv3d_fwd_convscale_fp8 @@ -64,11 +72,11 @@ add_executable(client_conv3d_fwd_convscale_bf8 grouped_convnd_fwd_convscale/conv3d_fwd_convscale_bf8.cpp) target_link_libraries(client_conv3d_fwd_convscale_bf8 PRIVATE composable_kernel::device_conv_operations) -add_executable(client_conv3d_fwd_convscale_fp8_bf8 +add_executable(client_conv3d_fwd_convscale_fp8_bf8 grouped_convnd_fwd_convscale/conv3d_fwd_convscale_fp8_bf8.cpp) target_link_libraries(client_conv3d_fwd_convscale_fp8_bf8 PRIVATE composable_kernel::device_conv_operations) -add_executable(client_conv3d_fwd_convscale_bf8_fp8 +add_executable(client_conv3d_fwd_convscale_bf8_fp8 grouped_convnd_fwd_convscale/conv3d_fwd_convscale_bf8_fp8.cpp) target_link_libraries(client_conv3d_fwd_convscale_bf8_fp8 PRIVATE composable_kernel::device_conv_operations) # Bwd data bilinear diff --git a/client_example/24_grouped_conv_activation/grouped_convnd_fwd_convscale_reduce/common.hpp b/client_example/24_grouped_conv_activation/grouped_convnd_fwd_convscale_reduce/common.hpp index b76c5191ec..c78cacf266 100644 --- a/client_example/24_grouped_conv_activation/grouped_convnd_fwd_convscale_reduce/common.hpp +++ b/client_example/24_grouped_conv_activation/grouped_convnd_fwd_convscale_reduce/common.hpp @@ -15,21 +15,18 @@ #include "ck/tensor_operation/gpu/device/device_elementwise.hpp" #include "ck/tensor_operation/gpu/device/device_reduce.hpp" #include "ck/tensor_operation/gpu/device/reduction_operator_mapping.hpp" -#include "ck/tensor_operation/gpu/element/combined_element_wise_operation.hpp" -#include "ck/tensor_operation/gpu/element/unary_element_wise_operation.hpp" #include "ck/utility/tuple.hpp" #include "ck/utility/type.hpp" #include "ck/library/tensor_operation_instance/gpu/grouped_convolution_forward_convscale_relu.hpp" +#include "ck/library/tensor_operation_instance/gpu/grouped_convolution_forward_convscale.hpp" #include "ck/utility/reduction_enums.hpp" #include "ck/library/tensor_operation_instance/gpu/permute_scale.hpp" #include "ck/library/tensor_operation_instance/gpu/reduce/reduce.hpp" #include "ck/library/utility/host_tensor.hpp" -namespace ew = ck::tensor_operation::element_wise; - -using PassThrough = ew::PassThrough; -using ConvScaleRelu = ew::UnaryCombinedOp; -using ConvScale = ew::UnaryCombinedOp; +using PassThrough = ck::tensor_operation::element_wise::PassThrough; +using ConvScaleRelu = ck::tensor_operation::element_wise::ScaleScaleRelu; +using ConvScale = ck::tensor_operation::element_wise::ScaleScalePass; struct SimpleDeviceMem { @@ -221,7 +218,9 @@ bool run_grouped_conv_fwd_convscale_reduce( * FP8 Convolution with Scaling */ std::cout << "\n\nConvolution with scale Benchmarking:" << std::endl; - auto elementwise_op = ConvElementOp{ew::Scale{scale_in}, ew::Scale{scale_wei}, {}}; + auto elementwise_op = ConvElementOp{ck::tensor_operation::element_wise::Scale{scale_in}, + ck::tensor_operation::element_wise::Scale{scale_wei}, + {}}; auto conv_ok = ConvolutionScale, ck::Tuple, - ew::Scale, + ck::tensor_operation::element_wise::Scale, NumDimSpatial + NumNonSpatialDim>; // get device op instances @@ -483,7 +482,7 @@ bool TensorScaleConvert(SimpleDeviceMem& in, // profile device operation instances std::cout << "Run all DeviceScaleConvert instances and do timing" << std::endl; - auto scale_convert = ew::Scale{scale_out}; + auto scale_convert = ck::tensor_operation::element_wise::Scale{scale_out}; for(int i = 0; i < op_ptrs.size(); ++i) { diff --git a/client_example/24_grouped_conv_activation/grouped_convnd_fwd_convscale_reduce/conv3d_fwd_convscale_amax_fp8.cpp b/client_example/24_grouped_conv_activation/grouped_convnd_fwd_convscale_reduce/conv3d_fwd_convscale_amax_fp8.cpp new file mode 100644 index 0000000000..1c0299b841 --- /dev/null +++ b/client_example/24_grouped_conv_activation/grouped_convnd_fwd_convscale_reduce/conv3d_fwd_convscale_amax_fp8.cpp @@ -0,0 +1,58 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved. + +#include "common.hpp" + +#include "ck/ck.hpp" +#include "ck/tensor_operation/gpu/device/tensor_layout.hpp" + +using InDataType = ck::f8_t; +using WeiDataType = ck::f8_t; +using CShuffleDataType = float; +using ConvOutDataType = float; // data type of convolution result +using OutDataType = ck::f8_t; // data type of final result +using AComputeDataType = ck::f8_t; +using BComputeDataType = ck::f8_t; + +using ConvElementOp = ConvScale; + +using InLayout = ck::tensor_layout::convolution::NDHWGC; +using WeiLayout = ck::tensor_layout::convolution::GKZYXC; +using OutLayout = ck::tensor_layout::convolution::NDHWGK; + +constexpr auto ReduceOpId = ck::ReduceTensorOp::AMAX; + +static constexpr ck::index_t NumDimSpatial = 3; +static constexpr ck::index_t G = 1; +static constexpr ck::index_t N = 64; +static constexpr ck::index_t K = 128; +static constexpr ck::index_t C = 64; +static constexpr ck::index_t Z = 3; +static constexpr ck::index_t Y = 3; +static constexpr ck::index_t X = 3; +static constexpr ck::index_t Di = 28; +static constexpr ck::index_t Hi = 28; +static constexpr ck::index_t Wi = 3; +static constexpr ck::index_t Do = 28; +static constexpr ck::index_t Ho = 28; +static constexpr ck::index_t Wo = 3; + +int main() +{ + return run_grouped_conv_fwd_convscale_reduce( + {N, Di, Hi, Wi, G, C}, {G, K, Z, Y, X, C}, {N, Do, Ho, Wo, G, K}) + ? EXIT_SUCCESS + : EXIT_FAILURE; +} diff --git a/example/62_convnd_activ/convscale_reduce/CMakeLists.txt b/example/62_convnd_activ/convscale_reduce/CMakeLists.txt index b3c6621509..ff9020a707 100644 --- a/example/62_convnd_activ/convscale_reduce/CMakeLists.txt +++ b/example/62_convnd_activ/convscale_reduce/CMakeLists.txt @@ -4,7 +4,10 @@ foreach(gpu IN LISTS GPU_TARGETS) if(gpu IN_LIST gpu_list AND target EQUAL 0) add_custom_target(example_convnd_activ_xdl_convscale_reduce) add_example_executable(example_convnd_fwd_xdl_convscale_relu_amax_fp8 convnd_fwd_xdl_convscale_relu_amax_fp8.cpp) - add_example_dependencies(example_convnd_activ_xdl_convscale_reduce example_convnd_fwd_xdl_convscale_relu_amax_fp8 ) + add_example_dependencies(example_convnd_activ_xdl_convscale_reduce example_convnd_fwd_xdl_convscale_relu_amax_fp8) + + add_example_executable(example_convnd_fwd_xdl_convscale_amax_fp8 convnd_fwd_xdl_convscale_amax_fp8.cpp) + add_example_dependencies(example_convnd_activ_xdl_convscale_reduce example_convnd_fwd_xdl_convscale_amax_fp8) set(target 1) endif() diff --git a/example/62_convnd_activ/convscale_reduce/convnd_fwd_xdl_convscale_amax_fp8.cpp b/example/62_convnd_activ/convscale_reduce/convnd_fwd_xdl_convscale_amax_fp8.cpp new file mode 100644 index 0000000000..a8b4fdbead --- /dev/null +++ b/example/62_convnd_activ/convscale_reduce/convnd_fwd_xdl_convscale_amax_fp8.cpp @@ -0,0 +1,82 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved. + +#include "convnd_fwd_convscale_reduce_common.hpp" + +#include "ck/tensor_operation/gpu/device/impl/device_grouped_conv_fwd_multiple_abd_xdl_cshuffle.hpp" + +using InDataType = ck::f8_t; +using WeiDataType = ck::f8_t; +using AccDataType = float; +using CShuffleDataType = float; +using ConvOutDataType = float; // data type of convolution result +using OutDataType = ck::f8_t; // data type of final result +using AComputeDataType = ck::f8_t; +using BComputeDataType = ck::f8_t; + +template +using S = ck::Sequence; + +using InElementOp = PassThrough; +using WeiElementOp = PassThrough; +using OutElementOp = ConvScale; + +static constexpr auto ConvSpec = + ck::tensor_operation::device::ConvolutionForwardSpecialization::Default; + +static constexpr auto GemmSpec = ck::tensor_operation::device::GemmSpecialization::MNKPadding; + +template +using DeviceGroupedConvNDFwdInstance = + ck::tensor_operation::device::DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle< + NDimSpatial, + InLayout, + WeiLayout, + ck::Tuple<>, + OutLayout, + InDataType, + WeiDataType, + AccDataType, + CShuffleDataType, + ck::Tuple<>, + ConvOutDataType, + InElementOp, + WeiElementOp, + OutElementOp, + ConvSpec, // ConvForwardSpecialization + GemmSpec, // GemmSpecialization + 1, // + 256, // BlockSize + 128, // MPerBlock + 256, // NPerBlock + 32, // KPerBlock + 8, // AK1 + 8, // BK1 + 32, // MPerXdl + 32, // NPerXdl + 2, // MXdlPerWave + 4, // NXdlPerWave + S<4, 64, 1>, // ABlockTransferThreadClusterLengths_AK0_M_AK1 + S<1, 0, 2>, // ABlockTransferThreadClusterArrangeOrder + S<1, 0, 2>, // ABlockTransferSrcAccessOrder + 2, // ABlockTransferSrcVectorDim + 8, // ABlockTransferSrcScalarPerVector + 8, // ABlockTransferDstScalarPerVector_AK1 + 1, // 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 + 1, + 1, + S<1, 32, 1, 8>, + 8, + AComputeDataType, + BComputeDataType>; + +#include "run_convnd_fwd_example.inc" + +int main(int argc, char* argv[]) { return run_convnd_fwd_example(argc, argv) ? 0 : 1; } diff --git a/include/ck/tensor_operation/gpu/element/combined_element_wise_operation.hpp b/include/ck/tensor_operation/gpu/element/combined_element_wise_operation.hpp index d8bac8da7a..3cc1c3c42c 100644 --- a/include/ck/tensor_operation/gpu/element/combined_element_wise_operation.hpp +++ b/include/ck/tensor_operation/gpu/element/combined_element_wise_operation.hpp @@ -3,7 +3,6 @@ #pragma once -#include "ck/utility/data_type.hpp" #include "ck/tensor_operation/gpu/element/element_wise_operation.hpp" namespace ck { @@ -107,6 +106,9 @@ struct TrinaryWithUnaryCombinedOp UnaryOp2 unary_op2_{}; }; +using ScaleScalePass = UnaryCombinedOp; +using ScaleScaleRelu = UnaryCombinedOp; + } // namespace element_wise } // namespace tensor_operation } // namespace ck diff --git a/library/include/ck/library/tensor_operation_instance/gpu/grouped_convolution_forward_convscale.hpp b/library/include/ck/library/tensor_operation_instance/gpu/grouped_convolution_forward_convscale.hpp index 63dcdc6053..e070b249e1 100644 --- a/library/include/ck/library/tensor_operation_instance/gpu/grouped_convolution_forward_convscale.hpp +++ b/library/include/ck/library/tensor_operation_instance/gpu/grouped_convolution_forward_convscale.hpp @@ -8,9 +8,7 @@ #include "ck/ck.hpp" #include "ck/tensor_operation/gpu/device/device_grouped_conv_fwd_multiple_abd.hpp" -#include "ck/tensor_operation/gpu/device/tensor_layout.hpp" -#include "ck/tensor_operation/gpu/element/element_wise_operation.hpp" - +#include "ck/tensor_operation/gpu/element/combined_element_wise_operation.hpp" #include "ck/library/tensor_operation_instance/device_operation_instance_factory.hpp" namespace ck { @@ -177,6 +175,88 @@ struct DeviceOperationInstanceFactory, + NDHWGK, + F8, + F8, + ck::Tuple<>, + F32, + PassThrough, + PassThrough, + CombConvScale, + F8, + F8>>>& instances); +#endif + +template +struct DeviceOperationInstanceFactory< + ck::tensor_operation::device::DeviceGroupedConvFwdMultipleABD> +{ + using DeviceOp = DeviceGroupedConvFwdMultipleABD; + + static auto GetInstances() + { + std::vector> op_ptrs; + if constexpr(NumDimSpatial == 3 && is_same_v && + is_same_v && is_same_v) + { +#ifdef CK_ENABLE_FP8 + if constexpr(is_same_v && is_same_v && + is_same_v && is_same_v && + is_same_v) + { + add_device_grouped_conv3d_fwd_xdl_combconvscale_ndhwgc_gkzyxc_ndhwgk_f8_f8_f32_instances( + op_ptrs); + } +#endif + } + return op_ptrs; + } +}; + } // namespace instance } // namespace device } // namespace tensor_operation diff --git a/library/include/ck/library/tensor_operation_instance/gpu/grouped_convolution_forward_convscale_relu.hpp b/library/include/ck/library/tensor_operation_instance/gpu/grouped_convolution_forward_convscale_relu.hpp index 419f5a609a..a0651912d4 100644 --- a/library/include/ck/library/tensor_operation_instance/gpu/grouped_convolution_forward_convscale_relu.hpp +++ b/library/include/ck/library/tensor_operation_instance/gpu/grouped_convolution_forward_convscale_relu.hpp @@ -9,7 +9,6 @@ #include "ck/ck.hpp" #include "ck/tensor_operation/gpu/device/device_grouped_conv_fwd_multiple_abd.hpp" #include "ck/tensor_operation/gpu/element/combined_element_wise_operation.hpp" -#include "ck/tensor_operation/gpu/element/unary_element_wise_operation.hpp" #include "ck/library/tensor_operation_instance/device_operation_instance_factory.hpp" namespace ck { @@ -100,8 +99,7 @@ struct DeviceOperationInstanceFactory< } }; -namespace ew = ck::tensor_operation::element_wise; -using CombConvScaleRelu = ew::UnaryCombinedOp; +using CombConvScaleRelu = ck::tensor_operation::element_wise::ScaleScaleRelu; #ifdef CK_ENABLE_FP8 void add_device_grouped_conv3d_fwd_xdl_combconvscale_relu_ndhwgc_gkzyxc_ndhwgk_f8_f8_f32_instances( diff --git a/library/include/ck/library/tensor_operation_instance/gpu/permute_scale/device_permute_scale_instances.hpp b/library/include/ck/library/tensor_operation_instance/gpu/permute_scale/device_permute_scale_instances.hpp index 204c9a310d..1a70db3bf0 100644 --- a/library/include/ck/library/tensor_operation_instance/gpu/permute_scale/device_permute_scale_instances.hpp +++ b/library/include/ck/library/tensor_operation_instance/gpu/permute_scale/device_permute_scale_instances.hpp @@ -47,7 +47,7 @@ using device_permute_scale_f16_instances = #if 0 // Disabled instances to improve compilation time - // They listed here to show other possible combinations of parameters + // They listed here to show other possible combinations of parameters DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 256, 256, 256, 16, 16, ck::Sequence<1, 0>, ck::Sequence<16>, ck::Sequence<16>>, DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 128, 256, 128, 16, 16, ck::Sequence<1, 0>, ck::Sequence<16>, ck::Sequence<16>>, DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 128, 128, 256, 16, 16, ck::Sequence<1, 0>, ck::Sequence<16>, ck::Sequence<16>>, @@ -58,7 +58,7 @@ using device_permute_scale_f16_instances = DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 64, 128, 128, 16, 16, ck::Sequence<1, 0>, ck::Sequence<16>, ck::Sequence<16>>, DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 32, 128, 64, 16, 16, ck::Sequence<1, 0>, ck::Sequence<16>, ck::Sequence<16>>, DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 32, 64, 128, 16, 16, ck::Sequence<1, 0>, ck::Sequence<16>, ck::Sequence<16>>, - + DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 256, 64, 128, 4, 8, ck::Sequence<1, 0>, ck::Sequence<4>, ck::Sequence<4>>, DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 256, 128, 64, 4, 8, ck::Sequence<1, 0>, ck::Sequence<4>, ck::Sequence<4>>, DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 128, 64, 64, 4, 8, ck::Sequence<1, 0>, ck::Sequence<4>, ck::Sequence<4>>, @@ -98,7 +98,7 @@ using device_permute_scale_f16_instances = DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 64, 64, 16, 4, 4, ck::Sequence<1, 0>, ck::Sequence<1>, ck::Sequence<1>>, DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 32, 32, 16, 4, 4, ck::Sequence<1, 0>, ck::Sequence<1>, ck::Sequence<1>>, DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 32, 16, 32, 4, 4, ck::Sequence<1, 0>, ck::Sequence<1>, ck::Sequence<1>> - + >; template , ck::Tuple, ElementwiseOp, NDims, 256, 256, 256, 16, 16, ck::Sequence<1, 0>, ck::Sequence<16>, ck::Sequence<16>>, DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 128, 256, 128, 16, 16, ck::Sequence<1, 0>, ck::Sequence<16>, ck::Sequence<16>>, DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 128, 128, 256, 16, 16, ck::Sequence<1, 0>, ck::Sequence<16>, ck::Sequence<16>>, @@ -143,7 +143,7 @@ using device_permute_scale_f32_instances = std::tuple< DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 64, 128, 128, 16, 16, ck::Sequence<1, 0>, ck::Sequence<16>, ck::Sequence<16>>, DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 32, 128, 64, 16, 16, ck::Sequence<1, 0>, ck::Sequence<16>, ck::Sequence<16>>, DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 32, 64, 128, 16, 16, ck::Sequence<1, 0>, ck::Sequence<16>, ck::Sequence<16>>, - + DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 256, 64, 128, 4, 8, ck::Sequence<1, 0>, ck::Sequence<4>, ck::Sequence<4>>, DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 256, 128, 64, 4, 8, ck::Sequence<1, 0>, ck::Sequence<4>, ck::Sequence<4>>, DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 128, 64, 64, 4, 8, ck::Sequence<1, 0>, ck::Sequence<4>, ck::Sequence<4>>, @@ -169,7 +169,7 @@ using device_permute_scale_f32_instances = std::tuple< DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 64, 128, 16, 8, 4, ck::Sequence<1, 0>, ck::Sequence<4>, ck::Sequence<4>>, DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 32, 64, 16, 8, 4, ck::Sequence<1, 0>, ck::Sequence<4>, ck::Sequence<4>>, DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 32, 32, 32, 8, 4, ck::Sequence<1, 0>, ck::Sequence<4>, ck::Sequence<4>>, -#endif +#endif DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 256, 64, 64, 4, 4, ck::Sequence<1, 0>, ck::Sequence<1>, ck::Sequence<1>>, DeviceElementwiseImpl, ck::Tuple, ElementwiseOp, NDims, 256, 128, 32, 4, 4, ck::Sequence<1, 0>, ck::Sequence<1>, ck::Sequence<1>>, diff --git a/library/src/tensor_operation_instance/gpu/grouped_conv3d_fwd_convscale/CMakeLists.txt b/library/src/tensor_operation_instance/gpu/grouped_conv3d_fwd_convscale/CMakeLists.txt index c7f4a3527e..e20e3f49ed 100644 --- a/library/src/tensor_operation_instance/gpu/grouped_conv3d_fwd_convscale/CMakeLists.txt +++ b/library/src/tensor_operation_instance/gpu/grouped_conv3d_fwd_convscale/CMakeLists.txt @@ -3,6 +3,7 @@ set(GROUPED_CONV3D_FWD_CONVSCALE xdl/device_grouped_conv3d_fwd_xdl_convscale_ndhwgc_gkzyxc_ndhwgk_f8_instance.cpp xdl/device_grouped_conv3d_fwd_xdl_convscale_ndhwgc_gkzyxc_ndhwgk_bf8_instance.cpp xdl/device_grouped_conv3d_fwd_xdl_convscale_ndhwgc_gkzyxc_ndhwgk_f8_bf8_instance.cpp - xdl/device_grouped_conv3d_fwd_xdl_convscale_ndhwgc_gkzyxc_ndhwgk_bf8_f8_instance.cpp) + xdl/device_grouped_conv3d_fwd_xdl_convscale_ndhwgc_gkzyxc_ndhwgk_bf8_f8_instance.cpp + xdl/device_grouped_conv3d_fwd_xdl_combconvscale_ndhwgc_gkzyxc_ndhwgk_f8_f8_f32_instance.cpp) add_instance_library(device_grouped_conv3d_fwd_convscale_instance ${GROUPED_CONV3D_FWD_CONVSCALE}) diff --git a/library/src/tensor_operation_instance/gpu/grouped_conv3d_fwd_convscale/xdl/device_grouped_conv3d_fwd_xdl_combconvscale_ndhwgc_gkzyxc_ndhwgk_f8_f8_f32_instance.cpp b/library/src/tensor_operation_instance/gpu/grouped_conv3d_fwd_convscale/xdl/device_grouped_conv3d_fwd_xdl_combconvscale_ndhwgc_gkzyxc_ndhwgk_f8_f8_f32_instance.cpp new file mode 100644 index 0000000000..2d387f1034 --- /dev/null +++ b/library/src/tensor_operation_instance/gpu/grouped_conv3d_fwd_convscale/xdl/device_grouped_conv3d_fwd_xdl_combconvscale_ndhwgc_gkzyxc_ndhwgk_f8_f8_f32_instance.cpp @@ -0,0 +1,61 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved. + +#include "ck/library/tensor_operation_instance/gpu/grouped_conv_fwd/device_grouped_conv_fwd_xdl_outelementop_instance.hpp" +#include "ck/library/tensor_operation_instance/gpu/grouped_convolution_forward_convscale.hpp" +#include "ck/library/tensor_operation_instance/add_device_operation_instance.hpp" + +namespace ck { +namespace tensor_operation { +namespace device { +namespace instance { + +void add_device_grouped_conv3d_fwd_xdl_combconvscale_ndhwgc_gkzyxc_ndhwgk_f8_f8_f32_instances( + std::vector, + NDHWGK, + F8, + F8, + ck::Tuple<>, + F32, + PassThrough, + PassThrough, + CombConvScale, + F8, + F8>>>& instances) +{ + add_device_operation_instances( + instances, + device_grouped_conv_fwd_xdl_outelementop_f8_f8_f32_instances<3, + NDHWGC, + GKZYXC, + ck::Tuple<>, + NDHWGK, + ConvFwdDefault, + CombConvScale>{}); + add_device_operation_instances( + instances, + device_grouped_conv_fwd_xdl_outelementop_f8_f8_f32_instances<3, + NDHWGC, + GKZYXC, + ck::Tuple<>, + NDHWGK, + ConvFwd1x1P0, + CombConvScale>{}); + add_device_operation_instances( + instances, + device_grouped_conv_fwd_xdl_outelementop_f8_f8_f32_instances<3, + NDHWGC, + GKZYXC, + ck::Tuple<>, + NDHWGK, + ConvFwd1x1S1P0, + CombConvScale>{}); +} + +} // namespace instance +} // namespace device +} // namespace tensor_operation +} // namespace ck diff --git a/library/src/tensor_operation_instance/gpu/grouped_conv3d_fwd_convscale_relu/CMakeLists.txt b/library/src/tensor_operation_instance/gpu/grouped_conv3d_fwd_convscale_relu/CMakeLists.txt index c60df5a733..8ba52adcb8 100644 --- a/library/src/tensor_operation_instance/gpu/grouped_conv3d_fwd_convscale_relu/CMakeLists.txt +++ b/library/src/tensor_operation_instance/gpu/grouped_conv3d_fwd_convscale_relu/CMakeLists.txt @@ -1,5 +1,6 @@ # ONLY XDL_KERNELS set(GROUPED_CONV3D_FWD_CONVSCALE_RELU - xdl/device_grouped_conv3d_fwd_xdl_convscale_relu_ndhwgc_gkzyxc_ndhwgk_f8_instance.cpp) + xdl/device_grouped_conv3d_fwd_xdl_convscale_relu_ndhwgc_gkzyxc_ndhwgk_f8_instance.cpp + xdl/device_grouped_conv3d_fwd_xdl_combconvscale_relu_ndhwgc_gkzyxc_ndhwgk_f8_f8_f32_instance.cpp) add_instance_library(device_grouped_conv3d_fwd_convscale_relu_instance ${GROUPED_CONV3D_FWD_CONVSCALE_RELU}) diff --git a/library/src/tensor_operation_instance/gpu/grouped_conv3d_fwd_convscale_relu/xdl/device_grouped_conv3d_fwd_xdl_combconvscale_relu_ndhwgc_gkzyxc_ndhwgk_f8_f8_f32_instance.cpp b/library/src/tensor_operation_instance/gpu/grouped_conv3d_fwd_convscale_relu/xdl/device_grouped_conv3d_fwd_xdl_combconvscale_relu_ndhwgc_gkzyxc_ndhwgk_f8_f8_f32_instance.cpp new file mode 100644 index 0000000000..1a27e64d31 --- /dev/null +++ b/library/src/tensor_operation_instance/gpu/grouped_conv3d_fwd_convscale_relu/xdl/device_grouped_conv3d_fwd_xdl_combconvscale_relu_ndhwgc_gkzyxc_ndhwgk_f8_f8_f32_instance.cpp @@ -0,0 +1,61 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved. + +#include "ck/library/tensor_operation_instance/gpu/grouped_conv_fwd/device_grouped_conv_fwd_xdl_outelementop_instance.hpp" +#include "ck/library/tensor_operation_instance/gpu/grouped_convolution_forward_convscale_relu.hpp" +#include "ck/library/tensor_operation_instance/add_device_operation_instance.hpp" + +namespace ck { +namespace tensor_operation { +namespace device { +namespace instance { + +void add_device_grouped_conv3d_fwd_xdl_combconvscale_relu_ndhwgc_gkzyxc_ndhwgk_f8_f8_f32_instances( + std::vector, + NDHWGK, + F8, + F8, + ck::Tuple<>, + F32, + PassThrough, + PassThrough, + CombConvScaleRelu, + F8, + F8>>>& instances) +{ + add_device_operation_instances( + instances, + device_grouped_conv_fwd_xdl_outelementop_f8_f8_f32_instances<3, + NDHWGC, + GKZYXC, + ck::Tuple<>, + NDHWGK, + ConvFwdDefault, + CombConvScaleRelu>{}); + add_device_operation_instances( + instances, + device_grouped_conv_fwd_xdl_outelementop_f8_f8_f32_instances<3, + NDHWGC, + GKZYXC, + ck::Tuple<>, + NDHWGK, + ConvFwd1x1P0, + CombConvScaleRelu>{}); + add_device_operation_instances( + instances, + device_grouped_conv_fwd_xdl_outelementop_f8_f8_f32_instances<3, + NDHWGC, + GKZYXC, + ck::Tuple<>, + NDHWGK, + ConvFwd1x1S1P0, + CombConvScaleRelu>{}); +} + +} // namespace instance +} // namespace device +} // namespace tensor_operation +} // namespace ck diff --git a/library/src/tensor_operation_instance/gpu/grouped_conv3d_fwd_convscale_relu/xdl/device_grouped_conv3d_fwd_xdl_convscale_relu_ndhwgc_gkzyxc_ndhwgk_f8_instance.cpp b/library/src/tensor_operation_instance/gpu/grouped_conv3d_fwd_convscale_relu/xdl/device_grouped_conv3d_fwd_xdl_convscale_relu_ndhwgc_gkzyxc_ndhwgk_f8_instance.cpp index 1fda1f4ee6..91bfdda0d8 100644 --- a/library/src/tensor_operation_instance/gpu/grouped_conv3d_fwd_convscale_relu/xdl/device_grouped_conv3d_fwd_xdl_convscale_relu_ndhwgc_gkzyxc_ndhwgk_f8_instance.cpp +++ b/library/src/tensor_operation_instance/gpu/grouped_conv3d_fwd_convscale_relu/xdl/device_grouped_conv3d_fwd_xdl_convscale_relu_ndhwgc_gkzyxc_ndhwgk_f8_instance.cpp @@ -3,16 +3,13 @@ #include "ck/library/tensor_operation_instance/gpu/grouped_conv_fwd/device_grouped_conv_fwd_xdl_outelementop_instance.hpp" #include "ck/library/tensor_operation_instance/add_device_operation_instance.hpp" -#include "ck/tensor_operation/gpu/element/combined_element_wise_operation.hpp" -#include "ck/tensor_operation/gpu/element/unary_element_wise_operation.hpp" +#include "ck/library/tensor_operation_instance/gpu/grouped_convolution_forward_convscale_relu.hpp" namespace ck { namespace tensor_operation { namespace device { namespace instance { -using ConvScaleRelu = ck::tensor_operation::element_wise::ConvScaleRelu; - void add_device_grouped_conv3d_fwd_xdl_convscale_relu_ndhwgc_gkzyxc_ndhwgk_f8_instances( std::vector{}); } - -namespace ew = ck::tensor_operation::element_wise; -using CombConvScaleRelu = ew::UnaryCombinedOp; - -void add_device_grouped_conv3d_fwd_xdl_combconvscale_relu_ndhwgc_gkzyxc_ndhwgk_f8_f8_f32_instances( - std::vector, - NDHWGK, - F8, - F8, - ck::Tuple<>, - F32, - PassThrough, - PassThrough, - CombConvScaleRelu, - F8, - F8>>>& instances) -{ - add_device_operation_instances( - instances, - device_grouped_conv_fwd_xdl_outelementop_f8_f8_f32_instances<3, - NDHWGC, - GKZYXC, - ck::Tuple<>, - NDHWGK, - ConvFwdDefault, - CombConvScaleRelu>{}); - add_device_operation_instances( - instances, - device_grouped_conv_fwd_xdl_outelementop_f8_f8_f32_instances<3, - NDHWGC, - GKZYXC, - ck::Tuple<>, - NDHWGK, - ConvFwd1x1P0, - CombConvScaleRelu>{}); - add_device_operation_instances( - instances, - device_grouped_conv_fwd_xdl_outelementop_f8_f8_f32_instances<3, - NDHWGC, - GKZYXC, - ck::Tuple<>, - NDHWGK, - ConvFwd1x1S1P0, - CombConvScaleRelu>{}); -} - } // namespace instance } // namespace device } // namespace tensor_operation From 1925b322eb507feffc31a5433f449d81a38812c0 Mon Sep 17 00:00:00 2001 From: Illia Silin <98187287+illsilin@users.noreply.github.com> Date: Wed, 21 Aug 2024 21:29:48 -0700 Subject: [PATCH 15/20] fix the build errors with clang20 (#1478) --- codegen/test/CMakeLists.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/codegen/test/CMakeLists.txt b/codegen/test/CMakeLists.txt index 5aad1ef877..9716b7c681 100644 --- a/codegen/test/CMakeLists.txt +++ b/codegen/test/CMakeLists.txt @@ -9,7 +9,7 @@ foreach(TEST_SRC ${TEST_SRCS}) add_test(NAME codegen_test_${BASE_NAME} COMMAND test_host_${BASE_NAME}) target_link_libraries(test_host_${BASE_NAME} ck_rtc ck_host) # target_link_libraries(test_host_${BASE_NAME} ${CK_ROOT}/build/lib/libutility.a) - target_include_directories(test_host_${BASE_NAME} PUBLIC include()) + target_include_directories(test_host_${BASE_NAME} PUBLIC ${CK_ROOT}/codegen/test/include) target_include_directories(test_host_${BASE_NAME} PUBLIC ${CK_ROOT}/include) target_include_directories(test_host_${BASE_NAME} PUBLIC ${CK_ROOT}/library/include) endforeach() From 0d9bf9f154396b7f5fb2cf40691fda3f0a4e151a Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Wed, 21 Aug 2024 22:40:49 -0700 Subject: [PATCH 16/20] Bump rocm-docs-core from 1.7.1 to 1.7.2 in /docs/sphinx (#1479) Bumps [rocm-docs-core](https://github.com/ROCm/rocm-docs-core) from 1.7.1 to 1.7.2. - [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.7.1...v1.7.2) --- updated-dependencies: - dependency-name: rocm-docs-core dependency-type: direct:production update-type: version-update:semver-patch ... 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 6bc6e07eea..246439bcdb 100644 --- a/docs/sphinx/requirements.in +++ b/docs/sphinx/requirements.in @@ -1,2 +1,2 @@ -rocm-docs-core==1.7.1 +rocm-docs-core==1.7.2 sphinxcontrib-bibtex==2.6.2 diff --git a/docs/sphinx/requirements.txt b/docs/sphinx/requirements.txt index 0e02dbb727..185acda5b2 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.7.1 +rocm-docs-core==1.7.2 # via -r requirements.in six==1.16.0 # via pybtex From 967b1f0fda8bbbae5f1a9de71d3100380d04a068 Mon Sep 17 00:00:00 2001 From: arai713 <67439843+arai713@users.noreply.github.com> Date: Thu, 22 Aug 2024 07:24:55 -0700 Subject: [PATCH 17/20] Codegen INSTANCES_ONLY build (#1468) * initial push - altering codegen build * fix the codegen cmake * enable codegen build for gfx908 and gfx90a * enable building codegen with INSTANCES_ONLY=ON * updating ck_rtc * remove gpu targets for codegen and rename tests * make codegen tests dependencies of tests and check targets --------- Co-authored-by: illsilin Co-authored-by: Illia Silin <98187287+illsilin@users.noreply.github.com> --- CMakeLists.txt | 7 ++++--- codegen/CMakeLists.txt | 2 ++ codegen/test/CMakeLists.txt | 18 ++++++++++-------- codegen/test/rtc/CMakeLists.txt | 2 -- 4 files changed, 16 insertions(+), 13 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 8a08ddd19c..5012635d3a 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -553,9 +553,6 @@ if(NOT DEFINED INSTANCES_ONLY) PACKAGE_NAME examples ) add_subdirectory(example) - if(GPU_TARGETS MATCHES "gfx9" AND NOT INSTANCES_ONLY) - add_subdirectory(codegen) - endif() if(BUILD_TESTING) add_subdirectory(test) endif() @@ -575,6 +572,10 @@ if(NOT DEFINED INSTANCES_ONLY) endif() endif() +if(NOT DEFINED PROFILER_ONLY AND (GPU_TARGETS MATCHES "gfx9" OR DEFINED INSTANCES_ONLY)) + add_subdirectory(codegen) +endif() + #Create an interface target for the include only files and call it "composablekernels" include(CMakePackageConfigHelpers) diff --git a/codegen/CMakeLists.txt b/codegen/CMakeLists.txt index d08fe2380b..1d8167b4be 100644 --- a/codegen/CMakeLists.txt +++ b/codegen/CMakeLists.txt @@ -27,6 +27,8 @@ file(GLOB_RECURSE KERNEL_FILES CONFIGURE_DEPENDS add_embed_library(ck_headers ${KERNEL_FILES} RELATIVE ${CK_ROOT}/include) file(GLOB SOURCES CONFIGURE_DEPENDS src/*.cpp) + +##message(STATUS "SOURCE_FILES: ${SOURCES}") # TODO: Use object library add_library(ck_host STATIC ${SOURCES}) target_link_libraries(ck_host PRIVATE ck_headers) diff --git a/codegen/test/CMakeLists.txt b/codegen/test/CMakeLists.txt index 9716b7c681..4841ca7e15 100644 --- a/codegen/test/CMakeLists.txt +++ b/codegen/test/CMakeLists.txt @@ -4,12 +4,14 @@ file(GLOB TEST_SRCS CONFIGURE_DEPENDS *.cpp) foreach(TEST_SRC ${TEST_SRCS}) set_source_files_properties(${TEST_SRC} PROPERTIES LANGUAGE HIP) get_filename_component(BASE_NAME ${TEST_SRC} NAME_WE) - add_executable(test_host_${BASE_NAME} ${TEST_SRC}) - add_dependencies(codegen test_host_${BASE_NAME}) - add_test(NAME codegen_test_${BASE_NAME} COMMAND test_host_${BASE_NAME}) - target_link_libraries(test_host_${BASE_NAME} ck_rtc ck_host) - # target_link_libraries(test_host_${BASE_NAME} ${CK_ROOT}/build/lib/libutility.a) - target_include_directories(test_host_${BASE_NAME} PUBLIC ${CK_ROOT}/codegen/test/include) - target_include_directories(test_host_${BASE_NAME} PUBLIC ${CK_ROOT}/include) - target_include_directories(test_host_${BASE_NAME} PUBLIC ${CK_ROOT}/library/include) + add_executable(codegen_test_${BASE_NAME} ${TEST_SRC}) + add_dependencies(codegen codegen_test_${BASE_NAME}) + add_dependencies(tests codegen_test_${BASE_NAME}) + add_dependencies(check codegen_test_${BASE_NAME}) + add_test(NAME codegen_test_${BASE_NAME} COMMAND codegen_test_${BASE_NAME}) + message("adding test codegen_test_${BASE_NAME}") + target_link_libraries(codegen_test_${BASE_NAME} ck_rtc ck_host) + target_include_directories(codegen_test_${BASE_NAME} PUBLIC include()) + target_include_directories(codegen_test_${BASE_NAME} PUBLIC ${CK_ROOT}/include) + target_include_directories(codegen_test_${BASE_NAME} PUBLIC ${CK_ROOT}/library/include) endforeach() diff --git a/codegen/test/rtc/CMakeLists.txt b/codegen/test/rtc/CMakeLists.txt index 441e60ca88..39497f1a21 100644 --- a/codegen/test/rtc/CMakeLists.txt +++ b/codegen/test/rtc/CMakeLists.txt @@ -1,5 +1,3 @@ - -find_package(hip) file(GLOB RTC_SOURCES CONFIGURE_DEPENDS src/*.cpp) add_library(ck_rtc ${RTC_SOURCES}) target_include_directories(ck_rtc PUBLIC include) From d3fa00f14c24270e4b25a09686d75a80c3073366 Mon Sep 17 00:00:00 2001 From: Illia Silin <98187287+illsilin@users.noreply.github.com> Date: Thu, 22 Aug 2024 09:50:17 -0700 Subject: [PATCH 18/20] disabel codegen tests when INSTANCES_ONLY is set (#1480) --- codegen/test/CMakeLists.txt | 30 ++++++++++++++++-------------- 1 file changed, 16 insertions(+), 14 deletions(-) diff --git a/codegen/test/CMakeLists.txt b/codegen/test/CMakeLists.txt index 4841ca7e15..e943d4ff5b 100644 --- a/codegen/test/CMakeLists.txt +++ b/codegen/test/CMakeLists.txt @@ -1,17 +1,19 @@ list(APPEND CMAKE_PREFIX_PATH /opt/rocm) add_subdirectory(rtc) file(GLOB TEST_SRCS CONFIGURE_DEPENDS *.cpp) -foreach(TEST_SRC ${TEST_SRCS}) - set_source_files_properties(${TEST_SRC} PROPERTIES LANGUAGE HIP) - get_filename_component(BASE_NAME ${TEST_SRC} NAME_WE) - add_executable(codegen_test_${BASE_NAME} ${TEST_SRC}) - add_dependencies(codegen codegen_test_${BASE_NAME}) - add_dependencies(tests codegen_test_${BASE_NAME}) - add_dependencies(check codegen_test_${BASE_NAME}) - add_test(NAME codegen_test_${BASE_NAME} COMMAND codegen_test_${BASE_NAME}) - message("adding test codegen_test_${BASE_NAME}") - target_link_libraries(codegen_test_${BASE_NAME} ck_rtc ck_host) - target_include_directories(codegen_test_${BASE_NAME} PUBLIC include()) - target_include_directories(codegen_test_${BASE_NAME} PUBLIC ${CK_ROOT}/include) - target_include_directories(codegen_test_${BASE_NAME} PUBLIC ${CK_ROOT}/library/include) -endforeach() +if(NOT INSTANCES_ONLY) + foreach(TEST_SRC ${TEST_SRCS}) + set_source_files_properties(${TEST_SRC} PROPERTIES LANGUAGE HIP) + get_filename_component(BASE_NAME ${TEST_SRC} NAME_WE) + add_executable(codegen_test_${BASE_NAME} ${TEST_SRC}) + add_dependencies(codegen codegen_test_${BASE_NAME}) + add_dependencies(tests codegen_test_${BASE_NAME}) + add_dependencies(check codegen_test_${BASE_NAME}) + add_test(NAME codegen_test_${BASE_NAME} COMMAND codegen_test_${BASE_NAME}) + message("adding test codegen_test_${BASE_NAME}") + target_link_libraries(codegen_test_${BASE_NAME} ck_rtc ck_host) + target_include_directories(codegen_test_${BASE_NAME} PUBLIC include()) + target_include_directories(codegen_test_${BASE_NAME} PUBLIC ${CK_ROOT}/include) + target_include_directories(codegen_test_${BASE_NAME} PUBLIC ${CK_ROOT}/library/include) + endforeach() +endif() From 0056e0bf4b270d1eb78807f64f94f3a761048a90 Mon Sep 17 00:00:00 2001 From: Illia Silin <98187287+illsilin@users.noreply.github.com> Date: Thu, 22 Aug 2024 15:05:20 -0700 Subject: [PATCH 19/20] disable bad fp8 test on gfx12 (#1481) --- test/data_type/CMakeLists.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/data_type/CMakeLists.txt b/test/data_type/CMakeLists.txt index 75e098ce48..3b183222bf 100644 --- a/test/data_type/CMakeLists.txt +++ b/test/data_type/CMakeLists.txt @@ -1,5 +1,5 @@ if (GPU_TARGETS) - if (GPU_TARGETS MATCHES "gfx10" OR GPU_TARGETS MATCHES "gfx11") + if (GPU_TARGETS MATCHES "gfx10" OR GPU_TARGETS MATCHES "gfx11" OR GPU_TARGETS MATCHES "gfx12") add_definitions(-DCK_SKIP_FLAKY_F8_TEST) set(CK_SKIP_FLAKY_F8_TEST "ON") endif() From 25935b57a0eac3284b450c2a34a1914a49eca077 Mon Sep 17 00:00:00 2001 From: Illia Silin <98187287+illsilin@users.noreply.github.com> Date: Fri, 23 Aug 2024 15:11:47 -0700 Subject: [PATCH 20/20] fix codegen rtc lib build issue (#1485) --- CMakeLists.txt | 4 +--- codegen/CMakeLists.txt | 2 -- codegen/test/CMakeLists.txt | 2 +- 3 files changed, 2 insertions(+), 6 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 5012635d3a..697eb3b745 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -553,9 +553,7 @@ if(NOT DEFINED INSTANCES_ONLY) PACKAGE_NAME examples ) add_subdirectory(example) - if(BUILD_TESTING) - add_subdirectory(test) - endif() + add_subdirectory(test) rocm_package_setup_component(profiler LIBRARY_NAME composablekernel diff --git a/codegen/CMakeLists.txt b/codegen/CMakeLists.txt index 1d8167b4be..3b3e9f06ee 100644 --- a/codegen/CMakeLists.txt +++ b/codegen/CMakeLists.txt @@ -50,6 +50,4 @@ rocm_install( ) rocm_install(DIRECTORY include/ck DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}) -if(BUILD_TESTING) add_subdirectory(test) -endif() diff --git a/codegen/test/CMakeLists.txt b/codegen/test/CMakeLists.txt index e943d4ff5b..3ca3f39a30 100644 --- a/codegen/test/CMakeLists.txt +++ b/codegen/test/CMakeLists.txt @@ -12,7 +12,7 @@ if(NOT INSTANCES_ONLY) add_test(NAME codegen_test_${BASE_NAME} COMMAND codegen_test_${BASE_NAME}) message("adding test codegen_test_${BASE_NAME}") target_link_libraries(codegen_test_${BASE_NAME} ck_rtc ck_host) - target_include_directories(codegen_test_${BASE_NAME} PUBLIC include()) + target_include_directories(codegen_test_${BASE_NAME} PUBLIC ${CK_ROOT}/codegen/test/include) target_include_directories(codegen_test_${BASE_NAME} PUBLIC ${CK_ROOT}/include) target_include_directories(codegen_test_${BASE_NAME} PUBLIC ${CK_ROOT}/library/include) endforeach()