From 749a63c1b0a8121f9b0055d7f12929b1f4c0d9ad Mon Sep 17 00:00:00 2001 From: Omar Muhammad Date: Fri, 24 Jul 2026 16:48:18 +0000 Subject: [PATCH] Enabled coexec llvm for Flydsl --- lib/Conversion/FlyToROCDL/FlyToROCDL.cpp | 4 +- lib/Dialect/Fly/IR/FlyOps.cpp | 2 +- .../Transforms/PromoteRegMemToVectorSSA.cpp | 14 - lib/Dialect/FlyROCDL/CDNA3/CopyAtom.cpp | 14 +- lib/Dialect/FlyROCDL/CDNA3/MmaAtom.cpp | 4 +- lib/Dialect/FlyROCDL/CDNA4/MmaAtom.cpp | 16 +- lib/Dialect/FlyROCDL/GFX1250/CopyAtom.cpp | 2 +- lib/Dialect/FlyROCDL/GFX1250/MmaAtom.cpp | 7 +- lib/Dialect/FlyROCDL/GFX1250/MmaAtomScale.cpp | 28 +- lib/Dialect/FlyROCDL/Ops.cpp | 2 +- python/flydsl/compiler/backends/base.py | 3 + python/flydsl/compiler/backends/rocm.py | 21 +- python/flydsl/compiler/jit_function.py | 30 +- python/flydsl/expr/buffer_ops.py | 8 +- python/flydsl/expr/rocdl/__init__.py | 58 ++- python/flydsl/expr/rocdl/_rocdl_enums.py | 456 ++++++++++++++++++ python/flydsl/expr/rocdl/tdm_ops.py | 12 +- tests/mlir/Conversion/mma_atom_stateful.mlir | 6 +- tests/mlir/Conversion/mma_scale_gfx1250.mlir | 16 +- tests/mlir/Conversion/tdm_gfx1250.mlir | 4 +- tests/unit/test_external_llvm_codegen.py | 1 + 21 files changed, 611 insertions(+), 97 deletions(-) create mode 100644 python/flydsl/expr/rocdl/_rocdl_enums.py diff --git a/lib/Conversion/FlyToROCDL/FlyToROCDL.cpp b/lib/Conversion/FlyToROCDL/FlyToROCDL.cpp index 698564cc9..a9d4d032f 100644 --- a/lib/Conversion/FlyToROCDL/FlyToROCDL.cpp +++ b/lib/Conversion/FlyToROCDL/FlyToROCDL.cpp @@ -381,7 +381,7 @@ class PtrLoadOpLowering : public OpConversionPattern { ArrayAttr noAttrs; Value loaded = ROCDL::RawPtrBufferLoadOp::create( rewriter, loc, loadTy, bp.bufferRsrc(rewriter, loc), bp.swizzleByteOffset(rewriter, loc), - zero, zero, noAttrs, noAttrs, noAttrs); + zero, rewriter.getI32IntegerAttr(0), noAttrs, noAttrs, noAttrs); rewriter.replaceOp(op, loaded); return success(); } else { @@ -425,7 +425,7 @@ class PtrStoreOpLowering : public OpConversionPattern { Value zero = arith::ConstantIntOp::create(rewriter, loc, 0, 32); ArrayAttr noAttrs; ROCDL::RawPtrBufferStoreOp::create(rewriter, loc, value, bp.bufferRsrc(rewriter, loc), - bp.swizzleByteOffset(rewriter, loc), zero, zero, noAttrs, + bp.swizzleByteOffset(rewriter, loc), zero, rewriter.getI32IntegerAttr(0), noAttrs, noAttrs, noAttrs); rewriter.eraseOp(op); return success(); diff --git a/lib/Dialect/Fly/IR/FlyOps.cpp b/lib/Dialect/Fly/IR/FlyOps.cpp index 41b02b1b9..9da5a6600 100644 --- a/lib/Dialect/Fly/IR/FlyOps.cpp +++ b/lib/Dialect/Fly/IR/FlyOps.cpp @@ -191,7 +191,7 @@ Type applyOffsetOnTensorLike(LayoutBuilder &builder, Type tensorLike llvm::LogicalResult OP::inferReturnTypes( \ mlir::MLIRContext *context, std::optional<::mlir::Location> location, \ mlir::ValueRange operands, mlir::DictionaryAttr attributes, \ - mlir::OpaqueProperties properties, mlir::RegionRange regions, \ + mlir::PropertyRef properties, mlir::RegionRange regions, \ llvm::SmallVectorImpl &inferredReturnTypes) //===----------------------------------------------------------------------===// diff --git a/lib/Dialect/Fly/Transforms/PromoteRegMemToVectorSSA.cpp b/lib/Dialect/Fly/Transforms/PromoteRegMemToVectorSSA.cpp index 7ef326586..6b553a8d9 100644 --- a/lib/Dialect/Fly/Transforms/PromoteRegMemToVectorSSA.cpp +++ b/lib/Dialect/Fly/Transforms/PromoteRegMemToVectorSSA.cpp @@ -25,20 +25,6 @@ using namespace mlir; using namespace mlir::fly; -namespace llvm { - -template <> struct DenseMapInfo : DenseMapInfo { - using Base = DenseMapInfo; - - static mlir::fly::MakePtrOp getEmptyKey() { return mlir::fly::MakePtrOp(Base::getEmptyKey()); } - - static mlir::fly::MakePtrOp getTombstoneKey() { - return mlir::fly::MakePtrOp(Base::getTombstoneKey()); - } -}; - -} // namespace llvm - namespace mlir { namespace fly { #define GEN_PASS_DEF_FLYPROMOTEREGMEMTOVECTORSSAPASS diff --git a/lib/Dialect/FlyROCDL/CDNA3/CopyAtom.cpp b/lib/Dialect/FlyROCDL/CDNA3/CopyAtom.cpp index 4b16116cb..25f4a4e46 100644 --- a/lib/Dialect/FlyROCDL/CDNA3/CopyAtom.cpp +++ b/lib/Dialect/FlyROCDL/CDNA3/CopyAtom.cpp @@ -89,7 +89,7 @@ FailureOr CopyOpCDNA3BufferCopyType::emitAtomCallSSA(OpBuilder &builder, }; // raw buffer load/store cachepolicy (0=cached, 2=nt) - Value aux = arith::ConstantIntOp::create(builder, loc, getCacheModifier(), 32); + auto aux = builder.getI32IntegerAttr(getCacheModifier()); ArrayAttr noAttrs; auto srcMemTy = srcTyArg ? dyn_cast(srcTyArg) : fly::MemRefType(); @@ -313,7 +313,7 @@ LogicalResult CopyOpCDNA3BufferCopyLDSType::emitAtomCall(OpBuilder &builder, Loc soffset = arith::DivUIOp::create(builder, loc, bits, eight); } - Value zero = arith::ConstantIntOp::create(builder, loc, 0, 32); + auto zero = builder.getI32IntegerAttr(0); Value size = arith::ConstantIntOp::create(builder, loc, sizeBytes, 32); BufferFatPtr bp(srcMemTy.getPointerType(), src); @@ -428,7 +428,7 @@ FailureOr CopyOpCDNA3BufferAtomicType::emitAtomCallSSA(OpBuilder &builder soffset = arith::DivUIOp::create(builder, loc, bits, eight); } - Value zero = arith::ConstantIntOp::create(builder, loc, 0, 32); + auto zero = builder.getI32IntegerAttr(0); ArrayAttr noAttrs; AtomicOp op = getAtomicOp().getValue(); @@ -437,21 +437,21 @@ FailureOr CopyOpCDNA3BufferAtomicType::emitAtomCallSSA(OpBuilder &builder case AtomicOp::Add: if (!isFloat) return failure(); - ROCDL::RawPtrBufferAtomicFaddOp::create(builder, loc, src, dstRsrc, dstOff, soffset, zero, + ROCDL::RawPtrBufferAtomicFaddOp::create(builder, loc, src.getType(), src, dstRsrc, dstOff, soffset, zero, noAttrs, noAttrs, noAttrs); break; case AtomicOp::Max: if (isFloat) - ROCDL::RawPtrBufferAtomicFmaxOp::create(builder, loc, src, dstRsrc, dstOff, soffset, zero, + ROCDL::RawPtrBufferAtomicFmaxOp::create(builder, loc, src.getType(), src, dstRsrc, dstOff, soffset, zero, noAttrs, noAttrs, noAttrs); else - ROCDL::RawPtrBufferAtomicSmaxOp::create(builder, loc, src, dstRsrc, dstOff, soffset, zero, + ROCDL::RawPtrBufferAtomicSmaxOp::create(builder, loc, src.getType(), src, dstRsrc, dstOff, soffset, zero, noAttrs, noAttrs, noAttrs); break; case AtomicOp::Min: if (isFloat) return failure(); - ROCDL::RawPtrBufferAtomicUminOp::create(builder, loc, src, dstRsrc, dstOff, soffset, zero, + ROCDL::RawPtrBufferAtomicUminOp::create(builder, loc, src.getType(), src, dstRsrc, dstOff, soffset, zero, noAttrs, noAttrs, noAttrs); break; default: diff --git a/lib/Dialect/FlyROCDL/CDNA3/MmaAtom.cpp b/lib/Dialect/FlyROCDL/CDNA3/MmaAtom.cpp index 25be7ceab..21199edc7 100644 --- a/lib/Dialect/FlyROCDL/CDNA3/MmaAtom.cpp +++ b/lib/Dialect/FlyROCDL/CDNA3/MmaAtom.cpp @@ -169,8 +169,8 @@ FailureOr MmaOpCDNA3_MFMAType::emitAtomCallSSA(OpBuilder &builder, Locati #define DISPATCH_MFMA_SSA(M_, K_, PRED, OP) \ if (m == M_ && n == M_ && k == K_ && (PRED)) { \ - auto zeroAttr = builder.getI32IntegerAttr(0); \ - return ROCDL::OP::create(builder, loc, accTy, a, b, c, zeroAttr, zeroAttr, zeroAttr) \ + return ROCDL::OP::create(builder, loc, accTy, a, b, c, 0u, 0u, \ + ROCDL::MFMAPermB::none) \ .getResult(); \ } diff --git a/lib/Dialect/FlyROCDL/CDNA4/MmaAtom.cpp b/lib/Dialect/FlyROCDL/CDNA4/MmaAtom.cpp index c96cca662..13c9c9a7d 100644 --- a/lib/Dialect/FlyROCDL/CDNA4/MmaAtom.cpp +++ b/lib/Dialect/FlyROCDL/CDNA4/MmaAtom.cpp @@ -237,20 +237,20 @@ FailureOr MmaOpCDNA4_MFMAScaleType::emitAtomCallSSA(OpBuilder &builder, L Value scaleB = LLVM::ExtractValueOp::create( builder, loc, atomVal, ArrayRef{*getFieldIndex(AtomStateField::ScaleB)}); - auto cbszAttr = builder.getI32IntegerAttr(*aTypeCode); - auto blgpAttr = builder.getI32IntegerAttr(*bTypeCode); - auto opselAAttr = builder.getI32IntegerAttr(getOpselA()); - auto opselBAttr = builder.getI32IntegerAttr(getOpselB()); + auto cbszFmt = static_cast(*aTypeCode); + auto blgpFmt = static_cast(*bTypeCode); + uint32_t opselA = getOpselA(); + uint32_t opselB = getOpselB(); if (m == 16 && n == 16 && k == 128) { - return ROCDL::mfma_scale_f32_16x16x128_f8f6f4::create(builder, loc, accTy, a, b, c, cbszAttr, - blgpAttr, opselAAttr, scaleA, opselBAttr, + return ROCDL::mfma_scale_f32_16x16x128_f8f6f4::create(builder, loc, accTy, a, b, c, cbszFmt, + blgpFmt, opselA, scaleA, opselB, scaleB) .getResult(); } if (m == 32 && n == 32 && k == 64) { - return ROCDL::mfma_scale_f32_32x32x64_f8f6f4::create(builder, loc, accTy, a, b, c, cbszAttr, - blgpAttr, opselAAttr, scaleA, opselBAttr, + return ROCDL::mfma_scale_f32_32x32x64_f8f6f4::create(builder, loc, accTy, a, b, c, cbszFmt, + blgpFmt, opselA, scaleA, opselB, scaleB) .getResult(); } diff --git a/lib/Dialect/FlyROCDL/GFX1250/CopyAtom.cpp b/lib/Dialect/FlyROCDL/GFX1250/CopyAtom.cpp index 676be7155..14ff05ec2 100644 --- a/lib/Dialect/FlyROCDL/GFX1250/CopyAtom.cpp +++ b/lib/Dialect/FlyROCDL/GFX1250/CopyAtom.cpp @@ -503,7 +503,7 @@ LogicalResult CopyOpGFX1250TDMType::emitAtomCall(OpBuilder &builder, Location lo vector::FromElementsOp::create(builder, loc, VectorType::get({8}, builder.getI32Type()), ValueRange{zero, zero, zero, zero, zero, zero, zero, zero}); - uint32_t cachePolicy = static_cast(getCacheModifier()); + auto cachePolicy = builder.getI32IntegerAttr(static_cast(getCacheModifier())); ArrayAttr noAliasScopes; if (isLoad) ROCDL::TensorLoadToLDSOp::create(builder, loc, dgroup0, dgroup1, dg2, dg3, dg4, cachePolicy, diff --git a/lib/Dialect/FlyROCDL/GFX1250/MmaAtom.cpp b/lib/Dialect/FlyROCDL/GFX1250/MmaAtom.cpp index 224cf74c9..6c9485bce 100644 --- a/lib/Dialect/FlyROCDL/GFX1250/MmaAtom.cpp +++ b/lib/Dialect/FlyROCDL/GFX1250/MmaAtom.cpp @@ -239,14 +239,13 @@ static FailureOr emitWmmaSSA(OpBuilder &builder, Location loc, VectorType Value res; if constexpr (Variant == WmmaVariant::ModsAllReuse) { // Float path: no sign/clamp operands. - res = WmmaOp::create(builder, loc, accTy, - /*signA=*/false, a, /*signB=*/false, b, - /*modC=*/(uint16_t)0, c) + res = WmmaOp::create(builder, loc, accTy, a, b, + /*modC=*/ROCDL::WMMACModifier::none, c) .getResult(); } else if constexpr (Variant == WmmaVariant::ModsC) { // fp8 path: no sign/clamp operands. res = WmmaOp::create(builder, loc, accTy, a, b, - /*modC=*/(uint16_t)0, c, + /*modC=*/ROCDL::WMMACModifier::none, c, /*reuseA=*/false, /*reuseB=*/false) .getResult(); } else if constexpr (Variant == WmmaVariant::ModsABClamp) { diff --git a/lib/Dialect/FlyROCDL/GFX1250/MmaAtomScale.cpp b/lib/Dialect/FlyROCDL/GFX1250/MmaAtomScale.cpp index fa2bdce5b..b9c2ede56 100644 --- a/lib/Dialect/FlyROCDL/GFX1250/MmaAtomScale.cpp +++ b/lib/Dialect/FlyROCDL/GFX1250/MmaAtomScale.cpp @@ -244,15 +244,15 @@ FailureOr MmaOpGFX1250_WMMAScaleType::emitAtomCallSSA(OpBuilder &builder, // fp4-only form; no fmtA/fmtB operands. if (block16) return ROCDL::wmma_scale16_f32_32x16x128_f4::create( - builder, loc, accTy, a, b, /*modC=*/(uint16_t)getModC(), c, - /*scaleAType=*/(uint32_t)getOpselA(), /*fmtScaleA=*/(uint32_t)0, scaleA, - /*scaleBType=*/(uint32_t)getOpselB(), /*fmtScaleB=*/(uint32_t)0, scaleB, + builder, loc, accTy, a, b, /*modC=*/static_cast((uint16_t)getModC()), c, + /*scaleAType=*/static_cast((uint32_t)getOpselA()), /*fmtScaleA=*/ROCDL::WMMAMatrixScaleFormat::e8, scaleA, + /*scaleBType=*/static_cast((uint32_t)getOpselB()), /*fmtScaleB=*/ROCDL::WMMAMatrixScaleFormat::e8, scaleB, /*reuseA=*/getReuseA(), /*reuseB=*/getReuseB()) .getResult(); return ROCDL::wmma_scale_f32_32x16x128_f4::create( - builder, loc, accTy, a, b, /*modC=*/(uint16_t)getModC(), c, - /*scaleAType=*/(uint32_t)getOpselA(), /*fmtScaleA=*/(uint32_t)0, scaleA, - /*scaleBType=*/(uint32_t)getOpselB(), /*fmtScaleB=*/(uint32_t)0, scaleB, + builder, loc, accTy, a, b, /*modC=*/static_cast((uint16_t)getModC()), c, + /*scaleAType=*/static_cast((uint32_t)getOpselA()), /*fmtScaleA=*/ROCDL::WMMAMatrixScaleFormat::e8, scaleA, + /*scaleBType=*/static_cast((uint32_t)getOpselB()), /*fmtScaleB=*/ROCDL::WMMAMatrixScaleFormat::e8, scaleB, /*reuseA=*/getReuseA(), /*reuseB=*/getReuseB()) .getResult(); } @@ -264,17 +264,17 @@ FailureOr MmaOpGFX1250_WMMAScaleType::emitAtomCallSSA(OpBuilder &builder, if (block16) return ROCDL::wmma_scale16_f32_16x16x128_f8f6f4::create( - builder, loc, accTy, /*fmtA=*/*aFmt, a, /*fmtB=*/*bFmt, b, - /*modC=*/(uint16_t)getModC(), c, - /*scaleAType=*/(uint32_t)getOpselA(), /*fmtScaleA=*/(uint32_t)0, scaleA, - /*scaleBType=*/(uint32_t)getOpselB(), /*fmtScaleB=*/(uint32_t)0, scaleB, + builder, loc, accTy, /*fmtA=*/static_cast(*aFmt), a, /*fmtB=*/static_cast(*bFmt), b, + /*modC=*/static_cast((uint16_t)getModC()), c, + /*scaleAType=*/static_cast((uint32_t)getOpselA()), /*fmtScaleA=*/ROCDL::WMMAMatrixScaleFormat::e8, scaleA, + /*scaleBType=*/static_cast((uint32_t)getOpselB()), /*fmtScaleB=*/ROCDL::WMMAMatrixScaleFormat::e8, scaleB, /*reuseA=*/getReuseA(), /*reuseB=*/getReuseB()) .getResult(); return ROCDL::wmma_scale_f32_16x16x128_f8f6f4::create( - builder, loc, accTy, /*fmtA=*/*aFmt, a, /*fmtB=*/*bFmt, b, - /*modC=*/(uint16_t)getModC(), c, - /*scaleAType=*/(uint32_t)getOpselA(), /*fmtScaleA=*/(uint32_t)0, scaleA, - /*scaleBType=*/(uint32_t)getOpselB(), /*fmtScaleB=*/(uint32_t)0, scaleB, + builder, loc, accTy, /*fmtA=*/static_cast(*aFmt), a, /*fmtB=*/static_cast(*bFmt), b, + /*modC=*/static_cast((uint16_t)getModC()), c, + /*scaleAType=*/static_cast((uint32_t)getOpselA()), /*fmtScaleA=*/ROCDL::WMMAMatrixScaleFormat::e8, scaleA, + /*scaleBType=*/static_cast((uint32_t)getOpselB()), /*fmtScaleB=*/ROCDL::WMMAMatrixScaleFormat::e8, scaleB, /*reuseA=*/getReuseA(), /*reuseB=*/getReuseB()) .getResult(); } diff --git a/lib/Dialect/FlyROCDL/Ops.cpp b/lib/Dialect/FlyROCDL/Ops.cpp index a7440b4bc..c33db625d 100644 --- a/lib/Dialect/FlyROCDL/Ops.cpp +++ b/lib/Dialect/FlyROCDL/Ops.cpp @@ -15,7 +15,7 @@ using namespace mlir::fly_rocdl; LogicalResult GetBufferRsrcOp::inferReturnTypes(MLIRContext *context, std::optional location, ValueRange operands, DictionaryAttr attributes, - OpaqueProperties properties, RegionRange regions, + PropertyRef properties, RegionRange regions, SmallVectorImpl &inferredReturnTypes) { auto ptrTy = dyn_cast(operands[0].getType()); if (!ptrTy) diff --git a/python/flydsl/compiler/backends/base.py b/python/flydsl/compiler/backends/base.py index 0f6da1181..f09dc2fc4 100644 --- a/python/flydsl/compiler/backends/base.py +++ b/python/flydsl/compiler/backends/base.py @@ -73,6 +73,9 @@ def lower_compile_hints(self, module, *, compile_hints: dict) -> None: """Optionally materialize backend-specific hints before the pipeline.""" return None + def finalize_llvm_compile_hints(self, module, *, compile_hints: dict) -> None: + return None + def external_binary_pipeline_fragments(self, *, compile_hints: dict) -> Tuple[List[str], str]: """Split the pipeline for external device binary code generation. diff --git a/python/flydsl/compiler/backends/rocm.py b/python/flydsl/compiler/backends/rocm.py index 9d6113af2..9e4518faa 100644 --- a/python/flydsl/compiler/backends/rocm.py +++ b/python/flydsl/compiler/backends/rocm.py @@ -120,6 +120,14 @@ def lower_compile_hints(self, module, *, compile_hints: dict) -> None: del func_op.attributes["rocdl.waves_per_eu"] _set_passthrough(func_op, "amdgpu-waves-per-eu", f"{waves_per_eu},{waves_per_eu}") + def finalize_llvm_compile_hints(self, module, *, compile_hints: dict) -> None: + waves_per_eu = compile_hints.get("waves_per_eu") + if not waves_per_eu: + return + with module.context: + for func_op in _iter_llvm_kernel_funcs(module): + _set_passthrough(func_op, "amdgpu-waves-per-eu", f"{waves_per_eu},{waves_per_eu}") + def gpu_module_targets(self) -> List[str]: chip = self.target.arch return [f'#rocdl.target'] @@ -148,7 +156,18 @@ def _iter_gpu_kernel_funcs(module): if top.operation.name != "gpu.module": continue for op in top.regions[0].blocks[0].operations: - if op.operation.name == "gpu.func" and "gpu.kernel" in op.attributes: + if op.operation.name == "gpu.func" and "kernel" in op.attributes: + yield op + + +def _iter_llvm_kernel_funcs(module): + for top in module.body.operations: + if top.operation.name != "gpu.module": + continue + for op in top.regions[0].blocks[0].operations: + if op.operation.name == "llvm.func" and ( + "rocdl.kernel" in op.attributes or "gpu.kernel" in op.attributes + ): yield op diff --git a/python/flydsl/compiler/jit_function.py b/python/flydsl/compiler/jit_function.py index a213d9780..24e21ee10 100644 --- a/python/flydsl/compiler/jit_function.py +++ b/python/flydsl/compiler/jit_function.py @@ -845,6 +845,7 @@ def compile( dump_fragments = pre_binary_fragments if external_binary else fragments for idx, frag in enumerate(dump_fragments): if frag.strip().startswith("gpu-module-to-binary"): + backend.finalize_llvm_compile_hints(module, compile_hints=compile_hints) llir = _extract_llvm_ir(module) stage_num = stage_num_base + idx @@ -868,6 +869,7 @@ def compile( if external_binary: from .external_llvm import run_external_binary_codegen + backend.finalize_llvm_compile_hints(module, compile_hints=compile_hints) llir = _extract_llvm_ir(module) stage_name = f"{next_stage:02d}_external_binary" run_external_binary_codegen( @@ -916,6 +918,7 @@ def compile( verifier=env.debug.enable_verifier, print_after_all=env.debug.print_after_all, ) + backend.finalize_llvm_compile_hints(module, compile_hints=compile_hints) if env.debug.dump_asm: raise RuntimeError( @@ -931,12 +934,27 @@ def compile( ) module.operation.verify() else: - _run_pipeline( - module, - fragments, - verifier=env.debug.enable_verifier, - print_after_all=env.debug.print_after_all, - ) + if compile_hints.get("waves_per_eu"): + _run_pipeline( + module, + fragments[:-1], + verifier=env.debug.enable_verifier, + print_after_all=env.debug.print_after_all, + ) + backend.finalize_llvm_compile_hints(module, compile_hints=compile_hints) + _run_pipeline( + module, + fragments[-1:], + verifier=env.debug.enable_verifier, + print_after_all=env.debug.print_after_all, + ) + else: + _run_pipeline( + module, + fragments, + verifier=env.debug.enable_verifier, + print_after_all=env.debug.print_after_all, + ) return module diff --git a/python/flydsl/expr/buffer_ops.py b/python/flydsl/expr/buffer_ops.py index 6a69ec3a9..d817cc2eb 100644 --- a/python/flydsl/expr/buffer_ops.py +++ b/python/flydsl/expr/buffer_ops.py @@ -555,11 +555,11 @@ def buffer_load( if not isinstance(soffset.type, ir.IntegerType) or soffset.type.width != 32: op = std_arith.IndexCastOp(T.i32(), soffset) soffset = _unwrap_value(op.result) - aux_flags = _create_i32_constant(cache_modifier) + aux_flags = ir.IntegerAttr.get(ir.IntegerType.get_signless(32), cache_modifier) # Emit buffer load load_op = rocdl.RawPtrBufferLoadOp( - result_type, rsrc, offset, soffset, aux_flags # soffset (scalar byte offset) # aux (cache modifiers) + result_type, rsrc, offset, soffset, aux=aux_flags # soffset (scalar byte offset) # aux (cache modifiers) ) return load_op.result @@ -642,9 +642,9 @@ def buffer_store( if not isinstance(soffset.type, ir.IntegerType) or soffset.type.width != 32: op = std_arith.IndexCastOp(T.i32(), soffset) soffset = _unwrap_value(op.result) - aux_flags = _create_i32_constant(cache_modifier) + aux_flags = ir.IntegerAttr.get(ir.IntegerType.get_signless(32), cache_modifier) # Emit buffer store rocdl.RawPtrBufferStoreOp( - data, rsrc, offset, soffset, aux_flags # soffset (scalar byte offset) # aux (cache modifiers) + data, rsrc, offset, soffset, aux=aux_flags # soffset (scalar byte offset) # aux (cache modifiers) ) diff --git a/python/flydsl/expr/rocdl/__init__.py b/python/flydsl/expr/rocdl/__init__.py index 4b18f15cc..120e8f492 100644 --- a/python/flydsl/expr/rocdl/__init__.py +++ b/python/flydsl/expr/rocdl/__init__.py @@ -20,6 +20,14 @@ from . import cdna5 as cdna5 from .enum import SyncScope as SyncScope +from ._rocdl_enums import ( + MFMAPermB as MFMAPermB, + MatrixFormat as MatrixFormat, + WMMACModifier as WMMACModifier, + WMMAMatrixScale as WMMAMatrixScale, + WMMAMatrixScaleFormat as WMMAMatrixScaleFormat, +) + # Keep references to ODS-generated builders so we can wrap them without losing access. _ods_wmma_scale_f32_16x16x128_f8f6f4 = globals().get("wmma_scale_f32_16x16x128_f8f6f4", None) _ods_wmma_scale_f32_32x16x128_f4 = globals().get("wmma_scale_f32_32x16x128_f4", None) @@ -247,13 +255,13 @@ def wmma_scale_f32_16x16x128_f8f6f4( c_v, sA, sB, - fmtA=fmtA, - fmtB=fmtB, - modC=modC, - scaleAType=scaleAType, - fmtScaleA=fmtScaleA, - scaleBType=scaleBType, - fmtScaleB=fmtScaleB, + fmtA=MatrixFormat(fmtA), + fmtB=MatrixFormat(fmtB), + modC=WMMACModifier(modC), + scaleAType=WMMAMatrixScale(scaleAType), + fmtScaleA=WMMAMatrixScaleFormat(fmtScaleA), + scaleBType=WMMAMatrixScale(scaleBType), + fmtScaleB=WMMAMatrixScaleFormat(fmtScaleB), reuseA=reuseA, reuseB=reuseB, ).result @@ -299,11 +307,11 @@ def wmma_scale_f32_32x16x128_f4( c_v, sA, sB, - modC=modC, - scaleAType=scaleAType, - fmtScaleA=fmtScaleA, - scaleBType=scaleBType, - fmtScaleB=fmtScaleB, + modC=WMMACModifier(modC), + scaleAType=WMMAMatrixScale(scaleAType), + fmtScaleA=WMMAMatrixScaleFormat(fmtScaleA), + scaleBType=WMMAMatrixScale(scaleBType), + fmtScaleB=WMMAMatrixScaleFormat(fmtScaleB), reuseA=reuseA, reuseB=reuseB, ).result @@ -323,7 +331,9 @@ def wmma_f32_16x16x128_fp8_fp8(result_type, a, b, c, *, modC=0, reuseA=False, re a_v = _unwrap_mfma_operand(a) b_v = _unwrap_mfma_operand(b) c_v = _unwrap_mfma_operand(c) - return _ods_wmma_f32_16x16x128_fp8_fp8(result_type, a_v, b_v, c_v, modC=modC, reuseA=reuseA, reuseB=reuseB).result + return _ods_wmma_f32_16x16x128_fp8_fp8( + result_type, a_v, b_v, c_v, modC=WMMACModifier(modC), reuseA=reuseA, reuseB=reuseB + ).result @dsl_loc_tracing @@ -480,18 +490,36 @@ def _to_ir(v): return v +def _as_cache_policy_attr(aux): + from ..._mlir import ir as _ir + + if isinstance(aux, _ir.Attribute): + return aux + if isinstance(aux, int): + return _ir.IntegerAttr.get(_ir.IntegerType.get_signless(32), aux) + v = _to_ir(aux) + try: # legacy: a constant i32 Value (e.g. arith.constant 0) -> reuse its attr + return v.owner.attributes["value"] + except Exception: + return _ir.IntegerAttr.get(_ir.IntegerType.get_signless(32), 0) + + @dsl_loc_tracing def raw_ptr_buffer_atomic_fadd(vdata, rsrc, offset, soffset, aux, **kw): from ..._mlir.dialects.rocdl import raw_ptr_buffer_atomic_fadd as _op - return _op(_to_ir(vdata), _to_ir(rsrc), _to_ir(offset), _to_ir(soffset), _to_ir(aux), **kw) + vd = _to_ir(vdata) + return _op(vd.type, vd, _to_ir(rsrc), _to_ir(offset), _to_ir(soffset), + aux=_as_cache_policy_attr(aux), **kw) @dsl_loc_tracing def raw_ptr_buffer_atomic_fmax(vdata, rsrc, offset, soffset, aux, **kw): from ..._mlir.dialects.rocdl import raw_ptr_buffer_atomic_fmax as _op - return _op(_to_ir(vdata), _to_ir(rsrc), _to_ir(offset), _to_ir(soffset), _to_ir(aux), **kw) + vd = _to_ir(vdata) + return _op(vd.type, vd, _to_ir(rsrc), _to_ir(offset), _to_ir(soffset), + aux=_as_cache_policy_attr(aux), **kw) @dsl_loc_tracing diff --git a/python/flydsl/expr/rocdl/_rocdl_enums.py b/python/flydsl/expr/rocdl/_rocdl_enums.py new file mode 100644 index 000000000..986ba75fe --- /dev/null +++ b/python/flydsl/expr/rocdl/_rocdl_enums.py @@ -0,0 +1,456 @@ + +# Autogenerated by mlir-tblgen; don't manually edit. + +from enum import IntEnum, auto, IntFlag +from flydsl._mlir.dialects._ods_common import _cext as _ods_cext +from flydsl._mlir.ir import register_attribute_builder +_ods_ir = _ods_cext.ir + +class BufferOOBMode(IntEnum): + """ROCDL buffer out-of-bounds mode""" + + Any = 0 + Relaxed = 1 + Strict = 2 + + def __str__(self): + if self is BufferOOBMode.Any: + return "any" + if self is BufferOOBMode.Relaxed: + return "relaxed" + if self is BufferOOBMode.Strict: + return "strict" + raise ValueError("Unknown BufferOOBMode enum entry.") + + + +class Gfx12AtomicCachePolicy(IntFlag): + """gfx12 atomic cache policy bits""" + + none = 0 + keep_return = 1 + nt = 2 + cascade = 4 + scope_se = 8 + scope_dev = 16 + scope_sys = 24 + nv = 32 + volatile_op = 2147483648 + + def __iter__(self): + return iter([case for case in type(self) if (self & case) is case and self is not case]) + def __len__(self): + return bin(self).count("1") + + def __str__(self): + if len(self) > 1: + return "|".join(map(str, self)) + if self is Gfx12AtomicCachePolicy.none: + return "none" + if self is Gfx12AtomicCachePolicy.keep_return: + return "return" + if self is Gfx12AtomicCachePolicy.nt: + return "nt" + if self is Gfx12AtomicCachePolicy.cascade: + return "cascade" + if self is Gfx12AtomicCachePolicy.scope_se: + return "scope_se" + if self is Gfx12AtomicCachePolicy.scope_dev: + return "scope_dev" + if self is Gfx12AtomicCachePolicy.scope_sys: + return "scope_sys" + if self is Gfx12AtomicCachePolicy.nv: + return "nv" + if self is Gfx12AtomicCachePolicy.volatile_op: + return "volatile" + raise ValueError("Unknown Gfx12AtomicCachePolicy enum entry.") + + + +class Gfx12CachePolicy(IntFlag): + """gfx12 cache policy bits""" + + none = 0 + nt = 1 + ht = 2 + lu = 3 + nt_rt = 4 + rt_nt = 5 + nt_ht = 6 + nt_wb = 7 + scope_se = 8 + scope_dev = 16 + scope_sys = 24 + nv = 32 + swz = 64 + scal = 2048 + volatile_op = 2147483648 + + def __iter__(self): + return iter([case for case in type(self) if (self & case) is case and self is not case]) + def __len__(self): + return bin(self).count("1") + + def __str__(self): + if len(self) > 1: + return "|".join(map(str, self)) + if self is Gfx12CachePolicy.none: + return "none" + if self is Gfx12CachePolicy.nt: + return "nt" + if self is Gfx12CachePolicy.ht: + return "ht" + if self is Gfx12CachePolicy.lu: + return "lu" + if self is Gfx12CachePolicy.nt_rt: + return "nt_rt" + if self is Gfx12CachePolicy.rt_nt: + return "rt_nt" + if self is Gfx12CachePolicy.nt_ht: + return "nt_ht" + if self is Gfx12CachePolicy.nt_wb: + return "nt_wb" + if self is Gfx12CachePolicy.scope_se: + return "scope_se" + if self is Gfx12CachePolicy.scope_dev: + return "scope_dev" + if self is Gfx12CachePolicy.scope_sys: + return "scope_sys" + if self is Gfx12CachePolicy.nv: + return "nv" + if self is Gfx12CachePolicy.swz: + return "swz" + if self is Gfx12CachePolicy.scal: + return "scal" + if self is Gfx12CachePolicy.volatile_op: + return "volatile" + raise ValueError("Unknown Gfx12CachePolicy enum entry.") + + + +class Gfx942CachePolicy(IntFlag): + """gfx942 buffer cache policy bits""" + + none = 0 + sc0 = 1 + nt = 2 + swz = 8 + sc1 = 16 + volatile_op = 2147483648 + + def __iter__(self): + return iter([case for case in type(self) if (self & case) is case and self is not case]) + def __len__(self): + return bin(self).count("1") + + def __str__(self): + if len(self) > 1: + return "|".join(map(str, self)) + if self is Gfx942CachePolicy.none: + return "none" + if self is Gfx942CachePolicy.sc0: + return "sc0" + if self is Gfx942CachePolicy.nt: + return "nt" + if self is Gfx942CachePolicy.swz: + return "swz" + if self is Gfx942CachePolicy.sc1: + return "sc1" + if self is Gfx942CachePolicy.volatile_op: + return "volatile" + raise ValueError("Unknown Gfx942CachePolicy enum entry.") + + + +class MFMANegModifier(IntFlag): + """negation modifier bitfield for gfx94x double-precision MFMA""" + + none = 0 + neg_a = 1 + neg_b = 2 + neg_c = 4 + + def __iter__(self): + return iter([case for case in type(self) if (self & case) is case and self is not case]) + def __len__(self): + return bin(self).count("1") + + def __str__(self): + if len(self) > 1: + return "|".join(map(str, self)) + if self is MFMANegModifier.none: + return "none" + if self is MFMANegModifier.neg_a: + return "neg_a" + if self is MFMANegModifier.neg_b: + return "neg_b" + if self is MFMANegModifier.neg_c: + return "neg_c" + raise ValueError("Unknown MFMANegModifier enum entry.") + + + +class MFMAPermB(IntEnum): + """permutations of the lanes storing B in an MFMA""" + + none = 0 + bcast_first_32 = 1 + bcast_second_32 = 2 + rotate_16_right = 3 + bcast_first_16 = 4 + bcast_second_16 = 5 + bcast_third_16 = 6 + bcast_fourth_16 = 7 + + def __str__(self): + if self is MFMAPermB.none: + return "none" + if self is MFMAPermB.bcast_first_32: + return "bcast_first_32" + if self is MFMAPermB.bcast_second_32: + return "bcast_second_32" + if self is MFMAPermB.rotate_16_right: + return "rotate_16_right" + if self is MFMAPermB.bcast_first_16: + return "bcast_first_16" + if self is MFMAPermB.bcast_second_16: + return "bcast_second_16" + if self is MFMAPermB.bcast_third_16: + return "bcast_third_16" + if self is MFMAPermB.bcast_fourth_16: + return "bcast_fourth_16" + raise ValueError("Unknown MFMAPermB enum entry.") + + + +class MatrixFormat(IntEnum): + """matrix operand formats selected by scaled MFMA/WMMA format fields""" + + fp8_e4m3 = 0 + fp8_e5m2 = 1 + fp6_e2m3 = 2 + fp6_e3m2 = 3 + fp4_e2m1 = 4 + + def __str__(self): + if self is MatrixFormat.fp8_e4m3: + return "fp8_e4m3" + if self is MatrixFormat.fp8_e5m2: + return "fp8_e5m2" + if self is MatrixFormat.fp6_e2m3: + return "fp6_e2m3" + if self is MatrixFormat.fp6_e3m2: + return "fp6_e3m2" + if self is MatrixFormat.fp4_e2m1: + return "fp4_e2m1" + raise ValueError("Unknown MatrixFormat enum entry.") + + + +class PreGfx12CachePolicy(IntFlag): + """pre-gfx12 buffer cache policy bits""" + + none = 0 + glc = 1 + slc = 2 + dlc = 4 + swz = 8 + scc = 16 + all = 23 + volatile_op = 2147483648 + + def __iter__(self): + return iter([case for case in type(self) if (self & case) is case and self is not case]) + def __len__(self): + return bin(self).count("1") + + def __str__(self): + if len(self) > 1: + return "|".join(map(str, self)) + if self is PreGfx12CachePolicy.none: + return "none" + if self is PreGfx12CachePolicy.glc: + return "glc" + if self is PreGfx12CachePolicy.slc: + return "slc" + if self is PreGfx12CachePolicy.dlc: + return "dlc" + if self is PreGfx12CachePolicy.swz: + return "swz" + if self is PreGfx12CachePolicy.scc: + return "scc" + if self is PreGfx12CachePolicy.all: + return "all" + if self is PreGfx12CachePolicy.volatile_op: + return "volatile" + raise ValueError("Unknown PreGfx12CachePolicy enum entry.") + + + +class SchedGroupMask(IntFlag): + """instruction type mask for scheduling barriers""" + + none = 0 + non_mem_non_sideeffect = 1 + valu = 2 + salu = 4 + mfma_wmma = 8 + all_vmem = 16 + vmem_read = 32 + vmem_write = 64 + all_ds = 128 + ds_read = 256 + ds_write = 512 + transcendental = 1024 + ldsdma = 2048 + all = 4095 + + def __iter__(self): + return iter([case for case in type(self) if (self & case) is case and self is not case]) + def __len__(self): + return bin(self).count("1") + + def __str__(self): + if len(self) > 1: + return "|".join(map(str, self)) + if self is SchedGroupMask.none: + return "none" + if self is SchedGroupMask.non_mem_non_sideeffect: + return "non_mem_non_sideeffect" + if self is SchedGroupMask.valu: + return "valu" + if self is SchedGroupMask.salu: + return "salu" + if self is SchedGroupMask.mfma_wmma: + return "mfma_wmma" + if self is SchedGroupMask.all_vmem: + return "all_vmem" + if self is SchedGroupMask.vmem_read: + return "vmem_read" + if self is SchedGroupMask.vmem_write: + return "vmem_write" + if self is SchedGroupMask.all_ds: + return "all_ds" + if self is SchedGroupMask.ds_read: + return "ds_read" + if self is SchedGroupMask.ds_write: + return "ds_write" + if self is SchedGroupMask.transcendental: + return "transcendental" + if self is SchedGroupMask.ldsdma: + return "ldsdma" + if self is SchedGroupMask.all: + return "all" + raise ValueError("Unknown SchedGroupMask enum entry.") + + + +class WMMACModifier(IntFlag): + """WMMA C operand modifiers""" + + none = 0 + neg = 1 + abs = 2 + neg_abs = 3 + + def __iter__(self): + return iter([case for case in type(self) if (self & case) is case and self is not case]) + def __len__(self): + return bin(self).count("1") + + def __str__(self): + if len(self) > 1: + return "|".join(map(str, self)) + if self is WMMACModifier.none: + return "none" + if self is WMMACModifier.neg: + return "neg" + if self is WMMACModifier.abs: + return "abs" + if self is WMMACModifier.neg_abs: + return "neg_abs" + raise ValueError("Unknown WMMACModifier enum entry.") + + + +class WMMAMatrixScale(IntEnum): + """matrix scale row selector""" + + row0 = 0 + row1 = 1 + + def __str__(self): + if self is WMMAMatrixScale.row0: + return "row0" + if self is WMMAMatrixScale.row1: + return "row1" + raise ValueError("Unknown WMMAMatrixScale enum entry.") + + + +class WMMAMatrixScaleFormat(IntEnum): + """matrix scale exponent formats""" + + e8 = 0 + e5m3 = 1 + e4m3 = 2 + + def __str__(self): + if self is WMMAMatrixScaleFormat.e8: + return "e8" + if self is WMMAMatrixScaleFormat.e5m3: + return "e5m3" + if self is WMMAMatrixScaleFormat.e4m3: + return "e4m3" + raise ValueError("Unknown WMMAMatrixScaleFormat enum entry.") + + + +@register_attribute_builder("rocdl.ROCDL_BufferOOBModeAttr") +def _rocdl_bufferoobmodeattr(x, context): + return _ods_ir.Attribute.parse(f'#rocdl.buffer_oob_mode<{str(x)}>', context=context) + +@register_attribute_builder("rocdl.ROCDL_Gfx12AtomicCachePolicyAttr") +def _rocdl_gfx12atomiccachepolicyattr(x, context): + return _ods_ir.Attribute.parse(f'#rocdl', context=context) + +@register_attribute_builder("rocdl.ROCDL_Gfx12CachePolicyAttr") +def _rocdl_gfx12cachepolicyattr(x, context): + return _ods_ir.Attribute.parse(f'#rocdl', context=context) + +@register_attribute_builder("rocdl.ROCDL_Gfx942CachePolicyAttr") +def _rocdl_gfx942cachepolicyattr(x, context): + return _ods_ir.Attribute.parse(f'#rocdl', context=context) + +@register_attribute_builder("rocdl.ROCDL_MFMANegModifierAttr") +def _rocdl_mfmanegmodifierattr(x, context): + return _ods_ir.Attribute.parse(f'#rocdl', context=context) + +@register_attribute_builder("rocdl.ROCDL_MFMAPermBAttr") +def _rocdl_mfmapermbattr(x, context): + return _ods_ir.Attribute.parse(f'#rocdl', context=context) + +@register_attribute_builder("rocdl.ROCDL_MatrixFormatAttr") +def _rocdl_matrixformatattr(x, context): + return _ods_ir.Attribute.parse(f'#rocdl', context=context) + +@register_attribute_builder("rocdl.ROCDL_PreGfx12CachePolicyAttr") +def _rocdl_pregfx12cachepolicyattr(x, context): + return _ods_ir.Attribute.parse(f'#rocdl', context=context) + +@register_attribute_builder("rocdl.ROCDL_SchedGroupMaskAttr") +def _rocdl_schedgroupmaskattr(x, context): + return _ods_ir.Attribute.parse(f'#rocdl', context=context) + +@register_attribute_builder("rocdl.ROCDL_WMMACModifierAttr") +def _rocdl_wmmacmodifierattr(x, context): + return _ods_ir.Attribute.parse(f'#rocdl', context=context) + +@register_attribute_builder("rocdl.ROCDL_WMMAMatrixScaleAttr") +def _rocdl_wmmamatrixscaleattr(x, context): + return _ods_ir.Attribute.parse(f'#rocdl', context=context) + +@register_attribute_builder("rocdl.ROCDL_WMMAMatrixScaleFormatAttr") +def _rocdl_wmmamatrixscaleformatattr(x, context): + return _ods_ir.Attribute.parse(f'#rocdl', context=context) + diff --git a/python/flydsl/expr/rocdl/tdm_ops.py b/python/flydsl/expr/rocdl/tdm_ops.py index a31006418..ec2fe7422 100644 --- a/python/flydsl/expr/rocdl/tdm_ops.py +++ b/python/flydsl/expr/rocdl/tdm_ops.py @@ -194,6 +194,10 @@ def _i32_const(v: int) -> ir.Value: return _unwrap(std_arith.ConstantOp(i32, ir.IntegerAttr.get(i32, v)).result) +def _cache_policy_attr(cache_policy: int): + return ir.IntegerAttr.get(ir.IntegerType.get_signless(32), cache_policy) + + # --------------------------------------------------------------------------- # Public API # --------------------------------------------------------------------------- @@ -764,7 +768,7 @@ def tensor_load_gather( _raw(desc.dgroup2), _raw(desc.dgroup3), dg4, - cache_policy, + cache_policy=_cache_policy_attr(cache_policy), ) @@ -789,7 +793,7 @@ def tensor_store_gather( _raw(desc.dgroup2), _raw(desc.dgroup3), dg4, - cache_policy, + cache_policy=_cache_policy_attr(cache_policy), ) @@ -1103,7 +1107,7 @@ def tensor_load_2d( dg2 = _raw(_zero_dgroup_v4i32()) dg3 = _raw(_zero_dgroup_v4i32()) dg4 = _raw(_zero_dgroup_v8i32()) - rocdl.tensor_load_to_lds(_raw(desc.dgroup0), _raw(desc.dgroup1), dg2, dg3, dg4, cache_policy) + rocdl.tensor_load_to_lds(_raw(desc.dgroup0), _raw(desc.dgroup1), dg2, dg3, dg4, cache_policy=_cache_policy_attr(cache_policy)) @dsl_loc_tracing @@ -1123,7 +1127,7 @@ def tensor_store_2d( dg2 = _raw(_zero_dgroup_v4i32()) dg3 = _raw(_zero_dgroup_v4i32()) dg4 = _raw(_zero_dgroup_v8i32()) - rocdl.tensor_store_from_lds(_raw(desc.dgroup0), _raw(desc.dgroup1), dg2, dg3, dg4, cache_policy) + rocdl.tensor_store_from_lds(_raw(desc.dgroup0), _raw(desc.dgroup1), dg2, dg3, dg4, cache_policy=_cache_policy_attr(cache_policy)) @dsl_loc_tracing diff --git a/tests/mlir/Conversion/mma_atom_stateful.mlir b/tests/mlir/Conversion/mma_atom_stateful.mlir index 14c64c2fc..ff57ec3af 100644 --- a/tests/mlir/Conversion/mma_atom_stateful.mlir +++ b/tests/mlir/Conversion/mma_atom_stateful.mlir @@ -79,7 +79,7 @@ func.func @test_mma_scale_atom_call_16x16x128( // CHECK-DAG: %[[C_VAL:.*]] = llvm.load %[[C]] : !llvm.ptr<5> -> vector<4xf32> // CHECK-DAG: %[[SA_VAL:.*]] = llvm.extractvalue %[[A2]][0] // CHECK-DAG: %[[SB_VAL:.*]] = llvm.extractvalue %[[A2]][1] - // CHECK: %[[RES:.*]] = rocdl.mfma.scale.f32.16x16x128.f8f6f4 %[[A_VAL]], %[[B_VAL]], %[[C_VAL]], 0, 0, 0, %[[SA_VAL]], 0, %[[SB_VAL]] + // CHECK: %[[RES:.*]] = rocdl.mfma.scale.f32.16x16x128.f8f6f4 %[[A_VAL]], %[[B_VAL]], %[[C_VAL]], fp8_e4m3, fp8_e4m3, 0, %[[SA_VAL]], 0, %[[SB_VAL]] // CHECK: llvm.store %[[RES]], %[[D]] : vector<4xf32>, !llvm.ptr<5> fly.mma_atom_call(%atom_ab, %d, %a, %b, %c) : (!fly.mma_atom f32, opselA = 0, opselB = 0>>, !fly.memref, !fly.memref, !fly.memref, !fly.memref) -> () return @@ -100,7 +100,7 @@ func.func @test_mma_scale_atom_call_32x32x64_mixed( %b = fly.memref.alloca(%lay_ab) : (!fly.layout<32:1>) -> !fly.memref %c = fly.memref.alloca(%lay_cd) : (!fly.layout<16:1>) -> !fly.memref - // CHECK: rocdl.mfma.scale.f32.32x32x64.f8f6f4 %{{.*}}, %{{.*}}, %{{.*}}, 4, 1, 0, %{{.*}}, 0, %{{.*}} + // CHECK: rocdl.mfma.scale.f32.32x32x64.f8f6f4 %{{.*}}, %{{.*}}, %{{.*}}, fp4_e2m1, fp8_e5m2, 0, %{{.*}}, 0, %{{.*}} fly.mma_atom_call(%atom, %d, %a, %b, %c) : (!fly.mma_atom f32, opselA = 0, opselB = 0>>, !fly.memref, !fly.memref, !fly.memref, !fly.memref) -> () return } @@ -122,7 +122,7 @@ func.func @test_mma_scale_atom_call_with_opsel( %b = fly.memref.alloca(%lay_ab) : (!fly.layout<32:1>) -> !fly.memref %c = fly.memref.alloca(%lay_cd) : (!fly.layout<4:1>) -> !fly.memref - // CHECK: rocdl.mfma.scale.f32.16x16x128.f8f6f4 %{{.*}}, %{{.*}}, %{{.*}}, 0, 0, 2, %{{.*}}, 3, %{{.*}} + // CHECK: rocdl.mfma.scale.f32.16x16x128.f8f6f4 %{{.*}}, %{{.*}}, %{{.*}}, fp8_e4m3, fp8_e4m3, 2, %{{.*}}, 3, %{{.*}} fly.mma_atom_call(%atom, %d, %a, %b, %c) : (!fly.mma_atom f32, opselA = 2, opselB = 3>>, !fly.memref, !fly.memref, !fly.memref, !fly.memref) -> () return } diff --git a/tests/mlir/Conversion/mma_scale_gfx1250.mlir b/tests/mlir/Conversion/mma_scale_gfx1250.mlir index 76565fcda..e47654b14 100644 --- a/tests/mlir/Conversion/mma_scale_gfx1250.mlir +++ b/tests/mlir/Conversion/mma_scale_gfx1250.mlir @@ -66,7 +66,7 @@ func.func @test_wmma_scale_call_fp8( // CHECK-DAG: %[[C_VAL:.*]] = llvm.load %{{.*}} : !llvm.ptr<5> -> vector<8xf32> // CHECK-DAG: %[[SA_VAL:.*]] = llvm.extractvalue %[[A2]][0] // CHECK-DAG: %[[SB_VAL:.*]] = llvm.extractvalue %[[A2]][1] - // CHECK: %[[RES:.*]] = rocdl.wmma.scale.f32.16x16x128.f8f6f4 %[[A_VAL]], %[[B_VAL]], %[[C_VAL]], %[[SA_VAL]], %[[SB_VAL]] : (vector<16xi32>, vector<16xi32>, vector<8xf32>, i32, i32) -> vector<8xf32> + // CHECK: %[[RES:.*]] = rocdl.wmma.scale.f32.16x16x128.f8f6f4 %[[A_VAL]], %[[B_VAL]], %[[C_VAL]], %[[SA_VAL]], %[[SB_VAL]], fmtA = fp8_e4m3, fmtB = fp8_e4m3, modC = none, scaleAType = row0, fmtScaleA = e8, scaleBType = row0, fmtScaleB = e8 : (vector<16xi32>, vector<16xi32>, vector<8xf32>, i32, i32) -> vector<8xf32> // CHECK: llvm.store %[[RES]], %{{.*}} : vector<8xf32>, !llvm.ptr<5> fly.mma_atom_call(%atom_ab, %d, %a, %b, %c) : (!fly.mma_atom f32, opselA = 0, opselB = 0, modC = 0, reuseA = false, reuseB = false, blockSize = 32>>, !fly.memref, !fly.memref, !fly.memref, !fly.memref) -> () return @@ -78,7 +78,7 @@ func.func @test_wmma_scale_call_fp8( // CHECK-LABEL: @test_wmma_scale_call_fp4_opsel func.func @test_wmma_scale_call_fp4_opsel( - %atom: !fly.mma_atom f32, opselA = 1, opselB = 2, modC = 0, reuseA = false, reuseB = false, blockSize = 32>>) { + %atom: !fly.mma_atom f32, opselA = 1, opselB = 0, modC = 0, reuseA = false, reuseB = false, blockSize = 32>>) { %lay_ab = fly.static : !fly.layout<32:1> %lay_cd = fly.static : !fly.layout<8:1> %d = fly.memref.alloca(%lay_cd) : (!fly.layout<8:1>) -> !fly.memref @@ -86,8 +86,8 @@ func.func @test_wmma_scale_call_fp4_opsel( %b = fly.memref.alloca(%lay_ab) : (!fly.layout<32:1>) -> !fly.memref %c = fly.memref.alloca(%lay_cd) : (!fly.layout<8:1>) -> !fly.memref - // CHECK: rocdl.wmma.scale.f32.16x16x128.f8f6f4 %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} {fmtA = 4 : i32, fmtB = 4 : i32, scaleAType = 1 : i32, scaleBType = 2 : i32} : (vector<8xi32>, vector<8xi32>, vector<8xf32>, i32, i32) -> vector<8xf32> - fly.mma_atom_call(%atom, %d, %a, %b, %c) : (!fly.mma_atom f32, opselA = 1, opselB = 2, modC = 0, reuseA = false, reuseB = false, blockSize = 32>>, !fly.memref, !fly.memref, !fly.memref, !fly.memref) -> () + // CHECK: rocdl.wmma.scale.f32.16x16x128.f8f6f4 %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, fmtA = fp4_e2m1, fmtB = fp4_e2m1, modC = none, scaleAType = row1, fmtScaleA = e8, scaleBType = row0, fmtScaleB = e8 : (vector<8xi32>, vector<8xi32>, vector<8xf32>, i32, i32) -> vector<8xf32> + fly.mma_atom_call(%atom, %d, %a, %b, %c) : (!fly.mma_atom f32, opselA = 1, opselB = 0, modC = 0, reuseA = false, reuseB = false, blockSize = 32>>, !fly.memref, !fly.memref, !fly.memref, !fly.memref) -> () return } @@ -108,7 +108,7 @@ func.func @test_wmma_scale_call_fp4_32x16( %b = fly.memref.alloca(%lay_b) : (!fly.layout<64:1>) -> !fly.memref %c = fly.memref.alloca(%lay_c) : (!fly.layout<16:1>) -> !fly.memref - // CHECK: rocdl.wmma.scale.f32.32x16x128.f4 %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} {scaleBType = 1 : i32} : (vector<16xi32>, vector<8xi32>, vector<16xf32>, i32, i32) -> vector<16xf32> + // CHECK: rocdl.wmma.scale.f32.32x16x128.f4 %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, modC = none, scaleAType = row0, fmtScaleA = e8, scaleBType = row1, fmtScaleB = e8 : (vector<16xi32>, vector<8xi32>, vector<16xf32>, i32, i32) -> vector<16xf32> fly.mma_atom_call(%atom, %d, %a, %b, %c) : (!fly.mma_atom f32, opselA = 0, opselB = 1, modC = 0, reuseA = false, reuseB = false, blockSize = 32>>, !fly.memref, !fly.memref, !fly.memref, !fly.memref) -> () return } @@ -129,7 +129,7 @@ func.func @test_wmma_scale_call_fp8_modc_reuse( %c = fly.memref.alloca(%lay_cd) : (!fly.layout<8:1>) -> !fly.memref // CHECK: rocdl.wmma.scale.f32.16x16x128.f8f6f4 - // CHECK-SAME: modC = 1 : i16 + // CHECK-SAME: modC = neg // CHECK-SAME: reuseA = true // CHECK-SAME: reuseB = true fly.mma_atom_call(%atom, %d, %a, %b, %c) : (!fly.mma_atom f32, opselA = 0, opselB = 0, modC = 1, reuseA = true, reuseB = true, blockSize = 32>>, !fly.memref, !fly.memref, !fly.memref, !fly.memref) -> () @@ -170,7 +170,7 @@ func.func @test_wmma_scale16_call_fp8( // CHECK-DAG: %[[SA_VAL:.*]] = llvm.extractvalue %[[A2]][0] // CHECK-DAG: %[[SB_VAL:.*]] = llvm.extractvalue %[[A2]][1] - // CHECK: rocdl.wmma.scale16.f32.16x16x128.f8f6f4 %{{.*}}, %{{.*}}, %{{.*}}, %[[SA_VAL]], %[[SB_VAL]] : (vector<16xi32>, vector<16xi32>, vector<8xf32>, i64, i64) -> vector<8xf32> + // CHECK: rocdl.wmma.scale16.f32.16x16x128.f8f6f4 %{{.*}}, %{{.*}}, %{{.*}}, %[[SA_VAL]], %[[SB_VAL]], fmtA = fp8_e4m3, fmtB = fp8_e4m3, modC = none, scaleAType = row0, fmtScaleA = e8, scaleBType = row0, fmtScaleB = e8 : (vector<16xi32>, vector<16xi32>, vector<8xf32>, i64, i64) -> vector<8xf32> fly.mma_atom_call(%atom_ab, %d, %a, %b, %c) : (!fly.mma_atom f32, opselA = 0, opselB = 0, modC = 0, reuseA = false, reuseB = false, blockSize = 16>>, !fly.memref, !fly.memref, !fly.memref, !fly.memref) -> () return } @@ -192,7 +192,7 @@ func.func @test_wmma_scale16_call_fp4_32x16( %b = fly.memref.alloca(%lay_b) : (!fly.layout<64:1>) -> !fly.memref %c = fly.memref.alloca(%lay_c) : (!fly.layout<16:1>) -> !fly.memref - // CHECK: rocdl.wmma.scale16.f32.32x16x128.f4 %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} {scaleBType = 1 : i32} : (vector<16xi32>, vector<8xi32>, vector<16xf32>, i64, i64) -> vector<16xf32> + // CHECK: rocdl.wmma.scale16.f32.32x16x128.f4 %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, modC = none, scaleAType = row0, fmtScaleA = e8, scaleBType = row1, fmtScaleB = e8 : (vector<16xi32>, vector<8xi32>, vector<16xf32>, i64, i64) -> vector<16xf32> fly.mma_atom_call(%atom, %d, %a, %b, %c) : (!fly.mma_atom f32, opselA = 0, opselB = 1, modC = 0, reuseA = false, reuseB = false, blockSize = 16>>, !fly.memref, !fly.memref, !fly.memref, !fly.memref) -> () return } diff --git a/tests/mlir/Conversion/tdm_gfx1250.mlir b/tests/mlir/Conversion/tdm_gfx1250.mlir index 92364d062..dda4ea523 100644 --- a/tests/mlir/Conversion/tdm_gfx1250.mlir +++ b/tests/mlir/Conversion/tdm_gfx1250.mlir @@ -44,7 +44,7 @@ func.func @test_tdm_load( // OOB clamp from the extent state fields. // CHECK: arith.subi // CHECK: arith.maxsi - // CHECK: rocdl.tensor.load.to.lds %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} cachepolicy 0 : vector<4xi32>, vector<8xi32> + // CHECK: rocdl.tensor.load.to.lds %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, 0 : vector<4xi32>, vector<8xi32> fly.copy_atom_call(%a3, %src, %dst) : (!fly.copy_atom, 0>, !fly.memref, !fly.memref) -> () return } @@ -58,7 +58,7 @@ func.func @test_tdm_store( %atom: !fly.copy_atom, 0>, %src: !fly.memref, %dst: !fly.memref) { - // CHECK: rocdl.tensor.store.from.lds %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} cachepolicy 0 : vector<4xi32>, vector<8xi32> + // CHECK: rocdl.tensor.store.from.lds %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, 0 : vector<4xi32>, vector<8xi32> fly.copy_atom_call(%atom, %src, %dst) : (!fly.copy_atom, 0>, !fly.memref, !fly.memref) -> () return } diff --git a/tests/unit/test_external_llvm_codegen.py b/tests/unit/test_external_llvm_codegen.py index f1460885f..ee6816095 100644 --- a/tests/unit/test_external_llvm_codegen.py +++ b/tests/unit/test_external_llvm_codegen.py @@ -144,6 +144,7 @@ def test_rocm_wpe_reaches_native_llvm_as_exact_constraint(): backend.lower_compile_hints(module, compile_hints={"waves_per_eu": 2}) pre_binary, _ = backend.external_binary_pipeline_fragments(compile_hints={}) PassManager.parse(f"builtin.module({','.join(pre_binary)})", ctx).run(module.operation) + backend.finalize_llvm_compile_hints(module, compile_hints={"waves_per_eu": 2}) llvm_ir = translate_module_to_llvmir(module.body.operations[0].operation) assert '"amdgpu-waves-per-eu"="2,2"' in llvm_ir