[Rock] Add i4 (iu4 16x16x32) WMMA GEMM support for gfx12

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.

mlir/include/mlir/Dialect/Rock/IR/RockOps.td

@@ -86,8 +86,8 @@ class Rock_ConvOpBase<string mnemonic, list<Type> inputTypes = [F32, F16, BF16],
   let results = (outs Optional<AnyRankedTensor>:$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];

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}},

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<RockGemmWrapperInterface>(*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();
 

mlir/lib/Dialect/Rock/IR/WmmaInsnGroup.cpp

@@ -94,6 +94,10 @@ static const llvm::DenseMap<WmmaInsnKey, WmmaInsnInfo> &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> 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");
     }

mlir/lib/Dialect/Rock/Tuning/ParamLookupTable.cpp

@@ -136,6 +136,7 @@ std::string ParamLookupTable<ParamsType>::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:

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 [<Unmerge{64, 64} ["k", "c"] at [1, 4] -> ["raw"] at [0]>, <AddDim{1} ["g"] at [0] -> [] at []>, <AddDim{1} ["0"] at [2] -> [] at []>, <AddDim{1} ["1"] at [3] -> [] at []>] bounds = [1, 64, 1, 1, 64] -> [4096]>
+#transform_map1 = #rock.transform_map<#map1 by [<Unmerge{8, 64, 8, 8} ["ni", "ci", "0i", "1i"] at [0, 2, 3, 4] -> ["raw"] at [0]>, <AddDim{1} ["gi"] at [1] -> [] at []>] bounds = [8, 1, 64, 8, 8] -> [32768]>
+#transform_map2 = #rock.transform_map<#map1 by [<Unmerge{8, 64, 8, 8} ["no", "ko", "0o", "1o"] at [0, 2, 3, 4] -> ["raw"] at [0]>, <AddDim{1} ["go"] at [1] -> [] 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
+  }
+}

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

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<std::string> 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<void, std::string>([](std::string v) {
@@ -554,7 +554,7 @@ static llvm::cl::opt<std::string> 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))