diff --git a/README.md b/README.md index 3aadfa397999..39758c3a14ef 100644 --- a/README.md +++ b/README.md @@ -67,7 +67,7 @@ build/bin/rocmlir-gen -pv -operation gemm -t f16 -out_datatype f32 \ Useful `rocmlir-gen` flags: - `--arch` -- target AMDGPU architecture (e.g. `gfx942`, `gfx950`, `gfx1100`); MFMA/WMMA support is inferred from the chosen architecture. -- `-t` / `--dtype` -- data type selector (e.g. `f16`, `f32`, `bf16`, `i8`, `fp8_fp8`). +- `-t` / `--dtype` -- data type selector (e.g. `f16`, `f32`, `bf16`, `i4`, `i8`, `fp8_fp8`). - `-out_datatype` / `--out_dtype` / `-tc` -- override the output data type independently of `-t` (e.g. f16 input with f32 output). - `--perf_config` -- supply a serialized tuning configuration. - `-ph` -- emit host code alongside the kernel. diff --git a/mlir/include/mlir/Dialect/Rock/IR/RockOps.td b/mlir/include/mlir/Dialect/Rock/IR/RockOps.td index c8fae96ff15a..7b8ab5ca2d8b 100644 --- a/mlir/include/mlir/Dialect/Rock/IR/RockOps.td +++ b/mlir/include/mlir/Dialect/Rock/IR/RockOps.td @@ -86,8 +86,8 @@ class Rock_ConvOpBase inputTypes = [F32, F16, BF16], let results = (outs Optional:$result); } -defvar GemmInputTypes = [F32, F16, BF16, I8, F8E5M2FNUZ, F8E4M3FNUZ, F8E5M2, - F8E4M3FN, F4E2M1FN, F8E8M0FNU]; +defvar GemmInputTypes = [F32, F16, BF16, I<4>, I8, F8E5M2FNUZ, F8E4M3FNUZ, + F8E5M2, F8E4M3FN, F4E2M1FN, F8E8M0FNU]; // This can be extended for quantization. defvar GemmOutputTypes = [F32, F16, BF16, I32, I8, F8E5M2FNUZ, F8E4M3FNUZ, F8E5M2, F8E4M3FN]; diff --git a/mlir/include/mlir/Dialect/Rock/Tuning/QuickTuningPerfconfigs.inc b/mlir/include/mlir/Dialect/Rock/Tuning/QuickTuningPerfconfigs.inc index 9024879a1b68..51e4719afbf7 100644 --- a/mlir/include/mlir/Dialect/Rock/Tuning/QuickTuningPerfconfigs.inc +++ b/mlir/include/mlir/Dialect/Rock/Tuning/QuickTuningPerfconfigs.inc @@ -1485,6 +1485,23 @@ const StringRef PopulateParamsWmma::initParametersI8GemmGfx1201[] = { }; // END_GEMM_Wmma_i8_gfx1201_DEFS +// BEGIN_GEMM_Wmma_i4_gfx1201_DEFS +// HAND-ADDED ENTRY — this file is generated by quickTuningGen.py from tuning-sweep +// TSVs, but the autotuner (rocmlir-tuning-driver) cannot yet tune i4 (its +// applicability pipeline rejects all i4 configs), so this entry is placed by hand +// and a regenerate/--update would DROP it. It is REQUIRED (not perf-only): the +// quick-tuning lookup fatal-errors on a miss and falls back only across +// architectures with the same op+dtype, so i4 has no fallback — without this, +// i4 GEMM does not lower. Correctness-verified (lowers to +// rocdl.wmma.i32.16x16x32.iu4, computes C==K on gfx1201). The config below was +// selected from an offline gfx1201 sweep as a robust single config; see the PR +// description for methodology and numbers. iu4 needs kpackPerBlock*kpack >= +// inputVectorLen (16); here 4*16 = 64. +const StringRef PopulateParamsWmma::initParametersI4GemmGfx1201[] = { + "v4:256,32,4,16,32,16,16,1,1,2,0,0,1,1" +}; +// END_GEMM_Wmma_i4_gfx1201_DEFS + // BEGIN_CONV_Wmma_f16_gfx1201_DEFS const StringRef PopulateParamsWmma::initParametersF16ConvGfx1201[] = { "v4:128,128,4,128,32,16,8,1,1,2,0,0,1,1", @@ -2207,6 +2224,11 @@ static constexpr size_t nInitParametersI8GemmGfx1201 = 13; static const StringRef initParametersI8GemmGfx1201[nInitParametersI8GemmGfx1201]; // END_GEMM_Wmma_i8_gfx1201_DECS +// BEGIN_GEMM_Wmma_i4_gfx1201_DECS +static constexpr size_t nInitParametersI4GemmGfx1201 = 1; +static const StringRef initParametersI4GemmGfx1201[nInitParametersI4GemmGfx1201]; +// END_GEMM_Wmma_i4_gfx1201_DECS + // BEGIN_CONV_Wmma_f16_gfx1201_DECS static constexpr size_t nInitParametersF16ConvGfx1201 = 36; static const StringRef initParametersF16ConvGfx1201[nInitParametersF16ConvGfx1201]; @@ -3220,6 +3242,8 @@ static const StringRef initParametersI8AttentionGfx1103[nInitParametersI8Attenti {"gfx1201_gemm_f16", {PopulateParamsWmma::initParametersF16GemmGfx1201, PopulateParamsWmma::nInitParametersF16GemmGfx1201}}, +{"gfx1201_gemm_i4", {PopulateParamsWmma::initParametersI4GemmGfx1201, PopulateParamsWmma::nInitParametersI4GemmGfx1201}}, + {"gfx1201_gemm_i8", {PopulateParamsWmma::initParametersI8GemmGfx1201, PopulateParamsWmma::nInitParametersI8GemmGfx1201}}, {"gfx1201_conv_f16", {PopulateParamsWmma::initParametersF16ConvGfx1201, PopulateParamsWmma::nInitParametersF16ConvGfx1201}}, diff --git a/mlir/lib/Dialect/Rock/IR/RockDialect.cpp b/mlir/lib/Dialect/Rock/IR/RockDialect.cpp index 9dee200e73b7..716507caaecc 100644 --- a/mlir/lib/Dialect/Rock/IR/RockDialect.cpp +++ b/mlir/lib/Dialect/Rock/IR/RockDialect.cpp @@ -778,6 +778,10 @@ static LogicalResult verifyGemmTypes(Operation *op, AmdArchInfo archInfo, if (isGfx1250) isValidTypeA = isValidTypeA || elemTypeA.isF32(); + // gfx12 (RDNA4) additionally supports int4 (iu4 16x16x32) + if (isRdna4) + isValidTypeA = isValidTypeA || elemTypeA.isInteger(4); + // gfx11 doesn't support float8 types if (isGfx11 && isFloat8Type(elemTypeA)) isValidTypeA = false; @@ -788,6 +792,9 @@ static LogicalResult verifyGemmTypes(Operation *op, AmdArchInfo archInfo, if (isGfx1250) return op->emitOpError( "Wmma supports only F32/F16/BF16/int8/E4M3/E5M2 data types"); + if (isRdna4) + return op->emitOpError( + "Wmma supports only F16/BF16/int4/int8/E4M3/E5M2 data types"); return op->emitOpError( "Wmma supports only F16/BF16/int8/E4M3/E5M2 data types"); } @@ -854,6 +861,13 @@ static LogicalResult verifyConvOp(RockConvInterface convOp) { Operation *op = convOp.getOperation(); RockGemmWrapperInterface gemmOp = cast(*convOp); + // i4 is currently wired for GEMM only. It is allowed in GemmInputTypes + // (shared with conv), so reject it here cleanly rather than fatal-erroring + // later in tuning-parameter selection (which has no conv i4 entry). + if (gemmOp.getAType().isInteger(4)) + return op->emitOpError( + "i4 convolution is not yet supported; i4 WMMA is GEMM-only"); + if (failed(verifyGemmTypes(gemmOp))) return failure(); diff --git a/mlir/lib/Dialect/Rock/IR/WmmaInsnGroup.cpp b/mlir/lib/Dialect/Rock/IR/WmmaInsnGroup.cpp index 5f21353e7ba9..207244709031 100644 --- a/mlir/lib/Dialect/Rock/IR/WmmaInsnGroup.cpp +++ b/mlir/lib/Dialect/Rock/IR/WmmaInsnGroup.cpp @@ -94,6 +94,10 @@ static const llvm::DenseMap &getWmmaInsnMapGfx12() { {ROCDL::wmma_f32_16x16x16_bf16::getOperationName(), 8, 8, 16, 16}}, {{WmmaTypeId::I8_To_I32_TyId, 16}, {ROCDL::wmma_i32_16x16x16_iu8::getOperationName(), 8, 8, 16, 16}}, + // iu4 is 16x16x32 (K=32, doubled vs i8) and packs 16 nibbles per lane, + // so inputVectorLen is 16 (vs 8 for i8). + {{WmmaTypeId::I4_To_I32_TyId, 32}, + {ROCDL::wmma_i32_16x16x32_iu4::getOperationName(), 16, 8, 16, 16}}, // FP8/BF8 {{WmmaTypeId::Fp8Fp8_To_F32_TyId, 16}, @@ -252,10 +256,13 @@ FailureOr WmmaInsn::select(mlir::Type elementTypeA, // Use gfx12 only if we don't have a selected instruction and not gfx11 if (!insnInfo && !isGfx11) { auto &gfx12Map = getWmmaInsnMapGfx12(); - auto it = gfx12Map.find({typeId, 16}); + // iu4 uses the 16x16x32 instruction (K=32); all other gfx12 wmma ops use + // K=16. + int64_t gfx12K = (typeId == WmmaTypeId::I4_To_I32_TyId) ? 32 : 16; + auto it = gfx12Map.find({typeId, gfx12K}); if (it != gfx12Map.end()) { insnInfo = &it->second; - selectedKDim = 16; + selectedKDim = gfx12K; LLVM_DEBUG(llvm::dbgs() << "Selected gfx12 instruction: " << insnInfo->insn << "\n"); } diff --git a/mlir/lib/Dialect/Rock/Tuning/ParamLookupTable.cpp b/mlir/lib/Dialect/Rock/Tuning/ParamLookupTable.cpp index 2dad1515231f..fc990c4e5d84 100644 --- a/mlir/lib/Dialect/Rock/Tuning/ParamLookupTable.cpp +++ b/mlir/lib/Dialect/Rock/Tuning/ParamLookupTable.cpp @@ -136,6 +136,7 @@ std::string ParamLookupTable::getDataTypeString(Type dataType) { // Normalize integer types by bitwidth unsigned bitwidth = dataType.getIntOrFloatBitWidth(); switch (bitwidth) { + case 4: case 8: return "i" + std::to_string(bitwidth); default: diff --git a/mlir/test/Dialect/Rock/conv_i4_unsupported.mlir b/mlir/test/Dialect/Rock/conv_i4_unsupported.mlir new file mode 100644 index 000000000000..ece09836fd0f --- /dev/null +++ b/mlir/test/Dialect/Rock/conv_i4_unsupported.mlir @@ -0,0 +1,21 @@ +// i4 is added to the shared GemmInputTypes, so rock.conv accepts i4 operands. +// Since i4 WMMA is wired for GEMM only (no conv tuning/lowering yet), the conv +// verifier must reject i4 cleanly rather than fatal-erroring later in tuning. +// RUN: not rocmlir-opt %s 2>&1 | FileCheck %s + +// CHECK: 'rock.conv' op i4 convolution is not yet supported + +#map = affine_map<(d0, d1, d2, d3, d4) -> (d1 * 64 + d4)> +#map1 = affine_map<(d0, d1, d2, d3, d4) -> (((d0 * 64 + d2) * 8 + d3) * 8 + d4)> +#transform_map = #rock.transform_map<#map by [ ["raw"] at [0]>, [] at []>, [] at []>, [] at []>] bounds = [1, 64, 1, 1, 64] -> [4096]> +#transform_map1 = #rock.transform_map<#map1 by [ ["raw"] at [0]>, [] at []>] bounds = [8, 1, 64, 8, 8] -> [32768]> +#transform_map2 = #rock.transform_map<#map1 by [ ["raw"] at [0]>, [] at []>] bounds = [8, 1, 64, 8, 8] -> [32768]> +module attributes {mhal.arch = "amdgcn-amd-amdhsa:gfx1201"} { + func.func @rock_conv_gk01c_ngc01_ngk01_0(%arg0: memref<4096xi4>, %arg1: memref<32768xi4>, %arg2: memref<32768xi32>) attributes {mhal.arch = "amdgcn-amd-amdhsa:gfx1201", rock.enable_splitk_for_tuning, rock.kernel = 0 : i32, rock.num_chiplets = 1 : i64, rock.num_cu = 12 : i32} { + %0 = rock.transform %arg0 by #transform_map : memref<4096xi4> to memref<1x64x1x1x64xi4> + %1 = rock.transform %arg1 by #transform_map1 : memref<32768xi4> to memref<8x1x64x8x8xi4> + %2 = rock.transform %arg2 by #transform_map2 : memref<32768xi32> to memref<8x1x64x8x8xi32> + rock.conv(%0, %1, %2) features = wmma|dot|atomic_add|atomic_add_bf16|atomic_add_f16|atomic_fmax_f32 {dilations = [1 : index, 1 : index], filter_layout = ["g", "k", "0", "1", "c"], input_layout = ["ni", "gi", "ci", "0i", "1i"], output_layout = ["no", "go", "ko", "0o", "1o"], padding = [0 : index, 0 : index, 0 : index, 0 : index], strides = [1 : index, 1 : index]} : memref<1x64x1x1x64xi4>, memref<8x1x64x8x8xi4>, memref<8x1x64x8x8xi32> + return + } +} diff --git a/mlir/test/Dialect/Rock/wmma_gemm_i4.mlir b/mlir/test/Dialect/Rock/wmma_gemm_i4.mlir new file mode 100644 index 000000000000..ae592ec392e5 --- /dev/null +++ b/mlir/test/Dialect/Rock/wmma_gemm_i4.mlir @@ -0,0 +1,10 @@ +// Checks that an i4 (iu4) GEMM on gfx1201 lowers through rock to the +// 16x16x32 iu4 WMMA instruction. +// RUN: rocmlir-gen --operation gemm --arch gfx1201 -t i4 -wmma=on -g 1 -m 256 -n 256 -k 256 | rocmlir-driver --arch gfx1201 --kernel-pipeline=gpu | FileCheck %s --check-prefix=AMDGPU +// RUN: rocmlir-gen --operation gemm --arch gfx1201 -t i4 -wmma=on -g 1 -m 256 -n 256 -k 256 | rocmlir-driver --arch gfx1201 --kernel-pipeline=gpu,rocdl | FileCheck %s --check-prefix=ROCDL +// iu4 WMMA is gfx12-only: gfx11 must reject it. +// RUN: not rocmlir-gen --operation gemm --arch gfx1100 -t i4 -wmma=on -g 1 -m 128 -n 128 -k 128 2>&1 | FileCheck %s --check-prefix=GFX11 + +// AMDGPU: amdgpu.wmma 16x16x32 {{.*}} : vector<16xi4>, vector<16xi4>, vector<8xi32> +// ROCDL: rocdl.wmma.i32.16x16x32.iu4 {{.*}} : (vector<2xi32>, vector<2xi32>, vector<8xi32>) -> vector<8xi32> +// GFX11: Wmma supports only F16/BF16/int8 diff --git a/mlir/tools/rocmlir-gen/rocmlir-gen.cpp b/mlir/tools/rocmlir-gen/rocmlir-gen.cpp index 72be9954bda7..4e29f66c9265 100644 --- a/mlir/tools/rocmlir-gen/rocmlir-gen.cpp +++ b/mlir/tools/rocmlir-gen/rocmlir-gen.cpp @@ -534,11 +534,11 @@ static llvm::cl::alias outTypeAliasShortC("tc", llvm::cl::aliasopt(outputDataType)); // Convenience setter for when you need all the data types the same or when you -// want the default (32-bit output) behavior for i8 or 8-bit floats. Also allows -// [a]_[b] syntax for mixed-type operations. +// want the default (32-bit output) behavior for i4/i8 or 8-bit floats. Also +// allows [a]_[b] syntax for mixed-type operations. static llvm::cl::opt dataTypeAlias( "t", - llvm::cl::desc("Data type selector. Extends i8 to i32 output and 8-bit " + llvm::cl::desc("Data type selector. Extends i4/i8 to i32 output and 8-bit " "floats to f32 output"), llvm::cl::value_desc("Type or Type_Type for mixed-type kernels."), llvm::cl::cb([](std::string v) { @@ -554,7 +554,7 @@ static llvm::cl::opt dataTypeAlias( } if (outputDataType.getNumOccurrences() == 0 || outputDataType.empty()) { - if (val == "i8") + if (val == "i8" || val == "i4") outputDataType = "i32"; else if (val.starts_with("f8") || val.starts_with("fp8") || val.starts_with("bf8") || val.starts_with("f4E2M1FN")) @@ -1618,6 +1618,7 @@ static Type typeFromString(StringRef name, MLIRContext *ctx) { .Case("f16", Float16Type::get(ctx)) .Case("fp16", Float16Type::get(ctx)) .Case("bf16", BFloat16Type::get(ctx)) + .Case("i4", IntegerType::get(ctx, 4)) .Case("i8", IntegerType::get(ctx, 8)) .Case("i32", IntegerType::get(ctx, 32)) .Case("f4E2M1FN", Float4E2M1FNType::get(ctx))