diff --git a/CMakeLists.txt b/CMakeLists.txt index acae1f5ece..6017f4ee38 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -621,7 +621,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(STATUS "CMAKE_CXX_FLAGS: ${CMAKE_CXX_FLAGS}") diff --git a/example/ck_tile/18_flatmm/mxgemm/mx_flatmm.cpp b/example/ck_tile/18_flatmm/mxgemm/mx_flatmm.cpp index d6c84f3064..2c12119651 100644 --- a/example/ck_tile/18_flatmm/mxgemm/mx_flatmm.cpp +++ b/example/ck_tile/18_flatmm/mxgemm/mx_flatmm.cpp @@ -179,10 +179,11 @@ auto preShuffleWeight(ck_tile::HostTensor& src) const int K = src_lengths[0]; const int N = src_lengths[1]; constexpr int packed_size = ck_tile::numeric_traits::PackedSize; - int KPack = 16 * packed_size; // fp4:32 or fp8:16 - int NLane = N_Warp_Tile; - int KLane = 64 / NLane; - int K0 = K / (KLane * KPack); + int KPack = + std::is_same_v ? 32 : 16 * packed_size; // fp4/fp6:32 or fp8:16 + int NLane = N_Warp_Tile; + int KLane = 64 / NLane; + int K0 = K / (KLane * KPack); ck_tile::HostTensor shuffled(ck_tile::HostTensorDescriptor({N * K}, {1})); @@ -204,6 +205,7 @@ auto preShuffleWeight(ck_tile::HostTensor& src) int outputIndex = n0 * KPack * NLane * KLane * K0 + k0 * KPack * NLane * KLane + k1 * KPack * NLane + n1 * KPack + k2; + std::cout << k << " " << n << " " << outputIndex << std::endl; shuffled(outputIndex) = src(k, n); } } @@ -343,6 +345,48 @@ int run_mx_flatmm_example(int argc, char* argv[]) int main(int argc, char* argv[]) { + + { + // using BDataType = ck_tile::pk_fp4_t; + using BDataType = ck_tile::f6x16_pk_t; + + ck_tile::index_t stride_B = 0; + ck_tile::index_t N = 32; + ck_tile::index_t K = 256; + stride_B = ck_tile::get_default_stride(K, N, stride_B, ck_tile::bool_constant{}); + // is_row_major + ck_tile::HostTensor b_origin_host( + ck_tile::host_tensor_descriptor(K, N, stride_B, ck_tile::bool_constant{})); + std::cout << b_origin_host.get_element_space_size_in_bytes() << std::endl; + auto try_pack_unpack = [&] { + int pack_k = 16; + for(int n = 0; n < N; n++) + { + for(int k = 0; k < K; k += pack_k) + { + for(int k_ = 0; k_ < pack_k; k_++) + { + int value = n * K + k + k_; + b_origin_host(n, k).pack(value, k_); + } + } + } + for(int n = 0; n < N; n++) + { + for(int k = 0; k < K; k += pack_k) + { + for(int k_ = 0; k_ < pack_k; k_++) + { + std::cout << b_origin_host(n, k).unpack(k_) << std::endl; + } + } + } + }; + try_pack_unpack(); + auto shuf_b = preShuffleWeight<16>(b_origin_host); + return 0; + } + auto [result, arg_parser] = create_args(argc, argv); if(!result) return EXIT_FAILURE; diff --git a/include/ck/utility/amd_xdlops.hpp b/include/ck/utility/amd_xdlops.hpp index f8b1736801..ef129f9bc0 100644 --- a/include/ck/utility/amd_xdlops.hpp +++ b/include/ck/utility/amd_xdlops.hpp @@ -639,7 +639,7 @@ struct intrin_mfma_f32_32x32x64f8f6f4<32, 32> arg_type{arg_a[0], arg_a[1], arg_a[2], arg_a[3], arg_a[4], arg_a[5], 0, 0}, arg_type{arg_b[0], arg_b[1], arg_b[2], arg_b[3], arg_b[4], arg_b[5], 0, 0}, reg_c.template AsType()[Number<0>{}], - 2, // cbsz {0 FP8 E4M3; 1 FP8 E5M2; 2 FP6 E2M3; 3 FP6 E3M2; 4 FP4 E2M1} + 2, // cbsz {0 FP8 E4M3; 1 FP8 E5M2; 2 FP6 ; 3 FP6 E3M2; 4 FP4 E2M1} 2, // blgp 0, // OPSEL 0, diff --git a/include/ck_tile/core.hpp b/include/ck_tile/core.hpp index d28d29a0ef..113508808b 100644 --- a/include/ck_tile/core.hpp +++ b/include/ck_tile/core.hpp @@ -52,6 +52,7 @@ #include "ck_tile/core/numeric/mxfp_convert.hpp" #include "ck_tile/core/numeric/null_type.hpp" #include "ck_tile/core/numeric/numeric.hpp" +#include "ck_tile/core/numeric/pk_fp6.hpp" #include "ck_tile/core/numeric/pk_fp4.hpp" #include "ck_tile/core/numeric/pk_int4.hpp" #include "ck_tile/core/numeric/type_convert.hpp" diff --git a/include/ck_tile/core/numeric/pk_fp6.hpp b/include/ck_tile/core/numeric/pk_fp6.hpp new file mode 100644 index 0000000000..31e7ceb17a --- /dev/null +++ b/include/ck_tile/core/numeric/pk_fp6.hpp @@ -0,0 +1,74 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#pragma once + +#include +#include "ck_tile/core/config.hpp" +#include "ck_tile/core/numeric/half.hpp" +#include "ck_tile/core/numeric/mxfp_convert.hpp" + +namespace ck_tile { +template +struct pk_f6_t +{ + static constexpr index_t num_bits_elem = 6; + using element_type = uint32_t; // element storage fundamental type + static constexpr index_t packed_size = pk_size; + static constexpr index_t num_bits_vec_elem = + sizeof(element_type) * 8; // 32-bit uint for storage + static_assert((packed_size * num_bits_elem) % num_bits_vec_elem == 0, + "Packed elements must fit exactly into the element storage."); + static constexpr index_t vector_size = (packed_size * num_bits_elem) / num_bits_vec_elem; + // using storage_type = element_type __attribute__((ext_vector_type(vector_size))); + // storage_type data_{storage_type(0)}; // packed data + element_type data_[3]; // packed data + using type = pk_f6_t; + void pack(const uint32_t x, const index_t i) + { + uint32_t bits = static_cast(x) & 0x3F; + const int bit_pos = i * num_bits_elem; + const int arr_index = bit_pos / num_bits_vec_elem; + const int bit_offset = bit_pos % num_bits_vec_elem; + const int overhang = bit_offset + num_bits_elem - num_bits_vec_elem; + uint32_t old_value = data_[arr_index]; + + // insert bits into the current 32-bit block + old_value |= (bits << bit_offset); + data_[arr_index] = old_value; + + // if it crosses into the next block, shift the remainder + if(overhang > 0 && (arr_index + 1) < vector_size) + { + uint32_t next_value = data_[arr_index + 1]; + next_value |= (bits >> (num_bits_elem - overhang)); + data_[arr_index + 1] = next_value; + } + } + + template + static inline uint32_t unpack(const type& pk, const index_t i) + { + const int bit_pos = i * num_bits_elem; + const int arr_idx = bit_pos / num_bits_vec_elem; + const int bit_offset = bit_pos % num_bits_vec_elem; + const int overhang = bit_offset + num_bits_elem - num_bits_vec_elem; + + uint32_t bits = pk.data_[arr_idx] >> bit_offset; + if(overhang > 0 && (arr_idx + 1) < vector_size) + { + bits |= (pk.data_[arr_idx + 1] & ((1u << overhang) - 1)) << (num_bits_elem - overhang); + } + + return bits & 0x3F; + } + + inline uint32_t unpack(const index_t i) const { return unpack(*this, i); } +}; +using f6x16_pk_t = pk_f6_t<16>; +template <> +struct numeric_traits +{ + static constexpr int PackedSize = 16; +}; +} // namespace ck_tile diff --git a/include/ck_tile/ops/flatmm/kernel/mx_flatmm_kernel.hpp b/include/ck_tile/ops/flatmm/kernel/mx_flatmm_kernel.hpp index 799f8f26a9..65c17b2615 100644 --- a/include/ck_tile/ops/flatmm/kernel/mx_flatmm_kernel.hpp +++ b/include/ck_tile/ops/flatmm/kernel/mx_flatmm_kernel.hpp @@ -314,7 +314,9 @@ struct MXFlatmmKernel : FlatmmKernel, "A tensor for mx must be RowMajor"); diff --git a/include/ck_tile/ops/flatmm/pipeline/mx_flatmm_pipeline_agmem_bgmem_creg_v1.hpp b/include/ck_tile/ops/flatmm/pipeline/mx_flatmm_pipeline_agmem_bgmem_creg_v1.hpp index c3843f1044..2862ee3fc3 100644 --- a/include/ck_tile/ops/flatmm/pipeline/mx_flatmm_pipeline_agmem_bgmem_creg_v1.hpp +++ b/include/ck_tile/ops/flatmm/pipeline/mx_flatmm_pipeline_agmem_bgmem_creg_v1.hpp @@ -115,7 +115,7 @@ struct MXFlatmmPipelineAGmemBGmemCRegV1 : FlatmmPipelineAGmemBGmemCRegV1(); static constexpr index_t MIterPerWarp = kMPerBlock / (MWarp * WG::kM); - static constexpr index_t NIterPerWarp = kNPerBlock / (NWarp * WG::kN); + static constexpr index_t NIterPerWarp = kNPerBlock / (NWarp * WG::kN);//512/(4*16) static constexpr index_t KIterPerWarp = kKPerBlock / WG::kK; static constexpr index_t KFlatBytesPerBlockPerIter = flatKPerWarp / BPackedSize; @@ -569,7 +569,7 @@ struct MXFlatmmPipelineAGmemBGmemCRegV1 : FlatmmPipelineAGmemBGmemCRegV1{}; constexpr auto packed_n_rank = nIter % number{}; return b_flat_dram_window.get_load_offset( - tuple, + tuple, number<0>>{}) + b_flat_dram_window.get_load_offset( tuple, number<0>>{}); @@ -659,7 +659,7 @@ struct MXFlatmmPipelineAGmemBGmemCRegV1 : FlatmmPipelineAGmemBGmemCRegV1( c_warp_tensors(number{})(number{}), bit_cast( - a_warp_tensor(number{})), + a_warp_tensor(number{})),//为什么这里不通过miter和kiter索引？ bit_cast( b_warp_tensor_ping(number{})(number{})), scale_a_tile_tensor_ping(mIter_pack)(kIter_pack) @@ -754,6 +754,7 @@ struct MXFlatmmPipelineAGmemBGmemCRegV1 : FlatmmPipelineAGmemBGmemCRegV1{}) = load_tile_with_offset( a_warp_window_ping, tuple, number>{}); @@ -951,6 +952,8 @@ struct MXFlatmmPipelineAGmemBGmemCRegV1 : FlatmmPipelineAGmemBGmemCRegV1{}) = load_tile_with_offset( a_warp_window_ping, tuple, number>{}); @@ -1052,6 +1055,8 @@ struct MXFlatmmPipelineAGmemBGmemCRegV1 : FlatmmPipelineAGmemBGmemCRegV1{}) = load_tile_with_offset( a_warp_window_ping, tuple, number>{}); diff --git a/include/ck_tile/ops/flatmm/pipeline/mx_flatmm_pipeline_agmem_bgmem_creg_v1_policy.hpp b/include/ck_tile/ops/flatmm/pipeline/mx_flatmm_pipeline_agmem_bgmem_creg_v1_policy.hpp index ed27386f48..2cb20c9f53 100644 --- a/include/ck_tile/ops/flatmm/pipeline/mx_flatmm_pipeline_agmem_bgmem_creg_v1_policy.hpp +++ b/include/ck_tile/ops/flatmm/pipeline/mx_flatmm_pipeline_agmem_bgmem_creg_v1_policy.hpp @@ -212,8 +212,8 @@ struct MXFlatmmPipelineAgBgCrPolicy : UniversalFlatmmPipelineAgBgCrPolicy if constexpr(K_Thread == AK1) return make_static_tile_distribution( tile_distribution_encoding< // - sequence, - tuple, sequence>, + sequence,//4 + tuple, sequence>,//1,2,16 | 4,32 tuple, sequence<2, 1>>, tuple, sequence<0, 2>>, sequence<2>, @@ -339,9 +339,9 @@ struct MXFlatmmPipelineAgBgCrPolicy : UniversalFlatmmPipelineAgBgCrPolicy CK_TILE_HOST_DEVICE static constexpr auto MakeMX_ScaleA_FlatDramTileDistribution() { return make_static_tile_distribution( - tile_distribution_encoding, // ? - tuple, // second direction - sequence>, // first direction + tile_distribution_encoding,//4 // ? + tuple,//1,16 // second direction + sequence>,//4,1 // first direction tuple, sequence<2, 1>>, // which direction tuple, sequence<0, 1>>, // which index //