From d12b750bb089fe31ab46f39fcdb2c66e8151557b Mon Sep 17 00:00:00 2001 From: "Ding, Yi" Date: Wed, 7 May 2025 08:17:07 +0000 Subject: [PATCH] wip --- CMakeLists.txt | 2 +- cmake/EnableCompilerWarnings.cmake | 4 +- .../67_gemm_microscaling/gemm_mx_common.hpp | 4 +- example/67_gemm_microscaling/gemm_mx_fp8.cpp | 6 +- include/ck/config.h.in | 7 ++ ...blockwise_gemm_mx_pipeline_xdlops_base.hpp | 2 +- ...kwise_gemm_pipeline_xdlops_mx_selector.hpp | 5 +- .../blockwise_gemm_pipeline_xdlops_v1_mx.hpp | 74 ++++++++++--------- .../element/unary_element_wise_operation.hpp | 7 ++ .../grid/gridwise_gemm_xdl_cshuffle_v3_mx.hpp | 16 ++-- .../threadwise_tensor_slice_transfer_v3r1.hpp | 8 +- .../tensor_operation/gpu/warp/xdlops_gemm.hpp | 13 +++- include/ck/utility/amd_buffer_addressing.hpp | 6 +- include/ck/utility/amd_xdlops.hpp | 11 +-- include/ck/utility/data_type.hpp | 45 ++++++++++- include/ck/utility/dtype_vector.hpp | 6 ++ 16 files changed, 154 insertions(+), 62 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 4e12462a41..6d0bfe8b36 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -542,7 +542,7 @@ include_directories(BEFORE SET(BUILD_DEV ON CACHE BOOL "BUILD_DEV") if(BUILD_DEV) - add_compile_options(-Werror) + # add_compile_options(-Werror) add_compile_options(-Weverything) endif() message("CMAKE_CXX_FLAGS: ${CMAKE_CXX_FLAGS}") diff --git a/cmake/EnableCompilerWarnings.cmake b/cmake/EnableCompilerWarnings.cmake index fb2b38d688..c60d9ea0c4 100644 --- a/cmake/EnableCompilerWarnings.cmake +++ b/cmake/EnableCompilerWarnings.cmake @@ -66,7 +66,7 @@ else() -Wunreachable-code -Wunused -Wno-reserved-identifier - -Werror + # -Werror -Wno-option-ignored -Wsign-compare -Wno-extra-semi-stmt @@ -108,7 +108,7 @@ else() endif() list(APPEND CMAKE_COMPILER_WARNINGS -Wno-missing-field-initializers - -Wno-deprecated-declarations + # -Wno-deprecated-declarations ) endif() add_definitions(${CMAKE_COMPILER_WARNINGS}) diff --git a/example/67_gemm_microscaling/gemm_mx_common.hpp b/example/67_gemm_microscaling/gemm_mx_common.hpp index 99ed2a23b9..7e28919a2a 100644 --- a/example/67_gemm_microscaling/gemm_mx_common.hpp +++ b/example/67_gemm_microscaling/gemm_mx_common.hpp @@ -202,8 +202,8 @@ bool run_mx_gemm(const ProblemSizeSplitK& problem_size, const ExecutionConfig& c case 0: // Initializations for development and debugging ck::utils::FillConstant{ck::type_convert(1.0f)}(a_m_k); ck::utils::FillConstant{ck::type_convert(2.0f)}(a_m_k_scale); - ck::utils::FillConstant{ck::type_convert(0.5f)}(b_k_n); - ck::utils::FillConstant{ck::type_convert(1.0f)}(b_k_n_scale); + ck::utils::FillConstant{ck::type_convert(1.f)}(b_k_n); + ck::utils::FillConstant{ck::type_convert(2.0f)}(b_k_n_scale); if(config.verbosity > 0) { std::cout << "Init A = {1}" << std::endl; diff --git a/example/67_gemm_microscaling/gemm_mx_fp8.cpp b/example/67_gemm_microscaling/gemm_mx_fp8.cpp index 9fc5666197..e40f9d0391 100644 --- a/example/67_gemm_microscaling/gemm_mx_fp8.cpp +++ b/example/67_gemm_microscaling/gemm_mx_fp8.cpp @@ -3,8 +3,10 @@ #include "gemm_mx_common.hpp" -using ADataType = ck::f8_t; -using BDataType = ck::f8_t; +using ADataType = ck::f4x2_pk_t; +using BDataType = ck::f4x2_pk_t; +// using ADataType = ck::f4_t; +// using BDataType = ck::f4_t; using XDataType = ck::e8m0_bexp_t; diff --git a/include/ck/config.h.in b/include/ck/config.h.in index 306a6c2ff1..b05eddf03f 100644 --- a/include/ck/config.h.in +++ b/include/ck/config.h.in @@ -26,6 +26,13 @@ #ifndef CK_CONFIG_H_IN #define CK_CONFIG_H_IN +template +[[deprecated("Help function to print value")]] +inline constexpr void CK_TILE_PRINT() {} +template +[[deprecated("Help function to print value")]] +inline constexpr void CK_TILE_PRINT() {} + // clang-format off // // DataType supports in the current CK build diff --git a/include/ck/tensor_operation/gpu/block/blockwise_gemm_mx_pipeline_xdlops_base.hpp b/include/ck/tensor_operation/gpu/block/blockwise_gemm_mx_pipeline_xdlops_base.hpp index ebe075b55d..77c13ef2cc 100644 --- a/include/ck/tensor_operation/gpu/block/blockwise_gemm_mx_pipeline_xdlops_base.hpp +++ b/include/ck/tensor_operation/gpu/block/blockwise_gemm_mx_pipeline_xdlops_base.hpp @@ -60,7 +60,7 @@ struct BlockwiseGemmXdlops_mx_pipeline_base //> e.g. [k0,...,k15,k64,...,k79] or [k0,...,k15,k32,...,k47] static constexpr index_t KThreadChunk = 16; - static constexpr index_t KPerThread = KPerBlock / xdlops_gemm.K0PerXdlops; + static constexpr index_t KPerThread = KPerBlock / 2 / xdlops_gemm.K0PerXdlops; static constexpr index_t KRepeat = KPerThread / KPack; static constexpr index_t KPerInnerLoop = KPack; diff --git a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_mx_selector.hpp b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_mx_selector.hpp index c1433659d6..471d7a581f 100644 --- a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_mx_selector.hpp +++ b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_mx_selector.hpp @@ -13,8 +13,9 @@ namespace ck { template static constexpr bool is_scale_mfma_data_type() { - return is_same_v || is_same_v || is_same_v || - is_same_v || is_same_v; + using U = element_type_t; + return is_same_v || is_same_v || is_same_v || + is_same_v || is_same_v; } /** diff --git a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_v1_mx.hpp b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_v1_mx.hpp index 9acf401410..14c2a11347 100644 --- a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_v1_mx.hpp +++ b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_v1_mx.hpp @@ -156,7 +156,8 @@ struct BlockwiseGemmXdlops_pipeline_v1_mx How many mx-vectors in each row/col is processed in one call to xdlops_gemm.Run() - static constexpr auto ScalesPerXdlopsRun = (KPack * xdlops_gemm.K0PerXdlops) / ScaleBlockSize; + static constexpr auto ScalesPerXdlopsRun = + (KPack * 2 * xdlops_gemm.K0PerXdlops) / ScaleBlockSize; //> How many scales a thread must read to accommodate one call to xdlops_gemm.Run() static constexpr auto ScalesPerXdlopsRunPerThread = @@ -338,40 +339,42 @@ struct BlockwiseGemmXdlops_pipeline_v1_mx{}([&](auto m0) { - static_for<0, xdlops_gemm.K1PerXdlops / KThreadChunk, 1>{}([&](auto chunk) { - constexpr auto a_k_step_chunk = - k_step + - chunk * KThreadChunk * xdlops_gemm.mfma_instr.num_input_blks; - a_thread_copy_.Run( - a_block_desc_m0_m1_m2_k, - make_tuple(m0, I0, I0, Number{}), - a_block_buf, - a_thread_desc_, - make_tuple(m0, I0, k, Number{}), - a_thread_buf); - }); + static_for<0, xdlops_gemm.K1PerXdlops / 2 / KThreadChunk, 1>{}( + [&](auto chunk) { + constexpr auto a_k_step_chunk = + k_step + + chunk * KThreadChunk * xdlops_gemm.mfma_instr.num_input_blks; + // a_thread_copy_.Run( + // a_block_desc_m0_m1_m2_k, + // make_tuple(m0, I0, I0, Number{}), + // a_block_buf, + // a_thread_desc_, + // make_tuple(m0, I0, k, Number{}), + // a_thread_buf); + }); }); static_for<0, NRepeat, 1>{}([&](auto n0) { // read block data in chunks to assemble correct thread vectors - static_for<0, xdlops_gemm.K1PerXdlops / KThreadChunk, 1>{}([&](auto chunk) { - constexpr auto b_k_step_chunk = - k_step + - chunk * KThreadChunk * xdlops_gemm.mfma_instr.num_input_blks; - b_thread_copy_.Run( - b_block_desc_n0_n1_n2_k, - make_tuple(n0, I0, I0, Number{}), - b_block_buf, - b_thread_desc_, - make_tuple(n0, I0, k, Number{}), - b_thread_buf); - }); + static_for<0, xdlops_gemm.K1PerXdlops / 2 / KThreadChunk, 1>{}( + [&](auto chunk) { + constexpr auto b_k_step_chunk = + k_step + + chunk * KThreadChunk * xdlops_gemm.mfma_instr.num_input_blks; + b_thread_copy_.Run( + b_block_desc_n0_n1_n2_k, + make_tuple(n0, I0, I0, Number{}), + b_block_buf, + b_thread_desc_, + make_tuple(n0, I0, k, Number{}), + b_thread_buf); + }); }); }); static_for<0, MRepeat, 1>{}([&](auto m0) { static_for<0, NRepeat, 1>{}([&](auto n0) { static_for<0, KRepeat, 1>{}([&](auto k0) { - vector_type a_thread_vec; + vector_type a_thread_vec; // = vec: pk_i4_t, 32 vector_type b_thread_vec; static_for<0, KPack, 1>{}([&](auto ik) { @@ -403,11 +406,16 @@ struct BlockwiseGemmXdlops_pipeline_v1_mx()(s) = b_scale_thread_buf[Number{}]; }); - + CK_TILE_PRINT(); + CK_TILE_PRINT(); using mfma_input_type_a = - typename vector_type::type; + typename vector_type::type; + // mfma input type = pk_f4_t, 32 + CK_TILE_PRINT(); using mfma_input_type_b = - typename vector_type::type; + typename vector_type::type; constexpr index_t c_offset = c_thread_desc_.CalculateOffset(make_tuple(m0, n0, 0)); @@ -505,7 +513,7 @@ struct BlockwiseGemmXdlops_pipeline_v1_mx{}([&](auto m0) { // read block data in chunks to assemble correct thread - static_for<0, xdlops_gemm.K1PerXdlops / KThreadChunk, 1>{}([&](auto chunk) { + static_for<0, xdlops_gemm.K1PerXdlops / 2 / KThreadChunk, 1>{}([&](auto chunk) { constexpr auto a_k_step_chunk = k_step + chunk * KThreadChunk * xdlops_gemm.mfma_instr.num_input_blks; a_thread_copy_.Run(a_block_desc_m0_m1_m2_k, @@ -518,7 +526,7 @@ struct BlockwiseGemmXdlops_pipeline_v1_mx{}([&](auto n0) { // read block data in chunks to assemble correct thread - static_for<0, xdlops_gemm.K1PerXdlops / KThreadChunk, 1>{}([&](auto chunk) { + static_for<0, xdlops_gemm.K1PerXdlops / 2 / KThreadChunk, 1>{}([&](auto chunk) { constexpr auto b_k_step_chunk = k_step + chunk * KThreadChunk * xdlops_gemm.mfma_instr.num_input_blks; b_thread_copy_.Run(b_block_desc_n0_n1_n2_k, @@ -564,9 +572,9 @@ struct BlockwiseGemmXdlops_pipeline_v1_mx::type; + typename vector_type::type; using mfma_input_type_b = - typename vector_type::type; + typename vector_type::type; constexpr index_t c_offset = c_thread_desc_.CalculateOffset(make_tuple(m0, n0, 0)); diff --git a/include/ck/tensor_operation/gpu/element/unary_element_wise_operation.hpp b/include/ck/tensor_operation/gpu/element/unary_element_wise_operation.hpp index 672998d811..047ff3bd06 100644 --- a/include/ck/tensor_operation/gpu/element/unary_element_wise_operation.hpp +++ b/include/ck/tensor_operation/gpu/element/unary_element_wise_operation.hpp @@ -315,6 +315,13 @@ struct PassThrough y = x; } + template <> + __host__ __device__ void operator()(f4x2_pk_t& y, + const f4x2_pk_t& x) const + { + y = x; + } + template <> __host__ __device__ void operator()(double& y, const double& x) const { diff --git a/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_v3_mx.hpp b/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_v3_mx.hpp index 1154fa2aa3..b6b6f92135 100644 --- a/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_v3_mx.hpp +++ b/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_v3_mx.hpp @@ -167,26 +167,29 @@ struct GridwiseGemmMX_xdl_cshuffle_v3 // // Should be a multiple of k_per_blk. // TODO: Move this to blockwise pipeline base - static constexpr index_t KPack = - math::max(lcm_AK1_BK1, + static constexpr index_t KPack = // = num of pk_f4 + math::max(lcm_AK1_BK1, // num of pk_f4 MfmaSelector::selected_mfma.k_per_blk); + is_scale_mfma>::selected_mfma.k_per_blk / + 2); // num of f4 using ThisThreadBlock = ThisThreadBlock; static constexpr index_t APackedSize = []() { - if constexpr(is_same_v, pk_i4_t>) + if constexpr(is_same_v, pk_i4_t> || + is_same_v, f4x2_pk_t>) return 2; else return 1; }(); static constexpr index_t BPackedSize = []() { - if constexpr(is_same_v, pk_i4_t>) + if constexpr(is_same_v, pk_i4_t> || + is_same_v, f4x2_pk_t>) return 2; else return 1; @@ -363,6 +366,9 @@ struct GridwiseGemmMX_xdl_cshuffle_v3 static_assert(!(is_same_v, pk_i4_t> && GemmSpec != GemmSpecialization::Default), "pk_i4_t does not support padding"); + static_assert(!(is_same_v, f4x2_pk_t> && + GemmSpec != GemmSpecialization::Default), + "f4x2_pk_t does not support padding"); if constexpr(GemmSpec == GemmSpecialization::NKPadding || GemmSpec == GemmSpecialization::MNKPadding) diff --git a/include/ck/tensor_operation/gpu/thread/threadwise_tensor_slice_transfer_v3r1.hpp b/include/ck/tensor_operation/gpu/thread/threadwise_tensor_slice_transfer_v3r1.hpp index 7ccea96dda..9305812352 100644 --- a/include/ck/tensor_operation/gpu/thread/threadwise_tensor_slice_transfer_v3r1.hpp +++ b/include/ck/tensor_operation/gpu/thread/threadwise_tensor_slice_transfer_v3r1.hpp @@ -69,7 +69,8 @@ struct ThreadwiseTensorSliceTransfer_v3r1 static constexpr auto I16 = Number<16>{}; static constexpr index_t PackedSize = []() { - if constexpr(is_same_v, pk_i4_t>) + if constexpr(is_same_v, pk_i4_t> || + is_same_v, f4x2_pk_t>) return 2; else return 1; @@ -90,7 +91,8 @@ struct ThreadwiseTensorSliceTransfer_v3r1 src_element_op_(src_element_op), dst_element_op_(dst_element_op) { - if constexpr(is_same_v, pk_i4_t>) + if constexpr(is_same_v, pk_i4_t> || + is_same_v, f4x2_pk_t>) { static_assert(is_same_v, remove_cvref_t>, "SrcData != DstData"); @@ -444,6 +446,8 @@ struct ThreadwiseTensorSliceTransfer_v3r1 { static_assert(!is_same_v, pk_i4_t>, "in-register transpose is not supported for pk_i4_t"); + static_assert(!is_same_v, f4x2_pk_t>, + "in-register transpose is not supported for f4x2_pk_t"); // each transpose does // DstScalarPerVector # of src vectors in src_thread_scratch_ // SrcScalarPerVector # of dst vectors in dst_thread_scratch_ diff --git a/include/ck/tensor_operation/gpu/warp/xdlops_gemm.hpp b/include/ck/tensor_operation/gpu/warp/xdlops_gemm.hpp index 06268f3cfb..15b2661c8d 100644 --- a/include/ck/tensor_operation/gpu/warp/xdlops_gemm.hpp +++ b/include/ck/tensor_operation/gpu/warp/xdlops_gemm.hpp @@ -858,6 +858,7 @@ struct mfma_type const ScaleB& scale_b, FloatC& reg_c) const { + // CK_TILE_PRINT(); static_assert(scalar_type::vector_size == 1, "Expect single scale at this point."); static_assert(scalar_type::vector_size == 1, "Expect single scale at this point."); @@ -1134,6 +1135,11 @@ struct MfmaSelector { return MfmaInstr::mfma_scale_f32_32x32x64f8f6f4; } + template <> + constexpr auto GetMfma() + { + return MfmaInstr::mfma_scale_f32_32x32x64f8f6f4; + } template <> constexpr auto GetMfma() @@ -1201,10 +1207,10 @@ struct MfmaSelector return MfmaInstr::mfma_f32_16x16x32bf8f8; } - static constexpr auto selected_mfma = mfma_type, MPerXdlops, NPerXdlops, - additional_type, + element_type_t, is_single_rate_mfma, is_scale_mfma>()>{}; @@ -1286,7 +1292,8 @@ struct XdlopsGemm MPerXdlops == 64, "Only support GemmMPerXdlops == 4, 8, 16, 32 or 64 for xdlops"); - static_assert(KPack % mfma_instr.k_per_blk == 0, "KPack should be a multiple of k_per_blk"); + static_assert(KPack * 2 % mfma_instr.k_per_blk == 0, + "KPack should be a multiple of k_per_blk"); } // XDL output supporting C = A * B diff --git a/include/ck/utility/amd_buffer_addressing.hpp b/include/ck/utility/amd_buffer_addressing.hpp index 62e3220b5a..30f734faa5 100644 --- a/include/ck/utility/amd_buffer_addressing.hpp +++ b/include/ck/utility/amd_buffer_addressing.hpp @@ -430,7 +430,9 @@ __device__ typename vector_type::type amd_buffer_load_impl(int32x4_t src_w (is_same::value && (N == 1 || N == 2 || N == 4 || N == 8 || N == 16)) || (is_same::value && (N == 1 || N == 2 || N == 4 || N == 8 || N == 16)) || (is_same::value && (N == 1 || N == 2 || N == 4 || N == 8 || N == 16)) || - (is_same::value && (N == 1 || N == 2 || N == 4 || N == 8 || N == 16)), + (is_same::value && (N == 1 || N == 2 || N == 4 || N == 8 || N == 16)) || + (is_same::value && + (N == 1 || N == 2 || N == 4 || N == 8 || N == 16)), "wrong! not implemented"); using r_t = typename vector_type::type; @@ -845,7 +847,7 @@ amd_buffer_load_invalid_element_return_zero(const T* p_src_wave, src_wave_buffer_resource, src_addr_shift + src_thread_addr_offset, 0); #else - + CK_TILE_PRINT(); vector_t tmp{amd_buffer_load_impl( src_wave_buffer_resource, src_thread_addr_offset, 0)}; return src_thread_element_valid ? tmp : vector_t(0); diff --git a/include/ck/utility/amd_xdlops.hpp b/include/ck/utility/amd_xdlops.hpp index 66c4958e1d..85be91a71e 100644 --- a/include/ck/utility/amd_xdlops.hpp +++ b/include/ck/utility/amd_xdlops.hpp @@ -749,11 +749,12 @@ struct intrin_mfma_scale_f32_16x16x128f8f6f4<16, 16> } template - __device__ static void Run(const f4x32_t& reg_a, - const int32_t scale_a, - const f4x32_t& reg_b, - const int32_t scale_b, - FloatC& reg_c) + __device__ static void + Run(const f4x32_t& reg_a, // misalignment between pk_f4_t, 32 and f4_t, 32 + const int32_t scale_a, + const f4x32_t& reg_b, + const int32_t scale_b, + FloatC& reg_c) { #if defined(__gfx950__) int32x4_t arg_a = bit_cast(reg_a); diff --git a/include/ck/utility/data_type.hpp b/include/ck/utility/data_type.hpp index a6106bb146..20054e5ac0 100644 --- a/include/ck/utility/data_type.hpp +++ b/include/ck/utility/data_type.hpp @@ -36,8 +36,8 @@ struct f4x2_pk_t { using type = uint8_t; type data; - __host__ __device__ f4x2_pk_t() : data{type{}} {} - __host__ __device__ f4x2_pk_t(type init) : data{init} {} + __host__ __device__ constexpr f4x2_pk_t() : data{type{}} {} + __host__ __device__ constexpr f4x2_pk_t(const type init) : data{init} {} template __host__ __device__ inline type unpack(Number) const @@ -346,6 +346,18 @@ inline constexpr bool is_native_type() is_same::value || is_same::value; } +template +struct is_f8f6f4 +{ + static constexpr bool value = + 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; +}; +template +inline constexpr bool is_f8f6f4_v = is_f8f6f4::value; + // scalar_type template struct scalar_type; @@ -484,6 +496,35 @@ struct scalar_type static constexpr index_t vector_size = 1; }; +template +struct element_type +{ + private: + static constexpr auto get_element_type() + { + using U = remove_cvref_t; + if constexpr(std::is_same_v) + return int4_t{}; + else if constexpr(std::is_same_v) + return f4_t{}; + else if constexpr(std::is_same_v) + return f6_t{}; + else if constexpr(std::is_same_v) + return bf6_t{}; + else if constexpr(std::is_same_v) + return f6_t{}; + else if constexpr(std::is_same_v) + return bf6_t{}; + else + return T{}; + } + + public: + using type = decltype(get_element_type()); +}; +template +using element_type_t = typename element_type::type; + #if defined(_WIN32) using int64_t = long long; #else diff --git a/include/ck/utility/dtype_vector.hpp b/include/ck/utility/dtype_vector.hpp index 9c40d923d3..66999f87b6 100644 --- a/include/ck/utility/dtype_vector.hpp +++ b/include/ck/utility/dtype_vector.hpp @@ -1248,6 +1248,12 @@ struct nnvb_data_t_selector using type = pk_i4_t::type; }; +template <> +struct nnvb_data_t_selector +{ + using type = f4x2_pk_t::type; +}; + template struct non_native_vector_base< T,