From fc1da1423cf97debe3fb45ab9cdfba5676a6d77b Mon Sep 17 00:00:00 2001 From: root Date: Fri, 14 Aug 2026 10:58:27 +0000 Subject: [PATCH] [MFMA] Add 16x16x16 bf16/f16 support with fly-fix-bitcast-width pass MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fix MFMA 16x16x16 bf16/f16 lowering crash when BufferCopy128b loads 128-bit values that feed 64-bit MFMA operands. The MLIR canonicalizer folds extract_strided_slice + bitcast chains back to the wider source, producing invalid width-changing bitcasts (e.g. i128 → vector<4xi16>) that LLVM rejects. Three layers of defense: 1. emitAtomCallSSA matchWidth (CDNA3/CDNA4 MmaAtom.cpp): handles the SSA path with vector extract_strided_slice when source is wider than the MFMA operand. 2. ConvertAtomCallToSSAForm narrowToMmaWidth: narrows register values to match the MMA atom's expected operand width at pass 06. 3. fly-fix-bitcast-width pass (new): runs before the canonicalizer, inserts llvm.freeze on narrowing extract_strided_slice results whose source traces back to a wider integer type, blocking the canonicalizer from folding the chain into an invalid bitcast. Verified correct for mma_k=32 block_n=32 (no regression), mma_k=16 block_n=32, and mma_k=16 block_n=64. Co-Authored-By: Claude --- .../flydsl/Dialect/Fly/Transforms/Passes.td | 16 ++ lib/Dialect/Fly/CMakeLists.txt | 1 + .../Transforms/ConvertAtomCallToSSAForm.cpp | 17 +++ .../Fly/Transforms/FixBitcastWidth.cpp | 141 ++++++++++++++++++ lib/Dialect/FlyROCDL/CDNA3/MmaAtom.cpp | 34 ++++- lib/Dialect/FlyROCDL/CDNA4/MmaAtom.cpp | 33 +++- python/flydsl/compiler/backends/rocm.py | 1 + tests/mlir/Conversion/mma_atom.mlir | 72 +++++++++ tests/mlir/Conversion/mma_atom_16x16x16.mlir | 65 ++++++++ 9 files changed, 368 insertions(+), 12 deletions(-) create mode 100644 lib/Dialect/Fly/Transforms/FixBitcastWidth.cpp create mode 100644 tests/mlir/Conversion/mma_atom_16x16x16.mlir diff --git a/include/flydsl/Dialect/Fly/Transforms/Passes.td b/include/flydsl/Dialect/Fly/Transforms/Passes.td index 9661b173e..b09d7e065 100644 --- a/include/flydsl/Dialect/Fly/Transforms/Passes.td +++ b/include/flydsl/Dialect/Fly/Transforms/Passes.td @@ -92,6 +92,22 @@ def FlyIntSwizzleSimplifyPass : Pass<"fly-int-swizzle-simplify"> { ]; } +def FlyFixBitcastWidthPass : Pass<"fly-fix-bitcast-width"> { + let summary = "Fix width-mismatched llvm.bitcast ops after canonicalization"; + let description = [{ + Rewrites llvm.bitcast ops where the source is wider than the destination + (e.g. i128 -> vector<4xi16>) into llvm.trunc + llvm.bitcast with matching + widths. This prevents LLVM codegen assertion failures when the MLIR + canonicalizer folds register load/store chains into direct bitcasts that + cross bit-width boundaries (e.g. BufferCopy128b load feeding a 16x16x16 + MFMA operand). + }]; + + let dependentDialects = [ + "LLVM::LLVMDialect" + ]; +} + def FlyPromoteRegMemToVectorSSAPass : Pass<"fly-promote-regmem-to-vectorssa"> { let summary = "Promote register memory to vector SSA values"; let description = [{ diff --git a/lib/Dialect/Fly/CMakeLists.txt b/lib/Dialect/Fly/CMakeLists.txt index e33ebaaf7..9a82e5848 100644 --- a/lib/Dialect/Fly/CMakeLists.txt +++ b/lib/Dialect/Fly/CMakeLists.txt @@ -13,6 +13,7 @@ add_mlir_dialect_library(MLIRFlyDialect Transforms/Canonicalize.cpp Transforms/RewriteFuncSignature.cpp Transforms/ConvertAtomCallToSSAForm.cpp + Transforms/FixBitcastWidth.cpp Transforms/PromoteRegMemToVectorSSA.cpp Transforms/IntSwizzleSimplify.cpp diff --git a/lib/Dialect/Fly/Transforms/ConvertAtomCallToSSAForm.cpp b/lib/Dialect/Fly/Transforms/ConvertAtomCallToSSAForm.cpp index 6b9c81cf0..9824b0c9a 100644 --- a/lib/Dialect/Fly/Transforms/ConvertAtomCallToSSAForm.cpp +++ b/lib/Dialect/Fly/Transforms/ConvertAtomCallToSSAForm.cpp @@ -123,13 +123,30 @@ class FlyConvertAtomCallToSSAFormPass Value bVal = mmaOp.getB(); Value cVal = mmaOp.getC(); + auto mmaAtomTy = cast(mmaOp.getMmaAtom().getType()); + + auto narrowToMmaWidth = [&](Value &val, LayoutAttr mmaLayout) { + auto regVecTy = dyn_cast(val.getType()); + if (!regVecTy) + return; + LayoutBuilder lb(builder.getContext()); + int32_t mmaCosize = layoutCosize(lb, mmaLayout).getLeafAsInt().getValue(); + if (regVecTy.getNumElements() > mmaCosize) { + val = vector::ExtractStridedSliceOp::create( + builder, loc, val, /*offsets=*/{0}, + /*sizes=*/{mmaCosize}, /*strides=*/{1}); + } + }; + if (aEligible) { Value aIter = aVal.getDefiningOp().getIter(); aVal = PtrLoadOp::create(builder, loc, RegMem2SSAType(aTy, true), aIter).getResult(); + narrowToMmaWidth(aVal, cast(mmaAtomTy.getThrValLayoutA())); } if (bEligible) { Value bIter = bVal.getDefiningOp().getIter(); bVal = PtrLoadOp::create(builder, loc, RegMem2SSAType(bTy, true), bIter).getResult(); + narrowToMmaWidth(bVal, cast(mmaAtomTy.getThrValLayoutB())); } if (cEligible) { Value cIter = cVal.getDefiningOp().getIter(); diff --git a/lib/Dialect/Fly/Transforms/FixBitcastWidth.cpp b/lib/Dialect/Fly/Transforms/FixBitcastWidth.cpp new file mode 100644 index 000000000..e54ac85f7 --- /dev/null +++ b/lib/Dialect/Fly/Transforms/FixBitcastWidth.cpp @@ -0,0 +1,141 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2025 FlyDSL Project Contributors + +#include "mlir/Dialect/LLVMIR/LLVMDialect.h" +#include "mlir/Dialect/Vector/IR/VectorOps.h" +#include "mlir/IR/Builders.h" +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/Pass/Pass.h" + +#include "flydsl/Dialect/Fly/Transforms/Passes.h" + +using namespace mlir; + +namespace mlir { +namespace fly { +#define GEN_PASS_DEF_FLYFIXBITCASTWIDTHPASS +#include "flydsl/Dialect/Fly/Transforms/Passes.h.inc" +} // namespace fly +} // namespace mlir + +namespace { + +class FlyFixBitcastWidthPass + : public mlir::fly::impl::FlyFixBitcastWidthPassBase { +public: + using mlir::fly::impl::FlyFixBitcastWidthPassBase< + FlyFixBitcastWidthPass>::FlyFixBitcastWidthPassBase; + + void runOnOperation() override { + // This pass runs BEFORE the canonicalizer (between convert-fly-to-rocdl and + // canonicalize). It inserts llvm.freeze on vector.extract_strided_slice results + // that feed llvm.bitcast ops, preventing the canonicalizer from folding the + // extract+bitcast chain back to the wider source value. + // + // Without this, the canonicalizer traces: + // buffer_load(i128) → bitcast(vector<8xbf16>) → extract[0:4](vector<4xbf16>) + // → bitcast(vector<4xi16>) + // and folds it into: + // buffer_load(i128) → bitcast(vector<4xi16>) [INVALID: 128 != 64 bits] + auto moduleOp = getOperation(); + + // Protect extract_strided_slice ops whose source traces back (through + // bitcasts) to a wider integer type (e.g. i128 from buffer_load). Only + // these would produce width-mismatched bitcasts after canonicalization. + SmallVector extractsToProtect; + moduleOp->walk([&](vector::ExtractStridedSliceOp op) { + auto srcVecTy = cast(op->getOperand(0).getType()); + auto dstVecTy = cast(op.getResult().getType()); + if (srcVecTy.getNumElements() <= dstVecTy.getNumElements()) + return; + // Trace through bitcasts to find the ultimate source type. + Value src = op->getOperand(0); + while (auto bc = dyn_cast_or_null(src.getDefiningOp())) + src = bc.getArg(); + // Only protect if the ultimate source is a wider integer type. + if (auto srcIntTy = dyn_cast(src.getType())) { + int64_t dstBits = dstVecTy.getNumElements() * + dstVecTy.getElementType().getIntOrFloatBitWidth(); + if (srcIntTy.getWidth() > dstBits) + extractsToProtect.push_back(op); + } + }); + + // Dummy to keep the toProtect interface (unused now). + SmallVector toProtect; + + // Also find direct width-mismatched bitcasts (not from extract chains) + // that need to be rewritten with freeze to block canonicalization. + SmallVector directMismatch; + moduleOp->walk([&](LLVM::BitcastOp op) { + Type srcTy = op.getArg().getType(); + Type dstTy = op.getResult().getType(); + auto getBits = [](Type ty) -> int64_t { + if (auto intTy = dyn_cast(ty)) + return intTy.getWidth(); + if (auto vecTy = dyn_cast(ty)) + return vecTy.getNumElements() * vecTy.getElementType().getIntOrFloatBitWidth(); + return 0; + }; + int64_t srcBits = getBits(srcTy); + int64_t dstBits = getBits(dstTy); + if (srcBits > 0 && dstBits > 0 && srcBits != dstBits) + directMismatch.push_back(op); + }); + + if (extractsToProtect.empty() && directMismatch.empty()) + return; + + OpBuilder builder(moduleOp->getContext()); + + // Insert freeze after each narrowing extract to block canonicalization + // from folding the extract chain back to the wider source. + for (auto extractOp : extractsToProtect) { + builder.setInsertionPointAfter(extractOp); + Value result = extractOp.getResult(); + Value frozen = LLVM::FreezeOp::create(builder, extractOp.getLoc(), + result.getType(), result); + result.replaceAllUsesExcept(frozen, frozen.getDefiningOp()); + } + + // For direct width-mismatched bitcasts (e.g. i128 → vector<4xbf16>), + // insert freeze on the source to prevent canonicalization from + // reintroducing them after other folds. + for (LLVM::BitcastOp op : directMismatch) { + // Skip if already protected by the extract chain fix above. + if (op->getOperand(0).getDefiningOp() && + isa(op->getOperand(0).getDefiningOp())) + continue; + + builder.setInsertionPoint(op); + Location loc = op.getLoc(); + Value src = op.getArg(); + Type dstTy = op.getResult().getType(); + + auto getBits = [](Type ty) -> int64_t { + if (auto intTy = dyn_cast(ty)) + return intTy.getWidth(); + if (auto vecTy = dyn_cast(ty)) + return vecTy.getNumElements() * vecTy.getElementType().getIntOrFloatBitWidth(); + return 0; + }; + int64_t srcBits = getBits(src.getType()); + int64_t dstBits = getBits(dstTy); + + // Bitcast to integer, truncate, then bitcast to destination. + Type srcIntTy = IntegerType::get(builder.getContext(), srcBits); + Type dstIntTy = IntegerType::get(builder.getContext(), dstBits); + + Value intVal = src; + if (src.getType() != srcIntTy) + intVal = LLVM::BitcastOp::create(builder, loc, srcIntTy, src); + Value truncated = LLVM::TruncOp::create(builder, loc, dstIntTy, intVal); + Value result = LLVM::BitcastOp::create(builder, loc, dstTy, truncated); + + op.getResult().replaceAllUsesWith(result); + op->erase(); + } + } +}; + +} // namespace diff --git a/lib/Dialect/FlyROCDL/CDNA3/MmaAtom.cpp b/lib/Dialect/FlyROCDL/CDNA3/MmaAtom.cpp index 25be7ceab..c69c548d8 100644 --- a/lib/Dialect/FlyROCDL/CDNA3/MmaAtom.cpp +++ b/lib/Dialect/FlyROCDL/CDNA3/MmaAtom.cpp @@ -3,6 +3,7 @@ #include "mlir/Dialect/LLVMIR/LLVMDialect.h" #include "mlir/Dialect/LLVMIR/ROCDLDialect.h" +#include "mlir/Dialect/Vector/IR/VectorOps.h" #include "mlir/IR/BuiltinTypes.h" #include "flydsl/Dialect/Fly/IR/FlyDialect.h" @@ -160,12 +161,33 @@ FailureOr MmaOpCDNA3_MFMAType::emitAtomCallSSA(OpBuilder &builder, Locati Type accElemTy = getElemTyAcc(); VectorType accTy = VectorType::get({accVecSize}, accElemTy); - if (a.getType() != abTyA) - a = LLVM::BitcastOp::create(builder, loc, abTyA, a); - if (b.getType() != abTyB) - b = LLVM::BitcastOp::create(builder, loc, abTyB, b); - if (c.getType() != accTy) - c = LLVM::BitcastOp::create(builder, loc, accTy, c); + auto matchWidth = [&](Value &val, Type targetTy) { + if (val.getType() == targetTy) + return; + auto srcVecTy = dyn_cast(val.getType()); + auto dstVecTy = dyn_cast(targetTy); + if (srcVecTy && dstVecTy) { + int64_t srcBits = srcVecTy.getNumElements() * + srcVecTy.getElementType().getIntOrFloatBitWidth(); + int64_t dstBits = dstVecTy.getNumElements() * + dstVecTy.getElementType().getIntOrFloatBitWidth(); + if (srcBits > dstBits && srcBits % dstBits == 0) { + int64_t fullCount = + srcBits / dstVecTy.getElementType().getIntOrFloatBitWidth(); + auto wideTy = + VectorType::get({fullCount}, dstVecTy.getElementType()); + val = LLVM::BitcastOp::create(builder, loc, wideTy, val); + val = vector::ExtractStridedSliceOp::create( + builder, loc, val, /*offsets=*/{0}, + /*sizes=*/{dstVecTy.getNumElements()}, /*strides=*/{1}); + return; + } + } + val = LLVM::BitcastOp::create(builder, loc, targetTy, val); + }; + matchWidth(a, abTyA); + matchWidth(b, abTyB); + matchWidth(c, accTy); #define DISPATCH_MFMA_SSA(M_, K_, PRED, OP) \ if (m == M_ && n == M_ && k == K_ && (PRED)) { \ diff --git a/lib/Dialect/FlyROCDL/CDNA4/MmaAtom.cpp b/lib/Dialect/FlyROCDL/CDNA4/MmaAtom.cpp index c96cca662..2f2bcd277 100644 --- a/lib/Dialect/FlyROCDL/CDNA4/MmaAtom.cpp +++ b/lib/Dialect/FlyROCDL/CDNA4/MmaAtom.cpp @@ -225,12 +225,33 @@ FailureOr MmaOpCDNA4_MFMAScaleType::emitAtomCallSSA(OpBuilder &builder, L Type accElemTy = getElemTyAcc(); VectorType accTy = VectorType::get({accVecSize}, accElemTy); - if (a.getType() != abTyA) - a = LLVM::BitcastOp::create(builder, loc, abTyA, a); - if (b.getType() != abTyB) - b = LLVM::BitcastOp::create(builder, loc, abTyB, b); - if (c.getType() != accTy) - c = LLVM::BitcastOp::create(builder, loc, accTy, c); + auto matchWidth = [&](Value &val, Type targetTy) { + if (val.getType() == targetTy) + return; + auto srcVecTy = dyn_cast(val.getType()); + auto dstVecTy = dyn_cast(targetTy); + if (srcVecTy && dstVecTy) { + int64_t srcBits = srcVecTy.getNumElements() * + srcVecTy.getElementType().getIntOrFloatBitWidth(); + int64_t dstBits = dstVecTy.getNumElements() * + dstVecTy.getElementType().getIntOrFloatBitWidth(); + if (srcBits > dstBits && srcBits % dstBits == 0) { + int64_t fullCount = + srcBits / dstVecTy.getElementType().getIntOrFloatBitWidth(); + auto wideTy = + VectorType::get({fullCount}, dstVecTy.getElementType()); + val = LLVM::BitcastOp::create(builder, loc, wideTy, val); + val = vector::ExtractStridedSliceOp::create( + builder, loc, val, /*offsets=*/{0}, + /*sizes=*/{dstVecTy.getNumElements()}, /*strides=*/{1}); + return; + } + } + val = LLVM::BitcastOp::create(builder, loc, targetTy, val); + }; + matchWidth(a, abTyA); + matchWidth(b, abTyB); + matchWidth(c, accTy); Value scaleA = LLVM::ExtractValueOp::create( builder, loc, atomVal, ArrayRef{*getFieldIndex(AtomStateField::ScaleA)}); diff --git a/python/flydsl/compiler/backends/rocm.py b/python/flydsl/compiler/backends/rocm.py index a9d865d29..ccfb9d3fa 100644 --- a/python/flydsl/compiler/backends/rocm.py +++ b/python/flydsl/compiler/backends/rocm.py @@ -86,6 +86,7 @@ def _pipeline_parts(self, *, compile_hints: dict) -> Tuple[List[str], str]: "fly-convert-atom-call-to-ssa-form", "fly-promote-regmem-to-vectorssa", "convert-fly-to-rocdl", + "fly-fix-bitcast-width", "canonicalize", f"gpu.module(convert-scf-to-cf,cse," f"convert-gpu-to-rocdl{{chipset={chip} index-bitwidth=0 runtime=HIP use-bare-ptr-memref-call-conv=true}}," diff --git a/tests/mlir/Conversion/mma_atom.mlir b/tests/mlir/Conversion/mma_atom.mlir index 02a34cf8e..d2bfbd3de 100644 --- a/tests/mlir/Conversion/mma_atom.mlir +++ b/tests/mlir/Conversion/mma_atom.mlir @@ -63,3 +63,75 @@ func.func @test_mma_atom_call_ssa_bf16_32x32x8( %res = fly.mma_atom_call_ssa(%atom, %a, %b, %c) : (!fly.mma_atom f32>>, vector<4xbf16>, vector<4xbf16>, vector<16xf32>) -> vector<16xf32> return %res : vector<16xf32> } + +// ---- 16x16x16 bf16: same-width input (regression guard) ---- + +// CHECK-LABEL: @test_mma_atom_call_ssa_bf16_16x16x16 +// CHECK-SAME: (%[[A:.*]]: vector<4xbf16>, %[[B:.*]]: vector<4xbf16>, %[[C:.*]]: vector<4xf32>) +func.func @test_mma_atom_call_ssa_bf16_16x16x16( + %a: vector<4xbf16>, + %b: vector<4xbf16>, + %c: vector<4xf32>) -> vector<4xf32> { + %atom = fly.make_mma_atom : !fly.mma_atom f32>> + // CHECK: %[[A_CAST:.*]] = llvm.bitcast %[[A]] : vector<4xbf16> to vector<4xi16> + // CHECK: %[[B_CAST:.*]] = llvm.bitcast %[[B]] : vector<4xbf16> to vector<4xi16> + // CHECK: %[[RES:.*]] = rocdl.mfma.f32.16x16x16bf16.1k %[[A_CAST]], %[[B_CAST]], %[[C]] + %res = fly.mma_atom_call_ssa(%atom, %a, %b, %c) : ( + !fly.mma_atom f32>>, + vector<4xbf16>, vector<4xbf16>, vector<4xf32>) -> vector<4xf32> + return %res : vector<4xf32> +} + +// ---- 16x16x16 bf16: wider input (the bug fix path) ---- + +// CHECK-LABEL: @test_mma_atom_call_ssa_bf16_16x16x16_wide +// CHECK-SAME: (%[[A:.*]]: vector<8xbf16>, %[[B:.*]]: vector<8xbf16>, %[[C:.*]]: vector<4xf32>) +func.func @test_mma_atom_call_ssa_bf16_16x16x16_wide( + %a: vector<8xbf16>, + %b: vector<8xbf16>, + %c: vector<4xf32>) -> vector<4xf32> { + %atom = fly.make_mma_atom : !fly.mma_atom f32>> + // CHECK: %[[A_I16:.*]] = llvm.bitcast %[[A]] : vector<8xbf16> to vector<8xi16> + // CHECK: %[[A_SLICE:.*]] = vector.extract_strided_slice %[[A_I16]] {offsets = [0], sizes = [4], strides = [1]} + // CHECK: %[[B_I16:.*]] = llvm.bitcast %[[B]] : vector<8xbf16> to vector<8xi16> + // CHECK: %[[B_SLICE:.*]] = vector.extract_strided_slice %[[B_I16]] {offsets = [0], sizes = [4], strides = [1]} + // CHECK: %[[RES:.*]] = rocdl.mfma.f32.16x16x16bf16.1k %[[A_SLICE]], %[[B_SLICE]], %[[C]] + %res = fly.mma_atom_call_ssa(%atom, %a, %b, %c) : ( + !fly.mma_atom f32>>, + vector<8xbf16>, vector<8xbf16>, vector<4xf32>) -> vector<4xf32> + return %res : vector<4xf32> +} + +// ---- 16x16x16 f16: same-width input (regression guard) ---- + +// CHECK-LABEL: @test_mma_atom_call_ssa_f16_16x16x16 +// CHECK-SAME: (%[[A:.*]]: vector<4xf16>, %[[B:.*]]: vector<4xf16>, %[[C:.*]]: vector<4xf32>) +func.func @test_mma_atom_call_ssa_f16_16x16x16( + %a: vector<4xf16>, + %b: vector<4xf16>, + %c: vector<4xf32>) -> vector<4xf32> { + %atom = fly.make_mma_atom : !fly.mma_atom f32>> + // CHECK: %[[RES:.*]] = rocdl.mfma.f32.16x16x16f16 %[[A]], %[[B]], %[[C]] + %res = fly.mma_atom_call_ssa(%atom, %a, %b, %c) : ( + !fly.mma_atom f32>>, + vector<4xf16>, vector<4xf16>, vector<4xf32>) -> vector<4xf32> + return %res : vector<4xf32> +} + +// ---- 16x16x16 f16: wider input ---- + +// CHECK-LABEL: @test_mma_atom_call_ssa_f16_16x16x16_wide +// CHECK-SAME: (%[[A:.*]]: vector<8xf16>, %[[B:.*]]: vector<8xf16>, %[[C:.*]]: vector<4xf32>) +func.func @test_mma_atom_call_ssa_f16_16x16x16_wide( + %a: vector<8xf16>, + %b: vector<8xf16>, + %c: vector<4xf32>) -> vector<4xf32> { + %atom = fly.make_mma_atom : !fly.mma_atom f32>> + // CHECK: %[[A_SLICE:.*]] = vector.extract_strided_slice %[[A]] {offsets = [0], sizes = [4], strides = [1]} + // CHECK: %[[B_SLICE:.*]] = vector.extract_strided_slice %[[B]] {offsets = [0], sizes = [4], strides = [1]} + // CHECK: %[[RES:.*]] = rocdl.mfma.f32.16x16x16f16 %[[A_SLICE]], %[[B_SLICE]], %[[C]] + %res = fly.mma_atom_call_ssa(%atom, %a, %b, %c) : ( + !fly.mma_atom f32>>, + vector<8xf16>, vector<8xf16>, vector<4xf32>) -> vector<4xf32> + return %res : vector<4xf32> +} diff --git a/tests/mlir/Conversion/mma_atom_16x16x16.mlir b/tests/mlir/Conversion/mma_atom_16x16x16.mlir new file mode 100644 index 000000000..4444baab7 --- /dev/null +++ b/tests/mlir/Conversion/mma_atom_16x16x16.mlir @@ -0,0 +1,65 @@ +// RUN: %fly-opt %s --fly-rewrite-func-signature --fly-canonicalize --fly-layout-lowering --convert-fly-to-rocdl | FileCheck %s + +// ---- 16x16x16 bf16: same-width input ---- + +// CHECK-LABEL: @test_bf16_16x16x16 +// CHECK: llvm.bitcast %{{.*}} : vector<4xbf16> to vector<4xi16> +// CHECK: rocdl.mfma.f32.16x16x16bf16.1k +func.func @test_bf16_16x16x16( + %a: vector<4xbf16>, + %b: vector<4xbf16>, + %c: vector<4xf32>) -> vector<4xf32> { + %atom = fly.make_mma_atom : !fly.mma_atom f32>> + %res = fly.mma_atom_call_ssa(%atom, %a, %b, %c) : ( + !fly.mma_atom f32>>, + vector<4xbf16>, vector<4xbf16>, vector<4xf32>) -> vector<4xf32> + return %res : vector<4xf32> +} + +// ---- 16x16x16 bf16: wider input (the bug fix) ---- + +// CHECK-LABEL: @test_bf16_16x16x16_wide +// CHECK: llvm.bitcast %{{.*}} : vector<8xbf16> to vector<8xi16> +// CHECK: vector.extract_strided_slice +// CHECK: rocdl.mfma.f32.16x16x16bf16.1k +func.func @test_bf16_16x16x16_wide( + %a: vector<8xbf16>, + %b: vector<8xbf16>, + %c: vector<4xf32>) -> vector<4xf32> { + %atom = fly.make_mma_atom : !fly.mma_atom f32>> + %res = fly.mma_atom_call_ssa(%atom, %a, %b, %c) : ( + !fly.mma_atom f32>>, + vector<8xbf16>, vector<8xbf16>, vector<4xf32>) -> vector<4xf32> + return %res : vector<4xf32> +} + +// ---- 16x16x16 f16: same-width ---- + +// CHECK-LABEL: @test_f16_16x16x16 +// CHECK: rocdl.mfma.f32.16x16x16f16 +func.func @test_f16_16x16x16( + %a: vector<4xf16>, + %b: vector<4xf16>, + %c: vector<4xf32>) -> vector<4xf32> { + %atom = fly.make_mma_atom : !fly.mma_atom f32>> + %res = fly.mma_atom_call_ssa(%atom, %a, %b, %c) : ( + !fly.mma_atom f32>>, + vector<4xf16>, vector<4xf16>, vector<4xf32>) -> vector<4xf32> + return %res : vector<4xf32> +} + +// ---- 16x16x16 f16: wider input ---- + +// CHECK-LABEL: @test_f16_16x16x16_wide +// CHECK: vector.extract_strided_slice +// CHECK: rocdl.mfma.f32.16x16x16f16 +func.func @test_f16_16x16x16_wide( + %a: vector<8xf16>, + %b: vector<8xf16>, + %c: vector<4xf32>) -> vector<4xf32> { + %atom = fly.make_mma_atom : !fly.mma_atom f32>> + %res = fly.mma_atom_call_ssa(%atom, %a, %b, %c) : ( + !fly.mma_atom f32>>, + vector<8xf16>, vector<8xf16>, vector<4xf32>) -> vector<4xf32> + return %res : vector<4xf32> +}