From 5ef30005fa5f5bef43910155719b16dd01fa3a92 Mon Sep 17 00:00:00 2001 From: wozeparrot Date: Mon, 8 Jun 2026 21:53:25 -0400 Subject: [PATCH] update hipkittens (#16544) --- extra/thunder/amd/fa_bwd_causal.cpp | 26 +- extra/thunder/amd/include/common/base_ops.cuh | 53 ++ .../thunder/amd/include/common/base_types.cuh | 101 +++- extra/thunder/amd/include/common/macros.cuh | 560 +++++++++++++++--- extra/thunder/amd/include/common/util.cuh | 4 + .../tile/assembly/global_to_register.cuh | 51 +- .../ops/warp/memory/tile/global_to_shared.cuh | 5 +- .../ops/warp/register/tile/assembly/maps.cuh | 60 +- .../ops/warp/register/tile/assembly/mma.cuh | 66 ++- .../include/ops/warp/register/tile/maps.cuh | 22 + .../include/ops/warp/register/tile/mma.cuh | 118 ++++ .../include/ops/warp/register/tile/scales.cuh | 74 +++ .../include/ops/warp/register/tile/tile.cuh | 1 + extra/thunder/amd/include/types/global/gl.cuh | 8 +- .../amd/include/types/register/art_base.cuh | 4 +- .../amd/include/types/register/rt_shape.cuh | 14 +- extra/thunder/amd/include/types/shared/sv.cuh | 1 + 17 files changed, 1026 insertions(+), 142 deletions(-) create mode 100644 extra/thunder/amd/include/ops/warp/register/tile/scales.cuh diff --git a/extra/thunder/amd/fa_bwd_causal.cpp b/extra/thunder/amd/fa_bwd_causal.cpp index 2de20495cf..558ffcee43 100644 --- a/extra/thunder/amd/fa_bwd_causal.cpp +++ b/extra/thunder/amd/fa_bwd_causal.cpp @@ -133,7 +133,7 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr constexpr int neg_inf_v = 29; // Move -inf to VGPR neg_inf_v kittens::macros::clobber_gpr(); - kittens::macros::v_mov_b32(0xff800000); + kittens::macros::v_mov_b32_up2p(0xff800000); art P_ij; // 16 registers art dP_ij; // 16 registers @@ -330,7 +330,7 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr load<0, 3>(dO_i_col, subtile_inplace(dO_i_smem[tic][0], {0, 0}), dO_i_col_addr); mma_ABt<0, 3, 3>(P_ij, Q_i, K_j, P_ij); // Dot slice 0 - kittens::macros::v_mov_b32(0xff800000); if constexpr (causal) { + kittens::macros::v_mov_b32_up2p(0xff800000); if constexpr (causal) { // If the query position is less than the key position, set P_ij to -inf if (q_pos < k_pos) { mov(P_ij); @@ -588,7 +588,7 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr load<0, 3>(dO_i_col, subtile_inplace(dO_i_smem[tic][0], {0, 0}), dO_i_col_addr); mma_ABt<0, 3, 3>(P_ij, Q_i, K_j, P_ij); // Dot slice 1 - kittens::macros::v_mov_b32(0xff800000); if constexpr (causal) { + kittens::macros::v_mov_b32_up2p(0xff800000); if constexpr (causal) { // If the query position is less than the key position, set P_ij to -inf if (q_pos < k_pos) { mov(P_ij); @@ -845,7 +845,7 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr load<0, 3>(dO_i_col, subtile_inplace(dO_i_smem[tic][0], {0, 0}), dO_i_col_addr); mma_ABt<0, 3, 3>(P_ij, Q_i, K_j, P_ij); // Dot slice 2 - kittens::macros::v_mov_b32(0xff800000); if constexpr (causal) { + kittens::macros::v_mov_b32_up2p(0xff800000); if constexpr (causal) { // If the query position is less than the key position, set P_ij to -inf if (q_pos < k_pos) { mov(P_ij); @@ -1101,7 +1101,7 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr load<0, 3>(dO_i_col, subtile_inplace(dO_i_smem[tic][0], {0, 0}), dO_i_col_addr); mma_ABt<0, 3, 3>(P_ij, Q_i, K_j, P_ij); // Dot slice 3 - kittens::macros::v_mov_b32(0xff800000); if constexpr (causal) { + kittens::macros::v_mov_b32_up2p(0xff800000); if constexpr (causal) { // If the query position is less than the key position, set P_ij to -inf if (q_pos < k_pos) { mov(P_ij); @@ -1371,7 +1371,7 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr load<0, 3>(dO_i_col, subtile_inplace(dO_i_smem[tic][0], {0, 0}), dO_i_col_addr); mma_ABt<0, 3, 3>(P_ij, Q_i, K_j, P_ij); // Dot slice 0 - kittens::macros::v_mov_b32(0xff800000); if constexpr (causal) { + kittens::macros::v_mov_b32_up2p(0xff800000); if constexpr (causal) { // If the query position is less than the key position, set P_ij to -inf if (q_pos < k_pos) { mov(P_ij); @@ -1632,7 +1632,7 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr load<0, 3>(dO_i_col, subtile_inplace(dO_i_smem[tic][0], {0, 0}), dO_i_col_addr); mma_ABt<0, 3, 3>(P_ij, Q_i, K_j, P_ij); // Dot slice 1 - kittens::macros::v_mov_b32(0xff800000); if constexpr (causal) { + kittens::macros::v_mov_b32_up2p(0xff800000); if constexpr (causal) { // If the query position is less than the key position, set P_ij to -inf if (q_pos < k_pos) { mov(P_ij); @@ -1889,7 +1889,7 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr load<0, 3>(dO_i_col, subtile_inplace(dO_i_smem[tic][0], {0, 0}), dO_i_col_addr); mma_ABt<0, 3, 3>(P_ij, Q_i, K_j, P_ij); // Dot slice 2 - kittens::macros::v_mov_b32(0xff800000); if constexpr (causal) { + kittens::macros::v_mov_b32_up2p(0xff800000); if constexpr (causal) { // If the query position is less than the key position, set P_ij to -inf if (q_pos < k_pos) { mov(P_ij); @@ -2145,7 +2145,7 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr load<0, 3>(dO_i_col, subtile_inplace(dO_i_smem[tic][0], {0, 0}), dO_i_col_addr); mma_ABt<0, 3, 3>(P_ij, Q_i, K_j, P_ij); // Dot slice 3 - kittens::macros::v_mov_b32(0xff800000); if constexpr (causal) { + kittens::macros::v_mov_b32_up2p(0xff800000); if constexpr (causal) { // If the query position is less than the key position, set P_ij to -inf if (q_pos < k_pos) { mov(P_ij); @@ -2410,7 +2410,7 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr load<0, 3>(dO_i_col, subtile_inplace(dO_i_smem[tic][0], {0, 0}), dO_i_col_addr); mma_ABt<0, 3, 3>(P_ij, Q_i, K_j, P_ij); // Dot slice 0 - kittens::macros::v_mov_b32(0xff800000); if constexpr (causal) { + kittens::macros::v_mov_b32_up2p(0xff800000); if constexpr (causal) { // If the query position is less than the key position, set P_ij to -inf if (q_pos < k_pos) { mov(P_ij); @@ -2671,7 +2671,7 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr load<0, 3>(dO_i_col, subtile_inplace(dO_i_smem[tic][0], {0, 0}), dO_i_col_addr); mma_ABt<0, 3, 3>(P_ij, Q_i, K_j, P_ij); // Dot slice 1 - kittens::macros::v_mov_b32(0xff800000); if constexpr (causal) { + kittens::macros::v_mov_b32_up2p(0xff800000); if constexpr (causal) { // If the query position is less than the key position, set P_ij to -inf if (q_pos < k_pos) { mov(P_ij); @@ -2927,7 +2927,7 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr load<0, 3>(dO_i_col, subtile_inplace(dO_i_smem[tic][0], {0, 0}), dO_i_col_addr); mma_ABt<0, 3, 3>(P_ij, Q_i, K_j, P_ij); // Dot slice 2 - kittens::macros::v_mov_b32(0xff800000); if constexpr (causal) { + kittens::macros::v_mov_b32_up2p(0xff800000); if constexpr (causal) { // If the query position is less than the key position, set P_ij to -inf if (q_pos < k_pos) { mov(P_ij); @@ -3183,7 +3183,7 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr load<0, 3>(dO_i_col, subtile_inplace(dO_i_smem[tic][0], {0, 0}), dO_i_col_addr); mma_ABt<0, 3, 3>(P_ij, Q_i, K_j, P_ij); // Dot slice 3 - kittens::macros::v_mov_b32(0xff800000); if constexpr (causal) { + kittens::macros::v_mov_b32_up2p(0xff800000); if constexpr (causal) { // If the query position is less than the key position, set P_ij to -inf if (q_pos < k_pos) { mov(P_ij); diff --git a/extra/thunder/amd/include/common/base_ops.cuh b/extra/thunder/amd/include/common/base_ops.cuh index d72ab47e3d..10d2d0e9d1 100644 --- a/extra/thunder/amd/include/common/base_ops.cuh +++ b/extra/thunder/amd/include/common/base_ops.cuh @@ -216,6 +216,59 @@ template<> __device__ inline bf16_2 relu::op(const bf16_2 &x) { return _ template<> __device__ inline half relu::op (const half &x ) { return __hmax(x, base_types::constants::zero()); } template<> __device__ inline half_2 relu::op(const half_2 &x) { return half_2{__hmax(x.x, base_types::constants::zero()), __hmax(x.y, base_types::constants::zero())}; } + + +constexpr float SQRT_2_OVER_PI = 0.7978845608028654f; +constexpr float GELU_COEFF = 0.044715f; +constexpr float GELU_INNER_COEFF = GELU_COEFF * SQRT_2_OVER_PI; +constexpr float DGELU_COEFF = 3.0f * GELU_COEFF * SQRT_2_OVER_PI; + +static __device__ inline float fast_tanh(float x) { + x = fmaxf(fminf(x, 20.f), -20.f); + float e2x = __builtin_amdgcn_exp2f(x * 2.8853900817779268f); + return (e2x - 1.0f) * __frcp_rn(e2x + 1.0f); +} + +/** + * @brief Gaussian Error Linear Unit (GELU) activation. + * + * Computes the GELU approximation: x * 0.5 * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3))). + * + * @tparam T The data type of the input and output values. + * @param x[in] The input value. + * @return The GELU activation applied to the input. + */ +struct gelu { + template static __device__ inline T op(const T &x); +}; +template<> __device__ inline float gelu::op(const float &x) { + return x * (0.5f + 0.5f * fast_tanh(x * (SQRT_2_OVER_PI + GELU_INNER_COEFF * x * x))); +} +template<> __device__ inline float2 gelu::op(const float2 &x) { + return float2{gelu::op(x.x), gelu::op(x.y)}; +} + +/** + * @brief Derivative of the GELU activation. + * + * Computes the derivative of the GELU approximation with respect to the input. + * + * @tparam T The data type of the input and output values. + * @param x[in] The input value. + * @return The derivative of GELU evaluated at the input. + */ +struct dgelu { + template static __device__ inline T op(const T &x); +}; +template<> __device__ inline float dgelu::op(const float &x) { + float tanh_out = fast_tanh(SQRT_2_OVER_PI * x * (1.f + GELU_COEFF * x * x)); + return 0.5f * x * ((1.f - tanh_out * tanh_out) * (SQRT_2_OVER_PI + DGELU_COEFF * x * x)) + + 0.5f * (1.f + tanh_out); +} +template<> __device__ inline float2 dgelu::op(const float2 &x) { + return float2{dgelu::op(x.x), dgelu::op(x.y)}; +} + /** * @brief Copy operation. * diff --git a/extra/thunder/amd/include/common/base_types.cuh b/extra/thunder/amd/include/common/base_types.cuh index cb6753a080..7587b4f69c 100644 --- a/extra/thunder/amd/include/common/base_types.cuh +++ b/extra/thunder/amd/include/common/base_types.cuh @@ -10,14 +10,16 @@ #pragma once -#include -#include -#include +#include +#include #include +#include +#include #include #include #include +typedef uint32_t __amd_fp8x4_storage_t; namespace kittens { @@ -37,7 +39,6 @@ using bf16_2 = __hip_bfloat162; * @brief Packed word of two half-precision floating-point values. */ using half_2 = __half2; -#ifdef KITTENS_CDNA4 /** * @brief float8 floating-point type. */ @@ -50,20 +51,30 @@ using fp8e4m3_2 = __hip_fp8x2_e4m3; * @brief Packed word of four float8 floating-point values. */ using fp8e4m3_4 = __hip_fp8x4_e4m3; -#else /** - * @brief float8 floating-point type. + * @brief 8-bit exponent-only block-scaling scale type. */ -using fp8e4m3 = __hip_fp8_e4m3_fnuz; +using fp8e8m0 = __amd_scale_t; /** - * @brief Packed word of two float8 floating-point values. + * @brief Packed word of two 8-bit exponent-only block-scaling scale values. */ -using fp8e4m3_2 = __hip_fp8x2_e4m3_fnuz; +using fp8e8m0_2 = __amd_fp8x2_storage_t; /** - * @brief Packed word of four float8 floating-point values. + * @brief Packed word of four 8-bit exponent-only block-scaling scale values. */ -using fp8e4m3_4 = __hip_fp8x4_e4m3_fnuz; -#endif +using fp8e8m0_4 = __amd_fp8x4_storage_t; +/** + * @brief FP4 E2M1 floating-point type. + */ +using fp4e2m1 = __hip_fp4_e2m1; +/** + * @brief Packed word of two FP4 E2M1 floating-point values. + */ +using fp4e2m1_2 = __hip_fp4x2_e2m1; +/** + * @brief Packed word of four FP4 E2M1 floating-point values. + */ +using fp4e2m1_4 = __hip_fp4x4_e2m1; namespace ducks { /** @@ -74,9 +85,11 @@ namespace ducks { namespace base_types { template -concept T2 = std::is_same_v || std::is_same_v || std::is_same_v || std::is_same_v; +concept T2 = std::is_same_v || std::is_same_v || std::is_same_v || std::is_same_v + || std::is_same_v; template -concept T1 = std::is_same_v || std::is_same_v || std::is_same_v || std::is_same_v; +concept T1 = std::is_same_v || std::is_same_v || std::is_same_v || std::is_same_v + || std::is_same_v; } // namespace base_types } // namespace ducks @@ -157,6 +170,26 @@ template<> struct constants { static __device__ inline constexpr fp8e4m3_4 zero() { return std::bit_cast(uint32_t(0x00000000)); } static __device__ inline constexpr fp8e4m3_4 one() { return std::bit_cast(uint32_t(0x38383838)); } }; +template<> struct constants { + static __device__ inline constexpr fp8e8m0 zero() { return std::bit_cast(uint8_t(0x00)); } // not actually 0 + static __device__ inline constexpr fp8e8m0 one() { return std::bit_cast(uint8_t(0x7F)); } +}; +template<> struct constants { + static __device__ inline constexpr fp8e8m0_2 zero() { return std::bit_cast(uint16_t(0x0000)); } // not actually 0 + static __device__ inline constexpr fp8e8m0_2 one() { return std::bit_cast(uint16_t(0x7F7F)); } +}; +template<> struct constants { + static __device__ inline constexpr fp8e8m0_4 zero() { return std::bit_cast(uint32_t(0x00000000)); } // not actually 0 + static __device__ inline constexpr fp8e8m0_4 one() { return std::bit_cast(uint32_t(0x7F7F7F7F)); } +}; +template<> struct constants { + static __device__ inline constexpr fp4e2m1 zero() { return std::bit_cast(uint8_t(0x00)); } + static __device__ inline constexpr fp4e2m1 one() { return std::bit_cast(uint8_t(0x02)); } +}; +template<> struct constants { + static __device__ inline constexpr fp4e2m1_4 zero() { return std::bit_cast(uint16_t(0x0000)); } + static __device__ inline constexpr fp4e2m1_4 one() { return std::bit_cast(uint16_t(0x2222)); } +}; template<> struct constants { static __device__ inline constexpr int zero() { return 0; } static __device__ inline constexpr int ones() { return 1; } @@ -250,6 +283,26 @@ template<> struct packing { using unpacked_type = fp8e4m3; using packed_type = fp8e4m3_4; }; +template<> struct packing { + static __host__ __device__ inline constexpr int num() { return 1; } + using unpacked_type = fp8e8m0; + using packed_type = fp8e8m0_4; +}; +template<> struct packing { + static __host__ __device__ inline constexpr int num() { return 4; } + using unpacked_type = fp8e8m0; + using packed_type = fp8e8m0_4; +}; +template<> struct packing { + static __host__ __device__ inline constexpr int num() { return 1; } + using unpacked_type = fp4e2m1; + using packed_type = fp4e2m1_4; +}; +template<> struct packing { + static __host__ __device__ inline constexpr int num() { return 4; } + using unpacked_type = fp4e2m1; + using packed_type = fp4e2m1_4; +}; /** * @brief Provides templated functionality to convert between different types. @@ -377,5 +430,25 @@ template<> struct convertor { return float(u); } }; +template<> struct convertor { + static __host__ __device__ inline fp4e2m1 convert(const float & u) { + return fp4e2m1(u); + } +}; +template<> struct convertor { + static __host__ __device__ inline float convert(const fp4e2m1 & u) { + return float(u); + } +}; +template<> struct convertor { + static __host__ __device__ inline fp4e2m1_4 convert(const float4& u) { + return fp4e2m1_4(u); + } +}; +template<> struct convertor { + static __host__ __device__ inline float4 convert(const fp4e2m1_4& u) { + return float4(u); + } +}; } } diff --git a/extra/thunder/amd/include/common/macros.cuh b/extra/thunder/amd/include/common/macros.cuh index bfb34b51ad..43ef1b956f 100644 --- a/extra/thunder/amd/include/common/macros.cuh +++ b/extra/thunder/amd/include/common/macros.cuh @@ -158,152 +158,405 @@ __device__ __forceinline__ void clobber_gpr() { #undef CLOBBER_AREG_CASE #undef CLOBBER_VREG_CASE -template -__device__ __forceinline__ void ds_read_b128(const uint32_t smem_ptr, const int offset) { +__device__ __forceinline__ constexpr uint32_t max_ds_inst_offset() +{ + // DS ops contain 2 8-bits instruction offset. + // For non-pk2 instructions like ds_read_b32, the 2 fields are regarded as 1. + // For pk2 instructions like ds_read2_b32, max offset is limited by 8 bits. + return (1u << 16) - 1; +} +__device__ __forceinline__ constexpr uint32_t max_ds_pk2_inst_offset() +{ + // DS ops contain 2 8-bits instruction offset. + // For non-pk2 instructions like ds_read_b32, the 2 fields are regarded as a whole. + // For pk2 instructions like ds_read2_b32, max offset is limited by 8 bits. + return (1u << 8) - 1; +} + +__device__ __forceinline__ constexpr uint32_t max_mubuf_inst_offset() +{ + // MUBUF ops contain 1 12-bits instruction offset. + return (1u << 12) - 1; +} + +template +__device__ __forceinline__ void ds_read_b32(const uint32_t smem_ptr, const int i_offset) { + // AGPRS + if constexpr (GPR_START >= 256) { + asm volatile("ds_read_b32 a[%0], %1 offset:%2" + : + : "n"(GPR_START - 256), "v"(smem_ptr), "i"(i_offset) + : "memory"); + // VGPRS + } else { + asm volatile("ds_read_b32 v[%0], %1 offset:%2" + : + : "n"(GPR_START), "v"(smem_ptr), "i"(i_offset) + : "memory"); + } +} + +template +__device__ __forceinline__ void ds_read_b32(T& dst, const uint32_t smem_ptr, const int i_offset) { + static_assert(sizeof(T) == sizeof(uint32_t)); + asm volatile("ds_read_b32 %0, %1 offset:%2" + : "=v"(dst) + : "v"(smem_ptr), "i"(i_offset) + : "memory"); +} + +template +__device__ __forceinline__ T ds_read_b64(const uint32_t smem_ptr, const int i_offset) { + static_assert(sizeof(T) == sizeof(uint32_t) * 2); + T result; + asm volatile("ds_read_b64 %0, %1 offset:%2" + : "=v"(result) + : "v"(smem_ptr), "i"(i_offset) + : "memory"); + return result; +} + +template +__device__ __forceinline__ void ds_read_b64(const uint32_t smem_ptr, const int i_offset) { + constexpr int GPR_END = GPR_START + 1; + // AGPRS + if constexpr (GPR_START >= 256) { + asm volatile("ds_read_b64 a[%0:%1], %2 offset:%3" + : + : "n"(GPR_START - 256), "n"(GPR_END - 256), "v"(smem_ptr), "i"(i_offset) + : "memory"); + // VGPRS + } else { + asm volatile("ds_read_b64 v[%0:%1], %2 offset:%3" + : + : "n"(GPR_START), "n"(GPR_END), "v"(smem_ptr), "i"(i_offset) + : "memory"); + } +} + +template +__device__ __forceinline__ void ds_read_b64_tr_b4(const uint32_t smem_ptr, const int i_offset) { + constexpr int GPR_END = GPR_START + 1; + if constexpr (GPR_START >= 256) { + asm volatile("ds_read_b64_tr_b4 a[%0:%1], %2 offset:%3" + : + : "n"(GPR_START - 256), "n"(GPR_END - 256), "v"(smem_ptr), "i"(i_offset) + : "memory"); + } else { + asm volatile("ds_read_b64_tr_b4 v[%0:%1], %2 offset:%3" + : + : "n"(GPR_START), "n"(GPR_END), "v"(smem_ptr), "i"(i_offset) + : "memory"); + } +} + +template +__device__ __forceinline__ void ds_read_b64_tr_b8(const uint32_t smem_ptr, const int i_offset) { + constexpr int GPR_END = GPR_START + 1; + if constexpr (GPR_START >= 256) { + asm volatile("ds_read_b64_tr_b8 a[%0:%1], %2 offset:%3" + : + : "n"(GPR_START - 256), "n"(GPR_END - 256), "v"(smem_ptr), "i"(i_offset) + : "memory"); + } else { + asm volatile("ds_read_b64_tr_b8 v[%0:%1], %2 offset:%3" + : + : "n"(GPR_START), "n"(GPR_END), "v"(smem_ptr), "i"(i_offset) + : "memory"); + } +} + +template +__device__ __forceinline__ void ds_read_b64_tr_b16(const uint32_t smem_ptr, const int i_offset) { + constexpr int GPR_END = GPR_START + 1; + if constexpr (GPR_START >= 256) { + asm volatile("ds_read_b64_tr_b16 a[%0:%1], %2 offset:%3" + : + : "n"(GPR_START - 256), "n"(GPR_END - 256), "v"(smem_ptr), "i"(i_offset) + : "memory"); + } else { + asm volatile("ds_read_b64_tr_b16 v[%0:%1], %2 offset:%3" + : + : "n"(GPR_START), "n"(GPR_END), "v"(smem_ptr), "i"(i_offset) + : "memory"); + } +} + +template +__device__ __forceinline__ T ds_read_b128(const uint32_t smem_ptr, const int i_offset) { + static_assert(sizeof(T) == sizeof(uint32_t) * 4); + T result; + asm volatile("ds_read_b128 %0, %1 offset:%2" + : "=v"(result) + : "v"(smem_ptr), "i"(i_offset) + : "memory"); + return result; +} + +template +__device__ __forceinline__ void ds_read_b128(const uint32_t smem_ptr, const int i_offset) { constexpr int GPR_END = GPR_START + 3; // AGPRS if constexpr (GPR_START >= 256) { asm volatile("ds_read_b128 a[%0:%1], %2 offset:%3" : - : "n"(GPR_START - 256), "n"(GPR_END - 256), "v"(smem_ptr), "i"(offset) + : "n"(GPR_START - 256), "n"(GPR_END - 256), "v"(smem_ptr), "i"(i_offset) : "memory"); // VGPRS } else { asm volatile("ds_read_b128 v[%0:%1], %2 offset:%3" : - : "n"(GPR_START), "n"(GPR_END), "v"(smem_ptr), "i"(offset) + : "n"(GPR_START), "n"(GPR_END), "v"(smem_ptr), "i"(i_offset) : "memory"); } } template -__device__ __forceinline__ void ds_read_b64_tr_b16(const uint32_t smem_ptr, const int offset) { - constexpr int GPR_END = GPR_START + 1; - +__device__ __forceinline__ void ds_write_b32(const uint32_t smem_ptr, const int i_offset) { if constexpr (GPR_START >= 256) { - asm volatile("ds_read_b64_tr_b16 a[%0:%1], %2 offset:%3" + asm volatile("ds_write_b32 %0, a[%1], offset:%2" : - : "n"(GPR_START - 256), "n"(GPR_END - 256), "v"(smem_ptr), "i"(offset) + : "v"(smem_ptr), "n"(GPR_START - 256), "i"(i_offset) : "memory"); } else { - asm volatile("ds_read_b64_tr_b16 v[%0:%1], %2 offset:%3" + asm volatile("ds_write_b32 %0, v[%1], offset:%2" : - : "n"(GPR_START), "n"(GPR_END), "v"(smem_ptr), "i"(offset) + : "v"(smem_ptr), "n"(GPR_START), "i"(i_offset) : "memory"); } } +template +__device__ __forceinline__ void ds_write_b32(const T& val, const uint32_t smem_ptr, const int i_offset = 0) { + static_assert(sizeof(T) == sizeof(uint32_t)); + asm volatile("ds_write_b32 %0, %1 offset:%2" + : + : "v"(smem_ptr), "v"(val), "i"(i_offset) + : "memory"); +} + template -__device__ __forceinline__ void ds_write_b64(const uint32_t smem_ptr, const int offset) { +__device__ __forceinline__ void ds_write_b64(const uint32_t smem_ptr, const int i_offset) { if constexpr (GPR_START >= 256) { asm volatile("ds_write_b64 %0, a[%1:%2], offset:%3" : - : "v"(smem_ptr), "n"(GPR_START - 256), "n"(GPR_START + 1 - 256), "i"(offset) + : "v"(smem_ptr), "n"(GPR_START - 256), "n"(GPR_START + 1 - 256), "i"(i_offset) : "memory"); } else { asm volatile("ds_write_b64 %0, v[%1:%2], offset:%3" : - : "v"(smem_ptr), "n"(GPR_START), "n"(GPR_START + 1), "i"(offset) + : "v"(smem_ptr), "n"(GPR_START), "n"(GPR_START + 1), "i"(i_offset) : "memory"); } } +template +__device__ __forceinline__ void ds_write_b64(const T& val, const uint32_t smem_ptr, const int i_offset = 0) { + static_assert(sizeof(T) == 2 * sizeof(uint32_t)); + asm volatile("ds_write_b64 %0, %1 offset:%2" + : + : "v"(smem_ptr), "v"(val), "i"(i_offset) + : "memory"); +} + +template +__device__ __forceinline__ void ds_write_b128(const uint32_t smem_ptr, const int i_offset = 0) { + if constexpr (GPR_START >= 256) { + asm volatile("ds_write_b128 %0, a[%1:%2], offset:%3" + : + : "v"(smem_ptr), "n"(GPR_START - 256), "n"(GPR_START + 3 - 256), "i"(i_offset) + : "memory"); + } else { + asm volatile("ds_write_b128 %0, v[%1:%2], offset:%3" + : + : "v"(smem_ptr), "n"(GPR_START), "n"(GPR_START + 3), "i"(i_offset) + : "memory"); + } +} + +template +__device__ __forceinline__ void ds_write_b128(const T& value, const uint32_t smem_ptr, const int i_offset = 0) { + static_assert(sizeof(T) == sizeof(u32x4)); + asm volatile("ds_write_b128 %0, %1 offset:%2" + : + : "v"(smem_ptr), "v"(value), "i"(i_offset) + : "memory"); +} + +template +__device__ __forceinline__ void buffer_load_dword(const buffer_resource& br, const uint32_t v_offset, const uint32_t s_offset = 0, const int i_offset = 0) { + if constexpr (GPR_START >= 256) { + asm volatile("buffer_load_dword a[%0], %1, %2, %3 offen offset:%4" + : + : "n"(GPR_START - 256), "v"(v_offset), "s"(*(const i32x4*)&br), "s"(s_offset), "i"(i_offset) + : "memory"); + } else { + asm volatile("buffer_load_dword v[%0], %1, %2, %3 offen offset:%4" + : + : "n"(GPR_START), "v"(v_offset), "s"(*(const i32x4*)&br), "s"(s_offset), "i"(i_offset) + : "memory"); + } +} + +template +__device__ __forceinline__ T buffer_load_dword( + const buffer_resource& br, const uint32_t v_offset, const uint32_t s_offset = 0, const uint32_t i_offset = 0) { + static_assert(sizeof(T) == sizeof(uint32_t)); + T result; + asm volatile("buffer_load_dword %0, %1, %2, %3 offen offset:%4" + : "=v"(result) + : "v"(v_offset), "s"(*(const i32x4*)&br), "s"(s_offset), "i"(i_offset) + : "memory"); + return result; +} + +template +__device__ __forceinline__ void buffer_load_dwordx2(const buffer_resource& br, const uint32_t v_offset, const uint32_t s_offset = 0, const int i_offset = 0) { + if constexpr (GPR_START >= 256) { + asm volatile("buffer_load_dwordx2 a[%0:%1], %2, %3, %4 offen offset:%5" + : + : "n"(GPR_START - 256), "n"(GPR_START + 1 - 256), "v"(v_offset), "s"(*(const i32x4*)&br), "s"(s_offset), "i"(i_offset) + : "memory"); + } else { + asm volatile("buffer_load_dwordx2 v[%0:%1], %2, %3, %4 offen offset:%5" + : + : "n"(GPR_START), "n"(GPR_START + 1), "v"(v_offset), "s"(*(const i32x4*)&br), "s"(s_offset), "i"(i_offset) + : "memory"); + } +} + +template +__device__ __forceinline__ T buffer_load_dwordx2( + const buffer_resource& br, const uint32_t v_offset, const uint32_t s_offset = 0, const uint32_t i_offset = 0) { + static_assert(sizeof(T) == sizeof(uint32_t) * 2); + T result; + asm volatile("buffer_load_dwordx2 %0, %1, %2, %3 offen offset:%4" + : "=v"(result) + : "v"(v_offset), "s"(*(const i32x4*)&br), "s"(s_offset), "i"(i_offset) + : "memory"); + return result; +} + +template +__device__ __forceinline__ void buffer_load_dwordx4(const buffer_resource& br, const uint32_t v_offset, const uint32_t s_offset = 0, const int i_offset = 0) { + if constexpr (GPR_START >= 256) { + asm volatile("buffer_load_dwordx4 a[%0:%1], %2, %3, %4 offen offset:%5" + : + : "n"(GPR_START - 256), "n"(GPR_START + 3 - 256), "v"(v_offset), "s"(*(const i32x4*)&br), "s"(s_offset), "i"(i_offset) + : "memory"); + } else { + asm volatile("buffer_load_dwordx4 v[%0:%1], %2, %3, %4 offen offset:%5" + : + : "n"(GPR_START), "n"(GPR_START + 3), "v"(v_offset), "s"(*(const i32x4*)&br), "s"(s_offset), "i"(i_offset) + : "memory"); + } +} + +template +__device__ __forceinline__ T buffer_load_dwordx4( + const buffer_resource& br, const uint32_t v_offset, const uint32_t s_offset = 0, const uint32_t i_offset = 0) { + static_assert(sizeof(T) == sizeof(uint32_t) * 4); + T result; + asm volatile("buffer_load_dwordx4 %0, %1, %2, %3 offen offset:%4" + : "=v"(result) + : "v"(v_offset), "s"(*(const i32x4*)&br), "s"(s_offset), "i"(i_offset) + : "memory"); + return result; +} + template -__device__ __forceinline__ void buffer_store_dword(buffer_resource& br, const uint32_t byte_offset) { - +__device__ __forceinline__ void buffer_store_dword(const buffer_resource& br, const uint32_t v_offset, const uint32_t s_offset = 0, const int i_offset = 0) { // AGPRS if constexpr (GPR >= 256) { - asm volatile("buffer_store_dword a[%0], %1, %2, 0 offen" + asm volatile("buffer_store_dword a[%0], %1, %2, %3 offen offset:%4" : - : "n"(GPR - 256), "v"(byte_offset), "s"(*(i32x4*)&br) + : "n"(GPR - 256), "v"(v_offset), "s"(*(const i32x4*)&br), "s"(s_offset), "i"(i_offset) : "memory"); // VGPRS } else { - asm volatile("buffer_store_dword v[%0], %1, %2, 0 offen" + asm volatile("buffer_store_dword v[%0], %1, %2, %3 offen offset:%4" : - : "n"(GPR), "v"(byte_offset), "s"(*(i32x4*)&br) + : "n"(GPR), "v"(v_offset), "s"(*(const i32x4*)&br), "s"(s_offset), "i"(i_offset) : "memory"); } } -template -__device__ __forceinline__ void buffer_store_dwordx2(buffer_resource& br, const uint32_t byte_offset) { +template +__device__ __forceinline__ void buffer_store_dword( + const T& value, const buffer_resource& br, const uint32_t v_offset, const uint32_t s_offset = 0, const uint32_t i_offset = 0) { + static_assert(sizeof(T) == sizeof(uint32_t) * 2); + asm volatile("buffer_store_dword %0, %1, %2, %3 offen offset:%4" + : + : "v"(value), "v"(v_offset), "s"(*(const i32x4*)&br), "s"(s_offset), "i"(i_offset) + : "memory"); +} +template +__device__ __forceinline__ void buffer_store_dwordx2(const buffer_resource& br, const uint32_t v_offset, const uint32_t s_offset = 0, const int i_offset = 0) { // AGPRS if constexpr (GPR_START >= 256) { - asm volatile("buffer_store_dwordx2 a[%0:%1], %2, %3, 0 offen" + asm volatile("buffer_store_dwordx2 a[%0:%1], %2, %3, %4 offen offset:%5" : - : "n"(GPR_START - 256), "n"(GPR_START + 1 - 256), "v"(byte_offset), "s"(*(i32x4*)&br) + : "n"(GPR_START - 256), "n"(GPR_START + 1 - 256), "v"(v_offset), "s"(*(const i32x4*)&br), "s"(s_offset), "i"(i_offset) : "memory"); // VGPRS } else { - asm volatile("buffer_store_dwordx2 v[%0:%1], %2, %3, 0 offen" + asm volatile("buffer_store_dwordx2 v[%0:%1], %2, %3, %4 offen offset:%5" : - : "n"(GPR_START), "n"(GPR_START + 1), "v"(byte_offset), "s"(*(i32x4*)&br) + : "n"(GPR_START), "n"(GPR_START + 1), "v"(v_offset), "s"(*(const i32x4*)&br), "s"(s_offset), "i"(i_offset) : "memory"); } } -template -__device__ __forceinline__ void buffer_store_dwordx4(buffer_resource& br, const uint32_t byte_offset) { +template +__device__ __forceinline__ void buffer_store_dwordx2( + const T& value, const buffer_resource& br, const uint32_t v_offset, const uint32_t s_offset = 0, const uint32_t i_offset = 0) { + static_assert(sizeof(T) == sizeof(uint32_t) * 2); + asm volatile("buffer_store_dwordx2 %0, %1, %2, %3 offen offset:%4" + : + : "v"(value), "v"(v_offset), "s"(*(const i32x4*)&br), "s"(s_offset), "i"(i_offset) + : "memory"); +} +template +__device__ __forceinline__ void buffer_store_dwordx4(const buffer_resource& br, const uint32_t v_offset, const uint32_t s_offset = 0, const int i_offset = 0) { // AGPRS if constexpr (GPR_START >= 256) { - asm volatile("buffer_store_dwordx4 a[%0:%1], %2, %3, 0 offen" + asm volatile("buffer_store_dwordx4 a[%0:%1], %2, %3, %4 offen offset:%5" : - : "n"(GPR_START - 256), "n"(GPR_START + 3 - 256), "v"(byte_offset), "s"(*(i32x4*)&br) + : "n"(GPR_START - 256), "n"(GPR_START + 3 - 256), "v"(v_offset), "s"(*(const i32x4*)&br), "s"(s_offset), "i"(i_offset) : "memory"); // VGPRS } else { - asm volatile("buffer_store_dwordx4 v[%0:%1], %2, %3, 0 offen" + asm volatile("buffer_store_dwordx4 v[%0:%1], %2, %3, %4 offen offset:%5" : - : "n"(GPR_START), "n"(GPR_START + 3), "v"(byte_offset), "s"(*(i32x4*)&br) + : "n"(GPR_START), "n"(GPR_START + 3), "v"(v_offset), "s"(*(const i32x4*)&br), "s"(s_offset), "i"(i_offset) : "memory"); } } -template -__device__ __forceinline__ void buffer_load_dwordx4(buffer_resource& br, const uint32_t byte_offset) { - if constexpr (GPR_START >= 256) { - asm volatile("buffer_load_dwordx4 a[%0:%1], %2, %3, 0 offen offset:%4" - : - : "n"(GPR_START - 256), "n"(GPR_START + 3 - 256), "v"(byte_offset), "s"(*(i32x4*)&br), "i"(0) - : "memory"); - } else { - asm volatile("buffer_load_dwordx4 v[%0:%1], %2, %3, 0 offen offset:%4" - : - : "n"(GPR_START), "n"(GPR_START + 3), "v"(byte_offset), "s"(*(i32x4*)&br), "i"(0) - : "memory"); - } -} - -template -__device__ __forceinline__ void buffer_load_dwordx2(buffer_resource& br, const uint32_t byte_offset) { - if constexpr (GPR_START >= 256) { - asm volatile("buffer_load_dwordx2 a[%0:%1], %2, %3, 0 offen offset:%4" - : - : "n"(GPR_START - 256), "n"(GPR_START + 1 - 256), "v"(byte_offset), "s"(*(i32x4*)&br), "i"(0) - : "memory"); - } else { - asm volatile("buffer_load_dwordx2 v[%0:%1], %2, %3, 0 offen offset:%4" - : - : "n"(GPR_START), "n"(GPR_START + 1), "v"(byte_offset), "s"(*(i32x4*)&br), "i"(0) - : "memory"); - } +template +__device__ __forceinline__ void buffer_store_dwordx4( + const T& value, const buffer_resource& br, const uint32_t v_offset, const uint32_t s_offset = 0, const uint32_t i_offset = 0) { + static_assert(sizeof(T) == sizeof(uint32_t) * 4); + asm volatile("buffer_store_dwordx4 %0, %1, %2, %3 offen offset:%4" + : + : "v"(value), "v"(v_offset), "s"(*(const i32x4*)&br), "s"(s_offset), "i"(i_offset) + : "memory"); } template -__device__ __forceinline__ void buffer_atomic_pk_add_bf16(buffer_resource& br, const uint32_t byte_offset) { +__device__ __forceinline__ void buffer_atomic_pk_add_bf16(const buffer_resource& br, const uint32_t v_offset, const uint32_t s_offset = 0, const int i_offset = 0) { if constexpr (GPR >= 256) { - asm volatile("buffer_atomic_pk_add_bf16 a[%0], %1, %2, 0 offen" + asm volatile("buffer_atomic_pk_add_bf16 a[%0], %1, %2, %3 offen offset:%4" : - : "n"(GPR - 256), "v"(byte_offset), "s"(*(i32x4*)&br) + : "n"(GPR - 256), "v"(v_offset), "s"(*(const i32x4*)&br), "s"(s_offset), "i"(i_offset) : "memory"); } else { - asm volatile("buffer_atomic_pk_add_bf16 v[%0], %1, %2, 0 offen" + asm volatile("buffer_atomic_pk_add_bf16 v[%0], %1, %2, %3 offen offset:%4" : - : "n"(GPR), "v"(byte_offset), "s"(*(i32x4*)&br) + : "n"(GPR), "v"(v_offset), "s"(*(const i32x4*)&br), "s"(s_offset), "i"(i_offset) : "memory"); } } @@ -468,6 +721,75 @@ __device__ __forceinline__ void mfma_f32_32x32x16_bf16() { } } +template +__device__ __forceinline__ void mfma_f32_16x16x32_fp8_fp8() { + if constexpr (GPR_START_D >= 256 && GPR_START_A >= 256 && GPR_START_B >= 256 && GPR_START_C >= 256) { + asm volatile("v_mfma_f32_16x16x32_fp8_fp8 a[%0:%1], a[%2:%3], a[%4:%5], a[%6:%7]" + : + : "n"(GPR_START_D - 256), "n"(GPR_START_D + 3 - 256), "n"(GPR_START_A - 256), "n"(GPR_START_A + 1 - 256), "n"(GPR_START_B - 256), "n"(GPR_START_B + 1 - 256), "n"(GPR_START_C - 256), "n"(GPR_START_C + 3 - 256)); + } else if constexpr (GPR_START_D >= 256 && GPR_START_A >= 256 && GPR_START_B >= 256 && GPR_START_C < 256) { + asm volatile("v_mfma_f32_16x16x32_fp8_fp8 a[%0:%1], a[%2:%3], a[%4:%5], v[%6:%7]" + : + : "n"(GPR_START_D - 256), "n"(GPR_START_D + 3 - 256), "n"(GPR_START_A - 256), "n"(GPR_START_A + 1 - 256), "n"(GPR_START_B - 256), "n"(GPR_START_B + 1 - 256), "n"(GPR_START_C), "n"(GPR_START_C + 3)); + } else if constexpr (GPR_START_D >= 256 && GPR_START_A >= 256 && GPR_START_B < 256 && GPR_START_C >= 256) { + asm volatile("v_mfma_f32_16x16x32_fp8_fp8 a[%0:%1], a[%2:%3], v[%4:%5], a[%6:%7]" + : + : "n"(GPR_START_D - 256), "n"(GPR_START_D + 3 - 256), "n"(GPR_START_A - 256), "n"(GPR_START_A + 1 - 256), "n"(GPR_START_B), "n"(GPR_START_B + 1), "n"(GPR_START_C - 256), "n"(GPR_START_C + 3 - 256)); + } else if constexpr (GPR_START_D >= 256 && GPR_START_A < 256 && GPR_START_B >= 256 && GPR_START_C >= 256) { + asm volatile("v_mfma_f32_16x16x32_fp8_fp8 a[%0:%1], v[%2:%3], a[%4:%5], a[%6:%7]" + : + : "n"(GPR_START_D - 256), "n"(GPR_START_D + 3 - 256), "n"(GPR_START_A), "n"(GPR_START_A + 1), "n"(GPR_START_B - 256), "n"(GPR_START_B + 1 - 256), "n"(GPR_START_C - 256), "n"(GPR_START_C + 3 - 256)); + } else if constexpr (GPR_START_D < 256 && GPR_START_A >= 256 && GPR_START_B >= 256 && GPR_START_C >= 256) { + asm volatile("v_mfma_f32_16x16x32_fp8_fp8 v[%0:%1], a[%2:%3], a[%4:%5], a[%6:%7]" + : + : "n"(GPR_START_D), "n"(GPR_START_D + 3), "n"(GPR_START_A - 256), "n"(GPR_START_A + 1 - 256), "n"(GPR_START_B - 256), "n"(GPR_START_B + 1 - 256), "n"(GPR_START_C - 256), "n"(GPR_START_C + 3 - 256)); + } else if constexpr (GPR_START_D < 256 && GPR_START_A >= 256 && GPR_START_B >= 256 && GPR_START_C < 256) { + asm volatile("v_mfma_f32_16x16x32_fp8_fp8 v[%0:%1], a[%2:%3], a[%4:%5], v[%6:%7]" + : + : "n"(GPR_START_D), "n"(GPR_START_D + 3), "n"(GPR_START_A - 256), "n"(GPR_START_A + 1 - 256), "n"(GPR_START_B - 256), "n"(GPR_START_B + 1 - 256), "n"(GPR_START_C), "n"(GPR_START_C + 3)); + } else if constexpr (GPR_START_D < 256 && GPR_START_A >= 256 && GPR_START_B < 256 && GPR_START_C >= 256) { + asm volatile("v_mfma_f32_16x16x32_fp8_fp8 v[%0:%1], a[%2:%3], v[%4:%5], a[%6:%7]" + : + : "n"(GPR_START_D), "n"(GPR_START_D + 3), "n"(GPR_START_A - 256), "n"(GPR_START_A + 1 - 256), "n"(GPR_START_B), "n"(GPR_START_B + 1), "n"(GPR_START_C - 256), "n"(GPR_START_C + 3 - 256)); + } else if constexpr (GPR_START_D < 256 && GPR_START_A < 256 && GPR_START_B >= 256 && GPR_START_C >= 256) { + asm volatile("v_mfma_f32_16x16x32_fp8_fp8 v[%0:%1], v[%2:%3], a[%4:%5], a[%6:%7]" + : + : "n"(GPR_START_D), "n"(GPR_START_D + 3), "n"(GPR_START_A), "n"(GPR_START_A + 1), "n"(GPR_START_B - 256), "n"(GPR_START_B + 1 - 256), "n"(GPR_START_C - 256), "n"(GPR_START_C + 3 - 256)); + } else if constexpr (GPR_START_D >= 256 && GPR_START_A < 256 && GPR_START_B >= 256 && GPR_START_C < 256) { + asm volatile("v_mfma_f32_16x16x32_fp8_fp8 a[%0:%1], v[%2:%3], a[%4:%5], v[%6:%7]" + : + : "n"(GPR_START_D - 256), "n"(GPR_START_D + 3 - 256), "n"(GPR_START_A), "n"(GPR_START_A + 1), "n"(GPR_START_B - 256), "n"(GPR_START_B + 1 - 256), "n"(GPR_START_C), "n"(GPR_START_C + 3)); + } else if constexpr (GPR_START_D >= 256 && GPR_START_A < 256 && GPR_START_B < 256 && GPR_START_C >= 256) { + asm volatile("v_mfma_f32_16x16x32_fp8_fp8 a[%0:%1], v[%2:%3], v[%4:%5], a[%6:%7]" + : + : "n"(GPR_START_D - 256), "n"(GPR_START_D + 3 - 256), "n"(GPR_START_A), "n"(GPR_START_A + 1), "n"(GPR_START_B), "n"(GPR_START_B + 1), "n"(GPR_START_C - 256), "n"(GPR_START_C + 3 - 256)); + } else if constexpr (GPR_START_D >= 256 && GPR_START_A >= 256 && GPR_START_B < 256 && GPR_START_C < 256) { + asm volatile("v_mfma_f32_16x16x32_fp8_fp8 a[%0:%1], a[%2:%3], v[%4:%5], v[%6:%7]" + : + : "n"(GPR_START_D - 256), "n"(GPR_START_D + 3 - 256), "n"(GPR_START_A - 256), "n"(GPR_START_A + 1 - 256), "n"(GPR_START_B), "n"(GPR_START_B + 1), "n"(GPR_START_C), "n"(GPR_START_C + 3)); + } else if constexpr (GPR_START_D < 256 && GPR_START_A >= 256 && GPR_START_B < 256 && GPR_START_C < 256) { + asm volatile("v_mfma_f32_16x16x32_fp8_fp8 v[%0:%1], a[%2:%3], v[%4:%5], v[%6:%7]" + : + : "n"(GPR_START_D), "n"(GPR_START_D + 3), "n"(GPR_START_A - 256), "n"(GPR_START_A + 1 - 256), "n"(GPR_START_B), "n"(GPR_START_B + 1), "n"(GPR_START_C), "n"(GPR_START_C + 3)); + } else if constexpr (GPR_START_D < 256 && GPR_START_A < 256 && GPR_START_B >= 256 && GPR_START_C < 256) { + asm volatile("v_mfma_f32_16x16x32_fp8_fp8 v[%0:%1], v[%2:%3], a[%4:%5], v[%6:%7]" + : + : "n"(GPR_START_D), "n"(GPR_START_D + 3), "n"(GPR_START_A), "n"(GPR_START_A + 1), "n"(GPR_START_B - 256), "n"(GPR_START_B + 1 - 256), "n"(GPR_START_C), "n"(GPR_START_C + 3)); + } else if constexpr (GPR_START_D < 256 && GPR_START_A < 256 && GPR_START_B < 256 && GPR_START_C >= 256) { + asm volatile("v_mfma_f32_16x16x32_fp8_fp8 v[%0:%1], v[%2:%3], v[%4:%5], a[%6:%7]" + : + : "n"(GPR_START_D), "n"(GPR_START_D + 3), "n"(GPR_START_A), "n"(GPR_START_A + 1), "n"(GPR_START_B), "n"(GPR_START_B + 1), "n"(GPR_START_C - 256), "n"(GPR_START_C + 3 - 256)); + } else if constexpr (GPR_START_D >= 256 && GPR_START_A < 256 && GPR_START_B < 256 && GPR_START_C < 256) { + asm volatile("v_mfma_f32_16x16x32_fp8_fp8 a[%0:%1], v[%2:%3], v[%4:%5], v[%6:%7]" + : + : "n"(GPR_START_D - 256), "n"(GPR_START_D + 3 - 256), "n"(GPR_START_A), "n"(GPR_START_A + 1), "n"(GPR_START_B), "n"(GPR_START_B + 1), "n"(GPR_START_C), "n"(GPR_START_C + 3)); + } else { + asm volatile("v_mfma_f32_16x16x32_fp8_fp8 v[%0:%1], v[%2:%3], v[%4:%5], v[%6:%7]" + : + : "n"(GPR_START_D), "n"(GPR_START_D + 3), "n"(GPR_START_A), "n"(GPR_START_A + 1), "n"(GPR_START_B), "n"(GPR_START_B + 1), "n"(GPR_START_C), "n"(GPR_START_C + 3)); + } +} + template __device__ __forceinline__ void mfma_f32_16x16x32_bf16_zero_accum() { if constexpr (GPR_START_D >= 256 && GPR_START_A >= 256 && GPR_START_B >= 256) { @@ -542,6 +864,43 @@ __device__ __forceinline__ void mfma_f32_32x32x16_bf16_zero_accum() { } } +template +__device__ __forceinline__ void mfma_f32_16x16x32_fp8_fp8_zero_accum() { + if constexpr (GPR_START_D >= 256 && GPR_START_A >= 256 && GPR_START_B >= 256) { + asm volatile("v_mfma_f32_16x16x32_fp8_fp8 a[%0:%1], a[%2:%3], a[%4:%5], 0" + : + : "n"(GPR_START_D - 256), "n"(GPR_START_D + 3 - 256), "n"(GPR_START_A - 256), "n"(GPR_START_A + 1 - 256), "n"(GPR_START_B - 256), "n"(GPR_START_B + 1 - 256)); + } else if constexpr (GPR_START_D < 256 && GPR_START_A >= 256 && GPR_START_B >= 256) { + asm volatile("v_mfma_f32_16x16x32_fp8_fp8 v[%0:%1], a[%2:%3], a[%4:%5], 0" + : + : "n"(GPR_START_D), "n"(GPR_START_D + 3), "n"(GPR_START_A - 256), "n"(GPR_START_A + 1 - 256), "n"(GPR_START_B - 256), "n"(GPR_START_B + 1 - 256)); + } else if constexpr (GPR_START_D >= 256 && GPR_START_A < 256 && GPR_START_B >= 256) { + asm volatile("v_mfma_f32_16x16x32_fp8_fp8 a[%0:%1], v[%2:%3], a[%4:%5], 0" + : + : "n"(GPR_START_D - 256), "n"(GPR_START_D + 3 - 256), "n"(GPR_START_A), "n"(GPR_START_A + 1), "n"(GPR_START_B - 256), "n"(GPR_START_B + 1 - 256)); + } else if constexpr (GPR_START_D >= 256 && GPR_START_A >= 256 && GPR_START_B < 256) { + asm volatile("v_mfma_f32_16x16x32_fp8_fp8 a[%0:%1], a[%2:%3], v[%4:%5], 0" + : + : "n"(GPR_START_D - 256), "n"(GPR_START_D + 3 - 256), "n"(GPR_START_A - 256), "n"(GPR_START_A + 1 - 256), "n"(GPR_START_B), "n"(GPR_START_B + 1)); + } else if constexpr (GPR_START_D < 256 && GPR_START_A >= 256 && GPR_START_B < 256) { + asm volatile("v_mfma_f32_16x16x32_fp8_fp8 v[%0:%1], a[%2:%3], v[%4:%5], 0" + : + : "n"(GPR_START_D), "n"(GPR_START_D + 3), "n"(GPR_START_A - 256), "n"(GPR_START_A + 1 - 256), "n"(GPR_START_B), "n"(GPR_START_B + 1)); + } else if constexpr (GPR_START_D < 256 && GPR_START_A < 256 && GPR_START_B >= 256) { + asm volatile("v_mfma_f32_16x16x32_fp8_fp8 v[%0:%1], v[%2:%3], a[%4:%5], 0" + : + : "n"(GPR_START_D), "n"(GPR_START_D + 3), "n"(GPR_START_A), "n"(GPR_START_A + 1), "n"(GPR_START_B - 256), "n"(GPR_START_B + 1 - 256)); + } else if constexpr (GPR_START_D >= 256 && GPR_START_A < 256 && GPR_START_B < 256) { + asm volatile("v_mfma_f32_16x16x32_fp8_fp8 a[%0:%1], v[%2:%3], v[%4:%5], 0" + : + : "n"(GPR_START_D - 256), "n"(GPR_START_D + 3 - 256), "n"(GPR_START_A), "n"(GPR_START_A + 1), "n"(GPR_START_B), "n"(GPR_START_B + 1)); + } else { + asm volatile("v_mfma_f32_16x16x32_fp8_fp8 v[%0:%1], v[%2:%3], v[%4:%5], 0" + : + : "n"(GPR_START_D), "n"(GPR_START_D + 3), "n"(GPR_START_A), "n"(GPR_START_A + 1), "n"(GPR_START_B), "n"(GPR_START_B + 1)); + } +} + template __device__ __forceinline__ void v_subrev_f32_dpp() { @@ -592,11 +951,29 @@ __device__ __forceinline__ void v_accvgpr_read_b32() { : "n"(GPR0), "n"(GPR1 - 256)); } -template -__device__ __forceinline__ void v_mov_b32(const uint32_t value) { +template +__device__ __forceinline__ void v_mov_b32_up2p(const T value) { + static_assert(sizeof(T) == sizeof(uint32_t)); asm volatile("v_mov_b32 v[%0], %1" : - : "n"(GPR), "i"(value)); + : "n"(GPR), "v"(value)); +} + +template +__device__ __forceinline__ T v_mov_b32_p2up() { + static_assert(sizeof(T) == sizeof(uint32_t)); + T r; + if constexpr (GPR < 256) { + asm volatile("v_mov_b32 %0, v[%1]" + : "=v"(r) + : "n"(GPR)); + } + else { + asm volatile("v_accvgpr_read_b32 %0, a[%1]" + : "=v"(r) + : "n"(GPR - 256)); + } + return r; } template @@ -612,8 +989,9 @@ __device__ __forceinline__ void v_cndmask_b32_e64(uint64_t mask) { : : "n"(GPR0), "n"(GPR1), "n"(GPR2), "s"(mask)); } + /** - * @brief Multiplication operation on explicit registers. + * @brief Multiplication operation on explicit registers and immediate operand. */ struct mul { template @@ -628,6 +1006,12 @@ struct mul { } } + template + static __device__ inline void op_pk2(const float ¶m) { + op(param); + op(param); + } + template static __device__ inline void op() { if constexpr (GPR0 < 256 && GPR1 < 256 && GPR2 < 256) { @@ -638,8 +1022,44 @@ struct mul { static_assert(false, "Invalid operand for instruction: v_mul_f32_e32"); } } + + template + static __device__ inline void op_pk2() { + if constexpr (GPR0 < (256 - 1) && GPR1 < (256 - 1) && GPR2 < (256 - 1)) { + asm volatile("v_pk_mul_f32 v[%0:%1], v[%4:%5], v[%2:%3]" + : + : "n"(GPR0), "n"(GPR0 + 1), "n"(GPR1), "n"(GPR1 + 1), "n"(GPR2), "n"(GPR2 + 1)); + } else { + static_assert(false, "Invalid operand for instruction: v_pk_mul_f32"); + } + } }; + struct mul_vgpr { + template + static __device__ inline void op(const float ¶m) { + if constexpr (GPR0 < 256 && GPR1 < 256) { + asm volatile("v_mul_f32_e32 v[%0], %2, v[%1]" + : + : "n"(GPR0), "n"(GPR1), "v"(param)); + } else { + static_assert(false, "Invalid operand for instruction: v_mul_f32_e32"); + } + } + + template + static __device__ inline void op_pk2(const float ¶m) { + if constexpr (GPR0 < (256 - 1) && GPR1 < (256 - 1)) { + const float2 param2 = {param, param}; + asm volatile("v_pk_mul_f32 v[%0:%1], %4, v[%2:%3]" + : + : "n"(GPR0), "n"(GPR0 + 1), "n"(GPR1), "n"(GPR1 + 1), "v"(param2)); + } else { + static_assert(false, "Invalid operand for instruction: v_pk_mul_f32"); + } + } +}; + struct exp2 { template static __device__ inline void op() { @@ -669,4 +1089,4 @@ struct zero { }; } // namespace macros -} // namespace kittens \ No newline at end of file +} // namespace kittens diff --git a/extra/thunder/amd/include/common/util.cuh b/extra/thunder/amd/include/common/util.cuh index 80dfd55d6e..c26d0b93e3 100644 --- a/extra/thunder/amd/include/common/util.cuh +++ b/extra/thunder/amd/include/common/util.cuh @@ -50,7 +50,11 @@ __device__ __forceinline__ int warpid() { return threadIdx.x >> 6; } */ __device__ __forceinline__ int laneid() { return threadIdx.x & 0x3f; } +using i32x2 = int32_t __attribute__((ext_vector_type(2))); +using u32x2 = uint32_t __attribute__((ext_vector_type(2))); using i32x4 = int32_t __attribute__((ext_vector_type(4))); +using u32x4 = uint32_t __attribute__((ext_vector_type(4))); + struct buffer_resource { uint64_t ptr; uint32_t range; diff --git a/extra/thunder/amd/include/ops/warp/memory/tile/assembly/global_to_register.cuh b/extra/thunder/amd/include/ops/warp/memory/tile/assembly/global_to_register.cuh index a285340d81..e901c87402 100644 --- a/extra/thunder/amd/include/ops/warp/memory/tile/assembly/global_to_register.cuh +++ b/extra/thunder/amd/include/ops/warp/memory/tile/assembly/global_to_register.cuh @@ -21,7 +21,7 @@ * @param idx[in] The index of the tile to load data from. */ - template> + template> __device__ inline static void load(RT &dst, const GL &src, const COORD &idx, const COORD &warp_idx) { using T2 = RT::dtype; constexpr int packing = base_types::packing::num(); @@ -42,22 +42,48 @@ buffer_resource br = make_buffer_resource(as_u64, buffer_size, 0x00020000); int warp_offset = src.idx(warp_idx.template unit_coord()); + int thr_offset = (row_offset * row_stride + col_offset + warp_offset) * sizeof(U); // Compile-time loop to load data into the tile auto perform_load_at = [&]() { using tile_range = ducks::art::get_nth_range_t; const int register_offset = K * RT::registers_per_stride; + + constexpr int col = RT::base_tile_cols*M + K * RT::base_tile_elements_per_stride_group; + constexpr int row = RT::base_tile_rows*N; + const int k_row_offset = row * row_stride * sizeof(U); - const int col = RT::base_tile_cols*M + col_offset + K * RT::base_tile_elements_per_stride_group; - const int row = RT::base_tile_rows*N + row_offset; - const int offset = (row*row_stride + col + warp_offset) * sizeof(U); - - if constexpr (std::is_same_v) { - if constexpr (RT::base_tile_stride == 8) { - macros::buffer_load_dwordx4(br, offset); - } else if constexpr (RT::base_tile_stride == 4) { - macros::buffer_load_dwordx2(br, offset); - } + constexpr int stride_in_bytes = RT::base_tile_stride * sizeof(U); + constexpr int offset_in_bytes = (elem_offset + col) * sizeof(U); + constexpr int start_gpr = tile_range::lo + register_offset; + + if constexpr (offset_in_bytes <= macros::max_mubuf_inst_offset()) { + if constexpr (stride_in_bytes == (sizeof(int32_t) * 4)) { + macros::buffer_load_dwordx4(br, thr_offset + k_row_offset, 0, offset_in_bytes); + } + else if constexpr (stride_in_bytes == (sizeof(int32_t) * 2)) { + macros::buffer_load_dwordx2(br, thr_offset + k_row_offset, 0, offset_in_bytes); + } + else if constexpr (stride_in_bytes == sizeof(int32_t)) { + macros::buffer_load_dword(br, thr_offset + k_row_offset, 0, offset_in_bytes); + } + else { + static_assert(false, "Encounter unsupported format in ops/warp/memory/tile/assembly/global_to_register.cuh\n"); + } + } + else { + if constexpr (stride_in_bytes == (sizeof(int32_t) * 4)) { + macros::buffer_load_dwordx4(br, thr_offset + offset_in_bytes + k_row_offset, 0, 0); + } + else if constexpr (stride_in_bytes == (sizeof(int32_t) * 2)) { + macros::buffer_load_dwordx2(br, thr_offset + offset_in_bytes + k_row_offset, 0, 0); + } + else if constexpr (stride_in_bytes == sizeof(int32_t)) { + macros::buffer_load_dword(br, thr_offset + offset_in_bytes + k_row_offset, 0, 0); + } + else { + static_assert(false, "Encounter unsupported format in ops/warp/memory/tile/assembly/global_to_register.cuh\n"); + } } }; @@ -74,12 +100,11 @@ }(std::make_index_sequence{}); }.template operator()(), ...); }(std::make_index_sequence{}); - } template> __device__ inline static void load(RT &dst, const GL &src, const COORD &idx, const COORD &warp_idx) { - load<2, RT, GL>(dst, src, idx, warp_idx); + load<2, 0, RT, GL>(dst, src, idx, warp_idx); } /** diff --git a/extra/thunder/amd/include/ops/warp/memory/tile/global_to_shared.cuh b/extra/thunder/amd/include/ops/warp/memory/tile/global_to_shared.cuh index 627691925e..45baadd082 100644 --- a/extra/thunder/amd/include/ops/warp/memory/tile/global_to_shared.cuh +++ b/extra/thunder/amd/include/ops/warp/memory/tile/global_to_shared.cuh @@ -226,7 +226,8 @@ __device__ inline void load(ST& dst, const GL& src, const COORD& idx, const uint if (warpid < leftover_warps) { - uintptr_t lds_addr = lds_base + (memcpy_per_tile * num_warps * bytes_per_warp); + const T* lds_elem_ptr = lds_base + (memcpy_per_tile * num_warps * elements_per_warp); + uintptr_t lds_addr = reinterpret_cast(lds_elem_ptr); as3_uint32_ptr lds_ptr = (as3_uint32_ptr)(lds_addr); llvm_amdgcn_raw_buffer_load_lds( @@ -414,4 +415,4 @@ template(dst, src, idx); } -} \ No newline at end of file +} diff --git a/extra/thunder/amd/include/ops/warp/register/tile/assembly/maps.cuh b/extra/thunder/amd/include/ops/warp/register/tile/assembly/maps.cuh index 9318d53dc7..36a9367d03 100644 --- a/extra/thunder/amd/include/ops/warp/register/tile/assembly/maps.cuh +++ b/extra/thunder/amd/include/ops/warp/register/tile/assembly/maps.cuh @@ -134,7 +134,17 @@ __device__ static inline void bin_map(T0 &dst, const T1 &src, const typename bas [&](std::index_sequence) { ([&]() { - op::template op::lo, ducks::art::get_nth_range_t::lo>(param); + constexpr int GPR0 = ducks::art::get_nth_range_t::lo; + constexpr int GPR1 = ducks::art::get_nth_range_t::lo; + if constexpr ((R % 2) == 0) { + if constexpr ((R + 1) < registers_T0::size) { + op::template op_pk2(param); + } + else { + op::template op(param); + } + } + // Odd indices are skipped because they're processed by op_pk2 in the previous even iteration }.template operator()(), ...); }(std::make_index_sequence{}); } @@ -156,7 +166,17 @@ __device__ static inline void bin_map(T0 &dst, const T1 &src, const typename bas [&](std::index_sequence) { ([&]() { - op::template op::lo, ducks::art::get_nth_range_t::lo>(param); + constexpr int GPR0 = ducks::art::get_nth_range_t::lo; + constexpr int GPR1 = ducks::art::get_nth_range_t::lo; + if constexpr ((R % 2) == 0) { + if constexpr ((R + 1) < registers_T0::size) { + op::template op_pk2(param); + } + else { + op::template op(param); + } + } + // Odd indices are skipped because they're processed by op_pk2 in the previous even iteration }.template operator()(), ...); }(std::make_index_sequence{}); }; @@ -205,7 +225,18 @@ __device__ static inline void bin_map(T0 &dst, const T1 &lhs, const T2 &rhs) { [&](std::index_sequence) { ([&]() { - op::template op::lo, ducks::art::get_nth_range_t::lo, ducks::art::get_nth_range_t::lo>(); + constexpr int GPR0 = ducks::art::get_nth_range_t::lo; + constexpr int GPR1 = ducks::art::get_nth_range_t::lo; + constexpr int GPR2 = ducks::art::get_nth_range_t::lo; + if constexpr ((R % 2) == 0) { + if constexpr ((R + 1) < registers_T0::size) { + op::template op_pk2(); + } + else { + op::template op(); + } + } + // Odd indices are skipped because they're processed by op_pk2 in the previous even iteration }.template operator()(), ...); }(std::make_index_sequence{}); } @@ -234,7 +265,18 @@ __device__ static inline void bin_map(T0 &dst, const T1 &lhs, const T2 &rhs) { [&](std::index_sequence) { ([&]() { - op::template op::lo, ducks::art::get_nth_range_t::lo, ducks::art::get_nth_range_t::lo>(); + constexpr int GPR0 = ducks::art::get_nth_range_t::lo; + constexpr int GPR1 = ducks::art::get_nth_range_t::lo; + constexpr int GPR2 = ducks::art::get_nth_range_t::lo; + if constexpr ((R % 2) == 0) { + if constexpr ((R + 1) < registers_T0::size) { + op::template op_pk2(); + } + else { + op::template op(); + } + } + // Odd indices are skipped because they're processed by op_pk2 in the previous even iteration }.template operator()(), ...); }(std::make_index_sequence{}); }; @@ -364,6 +406,16 @@ __device__ static inline void mul(T0 &dst, const T1 &lhs, const U &rhs) { bin_map(dst, lhs, rhs); } +template +__device__ static inline void mul_vgpr(T0 &dst, const T1 &lhs, const U &rhs) { + bin_map(dst, lhs, rhs); +} + +template +__device__ static inline void mul_vgpr(T0 &dst, const T1 &lhs, const U &rhs) { + bin_map(dst, lhs, rhs); +} + /** * @brief Subtracts row values from each row of a tile. * diff --git a/extra/thunder/amd/include/ops/warp/register/tile/assembly/mma.cuh b/extra/thunder/amd/include/ops/warp/register/tile/assembly/mma.cuh index 061c19b366..eea68acf7d 100644 --- a/extra/thunder/amd/include/ops/warp/register/tile/assembly/mma.cuh +++ b/extra/thunder/amd/include/ops/warp/register/tile/assembly/mma.cuh @@ -20,14 +20,44 @@ namespace kittens { * @param[in] b The second input rt_base matrix in row-major mode. * @param[in] c The input rt_base accumulator matrix. */ -template +template __device__ static inline void mma_ABt_base() { - macros::mfma_f32_16x16x32_bf16(); + if constexpr (std::is_same_v) + { + if constexpr (std::is_same_v) + { + macros::mfma_f32_16x16x32_fp8_fp8(); + } + else + { + macros::mfma_f32_16x16x32_bf16(); + } + } + else + { + macros::mfma_f32_16x16x32_bf16(); + } } -template + +template __device__ static inline void mma_ABt_base_zero_accum() { - macros::mfma_f32_16x16x32_bf16_zero_accum(); + if constexpr (std::is_same_v) + { + if constexpr (std::is_same_v) + { + macros::mfma_f32_16x16x32_fp8_fp8_zero_accum(); + } + else + { + macros::mfma_f32_16x16x32_bf16_zero_accum(); + } + } + else + { + macros::mfma_f32_16x16x32_bf16_zero_accum(); + } } + /** * @brief Base matrix multiply-accumulate operation for row layout with transposed A. * @@ -87,7 +117,9 @@ __device__ static inline void mma_ABt(D &d, (std::is_same_v && std::is_same_v && std::is_same_v) || (std::is_same_v && std::is_same_v && - std::is_same_v) + std::is_same_v) || + (std::is_same_v && std::is_same_v && + std::is_same_v) ); // Helper function template for compile-time MMA operations @@ -95,7 +127,7 @@ __device__ static inline void mma_ABt(D &d, using range_type_B = ducks::art::get_nth_range_t; using range_type_C = ducks::art::get_nth_range_t; using range_type_D = ducks::art::get_nth_range_t; - mma_ABt_base(); + mma_ABt_base(); } template @@ -117,7 +149,9 @@ __device__ static inline void mma_ABt(D &d, (std::is_same_v && std::is_same_v && std::is_same_v) || (std::is_same_v && std::is_same_v && - std::is_same_v) + std::is_same_v) || + (std::is_same_v && std::is_same_v && + std::is_same_v) ); // Helper function template for compile-time MMA operations @@ -127,7 +161,7 @@ __device__ static inline void mma_ABt(D &d, using range_type_B = ducks::art::get_nth_range_t; using range_type_C = ducks::art::get_nth_range_t; using range_type_D = ducks::art::get_nth_range_t; - mma_ABt_base(); + mma_ABt_base(); // Subsequent MMA operations for k=1 to A::width-1 [&](std::index_sequence) { @@ -138,7 +172,7 @@ __device__ static inline void mma_ABt(D &d, using range_type_B = ducks::art::get_nth_range_t; using range_type_C = ducks::art::get_nth_range_t; using range_type_D = ducks::art::get_nth_range_t; - mma_ABt_base(); + mma_ABt_base(); } }(), ...); }(std::make_index_sequence{}); @@ -172,7 +206,9 @@ __device__ static inline void mma_ABt(D &d, (std::is_same_v && std::is_same_v && std::is_same_v) || (std::is_same_v && std::is_same_v && - std::is_same_v) + std::is_same_v) || + (std::is_same_v && std::is_same_v && + std::is_same_v) ); // Helper function template for compile-time MMA operations @@ -180,7 +216,7 @@ __device__ static inline void mma_ABt(D &d, using range_type_A = ducks::art::get_nth_range_t; using range_type_B = ducks::art::get_nth_range_t; using range_type_D = ducks::art::get_nth_range_t; - mma_ABt_base_zero_accum(); + mma_ABt_base_zero_accum(); } template @@ -199,7 +235,9 @@ __device__ static inline void mma_ABt(D &d, (std::is_same_v && std::is_same_v && std::is_same_v) || (std::is_same_v && std::is_same_v && - std::is_same_v) + std::is_same_v) || + (std::is_same_v && std::is_same_v && + std::is_same_v) ); // Helper function template for compile-time MMA operations @@ -208,7 +246,7 @@ __device__ static inline void mma_ABt(D &d, using range_type_A = ducks::art::get_nth_range_t; using range_type_B = ducks::art::get_nth_range_t; using range_type_D = ducks::art::get_nth_range_t; - mma_ABt_base_zero_accum(); + mma_ABt_base_zero_accum(); // Subsequent MMA operations for k=1 to A::width-1 [&](std::index_sequence) { @@ -218,7 +256,7 @@ __device__ static inline void mma_ABt(D &d, using range_type_A = ducks::art::get_nth_range_t; using range_type_B = ducks::art::get_nth_range_t; using range_type_D = ducks::art::get_nth_range_t; - mma_ABt_base(); + mma_ABt_base(); } }(), ...); }(std::make_index_sequence{}); diff --git a/extra/thunder/amd/include/ops/warp/register/tile/maps.cuh b/extra/thunder/amd/include/ops/warp/register/tile/maps.cuh index 925f40a7dd..73fc82ae52 100644 --- a/extra/thunder/amd/include/ops/warp/register/tile/maps.cuh +++ b/extra/thunder/amd/include/ops/warp/register/tile/maps.cuh @@ -472,6 +472,28 @@ template __device__ static inline void relu(T &dst, const T &src) { unary_map(dst, src); } +/** + * @brief Applies the GELU function (tanh approximation) to each element of a tile. + * + * @tparam T Tile type. + * @param dst[out] Destination tile where the result is stored. + * @param src[in] Source tile to apply the GELU function on. + */ +template +__device__ static inline void gelu(T &dst, const T &src) { + unary_map(dst, src); +} +/** + * @brief Applies the GELU derivative to each element of a tile. + * + * @tparam T Tile type. + * @param dst[out] Destination tile where the result is stored. + * @param src[in] Source tile to apply the GELU derivative on. + */ +template +__device__ static inline void dgelu(T &dst, const T &src) { + unary_map(dst, src); +} /** * @brief Copies the elements from one tile to another. * diff --git a/extra/thunder/amd/include/ops/warp/register/tile/mma.cuh b/extra/thunder/amd/include/ops/warp/register/tile/mma.cuh index 39b5353765..7508d97626 100644 --- a/extra/thunder/amd/include/ops/warp/register/tile/mma.cuh +++ b/extra/thunder/amd/include/ops/warp/register/tile/mma.cuh @@ -124,6 +124,28 @@ __device__ static inline void mfma1616128( float2 (&D)[2], )}; } +template +__device__ static inline void mfma1616128_scaled( float2 (&D)[2], + const fp8e4m3_4 (&A)[8], + const fp8e4m3_4 (&B)[8], + const float2 (&C)[2], + const fp8e8m0_4 *scale_a, + const fp8e8m0_4 *scale_b) { + typedef __attribute__((__vector_size__(8 * sizeof(int)))) int intx8_t; + typedef __attribute__((__vector_size__(4 * sizeof(float)))) float floatx4_t; + + *(floatx4_t*)D = {__builtin_amdgcn_mfma_scale_f32_16x16x128_f8f6f4( + *(intx8_t*)A, + *(intx8_t*)B, + *(floatx4_t*)C, + cbsz, // cbsz: 0=fp8(e4m3) A, 1=bf8(e5m2) A + blgp, // blgp: 0=fp8(e4m3) B, 1=bf8(e5m2) B + opsel_a, // opsel_a + *scale_a, // scale_a + opsel_b, // opsel_b + *scale_b // scale_b + )}; +} /** * @brief Base matrix multiply-accumulate operation for row layout. @@ -221,6 +243,46 @@ __device__ static inline void mma_ABt_base(rt_base accumulator. + * @param[in] a The first input rt_base matrix. + * @param[in] b The second input rt_base matrix. + * @param[in] c The input rt_base accumulator matrix. + */ +template +__device__ static inline void mma_ABt_base_scaled(rt_base &d, + const rt_base &a, + const rt_base &b, + const rt_base &c, + const fp8e8m0_4 *scale_a, + const fp8e8m0_4 *scale_b) { + + static_assert(std::is_same_v, "D and C must have the same shape"); + + constexpr int A_rows = A_shape::rows; + constexpr int A_cols = A_shape::cols; + constexpr int B_rows = B_shape::rows; + constexpr int B_cols = B_shape::cols; + + constexpr int A_stride = A_shape::stride; + constexpr int B_stride = B_shape::stride; + static_assert(A_stride == B_stride, "A and B must have the same stride"); + + if constexpr (std::is_same_v && + A_rows == 16 && A_cols == 128 && + B_rows == 16 && B_cols == 128 && + std::is_same_v) { + mfma1616128_scaled(d.data, a.data, b.data, c.data, scale_a, scale_b); + } else { + static_assert(false, "Unsupported shape combination"); + } +} + /** * @brief Base matrix multiply-accumulate operation for row layout with transposed A. * @@ -420,6 +482,62 @@ __device__ static inline void mma_ABt(D &d, } } } + +/** + * @brief Block scaled dot product operation for row layout. + * + * This function performs the dot product operation + * for block-scaled matrices in row layout. + * + * @tparam N The number of row tiles. + * @tparam K The number of column tiles for the A matrix and row tiles for the B matrix. + * @tparam M The number of column tiles for the B matrix. + * @param[out] d The output rt_fl accumulator. + * @param[in] a The first input rt_bf matrix. + * @param[in] b The second input rt_bf matrix in row-major mode. + * @param[in] c The input rt_fl accumulator matrix. + * @param[in] scale_a Pointer to the packed E8M0 scale for the A matrix. + * @param[in] scale_b Pointer to the packed E8M0 scale for the B matrix. + */ +template +__device__ static inline void mma_ABt_scaled(D &d, + const A &a, + const B &b, + const C &c, + const fp8e8m0_4 *scale_a, + const fp8e8m0_4 *scale_b) { + + static_assert(D::rows == A::rows && D::cols == B::rows); + static_assert(A::cols == B::cols); + static_assert(D::rows == C::rows && D::cols == C::cols); + + static_assert( + (std::is_same_v && std::is_same_v && + std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v && + std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v && + std::is_same_v && std::is_same_v) + ); + + [&](std::index_sequence) { + ([&]() { + [&](std::index_sequence) { + ([&]() { + mma_ABt_base_scaled( + d.tiles[N][M], + a.tiles[N][0], + b.tiles[M][0], + c.tiles[N][M], + scale_a, + scale_b + ); + }.template operator()(), ...); + }(std::make_index_sequence{}); + }.template operator()(), ...); + }(std::make_index_sequence{}); +} + /** * @brief Matrix multiply-accumulate operation with transposed A. * diff --git a/extra/thunder/amd/include/ops/warp/register/tile/scales.cuh b/extra/thunder/amd/include/ops/warp/register/tile/scales.cuh new file mode 100644 index 0000000000..b691811eed --- /dev/null +++ b/extra/thunder/amd/include/ops/warp/register/tile/scales.cuh @@ -0,0 +1,74 @@ +/** + * @file + * @brief MXFP8 block scale loading and packing utilities. + * + * Provides functions for staging E8M0 block scales in LDS and packing them + * into fp8e8m0_4 registers for use with scaled MFMA instructions. + */ + +#pragma once + +#include "../../../../common/common.cuh" + +namespace kittens { + +/** + * @brief Load iteration-major packed E8M0 scales from global memory into LDS. + * + * First 256 threads each load one uint32 (4 packed E8M0 bytes) for A and B. + * A scales are placed at smem[0..1023], B scales at smem[1024..2047]. + * + * @param smem_scales LDS buffer, must be >= 2048 bytes. + * @param scale_A_iter Iteration-major A scales: [k_iter * M + row] as uint32. + * @param scale_B_iter Iteration-major B scales: [k_iter * N + row] as uint32. + * @param block_m Starting row offset for A within the current block. + * @param block_n Starting row offset for B within the current block. + * @param k_iter Current K iteration index. + * @param M_dim M dimension of the matrix. + * @param N_dim N dimension of the matrix. + */ +__device__ __forceinline__ void load_scales_to_lds( + uint8_t *smem_scales, + const uint32_t *__restrict__ scale_A_iter, + const uint32_t *__restrict__ scale_B_iter, + int block_m, int block_n, int k_iter, int M_dim, int N_dim) { + int tid = threadIdx.x; + if (tid < 256) { + uint32_t sa = scale_A_iter[k_iter * M_dim + block_m + tid]; + uint32_t sb = scale_B_iter[k_iter * N_dim + block_n + tid]; + *(uint32_t *)&smem_scales[tid * 4] = sa; + *(uint32_t *)&smem_scales[1024 + tid * 4] = sb; + } +} + +/** + * @brief Pack 4 E8M0 scale bytes from LDS into one fp8e8m0_4 register. + * + * Each lane (r16 = laneid%16, k_sub = laneid/16) loads 4 dwords from + * consecutive 16-row groups, then uses v_perm_b32 to extract byte k_sub + * from each, producing the packed scale register for scaled MFMA. + * + * @param smem_scales LDS pointer to scale region. + * @param lds_base Byte offset within smem_scales (0 for A, 1024 for B). + * @param row_offset Starting row within the scale region (warp's tile offset). + * @return fp8e8m0_4 with 4 scale bytes packed for MFMA opsel. + */ +__device__ __forceinline__ fp8e8m0_4 pack_scales( + const uint8_t *smem_scales, int lds_base, int row_offset) { + int lid = laneid(); + int r16 = lid % 16; + int k_sub = lid / 16; + + const uint32_t *s32 = (const uint32_t *)(smem_scales + lds_base); + uint32_t w0 = s32[row_offset + 0 * 16 + r16]; + uint32_t w1 = s32[row_offset + 1 * 16 + r16]; + uint32_t w2 = s32[row_offset + 2 * 16 + r16]; + uint32_t w3 = s32[row_offset + 3 * 16 + r16]; + + uint32_t sel = 0x0C0C0000u | (k_sub << 8) | (4u + k_sub); + uint32_t lo = __builtin_amdgcn_perm(w0, w1, sel); + uint32_t hi = __builtin_amdgcn_perm(w2, w3, sel); + + return (fp8e8m0_4)(lo | (hi << 16)); +} +} // namespace kittens diff --git a/extra/thunder/amd/include/ops/warp/register/tile/tile.cuh b/extra/thunder/amd/include/ops/warp/register/tile/tile.cuh index 588226a5ac..5fa0ca3894 100644 --- a/extra/thunder/amd/include/ops/warp/register/tile/tile.cuh +++ b/extra/thunder/amd/include/ops/warp/register/tile/tile.cuh @@ -9,5 +9,6 @@ #include "maps.cuh" #include "reductions.cuh" #include "mma.cuh" +#include "scales.cuh" #include "assembly/tile.cuh" diff --git a/extra/thunder/amd/include/types/global/gl.cuh b/extra/thunder/amd/include/types/global/gl.cuh index 6fffa098f6..e66bcd6ee2 100644 --- a/extra/thunder/amd/include/types/global/gl.cuh +++ b/extra/thunder/amd/include/types/global/gl.cuh @@ -68,11 +68,11 @@ struct gl { } __host__ __device__ inline gl(const gl &other) : raw_ptr(other.raw_ptr), batch_internal(other.batch_internal), depth_internal(other.depth_internal), rows_internal(other.rows_internal), cols_internal(other.cols_internal), tma_descs(other.tma_descs) {} - __device__ inline T& operator[](const coord &idx) const { // yes I am abusing the const qualifier here a bit. - return raw_ptr[((idx.b*depth() + idx.d)*rows() + idx.r)*cols() + idx.c]; + __device__ inline T& operator[](const coord &idx) const { + return raw_ptr[((int64_t(idx.b)*depth() + idx.d)*rows() + idx.r)*cols() + idx.c]; } - __device__ inline int idx(const coord &idx) const { - return ((idx.b*depth() + idx.d)*rows() + idx.r)*cols() + idx.c; + __device__ inline int64_t idx(const coord &idx) const { + return ((int64_t(idx.b)*depth() + idx.d)*rows() + idx.r)*cols() + idx.c; } template __device__ inline size_t shape() const { static_assert(axis==0 || axis==1 || axis==2 || axis==3, "Axis must be 0, 1, 2, or 3."); diff --git a/extra/thunder/amd/include/types/register/art_base.cuh b/extra/thunder/amd/include/types/register/art_base.cuh index 151d8a3e02..deb7e3996a 100644 --- a/extra/thunder/amd/include/types/register/art_base.cuh +++ b/extra/thunder/amd/include/types/register/art_base.cuh @@ -56,7 +56,9 @@ using register_range = _register_range; ///< Register range for this tile. static_assert( - std::is_same_v || std::is_same_v || std::is_same_v, + std::is_same_v || std::is_same_v || std::is_same_v || + std::is_same_v || std::is_same_v || + std::is_same_v || std::is_same_v, "art_base was provided an unsupported type." ); diff --git a/extra/thunder/amd/include/types/register/rt_shape.cuh b/extra/thunder/amd/include/types/register/rt_shape.cuh index 7514e73140..816a67eee7 100644 --- a/extra/thunder/amd/include/types/register/rt_shape.cuh +++ b/extra/thunder/amd/include/types/register/rt_shape.cuh @@ -37,12 +37,12 @@ using rt_16x32_4 = rt_shape<16, 32, 4>; using rt_16x128 = rt_shape<16, 128, 16>; template -concept all = std::is_same_v || - std::is_same_v || - std::is_same_v || - std::is_same_v || - std::is_same_v || - std::is_same_v || +concept all = std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || std::is_same_v || std::is_same_v; @@ -59,4 +59,4 @@ concept all = std::is_same_v || template<> struct transpose { using type = rt_32x16_4; }; } // namespace rt_shape } // namespace ducks -} // namespace kittens \ No newline at end of file +} // namespace kittens diff --git a/extra/thunder/amd/include/types/shared/sv.cuh b/extra/thunder/amd/include/types/shared/sv.cuh index c4a687856b..439c5b9e98 100644 --- a/extra/thunder/amd/include/types/shared/sv.cuh +++ b/extra/thunder/amd/include/types/shared/sv.cuh @@ -90,5 +90,6 @@ concept all = requires { template using sv_bf = sv; template using sv_hf = sv; template using sv_fl = sv; +template using sv_fp8e4m3 = sv; } // namespace kittens \ No newline at end of file