Skip to content

Commit c733bf7

Browse files
borontionborontion
authored
[AMD] Add more gfx1250 wmma data types (#8312)
Previously we only have 2 wmma v3 instructions: - bf16 * bf16 -> fp32 - bf8 * bf8 -> fp32 k=64 This PR extends to cover the following wmma v3 instructions: - fp32 * fp32 -> fp32 - fp16 * fp16 -> fp32 - fp8 * fp8 -> fp32 k=64/128 - fp8 * bf8 -> fp32 k=64/128 - bf8 * fp8 -> fp32 k=64/128 - bf8 * bf8 -> fp32 k=128
1 parent 8ee5840 commit c733bf7

File tree

5 files changed

+151
-67
lines changed

5 files changed

+151
-67
lines changed

‎lib/Dialect/TritonGPU/IR/Dialect.cpp‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1378,8 +1378,8 @@ LogicalResult AMDWmmaEncodingAttr::verify(
13781378
if (version == 2 && !llvm::is_contained(validShapesV2, shape))
13791379
return emitError() << "invalid WMMA v2 instruction shape";
13801380

1381-
auto validShapesV3 =
1382-
std::vector<llvm::SmallVector<unsigned>>{{16, 16, 32}, {16, 16, 64}};
1381+
auto validShapesV3 = std::vector<llvm::SmallVector<unsigned>>{
1382+
{16, 16, 4}, {16, 16, 32}, {16, 16, 64}, {16, 16, 128}};
13831383
if (version == 3 && !llvm::is_contained(validShapesV3, shape))
13841384
return emitError() << "invalid WMMA v3 instruction shape";
13851385

@@ -2490,13 +2490,13 @@ LogicalResult DotOperandEncodingAttr::verify(
24902490
return emitError()
24912491
<< "ttg.dot_op kWidth parameter must be 8/16 for WMMA v1 "
24922492
"(including packed cases for `scaled_dot`)";
2493-
if (parentAttr.getVersion() == 2 &&
2494-
(kWidth != 4 && kWidth != 8 && kWidth != 16))
2493+
if (parentAttr.getVersion() == 2 && !llvm::is_contained({4, 8, 16}, kWidth))
24952494
return emitError()
24962495
<< "ttg.dot_op kWidth parameter must be 4/8/16 for WMMA v2 "
24972496
"(including packed cases for `scaled_dot`)";
2498-
if (parentAttr.getVersion() == 3 && (kWidth != 8))
2499-
return emitError() << "ttg.dot_op kWidth parameter must be 8 for WMMA v3";
2497+
if (parentAttr.getVersion() == 3 && !llvm::is_contained({2, 8, 16}, kWidth))
2498+
return emitError()
2499+
<< "ttg.dot_op kWidth parameter must be 2/8/16 for WMMA v3";
25002500
return success();
25012501
}
25022502

‎third_party/amd/include/TritonAMDGPUTransforms/WmmaGroup.h‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,10 @@ struct WmmaIntrinsic {
1313
unsigned nDim, unsigned inputKDim,
1414
Type aElemType, Type bElemType,
1515
Type dElemType);
16+
// Gets the wmma intrinsic based on exact match of all parameters.
17+
static FailureOr<WmmaIntrinsic> get(int version, unsigned mDim, unsigned nDim,
18+
unsigned kDim, Type aElemType,
19+
Type bElemType, Type dElemType);
1620

1721
WmmaIntrinsic(StringRef symbol, unsigned m, unsigned n, unsigned k,
1822
unsigned kB, Type aET, Type bET, Type dET)

‎third_party/amd/lib/TritonAMDGPUToLLVM/DotOpToLLVM/WMMA.cpp‎

Lines changed: 13 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -56,10 +56,13 @@ ValueTable getValuesFromDotOperandLayoutStruct(
5656
}
5757

5858
Value convertedElems;
59-
if (type.isF16() || (wmmaVer == 3 && type.isBF16())) {
59+
if (type.isF32() || type.isF16()) {
6060
convertedElems = rawElems;
6161
} else if (type.isBF16()) {
62-
convertedElems = tb.bitcast(rawElems, vec_ty(i16_ty, kBase));
62+
convertedElems = rawElems;
63+
// Before wmma v3, bf16 is converted to i16
64+
if (wmmaVer < 3)
65+
convertedElems = tb.bitcast(rawElems, vec_ty(i16_ty, kBase));
6366
} else {
6467
convertedElems = tb.bitcast(
6568
rawElems, vec_ty(i32_ty, kBase * type.getIntOrFloatBitWidth() /
@@ -101,22 +104,22 @@ Value generateWMMAIntrinsic(ConversionPatternRewriter &rewriter, Location loc,
101104
} else {
102105
assert(wmmaVer == 3 && "unexpected wmma version");
103106
// arguments for v3:
104-
// int: %A_mod, %A, %B_mod, %B, %C, %A_reuse, %B_reuse
105-
// fp16/bf16: %A_mod, %A, %B_mod, %B, %C_mod, %C, %A_reuse, %B_reuse
106-
// fp8/bf8: %A, %B, %C_mod, %C, %A_reuse, %B_reuse
107+
// int: %A_mod, %A, %B_mod, %B, %C, %A_reuse, %B_reuse
108+
// f32/f16/bf16: %A_mod, %A, %B_mod, %B, %C_mod, %C, %A_reuse, %B_reuse
109+
// f8/bf8: %A, %B, %C_mod, %C, %A_reuse, %B_reuse
107110
if (aElType.isInteger())
108111
operands.push_back(b.int_val(1, !aElType.isUnsignedInteger()));
109-
else if (aElType.isBF16() || aElType.isF16())
112+
else if (aElType.isFloat(16) || aElType.isF32())
110113
operands.push_back(b.int_val(1, 0));
111114
operands.push_back(valA);
112115

113116
if (bElType.isInteger())
114117
operands.push_back(b.int_val(1, !bElType.isUnsignedInteger()));
115-
else if (bElType.isBF16() || bElType.isF16())
118+
else if (bElType.isFloat(16) || bElType.isF32())
116119
operands.push_back(b.int_val(1, 0));
117120
operands.push_back(valB);
118121

119-
if ((bElType.isBF16() || bElType.isF16()) || aElType.isFloat(8))
122+
if (bElType.isFloat(16) || bElType.isF32() || aElType.isFloat(8))
120123
operands.push_back(b.int_val(16, 0));
121124
operands.push_back(valC);
122125

@@ -165,11 +168,9 @@ LogicalResult convertDot(DotOp op, DotOpAdaptor adaptor,
165168
const auto kDimOperandSize = aTensorTy.getShape().back();
166169

167170
std::string intrinsicName;
168-
FailureOr<WmmaIntrinsic> maybeWmmaIntrinsic =
169-
WmmaIntrinsic::selectFor(wmmaVer, mnkDim[0], mnkDim[1], kDimOperandSize,
170-
aElemTy, bElemTy, dElemTy);
171+
FailureOr<WmmaIntrinsic> maybeWmmaIntrinsic = WmmaIntrinsic::get(
172+
wmmaVer, mnkDim[0], mnkDim[1], mnkDim[2], aElemTy, bElemTy, dElemTy);
171173
if (failed(maybeWmmaIntrinsic)) {
172-
173174
return op.emitError(
174175
"no matching matrix core intrinsic due to unsupported element type");
175176
}

‎third_party/amd/lib/TritonAMDGPUTransforms/WmmaGroup.cpp‎

Lines changed: 76 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -83,68 +83,86 @@ WmmaDatabase::WmmaDatabase(MLIRContext *context) {
8383
auto ocpBf8T = b.getType<Float8E5M2Type>();
8484

8585
wmmaMap = {
86+
// f32 inputs
87+
// wmma_f32_16x16x4_f32
88+
TRITON_WMMA_v(3, 16, 16, f32T, f32T, 32, f32T,
89+
"llvm.amdgcn.wmma.f32.16x16x4.f32", 4, 2),
90+
91+
// f16 inputs
92+
// wmma_f32_16x16x16_f16
93+
TRITON_WMMA_v(1, 16, 16, f16T, f16T, 16, f32T,
94+
"llvm.amdgcn.wmma.f32.16x16x16.f16", 16, 16),
95+
TRITON_WMMA_v(2, 16, 16, f16T, f16T, 16, f32T,
96+
"llvm.amdgcn.wmma.f32.16x16x16.f16", 16, 8),
97+
// wmma_f32_16x16x32_f16
98+
TRITON_WMMA_v(3, 16, 16, f16T, f16T, 16, f32T,
99+
"llvm.amdgcn.wmma.f32.16x16x32.f16", 32, 16),
86100
// wmma_f16_16x16x16_f16
87101
TRITON_WMMA_v(1, 16, 16, f16T, f16T, 16, f16T,
88102
"llvm.amdgcn.wmma.f16.16x16x16.f16", 16, 16),
89103
TRITON_WMMA_v(2, 16, 16, f16T, f16T, 16, f16T,
90104
"llvm.amdgcn.wmma.f16.16x16x16.f16", 16, 8),
91105

106+
// bf16 inputs
92107
// wmma_f32_16x16x16_bf16
93108
TRITON_WMMA_v(1, 16, 16, bf16T, bf16T, 16, f32T,
94109
"llvm.amdgcn.wmma.f32.16x16x16.bf16", 16, 16),
95110
TRITON_WMMA_v(2, 16, 16, bf16T, bf16T, 16, f32T,
96111
"llvm.amdgcn.wmma.f32.16x16x16.bf16", 16, 8),
97-
98112
// wmma_f32_16x16x32_bf16
99113
TRITON_WMMA_v(3, 16, 16, bf16T, bf16T, 16, f32T,
100114
"llvm.amdgcn.wmma.f32.16x16x32.bf16", 32, 16),
101-
102-
// wmma_f32_16x16x16_f16
103-
TRITON_WMMA_v(1, 16, 16, f16T, f16T, 16, f32T,
104-
"llvm.amdgcn.wmma.f32.16x16x16.f16", 16, 16),
105-
TRITON_WMMA_v(2, 16, 16, f16T, f16T, 16, f32T,
106-
"llvm.amdgcn.wmma.f32.16x16x16.f16", 16, 8),
107-
108115
// wmma_bf16_16x16x16_bf16
109116
TRITON_WMMA_v(1, 16, 16, bf16T, bf16T, 16, bf16T,
110117
"llvm.amdgcn.wmma.bf16.16x16x16.bf16", 16, 16),
111118
TRITON_WMMA_v(2, 16, 16, bf16T, bf16T, 16, bf16T,
112119
"llvm.amdgcn.wmma.bf16.16x16x16.bf16", 16, 8),
113120

114-
// wmma_i32_16x16x16_iu4
115-
TRITON_WMMA_v(1, 16, 16, i4T, i4T, 4, i32T,
116-
"llvm.amdgcn.wmma.i32.16x16x16.iu4", 16, 16),
117-
118-
// wmma_i32_16x16x32_iu4 && wmma_i32_16x16x16_iu4
119-
TRITON_WMMA_v2_2case(16, 16, i4T, i4T, 4, i32T,
120-
"llvm.amdgcn.wmma.i32.16x16x32.iu4", 32, 16,
121-
"llvm.amdgcn.wmma.i32.16x16x16.iu4", 16, 8),
122-
123-
// wmma_i32_16x16x16_iu8
124-
TRITON_WMMA_v(1, 16, 16, i8T, i8T, 8, i32T,
125-
"llvm.amdgcn.wmma.i32.16x16x16.iu8", 16, 16),
126-
TRITON_WMMA_v(2, 16, 16, i8T, i8T, 8, i32T,
127-
"llvm.amdgcn.wmma.i32.16x16x16.iu8", 16, 8),
128-
121+
// fp8/bf8 inputs
129122
// wmma_f32_16x16x16_fp8_fp8
130123
TRITON_WMMA_v(2, 16, 16, ocpFp8T, ocpFp8T, 8, f32T,
131124
"llvm.amdgcn.wmma.f32.16x16x16.fp8.fp8", 16, 8),
132-
125+
// wmma_f32_16x16x128_fp8_fp8 & wmma_f32_16x16x64_fp8_fp8
126+
TRITON_WMMA_v_2case(3, 16, 16, ocpFp8T, ocpFp8T, 8, f32T,
127+
"llvm.amdgcn.wmma.f32.16x16x128.fp8.fp8", 128, 64,
128+
"llvm.amdgcn.wmma.f32.16x16x64.fp8.fp8", 64, 32),
133129
// wmma_f32_16x16x16_fp8_bf8
134130
TRITON_WMMA_v(2, 16, 16, ocpFp8T, ocpBf8T, 8, f32T,
135131
"llvm.amdgcn.wmma.f32.16x16x16.fp8.bf8", 16, 8),
136-
132+
// wmma_f32_16x16x128_fp8_bf8 & wmma_f32_16x16x64_fp8_bf8
133+
TRITON_WMMA_v_2case(3, 16, 16, ocpFp8T, ocpBf8T, 8, f32T,
134+
"llvm.amdgcn.wmma.f32.16x16x128.fp8.bf8", 128, 64,
135+
"llvm.amdgcn.wmma.f32.16x16x64.fp8.bf8", 64, 32),
137136
// wmma_f32_16x16x16_bf8_fp8
138137
TRITON_WMMA_v(2, 16, 16, ocpBf8T, ocpFp8T, 8, f32T,
139138
"llvm.amdgcn.wmma.f32.16x16x16.bf8.fp8", 16, 8),
140-
139+
// wmma_f32_16x16x128_bf8_fp8 & wmma_f32_16x16x64_bf8_fp8
140+
TRITON_WMMA_v_2case(3, 16, 16, ocpBf8T, ocpFp8T, 8, f32T,
141+
"llvm.amdgcn.wmma.f32.16x16x128.bf8.fp8", 128, 64,
142+
"llvm.amdgcn.wmma.f32.16x16x64.bf8.fp8", 64, 32),
141143
// wmma_f32_16x16x16_bf8_bf8
142144
TRITON_WMMA_v(2, 16, 16, ocpBf8T, ocpBf8T, 8, f32T,
143145
"llvm.amdgcn.wmma.f32.16x16x16.bf8.bf8", 16, 8),
146+
// wmma_f32_16x16x128_bf8_bf8 & wmma_f32_16x16x64_bf8_bf8
147+
TRITON_WMMA_v_2case(3, 16, 16, ocpBf8T, ocpBf8T, 8, f32T,
148+
"llvm.amdgcn.wmma.f32.16x16x128.bf8.bf8", 128, 64,
149+
"llvm.amdgcn.wmma.f32.16x16x64.bf8.bf8", 64, 32),
144150

145-
// wmma_f32_16x16x64_bf8_bf8
146-
TRITON_WMMA_v(3, 16, 16, ocpBf8T, ocpBf8T, 8, f32T,
147-
"llvm.amdgcn.wmma.f32.16x16x64.bf8.bf8", 64, 32),
151+
// iu8 inputs
152+
// wmma_i32_16x16x16_iu8
153+
TRITON_WMMA_v(1, 16, 16, i8T, i8T, 8, i32T,
154+
"llvm.amdgcn.wmma.i32.16x16x16.iu8", 16, 16),
155+
TRITON_WMMA_v(2, 16, 16, i8T, i8T, 8, i32T,
156+
"llvm.amdgcn.wmma.i32.16x16x16.iu8", 16, 8),
157+
158+
// iu4 inputs
159+
// wmma_i32_16x16x16_iu4
160+
TRITON_WMMA_v(1, 16, 16, i4T, i4T, 4, i32T,
161+
"llvm.amdgcn.wmma.i32.16x16x16.iu4", 16, 16),
162+
// wmma_i32_16x16x32_iu4 && wmma_i32_16x16x16_iu4
163+
TRITON_WMMA_v2_2case(16, 16, i4T, i4T, 4, i32T,
164+
"llvm.amdgcn.wmma.i32.16x16x32.iu4", 32, 16,
165+
"llvm.amdgcn.wmma.i32.16x16x16.iu4", 16, 8),
148166
};
149167
}
150168

@@ -187,4 +205,33 @@ WmmaIntrinsic::selectFor(int version, unsigned mDim, unsigned nDim,
187205
return WmmaIntrinsic(symbol, mDim, nDim, k, kBase, aElemType, bElemType,
188206
dElemType);
189207
}
208+
209+
FailureOr<WmmaIntrinsic> WmmaIntrinsic::get(int version, unsigned mDim,
210+
unsigned nDim, unsigned kDim,
211+
Type aElemType, Type bElemType,
212+
Type dElemType) {
213+
const WmmaMap &wmmaMap = WmmaDatabase::get(aElemType.getContext());
214+
WmmaKey key = {version,
215+
mDim,
216+
nDim,
217+
aElemType.getTypeID(),
218+
bElemType.getTypeID(),
219+
aElemType.getIntOrFloatBitWidth(),
220+
dElemType.getTypeID()};
221+
222+
auto it = wmmaMap.find(key);
223+
if (it == wmmaMap.end())
224+
return failure();
225+
226+
const SmallVector<WmmaMapValue, 2> &values = it->second;
227+
auto match = llvm::find_if(values, [&](const WmmaMapValue &val) {
228+
return std::get<1>(val) == kDim;
229+
});
230+
if (match == values.end())
231+
return failure();
232+
233+
auto [symbol, k, kBase] = *match;
234+
return WmmaIntrinsic(symbol, mDim, nDim, k, kBase, aElemType, bElemType,
235+
dElemType);
236+
}
190237
} // namespace mlir

‎third_party/amd/python/test/test_gluon_gfx1250.py‎

Lines changed: 52 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -62,11 +62,35 @@ def gemm_kernel(a_ptr, b_ptr, c_ptr, #
6262
ttgl.store(c_ptr + offs_c, accumulator, mask=mask_c)
6363

6464

65-
@pytest.mark.parametrize("BLOCK_M,BLOCK_N,BLOCK_K", [(32, 32, 32), (64, 64, 64), (128, 128, 64)])
66-
@pytest.mark.parametrize("a_dtype,b_dtype,k_dim", [
67-
("bfloat16", "bfloat16", 32),
68-
("float8_e5m2", "float8_e5m2", 64),
69-
])
65+
def get_test_gemm_block_mnk():
66+
return [
67+
(m, n, k) for (m, n) in [(32, 32), (64, 64)] \
68+
for k in [32, 64, 128, 256]
69+
]
70+
71+
72+
def get_test_gemm_variants():
73+
return [
74+
# float32 * float32 -> float32
75+
("float32", "float32", 4),
76+
# bfloat16/float16 * bfloat16/float16 -> float32
77+
*[(a, a, 32) for a in ["bfloat16", "float16"]],
78+
# float8e4m3/float8e5m2 * float8e4m3/float8e5m2 -> float32/float16
79+
*[(a, b, k) for a in ["float8_e4m3fn", "float8_e5m2"] \
80+
for b in ["float8_e4m3fn", "float8_e5m2"] \
81+
for k in [64, 128]],
82+
]
83+
84+
85+
def get_test_gemm_shapes():
86+
return [
87+
(256, 256, 256),
88+
(250, 250, 250),
89+
]
90+
91+
92+
@pytest.mark.parametrize("BLOCK_M,BLOCK_N,BLOCK_K", get_test_gemm_block_mnk())
93+
@pytest.mark.parametrize("a_dtype,b_dtype,k_dim", get_test_gemm_variants())
7094
def test_compile_gemm(BLOCK_M, BLOCK_N, BLOCK_K, a_dtype, b_dtype, k_dim):
7195
if BLOCK_K < k_dim:
7296
pytest.skip("Skip tests where BLOCK_K < k_dim")
@@ -86,39 +110,47 @@ def test_compile_gemm(BLOCK_M, BLOCK_N, BLOCK_K, a_dtype, b_dtype, k_dim):
86110
"INSTR_SHAPE_K": "constexpr", "K_WIDTH": "constexpr"
87111
}, constexprs={
88112
"BLOCK_M": BLOCK_M, "BLOCK_N": BLOCK_N, "BLOCK_K": BLOCK_K, #
89-
"INSTR_SHAPE_K": k_dim, "K_WIDTH": 8
113+
"INSTR_SHAPE_K": k_dim, "K_WIDTH": 2 if a_dtype == "fp32" else 8
90114
}), target=GPUTarget("hip", 'gfx1250', 32))
91115
amdgcn = k.asm["amdgcn"]
92116

93117
wmma_pattern = "v_wmma_"
94118
wmma_pattern += "f32_"
95119
wmma_pattern += "16x16x" + str(k_dim) + "_"
96-
if a_dtype == "bf16":
97-
wmma_pattern += "bf16"
98-
if a_dtype == "fp8e5":
99-
wmma_pattern += "bf8_bf8"
120+
if a_dtype == "fp32":
121+
wmma_pattern += "f32"
122+
if a_dtype in ("fp16", "bf16"):
123+
a_ty = "f16" if a_dtype == "fp16" else "bf16"
124+
wmma_pattern += a_ty
125+
if a_dtype in ("fp8e4nv", "fp8e5"):
126+
a_ty = "fp8" if a_dtype == "fp8e4nv" else "bf8"
127+
b_ty = "fp8" if b_dtype == "fp8e4nv" else "bf8"
128+
wmma_pattern += a_ty + "_" + b_ty
100129

101130
assert re.search(wmma_pattern, amdgcn), "The AMDGCN assembly does not contain the expected WMMA instruction."
102131

103132

104-
@pytest.mark.parametrize("M,N,K", [(256, 256, 128), (250, 250, 120)])
105-
@pytest.mark.parametrize("BLOCK_M,BLOCK_N,BLOCK_K", [(32, 32, 32), (64, 64, 64), (128, 128, 64)])
106-
@pytest.mark.parametrize("a_dtype,b_dtype,k_dim", [
107-
("bfloat16", "bfloat16", 32),
108-
("float8_e5m2", "float8_e5m2", 64),
109-
])
133+
@pytest.mark.parametrize("M,N,K", get_test_gemm_shapes())
134+
@pytest.mark.parametrize("BLOCK_M,BLOCK_N,BLOCK_K", get_test_gemm_block_mnk())
135+
@pytest.mark.parametrize("a_dtype,b_dtype,k_dim", get_test_gemm_variants())
110136
def test_runtime_gemm(M, N, K, BLOCK_M, BLOCK_N, BLOCK_K, a_dtype, b_dtype, k_dim):
111137
if BLOCK_K < k_dim:
112138
pytest.skip("Skip tests where BLOCK_K < k_dim")
139+
if a_dtype == 'float8_e4m3fn' or b_dtype == 'float8_e4m3fn':
140+
pytest.skip("Skip float8_e4m3fn tests for now due to accuracy issue")
113141

114142
torch.manual_seed(42)
115143

116144
def create_operand(shape, dtype):
117-
if dtype == torch.bfloat16:
145+
if dtype in (torch.float16, torch.bfloat16, torch.float32):
118146
return torch.randn(shape, dtype=dtype)
119-
else:
120-
assert dtype == torch.float8_e5m2
147+
elif dtype == torch.float8_e5m2:
148+
# range from min normal (0 00001 00) to max normal (0 11110 11)
121149
return torch.randint(0x04, 0x7B, shape, dtype=torch.uint8).view(dtype)
150+
else:
151+
# range from min normal (0 0001 000) to max normal (0 1110 111)
152+
assert dtype == torch.float8_e4m3fn
153+
return torch.randint(0x08, 0x77, shape, dtype=torch.uint8).view(dtype)
122154

123155
a_dtype = getattr(torch, a_dtype)
124156
b_dtype = getattr(torch, b_dtype)
@@ -141,7 +173,7 @@ def create_operand(shape, dtype):
141173
stride_bk, stride_bn, #
142174
stride_cm, stride_cn, #
143175
BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K, #
144-
INSTR_SHAPE_K=k_dim, K_WIDTH=8)
176+
INSTR_SHAPE_K=k_dim, K_WIDTH=2 if a_dtype == torch.float32 else 8)
145177

146178
c_triton = c_device.cpu()
147179
c_torch = a.to(torch.float32) @ b.to(torch.float32)

0 commit comments

Comments
 (0)