From 294cb823142a815170cf1faa63e01a431a557a04 Mon Sep 17 00:00:00 2001 From: BrianHarrisonAMD <169072757+BrianHarrisonAMD@users.noreply.github.com> Date: Wed, 2 Oct 2024 10:26:11 -0600 Subject: [PATCH 1/8] Add generating mha static library for gfx90a (#1540) * Add generating mha static library for gfx90a * Update comment to reflect changes --- library/src/tensor_operation_instance/gpu/CMakeLists.txt | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/library/src/tensor_operation_instance/gpu/CMakeLists.txt b/library/src/tensor_operation_instance/gpu/CMakeLists.txt index df3283b543..bc66fe0bed 100644 --- a/library/src/tensor_operation_instance/gpu/CMakeLists.txt +++ b/library/src/tensor_operation_instance/gpu/CMakeLists.txt @@ -64,9 +64,9 @@ function(add_instance_library INSTANCE_NAME) list(REMOVE_ITEM ARGN "${source}") endif() endforeach() - # Do not build mha instances if gfx94 targets are not on the target list + # Do not build mha instances if gfx94 or gfx90a targets are not on the target list foreach(source IN LISTS ARGN) - if(NOT INST_TARGETS MATCHES "gfx94" AND source MATCHES "mha") + if(NOT INST_TARGETS MATCHES "gfx94" AND NOT INST_TARGETS MATCHES "gfx90a" AND source MATCHES "mha") message("removing mha instance ${source} ") list(REMOVE_ITEM ARGN "${source}") endif() @@ -85,7 +85,7 @@ function(add_instance_library INSTANCE_NAME) elseif(ARGN MATCHES "_wmma") list(REMOVE_ITEM INST_TARGETS gfx900 gfx906 gfx908 gfx90a gfx940 gfx941 gfx942 gfx1030) elseif(ARGN MATCHES "mha") - list(REMOVE_ITEM INST_TARGETS gfx900 gfx906 gfx908 gfx90a gfx1030 gfx1100 gfx1101 gfx1102 gfx1103 gfx1200 gfx1201) + list(REMOVE_ITEM INST_TARGETS gfx900 gfx906 gfx908 gfx1030 gfx1100 gfx1101 gfx1102 gfx1103 gfx1200 gfx1201) endif() set(offload_targets) foreach(target IN LISTS INST_TARGETS) @@ -320,8 +320,7 @@ if(CK_DEVICE_CONV_INSTANCES) endif() if(CK_DEVICE_MHA_INSTANCES) set(gpu_list ${INST_TARGETS}) - list(FILTER gpu_list INCLUDE REGEX "^gfx94") - if(gpu_list) + if(gpu_list MATCHES "gfx94" OR gpu_list MATCHES "gfx90a") add_library(device_mha_operations STATIC ${CK_DEVICE_MHA_INSTANCES}) add_library(composablekernels::device_mha_operations ALIAS device_mha_operations) target_compile_features(device_mha_operations PUBLIC) From aeb7c91f48a0e8fa1e288d91f719415282c03f03 Mon Sep 17 00:00:00 2001 From: macurtis-amd Date: Wed, 2 Oct 2024 15:56:22 -0500 Subject: [PATCH 2/8] Fix compilation errors generated by forthcoming Clang changes (#1544) Without this change, the following diagnostic is generated: a template argument list is expected after a name prefixed by the template keyword [-Wmissing-template-arg-list-after-template-kw] See C++17 spec [temp.names] p5. --- ...ckwise_gemm_pipeline_xdlops_v1_ab_scale.hpp | 9 +++++---- ...ckwise_gemm_pipeline_xdlops_v2_ab_scale.hpp | 18 ++++++++++-------- ...ckwise_gemm_pipeline_xdlops_v3_ab_scale.hpp | 9 +++++---- 3 files changed, 20 insertions(+), 16 deletions(-) diff --git a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_v1_ab_scale.hpp b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_v1_ab_scale.hpp index 121593d3cc..821bbb0051 100644 --- a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_v1_ab_scale.hpp +++ b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_v1_ab_scale.hpp @@ -308,7 +308,7 @@ struct BlockwiseGemmXdlops_pipeline_v1_ab_scale::type; - xdlops_gemm.template Run( + xdlops_gemm.template Run<>( a_thread_vec.template AsType(), b_thread_vec.template AsType(), c_thread_buf_per_scale.GetVectorTypeReference(I0)); @@ -390,9 +390,10 @@ struct BlockwiseGemmXdlops_pipeline_v1_ab_scale::type; - xdlops_gemm.template Run(a_thread_vec.template AsType(), - b_thread_vec.template AsType(), - c_thread_buf_per_scale.GetVectorTypeReference(I0)); + xdlops_gemm.template Run<>( + a_thread_vec.template AsType(), + b_thread_vec.template AsType(), + c_thread_buf_per_scale.GetVectorTypeReference(I0)); }); static_for<0, xdlops_gemm.GetRegSizePerXdlops(), 1>{}([&](auto t) { constexpr index_t c_offset = diff --git a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_v2_ab_scale.hpp b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_v2_ab_scale.hpp index cb7cf605be..40fa776484 100644 --- a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_v2_ab_scale.hpp +++ b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_v2_ab_scale.hpp @@ -350,7 +350,7 @@ struct BlockwiseGemmXdlops_pipeline_v2_ab_scale::type; - xdlops_gemm.template Run( + xdlops_gemm.template Run<>( a_thread_vec.template AsType(), b_thread_vec.template AsType(), c_thread_buf_per_scale.GetVectorTypeReference(I0)); @@ -443,7 +443,7 @@ struct BlockwiseGemmXdlops_pipeline_v2_ab_scale::type; - xdlops_gemm.template Run( + xdlops_gemm.template Run<>( a_thread_vec.template AsType(), b_thread_vec.template AsType(), c_thread_buf_per_scale.GetVectorTypeReference(I0)); @@ -518,9 +518,10 @@ struct BlockwiseGemmXdlops_pipeline_v2_ab_scale::type; - xdlops_gemm.template Run(a_thread_vec.template AsType(), - b_thread_vec.template AsType(), - c_thread_buf_per_scale.GetVectorTypeReference(I0)); + xdlops_gemm.template Run<>( + a_thread_vec.template AsType(), + b_thread_vec.template AsType(), + c_thread_buf_per_scale.GetVectorTypeReference(I0)); }); static_for<0, xdlops_gemm.GetRegSizePerXdlops(), 1>{}([&](auto t) { constexpr index_t c_offset = @@ -575,9 +576,10 @@ struct BlockwiseGemmXdlops_pipeline_v2_ab_scale::type; - xdlops_gemm.template Run(a_thread_vec.template AsType(), - b_thread_vec.template AsType(), - c_thread_buf_per_scale.GetVectorTypeReference(I0)); + xdlops_gemm.template Run<>( + a_thread_vec.template AsType(), + b_thread_vec.template AsType(), + c_thread_buf_per_scale.GetVectorTypeReference(I0)); }); static_for<0, xdlops_gemm.GetRegSizePerXdlops(), 1>{}([&](auto t) { constexpr index_t c_offset = diff --git a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_v3_ab_scale.hpp b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_v3_ab_scale.hpp index 66c9a5c339..de542866a6 100644 --- a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_v3_ab_scale.hpp +++ b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_v3_ab_scale.hpp @@ -427,7 +427,7 @@ struct BlockwiseGemmXdlops_pipeline_v3_ab_scale::type; - xdlops_gemm.template Run( + xdlops_gemm.template Run<>( a_thread_vec.template AsType(), b_thread_vec.template AsType(), c_thread_buf_per_scale.GetVectorTypeReference(I0)); @@ -504,9 +504,10 @@ struct BlockwiseGemmXdlops_pipeline_v3_ab_scale::type; - xdlops_gemm.template Run(a_thread_vec.template AsType(), - b_thread_vec.template AsType(), - c_thread_buf_per_scale.GetVectorTypeReference(I0)); + xdlops_gemm.template Run<>( + a_thread_vec.template AsType(), + b_thread_vec.template AsType(), + c_thread_buf_per_scale.GetVectorTypeReference(I0)); }); static_for<0, xdlops_gemm.GetRegSizePerXdlops(), 1>{}([&](auto t) { constexpr index_t c_offset = From 6b54d2faf8d0b106fb31719654b6b4d5f18552f2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Bart=C5=82omiej=20Kocot?= Date: Fri, 4 Oct 2024 17:32:43 +0200 Subject: [PATCH 3/8] Fix grouped gemm check to avoid overflow (#1545) --- .../device/impl/device_grouped_gemm_multi_abd_xdl_fixed_nk.hpp | 2 +- ...ce_grouped_gemm_multiple_d_splitk_xdl_cshuffle_two_stage.hpp | 2 +- .../device_grouped_gemm_multiple_d_xdl_cshuffle_tile_loop.hpp | 2 +- .../gpu/device/impl/device_grouped_gemm_xdl_fixed_nk.hpp | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/include/ck/tensor_operation/gpu/device/impl/device_grouped_gemm_multi_abd_xdl_fixed_nk.hpp b/include/ck/tensor_operation/gpu/device/impl/device_grouped_gemm_multi_abd_xdl_fixed_nk.hpp index 1f60818e39..21afc06040 100644 --- a/include/ck/tensor_operation/gpu/device/impl/device_grouped_gemm_multi_abd_xdl_fixed_nk.hpp +++ b/include/ck/tensor_operation/gpu/device/impl/device_grouped_gemm_multi_abd_xdl_fixed_nk.hpp @@ -64,7 +64,7 @@ __global__ void const index_t N = gemm_desc_ptr[group_id].N; const index_t K = gemm_desc_ptr[group_id].K; - if(M * N * K == 0) + if(M == 0 || N == 0 || K == 0) return; const auto StrideAs = gemm_desc_ptr[group_id].StrideAs; diff --git a/include/ck/tensor_operation/gpu/device/impl/device_grouped_gemm_multiple_d_splitk_xdl_cshuffle_two_stage.hpp b/include/ck/tensor_operation/gpu/device/impl/device_grouped_gemm_multiple_d_splitk_xdl_cshuffle_two_stage.hpp index 8354335577..68c6dcc0f5 100644 --- a/include/ck/tensor_operation/gpu/device/impl/device_grouped_gemm_multiple_d_splitk_xdl_cshuffle_two_stage.hpp +++ b/include/ck/tensor_operation/gpu/device/impl/device_grouped_gemm_multiple_d_splitk_xdl_cshuffle_two_stage.hpp @@ -345,7 +345,7 @@ struct DeviceGroupedGemmMultipleDSplitKXdlCShuffleTwoStage const index_t N = gemm_descs[i].N_; const index_t K = gemm_descs[i].K_; - if(M * N * K == 0) + if(M == 0 || N == 0 || K == 0) { skipped_group_count_++; continue; diff --git a/include/ck/tensor_operation/gpu/device/impl/device_grouped_gemm_multiple_d_xdl_cshuffle_tile_loop.hpp b/include/ck/tensor_operation/gpu/device/impl/device_grouped_gemm_multiple_d_xdl_cshuffle_tile_loop.hpp index 70011124fc..2884e558cd 100644 --- a/include/ck/tensor_operation/gpu/device/impl/device_grouped_gemm_multiple_d_xdl_cshuffle_tile_loop.hpp +++ b/include/ck/tensor_operation/gpu/device/impl/device_grouped_gemm_multiple_d_xdl_cshuffle_tile_loop.hpp @@ -109,7 +109,7 @@ __global__ void N = gemm_desc_ptr[group_id].N; K = gemm_desc_ptr[group_id].K; - if(M * N * K == 0) + if(M == 0 || N == 0 || K == 0) { grid_size_grp = 0; continue; diff --git a/include/ck/tensor_operation/gpu/device/impl/device_grouped_gemm_xdl_fixed_nk.hpp b/include/ck/tensor_operation/gpu/device/impl/device_grouped_gemm_xdl_fixed_nk.hpp index c98ec6e2aa..ac05a0703f 100644 --- a/include/ck/tensor_operation/gpu/device/impl/device_grouped_gemm_xdl_fixed_nk.hpp +++ b/include/ck/tensor_operation/gpu/device/impl/device_grouped_gemm_xdl_fixed_nk.hpp @@ -68,7 +68,7 @@ __global__ void const index_t N = gemm_desc_ptr[group_id].N; const index_t K = gemm_desc_ptr[group_id].K; - if(M * N * K == 0) + if(M == 0 || N == 0 || K == 0) return; const auto StrideA = gemm_desc_ptr[group_id].StrideA; From b545de175a7a1410baaea70f1f32ecb819e1d056 Mon Sep 17 00:00:00 2001 From: arai713 <67439843+arai713@users.noreply.github.com> Date: Fri, 4 Oct 2024 10:51:50 -0700 Subject: [PATCH 4/8] Codegen build (#1526) * updating codegen build for MIOpen access: adding .cmake for codegen component (cherry picked from commit 652a7c046381526947f507a89299aa92d89dbd02) * updating CMake (cherry picked from commit a685822e361045f3ef02a2f60c1c0eadd9cc4c85) --- CMakeLists.txt | 4 +++- cmake/Embed.cmake | 4 +++- codegen/CMakeLists.txt | 9 ++++++++- 3 files changed, 14 insertions(+), 3 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index fd321f7722..dc73b5f4d4 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -569,7 +569,9 @@ if(NOT DEFINED INSTANCES_ONLY) PACKAGE_NAME examples ) add_subdirectory(example) - add_subdirectory(test) + if(BUILD_TESTING) + add_subdirectory(test) + endif() rocm_package_setup_component(profiler LIBRARY_NAME composablekernel diff --git a/cmake/Embed.cmake b/cmake/Embed.cmake index 4bc638b446..3946cf4e8d 100644 --- a/cmake/Embed.cmake +++ b/cmake/Embed.cmake @@ -233,6 +233,8 @@ function(add_embed_library EMBED_NAME) else() target_sources(${EMBED_NAME} INTERFACE $) endif() - target_include_directories(${EMBED_NAME} INTERFACE "${EMBED_DIR}/include") + target_include_directories(${EMBED_NAME} INTERFACE + $ + $) endfunction() diff --git a/codegen/CMakeLists.txt b/codegen/CMakeLists.txt index 3b3e9f06ee..2492804f28 100644 --- a/codegen/CMakeLists.txt +++ b/codegen/CMakeLists.txt @@ -39,6 +39,7 @@ set_target_properties(ck_host PROPERTIES target_include_directories(ck_host PUBLIC $ + $ ) add_executable(ck-template-driver driver/main.cpp) @@ -48,6 +49,12 @@ rocm_install( TARGETS ck_host ck_headers EXPORT ck_hostTargets ) +rocm_install(EXPORT ck_hostTargets + FILE composable_kernelck_hostTargets.cmake + NAMESPACE composable_kernel:: + DESTINATION ${CMAKE_INSTALL_LIBDIR}/cmake/composable_kernel) rocm_install(DIRECTORY include/ck DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}) -add_subdirectory(test) +if(BUILD_TESTING) + add_subdirectory(test) +endif() From c24fae234600aa2863e945d072e6f5b3aec2a6b2 Mon Sep 17 00:00:00 2001 From: kylasa Date: Fri, 4 Oct 2024 11:48:47 -0700 Subject: [PATCH 5/8] Adding seed and offset pointer support to the philox random number generator. (#1523) * Adding seed and offset pointer support to the philox random number generator. * Separating seed and offset pointer checks with different condition statements. * Changes include, adding support for device seed and offset pointers, union is used to store seed/offset values and device pointers to minimize device SGPRs. * Correcting a typo in the readme file * Re-format files using remod.py * Use STL type for API parameters * Use simpler struct design for drop_seed & drop_offset * Undo unnecessary changes * Sync kargs style for fmha_fwd.hpp/.cpp * Use templated union to reduce code * Use structured binding to make code more readable --------- Co-authored-by: Sudhir Kylasa Co-authored-by: Po Yen Chen --- example/ck_tile/01_fmha/README.md | 7 +- example/ck_tile/01_fmha/fmha_bwd.cpp | 23 ++++- example/ck_tile/01_fmha/fmha_bwd.hpp | 6 +- example/ck_tile/01_fmha/fmha_fwd.cpp | 23 ++++- example/ck_tile/01_fmha/fmha_fwd.hpp | 6 +- .../ops/fmha/kernel/fmha_bwd_kernel.hpp | 91 ++++++++++++++++--- .../ops/fmha/kernel/fmha_fwd_kernel.hpp | 83 ++++++++++++++--- 7 files changed, 205 insertions(+), 34 deletions(-) diff --git a/example/ck_tile/01_fmha/README.md b/example/ck_tile/01_fmha/README.md index 0bb5408772..0803d54d66 100644 --- a/example/ck_tile/01_fmha/README.md +++ b/example/ck_tile/01_fmha/README.md @@ -70,8 +70,13 @@ args: -seed random seed used for initializing input tensors. 0 for non-deterministic seed (default:11939) -warmup number of iterations before benchmark the kernel (default:5) -repeat number of iterations to benchmark the kernel (default:20) + -drop_seed seed for the random number generator for the dropout layer, default is 1 +-drop_offset offset for the dropout layer which is used during random number generation, default is 0 + -drop_prefs flag to indicate `drop_seed` and `drop_offset` values if present on the GPU, default is 0, 0 - host, 1 - GPU ``` -Example: `./bin/tile_example_fmha_fwd -b=1 -h=16 -s=16384 -d=128` will run a fmha case with batch=1, nhead=16, sequence length=16384, hdim=128, fp16 case. +Example 1: `./bin/tile_example_fmha_fwd -b=1 -h=16 -s=16384 -d=128` will run a fmha case with batch=1, nhead=16, sequence length=16384, hdim=128, fp16 case. +Example 2: `./bin/tile_example_fmha_fwd -b=1 -h=8 -s=16384 -d=64 -drop_prefs=1 -drop_seed=10 -drop_offset=1234` will run a fmha case with + batch=1, nhead=8, sequence length=16384, hdim=64, drop_seed=0 (in GPU memory), drop_offset=1234 (in GPU memory) fp16 case ## support features Currently we are still in rapid development stage, so more features/optimizations will be coming soon. diff --git a/example/ck_tile/01_fmha/fmha_bwd.cpp b/example/ck_tile/01_fmha/fmha_bwd.cpp index c2f554f6cc..2d76627a72 100644 --- a/example/ck_tile/01_fmha/fmha_bwd.cpp +++ b/example/ck_tile/01_fmha/fmha_bwd.cpp @@ -85,6 +85,9 @@ auto create_args(int argc, char* argv[]) .insert("p_drop", "0", "0~1 probability of dropout") .insert("drop_seed", "1", "seed for random number generator") .insert("drop_offset", "0", "offset for random number generator") + .insert("drop_prefs", + "0", + "seed and offset values are present on GPU; 0 - host, 1 - device/GPU") .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") @@ -158,6 +161,8 @@ bool run(const ck_tile::ArgParser& arg_parser) float p_drop = arg_parser.get_float("p_drop"); uint64_t drop_seed = arg_parser.get_uint64("drop_seed"); uint64_t drop_offset = arg_parser.get_uint64("drop_offset"); + bool drop_prefs = arg_parser.get_bool("drop_prefs"); + if(use_dbias && bias.type != bias_enum::elementwise_bias) { std::cerr << "dbias only exists when bias type is elementwise" << std::endl; @@ -381,6 +386,8 @@ bool run(const ck_tile::ArgParser& arg_parser) ck_tile::DeviceMem dbias_buf(dbias_host.get_element_space_size_in_bytes()); 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 drop_seed_buf(drop_prefs ? sizeof(uint64_t) : 0); + ck_tile::DeviceMem drop_offset_buf(drop_prefs ? sizeof(uint64_t) : 0); 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()); @@ -391,6 +398,8 @@ bool run(const ck_tile::ArgParser& arg_parser) do_buf.ToDevice(do_host.data()); seqstart_q.ToDevice(seqstart_q_host.data()); seqstart_k.ToDevice(seqstart_k_host.data()); + drop_seed_buf.ToDevice(drop_prefs ? &drop_seed : nullptr); + drop_offset_buf.ToDevice(drop_prefs ? &drop_offset : nullptr); alibi_slope_buf.ToDevice(alibi_slope_host.data()); // clang-format off @@ -472,6 +481,18 @@ bool run(const ck_tile::ArgParser& arg_parser) const ck_tile::index_t split_stride_dq_acc = (shape_batch * nhead * shape_seqlen_q * hdim_q); + const auto drop_seed_offset = [&]() -> decltype(fmha_bwd_args::drop_seed_offset) { + if(drop_prefs) + { + return std::make_pair(drop_seed_buf.GetDeviceBuffer(), + drop_offset_buf.GetDeviceBuffer()); + } + else + { + return std::make_pair(drop_seed, drop_offset); + } + }(); + return fmha_bwd_args{q_buf.GetDeviceBuffer(), k_buf.GetDeviceBuffer(), v_buf.GetDeviceBuffer(), @@ -545,7 +566,7 @@ bool run(const ck_tile::ArgParser& arg_parser) static_cast(mask.type), p_drop, p_undrop, - {drop_seed, drop_offset}}; + drop_seed_offset}; }(); float ave_time = fmha_bwd(fmha_traits, fmha_args, stream_config); diff --git a/example/ck_tile/01_fmha/fmha_bwd.hpp b/example/ck_tile/01_fmha/fmha_bwd.hpp index aea42515dc..3b21a3257f 100644 --- a/example/ck_tile/01_fmha/fmha_bwd.hpp +++ b/example/ck_tile/01_fmha/fmha_bwd.hpp @@ -9,7 +9,10 @@ #include "ck_tile/ops/epilogue.hpp" #include "mask.hpp" #include "bias.hpp" + #include +#include +#include template struct FmhaBwdTypeConfig; @@ -135,7 +138,8 @@ struct fmha_bwd_args ck_tile::index_t mask_type; float p_drop; float p_undrop; - std::tuple drop_seed_offset; + std::variant, std::pair> + drop_seed_offset; }; template diff --git a/example/ck_tile/01_fmha/fmha_fwd.cpp b/example/ck_tile/01_fmha/fmha_fwd.cpp index b9cb9a1ec2..6d519a7ea8 100644 --- a/example/ck_tile/01_fmha/fmha_fwd.cpp +++ b/example/ck_tile/01_fmha/fmha_fwd.cpp @@ -122,6 +122,9 @@ auto create_args(int argc, char* argv[]) .insert("p_drop", "0", "0~1 probability of dropout") .insert("drop_seed", "1", "seed for random number generator") .insert("drop_offset", "0", "offset for random number generator") + .insert("drop_prefs", + "0", + "seed and offset values are present on GPU; 0 - host, 1 - device/GPU") .insert("timer", "gpu", "gpu:gpu timer, cpu:cpu timer") .insert( "rotary_dim", "0", "RoPE rotary dimension. rotary_dim <= 0 means not apply RoPE at all") @@ -442,6 +445,8 @@ bool run(const ck_tile::ArgParser& arg_parser) float p_drop = arg_parser.get_float("p_drop"); uint64_t drop_seed = arg_parser.get_uint64("drop_seed"); uint64_t drop_offset = arg_parser.get_uint64("drop_offset"); + bool drop_prefs = arg_parser.get_bool("drop_prefs"); + if(p_drop < 0.0f || p_drop > 1.0f) { std::cerr << "The value of p_drop should be 0~1" << std::endl; @@ -756,6 +761,8 @@ bool run(const ck_tile::ArgParser& arg_parser) need_append_kvcache ? cache_seqlen_ks.size() * sizeof(int32_t) : 0); ck_tile::DeviceMem rotary_cos_buf(rotary_cos_host.get_element_space_size_in_bytes()); ck_tile::DeviceMem rotary_sin_buf(rotary_sin_host.get_element_space_size_in_bytes()); + ck_tile::DeviceMem drop_seed_buf(drop_prefs ? sizeof(uint64_t) : 0); + ck_tile::DeviceMem drop_offset_buf(drop_prefs ? sizeof(uint64_t) : 0); ck_tile::DeviceMem randval_buf(randval_host.get_element_space_size_in_bytes()); ck_tile::DeviceMem alibi_slope_buf(alibi_slope_host.get_element_space_size_in_bytes()); ck_tile::DeviceMem block_table_buf(block_table_host.get_element_space_size_in_bytes()); @@ -774,6 +781,8 @@ bool run(const ck_tile::ArgParser& arg_parser) cache_seqlen_k_buf.ToDevice(need_append_kvcache ? cache_seqlen_ks.data() : nullptr); rotary_cos_buf.ToDevice(rotary_cos_host.data()); rotary_sin_buf.ToDevice(rotary_sin_host.data()); + drop_seed_buf.ToDevice(drop_prefs ? &drop_seed : nullptr); + drop_offset_buf.ToDevice(drop_prefs ? &drop_offset : nullptr); alibi_slope_buf.ToDevice(alibi_slope_host.data()); block_table_buf.ToDevice(block_table_host.data()); cache_batch_idx_buf.ToDevice(cache_batch_idx_host.data()); @@ -1013,9 +1022,17 @@ bool run(const ck_tile::ArgParser& arg_parser) args.nhead_stride_randval = nhead_stride_randval; args.batch_stride_randval = batch_stride_randval; - args.p_drop = p_drop; - args.s_randval = s_randval; - args.drop_seed_offset = std::tie(drop_seed, drop_offset); + args.p_drop = p_drop; + args.s_randval = s_randval; + if(drop_prefs) + { + args.drop_seed_offset = std::make_pair(drop_seed_buf.GetDeviceBuffer(), + drop_offset_buf.GetDeviceBuffer()); + } + else + { + args.drop_seed_offset = std::make_pair(drop_seed, drop_offset); + } } else if constexpr(std::is_same_v>) { diff --git a/example/ck_tile/01_fmha/fmha_fwd.hpp b/example/ck_tile/01_fmha/fmha_fwd.hpp index 5dcad7907f..251e61bc76 100644 --- a/example/ck_tile/01_fmha/fmha_fwd.hpp +++ b/example/ck_tile/01_fmha/fmha_fwd.hpp @@ -13,6 +13,8 @@ #include "rotary.hpp" #include +#include +#include template struct FmhaFwdTypeConfig; @@ -144,7 +146,9 @@ struct fmha_fwd_args float p_drop; bool s_randval; - std::tuple drop_seed_offset; + + std::variant, std::pair> + drop_seed_offset; }; struct fmha_fwd_splitkv_args 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 167494b193..c5858a20f7 100644 --- a/include/ck_tile/ops/fmha/kernel/fmha_bwd_kernel.hpp +++ b/include/ck_tile/ops/fmha/kernel/fmha_bwd_kernel.hpp @@ -6,8 +6,11 @@ #include "ck_tile/core.hpp" #include "ck_tile/ops/common.hpp" #include "ck_tile/ops/fmha/block/block_attention_bias_enum.hpp" + #include #include +#include +#include // S[seqlen_q, seqlen_k] = Q[seqlen_q, hdim_q] @ K[seqlen_k, hdim_q] // S'[seqlen_q, seqlen_k] = S[seqlen_q, seqlen_k] * Scale[1] @@ -194,11 +197,23 @@ struct FmhaBwdDQDKDVKernel ck_tile::GenericAttentionMaskEnum mask_type; }; - struct FmhaBwdCommonDropoutKargs + struct FmhaBwdDropoutSeedOffset { - void init_dropout(const float p_drop, - const std::tuple& drop_seed_offset, - const float raw_scale) + template + union ValueOrPointer + { + T val; + const T* ptr; + }; + + ValueOrPointer drop_seed; + ValueOrPointer drop_offset; + bool is_drop_seed_offset_from_host; + }; + + struct FmhaBwdCommonDropoutKargs : FmhaBwdDropoutSeedOffset + { + void init_dropout(float p_drop, uint64_t seed, uint64_t offset, float raw_scale) { float p_undrop = 1.0 - p_drop; p_undrop_in_uint8_t = @@ -206,23 +221,41 @@ struct FmhaBwdDQDKDVKernel rp_undrop = 1.0 / p_undrop; scale_rp_undrop = rp_undrop * raw_scale; - drop_seed = std::get<0>(drop_seed_offset); - drop_offset = std::get<1>(drop_seed_offset); + this->drop_seed.val = seed; + this->drop_offset.val = offset; + this->is_drop_seed_offset_from_host = true; } + + void init_dropout(float p_drop, + const uint64_t* seed_ptr, + const uint64_t* offset_ptr, + float raw_scale) + { + float p_undrop = 1.0 - p_drop; + p_undrop_in_uint8_t = + uint8_t(std::floor(p_undrop * std::numeric_limits::max())); + rp_undrop = 1.0 / p_undrop; + scale_rp_undrop = rp_undrop * raw_scale; + + this->drop_seed.ptr = seed_ptr; + this->drop_offset.ptr = offset_ptr; + this->is_drop_seed_offset_from_host = false; + } + float rp_undrop = 1; float scale_rp_undrop = 1; uint8_t p_undrop_in_uint8_t = std::numeric_limits::max(); - uint64_t drop_seed = 1; - uint64_t drop_offset = 0; void* rand_val_ptr = nullptr; ck_tile::index_t stride_randval = 0; ck_tile::index_t nhead_stride_randval = 0; }; + struct FmhaBwdBatchModeDropoutKargs : FmhaBwdCommonDropoutKargs { ck_tile::index_t batch_stride_randval = 0; }; + struct FmhaBwdDeterministicKargs { ck_tile::index_t split_stride_dq_acc = 0; @@ -327,7 +360,8 @@ struct FmhaBwdDQDKDVKernel ck_tile::index_t window_size_right, ck_tile::index_t mask_type, float p_drop, - const std::tuple& drop_seed_offset) + std::variant, std::pair> + drop_seed_offset) { Kargs kargs{{q_ptr, k_ptr, @@ -405,7 +439,20 @@ struct FmhaBwdDQDKDVKernel if constexpr(kHasDropout) { - kargs.init_dropout(p_drop, drop_seed_offset, scale); + if(drop_seed_offset.index() == 0) // seed & offset come from host + { + const auto& [seed, offset] = std::get<0>(drop_seed_offset); + kargs.init_dropout(p_drop, seed, offset, scale); + } + else // seed & offset come from device + { + const auto& [seed_ptr, offset_ptr] = std::get<1>(drop_seed_offset); + kargs.init_dropout(p_drop, + reinterpret_cast(seed_ptr), + reinterpret_cast(offset_ptr), + scale); + } + if constexpr(kIsStoreRandval) { kargs.rand_val_ptr = rand_val_ptr; @@ -471,7 +518,8 @@ struct FmhaBwdDQDKDVKernel ck_tile::index_t window_size_right, ck_tile::index_t mask_type, float p_drop, - const std::tuple& drop_seed_offset) + std::variant, std::pair> + drop_seed_offset) { Kargs kargs{{q_ptr, k_ptr, @@ -539,7 +587,20 @@ struct FmhaBwdDQDKDVKernel } if constexpr(kHasDropout) { - kargs.init_dropout(p_drop, drop_seed_offset, scale); + if(drop_seed_offset.index() == 0) // seed & offset come from host + { + const auto& [seed, offset] = std::get<0>(drop_seed_offset); + kargs.init_dropout(p_drop, seed, offset, scale); + } + else // seed & offset come from device + { + const auto& [seed_ptr, offset_ptr] = std::get<1>(drop_seed_offset); + kargs.init_dropout(p_drop, + reinterpret_cast(seed_ptr), + reinterpret_cast(offset_ptr), + scale); + } + if constexpr(kIsStoreRandval) { kargs.rand_val_ptr = rand_val_ptr; @@ -958,8 +1019,10 @@ struct FmhaBwdDQDKDVKernel return FmhaDropout{i_batch_, i_nhead_, kargs.num_head_q, - kargs.drop_seed, - kargs.drop_offset, + kargs.is_drop_seed_offset_from_host ? kargs.drop_seed.val + : *kargs.drop_seed.ptr, + kargs.is_drop_seed_offset_from_host ? kargs.drop_offset.val + : *kargs.drop_offset.ptr, kargs.rp_undrop, kargs.p_undrop_in_uint8_t}; } 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 49ef7bf6d9..adabda165c 100644 --- a/include/ck_tile/ops/fmha/kernel/fmha_fwd_kernel.hpp +++ b/include/ck_tile/ops/fmha/kernel/fmha_fwd_kernel.hpp @@ -6,8 +6,11 @@ #include "ck_tile/core.hpp" #include "ck_tile/ops/common.hpp" #include "ck_tile/ops/fmha/block/block_attention_bias_enum.hpp" + #include #include +#include +#include // S[seqlen_q, seqlen_k] = Q[seqlen_q, hdim_q] @ K[seqlen_k, hdim_q] // S'[seqlen_q, seqlen_k] = S[seqlen_q, seqlen_k] * Scale[1] @@ -170,29 +173,55 @@ struct FmhaFwdKernel ck_tile::index_t batch_stride_lse = 0; }; - struct FmhaFwdCommonDropoutKargs + struct FmhaFwdDropoutSeedOffset { - void init_dropout(const float p_drop, - const std::tuple& drop_seed_offset) + template + union ValueOrPointer + { + T val; + const T* ptr; + }; + + ValueOrPointer drop_seed; + ValueOrPointer drop_offset; + bool is_drop_seed_offset_from_host; + }; + + struct FmhaFwdCommonDropoutKargs : FmhaFwdDropoutSeedOffset + { + void init_dropout(float p_drop, uint64_t seed, uint64_t offset) { float p_undrop = 1.0 - p_drop; p_undrop_in_uint8_t = uint8_t(std::floor(p_undrop * std::numeric_limits::max())); rp_undrop = 1.0 / p_undrop; - drop_seed = std::get<0>(drop_seed_offset); - drop_offset = std::get<1>(drop_seed_offset); + this->drop_seed.val = seed; + this->drop_offset.val = offset; + this->is_drop_seed_offset_from_host = true; } + + void init_dropout(float p_drop, const uint64_t* seed_ptr, const uint64_t* offset_ptr) + { + float p_undrop = 1.0 - p_drop; + p_undrop_in_uint8_t = + uint8_t(std::floor(p_undrop * std::numeric_limits::max())); + rp_undrop = 1.0 / p_undrop; + + this->drop_seed.ptr = seed_ptr; + this->drop_offset.ptr = offset_ptr; + this->is_drop_seed_offset_from_host = false; + } + float 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; ck_tile::index_t stride_randval = 0; ck_tile::index_t nhead_stride_randval = 0; }; + struct FmhaFwdBatchModeDropoutKargs : FmhaFwdCommonDropoutKargs { ck_tile::index_t batch_stride_randval = 0; @@ -278,7 +307,8 @@ struct FmhaFwdKernel ck_tile::index_t mask_type, float p_drop, bool s_randval, - const std::tuple& drop_seed_offset) + std::variant, std::pair> + drop_seed_offset) { Kargs kargs{{q_ptr, k_ptr, @@ -344,7 +374,19 @@ struct FmhaFwdKernel } if constexpr(kHasDropout) { - kargs.init_dropout(p_drop, drop_seed_offset); + if(drop_seed_offset.index() == 0) // seed & offset come from host + { + const auto& [seed, offset] = std::get<0>(drop_seed_offset); + kargs.init_dropout(p_drop, seed, offset); + } + else // seed & offset come from device + { + const auto& [seed_ptr, offset_ptr] = std::get<1>(drop_seed_offset); + kargs.init_dropout(p_drop, + reinterpret_cast(seed_ptr), + reinterpret_cast(offset_ptr)); + } + kargs.rand_val_ptr = rand_val_ptr; kargs.stride_randval = stride_randval; kargs.nhead_stride_randval = nhead_stride_randval; @@ -392,7 +434,8 @@ struct FmhaFwdKernel ck_tile::index_t mask_type, float p_drop, bool s_randval, - const std::tuple& drop_seed_offset) + std::variant, std::pair> + drop_seed_offset) { Kargs kargs{{q_ptr, k_ptr, @@ -455,7 +498,19 @@ struct FmhaFwdKernel } if constexpr(kHasDropout) { - kargs.init_dropout(p_drop, drop_seed_offset); + if(drop_seed_offset.index() == 0) // seed & offset come from host + { + const auto& [seed, offset] = std::get<0>(drop_seed_offset); + kargs.init_dropout(p_drop, seed, offset); + } + else // seed & offset come from device + { + const auto& [seed_ptr, offset_ptr] = std::get<1>(drop_seed_offset); + kargs.init_dropout(p_drop, + reinterpret_cast(seed_ptr), + reinterpret_cast(offset_ptr)); + } + kargs.rand_val_ptr = rand_val_ptr; kargs.stride_randval = stride_randval; kargs.nhead_stride_randval = nhead_stride_randval; @@ -748,8 +803,10 @@ struct FmhaFwdKernel return BlockDropout{i_batch_, i_nhead_, kargs.num_head_q, - kargs.drop_seed, - kargs.drop_offset, + kargs.is_drop_seed_offset_from_host ? kargs.drop_seed.val + : *kargs.drop_seed.ptr, + kargs.is_drop_seed_offset_from_host ? kargs.drop_offset.val + : *kargs.drop_offset.ptr, kargs.rp_undrop, kargs.p_undrop_in_uint8_t, kargs.is_store_randval}; From 0023f01ab02b9cc05a98ae1a7753df1481252e4d Mon Sep 17 00:00:00 2001 From: rocking Date: Mon, 7 Oct 2024 14:25:53 +0800 Subject: [PATCH 6/8] [Ck tile] Support layernorm one pass (#1512) * Fix compile error * Add one pass pipeline * Extract creating tile_window to operator() * clang format * reduce duplicated code * do not hardcode * Support padding in layernorm --------- Co-authored-by: Po Yen Chen --- .../02_layernorm2d/layernorm2d_fwd.cpp | 4 +- .../kernel/layernorm2d_fwd_kernel.hpp | 333 +++++++++++++----- .../block_layernorm2d_fwd_problem.hpp | 22 +- 3 files changed, 263 insertions(+), 96 deletions(-) diff --git a/example/ck_tile/02_layernorm2d/layernorm2d_fwd.cpp b/example/ck_tile/02_layernorm2d/layernorm2d_fwd.cpp index 9cbd286104..35f291e060 100644 --- a/example/ck_tile/02_layernorm2d/layernorm2d_fwd.cpp +++ b/example/ck_tile/02_layernorm2d/layernorm2d_fwd.cpp @@ -35,7 +35,9 @@ float layernorm2d_fwd(layernorm2d_fwd_traits t, YDataType, MeanDataType, InvStdDataType, - Shape>; + Shape, + true, + true>; using Kernel = ck_tile::Layernorm2dFwd; diff --git a/include/ck_tile/ops/layernorm2d/kernel/layernorm2d_fwd_kernel.hpp b/include/ck_tile/ops/layernorm2d/kernel/layernorm2d_fwd_kernel.hpp index 4be3e56874..468df793da 100644 --- a/include/ck_tile/ops/layernorm2d/kernel/layernorm2d_fwd_kernel.hpp +++ b/include/ck_tile/ops/layernorm2d/kernel/layernorm2d_fwd_kernel.hpp @@ -31,8 +31,14 @@ struct Layernorm2dFwd static constexpr ck_tile::index_t kMPerBlock = Problem::BlockShape::kMPerBlock; static constexpr ck_tile::index_t kNPerBlock = Problem::BlockShape::kNPerBlock; + static constexpr bool kPadM = Problem::kPadM; + static constexpr bool kPadN = Problem::kPadN; static constexpr ck_tile::index_t kNThreadPerWarp = Problem::BlockShape::kNThreadPerWarp; + static constexpr ck_tile::index_t kNPerThread = Problem::BlockShape::kNPerThread; + + static constexpr auto I0 = number<0>{}; + static constexpr auto I1 = number<1>{}; struct Kargs { @@ -96,19 +102,25 @@ struct Layernorm2dFwd sequence<2>>{}); } - template - CK_TILE_DEVICE static constexpr auto GetNPerThread(Dstr) + CK_TILE_DEVICE static int GetWelfordMaxCount(int N) { - constexpr auto nDstrSpan = Dstr::get_distributed_spans().template at<1>(); + constexpr ck_tile::index_t kNThreadPerBlock = kNPerBlock / kNPerThread; - using Lengths = decltype(nDstrSpan.impl_); + int thread_id_n = get_thread_id() % kNThreadPerBlock; + int max_count = + __builtin_amdgcn_readfirstlane(N < kNPerBlock ? 0 : kNPerThread * (N / kNPerBlock)); + int n_per_block_tail_loop = + __builtin_amdgcn_readfirstlane(N - max_count * kNThreadPerBlock); - ck_tile::index_t ret = 1; + if(n_per_block_tail_loop > 0) + { + int thread_max_n = (thread_id_n + 1) * kNPerThread; + int delta = thread_max_n - n_per_block_tail_loop; + delta = clamp(thread_max_n - n_per_block_tail_loop, 0, kNPerThread); + max_count += kNPerThread - delta; + } - ck_tile::static_for<0, Lengths::size(), 1>{}( - [&](auto idx) { ret *= Lengths::template at(idx); }); - - return ret; + return max_count; } template @@ -129,42 +141,29 @@ struct Layernorm2dFwd return out_dstr_tensor; } - template - CK_TILE_DEVICE std::enable_if_t TwoPassLayernorm2dFwd(const XDataType* p_x, - const GammaDataType* p_gamma, - const BetaDataType* p_beta, - YDataType* p_y, - MeanDataType* p_mean, - InvStdDataType* p_invStd, - const ComputeDataType epsilon, - ck_tile::index_t M, - ck_tile::index_t N) const + template + CK_TILE_DEVICE std::enable_if_t + TwoPassLayernorm2dFwd(XBlockWindow& x_block_window, + GammaBlockWindow& gamma_block_window, + BetaBlockWindow& beta_block_window, + YBlockWindow& y_block_window, + MeanBlockWindow& mean_block_window, + InvStdBlockWindow& inv_std_block_window, + ComputeDataType epsilon, + ck_tile::index_t N) const { - constexpr auto I0 = number<0>{}; - constexpr auto I1 = number<1>{}; + // TODO - Optimize tail loop to reduce move_tile_window() + index_t num_n_tile_iteration = + __builtin_amdgcn_readfirstlane(integer_divide_ceil(N, kNPerBlock)); - const auto x_m_n = make_naive_tensor_view( - p_x, make_tuple(M, N), make_tuple(N, 1), number<32>{}, number<1>{}); - - const auto gamma_n = make_naive_tensor_view( - p_gamma, make_tuple(N), make_tuple(1), number<32>{}, number<1>{}); - - const auto beta_n = make_naive_tensor_view( - p_beta, make_tuple(N), make_tuple(1), number<32>{}, number<1>{}); - - const auto iM = get_block_id() * kMPerBlock; - - constexpr auto xDstr = MakeXBlockTileDistribution(); - - auto x_block_window = make_tile_window( - x_m_n, make_tuple(number{}, number{}), {iM, 0}, xDstr); - - index_t num_n_tile_iteration = __builtin_amdgcn_readfirstlane(N / kNPerBlock); - - // TODO: padding - handle max_count if N % kNPerBlock != 0 - constexpr auto NPerThread = GetNPerThread(xDstr); - ThreadWelford thread_welford{ - type_convert(NPerThread * N / kNPerBlock)}; + int welford_max_count = GetWelfordMaxCount(N); + ThreadWelford thread_welford{welford_max_count}; using XTensorType = decltype(load_tile(x_block_window)); auto mean_compute_block_tensor = @@ -190,44 +189,14 @@ struct Layernorm2dFwd auto inv_std_compute_block_tensor = InvSqrt(var_compute_block_tensor, epsilon); if constexpr(kSaveMean) - { - const auto mean_m = make_naive_tensor_view_packed( - p_mean, make_tuple(M), number<32>{}); - - auto mean_block_window = - make_tile_window(mean_m, make_tuple(number{}), {iM}); - store_tile(mean_block_window, cast_tile(mean_compute_block_tensor)); - } if constexpr(kSaveInvStd) - { - const auto inv_std_m = make_naive_tensor_view_packed( - p_invStd, make_tuple(M), number<32>{}); - - auto inv_std_block_window = - make_tile_window(inv_std_m, make_tuple(number{}), {iM}); - - store_tile(inv_std_block_window, cast_tile(inv_std_compute_block_tensor)); - } - - // TODO: Extract normalize pipeline - const auto y_m_n = make_naive_tensor_view( - p_y, make_tuple(M, N), make_tuple(N, 1), number<32>{}, number<1>{}); - - auto y_block_window = make_tile_window( - y_m_n, make_tuple(number{}, number{}), {iM, 0}); - - constexpr auto gammaDstr = MakeGammaBetaBlockTileDistribution(); - constexpr auto betaDstr = gammaDstr; - - auto gamma_block_window = - make_tile_window(gamma_n, make_tuple(number{}), {0}, gammaDstr); - - auto beta_block_window = make_tile_window( - beta_n, make_tuple(number{}, number{}), {0}, betaDstr); + store_tile(inv_std_block_window, + cast_tile(inv_std_compute_block_tensor)); // reverse read x to reuse cache - ck_tile::index_t stride_to_right_most_window = N - kNPerBlock; + ck_tile::index_t stride_to_right_most_window = + N % kNPerBlock == 0 ? N - kNPerBlock : N - N % kNPerBlock; move_tile_window(x_block_window, {0, -kNPerBlock}); move_tile_window(gamma_block_window, {stride_to_right_most_window}); @@ -274,17 +243,209 @@ struct Layernorm2dFwd } } + template + CK_TILE_DEVICE std::enable_if_t + OnePassLayernorm2dFwd(XBlockWindow& x_block_window, + GammaBlockWindow& gamma_block_window, + BetaBlockWindow& beta_block_window, + YBlockWindow& y_block_window, + MeanBlockWindow& mean_block_window, + InvStdBlockWindow& inv_std_block_window, + ComputeDataType epsilon, + ck_tile::index_t N) const + { + int welford_max_count = GetWelfordMaxCount(N); + ThreadWelford thread_welford{welford_max_count}; + + using XTensorType = decltype(load_tile(x_block_window)); + auto mean_compute_block_tensor = + thread_welford.template MakeInitialMeanVarDistributedTensor(); + auto var_compute_block_tensor = + thread_welford.template MakeInitialMeanVarDistributedTensor(); + + clear_tile(mean_compute_block_tensor); + clear_tile(var_compute_block_tensor); + + const auto x_block_tensor = load_tile(x_block_window); + thread_welford(x_block_tensor, mean_compute_block_tensor, var_compute_block_tensor); + // TODO: support cross warp Welford + WarpMergeWelford{}( + mean_compute_block_tensor, var_compute_block_tensor, thread_welford.cur_count_); + + auto inv_std_compute_block_tensor = InvSqrt(var_compute_block_tensor, epsilon); + + if constexpr(kSaveMean) + store_tile(mean_block_window, cast_tile(mean_compute_block_tensor)); + if constexpr(kSaveInvStd) + store_tile(inv_std_block_window, + cast_tile(inv_std_compute_block_tensor)); + + // normalize + const auto gamma_block_tensor = load_tile(gamma_block_window); + const auto beta_block_tensor = load_tile(beta_block_window); + + constexpr auto x_spans = decltype(x_block_tensor)::get_distributed_spans(); + + auto y_block_tensor = + make_static_distributed_tensor(x_block_tensor.get_tile_distribution()); + + sweep_tile_span(x_spans[I1], [&](auto idx1) { + constexpr auto j_idx = make_tuple(idx1); + const auto gamma = type_convert(gamma_block_tensor[j_idx]); + const auto beta = type_convert(beta_block_tensor[j_idx]); + + sweep_tile_span(x_spans[I0], [&](auto idx0) { + constexpr auto i_idx = make_tuple(idx0); + constexpr auto i_j_idx = make_tuple(idx0, idx1); + + const auto mean = mean_compute_block_tensor[i_idx]; + const auto inv_std = inv_std_compute_block_tensor[i_idx]; + + const auto x = type_convert(x_block_tensor[i_j_idx]); + auto y = (x - mean) * inv_std * gamma + beta; + + y_block_tensor(i_j_idx) = type_convert(y); + }); + }); + + store_tile(y_block_window, y_block_tensor); + } + CK_TILE_DEVICE void operator()(Kargs kargs) const { - TwoPassLayernorm2dFwd(static_cast(kargs.p_x), - static_cast(kargs.p_gamma), - static_cast(kargs.p_beta), - static_cast(kargs.p_y), - static_cast(kargs.p_mean), - static_cast(kargs.p_invStd), - static_cast(kargs.epsilon), - kargs.M, - kargs.N); + const auto x_m_n = [&]() { + const auto x_dram_naive = make_naive_tensor_view( + static_cast(kargs.p_x), + make_tuple(kargs.M, kargs.N), + make_tuple(kargs.N, 1), + number{}, + number<1>{}); + + return pad_tensor_view(x_dram_naive, + make_tuple(number{}, number{}), + sequence{}); + }(); + + const auto gamma_n = [&]() { + const auto gamma_dram_naive = make_naive_tensor_view( + static_cast(kargs.p_gamma), + make_tuple(kargs.N), + make_tuple(1), + number{}, + number<1>{}); + + return pad_tensor_view( + gamma_dram_naive, make_tuple(number{}), sequence{}); + }(); + + const auto beta_n = [&]() { + const auto gamma_dram_naive = make_naive_tensor_view( + static_cast(kargs.p_beta), + make_tuple(kargs.N), + make_tuple(1), + number{}, + number<1>{}); + + return pad_tensor_view( + gamma_dram_naive, make_tuple(number{}), sequence{}); + }(); + + const auto iM = get_block_id() * kMPerBlock; + + constexpr auto xDstr = MakeXBlockTileDistribution(); + + auto x_block_window = make_tile_window( + x_m_n, make_tuple(number{}, number{}), {iM, 0}, xDstr); + + const auto y_m_n = [&]() { + const auto y_dram_naive = make_naive_tensor_view( + static_cast(kargs.p_y), + make_tuple(kargs.M, kargs.N), + make_tuple(kargs.N, 1), + number{}, + number<1>{}); + + return pad_tensor_view(y_dram_naive, + make_tuple(number{}, number{}), + sequence{}); + }(); + + auto y_block_window = make_tile_window( + y_m_n, make_tuple(number{}, number{}), {iM, 0}); + + constexpr auto gammaDstr = MakeGammaBetaBlockTileDistribution(); + constexpr auto betaDstr = gammaDstr; + + auto gamma_block_window = + make_tile_window(gamma_n, make_tuple(number{}), {0}, gammaDstr); + + auto beta_block_window = make_tile_window( + beta_n, make_tuple(number{}, number{}), {0}, betaDstr); + + auto mean_block_window = [&]() { + if constexpr(kSaveMean) + { + const auto mean_m = [&]() { + const auto mean_dram_naive = + make_naive_tensor_view_packed( + static_cast(kargs.p_mean), + make_tuple(kargs.M), + number<1>{}); + + return pad_tensor_view( + mean_dram_naive, make_tuple(number{}), sequence{}); + }(); + + return make_tile_window(mean_m, make_tuple(number{}), {iM}); + } + else + return make_null_tile_window(make_tuple(number{})); + }(); + + auto inv_std_block_window = [&]() { + if constexpr(kSaveInvStd) + { + const auto inv_std_m = [&]() { + const auto inv_std_dram_naive = + make_naive_tensor_view_packed( + static_cast(kargs.p_invStd), + make_tuple(kargs.M), + number<1>{}); + + return pad_tensor_view( + inv_std_dram_naive, make_tuple(number{}), sequence{}); + }(); + + return make_tile_window(inv_std_m, make_tuple(number{}), {iM}); + } + else + return make_null_tile_window(make_tuple(number{})); + }(); + + if(kargs.N <= kNPerBlock) + OnePassLayernorm2dFwd(x_block_window, + gamma_block_window, + beta_block_window, + y_block_window, + mean_block_window, + inv_std_block_window, + static_cast(kargs.epsilon), + kargs.N); + else + TwoPassLayernorm2dFwd(x_block_window, + gamma_block_window, + beta_block_window, + y_block_window, + mean_block_window, + inv_std_block_window, + static_cast(kargs.epsilon), + kargs.N); } }; diff --git a/include/ck_tile/ops/layernorm2d/pipeline/block_layernorm2d_fwd_problem.hpp b/include/ck_tile/ops/layernorm2d/pipeline/block_layernorm2d_fwd_problem.hpp index 5206d36d7d..707a38f621 100644 --- a/include/ck_tile/ops/layernorm2d/pipeline/block_layernorm2d_fwd_problem.hpp +++ b/include/ck_tile/ops/layernorm2d/pipeline/block_layernorm2d_fwd_problem.hpp @@ -14,17 +14,21 @@ template + typename BlockShape_, + bool kPadM_, + bool kPadN_> struct BlockLayernorm2dFwdProblem { - using XDataType = remove_cvref_t; - using GammaDataType = remove_cvref_t; - using BetaDataType = remove_cvref_t; - using ComputeDataType = remove_cvref_t; - using YDataType = remove_cvref_t; - using MeanDataType = remove_cvref_t; - using InvStdDataType = remove_cvref_t; - using BlockShape = remove_cvref_t; + using XDataType = remove_cvref_t; + using GammaDataType = remove_cvref_t; + using BetaDataType = remove_cvref_t; + using ComputeDataType = remove_cvref_t; + using YDataType = remove_cvref_t; + using MeanDataType = remove_cvref_t; + using InvStdDataType = remove_cvref_t; + using BlockShape = remove_cvref_t; + static constexpr bool kPadM = kPadM_; + static constexpr bool kPadN = kPadN_; }; } // namespace ck_tile From cc8f466a7ecbdca058aa9b8aeb2c75c57864d2ac Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Bart=C5=82omiej=20Kocot?= Date: Mon, 7 Oct 2024 15:21:21 +0200 Subject: [PATCH 7/8] [CK_TILE] Fix conv param multiple definition (#1550) Co-authored-by: Po Yen Chen --- include/ck_tile/host/convolution_parameter.hpp | 6 ------ 1 file changed, 6 deletions(-) diff --git a/include/ck_tile/host/convolution_parameter.hpp b/include/ck_tile/host/convolution_parameter.hpp index 741a25ad73..81ea51a94f 100644 --- a/include/ck_tile/host/convolution_parameter.hpp +++ b/include/ck_tile/host/convolution_parameter.hpp @@ -13,7 +13,6 @@ namespace conv { struct ConvParam { - ConvParam(); ConvParam(ck_tile::index_t n_dim, ck_tile::index_t group_count, ck_tile::index_t n_batch, @@ -199,11 +198,6 @@ struct ConvParam } }; -ConvParam::ConvParam() - : ConvParam::ConvParam(2, 1, 128, 256, 192, {3, 3}, {71, 71}, {2, 2}, {1, 1}, {1, 1}, {1, 1}) -{ -} - CK_TILE_HOST std::string get_conv_param_parser_helper_msg() { std::string msg; From 7d8ea5f08bfea303b978c3fcb4f5b7069985b0ff Mon Sep 17 00:00:00 2001 From: Illia Silin <98187287+illsilin@users.noreply.github.com> Date: Mon, 7 Oct 2024 08:18:23 -0700 Subject: [PATCH 8/8] Fix build logic using GRU_ARCHS. (#1536) * update build logic with GPU_ARCHS * fix the GPU_ARCHS build for codegen * unset GPU_TARGETS when GPU_ARCHS are set --- CMakeLists.txt | 101 +++++++----------- Jenkinsfile | 4 +- README.md | 11 +- codegen/test/CMakeLists.txt | 3 +- example/CMakeLists.txt | 13 +-- include/ck/config.h.in | 7 -- .../gpu/CMakeLists.txt | 19 +--- profiler/src/CMakeLists.txt | 12 +-- test/CMakeLists.txt | 16 +-- 9 files changed, 64 insertions(+), 122 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index dc73b5f4d4..989995d0f5 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -98,11 +98,6 @@ if(DL_KERNELS) set(CK_ENABLE_DL_KERNELS "ON") endif() -if(INSTANCES_ONLY) - add_definitions(-DINSTANCES_ONLY) - set(CK_ENABLE_INSTANCES_ONLY "ON") -endif() - include(getopt) # CK version file to record release version as well as git commit hash @@ -127,6 +122,12 @@ rocm_setup_version(VERSION ${version}) list(APPEND CMAKE_PREFIX_PATH ${CMAKE_INSTALL_PREFIX} ${CMAKE_INSTALL_PREFIX}/llvm ${CMAKE_INSTALL_PREFIX}/hip /opt/rocm /opt/rocm/llvm /opt/rocm/hip "$ENV{ROCM_PATH}" "$ENV{HIP_PATH}") message("GPU_TARGETS= ${GPU_TARGETS}") +message("GPU_ARCHS= ${GPU_ARCHS}") +if(GPU_ARCHS) + #disable GPU_TARGETS to avoid conflicts, this needs to happen before we call hip package + unset(GPU_TARGETS CACHE) + unset(AMDGPU_TARGETS CACHE) +endif() find_package(hip) # No assumption that HIP kernels are launched with uniform block size for backward compatibility @@ -135,55 +136,38 @@ math(EXPR hip_VERSION_FLAT "(${hip_VERSION_MAJOR} * 1000 + ${hip_VERSION_MINOR}) message("hip_version_flat=${hip_VERSION_FLAT}") message("checking which targets are supported") -#This is the list of targets to be used in case GPU_TARGETS is not set on command line -#These targets will be filtered and only supported ones will be used -#Setting GPU_TARGETS on command line will override this list -if(NOT PROFILER_ONLY) - if(NOT ENABLE_ASAN_PACKAGING) - #build CK for all supported targets - if(NOT WIN32 AND ${hip_VERSION_FLAT} LESS 600300000) - # WORKAROUND: compiler does not yet fully support gfx12 targets, need to fix version above - rocm_check_target_ids(DEFAULT_GPU_TARGETS - TARGETS "gfx908;gfx90a;gfx940;gfx941;gfx942;gfx1030;gfx1100;gfx1101;gfx1102") - else() - rocm_check_target_ids(DEFAULT_GPU_TARGETS - TARGETS "gfx908;gfx90a;gfx940;gfx941;gfx942;gfx1030;gfx1100;gfx1101;gfx1102;gfx1200;gfx1201") - endif() +#In order to build just the CK library (without tests and examples) for all supported GPU targets +#use -D GPU_ARCHS="gfx908;gfx90a;gfx940;gfx941;gfx942;gfx1030;gfx1100;gfx1101;gfx1102;gfx1200;gfx1201" +#the GPU_TARGETS flag will be reset in this case in order to avoid conflicts. +# +#In order to build CK along with all tests and examples it should be OK to set GPU_TARGETS to just 1 or 2 similar architectures. +if(NOT ENABLE_ASAN_PACKAGING) + if(NOT WIN32 AND ${hip_VERSION_FLAT} LESS 600300000) + # WORKAROUND: compiler does not yet fully support gfx12 targets, need to fix version above + set(CK_GPU_TARGETS "gfx908;gfx90a;gfx940;gfx941;gfx942;gfx1030;gfx1100;gfx1101;gfx1102") else() - #build CK only for xnack-supported targets - rocm_check_target_ids(DEFAULT_GPU_TARGETS - TARGETS "gfx908:xnack+;gfx90a:xnack+;gfx940:xnack+;gfx941:xnack+;gfx942:xnack+") - set(GPU_TARGETS "${DEFAULT_GPU_TARGETS}" CACHE STRING " " FORCE) + set(CK_GPU_TARGETS "gfx908;gfx90a;gfx940;gfx941;gfx942;gfx1030;gfx1100;gfx1101;gfx1102;gfx1200;gfx1201") endif() else() - add_definitions(-DPROFILER_ONLY) - set(GPU_TARGETS "" CACHE STRING "" FORCE) + #build CK only for xnack-supported targets when using ASAN + set(CK_GPU_TARGETS "gfx908:xnack+;gfx90a:xnack+;gfx940:xnack+;gfx941:xnack+;gfx942:xnack+") +endif() + +#if user set GPU_ARCHS on the cmake command line, overwrite default target list with user's list +#otherwise, if user set GPU_TARGETS, use that set of targets +if(GPU_ARCHS) + set(CK_GPU_TARGETS ${GPU_ARCHS}) +else() if(GPU_TARGETS) - message(FATAL_ERROR "For PROFILE_ONLY build, please do not set GPU_TARGETS, use GPU_ARCH = gfx90, gfx94, gfx10, gfx11 or gfx12") + set(CK_GPU_TARGETS ${GPU_TARGETS}) endif() - if(GPU_ARCH MATCHES "gfx90") - rocm_check_target_ids(DEFAULT_GPU_TARGETS TARGETS "gfx908;gfx90a") - elseif(GPU_ARCH MATCHES "gfx94") - rocm_check_target_ids(DEFAULT_GPU_TARGETS TARGETS "gfx940;gfx941;gfx942") - elseif(GPU_ARCH MATCHES "gfx10") - rocm_check_target_ids(DEFAULT_GPU_TARGETS TARGETS "gfx1030") - elseif(GPU_ARCH MATCHES "gfx11") - rocm_check_target_ids(DEFAULT_GPU_TARGETS TARGETS "gfx1100;gfx1101;gfx1102") - elseif(GPU_ARCH MATCHES "gfx12") - rocm_check_target_ids(DEFAULT_GPU_TARGETS TARGETS "gfx1200;gfx1201") - else() - message(FATAL_ERROR "For PROFILE_ONLY build, please specify GPU_ARCH as gfx90, gfx94, gfx10, gfx11 or gfx12") - endif() - set(GPU_TARGETS "${DEFAULT_GPU_TARGETS}" CACHE STRING " " FORCE) endif() -message("Supported GPU_TARGETS= ${DEFAULT_GPU_TARGETS}") +#make sure all the targets on the list are actually supported by the current compiler +rocm_check_target_ids(SUPPORTED_GPU_TARGETS + TARGETS ${CK_GPU_TARGETS}) -if(GPU_TARGETS) - message("Building CK for the following targets: ${GPU_TARGETS}") -else() - message("Building CK for the default targets: ${DEFAULT_GPU_TARGETS}") -endif() +message("Building CK for the following targets: ${SUPPORTED_GPU_TARGETS}") if (GPU_TARGETS) if (GPU_TARGETS MATCHES "gfx9") @@ -557,8 +541,7 @@ ENDFOREACH() add_custom_target(instances DEPENDS utility;${CK_DEVICE_INSTANCES} SOURCES ${INSTANCE_FILES}) add_subdirectory(library) -if(NOT DEFINED INSTANCES_ONLY) - if(NOT DEFINED PROFILER_ONLY) +if(NOT GPU_ARCHS) rocm_package_setup_component(tests LIBRARY_NAME composablekernel PACKAGE_NAME tests # Prevent -static suffix on package name @@ -572,23 +555,15 @@ if(NOT DEFINED INSTANCES_ONLY) if(BUILD_TESTING) add_subdirectory(test) endif() - - rocm_package_setup_component(profiler - LIBRARY_NAME composablekernel - PACKAGE_NAME ckprofiler - ) - add_subdirectory(profiler) - else() - #When building PROFILER_ONLY, label the package with GPU_ARCH - rocm_package_setup_component(profiler - LIBRARY_NAME composablekernel - PACKAGE_NAME ckprofiler_${GPU_ARCH} - ) - add_subdirectory(profiler) - endif() endif() -if(NOT DEFINED PROFILER_ONLY AND (GPU_TARGETS MATCHES "gfx9" OR DEFINED INSTANCES_ONLY)) +rocm_package_setup_component(profiler + LIBRARY_NAME composablekernel + PACKAGE_NAME ckprofiler +) +add_subdirectory(profiler) + +if(GPU_TARGETS MATCHES "gfx9" OR GPU_ARCHS) add_subdirectory(codegen) endif() diff --git a/Jenkinsfile b/Jenkinsfile index 22468401dc..e61fb71e8e 100644 --- a/Jenkinsfile +++ b/Jenkinsfile @@ -1138,8 +1138,8 @@ pipeline { execute_args = """ cmake -D CMAKE_PREFIX_PATH=/opt/rocm \ -D CMAKE_CXX_COMPILER="${build_compiler()}" \ -D CMAKE_BUILD_TYPE=Release \ - -D INSTANCES_ONLY=ON \ - -DCMAKE_CXX_FLAGS=" -O3 " .. && make -j64 """ + -D GPU_ARCHS="gfx908;gfx90a;gfx940;gfx941;gfx942;gfx1030;gfx1100;gfx1101;gfx1102;gfx1200;gfx1201" \ + -D CMAKE_CXX_FLAGS=" -O3 " .. && make -j64 """ } steps{ buildHipClangJobAndReboot(setup_cmd: "", build_cmd: "", no_reboot:true, build_type: 'Release', execute_cmd: execute_args) diff --git a/README.md b/README.md index 4889914691..34ac0919ae 100644 --- a/README.md +++ b/README.md @@ -90,7 +90,12 @@ Docker images are available on [DockerHub](https://hub.docker.com/r/rocm/composa ``` If you don't set `GPU_TARGETS` on the cmake command line, CK is built for all GPU targets - supported by the current compiler (this may take a long time). + supported by the current compiler (this may take a long time). + + NOTE: If you try setting `GPU_TARGETS` to a list of architectures, the build will only work if the + architectures are similar, e.g., `gfx908;gfx90a`, or `gfx1100;gfx1101;gfx11012`. Otherwise, if you + want to build the library for a list of different architectures, + you should use the `GPU_ARCHS` build argument, for example `GPU_ARCHS=gfx908;gfx1030;gfx1100;gfx942`. 4. Build the entire CK library: @@ -137,10 +142,6 @@ crash. In such cases, you can reduce the number of threads to 32 by using `-j32` Additional cmake flags can be used to significantly speed-up the build: -* `INSTANCES_ONLY` (default is OFF) must be set to ON in order to build only the instances and library - while skipping all tests, examples, and profiler. This is useful in cases when you plan to use CK as a - dependency and don't plan to run any examples or tests. - * `DTYPES` (default is not set) can be set to any subset of "fp64;fp32;fp16;fp8;bf16;int8" to build instances of select data types only. The main default data types are fp32 and fp16; you can safely skip other data types. diff --git a/codegen/test/CMakeLists.txt b/codegen/test/CMakeLists.txt index 6dd130bc3f..1de612e49a 100644 --- a/codegen/test/CMakeLists.txt +++ b/codegen/test/CMakeLists.txt @@ -1,7 +1,8 @@ list(APPEND CMAKE_PREFIX_PATH /opt/rocm) add_subdirectory(rtc) file(GLOB TEST_SRCS CONFIGURE_DEPENDS *.cpp) -if(NOT INSTANCES_ONLY) +# do not build the tests when we build the library for various targets +if(NOT GPU_ARCHS) foreach(TEST_SRC ${TEST_SRCS}) set_source_files_properties(${TEST_SRC} PROPERTIES LANGUAGE HIP) get_filename_component(BASE_NAME ${TEST_SRC} NAME_WE) diff --git a/example/CMakeLists.txt b/example/CMakeLists.txt index f9e62a2356..ad3f7c787f 100644 --- a/example/CMakeLists.txt +++ b/example/CMakeLists.txt @@ -45,11 +45,7 @@ function(add_example_executable EXAMPLE_NAME FILE_NAME) endforeach() endif() - if(INSTANCES_ONLY) - set(EX_TARGETS ${DEFAULT_GPU_TARGETS}) - else() - set(EX_TARGETS ${GPU_TARGETS}) - endif() + set(EX_TARGETS ${SUPPORTED_GPU_TARGETS}) #Do not build any DL examples if DL_KERNELS not set foreach(source IN LISTS FILE_NAME) @@ -147,11 +143,8 @@ function(add_example_executable_no_testing EXAMPLE_NAME FILE_NAME) endforeach() endif() - if(INSTANCES_ONLY) - set(EX_TARGETS ${DEFAULT_GPU_TARGETS}) - else() - set(EX_TARGETS ${GPU_TARGETS}) - endif() + set(EX_TARGETS ${SUPPORTED_GPU_TARGETS}) + #Do not build any DL examples if DL_KERNELS not set foreach(source IN LISTS FILE_NAME) if(NOT DEFINED DL_KERNELS AND source MATCHES "_dl") diff --git a/include/ck/config.h.in b/include/ck/config.h.in index eb9049b599..0f0b7bd607 100644 --- a/include/ck/config.h.in +++ b/include/ck/config.h.in @@ -97,13 +97,6 @@ #cmakedefine CK_ENABLE_DL_KERNELS @CK_ENABLE_DL_KERNELS@ #endif -// -// Instances supports in the current CK build -// -#ifndef CK_ENABLE_INSTANCES_ONLY -#cmakedefine CK_ENABLE_INSTANCES_ONLY @CK_ENABLE_INSTANCES_ONLY@ -#endif - // // CK kernels which support XDL (MI series) // diff --git a/library/src/tensor_operation_instance/gpu/CMakeLists.txt b/library/src/tensor_operation_instance/gpu/CMakeLists.txt index bc66fe0bed..f82176ffc6 100644 --- a/library/src/tensor_operation_instance/gpu/CMakeLists.txt +++ b/library/src/tensor_operation_instance/gpu/CMakeLists.txt @@ -37,11 +37,7 @@ function(add_instance_library INSTANCE_NAME) endforeach() endif() - if(INSTANCES_ONLY) - set(INST_TARGETS ${DEFAULT_GPU_TARGETS}) - else() - set(INST_TARGETS ${GPU_TARGETS}) - endif() + set(INST_TARGETS ${SUPPORTED_GPU_TARGETS}) # Do not build DL instances if DL_KERNELS macro is not set foreach(source IN LISTS ARGN) @@ -75,11 +71,7 @@ function(add_instance_library INSTANCE_NAME) if(ARGN) set(INST_OBJ) foreach(source IN LISTS ARGN) - if(INSTANCES_ONLY) - set(INST_TARGETS ${DEFAULT_GPU_TARGETS}) - else() - set(INST_TARGETS ${GPU_TARGETS}) - endif() + set(INST_TARGETS ${SUPPORTED_GPU_TARGETS}) if(source MATCHES "_xdl") list(REMOVE_ITEM INST_TARGETS gfx900 gfx906 gfx1030 gfx1100 gfx1101 gfx1102 gfx1103 gfx1200 gfx1201) elseif(ARGN MATCHES "_wmma") @@ -191,12 +183,7 @@ FOREACH(subdir_path ${dir_list}) set(add_inst 1) endif() - if(INSTANCES_ONLY) - set(INST_TARGETS ${DEFAULT_GPU_TARGETS}) - else() - set(INST_TARGETS ${GPU_TARGETS}) - endif() - + set(INST_TARGETS ${SUPPORTED_GPU_TARGETS}) if(("${cmake_instance}" MATCHES "quantization") AND (DEFINED DTYPES) AND (NOT DTYPES MATCHES "int8")) message("quantization instances will not be built!") diff --git a/profiler/src/CMakeLists.txt b/profiler/src/CMakeLists.txt index e9528baeb6..7d4df3cf9b 100644 --- a/profiler/src/CMakeLists.txt +++ b/profiler/src/CMakeLists.txt @@ -24,7 +24,7 @@ set(PROFILER_SOURCES profile_permute_scale.cpp ) -if(GPU_TARGETS MATCHES "gfx9") +if(SUPPORTED_GPU_TARGETS MATCHES "gfx9") if(DTYPES MATCHES "fp32" OR DTYPES MATCHES "fp64" OR NOT DEFINED DTYPES) list(APPEND PROFILER_SOURCES profile_contraction_bilinear.cpp) list(APPEND PROFILER_SOURCES profile_contraction_scale.cpp) @@ -49,7 +49,7 @@ if(GPU_TARGETS MATCHES "gfx9") list(APPEND PROFILER_SOURCES profile_grouped_gemm_multiply_tile_loop.cpp) endif() list(APPEND PROFILER_SOURCES profile_gemm_multiply_add.cpp) - if(GPU_TARGETS MATCHES "gfx94") + if(SUPPORTED_GPU_TARGETS MATCHES "gfx94") list(APPEND PROFILER_SOURCES profile_gemm_multiply_multiply.cpp) list(APPEND PROFILER_SOURCES profile_gemm_ab_scale.cpp) endif() @@ -69,7 +69,7 @@ if(GPU_TARGETS MATCHES "gfx9") endif() -if(GPU_TARGETS MATCHES "gfx11" OR GPU_TARGETS MATCHES "gfx12" OR GPU_TARGETS MATCHES "gfx9") +if(SUPPORTED_GPU_TARGETS MATCHES "gfx11" OR SUPPORTED_GPU_TARGETS MATCHES "gfx12" OR SUPPORTED_GPU_TARGETS MATCHES "gfx9") if(DTYPES MATCHES "fp16" OR NOT DEFINED DTYPES) list(APPEND PROFILER_SOURCES profile_gemm_bilinear.cpp) endif() @@ -111,7 +111,7 @@ target_link_libraries(${PROFILER_EXECUTABLE} PRIVATE device_column_to_image_inst target_link_libraries(${PROFILER_EXECUTABLE} PRIVATE device_transpose_instance) target_link_libraries(${PROFILER_EXECUTABLE} PRIVATE device_permute_scale_instance) -if(GPU_TARGETS MATCHES "gfx9") +if(SUPPORTED_GPU_TARGETS MATCHES "gfx9") if(DTYPES MATCHES "fp32" OR DTYPES MATCHES "fp64" OR NOT DEFINED DTYPES) target_link_libraries(${PROFILER_EXECUTABLE} PRIVATE device_contraction_bilinear_instance) target_link_libraries(${PROFILER_EXECUTABLE} PRIVATE device_contraction_scale_instance) @@ -135,7 +135,7 @@ if(GPU_TARGETS MATCHES "gfx9") target_link_libraries(${PROFILER_EXECUTABLE} PRIVATE device_batched_gemm_instance) target_link_libraries(${PROFILER_EXECUTABLE} PRIVATE device_batched_gemm_reduce_instance) target_link_libraries(${PROFILER_EXECUTABLE} PRIVATE device_gemm_multiply_add_instance) - if(GPU_TARGETS MATCHES "gfx94") + if(SUPPORTED_GPU_TARGETS MATCHES "gfx94") target_link_libraries(${PROFILER_EXECUTABLE} PRIVATE device_gemm_multiply_multiply_instance) target_link_libraries(${PROFILER_EXECUTABLE} PRIVATE device_gemm_ab_scale_instance) endif() @@ -159,7 +159,7 @@ if(GPU_TARGETS MATCHES "gfx9") target_link_libraries(${PROFILER_EXECUTABLE} PRIVATE device_grouped_conv3d_fwd_convinvscale_instance) endif() -if(GPU_TARGETS MATCHES "gfx9" OR GPU_TARGETS MATCHES "gfx11" OR GPU_TARGETS MATCHES "gfx12") +if(SUPPORTED_GPU_TARGETS MATCHES "gfx9" OR SUPPORTED_GPU_TARGETS MATCHES "gfx11" OR SUPPORTED_GPU_TARGETS MATCHES "gfx12") if(DTYPES MATCHES "fp16" OR NOT DEFINED DTYPES) target_link_libraries(${PROFILER_EXECUTABLE} PRIVATE device_gemm_bilinear_instance) endif() diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt index e61d937f08..b836dd687e 100644 --- a/test/CMakeLists.txt +++ b/test/CMakeLists.txt @@ -41,11 +41,7 @@ function(add_test_executable TEST_NAME) endforeach() endif() - if(INSTANCES_ONLY) - set(TEST_TARGETS ${DEFAULT_GPU_TARGETS}) - else() - set(TEST_TARGETS ${GPU_TARGETS}) - endif() + set(TEST_TARGETS ${SUPPORTED_GPU_TARGETS}) foreach(source IN LISTS ARGN) if(NOT DEFINED DL_KERNELS AND source MATCHES "_dl") @@ -122,11 +118,7 @@ function(add_gtest_executable TEST_NAME) endforeach() endif() - if(INSTANCES_ONLY) - set(TEST_TARGETS ${DEFAULT_GPU_TARGETS}) - else() - set(TEST_TARGETS ${GPU_TARGETS}) - endif() + set(TEST_TARGETS ${SUPPORTED_GPU_TARGETS}) foreach(source IN LISTS ARGN) if(NOT DEFINED DL_KERNELS AND source MATCHES "_dl") @@ -211,10 +203,10 @@ add_subdirectory(conv_tensor_rearrange) add_subdirectory(transpose) add_subdirectory(permute_scale) add_subdirectory(wrapper) -if(GPU_TARGETS MATCHES "gfx11") +if(SUPPORTED_GPU_TARGETS MATCHES "gfx11") add_subdirectory(wmma_op) endif() -if(GPU_TARGETS MATCHES "gfx942" AND CK_HIP_VERSION_MAJOR GREATER_EQUAL 6 AND CK_HIP_VERSION_MINOR GREATER_EQUAL 2) # smfmac needs ROCm6.2 +if(SUPPORTED_GPU_TARGETS MATCHES "gfx942" AND CK_HIP_VERSION_MAJOR GREATER_EQUAL 6 AND CK_HIP_VERSION_MINOR GREATER_EQUAL 2) # smfmac needs ROCm6.2 add_subdirectory(smfmac_op) endif() add_subdirectory(position_embedding)