ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-17 09:08:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

[rocm-libraries] ROCm/rocm-libraries#7852 (commit 8f6a245)

Browse Source 
feat: [CK Tile] Adding gfx1250 wrappers for dense and scale
 builtins (#7852)

## Motivation

This PR is part of the [WMMA/MFMA] unification work. It adds all the
necessary dense and scale MMA gfx1250 builtins as amdgcn_mma structs.

## Technical Details

This change adds 5 new specializations for WMMA gfx1250 scale builtins
and 25 new for dense. On top of that we:
- dispatch between scale and scale16 builtin call
- add scale abstraction for WMMA
- small refactor for scale traits (data type to flag will also be needed
for dense builtins with packed data types)
- fix layout test (we store scale values in every byte)
- create gfx1250 target family

## Test Plan

All the new wrappers were added to the test suite in
`test_amdgcn_mma_layout.inc`.

## Test Result

Test pass locally, waiting for the CI.

## Submission Checklist

- [x] Look over the contributing guidelines at
https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

Closes https://github.com/ROCm/rocm-libraries/issues/8957
This commit is contained in:
Wojciech Laskowski
2026-07-14 06:51:40 +00:00
committed by assistant-librarian[bot]
parent cb859854a7
commit 327dd55f46
14 changed files with 1630 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
include/ck_tile/core.hpp
View File

@@ -25,6 +25,7 @@
#include "ck_tile/core/arch/mma/mfma/mfma_traits.hpp"
#include "ck_tile/core/arch/mma/mfma/mfma_transforms.hpp"
#include "ck_tile/core/arch/mma/mma.hpp"
#include "ck_tile/core/arch/mma/mma_data_format.hpp"
#include "ck_tile/core/arch/mma/mma_op_family.hpp"
#include "ck_tile/core/arch/mma/mma_pipeline.hpp"
#include "ck_tile/core/arch/mma/mma_selector.hpp"
@@ -38,6 +39,8 @@
#include "ck_tile/core/arch/mma/scale/scale_selector.hpp"
#include "ck_tile/core/arch/mma/scale/scale_traits.hpp"
#include "ck_tile/core/arch/mma/scale/scale_transforms.hpp"
#include "ck_tile/core/arch/mma/scale/wmma/scale_gfx12.hpp"
#include "ck_tile/core/arch/mma/scale/wmma/selector.hpp"
#include "ck_tile/core/arch/mma/sparse/mfma/selector.hpp"
#include "ck_tile/core/arch/mma/sparse/mfma/sparse_gfx9.hpp"
#include "ck_tile/core/arch/mma/sparse/sparse.hpp"

71
include/ck_tile/core/arch/arch.hpp
View File

@@ -165,6 +165,10 @@ enum struct amdgcn_target_family_id
GFX10_3 = 0x10,
GFX11 = 0x11,
GFX12 = 0x12,
// GFX1250 is its own standalone family. Although it shares the RDNA architecture with the
// GFX12 family, its MMA builtins and data-type ABI differ, so it must not be treated as a
// GFX12-family device (which would incorrectly enable the legacy GFX12 WMMA specializations).
GFX1250 = 0x1250,
HOST = 0x00,
};
@@ -177,6 +181,7 @@ CK_TILE_HOST_DEVICE constexpr const char* to_string(amdgcn_target_family_id fami
case amdgcn_target_family_id::GFX10_3: return "GFX10_3";
case amdgcn_target_family_id::GFX11: return "GFX11";
case amdgcn_target_family_id::GFX12: return "GFX12";
case amdgcn_target_family_id::GFX1250: return "GFX1250";
case amdgcn_target_family_id::HOST: return "HOST";
}
__builtin_unreachable();
@@ -272,6 +277,15 @@ static constexpr auto make_amdgcn_gfx12_target()
amdgcn_target_wave_size_id::WAVE32>{};
}
template <amdgcn_target_id targetId>
static constexpr auto make_amdgcn_gfx1250_target()
{
return amdgcn_target<targetId,
amdgcn_target_family_id::GFX1250,
amdgcn_target_arch_id::RDNA,
amdgcn_target_wave_size_id::WAVE32>{};
}
template <typename CompilerTarget, amdgcn_target_id... TargetIds>
static constexpr auto is_target_id_any_of()
{
@@ -308,6 +322,12 @@ static constexpr bool is_target_family_gfx12()
return CompilerTarget::FAMILY_ID == amdgcn_target_family_id::GFX12;
}
template <typename CompilerTarget>
static constexpr bool is_target_family_gfx1250()
{
return CompilerTarget::FAMILY_ID == amdgcn_target_family_id::GFX1250;
}
template <typename CompilerTarget>
static constexpr bool is_target_arch_cdna()
{
@@ -362,6 +382,13 @@ static constexpr bool is_target_wave_size_64()
} \
else
#define MAP_COMPILER_STATE_TO_GFX1250_TARGET(COMPILER_STATE, TARGET_ID) \
if constexpr(amdgcn_compiler_target_state::COMPILER_STATE) \
{ \
return make_amdgcn_gfx1250_target<amdgcn_target_id::TARGET_ID>(); \
} \
else
/**
* @brief Returns the amdgcn_target of the current compiler pass.
* @note This is where we tie the compiler state to our internal target architecture representation
@@ -393,7 +420,7 @@ constexpr auto get_compiler_target()
MAP_COMPILER_STATE_TO_GFX12_TARGET(CK_TILE_ARCH_GFX1200, GFX1200);
MAP_COMPILER_STATE_TO_GFX12_TARGET(CK_TILE_ARCH_GFX1201, GFX1201);
MAP_COMPILER_STATE_TO_GFX12_TARGET(CK_TILE_ARCH_GFX12_GENERIC, GFX12_GENERIC);
MAP_COMPILER_STATE_TO_GFX12_TARGET(CK_TILE_ARCH_GFX1250, GFX1250);
MAP_COMPILER_STATE_TO_GFX1250_TARGET(CK_TILE_ARCH_GFX1250, GFX1250);
// Return HOST by default
if constexpr(amdgcn_compiler_target_state::CK_TILE_HOST_COMPILE)
@@ -449,7 +476,7 @@ static constexpr auto getCMakeCompilerTarget()
}
else if constexpr(id == amdgcn_target_id::GFX1250)
{
return make_amdgcn_gfx12_target<id>(); // TODO: This should not be a GFX12 target.
return make_amdgcn_gfx1250_target<id>();
}
else
{
@@ -472,6 +499,7 @@ static constexpr auto getCMakeCompilerTarget()
#undef MAP_COMPILER_STATE_TO_GFX10_3_TARGET
#undef MAP_COMPILER_STATE_TO_GFX11_TARGET
#undef MAP_COMPILER_STATE_TO_GFX12_TARGET
#undef MAP_COMPILER_STATE_TO_GFX1250_TARGET
// Sanity check: device compile must have a valid target architecture
static_assert(!amdgcn_compiler_target_state::CK_TILE_DEVICE_COMPILE ||
@@ -610,6 +638,13 @@ template <typename CompilerTarget>
using enable_if_target_family_gfx12_t =
enable_if_target_family_id_t<CompilerTarget, amdgcn_target_family_id::GFX12>;
/**
* @brief SFINAE enabler for GFX1250 target
* @tparam CompilerTarget The compiler target to check
*/
template <typename CompilerTarget>
using enable_if_target_gfx1250_t = enable_if_target_id_t<CompilerTarget, amdgcn_target_id::GFX1250>;
/**
* @brief SFINAE enabler for CDNA architecture targets
* @tparam CompilerTarget The compiler target to check
@@ -699,6 +734,14 @@ static constexpr auto make_amdgcn_gfx12_target(amdgcn_target_id targetId)
.WAVE_SIZE_ID = amdgcn_target_wave_size_id::WAVE32};
}
static constexpr auto make_amdgcn_gfx1250_target(amdgcn_target_id targetId)
{
return amdgcn_target{.TARGET_ID = targetId,
.FAMILY_ID = amdgcn_target_family_id::GFX1250,
.ARCH_ID = amdgcn_target_arch_id::RDNA,
.WAVE_SIZE_ID = amdgcn_target_wave_size_id::WAVE32};
}
static constexpr bool is_target_family_gfx9(amdgcn_target target)
{
return target.FAMILY_ID == amdgcn_target_family_id::GFX9;
@@ -719,6 +762,11 @@ static constexpr bool is_target_family_gfx12(amdgcn_target target)
return target.FAMILY_ID == amdgcn_target_family_id::GFX12;
}
static constexpr bool is_target_family_gfx1250(amdgcn_target target)
{
return target.FAMILY_ID == amdgcn_target_family_id::GFX1250;
}
static constexpr bool is_target_arch_cdna(amdgcn_target target)
{
return target.ARCH_ID == amdgcn_target_arch_id::CDNA;
@@ -764,6 +812,12 @@ static constexpr bool is_target_wave_size_64(amdgcn_target target)
return make_amdgcn_gfx12_target(amdgcn_target_id::TARGET_ID); \
}
#define MAP_COMPILER_STATE_TO_GFX1250_TARGET(COMPILER_STATE, TARGET_ID) \
if constexpr(amdgcn_compiler_target_state::COMPILER_STATE) \
{ \
return make_amdgcn_gfx1250_target(amdgcn_target_id::TARGET_ID); \
}
/*! @brief Returns the amdgcn_target of the current compiler pass.
* @note This is where we tie the compiler state to our internal target architecture representation
* at compile time.
@@ -794,7 +848,7 @@ CK_TILE_HOST_DEVICE constexpr auto get_compiler_target()
MAP_COMPILER_STATE_TO_GFX12_TARGET(CK_TILE_ARCH_GFX1200, GFX1200);
MAP_COMPILER_STATE_TO_GFX12_TARGET(CK_TILE_ARCH_GFX1201, GFX1201);
MAP_COMPILER_STATE_TO_GFX12_TARGET(CK_TILE_ARCH_GFX12_GENERIC, GFX12_GENERIC);
MAP_COMPILER_STATE_TO_GFX12_TARGET(CK_TILE_ARCH_GFX1250, GFX1250);
MAP_COMPILER_STATE_TO_GFX1250_TARGET(CK_TILE_ARCH_GFX1250, GFX1250);
// Default to HOST
return amdgcn_target{};
@@ -805,6 +859,7 @@ CK_TILE_HOST_DEVICE constexpr auto get_compiler_target()
#undef MAP_COMPILER_STATE_TO_GFX10_3_TARGET
#undef MAP_COMPILER_STATE_TO_GFX11_TARGET
#undef MAP_COMPILER_STATE_TO_GFX12_TARGET
#undef MAP_COMPILER_STATE_TO_GFX1250_TARGET
// Sanity check: device compile must have a valid target architecture
static_assert(!amdgcn_compiler_target_state::CK_TILE_DEVICE_COMPILE ||
@@ -844,6 +899,13 @@ static_assert(!amdgcn_compiler_target_state::CK_TILE_HOST_COMPILE ||
} \
else
#define MAP_HIP_DEVICE_PROP_GCN_ARCH_NAME_STRING_TO_GFX1250_TARGET(NAME_STRING, TARGET_ID) \
if constexpr(str.find(NAME_STRING) != std::string::npos) \
{ \
return make_amdgcn_gfx1250_target(amdgcn_target_id::TARGET_ID); \
} \
else
/**
* @brief Converts a lower-case string to the corresponding amdgcn_target_arch_id value.
* Returns amdgcn_target_arch_id::HOST if no match is found.
@@ -877,7 +939,7 @@ CK_TILE_HOST auto hip_device_prop_gcn_arch_name_to_amdgcn_target(char const* tes
MAP_HIP_DEVICE_PROP_GCN_ARCH_NAME_STRING_TO_GFX12_TARGET("gfx1200", GFX1200);
MAP_HIP_DEVICE_PROP_GCN_ARCH_NAME_STRING_TO_GFX12_TARGET("gfx1201", GFX1201);
MAP_HIP_DEVICE_PROP_GCN_ARCH_NAME_STRING_TO_GFX12_TARGET("gfx12_generic", GFX12_GENERIC);
MAP_HIP_DEVICE_PROP_GCN_ARCH_NAME_STRING_TO_GFX12_TARGET("gfx1250", GFX1250);
MAP_HIP_DEVICE_PROP_GCN_ARCH_NAME_STRING_TO_GFX1250_TARGET("gfx1250", GFX1250);
// Default case
return amdgcn_target{};
}
@@ -886,6 +948,7 @@ CK_TILE_HOST auto hip_device_prop_gcn_arch_name_to_amdgcn_target(char const* tes
#undef MAP_HIP_DEVICE_PROP_GCN_ARCH_NAME_STRING_TO_GFX10_3_TARGET
#undef MAP_HIP_DEVICE_PROP_GCN_ARCH_NAME_STRING_TO_GFX11_TARGET
#undef MAP_HIP_DEVICE_PROP_GCN_ARCH_NAME_STRING_TO_GFX12_TARGET
#undef MAP_HIP_DEVICE_PROP_GCN_ARCH_NAME_STRING_TO_GFX1250_TARGET
/**
* @brief SFINAE enabler for a compiler target if the target id is in the list of supported target

53
include/ck_tile/core/arch/mma/mma_data_format.hpp Normal file
View File

@@ -0,0 +1,53 @@
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
#pragma once
#include "ck_tile/core/numeric/float8.hpp"
#include "ck_tile/core/numeric/pk_f6.hpp"
#include "ck_tile/core/numeric/pk_fp4.hpp"
#include <cstdint>
namespace ck_tile::core::arch::mma {
/**
* @brief Maps a data type to its hardware matrix format code used by f8f6f4 builtins.
*/
template <typename T>
struct PackedDataTypeToFlag;
template <>
struct PackedDataTypeToFlag<fp8_t> // e4m3
{
static constexpr int32_t value = 0;
};
template <>
struct PackedDataTypeToFlag<bf8_t> // e5m2
{
static constexpr int32_t value = 1;
};
template <>
struct PackedDataTypeToFlag<pk_fp6x16_t> // e2m3
{
static constexpr int32_t value = 2;
};
template <>
struct PackedDataTypeToFlag<pk_bf6x16_t> // e3m2
{
static constexpr int32_t value = 3;
};
template <>
struct PackedDataTypeToFlag<pk_fp4_t> // e2m1
{
static constexpr int32_t value = 4;
};
template <typename T>
inline constexpr int32_t PackedDataTypeToFlag_v = PackedDataTypeToFlag<T>::value;
} // namespace ck_tile::core::arch::mma

10
include/ck_tile/core/arch/mma/mma_op_family.hpp
View File

@@ -18,8 +18,17 @@ enum struct MmaOpFamily
DENSE,
SPARSE,
SCALE,
SCALE16,
};
/**
* @brief Helper to check if an op family is any of the scale families.
*/
CK_TILE_HOST_DEVICE constexpr bool is_scale_op_family(MmaOpFamily f)
{
return f == MmaOpFamily::SCALE || f == MmaOpFamily::SCALE16;
}
/**
* @class is_ctrl_fis_mma_op_of_familylag_of_family
* @brief Meta-function to check if MmaOp is of the specified MmaOpFamily
@@ -58,6 +67,7 @@ CK_TILE_HOST_DEVICE constexpr const char* to_string(MmaOpFamily opFamily)
case MmaOpFamily::DENSE: return "DENSE";
case MmaOpFamily::SPARSE: return "SPARSE";
case MmaOpFamily::SCALE: return "SCALE";
case MmaOpFamily::SCALE16: return "SCALE16";
}
__builtin_unreachable();
}

2
include/ck_tile/core/arch/mma/mma_traits.hpp
View File

@@ -95,7 +95,7 @@ struct MmaOpTraits<amdgcn_mma<ADataType_,
constexpr static bool IsWmma = is_mma_op_wmma_v<MmaOp>;
constexpr static bool IsDense = OpFamily_ == MmaOpFamily::DENSE;
constexpr static bool IsSparse = OpFamily_ == MmaOpFamily::SPARSE;
constexpr static bool IsScale = OpFamily_ == MmaOpFamily::SCALE;
constexpr static bool IsScale = is_scale_op_family(OpFamily_);
constexpr static bool IsSupported =
is_mma_op_supported_v<MmaOp> && OpFamily_ != MmaOpFamily::UNDEFINED;
};

1
include/ck_tile/core/arch/mma/scale/scale.hpp
View File

@@ -5,6 +5,7 @@
// Include scale MFMA traits and architecture-specific implementations
#include "ck_tile/core/arch/mma/scale/mfma/scale_gfx9.hpp"
#include "ck_tile/core/arch/mma/scale/wmma/scale_gfx12.hpp"
#include "ck_tile/core/arch/mma/scale/scale_selector.hpp"
#include "ck_tile/core/arch/mma/scale/scale_traits.hpp"
#include "ck_tile/core/arch/mma/scale/scale_transforms.hpp"

1
include/ck_tile/core/arch/mma/scale/scale_selector.hpp
View File

@@ -4,3 +4,4 @@
#pragma once
#include "ck_tile/core/arch/mma/scale/mfma/selector.hpp"
#include "ck_tile/core/arch/mma/scale/wmma/selector.hpp"

565
include/ck_tile/core/arch/mma/scale/wmma/scale_gfx12.hpp Normal file
View File

@@ -0,0 +1,565 @@
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
#pragma once
#include "ck_tile/core/arch/arch.hpp"
#include "ck_tile/core/arch/mma/amdgcn_mma.hpp"
#include "ck_tile/core/arch/mma/mma_data_format.hpp"
#include "ck_tile/core/arch/mma/mma_op_family.hpp"
#include "ck_tile/core/arch/mma/scale/scale_traits.hpp"
#include "ck_tile/core/arch/mma/wmma/wmma_traits.hpp"
#include "ck_tile/core/config.hpp"
#include "ck_tile/core/numeric/float8.hpp"
#include "ck_tile/core/numeric/pk_f6.hpp"
#include "ck_tile/core/numeric/pk_fp4.hpp"
#include "ck_tile/core/numeric/vector_type.hpp"
#include "ck_tile/core/utility/bit_cast.hpp"
namespace ck_tile::core::arch::mma {
/**
* @struct amdgcn_mma
* @brief Specialization for fp8_t scale32 WMMA on GFX1250.
* @tparam CompilerTarget Current compiler target
*/
template <typename CompilerTarget>
// clang-format off
// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar |
struct amdgcn_mma<fp8_t, fp8_t, fp32_t, 16u, 16u, 128u, CompilerTarget, MmaOpFamily::SCALE, enable_if_target_gfx1250_t<CompilerTarget>>
: amdgcn_mma_base<fp8_t, fp8_t, fp32_t, 16u, 16u, 128u, 32u, 64, 1, 1, 1, 1, 8, 1, WmmaOp, MmaOpFamily::SCALE>
// clang-format on
{
using ScaleType = int32_t;
static constexpr const char* instruction_name =
"__builtin_amdgcn_wmma_scale_f32_16x16x128_f8f6f4";
CK_TILE_DEVICE static CVecType exec(AVecType const& aVec,
BVecType const& bVec,
CVecType const& cVec,
ScaleType scaleA,
ScaleType scaleB)
{
return {__builtin_amdgcn_wmma_scale_f32_16x16x128_f8f6f4(PackedDataTypeToFlag_v<fp8_t>,
bit_cast<int32x16_t>(aVec),
PackedDataTypeToFlag_v<fp8_t>,
bit_cast<int32x16_t>(bVec),
0, // C_mod
cVec,
0, // OPSEL[0]
0, // Scale Type for A
scaleA,
0, // OPSEL[1]
0, // Scale Type for B
scaleB,
false, // matrix_a_reuse
false)}; // matrix_b_reuse
}
};
/**
* @struct amdgcn_mma
* @brief Specialization for bf8_t scale32 WMMA on GFX1250.
* @tparam CompilerTarget Current compiler target
*/
template <typename CompilerTarget>
// clang-format off
// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar |
struct amdgcn_mma<bf8_t, bf8_t, fp32_t, 16u, 16u, 128u, CompilerTarget, MmaOpFamily::SCALE, enable_if_target_gfx1250_t<CompilerTarget>>
: amdgcn_mma_base<bf8_t, bf8_t, fp32_t, 16u, 16u, 128u, 32u, 64, 1, 1, 1, 1, 8, 1, WmmaOp, MmaOpFamily::SCALE>
// clang-format on
{
using ScaleType = int32_t;
static constexpr const char* instruction_name =
"__builtin_amdgcn_wmma_scale_f32_16x16x128_f8f6f4";
CK_TILE_DEVICE static CVecType exec(AVecType const& aVec,
BVecType const& bVec,
CVecType const& cVec,
ScaleType scaleA,
ScaleType scaleB)
{
return {__builtin_amdgcn_wmma_scale_f32_16x16x128_f8f6f4(PackedDataTypeToFlag_v<bf8_t>,
bit_cast<int32x16_t>(aVec),
PackedDataTypeToFlag_v<bf8_t>,
bit_cast<int32x16_t>(bVec),
0, // C_mod
cVec,
0, // OPSEL[0]
0, // Scale Type for A
scaleA,
0, // OPSEL[1]
0, // Scale Type for B
scaleB,
false, // matrix_a_reuse
false)}; // matrix_b_reuse
}
};
/**
* @struct amdgcn_mma
* @brief Specialization for pk_fp6x16_t scale32 WMMA on GFX1250.
* @tparam CompilerTarget Current compiler target
*/
template <typename CompilerTarget>
// clang-format off
// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar |
struct amdgcn_mma<pk_fp6x16_t, pk_fp6x16_t, fp32_t, 16u, 16u, 128u, CompilerTarget, MmaOpFamily::SCALE, enable_if_target_gfx1250_t<CompilerTarget>>
: amdgcn_mma_base<pk_fp6x16_t, pk_fp6x16_t, fp32_t, 16u, 16u, 128u, 32u, 64, 1, 1, 1, 1, 8, 1, WmmaOp, MmaOpFamily::SCALE>
// clang-format on
{
using ScaleType = int32_t;
static constexpr const char* instruction_name =
"__builtin_amdgcn_wmma_scale_f32_16x16x128_f8f6f4";
CK_TILE_DEVICE static CVecType exec(AVecType const& aVec,
BVecType const& bVec,
CVecType const& cVec,
ScaleType scaleA,
ScaleType scaleB)
{
int32x16_t a_padded = {aVec.data[0],
aVec.data[1],
aVec.data[2],
aVec.data[3],
aVec.data[4],
aVec.data[5],
aVec.data[6],
aVec.data[7],
aVec.data[8],
aVec.data[9],
aVec.data[10],
aVec.data[11],
0,
0,
0,
0};
int32x16_t b_padded = {bVec.data[0],
bVec.data[1],
bVec.data[2],
bVec.data[3],
bVec.data[4],
bVec.data[5],
bVec.data[6],
bVec.data[7],
bVec.data[8],
bVec.data[9],
bVec.data[10],
bVec.data[11],
0,
0,
0,
0};
return {
__builtin_amdgcn_wmma_scale_f32_16x16x128_f8f6f4(PackedDataTypeToFlag_v<pk_fp6x16_t>,
a_padded,
PackedDataTypeToFlag_v<pk_fp6x16_t>,
b_padded,
0, // C_mod
cVec,
0, // OPSEL[0]
0, // Scale Type for A
scaleA,
0, // OPSEL[1]
0, // Scale Type for B
scaleB,
false, // matrix_a_reuse
false)}; // matrix_b_reuse
}
};
/**
* @struct amdgcn_mma
* @brief Specialization for pk_bf6x16_t scale32 WMMA on GFX1250.
* @tparam CompilerTarget Current compiler target
*/
template <typename CompilerTarget>
// clang-format off
// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar |
struct amdgcn_mma<pk_bf6x16_t, pk_bf6x16_t, fp32_t, 16u, 16u, 128u, CompilerTarget, MmaOpFamily::SCALE, enable_if_target_gfx1250_t<CompilerTarget>>
: amdgcn_mma_base<pk_bf6x16_t, pk_bf6x16_t, fp32_t, 16u, 16u, 128u, 32u, 64, 1, 1, 1, 1, 8, 1, WmmaOp, MmaOpFamily::SCALE>
// clang-format on
{
using ScaleType = int32_t;
static constexpr const char* instruction_name =
"__builtin_amdgcn_wmma_scale_f32_16x16x128_f8f6f4";
CK_TILE_DEVICE static CVecType exec(AVecType const& aVec,
BVecType const& bVec,
CVecType const& cVec,
ScaleType scaleA,
ScaleType scaleB)
{
int32x16_t a_padded = {aVec.data[0],
aVec.data[1],
aVec.data[2],
aVec.data[3],
aVec.data[4],
aVec.data[5],
aVec.data[6],
aVec.data[7],
aVec.data[8],
aVec.data[9],
aVec.data[10],
aVec.data[11],
0,
0,
0,
0};
int32x16_t b_padded = {bVec.data[0],
bVec.data[1],
bVec.data[2],
bVec.data[3],
bVec.data[4],
bVec.data[5],
bVec.data[6],
bVec.data[7],
bVec.data[8],
bVec.data[9],
bVec.data[10],
bVec.data[11],
0,
0,
0,
0};
return {
__builtin_amdgcn_wmma_scale_f32_16x16x128_f8f6f4(PackedDataTypeToFlag_v<pk_bf6x16_t>,
a_padded,
PackedDataTypeToFlag_v<pk_bf6x16_t>,
b_padded,
0, // C_mod
cVec,
0, // OPSEL[0]
0, // Scale Type for A
scaleA,
0, // OPSEL[1]
0, // Scale Type for B
scaleB,
false, // matrix_a_reuse
false)}; // matrix_b_reuse
}
};
/**
* @struct amdgcn_mma
* @brief Specialization for pk_fp4_t scale32 WMMA on GFX1250.
* @tparam CompilerTarget Current compiler target
*/
template <typename CompilerTarget>
// clang-format off
// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar |
struct amdgcn_mma<pk_fp4_t, pk_fp4_t, fp32_t, 16u, 16u, 128u, CompilerTarget, MmaOpFamily::SCALE, enable_if_target_gfx1250_t<CompilerTarget>>
: amdgcn_mma_base<pk_fp4_t, pk_fp4_t, fp32_t, 16u, 16u, 128u, 32u, 64, 1, 1, 1, 1, 8, 1, WmmaOp, MmaOpFamily::SCALE>
// clang-format on
{
using ScaleType = int32_t;
static constexpr const char* instruction_name =
"__builtin_amdgcn_wmma_scale_f32_16x16x128_f8f6f4";
CK_TILE_DEVICE static CVecType exec(AVecType const& aVec,
BVecType const& bVec,
CVecType const& cVec,
ScaleType scaleA,
ScaleType scaleB)
{
int32x8_t a8 = bit_cast<int32x8_t>(aVec);
int32x8_t b8 = bit_cast<int32x8_t>(bVec);
int32x16_t a_padded = {
a8[0], a8[1], a8[2], a8[3], a8[4], a8[5], a8[6], a8[7], 0, 0, 0, 0, 0, 0, 0, 0};
int32x16_t b_padded = {
b8[0], b8[1], b8[2], b8[3], b8[4], b8[5], b8[6], b8[7], 0, 0, 0, 0, 0, 0, 0, 0};
return {__builtin_amdgcn_wmma_scale_f32_16x16x128_f8f6f4(PackedDataTypeToFlag_v<pk_fp4_t>,
a_padded,
PackedDataTypeToFlag_v<pk_fp4_t>,
b_padded,
0, // C_mod
cVec,
0, // OPSEL[0]
0, // Scale Type for A
scaleA,
0, // OPSEL[1]
0, // Scale Type for B
scaleB,
false, // matrix_a_reuse
false)}; // matrix_b_reuse
}
};
/**
* @struct amdgcn_mma
* @brief Specialization for fp8_t scale16 WMMA on GFX1250.
* @tparam CompilerTarget Current compiler target
*/
template <typename CompilerTarget>
// clang-format off
// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar |
struct amdgcn_mma<fp8_t, fp8_t, fp32_t, 16u, 16u, 128u, CompilerTarget, MmaOpFamily::SCALE16, enable_if_target_gfx1250_t<CompilerTarget>>
: amdgcn_mma_base<fp8_t, fp8_t, fp32_t, 16u, 16u, 128u, 32u, 64, 1, 1, 1, 1, 8, 1, WmmaOp, MmaOpFamily::SCALE16>
// clang-format on
{
using ScaleType = int64_t;
static constexpr const char* instruction_name =
"__builtin_amdgcn_wmma_scale16_f32_16x16x128_f8f6f4";
CK_TILE_DEVICE static CVecType exec(AVecType const& aVec,
BVecType const& bVec,
CVecType const& cVec,
ScaleType scaleA,
ScaleType scaleB)
{
return {__builtin_amdgcn_wmma_scale16_f32_16x16x128_f8f6f4(PackedDataTypeToFlag_v<fp8_t>,
bit_cast<int32x16_t>(aVec),
PackedDataTypeToFlag_v<fp8_t>,
bit_cast<int32x16_t>(bVec),
0, // C_mod
cVec,
0, // OPSEL[0]
0, // Scale Type for A
scaleA,
0, // OPSEL[1]
0, // Scale Type for B
scaleB,
false, // matrix_a_reuse
false)}; // matrix_b_reuse
}
};
/**
* @struct amdgcn_mma
* @brief Specialization for bf8_t scale16 WMMA on GFX1250.
* @tparam CompilerTarget Current compiler target
*/
template <typename CompilerTarget>
// clang-format off
// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar |
struct amdgcn_mma<bf8_t, bf8_t, fp32_t, 16u, 16u, 128u, CompilerTarget, MmaOpFamily::SCALE16, enable_if_target_gfx1250_t<CompilerTarget>>
: amdgcn_mma_base<bf8_t, bf8_t, fp32_t, 16u, 16u, 128u, 32u, 64, 1, 1, 1, 1, 8, 1, WmmaOp, MmaOpFamily::SCALE16>
// clang-format on
{
using ScaleType = int64_t;
static constexpr const char* instruction_name =
"__builtin_amdgcn_wmma_scale16_f32_16x16x128_f8f6f4";
CK_TILE_DEVICE static CVecType exec(AVecType const& aVec,
BVecType const& bVec,
CVecType const& cVec,
ScaleType scaleA,
ScaleType scaleB)
{
return {__builtin_amdgcn_wmma_scale16_f32_16x16x128_f8f6f4(PackedDataTypeToFlag_v<bf8_t>,
bit_cast<int32x16_t>(aVec),
PackedDataTypeToFlag_v<bf8_t>,
bit_cast<int32x16_t>(bVec),
0, // C_mod
cVec,
0, // OPSEL[0]
0, // Scale Type for A
scaleA,
0, // OPSEL[1]
0, // Scale Type for B
scaleB,
false, // matrix_a_reuse
false)}; // matrix_b_reuse
}
};
/**
* @struct amdgcn_mma
* @brief Specialization for pk_fp6x16_t scale16 WMMA on GFX1250.
* @tparam CompilerTarget Current compiler target
*/
template <typename CompilerTarget>
// clang-format off
// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar |
struct amdgcn_mma<pk_fp6x16_t, pk_fp6x16_t, fp32_t, 16u, 16u, 128u, CompilerTarget, MmaOpFamily::SCALE16, enable_if_target_gfx1250_t<CompilerTarget>>
: amdgcn_mma_base<pk_fp6x16_t, pk_fp6x16_t, fp32_t, 16u, 16u, 128u, 32u, 64, 1, 1, 1, 1, 8, 1, WmmaOp, MmaOpFamily::SCALE16>
// clang-format on
{
using ScaleType = int64_t;
static constexpr const char* instruction_name =
"__builtin_amdgcn_wmma_scale16_f32_16x16x128_f8f6f4";
CK_TILE_DEVICE static CVecType exec(AVecType const& aVec,
BVecType const& bVec,
CVecType const& cVec,
ScaleType scaleA,
ScaleType scaleB)
{
int32x16_t a_padded = {aVec.data[0],
aVec.data[1],
aVec.data[2],
aVec.data[3],
aVec.data[4],
aVec.data[5],
aVec.data[6],
aVec.data[7],
aVec.data[8],
aVec.data[9],
aVec.data[10],
aVec.data[11],
0,
0,
0,
0};
int32x16_t b_padded = {bVec.data[0],
bVec.data[1],
bVec.data[2],
bVec.data[3],
bVec.data[4],
bVec.data[5],
bVec.data[6],
bVec.data[7],
bVec.data[8],
bVec.data[9],
bVec.data[10],
bVec.data[11],
0,
0,
0,
0};
return {
__builtin_amdgcn_wmma_scale16_f32_16x16x128_f8f6f4(PackedDataTypeToFlag_v<pk_fp6x16_t>,
a_padded,
PackedDataTypeToFlag_v<pk_fp6x16_t>,
b_padded,
0, // C_mod
cVec,
0, // OPSEL[0]
0, // Scale Type for A
scaleA,
0, // OPSEL[1]
0, // Scale Type for B
scaleB,
false, // matrix_a_reuse
false)}; // matrix_b_reuse
}
};
/**
* @struct amdgcn_mma
* @brief Specialization for pk_bf6x16_t scale16 WMMA on GFX1250.
* @tparam CompilerTarget Current compiler target
*/
template <typename CompilerTarget>
// clang-format off
// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar |
struct amdgcn_mma<pk_bf6x16_t, pk_bf6x16_t, fp32_t, 16u, 16u, 128u, CompilerTarget, MmaOpFamily::SCALE16, enable_if_target_gfx1250_t<CompilerTarget>>
: amdgcn_mma_base<pk_bf6x16_t, pk_bf6x16_t, fp32_t, 16u, 16u, 128u, 32u, 64, 1, 1, 1, 1, 8, 1, WmmaOp, MmaOpFamily::SCALE16>
// clang-format on
{
using ScaleType = int64_t;
static constexpr const char* instruction_name =
"__builtin_amdgcn_wmma_scale16_f32_16x16x128_f8f6f4";
CK_TILE_DEVICE static CVecType exec(AVecType const& aVec,
BVecType const& bVec,
CVecType const& cVec,
ScaleType scaleA,
ScaleType scaleB)
{
int32x16_t a_padded = {aVec.data[0],
aVec.data[1],
aVec.data[2],
aVec.data[3],
aVec.data[4],
aVec.data[5],
aVec.data[6],
aVec.data[7],
aVec.data[8],
aVec.data[9],
aVec.data[10],
aVec.data[11],
0,
0,
0,
0};
int32x16_t b_padded = {bVec.data[0],
bVec.data[1],
bVec.data[2],
bVec.data[3],
bVec.data[4],
bVec.data[5],
bVec.data[6],
bVec.data[7],
bVec.data[8],
bVec.data[9],
bVec.data[10],
bVec.data[11],
0,
0,
0,
0};
return {
__builtin_amdgcn_wmma_scale16_f32_16x16x128_f8f6f4(PackedDataTypeToFlag_v<pk_bf6x16_t>,
a_padded,
PackedDataTypeToFlag_v<pk_bf6x16_t>,
b_padded,
0, // C_mod
cVec,
0, // OPSEL[0]
0, // Scale Type for A
scaleA,
0, // OPSEL[1]
0, // Scale Type for B
scaleB,
false, // matrix_a_reuse
false)}; // matrix_b_reuse
}
};
/**
* @struct amdgcn_mma
* @brief Specialization for pk_fp4_t scale16 WMMA on GFX1250.
* @tparam CompilerTarget Current compiler target
*/
template <typename CompilerTarget>
// clang-format off
// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar |
struct amdgcn_mma<pk_fp4_t, pk_fp4_t, fp32_t, 16u, 16u, 128u, CompilerTarget, MmaOpFamily::SCALE16, enable_if_target_gfx1250_t<CompilerTarget>>
: amdgcn_mma_base<pk_fp4_t, pk_fp4_t, fp32_t, 16u, 16u, 128u, 32u, 64, 1, 1, 1, 1, 8, 1, WmmaOp, MmaOpFamily::SCALE16>
// clang-format on
{
using ScaleType = int64_t;
static constexpr const char* instruction_name =
"__builtin_amdgcn_wmma_scale16_f32_16x16x128_f8f6f4";
CK_TILE_DEVICE static CVecType exec(AVecType const& aVec,
BVecType const& bVec,
CVecType const& cVec,
ScaleType scaleA,
ScaleType scaleB)
{
int32x8_t a8 = bit_cast<int32x8_t>(aVec);
int32x8_t b8 = bit_cast<int32x8_t>(bVec);
int32x16_t a_padded = {
a8[0], a8[1], a8[2], a8[3], a8[4], a8[5], a8[6], a8[7], 0, 0, 0, 0, 0, 0, 0, 0};
int32x16_t b_padded = {
b8[0], b8[1], b8[2], b8[3], b8[4], b8[5], b8[6], b8[7], 0, 0, 0, 0, 0, 0, 0, 0};
return {__builtin_amdgcn_wmma_scale16_f32_16x16x128_f8f6f4(PackedDataTypeToFlag_v<pk_fp4_t>,
a_padded,
PackedDataTypeToFlag_v<pk_fp4_t>,
b_padded,
0, // C_mod
cVec,
0, // OPSEL[0]
0, // Scale Type for A
scaleA,
0, // OPSEL[1]
0, // Scale Type for B
scaleB,
false, // matrix_a_reuse
false)}; // matrix_b_reuse
}
};
} // namespace ck_tile::core::arch::mma

78
include/ck_tile/core/arch/mma/scale/wmma/selector.hpp Normal file
View File

@@ -0,0 +1,78 @@
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
#pragma once
#include "ck_tile/core/arch/arch.hpp"
#include "ck_tile/core/arch/mma/amdgcn_mma.hpp"
#include "ck_tile/core/arch/mma/mma_selector.hpp"
#include "ck_tile/core/arch/mma/mma_traits.hpp"
#include "ck_tile/core/arch/mma/scale/scale_traits.hpp"
#include "ck_tile/core/arch/mma/scale/wmma/scale_gfx12.hpp"
#include <cstdint>
#include <type_traits>
namespace ck_tile::core::arch::mma {
/**
* @struct MmaDefaultSelector
* @brief Implements the RDNA default MMA selector strategy for scale WMMA on GFX1250.
* If no supported instruction is found, falls back to an unsupported pass-through implementation.
* @tparam ADataType Data type of matrix A
* @tparam BDataType Data type of matrix B
* @tparam CDataType Data type of the accumulator
* @tparam WaveTileM Size of the M dimension of the WaveTile to decompose
* @tparam WaveTileN Size of the N dimension of the WaveTile to decompose
* @tparam WaveTileK Size of the K dimension of the WaveTile to decompose
* @tparam CompilerTarget The compiler target
* @tparam OpFamily The MMA operation family (SCALE or SCALE16)
*/
template <typename ADataType,
typename BDataType,
typename CDataType,
std::uint32_t WaveTileM,
std::uint32_t WaveTileN,
std::uint32_t WaveTileK,
typename CompilerTarget,
MmaOpFamily OpFamily>
// TODO: c++20 amdgcn_target_arch_id CompilerTarget>
// TODO: c++20 requires
struct MmaDefaultSelector<ADataType,
BDataType,
CDataType,
WaveTileM,
WaveTileN,
WaveTileK,
CompilerTarget,
OpFamily,
enable_if_all<enable_if_target_gfx1250_t<CompilerTarget>,
std::enable_if_t<OpFamily == MmaOpFamily::SCALE ||
OpFamily == MmaOpFamily::SCALE16>>>
{
private:
// Candidate 16x16x128 scale WMMA operation
using CandidateOp16x16 =
amdgcn_mma<ADataType, BDataType, CDataType, 16u, 16u, 128u, CompilerTarget, OpFamily>;
// Fall back to the unsupported pass-through implementation.
using UnsupportedOp = amdgcn_mma<ADataType,
BDataType,
CDataType,
WaveTileM,
WaveTileN,
WaveTileK,
amdgcn_target<>,
MmaOpFamily::UNDEFINED>;
// Check if the candidate is supported for the given WaveTile sizes
static constexpr bool IsSupported16x16 =
MmaOpTraits<CandidateOp16x16>::IsSupported && (WaveTileM % CandidateOp16x16::kM == 0u) &&
(WaveTileN % CandidateOp16x16::kN == 0u) && (WaveTileK % CandidateOp16x16::kK == 0u);
public:
// Select the largest supported WMMA operation for the given WaveTile shape
using SelectedOp = std::conditional_t<IsSupported16x16, CandidateOp16x16, UnsupportedOp>;
};
} // namespace ck_tile::core::arch::mma

746
include/ck_tile/core/arch/mma/wmma/wmma_gfx12.hpp
View File

@@ -5,6 +5,7 @@
#include "ck_tile/core/arch/arch.hpp"
#include "ck_tile/core/arch/mma/amdgcn_mma.hpp"
#include "ck_tile/core/arch/mma/mma_data_format.hpp"
#include "ck_tile/core/arch/mma/mma_op_family.hpp"
#include "ck_tile/core/arch/mma/wmma/wmma_traits.hpp"
#include "ck_tile/core/config.hpp"
@@ -335,4 +336,749 @@ struct amdgcn_mma<pk_int4_t, pk_int4_t, int32_t, 16u, 16u, 32u, CompilerTarget,
}
};
// GFX1250
/**
* @struct amdgcn_mma
* @brief Specialization of amdgcn_mma for fp32_t, fp32_t, fp32_t MMA operation on GFX1250
* architecture.
* @tparam CompilerTarget Current compiler target
*/
// TODO: c++20 template <amdgcn_target CompilerTarget>
// TODO: c++20 requires
template <typename CompilerTarget>
// clang-format off
// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar |
struct amdgcn_mma<fp32_t, fp32_t, fp32_t, 16u, 16u, 4u, CompilerTarget, MmaOpFamily::DENSE, enable_if_target_gfx1250_t<CompilerTarget>>
: amdgcn_mma_base<fp32_t, fp32_t, fp32_t, 16u, 16u, 4u, 32u, 2, 1, 1, 1, 1, 8, 1, WmmaOp, MmaOpFamily::DENSE>
// clang-format on
{
static constexpr const char* instruction_name = "__builtin_amdgcn_wmma_f32_16x16x4_f32";
CK_TILE_DEVICE static CVecType
exec(AVecType const& aVec, BVecType const& bVec, CVecType const& cVec)
{
return {__builtin_amdgcn_wmma_f32_16x16x4_f32(0, // A_mod
aVec,
0, // B_mod
bVec,
0, // C_mod
cVec,
0, // matrix_a_reuse
0)}; // matrix_b_reuse
}
};
/**
* @struct amdgcn_mma
* @brief Specialization of amdgcn_mma for bf16_t, bf16_t, fp32_t MMA operation on GFX1250
* architecture.
* @tparam CompilerTarget Current compiler target
*/
// TODO: c++20 template <amdgcn_target CompilerTarget>
// TODO: c++20 requires
template <typename CompilerTarget>
// clang-format off
// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar |
struct amdgcn_mma<bf16_t, bf16_t, fp32_t, 16u, 16u, 32u, CompilerTarget, MmaOpFamily::DENSE, enable_if_target_gfx1250_t<CompilerTarget>>
: amdgcn_mma_base<bf16_t, bf16_t, fp32_t, 16u, 16u, 32u, 32u, 16, 1, 1, 1, 1, 8, 1, WmmaOp, MmaOpFamily::DENSE>
// clang-format on
{
static constexpr const char* instruction_name = "__builtin_amdgcn_wmma_f32_16x16x32_bf16";
CK_TILE_DEVICE static CVecType
exec(AVecType const& aVec, BVecType const& bVec, CVecType const& cVec)
{
return {__builtin_amdgcn_wmma_f32_16x16x32_bf16(0, // A_mod
aVec,
0, // B_mod
bVec,
0, // C_mod
cVec,
0, // matrix_a_reuse
0)}; // matrix_b_reuse
}
};
/**
* @struct amdgcn_mma
* @brief Specialization of amdgcn_mma for bf16_t, bf16_t, bf16_t MMA operation on GFX1250
* architecture.
* @tparam CompilerTarget Current compiler target
*/
// TODO: c++20 template <amdgcn_target CompilerTarget>
// TODO: c++20 requires
template <typename CompilerTarget>
// clang-format off
// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar |
struct amdgcn_mma<bf16_t, bf16_t, bf16_t, 16u, 16u, 32u, CompilerTarget, MmaOpFamily::DENSE, enable_if_target_gfx1250_t<CompilerTarget>>
: amdgcn_mma_base<bf16_t, bf16_t, bf16_t, 16u, 16u, 32u, 32u, 16, 1, 1, 1, 1, 8, 1, WmmaOp, MmaOpFamily::DENSE>
// clang-format on
{
static constexpr const char* instruction_name = "__builtin_amdgcn_wmma_bf16_16x16x32_bf16";
CK_TILE_DEVICE static CVecType
exec(AVecType const& aVec, BVecType const& bVec, CVecType const& cVec)
{
return {__builtin_amdgcn_wmma_bf16_16x16x32_bf16(0, // A_mod
aVec,
0, // B_mod
bVec,
0, // C_mod
cVec,
0, // matrix_a_reuse
0)}; // matrix_b_reuse
}
};
/**
* @struct amdgcn_mma
* @brief Specialization of amdgcn_mma for fp8_t, fp8_t, fp32_t MMA operation on GFX1250
* architecture.
* @tparam CompilerTarget Current compiler target
*/
// TODO: c++20 template <amdgcn_target CompilerTarget>
// TODO: c++20 requires
template <typename CompilerTarget>
// clang-format off
// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar |
struct amdgcn_mma<fp8_t, fp8_t, fp32_t, 16u, 16u, 64u, CompilerTarget, MmaOpFamily::DENSE, enable_if_target_gfx1250_t<CompilerTarget>>
: amdgcn_mma_base<fp8_t, fp8_t, fp32_t, 16u, 16u, 64u, 32u, 32, 1, 1, 1, 1, 8, 1, WmmaOp, MmaOpFamily::DENSE>
// clang-format on
{
static constexpr const char* instruction_name = "__builtin_amdgcn_wmma_f32_16x16x64_fp8_fp8";
CK_TILE_DEVICE static CVecType
exec(AVecType const& aVec, BVecType const& bVec, CVecType const& cVec)
{
return {__builtin_amdgcn_wmma_f32_16x16x64_fp8_fp8(bit_cast<int32x8_t>(aVec),
bit_cast<int32x8_t>(bVec),
0, // C_mod
cVec,
0, // matrix_a_reuse
0)}; // matrix_b_reuse
}
};
/**
* @struct amdgcn_mma
* @brief Specialization of amdgcn_mma for fp8_t, bf8_t, fp32_t MMA operation on GFX1250
* architecture.
* @tparam CompilerTarget Current compiler target
*/
// TODO: c++20 template <amdgcn_target CompilerTarget>
// TODO: c++20 requires
template <typename CompilerTarget>
// clang-format off
// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar |
struct amdgcn_mma<fp8_t, bf8_t, fp32_t, 16u, 16u, 64u, CompilerTarget, MmaOpFamily::DENSE, enable_if_target_gfx1250_t<CompilerTarget>>
: amdgcn_mma_base<fp8_t, bf8_t, fp32_t, 16u, 16u, 64u, 32u, 32, 1, 1, 1, 1, 8, 1, WmmaOp, MmaOpFamily::DENSE>
// clang-format on
{
static constexpr const char* instruction_name = "__builtin_amdgcn_wmma_f32_16x16x64_fp8_bf8";
CK_TILE_DEVICE static CVecType
exec(AVecType const& aVec, BVecType const& bVec, CVecType const& cVec)
{
return {__builtin_amdgcn_wmma_f32_16x16x64_fp8_bf8(bit_cast<int32x8_t>(aVec),
bit_cast<int32x8_t>(bVec),
0, // C_mod
cVec,
0, // matrix_a_reuse
0)}; // matrix_b_reuse
}
};
/**
* @struct amdgcn_mma
* @brief Specialization of amdgcn_mma for bf8_t, fp8_t, fp32_t MMA operation on GFX1250
* architecture.
* @tparam CompilerTarget Current compiler target
*/
// TODO: c++20 template <amdgcn_target CompilerTarget>
// TODO: c++20 requires
template <typename CompilerTarget>
// clang-format off
// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar |
struct amdgcn_mma<bf8_t, fp8_t, fp32_t, 16u, 16u, 64u, CompilerTarget, MmaOpFamily::DENSE, enable_if_target_gfx1250_t<CompilerTarget>>
: amdgcn_mma_base<bf8_t, fp8_t, fp32_t, 16u, 16u, 64u, 32u, 32, 1, 1, 1, 1, 8, 1, WmmaOp, MmaOpFamily::DENSE>
// clang-format on
{
static constexpr const char* instruction_name = "__builtin_amdgcn_wmma_f32_16x16x64_bf8_fp8";
CK_TILE_DEVICE static CVecType
exec(AVecType const& aVec, BVecType const& bVec, CVecType const& cVec)
{
return {__builtin_amdgcn_wmma_f32_16x16x64_bf8_fp8(bit_cast<int32x8_t>(aVec),
bit_cast<int32x8_t>(bVec),
0, // C_mod
cVec,
0, // matrix_a_reuse
0)}; // matrix_b_reuse
}
};
/**
* @struct amdgcn_mma
* @brief Specialization of amdgcn_mma for bf8_t, bf8_t, fp32_t MMA operation on GFX1250
* architecture.
* @tparam CompilerTarget Current compiler target
*/
// TODO: c++20 template <amdgcn_target CompilerTarget>
// TODO: c++20 requires
template <typename CompilerTarget>
// clang-format off
// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar |
struct amdgcn_mma<bf8_t, bf8_t, fp32_t, 16u, 16u, 64u, CompilerTarget, MmaOpFamily::DENSE, enable_if_target_gfx1250_t<CompilerTarget>>
: amdgcn_mma_base<bf8_t, bf8_t, fp32_t, 16u, 16u, 64u, 32u, 32, 1, 1, 1, 1, 8, 1, WmmaOp, MmaOpFamily::DENSE>
// clang-format on
{
static constexpr const char* instruction_name = "__builtin_amdgcn_wmma_f32_16x16x64_bf8_bf8";
CK_TILE_DEVICE static CVecType
exec(AVecType const& aVec, BVecType const& bVec, CVecType const& cVec)
{
return {__builtin_amdgcn_wmma_f32_16x16x64_bf8_bf8(bit_cast<int32x8_t>(aVec),
bit_cast<int32x8_t>(bVec),
0, // C_mod
cVec,
0, // matrix_a_reuse
0)}; // matrix_b_reuse
}
};
/**
* @struct amdgcn_mma
* @brief Specialization of amdgcn_mma for fp8_t, fp8_t, fp16_t MMA operation on GFX1250
* architecture.
* @tparam CompilerTarget Current compiler target
*/
// TODO: c++20 template <amdgcn_target CompilerTarget>
// TODO: c++20 requires
template <typename CompilerTarget>
// clang-format off
// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar |
struct amdgcn_mma<fp8_t, fp8_t, fp16_t, 16u, 16u, 64u, CompilerTarget, MmaOpFamily::DENSE, enable_if_target_gfx1250_t<CompilerTarget>>
: amdgcn_mma_base<fp8_t, fp8_t, fp16_t, 16u, 16u, 64u, 32u, 32, 1, 1, 1, 1, 8, 1, WmmaOp, MmaOpFamily::DENSE>
// clang-format on
{
static constexpr const char* instruction_name = "__builtin_amdgcn_wmma_f16_16x16x64_fp8_fp8";
CK_TILE_DEVICE static CVecType
exec(AVecType const& aVec, BVecType const& bVec, CVecType const& cVec)
{
return {__builtin_amdgcn_wmma_f16_16x16x64_fp8_fp8(bit_cast<int32x8_t>(aVec),
bit_cast<int32x8_t>(bVec),
0, // C_mod
cVec,
0, // matrix_a_reuse
0)}; // matrix_b_reuse
}
};
/**
* @struct amdgcn_mma
* @brief Specialization of amdgcn_mma for fp8_t, bf8_t, fp16_t MMA operation on GFX1250
* architecture.
* @tparam CompilerTarget Current compiler target
*/
// TODO: c++20 template <amdgcn_target CompilerTarget>
// TODO: c++20 requires
template <typename CompilerTarget>
// clang-format off
// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar |
struct amdgcn_mma<fp8_t, bf8_t, fp16_t, 16u, 16u, 64u, CompilerTarget, MmaOpFamily::DENSE, enable_if_target_gfx1250_t<CompilerTarget>>
: amdgcn_mma_base<fp8_t, bf8_t, fp16_t, 16u, 16u, 64u, 32u, 32, 1, 1, 1, 1, 8, 1, WmmaOp, MmaOpFamily::DENSE>
// clang-format on
{
static constexpr const char* instruction_name = "__builtin_amdgcn_wmma_f16_16x16x64_fp8_bf8";
CK_TILE_DEVICE static CVecType
exec(AVecType const& aVec, BVecType const& bVec, CVecType const& cVec)
{
return {__builtin_amdgcn_wmma_f16_16x16x64_fp8_bf8(bit_cast<int32x8_t>(aVec),
bit_cast<int32x8_t>(bVec),
0, // C_mod
cVec,
0, // matrix_a_reuse
0)}; // matrix_b_reuse
}
};
/**
* @struct amdgcn_mma
* @brief Specialization of amdgcn_mma for bf8_t, fp8_t, fp16_t MMA operation on GFX1250
* architecture.
* @tparam CompilerTarget Current compiler target
*/
// TODO: c++20 template <amdgcn_target CompilerTarget>
// TODO: c++20 requires
template <typename CompilerTarget>
// clang-format off
// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar |
struct amdgcn_mma<bf8_t, fp8_t, fp16_t, 16u, 16u, 64u, CompilerTarget, MmaOpFamily::DENSE, enable_if_target_gfx1250_t<CompilerTarget>>
: amdgcn_mma_base<bf8_t, fp8_t, fp16_t, 16u, 16u, 64u, 32u, 32, 1, 1, 1, 1, 8, 1, WmmaOp, MmaOpFamily::DENSE>
// clang-format on
{
static constexpr const char* instruction_name = "__builtin_amdgcn_wmma_f16_16x16x64_bf8_fp8";
CK_TILE_DEVICE static CVecType
exec(AVecType const& aVec, BVecType const& bVec, CVecType const& cVec)
{
return {__builtin_amdgcn_wmma_f16_16x16x64_bf8_fp8(bit_cast<int32x8_t>(aVec),
bit_cast<int32x8_t>(bVec),
0, // C_mod
cVec,
0, // matrix_a_reuse
0)}; // matrix_b_reuse
}
};
/**
* @struct amdgcn_mma
* @brief Specialization of amdgcn_mma for bf8_t, bf8_t, fp16_t MMA operation on GFX1250
* architecture.
* @tparam CompilerTarget Current compiler target
*/
// TODO: c++20 template <amdgcn_target CompilerTarget>
// TODO: c++20 requires
template <typename CompilerTarget>
// clang-format off
// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar |
struct amdgcn_mma<bf8_t, bf8_t, fp16_t, 16u, 16u, 64u, CompilerTarget, MmaOpFamily::DENSE, enable_if_target_gfx1250_t<CompilerTarget>>
: amdgcn_mma_base<bf8_t, bf8_t, fp16_t, 16u, 16u, 64u, 32u, 32, 1, 1, 1, 1, 8, 1, WmmaOp, MmaOpFamily::DENSE>
// clang-format on
{
static constexpr const char* instruction_name = "__builtin_amdgcn_wmma_f16_16x16x64_bf8_bf8";
CK_TILE_DEVICE static CVecType
exec(AVecType const& aVec, BVecType const& bVec, CVecType const& cVec)
{
return {__builtin_amdgcn_wmma_f16_16x16x64_bf8_bf8(bit_cast<int32x8_t>(aVec),
bit_cast<int32x8_t>(bVec),
0, // C_mod
cVec,
0, // matrix_a_reuse
0)}; // matrix_b_reuse
}
};
/**
* @struct amdgcn_mma
* @brief Specialization of amdgcn_mma for int8_t, int8_t, int32_t MMA operation on GFX1250
* architecture.
* @tparam CompilerTarget Current compiler target
*/
// TODO: c++20 template <amdgcn_target CompilerTarget>
// TODO: c++20 requires
template <typename CompilerTarget>
// clang-format off
// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar |
struct amdgcn_mma<int8_t, int8_t, int32_t, 16u, 16u, 64u, CompilerTarget, MmaOpFamily::DENSE, enable_if_target_gfx1250_t<CompilerTarget>>
: amdgcn_mma_base<int8_t, int8_t, int32_t, 16u, 16u, 64u, 32u, 32, 1, 1, 1, 1, 8, 1, WmmaOp, MmaOpFamily::DENSE>
// clang-format on
{
static constexpr const char* instruction_name = "__builtin_amdgcn_wmma_i32_16x16x64_iu8";
CK_TILE_DEVICE static CVecType
exec(AVecType const& aVec, BVecType const& bVec, CVecType const& cVec)
{
return {__builtin_amdgcn_wmma_i32_16x16x64_iu8(true, // A signedness
bit_cast<int32x8_t>(aVec),
true, // B signedness
bit_cast<int32x8_t>(bVec),
cVec,
false,
0)};
}
};
/**
* @struct amdgcn_mma
* @brief Specialization of amdgcn_mma for fp8_t, fp8_t, fp16_t MMA operation on GFX1250
* architecture.
* @tparam CompilerTarget Current compiler target
*/
// TODO: c++20 template <amdgcn_target CompilerTarget>
// TODO: c++20 requires
template <typename CompilerTarget>
// clang-format off
// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar |
struct amdgcn_mma<fp8_t, fp8_t, fp16_t, 16u, 16u, 128u, CompilerTarget, MmaOpFamily::DENSE, enable_if_target_gfx1250_t<CompilerTarget>>
: amdgcn_mma_base<fp8_t, fp8_t, fp16_t, 16u, 16u, 128u, 32u, 64, 1, 1, 1, 1, 8, 1, WmmaOp, MmaOpFamily::DENSE>
// clang-format on
{
static constexpr const char* instruction_name = "__builtin_amdgcn_wmma_f16_16x16x128_fp8_fp8";
CK_TILE_DEVICE static CVecType
exec(AVecType const& aVec, BVecType const& bVec, CVecType const& cVec)
{
return {__builtin_amdgcn_wmma_f16_16x16x128_fp8_fp8(bit_cast<int32x16_t>(aVec),
bit_cast<int32x16_t>(bVec),
0, // C_mod
cVec,
0, // matrix_a_reuse
0)}; // matrix_b_reuse
}
};
/**
* @struct amdgcn_mma
* @brief Specialization of amdgcn_mma for fp8_t, bf8_t, fp16_t MMA operation on GFX1250
* architecture.
* @tparam CompilerTarget Current compiler target
*/
// TODO: c++20 template <amdgcn_target CompilerTarget>
// TODO: c++20 requires
template <typename CompilerTarget>
// clang-format off
// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar |
struct amdgcn_mma<fp8_t, bf8_t, fp16_t, 16u, 16u, 128u, CompilerTarget, MmaOpFamily::DENSE, enable_if_target_gfx1250_t<CompilerTarget>>
: amdgcn_mma_base<fp8_t, bf8_t, fp16_t, 16u, 16u, 128u, 32u, 64, 1, 1, 1, 1, 8, 1, WmmaOp, MmaOpFamily::DENSE>
// clang-format on
{
static constexpr const char* instruction_name = "__builtin_amdgcn_wmma_f16_16x16x128_fp8_bf8";
CK_TILE_DEVICE static CVecType
exec(AVecType const& aVec, BVecType const& bVec, CVecType const& cVec)
{
return {__builtin_amdgcn_wmma_f16_16x16x128_fp8_bf8(bit_cast<int32x16_t>(aVec),
bit_cast<int32x16_t>(bVec),
0, // C_mod
cVec,
0, // matrix_a_reuse
0)}; // matrix_b_reuse
}
};
/**
* @struct amdgcn_mma
* @brief Specialization of amdgcn_mma for bf8_t, fp8_t, fp16_t MMA operation on GFX1250
* architecture.
* @tparam CompilerTarget Current compiler target
*/
// TODO: c++20 template <amdgcn_target CompilerTarget>
// TODO: c++20 requires
template <typename CompilerTarget>
// clang-format off
// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar |
struct amdgcn_mma<bf8_t, fp8_t, fp16_t, 16u, 16u, 128u, CompilerTarget, MmaOpFamily::DENSE, enable_if_target_gfx1250_t<CompilerTarget>>
: amdgcn_mma_base<bf8_t, fp8_t, fp16_t, 16u, 16u, 128u, 32u, 64, 1, 1, 1, 1, 8, 1, WmmaOp, MmaOpFamily::DENSE>
// clang-format on
{
static constexpr const char* instruction_name = "__builtin_amdgcn_wmma_f16_16x16x128_bf8_fp8";
CK_TILE_DEVICE static CVecType
exec(AVecType const& aVec, BVecType const& bVec, CVecType const& cVec)
{
return {__builtin_amdgcn_wmma_f16_16x16x128_bf8_fp8(bit_cast<int32x16_t>(aVec),
bit_cast<int32x16_t>(bVec),
0, // C_mod
cVec,
0, // matrix_a_reuse
0)}; // matrix_b_reuse
}
};
/**
* @struct amdgcn_mma
* @brief Specialization of amdgcn_mma for bf8_t, bf8_t, fp16_t MMA operation on GFX1250
* architecture.
* @tparam CompilerTarget Current compiler target
*/
// TODO: c++20 template <amdgcn_target CompilerTarget>
// TODO: c++20 requires
template <typename CompilerTarget>
// clang-format off
// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar |
struct amdgcn_mma<bf8_t, bf8_t, fp16_t, 16u, 16u, 128u, CompilerTarget, MmaOpFamily::DENSE, enable_if_target_gfx1250_t<CompilerTarget>>
: amdgcn_mma_base<bf8_t, bf8_t, fp16_t, 16u, 16u, 128u, 32u, 64, 1, 1, 1, 1, 8, 1, WmmaOp, MmaOpFamily::DENSE>
// clang-format on
{
static constexpr const char* instruction_name = "__builtin_amdgcn_wmma_f16_16x16x128_bf8_bf8";
CK_TILE_DEVICE static CVecType
exec(AVecType const& aVec, BVecType const& bVec, CVecType const& cVec)
{
return {__builtin_amdgcn_wmma_f16_16x16x128_bf8_bf8(bit_cast<int32x16_t>(aVec),
bit_cast<int32x16_t>(bVec),
0, // C_mod
cVec,
0, // matrix_a_reuse
0)}; // matrix_b_reuse
}
};
/**
* @struct amdgcn_mma
* @brief Specialization of amdgcn_mma for fp6_t, fp6_t, fp32_t MMA operation on GFX1250
* architecture.
* @tparam CompilerTarget Current compiler target
*/
// TODO: c++20 template <amdgcn_target CompilerTarget>
// TODO: c++20 requires
template <typename CompilerTarget>
// clang-format off
// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar |
struct amdgcn_mma<pk_fp6x16_t, pk_fp6x16_t, fp32_t, 16u, 16u, 128u, CompilerTarget, MmaOpFamily::DENSE, enable_if_target_gfx1250_t<CompilerTarget>>
: amdgcn_mma_base<pk_fp6x16_t, pk_fp6x16_t, fp32_t, 16u, 16u, 128u, 32u, 64, 1, 1, 1, 1, 8, 1, WmmaOp, MmaOpFamily::DENSE>
// clang-format on
{
CK_TILE_DEVICE static CVecType
exec(AVecType const& aVec, BVecType const& bVec, CVecType const& cVec)
{
// fp6 format = 2, data is 12 dwords per operand, pad to 16 dwords for the builtin
int32x16_t a_padded = {aVec.data[0],
aVec.data[1],
aVec.data[2],
aVec.data[3],
aVec.data[4],
aVec.data[5],
aVec.data[6],
aVec.data[7],
aVec.data[8],
aVec.data[9],
aVec.data[10],
aVec.data[11],
0,
0,
0,
0};
int32x16_t b_padded = {bVec.data[0],
bVec.data[1],
bVec.data[2],
bVec.data[3],
bVec.data[4],
bVec.data[5],
bVec.data[6],
bVec.data[7],
bVec.data[8],
bVec.data[9],
bVec.data[10],
bVec.data[11],
0,
0,
0,
0};
return {__builtin_amdgcn_wmma_f32_16x16x128_f8f6f4(PackedDataTypeToFlag_v<pk_fp6x16_t>,
a_padded,
PackedDataTypeToFlag_v<pk_fp6x16_t>,
b_padded,
0,
cVec)};
}
};
/**
* @struct amdgcn_mma
* @brief Specialization of amdgcn_mma for fp4_t, fp4_t, fp32_t MMA operation on GFX1250
* architecture.
* @tparam CompilerTarget Current compiler target
*/
// TODO: c++20 template <amdgcn_target CompilerTarget>
// TODO: c++20 requires
template <typename CompilerTarget>
// clang-format off
// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar |
struct amdgcn_mma<pk_fp4_t, pk_fp4_t, fp32_t, 16u, 16u, 128u, CompilerTarget, MmaOpFamily::DENSE, enable_if_target_gfx1250_t<CompilerTarget>>
: amdgcn_mma_base<pk_fp4_t, pk_fp4_t, fp32_t, 16u, 16u, 128u, 32u, 64, 1, 1, 1, 1, 8, 1, WmmaOp, MmaOpFamily::DENSE>
// clang-format on
{
CK_TILE_DEVICE static CVecType
exec(AVecType const& aVec, BVecType const& bVec, CVecType const& cVec)
{
// fp4 format = 4, data is 8 dwords per operand, pad to 16 dwords for the builtin
int32x8_t a8 = bit_cast<int32x8_t>(aVec);
int32x8_t b8 = bit_cast<int32x8_t>(bVec);
int32x16_t a_padded = {
a8[0], a8[1], a8[2], a8[3], a8[4], a8[5], a8[6], a8[7], 0, 0, 0, 0, 0, 0, 0, 0};
int32x16_t b_padded = {
b8[0], b8[1], b8[2], b8[3], b8[4], b8[5], b8[6], b8[7], 0, 0, 0, 0, 0, 0, 0, 0};
return {__builtin_amdgcn_wmma_f32_16x16x128_f8f6f4(PackedDataTypeToFlag_v<pk_fp4_t>,
a_padded,
PackedDataTypeToFlag_v<pk_fp4_t>,
b_padded,
0,
cVec)};
}
};
/**
* @struct amdgcn_mma
* @brief Specialization of amdgcn_mma for fp8_t, fp8_t, fp32_t MMA operation on GFX1250
* architecture.
* @tparam CompilerTarget Current compiler target
*/
// TODO: c++20 template <amdgcn_target CompilerTarget>
// TODO: c++20 requires
template <typename CompilerTarget>
// clang-format off
// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar |
struct amdgcn_mma<fp8_t, fp8_t, fp32_t, 16u, 16u, 128u, CompilerTarget, MmaOpFamily::DENSE, enable_if_target_gfx1250_t<CompilerTarget>>
: amdgcn_mma_base<fp8_t, fp8_t, fp32_t, 16u, 16u, 128u, 32u, 64, 1, 1, 1, 1, 8, 1, WmmaOp, MmaOpFamily::DENSE>
// clang-format on
{
static constexpr const char* instruction_name = "__builtin_amdgcn_wmma_f32_16x16x128_fp8_fp8";
CK_TILE_DEVICE static CVecType
exec(AVecType const& aVec, BVecType const& bVec, CVecType const& cVec)
{
return {__builtin_amdgcn_wmma_f32_16x16x128_fp8_fp8(bit_cast<int32x16_t>(aVec),
bit_cast<int32x16_t>(bVec),
0, // C_mod
cVec,
0, // matrix_a_reuse
0)}; // matrix_b_reuse
}
};
/**
* @struct amdgcn_mma
* @brief Specialization of amdgcn_mma for fp8_t, bf8_t, fp32_t MMA operation on GFX1250
* architecture.
* @tparam CompilerTarget Current compiler target
*/
// TODO: c++20 template <amdgcn_target CompilerTarget>
// TODO: c++20 requires
template <typename CompilerTarget>
// clang-format off
// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar |
struct amdgcn_mma<fp8_t, bf8_t, fp32_t, 16u, 16u, 128u, CompilerTarget, MmaOpFamily::DENSE, enable_if_target_gfx1250_t<CompilerTarget>>
: amdgcn_mma_base<fp8_t, bf8_t, fp32_t, 16u, 16u, 128u, 32u, 64, 1, 1, 1, 1, 8, 1, WmmaOp, MmaOpFamily::DENSE>
// clang-format on
{
static constexpr const char* instruction_name = "__builtin_amdgcn_wmma_f32_16x16x128_fp8_bf8";
CK_TILE_DEVICE static CVecType
exec(AVecType const& aVec, BVecType const& bVec, CVecType const& cVec)
{
return {__builtin_amdgcn_wmma_f32_16x16x128_fp8_bf8(bit_cast<int32x16_t>(aVec),
bit_cast<int32x16_t>(bVec),
0, // C_mod
cVec,
0, // matrix_a_reuse
0)}; // matrix_b_reuse
}
};
/**
* @struct amdgcn_mma
* @brief Specialization of amdgcn_mma for bf8_t, fp8_t, fp32_t MMA operation on GFX1250
* architecture.
* @tparam CompilerTarget Current compiler target
*/
// TODO: c++20 template <amdgcn_target CompilerTarget>
// TODO: c++20 requires
template <typename CompilerTarget>
// clang-format off
// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar |
struct amdgcn_mma<bf8_t, fp8_t, fp32_t, 16u, 16u, 128u, CompilerTarget, MmaOpFamily::DENSE, enable_if_target_gfx1250_t<CompilerTarget>>
: amdgcn_mma_base<bf8_t, fp8_t, fp32_t, 16u, 16u, 128u, 32u, 64, 1, 1, 1, 1, 8, 1, WmmaOp, MmaOpFamily::DENSE>
// clang-format on
{
static constexpr const char* instruction_name = "__builtin_amdgcn_wmma_f32_16x16x128_bf8_fp8";
CK_TILE_DEVICE static CVecType
exec(AVecType const& aVec, BVecType const& bVec, CVecType const& cVec)
{
return {__builtin_amdgcn_wmma_f32_16x16x128_bf8_fp8(bit_cast<int32x16_t>(aVec),
bit_cast<int32x16_t>(bVec),
0, // C_mod
cVec,
0, // matrix_a_reuse
0)}; // matrix_b_reuse
}
};
/**
* @struct amdgcn_mma
* @brief Specialization of amdgcn_mma for bf8_t, bf8_t, fp32_t MMA operation on GFX1250
* architecture.
* @tparam CompilerTarget Current compiler target
*/
// TODO: c++20 template <amdgcn_target CompilerTarget>
// TODO: c++20 requires
template <typename CompilerTarget>
// clang-format off
// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar |
struct amdgcn_mma<bf8_t, bf8_t, fp32_t, 16u, 16u, 128u, CompilerTarget, MmaOpFamily::DENSE, enable_if_target_gfx1250_t<CompilerTarget>>
: amdgcn_mma_base<bf8_t, bf8_t, fp32_t, 16u, 16u, 128u, 32u, 64, 1, 1, 1, 1, 8, 1, WmmaOp, MmaOpFamily::DENSE>
// clang-format on
{
static constexpr const char* instruction_name = "__builtin_amdgcn_wmma_f32_16x16x128_bf8_bf8";
CK_TILE_DEVICE static CVecType
exec(AVecType const& aVec, BVecType const& bVec, CVecType const& cVec)
{
return {__builtin_amdgcn_wmma_f32_16x16x128_bf8_bf8(bit_cast<int32x16_t>(aVec),
bit_cast<int32x16_t>(bVec),
0, // C_mod
cVec,
0, // matrix_a_reuse
0)}; // matrix_b_reuse
}
};
/**
* @struct amdgcn_mma
* @brief Specialization of amdgcn_mma for fp16_t, fp16_t, fp32_t MMA operation on GFX1250
* architecture.
* @tparam CompilerTarget Current compiler target
*/
// TODO: c++20 template <amdgcn_target CompilerTarget>
// TODO: c++20 requires
template <typename CompilerTarget>
// clang-format off
// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar |
struct amdgcn_mma<fp16_t, fp16_t, fp32_t, 16u, 16u, 32u, CompilerTarget, MmaOpFamily::DENSE, enable_if_target_gfx1250_t<CompilerTarget>>
: amdgcn_mma_base<fp16_t, fp16_t, fp32_t, 16u, 16u, 32u, 32u, 16, 1, 1, 1, 1, 8, 1, WmmaOp, MmaOpFamily::DENSE>
// clang-format on
{
static constexpr const char* instruction_name = "__builtin_amdgcn_wmma_f32_16x16x32_f16";
CK_TILE_DEVICE static CVecType
exec(AVecType const& aVec, BVecType const& bVec, CVecType const& cVec)
{
return {__builtin_amdgcn_wmma_f32_16x16x32_f16(0, // A_mod
aVec,
0, // B_mod
bVec,
0, // C_mod
cVec,
0, // matrix_a_reuse
0)}; // matrix_b_reuse
}
};
/**
* @struct amdgcn_mma
* @brief Specialization of amdgcn_mma for fp16_t, fp16_t, fp16_t MMA operation on GFX1250
* architecture.
* @tparam CompilerTarget Current compiler target
*/
// TODO: c++20 template <amdgcn_target CompilerTarget>
// TODO: c++20 requires
template <typename CompilerTarget>
// clang-format off
// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar |
struct amdgcn_mma<fp16_t, fp16_t, fp16_t, 16u, 16u, 32u, CompilerTarget, MmaOpFamily::DENSE, enable_if_target_gfx1250_t<CompilerTarget>>
: amdgcn_mma_base<fp16_t, fp16_t, fp16_t, 16u, 16u, 32u, 32u, 16, 1, 1, 1, 1, 8, 1, WmmaOp, MmaOpFamily::DENSE>
// clang-format on
{
static constexpr const char* instruction_name = "__builtin_amdgcn_wmma_f16_16x16x32_f16";
CK_TILE_DEVICE static CVecType
exec(AVecType const& aVec, BVecType const& bVec, CVecType const& cVec)
{
return {__builtin_amdgcn_wmma_f16_16x16x32_f16(0, // A_mod
aVec,
0, // B_mod
bVec,
0, // C_mod
cVec,
0, // matrix_a_reuse
0)}; // matrix_b_reuse
}
};
} // namespace ck_tile::core::arch::mma

25
include/ck_tile/core/numeric/pk_f6.hpp
View File

@@ -1350,6 +1350,14 @@ struct impl::ext_vector<pk_fp6x16_t, 2>
using type = f6x16xN_tt<2, f6_kind::fp6>;
};
template <>
struct impl::ext_vector<pk_fp6x16_t, 4>
{
static constexpr index_t N = 4;
using value_type = f6x16xN_tt<4, f6_kind::fp6>;
using type = f6x16xN_tt<4, f6_kind::fp6>;
};
template <>
struct impl::ext_vector<pk_bf6x16_t, 1>
{
@@ -1358,6 +1366,14 @@ struct impl::ext_vector<pk_bf6x16_t, 1>
using type = f6x16xN_tt<1, f6_kind::bf6>;
};
template <>
struct impl::ext_vector<pk_bf6x16_t, 4>
{
static constexpr index_t N = 4;
using value_type = f6x16xN_tt<4, f6_kind::bf6>;
using type = f6x16xN_tt<4, f6_kind::bf6>;
};
template <>
struct impl::ext_vector<pk_bf6x16_t, 2>
{
@@ -1366,15 +1382,6 @@ struct impl::ext_vector<pk_bf6x16_t, 2>
using type = f6x16xN_tt<2, f6_kind::bf6>;
};
// Used as AVecType / BVecType for the gfx1250 16x16x128 mx-scale wmma kernel
template <>
struct impl::ext_vector<pk_fp6x16_t, 4>
{
static constexpr index_t N = 4;
using value_type = f6x16xN_tt<4, f6_kind::fp6>;
using type = f6x16xN_tt<4, f6_kind::fp6>;
};
// Arithmetic operations using float conversion
// Note: Arithmetic operations on packed types containing 32 elements
// may not be semantically meaningful for element-wise operations

5
test/ck_tile/core/arch/mma/CMakeLists.txt
View File

@@ -87,3 +87,8 @@ if(GPU_TARGETS MATCHES "gfx120")
target_compile_options(test_amdgcn_mma_layout_gfx12 PRIVATE ${EXAMPLE_GEMM_COMPILE_OPTIONS})
endif()
if(GPU_TARGETS MATCHES "gfx1250")
_add_mma_gtest(test_amdgcn_mma_layout_gfx1250 test_amdgcn_mma_layout_gfx1250.cpp)
target_compile_options(test_amdgcn_mma_layout_gfx1250 PRIVATE ${EXAMPLE_GEMM_COMPILE_OPTIONS})
set_mma_test_arch_define(test_amdgcn_mma_layout_gfx1250)
endif()

104
test/ck_tile/core/arch/mma/test_amdgcn_mma_layout.inc
View File

@@ -34,26 +34,27 @@ using namespace ck_tile;
using namespace ck_tile::core::arch;
using namespace mma;
using I4 = pk_int4_t;
using F4 = pk_fp4_t;
using F6 = pk_fp6x16_t;
using BF6 = pk_bf6x16_t;
using F8 = fp8_t;
using BF8 = bf8_t;
using I32 = int32_t;
using F16 = fp16_t;
using BF16 = bf16_t;
using F32 = fp32_t;
using TF32 = tf32_t;
using F64 = fp64_t;
using I8 = int8_t;
using I32 = int32_t;
using Target908 = decltype(make_amdgcn_gfx9_target<amdgcn_target_id::GFX908>());
using Target90a = decltype(make_amdgcn_gfx9_target<amdgcn_target_id::GFX90A>());
using Target942 = decltype(make_amdgcn_gfx9_target<amdgcn_target_id::GFX942>());
using Target950 = decltype(make_amdgcn_gfx9_target<amdgcn_target_id::GFX950>());
using Target11 = decltype(make_amdgcn_gfx11_target<amdgcn_target_id::GFX1100>());
using Target12 = decltype(make_amdgcn_gfx12_target<amdgcn_target_id::GFX1201>());
using I4 = pk_int4_t;
using F4 = pk_fp4_t;
using F6 = pk_fp6x16_t;
using BF6 = pk_bf6x16_t;
using F8 = fp8_t;
using BF8 = bf8_t;
using I32 = int32_t;
using F16 = fp16_t;
using BF16 = bf16_t;
using F32 = fp32_t;
using TF32 = tf32_t;
using F64 = fp64_t;
using I8 = int8_t;
using I32 = int32_t;
using Target908 = decltype(make_amdgcn_gfx9_target<amdgcn_target_id::GFX908>());
using Target90a = decltype(make_amdgcn_gfx9_target<amdgcn_target_id::GFX90A>());
using Target942 = decltype(make_amdgcn_gfx9_target<amdgcn_target_id::GFX942>());
using Target950 = decltype(make_amdgcn_gfx9_target<amdgcn_target_id::GFX950>());
using Target11 = decltype(make_amdgcn_gfx11_target<amdgcn_target_id::GFX1100>());
using Target12 = decltype(make_amdgcn_gfx12_target<amdgcn_target_id::GFX1201>());
using Target1250 = decltype(make_amdgcn_gfx1250_target<amdgcn_target_id::GFX1250>());
#if defined(CK_MMA_TEST_ARCH_GFX950)
using TestTarget = Target950;
#elif defined(CK_MMA_TEST_ARCH_GFX942)
@@ -83,6 +84,10 @@ using TestTarget = Target908;
// gather back into C matrix. The position of "1" in C is checked against the expected (m, n)
// location.
template <typename MmaOp>
using scale_value_type_t =
std::conditional_t<MmaOp::OpFamily == MmaOpFamily::SCALE16, int64_t, int32_t>;
// Helper function to place a single "1" value in any datatype. For packed types, "off"
// indicates the offset at which the "1" must be placed within the pack. The packed datatypes
// don't always have utilities to convert "1" to their type so we use raw bits instead.
@@ -252,12 +257,20 @@ struct MmaLayoutTestKernel
}
else if constexpr(MmaOpTraits<MmaOp>::IsScale)
{
// The actual scale is computed as pow(2, scale - 127), so:
// 125 -> 2^-2 and 129 -> 2^2.
int scale_A = 125;
int scale_B = 129;
c_frag = MmaOp::template exec<OpSelA<0>, OpSelB<0>>(
a_frag, b_frag, c_frag, scale_A, scale_B);
using ScaleT = scale_value_type_t<MmaOp>;
if constexpr(sizeof(ScaleT) == sizeof(int64_t))
{
// E8M0 = 125 / 129 packed into every byte of a 64-bit value
ScaleT scale_A = 0x7D7D7D7D7D7D7D7DLL;
ScaleT scale_B = 0x8181818181818181LL;
c_frag = MmaOp::exec(a_frag, b_frag, c_frag, scale_A, scale_B);
}
else
{
ScaleT scale_A = 0x7D7D7D7D; // E8M0 = 125 in every byte
ScaleT scale_B = 0x81818181; // E8M0 = 129 in every byte
c_frag = MmaOp::exec(a_frag, b_frag, c_frag, scale_A, scale_B);
}
}
else
{
@@ -554,6 +567,45 @@ using Gfx12Intrinsics = ::testing::Types<
amdgcn_mma<I4, I4, I32, 16u, 16u, 32u, Target12, MmaOpFamily::SPARSE>, // swmmac_i32_16x16x32_iu4_w32
amdgcn_mma<I4, I4, I32, 16u, 16u, 64u, Target12, MmaOpFamily::SPARSE> // swmmac_i32_16x16x64_iu4_w32
>;
using Gfx1250Intrinsics = ::testing::Types<
amdgcn_mma<F32, F32, F32, 16u, 16u, 4u, Target1250, MmaOpFamily::DENSE>, // wmma_f32_16x16x4_f32
amdgcn_mma<BF16, BF16, F32, 16u, 16u, 32u, Target1250, MmaOpFamily::DENSE>, // wmma_f32_16x16x32_bf16
amdgcn_mma<BF16, BF16, BF16, 16u, 16u, 32u, Target1250, MmaOpFamily::DENSE>, // wmma_bf16_16x16x32_bf16
amdgcn_mma<F8, F8, F32, 16u, 16u, 64u, Target1250, MmaOpFamily::DENSE>, // wmma_f32_16x16x64_fp8_fp8
amdgcn_mma<F8, BF8, F32, 16u, 16u, 64u, Target1250, MmaOpFamily::DENSE>, // wmma_f32_16x16x64_fp8_bf8
amdgcn_mma<BF8, F8, F32, 16u, 16u, 64u, Target1250, MmaOpFamily::DENSE>, // wmma_f32_16x16x64_bf8_fp8
amdgcn_mma<BF8, BF8, F32, 16u, 16u, 64u, Target1250, MmaOpFamily::DENSE>, // wmma_f32_16x16x64_bf8_bf8
amdgcn_mma<F8, F8, F16, 16u, 16u, 64u, Target1250, MmaOpFamily::DENSE>, // wmma_f16_16x16x64_fp8_fp8
amdgcn_mma<F8, BF8, F16, 16u, 16u, 64u, Target1250, MmaOpFamily::DENSE>, // wmma_f16_16x16x64_fp8_bf8
amdgcn_mma<BF8, F8, F16, 16u, 16u, 64u, Target1250, MmaOpFamily::DENSE>, // wmma_f16_16x16x64_bf8_fp8
amdgcn_mma<BF8, BF8, F16, 16u, 16u, 64u, Target1250, MmaOpFamily::DENSE>, // wmma_f16_16x16x64_bf8_bf8
amdgcn_mma<I8, I8, I32, 16u, 16u, 64u, Target1250, MmaOpFamily::DENSE>, // wmma_i32_16x16x64_iu8
amdgcn_mma<F8, F8, F16, 16u, 16u, 128u, Target1250, MmaOpFamily::DENSE>, // wmma_f16_16x16x128_fp8_fp8
amdgcn_mma<F8, BF8, F16, 16u, 16u, 128u, Target1250, MmaOpFamily::DENSE>, // wmma_f16_16x16x128_fp8_bf8
amdgcn_mma<BF8, F8, F16, 16u, 16u, 128u, Target1250, MmaOpFamily::DENSE>, // wmma_f16_16x16x128_bf8_fp8
amdgcn_mma<BF8, BF8, F16, 16u, 16u, 128u, Target1250, MmaOpFamily::DENSE>, // wmma_f16_16x16x128_bf8_bf8
amdgcn_mma<F8, F8, F32, 16u, 16u, 128u, Target1250, MmaOpFamily::DENSE>, // wmma_f32_16x16x128_fp8_fp8
amdgcn_mma<F8, BF8, F32, 16u, 16u, 128u, Target1250, MmaOpFamily::DENSE>, // wmma_f32_16x16x128_fp8_bf8
amdgcn_mma<BF8, F8, F32, 16u, 16u, 128u, Target1250, MmaOpFamily::DENSE>, // wmma_f32_16x16x128_bf8_fp8
amdgcn_mma<BF8, BF8, F32, 16u, 16u, 128u, Target1250, MmaOpFamily::DENSE>, // wmma_f32_16x16x128_bf8_bf8
amdgcn_mma<F16, F16, F32, 16u, 16u, 32u, Target1250, MmaOpFamily::DENSE>, // wmma_f32_16x16x32_f16
amdgcn_mma<F16, F16, F16, 16u, 16u, 32u, Target1250, MmaOpFamily::DENSE>, // wmma_f16_16x16x32_f16
amdgcn_mma<F6, F6, F32, 16u, 16u, 128u, Target1250, MmaOpFamily::DENSE>, // wmma_f32_16x16x128_f8f6f4
amdgcn_mma<F4, F4, F32, 16u, 16u, 128u, Target1250, MmaOpFamily::DENSE>, // wmma_f32_16x16x128_f8f6f4
// Scale WMMA (32-bit scale)
amdgcn_mma<F8, F8, F32, 16u, 16u, 128u, Target1250, MmaOpFamily::SCALE>, // wmma_scale_f32_16x16x128_fp8
amdgcn_mma<BF8, BF8, F32, 16u, 16u, 128u, Target1250, MmaOpFamily::SCALE>, // wmma_scale_f32_16x16x128_bf8
amdgcn_mma<F6, F6, F32, 16u, 16u, 128u, Target1250, MmaOpFamily::SCALE>, // wmma_scale_f32_16x16x128_fp6
amdgcn_mma<BF6, BF6, F32, 16u, 16u, 128u, Target1250, MmaOpFamily::SCALE>, // wmma_scale_f32_16x16x128_bf6
amdgcn_mma<F4, F4, F32, 16u, 16u, 128u, Target1250, MmaOpFamily::SCALE>, // wmma_scale_f32_16x16x128_fp4
// Scale16 WMMA (64-bit scale)
amdgcn_mma<F8, F8, F32, 16u, 16u, 128u, Target1250, MmaOpFamily::SCALE16>, // wmma_scale16_f32_16x16x128_fp8
amdgcn_mma<BF8, BF8, F32, 16u, 16u, 128u, Target1250, MmaOpFamily::SCALE16>, // wmma_scale16_f32_16x16x128_bf8
amdgcn_mma<F6, F6, F32, 16u, 16u, 128u, Target1250, MmaOpFamily::SCALE16>, // wmma_scale16_f32_16x16x128_fp6
amdgcn_mma<BF6, BF6, F32, 16u, 16u, 128u, Target1250, MmaOpFamily::SCALE16>, // wmma_scale16_f32_16x16x128_bf6
amdgcn_mma<F4, F4, F32, 16u, 16u, 128u, Target1250, MmaOpFamily::SCALE16> // wmma_scale16_f32_16x16x128_fp4
>;
// clang-format on
template <typename MmaOp>

6
test/ck_tile/core/arch/mma/test_amdgcn_mma_layout_gfx1250.cpp Normal file
View File

@@ -0,0 +1,6 @@
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
#include "test_amdgcn_mma_layout.inc"
TYPED_TEST_SUITE(TestMmaLayout, Gfx1250Intrinsics);
TYPED_TEST(TestMmaLayout, Gfx1250Intrinsics) { run_mma_layout_test<TypeParam>(); }