ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-15 11:34:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

wip

Browse Source
This commit is contained in:
Ding, Yi
2025-05-07 08:17:07 +00:00
parent 8a0d659f92
commit d12b750bb0
16 changed files with 154 additions and 62 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
CMakeLists.txt
View File

@@ -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}")

4
cmake/EnableCompilerWarnings.cmake
View File

@@ -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})

4
example/67_gemm_microscaling/gemm_mx_common.hpp
View File

@@ -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<ADataType>{ck::type_convert<ADataType>(1.0f)}(a_m_k);
ck::utils::FillConstant<XDataType>{ck::type_convert<XDataType>(2.0f)}(a_m_k_scale);
ck::utils::FillConstant<BDataType>{ck::type_convert<BDataType>(0.5f)}(b_k_n);
ck::utils::FillConstant<XDataType>{ck::type_convert<XDataType>(1.0f)}(b_k_n_scale);
ck::utils::FillConstant<BDataType>{ck::type_convert<BDataType>(1.f)}(b_k_n);
ck::utils::FillConstant<XDataType>{ck::type_convert<XDataType>(2.0f)}(b_k_n_scale);
if(config.verbosity > 0)
{
std::cout << "Init A = {1}" << std::endl;

6
example/67_gemm_microscaling/gemm_mx_fp8.cpp
View File

@@ -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;

7
include/ck/config.h.in
View File

@@ -26,6 +26,13 @@
#ifndef CK_CONFIG_H_IN
#define CK_CONFIG_H_IN
template <auto ...val>
[[deprecated("Help function to print value")]]
inline constexpr void CK_TILE_PRINT() {}
template <typename ...type>
[[deprecated("Help function to print value")]]
inline constexpr void CK_TILE_PRINT() {}
// clang-format off
//
// DataType supports in the current CK build

2
include/ck/tensor_operation/gpu/block/blockwise_gemm_mx_pipeline_xdlops_base.hpp
View File

@@ -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;

5
include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_mx_selector.hpp
View File

@@ -13,8 +13,9 @@ namespace ck {
template <typename T>
static constexpr bool is_scale_mfma_data_type()
{
return is_same_v<T, f8_ocp_t> || is_same_v<T, bf8_ocp_t> || is_same_v<T, f6_t> ||
is_same_v<T, bf6_t> || is_same_v<T, f4_t>;
using U = element_type_t<T>;
return is_same_v<U, f8_ocp_t> || is_same_v<U, bf8_ocp_t> || is_same_v<U, f6_t> ||
is_same_v<U, bf6_t> || is_same_v<U, f4_t>;
}
/**

74
include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_v1_mx.hpp
View File

@@ -156,7 +156,8 @@ struct BlockwiseGemmXdlops_pipeline_v1_mx<BlockGemmPipelineScheduler::Intrawave,
KPerBlock / ScaleBlockSize; // How many mx-vectors per K block
//> 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<BlockGemmPipelineScheduler::Intrawave,
k * xdlops_gemm.KPerXdlops * (KPack / xdlops_gemm.K1PerXdlops);
static_for<0, MRepeat, 1>{}([&](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_k_step_chunk>{}),
a_block_buf,
a_thread_desc_,
make_tuple(m0, I0, k, Number<chunk * KThreadChunk>{}),
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_k_step_chunk>{}),
// a_block_buf,
// a_thread_desc_,
// make_tuple(m0, I0, k, Number<chunk * KThreadChunk>{}),
// 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_k_step_chunk>{}),
b_block_buf,
b_thread_desc_,
make_tuple(n0, I0, k, Number<chunk * KThreadChunk>{}),
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_k_step_chunk>{}),
b_block_buf,
b_thread_desc_,
make_tuple(n0, I0, k, Number<chunk * KThreadChunk>{}),
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<ComputeTypeA, KPack> a_thread_vec;
vector_type<ComputeTypeA, KPack> a_thread_vec; // = vec: pk_i4_t, 32
vector_type<ComputeTypeB, KPack> b_thread_vec;
static_for<0, KPack, 1>{}([&](auto ik) {
@@ -403,11 +406,16 @@ struct BlockwiseGemmXdlops_pipeline_v1_mx<BlockGemmPipelineScheduler::Intrawave,
b_scale_thread_vec.template AsType<BScaleDataType>()(s) =
b_scale_thread_buf[Number<b_scale_offset + s>{}];
});
CK_TILE_PRINT<xdlops_gemm.K1PerXdlops>();
CK_TILE_PRINT<decltype(xdlops_gemm)>();
using mfma_input_type_a =
typename vector_type<ComputeTypeA, xdlops_gemm.K1PerXdlops>::type;
typename vector_type<ComputeTypeA,
xdlops_gemm.K1PerXdlops / 2>::type;
// mfma input type = pk_f4_t, 32
CK_TILE_PRINT<mfma_input_type_a>();
using mfma_input_type_b =
typename vector_type<ComputeTypeB, xdlops_gemm.K1PerXdlops>::type;
typename vector_type<ComputeTypeB,
xdlops_gemm.K1PerXdlops / 2>::type;
constexpr index_t c_offset =
c_thread_desc_.CalculateOffset(make_tuple(m0, n0, 0));
@@ -505,7 +513,7 @@ struct BlockwiseGemmXdlops_pipeline_v1_mx<BlockGemmPipelineScheduler::Intrawave,
static_for<0, MRepeat, 1>{}([&](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<BlockGemmPipelineScheduler::Intrawave,
});
static_for<0, NRepeat, 1>{}([&](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<BlockGemmPipelineScheduler::Intrawave,
});
using mfma_input_type_a =
typename vector_type<ComputeTypeA, xdlops_gemm.K1PerXdlops>::type;
typename vector_type<ComputeTypeA, xdlops_gemm.K1PerXdlops / 2>::type;
using mfma_input_type_b =
typename vector_type<ComputeTypeB, xdlops_gemm.K1PerXdlops>::type;
typename vector_type<ComputeTypeB, xdlops_gemm.K1PerXdlops / 2>::type;
constexpr index_t c_offset =
c_thread_desc_.CalculateOffset(make_tuple(m0, n0, 0));

7
include/ck/tensor_operation/gpu/element/unary_element_wise_operation.hpp
View File

@@ -315,6 +315,13 @@ struct PassThrough
y = x;
}
template <>
__host__ __device__ void operator()<f4x2_pk_t, f4x2_pk_t>(f4x2_pk_t& y,
const f4x2_pk_t& x) const
{
y = x;
}
template <>
__host__ __device__ void operator()<double, double>(double& y, const double& x) const
{

16
include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_v3_mx.hpp
View File

@@ -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<ComputeTypeA,
MPerXdl,
NPerXdl,
ComputeTypeB,
is_single_rate_mfma,
is_scale_mfma>::selected_mfma.k_per_blk);
is_scale_mfma>::selected_mfma.k_per_blk /
2); // num of f4
using ThisThreadBlock = ThisThreadBlock<BlockSize>;
static constexpr index_t APackedSize = []() {
if constexpr(is_same_v<remove_cvref_t<ADataType>, pk_i4_t>)
if constexpr(is_same_v<remove_cvref_t<ADataType>, pk_i4_t> ||
is_same_v<remove_cvref_t<ADataType>, f4x2_pk_t>)
return 2;
else
return 1;
}();
static constexpr index_t BPackedSize = []() {
if constexpr(is_same_v<remove_cvref_t<BDataType>, pk_i4_t>)
if constexpr(is_same_v<remove_cvref_t<BDataType>, pk_i4_t> ||
is_same_v<remove_cvref_t<BDataType>, f4x2_pk_t>)
return 2;
else
return 1;
@@ -363,6 +366,9 @@ struct GridwiseGemmMX_xdl_cshuffle_v3
static_assert(!(is_same_v<remove_cvref_t<ADataType>, pk_i4_t> &&
GemmSpec != GemmSpecialization::Default),
"pk_i4_t does not support padding");
static_assert(!(is_same_v<remove_cvref_t<ADataType>, f4x2_pk_t> &&
GemmSpec != GemmSpecialization::Default),
"f4x2_pk_t does not support padding");
if constexpr(GemmSpec == GemmSpecialization::NKPadding ||
GemmSpec == GemmSpecialization::MNKPadding)

8
include/ck/tensor_operation/gpu/thread/threadwise_tensor_slice_transfer_v3r1.hpp
View File

@@ -69,7 +69,8 @@ struct ThreadwiseTensorSliceTransfer_v3r1
static constexpr auto I16 = Number<16>{};
static constexpr index_t PackedSize = []() {
if constexpr(is_same_v<remove_cvref_t<SrcData>, pk_i4_t>)
if constexpr(is_same_v<remove_cvref_t<SrcData>, pk_i4_t> ||
is_same_v<remove_cvref_t<SrcData>, 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<remove_cvref_t<SrcData>, pk_i4_t>)
if constexpr(is_same_v<remove_cvref_t<SrcData>, pk_i4_t> ||
is_same_v<remove_cvref_t<SrcData>, f4x2_pk_t>)
{
static_assert(is_same_v<remove_cvref_t<SrcData>, remove_cvref_t<DstData>>,
"SrcData != DstData");
@@ -444,6 +446,8 @@ struct ThreadwiseTensorSliceTransfer_v3r1
{
static_assert(!is_same_v<remove_cvref_t<SrcData>, pk_i4_t>,
"in-register transpose is not supported for pk_i4_t");
static_assert(!is_same_v<remove_cvref_t<SrcData>, 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_

13
include/ck/tensor_operation/gpu/warp/xdlops_gemm.hpp
View File

@@ -858,6 +858,7 @@ struct mfma_type<MfmaInstr::mfma_scale_f32_32x32x64f8f6f4>
const ScaleB& scale_b,
FloatC& reg_c) const
{
// CK_TILE_PRINT<FloatA>();
static_assert(scalar_type<ScaleA>::vector_size == 1, "Expect single scale at this point.");
static_assert(scalar_type<ScaleB>::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<f4_t, 32, 32, f4_t, false, true>()
{
return MfmaInstr::mfma_scale_f32_32x32x64f8f6f4;
}
template <>
constexpr auto GetMfma<f8_t, 16, 16>()
@@ -1201,10 +1207,10 @@ struct MfmaSelector
return MfmaInstr::mfma_f32_16x16x32bf8f8;
}
static constexpr auto selected_mfma = mfma_type<GetMfma<base_type,
static constexpr auto selected_mfma = mfma_type<GetMfma<element_type_t<base_type>,
MPerXdlops,
NPerXdlops,
additional_type,
element_type_t<additional_type>,
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

6
include/ck/utility/amd_buffer_addressing.hpp
View File

@@ -430,7 +430,9 @@ __device__ typename vector_type<T, N>::type amd_buffer_load_impl(int32x4_t src_w
(is_same<T, bf8_t>::value && (N == 1 || N == 2 || N == 4 || N == 8 || N == 16)) ||
(is_same<T, int8_t>::value && (N == 1 || N == 2 || N == 4 || N == 8 || N == 16)) ||
(is_same<T, uint8_t>::value && (N == 1 || N == 2 || N == 4 || N == 8 || N == 16)) ||
(is_same<T, pk_i4_t>::value && (N == 1 || N == 2 || N == 4 || N == 8 || N == 16)),
(is_same<T, pk_i4_t>::value && (N == 1 || N == 2 || N == 4 || N == 8 || N == 16)) ||
(is_same<T, f4x2_pk_t::type>::value &&
(N == 1 || N == 2 || N == 4 || N == 8 || N == 16)),
"wrong! not implemented");
using r_t = typename vector_type<T, N>::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<T, vector_t, scalar_t>();
vector_t tmp{amd_buffer_load_impl<scalar_t, vector_size, coherence>(
src_wave_buffer_resource, src_thread_addr_offset, 0)};
return src_thread_element_valid ? tmp : vector_t(0);

11
include/ck/utility/amd_xdlops.hpp
View File

@@ -749,11 +749,12 @@ struct intrin_mfma_scale_f32_16x16x128f8f6f4<16, 16>
}
template <class FloatC>
__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<int32x4_t>(reg_a);

45
include/ck/utility/data_type.hpp
View File

@@ -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 <index_t I>
__host__ __device__ inline type unpack(Number<I>) const
@@ -346,6 +346,18 @@ inline constexpr bool is_native_type()
is_same<T, f6_t>::value || is_same<T, bf6_t>::value;
}
template <typename T>
struct is_f8f6f4
{
static constexpr bool value =
std::is_same_v<T, f8_t> || std::is_same_v<T, bf8_t> || std::is_same_v<T, f6_t> ||
std::is_same_v<T, bf6_t> || std::is_same_v<T, f6x16_pk_t> ||
std::is_same_v<T, f6x32_pk_t> || std::is_same_v<T, bf6x16_pk_t> ||
std::is_same_v<T, bf6x32_pk_t> || std::is_same_v<T, f4_t> || std::is_same_v<T, f4x2_pk_t>;
};
template <typename T>
inline constexpr bool is_f8f6f4_v = is_f8f6f4<T>::value;
// scalar_type
template <typename TV>
struct scalar_type;
@@ -484,6 +496,35 @@ struct scalar_type<bool>
static constexpr index_t vector_size = 1;
};
template <typename T>
struct element_type
{
private:
static constexpr auto get_element_type()
{
using U = remove_cvref_t<T>;
if constexpr(std::is_same_v<U, pk_i4_t>)
return int4_t{};
else if constexpr(std::is_same_v<U, f4x2_pk_t>)
return f4_t{};
else if constexpr(std::is_same_v<U, f6x16_pk_t>)
return f6_t{};
else if constexpr(std::is_same_v<U, bf6x16_pk_t>)
return bf6_t{};
else if constexpr(std::is_same_v<U, f6x32_pk_t>)
return f6_t{};
else if constexpr(std::is_same_v<U, bf6x32_pk_t>)
return bf6_t{};
else
return T{};
}
public:
using type = decltype(get_element_type());
};
template <typename T>
using element_type_t = typename element_type<T>::type;
#if defined(_WIN32)
using int64_t = long long;
#else

6
include/ck/utility/dtype_vector.hpp
View File

@@ -1248,6 +1248,12 @@ struct nnvb_data_t_selector<pk_i4_t>
using type = pk_i4_t::type;
};
template <>
struct nnvb_data_t_selector<f4x2_pk_t>
{
using type = f4x2_pk_t::type;
};
template <typename T, index_t N>
struct non_native_vector_base<
T,