[rocm-libraries] ROCm/rocm-libraries#8535 (commit a0f47eb)

[CK Tile] EightWaves pipeline int8 support ## Motivation EightWaves pipeline currently is supporting only FP types ## Technical Details - Enable 16x16x64 int8 instruction for gfx950 in dispatcher - Enable int8 in EightWaves pipeline - Add tests - Fix bug in `warp_gemm_attribute_mfma_impl.hpp` ## Test Plan Tests have been added for int8 GEMM using EightWaves pipeline ## Submission Checklist - [x] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
2026-07-07 07:37:01 +00:00 · 2026-06-18 12:59:59 +00:00
parent 60b276647b
commit 1762eaeaec
7 changed files with 77 additions and 29 deletions
--- a/include/ck_tile/ops/gemm/block/block_gemm_areg_breg_creg_eight_waves_v1.hpp
+++ b/include/ck_tile/ops/gemm/block/block_gemm_areg_breg_creg_eight_waves_v1.hpp
@@ -103,8 +103,6 @@ struct BlockGemmARegBRegCRegEightWavesV1
    using BWarpTensor = typename WarpGemm::BWarpTensor;
    using CWarpTensor = typename WarpGemm::CWarpTensor;

-    static_assert(std::is_same_v<typename WarpGemm::CDataType, float>);
-
    static constexpr auto a_warp_y_lengths =
        to_sequence(AWarpDstr{}.get_ys_to_d_descriptor().get_lengths());
    static constexpr auto b_warp_y_lengths =
--- a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_async_eight_waves_policy.hpp
+++ b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_async_eight_waves_policy.hpp
@@ -36,11 +36,12 @@ struct GemmPipelineAgBgCrCompAsyncEightWavesPolicy
    static_assert(std::is_same_v<BLayout, ck_tile::tensor_layout::gemm::ColumnMajor>,
                  "BLayout must be ColumnMajor!");
    static_assert(
-        is_any_of<AComputeDataType, fp8_t, bf8_t, pk_fp4_t, pk_fp6x16_t, fp16_t, bf16_t>::value);
+        is_any_of<AComputeDataType, fp8_t, bf8_t, pk_fp4_t, pk_fp6x16_t, fp16_t, bf16_t, int8_t>::
+            value);
    static_assert(
-        is_any_of<BComputeDataType, fp8_t, bf8_t, pk_fp4_t, pk_fp6x16_t, fp16_t, bf16_t>::value);
+        is_any_of<BComputeDataType, fp8_t, bf8_t, pk_fp4_t, pk_fp6x16_t, fp16_t, bf16_t, int8_t>::
+            value);
    static_assert(std::is_same_v<AComputeDataType, BComputeDataType>);
-    static_assert(std::is_same_v<CDataType, float>);

    static constexpr auto WGAccess =
        std::is_same_v<ComputeDataType, fp8_t> || std::is_same_v<ComputeDataType, bf8_t>
--- a/include/ck_tile/ops/gemm/warp/warp_gemm.hpp
+++ b/include/ck_tile/ops/gemm/warp/warp_gemm.hpp
@@ -511,6 +511,19 @@ using WarpGemmMfma_i32_32x32x16_i8_i8_CTransposed =
 using WarpGemmMfma_i32_16x16x32_i8_i8 = WarpGemmImpl<
    WarpGemmAttributeMfma<WarpGemmAttributeMfmaImpl_i32_16x16x32_i8<WGAttrCtlEnum::Default_>>>;

+template <WGAttrNumAccessEnum AttrNumAccessA = WGAttrNumAccessEnum::Single,
+          WGAttrNumAccessEnum AttrNumAccessB = AttrNumAccessA>
+using WarpGemmMfma_i32_16x16x64_i8_i8 = WarpGemmImpl<
+    WarpGemmAttributeMfma<WarpGemmAttributeMfmaImpl_i32_16x16x64_i8<WGAttrCtlEnum::Default_>,
+                          AttrNumAccessA,
+                          AttrNumAccessB>>;
+
+template <WGAttrNumAccessEnum AttrNumAccess>
+using WarpGemmMfma_i32_16x16x64_i8_i8_CTransposed =
+    WarpGemmImpl<WarpGemmAttributeMfmaTransposedCDistribution<
+        WarpGemmAttributeMfmaImpl_i32_16x16x64_i8<WGAttrCtlEnum::Default_>,
+        AttrNumAccess>>;
+
 using WarpGemmMfma_i32_16x16x32_i8_i8_CTransposed =
    WarpGemmImpl<WarpGemmAttributeMfmaTransposedCDistribution<
        WarpGemmAttributeMfmaImpl_i32_16x16x32_i8<WGAttrCtlEnum::Default_>>>;
--- a/include/ck_tile/ops/gemm/warp/warp_gemm_attribute_mfma_impl.hpp
+++ b/include/ck_tile/ops/gemm/warp/warp_gemm_attribute_mfma_impl.hpp
@@ -2112,9 +2112,8 @@ struct WarpGemmAttributeMfmaImpl_i32_16x16x64_i8
        DISPATCH_MFMA_CTRL_("v_mfma_i32_16x16x64_i8", Ctrl)
        else
        {
-#if defined(__gfx95__)
-            c_vec = __builtin_amdgcn_mfma_i32_16x16x64_i8(
-                bit_cast<int64_t>(a_vec), bit_cast<int64_t>(b_vec), c_vec, 0, 0, 0);
+#if defined(__gfx950__)
+            c_vec = __builtin_amdgcn_mfma_i32_16x16x64_i8(a_vec, b_vec, c_vec, 0, 0, 0);
 #else
            ck_tile::ignore = c_vec;
            ck_tile::ignore = a_vec;
--- a/include/ck_tile/ops/gemm/warp/warp_gemm_dispatcher.hpp
+++ b/include/ck_tile/ops/gemm/warp/warp_gemm_dispatcher.hpp
@@ -288,7 +288,20 @@ template<> struct Dispatcher<int8_t, int8_t, int32_t, 16, 16, 32, false> { using
 template<> struct Dispatcher<int8_t, int8_t, int32_t, 16, 16, 32,  true> { using Type = WarpGemmMfma_i32_16x16x32_i8_i8_CTransposed; };
 // WMMA cases
 template<bool TransposeC, WGAttrNumAccessEnum AttrNumAccess> struct Dispatcher<int8_t, int8_t, int32_t, 16, 16, 16, TransposeC, false, false, AttrNumAccess, AttrNumAccess> : WmmaTag { using Type = WarpGemmWmma_i32_16x16x16_i8_i8<TransposeC>; };
+#if defined(__gfx125__)
 template<bool TransposeC, WGAttrNumAccessEnum AttrNumAccess> struct Dispatcher<int8_t, int8_t, int32_t, 16, 16, 64, TransposeC, false, false, AttrNumAccess, AttrNumAccess> : WmmaTag { using Type = WarpGemmWmma_i32_16x16x64_i8_i8<TransposeC>; };
+#else
+template<WGAttrNumAccessEnum AttrNumAccessA, WGAttrNumAccessEnum AttrNumAccessB>
+struct Dispatcher<int8_t, int8_t, int32_t, 16, 16, 64, false, false, false, AttrNumAccessA, AttrNumAccessB, false, std::enable_if_t<AttrNumAccessA != AttrNumAccessB>>
+{ using Type = WarpGemmMfma_i32_16x16x64_i8_i8<AttrNumAccessA, AttrNumAccessB>; };
+template<WGAttrNumAccessEnum AttrNumAccess>
+struct Dispatcher<int8_t, int8_t, int32_t, 16, 16, 64, false, false, false, AttrNumAccess, AttrNumAccess, false, std::enable_if_t<AttrNumAccess != EDefault>>
+{ using Type = WarpGemmMfma_i32_16x16x64_i8_i8<AttrNumAccess>; };
+
+template<WGAttrNumAccessEnum AttrNumAccess>
+struct Dispatcher<int8_t, int8_t, int32_t, 16, 16, 64, true, false, false, AttrNumAccess, AttrNumAccess, false, std::enable_if_t<AttrNumAccess != EDefault>>
+{ using Type = WarpGemmMfma_i32_16x16x64_i8_i8_CTransposed<AttrNumAccess>; };
+#endif
 template<bool TransposeC, WGAttrNumAccessEnum AttrNumAccess> struct Dispatcher<uint8_t, uint8_t, int32_t, 16, 16, 64, TransposeC, false, false, AttrNumAccess, AttrNumAccess> : WmmaTag { using Type = WarpGemmWmma_i32_16x16x64_u8_u8<TransposeC>; };

 template <typename AType, typename BType, typename AccType,
--- a/test/ck_tile/gemm/test_gemm_pipeline_kernel_types.hpp
+++ b/test/ck_tile/gemm/test_gemm_pipeline_kernel_types.hpp
@@ -327,20 +327,36 @@ using CompAsyncEightWavesConfig4Bit = std::tuple<ALayout,
                                                 CompAsyncEightWaves>;

 template <typename ALayout, typename BLayout, typename CLayout, typename InputType>
-using CompAsyncEightWavesConfig8Bit = std::tuple<ALayout,
-                                                 BLayout,
-                                                 CLayout,
-                                                 InputType, // AType
-                                                 InputType, // BType
-                                                 F32,       // AccType
-                                                 F16,       // OutputType
-                                                 I128,      // MBlockTileSize
-                                                 I256,      // NBlockTileSize
-                                                 I128,      // KBlockTileSize
-                                                 I16,       // MWarpTileSize
-                                                 I16,       // NWarpTileSize
-                                                 Intrawave,
-                                                 CompAsyncEightWaves>;
+using CompAsyncEightWavesConfig8BitFP = std::tuple<ALayout,
+                                                   BLayout,
+                                                   CLayout,
+                                                   InputType, // AType
+                                                   InputType, // BType
+                                                   F32,       // AccType
+                                                   F16,       // OutputType
+                                                   I128,      // MBlockTileSize
+                                                   I256,      // NBlockTileSize
+                                                   I128,      // KBlockTileSize
+                                                   I16,       // MWarpTileSize
+                                                   I16,       // NWarpTileSize
+                                                   Intrawave,
+                                                   CompAsyncEightWaves>;
+
+template <typename ALayout, typename BLayout, typename CLayout, typename InputType>
+using CompAsyncEightWavesConfig8BitINT = std::tuple<ALayout,
+                                                    BLayout,
+                                                    CLayout,
+                                                    InputType, // AType
+                                                    InputType, // BType
+                                                    INT32,     // AccType
+                                                    INT32,     // OutputType
+                                                    I128,      // MBlockTileSize
+                                                    I256,      // NBlockTileSize
+                                                    I128,      // KBlockTileSize
+                                                    I16,       // MWarpTileSize
+                                                    I16,       // NWarpTileSize
+                                                    Intrawave,
+                                                    CompAsyncEightWaves>;

 template <typename ALayout, typename BLayout, typename CLayout, typename InputType>
 using CompAsyncEightWavesConfig16Bit = std::tuple<ALayout,
@@ -371,8 +387,9 @@ using KernelTypesCompAsync16x16x128 = ::testing::Types<CompAsyncConfig16x16x128<
                                                       CompAsyncConfig16x16x128<Row, Col, Row, F8>>;

 using KernelTypesCompAsyncEightWaves =
-    ::testing::Types<CompAsyncEightWavesConfig8Bit<Row, Col, Row, F8>,
-                     CompAsyncEightWavesConfig8Bit<Row, Col, Row, BF8>,
+    ::testing::Types<CompAsyncEightWavesConfig8BitINT<Row, Col, Row, INT8>,
+                     CompAsyncEightWavesConfig8BitFP<Row, Col, Row, F8>,
+                     CompAsyncEightWavesConfig8BitFP<Row, Col, Row, BF8>,
                     CompAsyncEightWavesConfig4Bit<Row, Col, Row, F4>,
                     CompAsyncEightWavesConfig16Bit<Row, Col, Row, F16>,
                     CompAsyncEightWavesConfig16Bit<Row, Col, Row, BF16>>;
--- a/test/ck_tile/gemm/test_gemm_pipeline_util.hpp
+++ b/test/ck_tile/gemm/test_gemm_pipeline_util.hpp
@@ -63,11 +63,18 @@ constexpr ck_tile::index_t get_k_warp_tile()
    return 16;
 #endif
 #else
-    if constexpr(PipelineType == GemmPipelineType::CompAsyncEightWaves && sizeof(PrecType) == 1)
-        return 128;
-    else if constexpr(PipelineType == GemmPipelineType::CompAsyncEightWaves &&
-                      sizeof(PrecType) == 2)
-        return 32;
+    if constexpr(PipelineType == GemmPipelineType::CompAsyncEightWaves)
+        if constexpr(std::is_same_v<PrecType, ck_tile::int8_t>)
+            return 64;
+        else if constexpr(ck_tile::is_any_of<PrecType,
+                                             ck_tile::fp8_t,
+                                             ck_tile::pk_fp4_t,
+                                             ck_tile::bf8_t>::value)
+            return 128;
+        else if constexpr(ck_tile::is_any_of<PrecType, ck_tile::fp16_t, ck_tile::bf16_t>::value)
+            return 32;
+        else
+            static_assert(false, "CompAsyncEightWaves: unsupported type");
    // CompAsyncConfig16x16x128
    else if constexpr(PipelineType == GemmPipelineType::CompAsync && M_Warp_Tile == 16)
        return 128;