[MFMA] Add 16x16x16 bf16/f16 support with fly-fix-bitcast-width pass

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 = [{

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
 

lib/Dialect/Fly/Transforms/ConvertAtomCallToSSAForm.cpp

@@ -123,13 +123,30 @@ class FlyConvertAtomCallToSSAFormPass
       Value bVal = mmaOp.getB();
       Value cVal = mmaOp.getC();
 
+      auto mmaAtomTy = cast<MmaAtomType>(mmaOp.getMmaAtom().getType());
+
+      auto narrowToMmaWidth = [&](Value &val, LayoutAttr mmaLayout) {
+        auto regVecTy = dyn_cast<VectorType>(val.getType());
+        if (!regVecTy)
+          return;
+        LayoutBuilder<LayoutAttr> 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<MakeViewOp>().getIter();
         aVal = PtrLoadOp::create(builder, loc, RegMem2SSAType(aTy, true), aIter).getResult();
+        narrowToMmaWidth(aVal, cast<LayoutAttr>(mmaAtomTy.getThrValLayoutA()));
       }
       if (bEligible) {
         Value bIter = bVal.getDefiningOp<MakeViewOp>().getIter();
         bVal = PtrLoadOp::create(builder, loc, RegMem2SSAType(bTy, true), bIter).getResult();
+        narrowToMmaWidth(bVal, cast<LayoutAttr>(mmaAtomTy.getThrValLayoutB()));
       }
       if (cEligible) {
         Value cIter = cVal.getDefiningOp<MakeViewOp>().getIter();

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<FlyFixBitcastWidthPass> {
+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<vector::ExtractStridedSliceOp> extractsToProtect;
+    moduleOp->walk([&](vector::ExtractStridedSliceOp op) {
+      auto srcVecTy = cast<VectorType>(op->getOperand(0).getType());
+      auto dstVecTy = cast<VectorType>(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<LLVM::BitcastOp>(src.getDefiningOp()))
+        src = bc.getArg();
+      // Only protect if the ultimate source is a wider integer type.
+      if (auto srcIntTy = dyn_cast<IntegerType>(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<LLVM::BitcastOp> toProtect;
+
+    // Also find direct width-mismatched bitcasts (not from extract chains)
+    // that need to be rewritten with freeze to block canonicalization.
+    SmallVector<LLVM::BitcastOp> 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<IntegerType>(ty))
+          return intTy.getWidth();
+        if (auto vecTy = dyn_cast<VectorType>(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<LLVM::FreezeOp>(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<IntegerType>(ty))
+          return intTy.getWidth();
+        if (auto vecTy = dyn_cast<VectorType>(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

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<Value> 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<VectorType>(val.getType());
+    auto dstVecTy = dyn_cast<VectorType>(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)) {                                                   \

lib/Dialect/FlyROCDL/CDNA4/MmaAtom.cpp

@@ -225,12 +225,33 @@ FailureOr<Value> 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<VectorType>(val.getType());
+    auto dstVecTy = dyn_cast<VectorType>(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<int64_t>{*getFieldIndex(AtomStateField::ScaleA)});

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

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<!fly_rocdl.cdna3.mfma<32x32x8, (bf16, bf16) -> 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<!fly_rocdl.cdna3.mfma<16x16x16, (bf16, bf16) -> 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<!fly_rocdl.cdna3.mfma<16x16x16, (bf16, bf16) -> 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<!fly_rocdl.cdna3.mfma<16x16x16, (bf16, bf16) -> 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<!fly_rocdl.cdna3.mfma<16x16x16, (bf16, bf16) -> 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<!fly_rocdl.cdna3.mfma<16x16x16, (f16, f16) -> f32>>
+  // CHECK: %[[RES:.*]] = rocdl.mfma.f32.16x16x16f16 %[[A]], %[[B]], %[[C]]
+  %res = fly.mma_atom_call_ssa(%atom, %a, %b, %c) : (
+    !fly.mma_atom<!fly_rocdl.cdna3.mfma<16x16x16, (f16, f16) -> 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<!fly_rocdl.cdna3.mfma<16x16x16, (f16, f16) -> 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<!fly_rocdl.cdna3.mfma<16x16x16, (f16, f16) -> f32>>,
+    vector<8xf16>, vector<8xf16>, vector<4xf32>) -> vector<4xf32>
+  return %res : vector<4xf32>
+}