diff --git a/include/ck/utility/amd_wmma.hpp b/include/ck/utility/amd_wmma.hpp index e14c0d62a8..f6f6712b9a 100644 --- a/include/ck/utility/amd_wmma.hpp +++ b/include/ck/utility/amd_wmma.hpp @@ -143,6 +143,33 @@ struct intrin_wmma_i32_16x16x16_iu8_w32<16, 16, neg_a, neg_b, clamp> } }; +// src: iu8, dst: i32 +template +struct intrin_wmma_i32_16x16x16_iu16_w32; + +template +struct intrin_wmma_i32_16x16x16_iu16_w32<16, 16, neg_a, neg_b, clamp> +{ + template + __device__ static void Run(const int16x16_t& reg_a, const int16x16_t& reg_b, FloatC& reg_c) + { +#if defined(__gfx11__) + reg_c.template AsType()(Number<0>{}) = + __builtin_amdgcn_wmma_i32_16x16x16_iu16_w32( + neg_a, + bit_cast(reg_a), + neg_b, + bit_cast(reg_b), + reg_c.template AsType()[Number<0>{}], + clamp); +#else + ignore = reg_a; + ignore = reg_b; + ignore = reg_c; +#endif + } +}; + /********************************WAVE64 MODE***********************************************/ template @@ -263,6 +290,33 @@ struct intrin_wmma_i32_16x16x16_iu8_w64<16, 16, neg_a, neg_b, clamp> } }; +// src: iu16, dst: i32 +template +struct intrin_wmma_i32_16x16x16_iu16_w64; + +template +struct intrin_wmma_i32_16x16x16_iu16_w64<16, 16, neg_a, neg_b, clamp> +{ + template + __device__ static void Run(const int16x16_t& reg_a, const int16x16_t& reg_b, FloatC& reg_c) + { +#if defined(__gfx11__) + reg_c.template AsType()(Number<0>{}) = + __builtin_amdgcn_wmma_i32_16x16x16_iu8_w64( + neg_a, + bit_cast(reg_a), + neg_b, + bit_cast(reg_b), + reg_c.template AsType()[Number<0>{}], + clamp); +#else + ignore = reg_a; + ignore = reg_b; + ignore = reg_c; +#endif + } +}; + // gfx12 /********************************WAVE32 MODE***********************************************/ @@ -341,94 +395,25 @@ struct intrin_wmma_i32_16x16x16_iu8_w32_gfx12<16, 16, neg_a, neg_b, clamp> } }; -// src: f8, f8, dst: fp32 -template -struct intrin_wmma_f32_16x16x16_f8f8_w32_gfx12; +// src: iu16, dst: i32 +template +struct intrin_wmma_i32_16x16x16_iu16_w32_gfx12; -template <> -struct intrin_wmma_f32_16x16x16_f8f8_w32_gfx12<16, 16> +template +struct intrin_wmma_i32_16x16x16_iu16_w32_gfx12<16, 16, neg_a, neg_b, clamp> { template - __device__ static void Run(const f8x8_t& reg_a, const f8x8_t& reg_b, FloatC& reg_c) + __device__ static void Run(const int16x8_t& reg_a, const int16x8_t& reg_b, FloatC& reg_c) { #if defined(__gfx12__) - reg_c.template AsType()(Number<0>{}) = - __builtin_amdgcn_wmma_f32_16x16x16_fp8_fp8_w32_gfx12( - bit_cast(reg_a), - bit_cast(reg_b), - reg_c.template AsType()[Number<0>{}]); -#else - ignore = reg_a; - ignore = reg_b; - ignore = reg_c; -#endif - } -}; - -// src: f8, bf8, dst: fp32 -template -struct intrin_wmma_f32_16x16x16_f8bf8_w32_gfx12; - -template <> -struct intrin_wmma_f32_16x16x16_f8bf8_w32_gfx12<16, 16> -{ - template - __device__ static void Run(const f8x8_t& reg_a, const bf8x8_t& reg_b, FloatC& reg_c) - { -#if defined(__gfx12__) - reg_c.template AsType()(Number<0>{}) = - __builtin_amdgcn_wmma_f32_16x16x16_fp8_bf8_w32_gfx12( - bit_cast(reg_a), - bit_cast(reg_b), - reg_c.template AsType()[Number<0>{}]); -#else - ignore = reg_a; - ignore = reg_b; - ignore = reg_c; -#endif - } -}; - -// src: bf8, f8, dst: fp32 -template -struct intrin_wmma_f32_16x16x16_bf8f8_w32_gfx12; - -template <> -struct intrin_wmma_f32_16x16x16_bf8f8_w32_gfx12<16, 16> -{ - template - __device__ static void Run(const bf8x8_t& reg_a, const f8x8_t& reg_b, FloatC& reg_c) - { -#if defined(__gfx12__) - reg_c.template AsType()(Number<0>{}) = - __builtin_amdgcn_wmma_f32_16x16x16_bf8_fp8_w32_gfx12( - bit_cast(reg_a), - bit_cast(reg_b), - reg_c.template AsType()[Number<0>{}]); -#else - ignore = reg_a; - ignore = reg_b; - ignore = reg_c; -#endif - } -}; - -// src: bf8, bf8, dst: fp32 -template -struct intrin_wmma_f32_16x16x16_bf8bf8_w32_gfx12; - -template <> -struct intrin_wmma_f32_16x16x16_bf8bf8_w32_gfx12<16, 16> -{ - template - __device__ static void Run(const bf8x8_t& reg_a, const bf8x8_t& reg_b, FloatC& reg_c) - { -#if defined(__gfx12__) - reg_c.template AsType()(Number<0>{}) = - __builtin_amdgcn_wmma_f32_16x16x16_bf8_bf8_w32_gfx12( - bit_cast(reg_a), - bit_cast(reg_b), - reg_c.template AsType()[Number<0>{}]); + reg_c.template AsType()(Number<0>{}) = + __builtin_amdgcn_wmma_i32_16x16x16_iu8_w32_gfx12( + neg_a, + bit_cast(reg_a), + neg_b, + bit_cast(reg_b), + reg_c.template AsType()[Number<0>{}], + clamp); #else ignore = reg_a; ignore = reg_b; diff --git a/include/ck/utility/dtype_vector.hpp b/include/ck/utility/dtype_vector.hpp index 4a91054fd9..3342e43a48 100644 --- a/include/ck/utility/dtype_vector.hpp +++ b/include/ck/utility/dtype_vector.hpp @@ -2093,6 +2093,8 @@ struct vector_type()>> } }; +using int64_t = long; + // fp32 using float2_t = typename vector_type::type; using float4_t = typename vector_type::type; diff --git a/library/src/tensor_operation_instance/gpu/gemm_add/device_gemm_add_wmma_c_shuffle_bf16_f16_bf16_bf16_mk_kn_mn_mn_instance.cpp b/library/src/tensor_operation_instance/gpu/gemm_add/device_gemm_add_wmma_c_shuffle_bf16_f16_bf16_bf16_mk_kn_mn_mn_instance.cpp index 50c90e781d..32d5ae580f 100644 --- a/library/src/tensor_operation_instance/gpu/gemm_add/device_gemm_add_wmma_c_shuffle_bf16_f16_bf16_bf16_mk_kn_mn_mn_instance.cpp +++ b/library/src/tensor_operation_instance/gpu/gemm_add/device_gemm_add_wmma_c_shuffle_bf16_f16_bf16_bf16_mk_kn_mn_mn_instance.cpp @@ -1,5 +1,5 @@ // SPDX-License-Identifier: MIT -// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved. +// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. #include "ck/library/tensor_operation_instance/add_device_operation_instance.hpp" #include "ck/library/tensor_operation_instance/device_operation_instance_factory.hpp" @@ -28,11 +28,11 @@ using device_gemm_add_wmma_c_shuffle_bf16_f16_bf16_bf16_mk_kn_mn_mn_generic_inst //################################| | | | | | | | | | | Operation| Operation| Operation| | | | | | | | | | | | Lengths_K0_M_K1| ArrangeOrder| | | PerVector| PerVector_K1| | Lengths_K0_N_K1| ArrangeOrder| | | PerVector| PerVector_K1| | PerShuffle| PerShuffle| _NBlock_NWaveNPerXdl| _NWaveNPerXdl| //################################| | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | // TODO: these template variables need to be adjusted - DeviceGemmMultipleD_Wmma_CShuffle< Row, Row, Row_Tuple, Row, BF16, BF16, I32, F32, BF16_Tuple, BF16, PassThrough, PassThrough, Add, GemmMNKPadding, 1, 256, 64, 128, 32, 8, 16, 16, 1, 2, S<4, 64, 1>, S<0, 2, 1>, S<1, 0, 2>, 2, 1, 8, 1, S<4, 64, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 1, 8, 1, 1, 1, S<1, 32, 1, 8>, 1> + DeviceGemmMultipleD_Wmma_CShuffle< Row, Row, Row_Tuple, Row, BF16, BF16, F32, F32, BF16_Tuple, BF16, PassThrough, PassThrough, Add, GemmMNKPadding, 1, 256, 128, 128, 64, 8, 16, 16, 4, 2, S<4, 64, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 2, 8, 1, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, 1, 2, S<1, 32, 1, 8>, 8> // clang-format on >; -using device_gemm_add_wmma_c_shuffle_bf16_f16_bf16_bf16_mk_kn_mn_mn_instances = std::tuple< +using add_device_gemm_add_wmma_c_shuffle_bf16_f16_bf16_bf16_mk_kn_mn_mn_instances = std::tuple< // clang-format off // M/N/K padding //################################| A| B| Ds| E| AData| BData| AccData| CShuffle| DsData| EData| A| B| CDE| GEMM| Prefetch| Block| MPer| NPer| K0Per| K1| MPer| NPer| MRepeat| NRepeat| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockLds| BBlockTransfer| BBlockTransfer| BBlockTransfer| BlockTransfer| BBlockTransfer| BBlockTransfer| BBlockLds| CShuffle| CShuffle| CBlockTransferClusterLengths| CBlockTransfer| @@ -40,13 +40,13 @@ using device_gemm_add_wmma_c_shuffle_bf16_f16_bf16_bf16_mk_kn_mn_mn_instances = //################################| | | | | | | | | | | Operation| Operation| Operation| | | | | | | | | | | | Lengths_K0_M_K1| ArrangeOrder| | | PerVector| PerVector_K1| | Lengths_K0_N_K1| ArrangeOrder| | | PerVector| PerVector_K1| | PerShuffle| PerShuffle| _NBlock_NWaveNPerXdl| _NWaveNPerXdl| //################################| | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | // TODO: these template variables need to be adjusted - DeviceGemmMultipleD_Wmma_CShuffle< Row, Row, Row_Tuple, Row, BF16, BF16, F32, F32, BF16_Tuple, BF16, PassThrough, PassThrough, Add, GemmMNKPadding, 1, 256, 16, 128, 32, 8, 16, 16, 1, 2, S<4, 16, 4>, S<1, 0, 2>, S<1, 0, 2>, 2, 2, 2, 1, S<4, 64, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 2, 8, 1, 1, 1, S<1, 16, 1, 4>, 1, LoopScheduler::Default, PipelineVersion::v1>, - DeviceGemmMultipleD_Wmma_CShuffle< Row, Row, Row_Tuple, Row, BF16, BF16, F32, F32, BF16_Tuple, BF16, PassThrough, PassThrough, Add, GemmMNKPadding, 1, 64, 16, 16, 32, 8, 16, 16, 1, 1, S<4, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S<4, 16, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 1, 1, 1, 1, 1, S<1, 16, 1, 4>, 1, LoopScheduler::Default, PipelineVersion::v1>, - DeviceGemmMultipleD_Wmma_CShuffle< Row, Row, Row_Tuple, Row, BF16, BF16, F32, F32, BF16_Tuple, BF16, PassThrough, PassThrough, Add, GemmMNKPadding, 1, 64, 16, 16, 64, 8, 16, 16, 1, 1, S<4, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S<4, 16, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 1, 1, 1, 1, 1, S<1, 16, 1, 4>, 1, LoopScheduler::Default, PipelineVersion::v1> + DeviceGemmMultipleD_Wmma_CShuffle< Row, Row, Row_Tuple, Row, BF16, BF16, F32, F32, BF16_Tuple, BF16, PassThrough, PassThrough, Add, GemmMNKPadding, 1, 256, 128, 128, 64, 8, 16, 16, 4, 2, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S<4, 64, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 2, 8, 1, 1, 2, S<1, 32, 1, 8>, 8, LoopScheduler::Default, PipelineVersion::v1>, + DeviceGemmMultipleD_Wmma_CShuffle< Row, Row, Row_Tuple, Row, BF16, BF16, F32, F32, BF16_Tuple, BF16, PassThrough, PassThrough, Add, GemmMNKPadding, 1, 128, 64, 64, 64, 8, 16, 16, 2, 2, S<4, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S<4, 32, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 2, 8, 1, 1, 2, S<1, 32, 1, 4>, 8, LoopScheduler::Default, PipelineVersion::v1>, + DeviceGemmMultipleD_Wmma_CShuffle< Row, Row, Row_Tuple, Row, BF16, BF16, F32, F32, BF16_Tuple, BF16, PassThrough, PassThrough, Add, GemmMNKPadding, 1, 128, 64, 64, 64, 8, 16, 16, 2, 2, S<4, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S<4, 32, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 2, 8, 1, 1, 2, S<1, 32, 1, 4>, 8, LoopScheduler::Default, PipelineVersion::v1> // clang-format on >; -void add_device_gemm_add_wmma_c_shuffle_bf16_f16_bf16_bf16_mk_kn_mn_mn_instances( +void device_gemm_add_wmma_c_shuffle_bf16_f16_bf16_bf16_mk_kn_mn_mn_instances( std::vector