diff --git a/example/67_gemm_microscaling/gemm_mx_common.hpp b/example/67_gemm_microscaling/gemm_mx_common.hpp index 66645bbef2..b649215bc6 100644 --- a/example/67_gemm_microscaling/gemm_mx_common.hpp +++ b/example/67_gemm_microscaling/gemm_mx_common.hpp @@ -103,6 +103,50 @@ bool parse_cmd_args(int argc, return true; } +#if 1 +void preShuffleScaleBuffer(const ck::e8m0_bexp_t* src, ck::e8m0_bexp_t* dst, int MN, int K) +{ + int MNXdlPack = 2; + int KXdlPack = 2; + + int XdlMNThread = 16; + int XdlKThread = 64 / XdlMNThread; + + int K0 = K / KXdlPack / XdlKThread; // KRepeat + + // The 4 16x128 building blocks will be packed into 1 32x256 for F4 + // The 8 16x16x128 mfma will be packed into 1 32x32x256 for F4 + + // unfold the MN32xK(256/32) scale buffer + // 4 16 2 2 + // To XdlKThread-> XdlMNThread -> KXdlPack -> MNXdlPack + // Then, MNRepeat->KRepeat + + for(int n = 0; n < MN; ++n) + { + for(int k = 0; k < K; ++k) + { + int n0 = n / (XdlMNThread * MNXdlPack); // i MNRepeat + int tempn = n % (XdlMNThread * MNXdlPack); + int n1 = tempn / MNXdlPack; // i XdlMNThread + int n2 = tempn % MNXdlPack; // i MNXdlPack + + int k0 = k / (XdlKThread * KXdlPack); // i KRepeat + int tempk = k % (XdlKThread * KXdlPack); + int k1 = tempk / KXdlPack; // i XdlKThread + int k2 = tempk % KXdlPack; // i KXdlPack + + int outputIndex = n0 * MNXdlPack * KXdlPack * XdlMNThread * XdlKThread * K0 + + k0 * MNXdlPack * KXdlPack * XdlMNThread * XdlKThread + + k1 * MNXdlPack * KXdlPack * XdlMNThread + n1 * MNXdlPack * KXdlPack + + k2 * MNXdlPack + n2; + + dst[outputIndex] = src[n * K + k]; + } + } +} +#endif + template b_k_n_scale(f_host_tensor_descriptor( K / ScaleBlockSize, N, Scale_Stride_BN, BScaleLayout{})); // scales for B + Tensor a_shuffled_scale(f_host_tensor_descriptor( + M, K / ScaleBlockSize, Scale_Stride_AM, AScaleLayout{})); // scales for A + Tensor b_shuffled_scale(f_host_tensor_descriptor( + K / ScaleBlockSize, N, Scale_Stride_BN, BScaleLayout{})); // scales for B + Tensor c_m_n_host_result( f_host_tensor_descriptor(M, N, StrideC, CLayout{})); // host verification Tensor c_m_n_device_result( @@ -283,6 +332,12 @@ bool run_mx_gemm(const ProblemSizeSplitK& problem_size, const ExecutionConfig& c std::cout << "NOTE: No input data initialization." << std::endl; } } +#if 1 + preShuffleScaleBuffer( + a_m_k_scale.mData.data(), a_shuffled_scale.mData.data(), M, K / ScaleBlockSize); + preShuffleScaleBuffer( + b_k_n_scale.mData.data(), b_shuffled_scale.mData.data(), N, K / ScaleBlockSize); +#endif if(config.verbosity > 0) std::cout << "Device memory allocation..." << std::endl; @@ -295,9 +350,18 @@ bool run_mx_gemm(const ProblemSizeSplitK& problem_size, const ExecutionConfig& c if(config.verbosity > 0) std::cout << "Upload data to device..." << std::endl; a_device_buf.ToDevice(a_m_k.mData.data()); - a_scale_device_buf.ToDevice(a_m_k_scale.mData.data()); + a_scale_device_buf.ToDevice(a_shuffled_scale.mData.data()); b_device_buf.ToDevice(b_k_n.mData.data()); - b_scale_device_buf.ToDevice(b_k_n_scale.mData.data()); + b_scale_device_buf.ToDevice(b_shuffled_scale.mData.data()); + // for (size_t i = 0; i < N; i++) + // { + // for (size_t j = 0; j < K / ScaleBlockSize; j++) + // { + // printf("%02x ", *reinterpret_cast(&b_shuffled_scale(j, i))); + // } + // printf("\n"); + // } + if(config.verbosity > 0) std::cout << "Done." << std::endl; 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 2d8def5114..0c3fc8ffdd 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 @@ -72,6 +72,12 @@ struct BlockwiseGemmXdlops_mx_pipeline_base static constexpr index_t MWaves = MPerBlock / (MRepeat * MPerXDL); static constexpr index_t NWaves = NPerBlock / (NRepeat * NPerXDL); + // Hardcode to 2, for better 8-bit access pattern + + static constexpr index_t MXdlPack = 2; + static constexpr index_t NXdlPack = 2; + static constexpr index_t KXdlPack = 2; + using HotLoopInstList = ck::BlockwiseGemmXdlops_pipeline_hotloop_inst< BlockSize, MPerBlock, 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 947b7105ec..a38efd3be7 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 @@ -234,7 +234,7 @@ struct BlockwiseGemmXdlops_pipeline_v1_mx(); + // CK_PRINT(); a_blockwise_copy.MoveSrcSliceWindow(a_grid_desc, a_block_copy_step); b_blockwise_copy.MoveSrcSliceWindow(b_grid_desc, b_block_copy_step); @@ -506,13 +506,13 @@ struct BlockwiseGemmXdlops_pipeline_v1_mx(); - CK_PRINT(); + // CK_PRINT(); + // CK_PRINT(); static_for<0, KRepeat, 1>{}([&](auto k) { constexpr auto k_step = k * xdlops_gemm.KPerXdlops * (KPack / xdlops_gemm.K1PerXdlops); diff --git a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_v3_mx.hpp b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_v3_mx.hpp index 04f4ee8879..b8b248c23d 100644 --- a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_v3_mx.hpp +++ b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_v3_mx.hpp @@ -146,6 +146,10 @@ struct BlockwiseGemmXdlops_pipeline_v3_mx{}([&](auto m0) { - static_for<0, KRepeat, 1>{}([&](auto k0) { - static_for<0, ScalesPerXdlopsRunPerThread, 1>{}([&](auto s) { - constexpr auto a_scale_offset = - a_scale_thread_desc.CalculateOffset(make_tuple(m0, k0, s)); - auto a_scale_thread_buf_copy = - make_static_buffer( - a_scale_thread_desc_copy.GetElementSpaceSize()); - a_scale_thread_copy.Run(a_scale_grid_desc, - a_scale_grid_buf, - a_scale_thread_desc_copy, - make_tuple(I0, I0), - a_scale_thread_buf_copy); + static_for<0, MRepeat / MXdlPack, 1>{}([&](auto m0) { + static_for<0, KRepeat / KXdlPack, 1>{}([&](auto k0) { + a_scale_thread_copy.Run(a_scale_grid_desc, + a_scale_grid_buf, + a_scale_thread_desc, + make_tuple(m0, k0, I0), + a_scale_thread_bufs(I0)); - a_scale_thread_bufs(I0)(Number{}) = - a_scale_thread_buf_copy[Number<0>{}]; - a_scale_thread_copy.MoveSrcSliceWindow( - a_scale_grid_desc, - make_multi_index(0, xdlops_gemm.KPerXdlops / ScaleBlockSize)); - }); + a_scale_thread_copy.MoveSrcSliceWindow(a_scale_grid_desc, + make_multi_index(0, I1, 0)); }); a_scale_thread_copy.MoveSrcSliceWindow( - a_scale_grid_desc, make_multi_index(MWaves * MPerXDL, -ScalesPerKBlockSize)); + a_scale_grid_desc, make_multi_index(MWaves, -KRepeat / KXdlPack, 0)); }); + if(get_thread_local_1d_id() == 0) + { + printf("Scale A: %02x %02x %02x %02x\n", + *reinterpret_cast(&a_scale_thread_bufs(I0)[Number<0>{}]), + *reinterpret_cast(&a_scale_thread_bufs(I0)[Number<1>{}]), + *reinterpret_cast(&a_scale_thread_bufs(I0)[Number<2>{}]), + *reinterpret_cast(&a_scale_thread_bufs(I0)[Number<3>{}])); + } + // restore row id and advance to the next set of scales - a_scale_thread_copy.MoveSrcSliceWindow(a_scale_grid_desc, - make_multi_index(-MPerBlock, ScalesPerKBlockSize)); + a_scale_thread_copy.MoveSrcSliceWindow( + a_scale_grid_desc, make_multi_index(-MWaves * MRepeat / MXdlPack, 0, 0)); // Prefetch b_scales - static_for<0, NRepeat, 1>{}([&](auto n0) { - static_for<0, KRepeat, 1>{}([&](auto k0) { - static_for<0, ScalesPerXdlopsRunPerThread, 1>{}([&](auto s) { - constexpr auto b_scale_offset = - b_scale_thread_desc.CalculateOffset(make_tuple(n0, k0, s)); - auto b_scale_thread_buf_copy = - make_static_buffer( - b_scale_thread_desc_copy.GetElementSpaceSize()); - b_scale_thread_copy.Run(b_scale_grid_desc, - b_scale_grid_buf, - b_scale_thread_desc_copy, - make_tuple(I0, I0), - b_scale_thread_buf_copy); + static_for<0, NRepeat / NXdlPack, 1>{}([&](auto n0) { + static_for<0, KRepeat / KXdlPack, 1>{}([&](auto k0) { + b_scale_thread_copy.Run(b_scale_grid_desc, + b_scale_grid_buf, + b_scale_thread_desc, + make_tuple(n0, k0, I0), + b_scale_thread_bufs(I0)); - b_scale_thread_bufs(I0)(Number{}) = - b_scale_thread_buf_copy[Number<0>{}]; - b_scale_thread_copy.MoveSrcSliceWindow( - b_scale_grid_desc, - make_multi_index(0, xdlops_gemm.KPerXdlops / ScaleBlockSize)); - }); + b_scale_thread_copy.MoveSrcSliceWindow(b_scale_grid_desc, + make_multi_index(0, I1, 0)); }); b_scale_thread_copy.MoveSrcSliceWindow( - b_scale_grid_desc, make_multi_index(NWaves * NPerXDL, -ScalesPerKBlockSize)); + b_scale_grid_desc, make_multi_index(NWaves, -KRepeat / KXdlPack, 0)); }); + if(get_thread_local_1d_id() == 0) + { + printf("Scale B: %02x %02x %02x %02x\n", + *reinterpret_cast(&b_scale_thread_bufs(I0)[Number<0>{}]), + *reinterpret_cast(&b_scale_thread_bufs(I0)[Number<1>{}]), + *reinterpret_cast(&b_scale_thread_bufs(I0)[Number<2>{}]), + *reinterpret_cast(&b_scale_thread_bufs(I0)[Number<3>{}])); + } + // restore col id and advance to the next set of scales // NWaves * NPerXDL * NRepeat == NPerBlock - b_scale_thread_copy.MoveSrcSliceWindow(b_scale_grid_desc, - make_multi_index(-NPerBlock, ScalesPerKBlockSize)); + b_scale_thread_copy.MoveSrcSliceWindow( + b_scale_grid_desc, make_multi_index(-NWaves * NRepeat / NXdlPack, 0, 0)); // Local prefill 1 a_blockwise_copy.RunWrite(a_block_desc, a_block_buf); @@ -464,66 +466,45 @@ struct BlockwiseGemmXdlops_pipeline_v3_mx{}([&](auto m0) { - static_for<0, KRepeat, 1>{}([&](auto k0) { - static_for<0, ScalesPerXdlopsRunPerThread, 1>{}([&](auto s) { - constexpr auto a_scale_offset = - a_scale_thread_desc.CalculateOffset(make_tuple(m0, k0, s)); - auto a_scale_thread_buf_copy = - make_static_buffer( - a_scale_thread_desc_copy.GetElementSpaceSize()); - a_scale_thread_copy.Run(a_scale_grid_desc, - a_scale_grid_buf, - a_scale_thread_desc_copy, - make_tuple(I0, I0), - a_scale_thread_buf_copy); + static_for<0, MRepeat / MXdlPack, 1>{}([&](auto m0) { + static_for<0, KRepeat / KXdlPack, 1>{}([&](auto k0) { + a_scale_thread_copy.Run(a_scale_grid_desc, + a_scale_grid_buf, + a_scale_thread_desc, + make_tuple(m0, k0, I0), + a_scale_thread_bufs(scale_mem_buf)); - a_scale_thread_bufs(scale_mem_buf)(Number{}) = - a_scale_thread_buf_copy[Number<0>{}]; - a_scale_thread_copy.MoveSrcSliceWindow( - a_scale_grid_desc, - make_multi_index(0, xdlops_gemm.KPerXdlops / ScaleBlockSize)); - }); + a_scale_thread_copy.MoveSrcSliceWindow(a_scale_grid_desc, + make_multi_index(0, I1, 0)); }); a_scale_thread_copy.MoveSrcSliceWindow( - a_scale_grid_desc, - make_multi_index(MWaves * MPerXDL, -ScalesPerKBlockSize)); + a_scale_grid_desc, make_multi_index(MWaves, -KRepeat / KXdlPack, 0)); }); // restore row id and advance to the next set of scales a_scale_thread_copy.MoveSrcSliceWindow( - a_scale_grid_desc, make_multi_index(-MPerBlock, ScalesPerKBlockSize)); + a_scale_grid_desc, make_multi_index(-MWaves * MRepeat / MXdlPack, 0, 0)); // Prefetch b_scales - static_for<0, NRepeat, 1>{}([&](auto n0) { - static_for<0, KRepeat, 1>{}([&](auto k0) { - static_for<0, ScalesPerXdlopsRunPerThread, 1>{}([&](auto s) { - constexpr auto b_scale_offset = - b_scale_thread_desc.CalculateOffset(make_tuple(n0, k0, s)); - auto b_scale_thread_buf_copy = - make_static_buffer( - b_scale_thread_desc_copy.GetElementSpaceSize()); - b_scale_thread_copy.Run(b_scale_grid_desc, - b_scale_grid_buf, - b_scale_thread_desc_copy, - make_tuple(I0, I0), - b_scale_thread_buf_copy); + static_for<0, NRepeat / NXdlPack, 1>{}([&](auto n0) { + static_for<0, KRepeat / KXdlPack, 1>{}([&](auto k0) { + b_scale_thread_copy.Run(b_scale_grid_desc, + b_scale_grid_buf, + b_scale_thread_desc, + make_tuple(n0, k0, I0), + b_scale_thread_bufs(scale_mem_buf)); - b_scale_thread_bufs(scale_mem_buf)(Number{}) = - b_scale_thread_buf_copy[Number<0>{}]; - b_scale_thread_copy.MoveSrcSliceWindow( - b_scale_grid_desc, - make_multi_index(0, xdlops_gemm.KPerXdlops / ScaleBlockSize)); - }); + b_scale_thread_copy.MoveSrcSliceWindow(b_scale_grid_desc, + make_multi_index(0, I1, 0)); }); b_scale_thread_copy.MoveSrcSliceWindow( - b_scale_grid_desc, - make_multi_index(NWaves * NPerXDL, -ScalesPerKBlockSize)); + b_scale_grid_desc, make_multi_index(NWaves, -KRepeat / KXdlPack, 0)); }); + // restore col id and advance to the next set of scales // NWaves * NPerXDL * NRepeat == NPerBlock b_scale_thread_copy.MoveSrcSliceWindow( - b_scale_grid_desc, make_multi_index(-NPerBlock, ScalesPerKBlockSize)); + b_scale_grid_desc, make_multi_index(-NWaves * NRepeat / NXdlPack, 0, 0)); // TODO: consider scheduling the scale load // ------------------------------------------------------------------------------------------- @@ -538,68 +519,92 @@ struct BlockwiseGemmXdlops_pipeline_v3_mx{}([&](auto m0) { - static_for<0, NRepeat, 1>{}([&](auto n0) { - static_for<0, KRepeat, 1>{}([&](auto k0) { - vector_type - a_thread_vec; // = vec: pk_i4_t, 32 - vector_type b_thread_vec; - - static_for<0, KPack / APackedSize, 1>{}([&](auto ik) { - a_thread_vec.template AsType()(ik) = - a_thread_buf[Number{}]; - b_thread_vec.template AsType()(ik) = - b_thread_buf[Number{}]; - }); - + static_for<0, MRepeat / MXdlPack, 1>{}([&](auto m0) { + static_for<0, NRepeat / NXdlPack, 1>{}([&](auto n0) { + static_for<0, KRepeat / KXdlPack, 1>{}([&](auto k0) { constexpr index_t a_scale_offset = a_scale_thread_desc.CalculateOffset(make_tuple(m0, k0, I0)); constexpr index_t b_scale_offset = b_scale_thread_desc.CalculateOffset(make_tuple(n0, k0, I0)); - static_assert( - 0 < ScalesPerXdlopsRunPerThread, - "Must have at least one scale per Xdlops per Thread."); + static_assert(0 < ScalesPerXdlopsRunPerThread, + "Must have at least one scale per Xdlops " + "per Thread."); - vector_type - a_scale_thread_vec; - vector_type - b_scale_thread_vec; + vector_type a_scale_thread_vec; + vector_type b_scale_thread_vec; // Pack scale_thread_buf into scale_thread_vec - static_for<0, ScalesPerXdlopsRunPerThread, 1>{}([&](auto s) { + static_for<0, KXdlPack * MXdlPack, 1>{}([&](auto s) { a_scale_thread_vec.template AsType()(s) = a_scale_thread_bufs( scale_comp_buf)[Number{}]; + }); + + static_for<0, KXdlPack * NXdlPack, 1>{}([&](auto s) { b_scale_thread_vec.template AsType()(s) = b_scale_thread_bufs( scale_comp_buf)[Number{}]; }); - // CK_PRINT(); - // CK_PRINT(); - using mfma_input_type_a = - typename vector_type::type; - // mfma input type = pk_f4_t, 32 - // CK_PRINT(); - using mfma_input_type_b = - typename vector_type::type; - constexpr index_t c_offset = - c_thread_desc_.CalculateOffset(make_tuple(m0, n0, 0)); + static_for<0, KXdlPack, 1>{}([&](auto ikxdl) { + static_for<0, MXdlPack, 1>{}([&](auto imxdl) { + static_for<0, NXdlPack, 1>{}([&](auto inxdl) { + constexpr auto kxdl = ikxdl + k0 * KXdlPack; + constexpr auto mxdl = imxdl + m0 * MXdlPack; + constexpr auto nxdl = inxdl + n0 * NXdlPack; - // MFMA accumulation - xdlops_gemm.template Run<>( - a_thread_vec.template AsType(), - a_scale_thread_vec.template AsType(), - b_thread_vec.template AsType(), - b_scale_thread_vec.template AsType(), - c_thread_buf.GetVectorTypeReference(Number{})); + vector_type + a_thread_vec; + vector_type + b_thread_vec; + + static_for<0, KPack / APackedSize, 1>{}([&](auto ik) { + a_thread_vec.template AsType()(ik) = + a_thread_buf + [Number{}]; + b_thread_vec.template AsType()(ik) = + b_thread_buf + [Number{}]; + }); + + using mfma_input_type_a = + typename vector_type::type; + + using mfma_input_type_b = + typename vector_type::type; + + using mfma_scale_input_type_a = + typename vector_type::type; + using mfma_scale_input_type_b = + typename vector_type::type; + + constexpr index_t c_offset = + c_thread_desc_.CalculateOffset( + make_tuple(mxdl, nxdl, 0)); + + // MFMA accumulation + xdlops_gemm.template Run( + a_thread_vec.template AsType(), + a_scale_thread_vec + .template AsType(), + b_thread_vec.template AsType(), + b_scale_thread_vec + .template AsType(), + c_thread_buf.GetVectorTypeReference( + Number{})); + }); + }); + }); }); }); }); @@ -667,111 +672,127 @@ struct BlockwiseGemmXdlops_pipeline_v3_mx{}([&](auto m0) { - static_for<0, KRepeat, 1>{}([&](auto k0) { - static_for<0, ScalesPerXdlopsRunPerThread, 1>{}([&](auto s) { - constexpr auto a_scale_offset = - a_scale_thread_desc.CalculateOffset(make_tuple(m0, k0, s)); - auto a_scale_thread_buf_copy = - make_static_buffer( - a_scale_thread_desc_copy.GetElementSpaceSize()); - a_scale_thread_copy.Run(a_scale_grid_desc, - a_scale_grid_buf, - a_scale_thread_desc_copy, - make_tuple(I0, I0), - a_scale_thread_buf_copy); + // Global prefetch 1 + a_blockwise_copy.RunRead(a_grid_desc, a_grid_buf); + b_blockwise_copy.RunRead(b_grid_desc, b_grid_buf); - a_scale_thread_bufs(I1)(Number{}) = - a_scale_thread_buf_copy[Number<0>{}]; - a_scale_thread_copy.MoveSrcSliceWindow( - a_scale_grid_desc, - make_multi_index(0, xdlops_gemm.KPerXdlops / ScaleBlockSize)); - }); + a_blockwise_copy.MoveSrcSliceWindow(a_grid_desc, a_block_copy_step); + b_blockwise_copy.MoveSrcSliceWindow(b_grid_desc, b_block_copy_step); + + // Prefetch a_scales + static_for<0, MRepeat / MXdlPack, 1>{}([&](auto m0) { + static_for<0, KRepeat / KXdlPack, 1>{}([&](auto k0) { + a_scale_thread_copy.Run(a_scale_grid_desc, + a_scale_grid_buf, + a_scale_thread_desc, + make_tuple(m0, k0, I0), + a_scale_thread_bufs(I1)); + + a_scale_thread_copy.MoveSrcSliceWindow(a_scale_grid_desc, + make_multi_index(0, I1, 0)); }); a_scale_thread_copy.MoveSrcSliceWindow( - a_scale_grid_desc, make_multi_index(MWaves * MPerXDL, -ScalesPerKBlockSize)); + a_scale_grid_desc, make_multi_index(MWaves, -KRepeat / KXdlPack, 0)); }); // Prefetch b_scales - static_for<0, NRepeat, 1>{}([&](auto n0) { - static_for<0, KRepeat, 1>{}([&](auto k0) { - static_for<0, ScalesPerXdlopsRunPerThread, 1>{}([&](auto s) { - constexpr auto b_scale_offset = - b_scale_thread_desc.CalculateOffset(make_tuple(n0, k0, s)); - auto b_scale_thread_buf_copy = - make_static_buffer( - b_scale_thread_desc_copy.GetElementSpaceSize()); - b_scale_thread_copy.Run(b_scale_grid_desc, - b_scale_grid_buf, - b_scale_thread_desc_copy, - make_tuple(I0, I0), - b_scale_thread_buf_copy); + static_for<0, NRepeat / NXdlPack, 1>{}([&](auto n0) { + static_for<0, KRepeat / KXdlPack, 1>{}([&](auto k0) { + b_scale_thread_copy.Run(b_scale_grid_desc, + b_scale_grid_buf, + b_scale_thread_desc, + make_tuple(n0, k0, I0), + b_scale_thread_bufs(I1)); - b_scale_thread_bufs(I1)(Number{}) = - b_scale_thread_buf_copy[Number<0>{}]; - b_scale_thread_copy.MoveSrcSliceWindow( - b_scale_grid_desc, - make_multi_index(0, xdlops_gemm.KPerXdlops / ScaleBlockSize)); - }); + b_scale_thread_copy.MoveSrcSliceWindow(b_scale_grid_desc, + make_multi_index(0, I1, 0)); }); b_scale_thread_copy.MoveSrcSliceWindow( - b_scale_grid_desc, make_multi_index(NWaves * NPerXDL, -ScalesPerKBlockSize)); + b_scale_grid_desc, make_multi_index(NWaves, -KRepeat / KXdlPack, 0)); }); block_sync_lds(); a_blockwise_copy.RunWrite(a_block_desc, a_block_buf); b_blockwise_copy.RunWrite(b_block_desc, b_block_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 b_thread_vec; - - static_for<0, KPack / APackedSize, 1>{}([&](auto ik) { - a_thread_vec.template AsType()(ik) = - a_thread_buf[Number{}]; - b_thread_vec.template AsType()(ik) = - b_thread_buf[Number{}]; - }); - + static_for<0, MRepeat / MXdlPack, 1>{}([&](auto m0) { + static_for<0, NRepeat / NXdlPack, 1>{}([&](auto n0) { + static_for<0, KRepeat / KXdlPack, 1>{}([&](auto k0) { constexpr index_t a_scale_offset = a_scale_thread_desc.CalculateOffset(make_tuple(m0, k0, I0)); - constexpr index_t b_scale_offset = b_scale_thread_desc.CalculateOffset(make_tuple(n0, k0, I0)); - vector_type a_scale_thread_vec; - vector_type b_scale_thread_vec; + static_assert(0 < ScalesPerXdlopsRunPerThread, + "Must have at least one scale per Xdlops " + "per Thread."); - // Pack b_scale_thread_buf into b_scale_thread_vec - static_for<0, ScalesPerXdlopsRunPerThread, 1>{}([&](auto s) { + vector_type a_scale_thread_vec; + vector_type b_scale_thread_vec; + + // Pack scale_thread_buf into scale_thread_vec + static_for<0, KXdlPack * MXdlPack, 1>{}([&](auto s) { a_scale_thread_vec.template AsType()(s) = a_scale_thread_bufs(I0)[Number{}]; + }); + + static_for<0, KXdlPack * NXdlPack, 1>{}([&](auto s) { b_scale_thread_vec.template AsType()(s) = b_scale_thread_bufs(I0)[Number{}]; }); - using mfma_input_type_a = - typename vector_type::type; - using mfma_input_type_b = - typename vector_type::type; + static_for<0, KXdlPack, 1>{}([&](auto ikxdl) { + static_for<0, MXdlPack, 1>{}([&](auto imxdl) { + static_for<0, NXdlPack, 1>{}([&](auto inxdl) { + constexpr auto kxdl = ikxdl + k0 * KXdlPack; + constexpr auto mxdl = imxdl + m0 * MXdlPack; + constexpr auto nxdl = inxdl + n0 * NXdlPack; - constexpr index_t c_offset = - c_thread_desc_.CalculateOffset(make_tuple(m0, n0, 0)); + vector_type a_thread_vec; + vector_type b_thread_vec; - // MFMA accumulation - xdlops_gemm.template Run<>( - a_thread_vec.template AsType(), - a_scale_thread_vec.template AsType(), - b_thread_vec.template AsType(), - b_scale_thread_vec.template AsType(), - c_thread_buf.GetVectorTypeReference(Number{})); + static_for<0, KPack / APackedSize, 1>{}([&](auto ik) { + a_thread_vec.template AsType()(ik) = + a_thread_buf[Number{}]; + b_thread_vec.template AsType()(ik) = + b_thread_buf[Number{}]; + }); + + using mfma_input_type_a = + typename vector_type::type; + + using mfma_input_type_b = + typename vector_type::type; + + using mfma_scale_input_type_a = + typename vector_type::type; + using mfma_scale_input_type_b = + typename vector_type::type; + + constexpr index_t c_offset = + c_thread_desc_.CalculateOffset(make_tuple(mxdl, nxdl, 0)); + + // MFMA accumulation + xdlops_gemm.template Run( + a_thread_vec.template AsType(), + a_scale_thread_vec + .template AsType(), + b_thread_vec.template AsType(), + b_scale_thread_vec + .template AsType(), + c_thread_buf.GetVectorTypeReference(Number{})); + }); + }); + }); }); }); }); @@ -809,110 +830,168 @@ struct BlockwiseGemmXdlops_pipeline_v3_mx{}([&](auto m0) { - static_for<0, NRepeat, 1>{}([&](auto n0) { - static_for<0, KRepeat, 1>{}([&](auto k0) { - vector_type a_thread_vec; - vector_type b_thread_vec; - - static_for<0, KPack / APackedSize, 1>{}([&](auto ik) { - a_thread_vec.template AsType()(ik) = - a_thread_buf[Number{}]; - b_thread_vec.template AsType()(ik) = - b_thread_buf[Number{}]; - }); - + static_for<0, MRepeat / MXdlPack, 1>{}([&](auto m0) { + static_for<0, NRepeat / NXdlPack, 1>{}([&](auto n0) { + static_for<0, KRepeat / KXdlPack, 1>{}([&](auto k0) { constexpr index_t a_scale_offset = a_scale_thread_desc.CalculateOffset(make_tuple(m0, k0, I0)); - constexpr index_t b_scale_offset = b_scale_thread_desc.CalculateOffset(make_tuple(n0, k0, I0)); - vector_type a_scale_thread_vec; - vector_type b_scale_thread_vec; + static_assert(0 < ScalesPerXdlopsRunPerThread, + "Must have at least one scale per Xdlops " + "per Thread."); - // Pack b_scale_thread_buf into b_scale_thread_vec - static_for<0, ScalesPerXdlopsRunPerThread, 1>{}([&](auto s) { + vector_type a_scale_thread_vec; + vector_type b_scale_thread_vec; + + // Pack scale_thread_buf into scale_thread_vec + static_for<0, KXdlPack * MXdlPack, 1>{}([&](auto s) { a_scale_thread_vec.template AsType()(s) = a_scale_thread_bufs(I1)[Number{}]; + }); + + static_for<0, KXdlPack * NXdlPack, 1>{}([&](auto s) { b_scale_thread_vec.template AsType()(s) = b_scale_thread_bufs(I1)[Number{}]; }); - using mfma_input_type_a = - typename vector_type::type; - using mfma_input_type_b = - typename vector_type::type; + static_for<0, KXdlPack, 1>{}([&](auto ikxdl) { + static_for<0, MXdlPack, 1>{}([&](auto imxdl) { + static_for<0, NXdlPack, 1>{}([&](auto inxdl) { + constexpr auto kxdl = ikxdl + k0 * KXdlPack; + constexpr auto mxdl = imxdl + m0 * MXdlPack; + constexpr auto nxdl = inxdl + n0 * NXdlPack; - constexpr index_t c_offset = - c_thread_desc_.CalculateOffset(make_tuple(m0, n0, 0)); + vector_type a_thread_vec; + vector_type b_thread_vec; - // MFMA accumulation - xdlops_gemm.template Run<>( - a_thread_vec.template AsType(), - a_scale_thread_vec.template AsType(), - b_thread_vec.template AsType(), - b_scale_thread_vec.template AsType(), - c_thread_buf.GetVectorTypeReference(Number{})); + static_for<0, KPack / APackedSize, 1>{}([&](auto ik) { + a_thread_vec.template AsType()(ik) = + a_thread_buf[Number{}]; + b_thread_vec.template AsType()(ik) = + b_thread_buf[Number{}]; + }); + + using mfma_input_type_a = + typename vector_type::type; + + using mfma_input_type_b = + typename vector_type::type; + + using mfma_scale_input_type_a = + typename vector_type::type; + using mfma_scale_input_type_b = + typename vector_type::type; + + constexpr index_t c_offset = + c_thread_desc_.CalculateOffset(make_tuple(mxdl, nxdl, 0)); + + // MFMA accumulation + xdlops_gemm.template Run( + a_thread_vec.template AsType(), + a_scale_thread_vec + .template AsType(), + b_thread_vec.template AsType(), + b_scale_thread_vec + .template AsType(), + c_thread_buf.GetVectorTypeReference(Number{})); + }); + }); + }); }); }); }); } else if constexpr(TailNum == TailNumber::Odd) { - 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 b_thread_vec; - - static_for<0, KPack / APackedSize, 1>{}([&](auto ik) { - a_thread_vec.template AsType()(ik) = - a_thread_buf[Number{}]; - b_thread_vec.template AsType()(ik) = - b_thread_buf[Number{}]; - }); - + static_for<0, MRepeat / MXdlPack, 1>{}([&](auto m0) { + static_for<0, NRepeat / NXdlPack, 1>{}([&](auto n0) { + static_for<0, KRepeat / KXdlPack, 1>{}([&](auto k0) { constexpr index_t a_scale_offset = a_scale_thread_desc.CalculateOffset(make_tuple(m0, k0, I0)); - constexpr index_t b_scale_offset = b_scale_thread_desc.CalculateOffset(make_tuple(n0, k0, I0)); - vector_type a_scale_thread_vec; - vector_type b_scale_thread_vec; + static_assert(0 < ScalesPerXdlopsRunPerThread, + "Must have at least one scale per Xdlops " + "per Thread."); - // Pack b_scale_thread_buf into b_scale_thread_vec - static_for<0, ScalesPerXdlopsRunPerThread, 1>{}([&](auto s) { + vector_type a_scale_thread_vec; + vector_type b_scale_thread_vec; + + // Pack scale_thread_buf into scale_thread_vec + static_for<0, KXdlPack * MXdlPack, 1>{}([&](auto s) { a_scale_thread_vec.template AsType()(s) = a_scale_thread_bufs(I0)[Number{}]; + }); + + static_for<0, KXdlPack * NXdlPack, 1>{}([&](auto s) { b_scale_thread_vec.template AsType()(s) = b_scale_thread_bufs(I0)[Number{}]; }); - using mfma_input_type_a = - typename vector_type::type; - using mfma_input_type_b = - typename vector_type::type; + static_for<0, KXdlPack, 1>{}([&](auto ikxdl) { + static_for<0, MXdlPack, 1>{}([&](auto imxdl) { + static_for<0, NXdlPack, 1>{}([&](auto inxdl) { + constexpr auto kxdl = ikxdl + k0 * KXdlPack; + constexpr auto mxdl = imxdl + m0 * MXdlPack; + constexpr auto nxdl = inxdl + n0 * NXdlPack; - constexpr index_t c_offset = - c_thread_desc_.CalculateOffset(make_tuple(m0, n0, 0)); + vector_type a_thread_vec; + vector_type b_thread_vec; - // MFMA accumulation - xdlops_gemm.template Run<>( - a_thread_vec.template AsType(), - a_scale_thread_vec.template AsType(), - b_thread_vec.template AsType(), - b_scale_thread_vec.template AsType(), - c_thread_buf.GetVectorTypeReference(Number{})); + static_for<0, KPack / APackedSize, 1>{}([&](auto ik) { + a_thread_vec.template AsType()(ik) = + a_thread_buf[Number{}]; + b_thread_vec.template AsType()(ik) = + b_thread_buf[Number{}]; + }); + + using mfma_input_type_a = + typename vector_type::type; + + using mfma_input_type_b = + typename vector_type::type; + + using mfma_scale_input_type_a = + typename vector_type::type; + using mfma_scale_input_type_b = + typename vector_type::type; + + constexpr index_t c_offset = + c_thread_desc_.CalculateOffset(make_tuple(mxdl, nxdl, 0)); + + // MFMA accumulation + xdlops_gemm.template Run( + a_thread_vec.template AsType(), + a_scale_thread_vec + .template AsType(), + b_thread_vec.template AsType(), + b_scale_thread_vec + .template AsType(), + c_thread_buf.GetVectorTypeReference(Number{})); + }); + }); + }); }); }); }); @@ -922,20 +1001,16 @@ struct BlockwiseGemmXdlops_pipeline_v3_mx{}, Number{}, Number{})); - - // Is used to copy data from a_scale_grid to a_scale_thread - static constexpr auto a_scale_thread_desc_copy = - make_naive_tensor_descriptor_packed(make_tuple(Number<1>{}, Number<1>{})); + make_tuple(Number{}, + Number{}, + Number{})); // TODO: make this field protected when b_scale_thread_copy_ is moved // here static constexpr auto b_scale_thread_desc = make_naive_tensor_descriptor_packed( - make_tuple(Number{}, Number{}, Number{})); - - // Is used to copy data from b_scale_grid to b_scale_thread_buf - static constexpr auto b_scale_thread_desc_copy = - make_naive_tensor_descriptor_packed(make_tuple(Number<1>{}, Number<1>{})); + make_tuple(Number{}, + Number{}, + Number{})); protected: using Base::a_thread_copy_; diff --git a/include/ck/tensor_operation/gpu/device/impl/device_gemm_xdl_cshuffle_v3_mx.hpp b/include/ck/tensor_operation/gpu/device/impl/device_gemm_xdl_cshuffle_v3_mx.hpp index 500013fc79..f71f47d0be 100644 --- a/include/ck/tensor_operation/gpu/device/impl/device_gemm_xdl_cshuffle_v3_mx.hpp +++ b/include/ck/tensor_operation/gpu/device/impl/device_gemm_xdl_cshuffle_v3_mx.hpp @@ -340,6 +340,7 @@ struct DeviceGemmMX_Xdl_CShuffleV3 : public DeviceGemmMX 1) { if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == TailNumber::Odd) @@ -386,6 +387,13 @@ struct DeviceGemmMX_Xdl_CShuffleV3 : public DeviceGemmMX; + Run(kernel); } else { 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 3e52940722..1927eae105 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 @@ -163,6 +163,10 @@ struct GridwiseGemmMX_xdl_cshuffle_v3 static constexpr bool is_single_rate_mfma = false; static constexpr auto is_scale_mfma = true; + static constexpr auto MXdlPack = 2; + static constexpr auto NXdlPack = 2; + static constexpr auto KXdlPack = 2; + //> KPack is at least the k_per_blk of selected mfma // // Should be a multiple of k_per_blk. @@ -1468,7 +1472,6 @@ struct GridwiseGemmMX_xdl_cshuffle_v3 auto a_block_buf = make_dynamic_buffer( static_cast(p_shared), a_block_desc_ak0_m_ak1.GetElementSpaceSize() / APackedSize); - CK_PRINT>(); auto b_block_buf = make_dynamic_buffer( reinterpret_cast(static_cast(p_shared) + a_block_space_size_aligned * @@ -1509,42 +1512,47 @@ struct GridwiseGemmMX_xdl_cshuffle_v3 const auto waveId_m = wave_idx[I0]; const auto waveId_n = wave_idx[I1]; - static constexpr auto mfma = BlockwiseGemmPipe::xdlops_gemm.mfma; + // static constexpr auto mfma = BlockwiseGemmPipe::xdlops_gemm.mfma; - auto thread_offset_k = (get_thread_local_1d_id() % BlockwiseGemmPipe::WaveSize) / - mfma.selected_mfma.num_threads_per_blk; + // auto thread_offset_k = (get_thread_local_1d_id() % BlockwiseGemmPipe::WaveSize) / + // mfma.selected_mfma.num_threads_per_blk; - auto a_thread_offset_m = get_thread_local_1d_id() % MPerXdl + waveId_m * MPerXdl; + // A wave access continuous memory + auto thread_offset_shuffled = get_thread_local_1d_id() % BlockwiseGemmPipe::WaveSize; + + auto a_thread_offset_m = waveId_m * MPerXdl * MXdlPack; auto a_scale_thread_copy = ThreadwiseTensorSliceTransfer_v2, // SliceLengths - Sequence<0, 1>, // DimAccessOrder - 1, // SrcVectorDim - 1, // SrcScalarPerVector - 1, // SrcScalarStrideInVector + decltype(BlockwiseGemmPipe::a_scale_thread_desc), + Sequence<1, 1, KXdlPack * MXdlPack>, // SliceLengths + Sequence<0, 1, 2>, // DimAccessOrder + 2, // SrcVectorDim + KXdlPack * MXdlPack, // SrcScalarPerVector + 1, // SrcScalarStrideInVector true>( a_scale_grid_desc_am_ak, - make_multi_index(block_m_id * MPerBlock + a_thread_offset_m, thread_offset_k)); + make_multi_index( + block_m_id * MPerBlock + a_thread_offset_m, 0, thread_offset_shuffled)); - auto b_thread_offset_n = get_thread_local_1d_id() % NPerXdl + waveId_n * NPerXdl; + auto b_thread_offset_n = waveId_n * NPerXdl * NXdlPack; auto b_scale_thread_copy = ThreadwiseTensorSliceTransfer_v2, // SliceLengths - Sequence<0, 1>, // DimAccessOrder - 1, // SrcVectorDim - 1, // SrcScalarPerVector + decltype(BlockwiseGemmPipe::b_scale_thread_desc), + Sequence<1, 1, KXdlPack * NXdlPack>, // SliceLengths + Sequence<0, 1, 2>, // DimAccessOrder + 2, // SrcVectorDim + KXdlPack * MXdlPack, // SrcScalarPerVector 1, true>( b_scale_grid_desc_bn_ak, - make_multi_index(block_n_id * NPerBlock + b_thread_offset_n, thread_offset_k)); + make_multi_index( + block_n_id * NPerBlock + b_thread_offset_n, 0, thread_offset_shuffled)); blockwise_gemm_pipeline.template Run(a_grid_desc_ak0_m_ak1, a_block_desc_ak0_m_ak1, @@ -1787,6 +1795,7 @@ struct GridwiseGemmMX_xdl_cshuffle_v3 MakeCGridDescriptor_MBlock_MPerBlock_NBlock_NPerBlock( c_grid_desc_m_n, problem.MBlock, problem.NBlock); +#if 0 // A Scale grid const auto a_scale_grid_desc_am_ak = make_naive_tensor_descriptor( make_tuple(problem.M, math::integer_divide_ceil(problem.K, ScaleBlockSize)), @@ -1796,6 +1805,18 @@ struct GridwiseGemmMX_xdl_cshuffle_v3 const auto b_scale_grid_desc_bn_ak = make_naive_tensor_descriptor( make_tuple(problem.N, math::integer_divide_ceil(problem.K, ScaleBlockSize)), make_tuple(problem.StrideScaleB, 1)); +#endif + // A/B shuffled scale for better 8-bit scale access pattern + // MNRepeat -> KRepeat -> KThreadPerXdl -> MNThreadPerXdl -> KXdlPack -> MNXdlPack + const auto a_scale_grid_desc_am_ak = make_naive_tensor_descriptor_packed(make_tuple( + problem.M / (MXdlPack * MPerXdl), + math::integer_divide_ceil(problem.K, ScaleBlockSize) / (KXdlPack * 64 / MPerXdl), + 64 * KXdlPack * MXdlPack)); + + const auto b_scale_grid_desc_bn_ak = make_naive_tensor_descriptor_packed(make_tuple( + problem.N / (NXdlPack * NPerXdl), + math::integer_divide_ceil(problem.K, ScaleBlockSize) / (KXdlPack * 64 / NPerXdl), + 64 * KXdlPack * NXdlPack)); Run(); - CK_PRINT(); + // CK_PRINT(); + // CK_PRINT(); static_ford{}([&](auto ordered_access_idx) { #if 0 // TODO: unable to compile 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 1592d77286..9e988f61c0 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 @@ -169,7 +169,7 @@ struct ThreadwiseTensorSliceTransfer_v3r1 }, Number{}); - CK_PRINT(); + // // CK_PRINT(); // loop over tensor and copy static_ford{}([&](auto ordered_src_access_idx) { // judge move forward or move backward @@ -282,7 +282,7 @@ struct ThreadwiseTensorSliceTransfer_v3r1 Sequence, Sequence>; - CK_PRINT>(); + // // CK_PRINT>(); static_for<0, tuple_element_t::Size(), 1>{}( [&](auto v_idx) { constexpr auto VectorLoadSize = @@ -292,7 +292,7 @@ struct ThreadwiseTensorSliceTransfer_v3r1 using src_vector_container = vector_type_maker_t; using src_vector_container_t = typename src_vector_container::type; - CK_PRINT(); + // CK_PRINT(); src_vector_container src_vector = src_vector_container{src_buf.template Get( @@ -553,7 +553,7 @@ struct ThreadwiseTensorSliceTransfer_v3r1 constexpr auto dst_dim_access_order = DstDimAccessOrder{}; - CK_PRINT(); + // CK_PRINT(); constexpr auto ordered_dst_access_lengths = container_reorder_given_new2old(dst_access_lengths, dst_dim_access_order); @@ -584,7 +584,7 @@ struct ThreadwiseTensorSliceTransfer_v3r1 Number{}); // loop over tensor and copy - CK_PRINT(); + // CK_PRINT(); static_ford{}([&](auto ordered_dst_access_idx) { // judge move forward or move backward constexpr auto forward_sweep = [&]() { diff --git a/include/ck/tensor_operation/gpu/warp/xdlops_gemm.hpp b/include/ck/tensor_operation/gpu/warp/xdlops_gemm.hpp index 84cb603b87..865a8f0fd0 100644 --- a/include/ck/tensor_operation/gpu/warp/xdlops_gemm.hpp +++ b/include/ck/tensor_operation/gpu/warp/xdlops_gemm.hpp @@ -886,6 +886,8 @@ struct mfma_type template const ScaleB& scale_b, FloatC& reg_c) const { - 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."); + if(get_thread_local_1d_id() == 0) + { + printf("Before BitCast: Scale A: %08x, Scale B: %08x\n", + *reinterpret_cast(&scale_a), + *reinterpret_cast(&scale_b)); + } + // 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."); - intrin_mfma_scale_f32_16x16x128f8f6f4::Run( - a, utils::get_exponent_value(scale_a), b, utils::get_exponent_value(scale_b), reg_c); + // intrin_mfma_scale_f32_16x16x128f8f6f4::Run( + // a, utils::get_exponent_value(scale_a), b, utils::get_exponent_value(scale_b), reg_c); + intrin_mfma_scale_f32_16x16x128f8f6f4::Run( + a, bit_cast(scale_a), b, bit_cast(scale_b), reg_c); } }; @@ -1441,7 +1452,13 @@ struct XdlopsGemm }); } - template + template __device__ void Run(const FloatA& p_a_wave, const ScaleA& a_scale_thread, const FloatB& p_b_wave, @@ -1451,12 +1468,12 @@ struct XdlopsGemm static_for<0, KPack / mfma_instr.k_per_blk, 1>{}([&](auto k) { if constexpr(!TransposeC) { - mfma_instr.template run( + mfma_instr.template run( p_a_wave[k], a_scale_thread[k], p_b_wave[k], b_scale_thread[k], p_c_thread); } else { - mfma_instr.template run( + mfma_instr.template run( p_b_wave[k], b_scale_thread[k], p_a_wave[k], a_scale_thread[k], p_c_thread); } }); diff --git a/include/ck/utility/amd_buffer_addressing.hpp b/include/ck/utility/amd_buffer_addressing.hpp index 921df569a3..32a1329448 100644 --- a/include/ck/utility/amd_buffer_addressing.hpp +++ b/include/ck/utility/amd_buffer_addressing.hpp @@ -847,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_PRINT(); + // CK_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 85be91a71e..f4ccf9d232 100644 --- a/include/ck/utility/amd_xdlops.hpp +++ b/include/ck/utility/amd_xdlops.hpp @@ -655,11 +655,11 @@ struct intrin_mfma_scale_f32_32x32x64f8f6f4<32, 32> } }; -template +template struct intrin_mfma_scale_f32_16x16x128f8f6f4; -template <> -struct intrin_mfma_scale_f32_16x16x128f8f6f4<16, 16> +template +struct intrin_mfma_scale_f32_16x16x128f8f6f4<16, 16, OpselA, OpselB> { template __device__ static void Run(const f8x32_t& reg_a, @@ -675,11 +675,11 @@ struct intrin_mfma_scale_f32_16x16x128f8f6f4<16, 16> reg_a, reg_b, reg_c.template AsType()[Number<0>{}], - 0, // cbsz {0 FP8 E4M3; 1 FP8 E5M2; 2 FP6 E2M3; 3 FP6 E3M2; 4 FP4 E2M1} - 0, // blgp - 0, // OPSEL + 0, // cbsz {0 FP8 E4M3; 1 FP8 E5M2; 2 FP6 E2M3; 3 FP6 E3M2; 4 FP4 E2M1} + 0, // blgp + OpselA, // OPSEL scale_a, - 0, // OPSEL + OpselB, // OPSEL scale_b); #else ignore = reg_a; @@ -704,11 +704,11 @@ struct intrin_mfma_scale_f32_16x16x128f8f6f4<16, 16> reg_a, reg_b, reg_c.template AsType()[Number<0>{}], - 1, // cbsz {0 FP8 E4M3; 1 FP8 E5M2; 2 FP6 E2M3; 3 FP6 E3M2; 4 FP4 E2M1} - 1, // blgp - 0, // OPSEL + 1, // cbsz {0 FP8 E4M3; 1 FP8 E5M2; 2 FP6 E2M3; 3 FP6 E3M2; 4 FP4 E2M1} + 1, // blgp + OpselA, // OPSEL scale_a, - 0, // OPSEL + OpselB, // OPSEL scale_b); #else ignore = reg_a; @@ -733,11 +733,11 @@ struct intrin_mfma_scale_f32_16x16x128f8f6f4<16, 16> reg_a, reg_b, reg_c.template AsType()[Number<0>{}], - 0, // cbsz {0 FP8 E4M3; 1 FP8 E5M2; 2 FP6 E2M3; 3 FP6 E3M2; 4 FP4 E2M1} - 1, // blgp - 0, // OPSEL + 0, // cbsz {0 FP8 E4M3; 1 FP8 E5M2; 2 FP6 E2M3; 3 FP6 E3M2; 4 FP4 E2M1} + 1, // blgp + OpselA, // OPSEL scale_a, - 0, // OPSEL + OpselB, // OPSEL scale_b); #else ignore = reg_a; @@ -756,6 +756,11 @@ struct intrin_mfma_scale_f32_16x16x128f8f6f4<16, 16> const int32_t scale_b, FloatC& reg_c) { + // if(get_thread_local_1d_id()){ + // printf("Scale A: %08x, Scale B: %08x\n", + // *reinterpret_cast(&scale_a), *reinterpret_cast(&scale_b)); + // } #if defined(__gfx950__) int32x4_t arg_a = bit_cast(reg_a); int32x4_t arg_b = bit_cast(reg_b); @@ -767,11 +772,11 @@ struct intrin_mfma_scale_f32_16x16x128f8f6f4<16, 16> arg_type{arg_a[0], arg_a[1], arg_a[2], arg_a[3], 0, 0, 0, 0}, arg_type{arg_b[0], arg_b[1], arg_b[2], arg_b[3], 0, 0, 0, 0}, reg_c.template AsType()[Number<0>{}], - 4, // cbsz - 4, // blgp - 0, // OPSEL + 4, // cbsz + 4, // blgp + OpselA, // OPSEL scale_a, - 0, // OPSEL + OpselB, // OPSEL scale_b); #else ignore = reg_a; diff --git a/library/include/ck/library/reference_tensor_operation/cpu/reference_mx_gemm.hpp b/library/include/ck/library/reference_tensor_operation/cpu/reference_mx_gemm.hpp index 89460f82d9..5faecab04f 100644 --- a/library/include/ck/library/reference_tensor_operation/cpu/reference_mx_gemm.hpp +++ b/library/include/ck/library/reference_tensor_operation/cpu/reference_mx_gemm.hpp @@ -86,9 +86,9 @@ struct ReferenceMXGemm : public device::BaseOperator Tensor b_k_n_scaled(HostTensorDescriptor({K, N}, {1, K})); // printf("K: %d\n", K); - for(size_t m = 0; m < M; m++) + for(int m = 0; m < M; m++) { - for(size_t k = 0; k < K; k++) + for(int k = 0; k < K; k++) { if constexpr(is_same_v) { @@ -105,14 +105,6 @@ struct ReferenceMXGemm : public device::BaseOperator a_m_k_scaled(m, k) = type_convert(a_f4_lo) * a_scale; a_m_k_scaled(m, k + 1) = type_convert(a_f4_hi) * a_scale; - if(m == 0 && 0) - { - printf("a_m_k_scaled(%zu, %zu): %f, %f\n", - m, - k, - a_m_k_scaled(m, k), - a_m_k_scaled(m, k + 1)); - } } else { @@ -123,9 +115,9 @@ struct ReferenceMXGemm : public device::BaseOperator } } - for(size_t n = 0; n < N; n++) + for(int n = 0; n < N; n++) { - for(size_t k = 0; k < K; k++) + for(int k = 0; k < K; k++) { if constexpr(is_same_v) { @@ -141,19 +133,6 @@ struct ReferenceMXGemm : public device::BaseOperator auto b_f4_hi = f4_t(b_pack.template unpack<>(Number<1>{})); b_k_n_scaled(k, n) = type_convert(b_f4_lo) * b_scale; b_k_n_scaled(k + 1, n) = type_convert(b_f4_hi) * b_scale; - if(n == 0 && 0) - { - printf("b_k_n(%zu, %zu): %2x\n", - n, - k, - *reinterpret_cast(&b_pack)); - // printf("b_k_n_scaled(%zu, %zu): %f, %f\n", - // n, - // k, - // b_k_n_scaled(k, n), - // b_k_n_scaled(k+1, n) - // ); - } } else {