Skip to content

Enabling coexec llvm for Flydsl #900

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
omuhamma wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

Enabling coexec llvm for Flydsl #900

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions lib/Conversion/FlyToROCDL/FlyToROCDL.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -381,7 +381,7 @@ class PtrLoadOpLowering : public OpConversionPattern<PtrLoadOp> {
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 {
Expand Down Expand Up @@ -425,7 +425,7 @@ class PtrStoreOpLowering : public OpConversionPattern<PtrStoreOp> {
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();
Expand Down
2 changes: 1 addition & 1 deletion lib/Dialect/Fly/IR/FlyOps.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -191,7 +191,7 @@ Type applyOffsetOnTensorLike(LayoutBuilder<LayoutAttr> &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<mlir::Type> &inferredReturnTypes)

//===----------------------------------------------------------------------===//
Expand Down
14 changes: 0 additions & 14 deletions lib/Dialect/Fly/Transforms/PromoteRegMemToVectorSSA.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,20 +25,6 @@
using namespace mlir;
using namespace mlir::fly;

namespace llvm {

template <> struct DenseMapInfo<mlir::fly::MakePtrOp> : DenseMapInfo<mlir::Operation *> {
using Base = DenseMapInfo<mlir::Operation *>;

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
Expand Down
14 changes: 7 additions & 7 deletions lib/Dialect/FlyROCDL/CDNA3/CopyAtom.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -89,7 +89,7 @@ FailureOr<Value> 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<fly::MemRefType>(srcTyArg) : fly::MemRefType();
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -428,7 +428,7 @@ FailureOr<Value> 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();
Expand All @@ -437,21 +437,21 @@ FailureOr<Value> 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:
Expand Down
4 changes: 2 additions & 2 deletions lib/Dialect/FlyROCDL/CDNA3/MmaAtom.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -169,8 +169,8 @@ FailureOr<Value> 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(); \
}

Expand Down
16 changes: 8 additions & 8 deletions lib/Dialect/FlyROCDL/CDNA4/MmaAtom.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -237,20 +237,20 @@ FailureOr<Value> MmaOpCDNA4_MFMAScaleType::emitAtomCallSSA(OpBuilder &builder, L
Value scaleB = LLVM::ExtractValueOp::create(
builder, loc, atomVal, ArrayRef<int64_t>{*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<ROCDL::MatrixFormat>(*aTypeCode);
auto blgpFmt = static_cast<ROCDL::MatrixFormat>(*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();
}
Expand Down
2 changes: 1 addition & 1 deletion lib/Dialect/FlyROCDL/GFX1250/CopyAtom.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<uint32_t>(getCacheModifier());
auto cachePolicy = builder.getI32IntegerAttr(static_cast<uint32_t>(getCacheModifier()));
ArrayAttr noAliasScopes;
if (isLoad)
ROCDL::TensorLoadToLDSOp::create(builder, loc, dgroup0, dgroup1, dg2, dg3, dg4, cachePolicy,
Expand Down
7 changes: 3 additions & 4 deletions lib/Dialect/FlyROCDL/GFX1250/MmaAtom.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -239,14 +239,13 @@ static FailureOr<Value> 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) {
Expand Down
28 changes: 14 additions & 14 deletions lib/Dialect/FlyROCDL/GFX1250/MmaAtomScale.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -244,15 +244,15 @@ FailureOr<Value> 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<ROCDL::WMMACModifier>((uint16_t)getModC()), c,
/*scaleAType=*/static_cast<ROCDL::WMMAMatrixScale>((uint32_t)getOpselA()), /*fmtScaleA=*/ROCDL::WMMAMatrixScaleFormat::e8, scaleA,
/*scaleBType=*/static_cast<ROCDL::WMMAMatrixScale>((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<ROCDL::WMMACModifier>((uint16_t)getModC()), c,
/*scaleAType=*/static_cast<ROCDL::WMMAMatrixScale>((uint32_t)getOpselA()), /*fmtScaleA=*/ROCDL::WMMAMatrixScaleFormat::e8, scaleA,
/*scaleBType=*/static_cast<ROCDL::WMMAMatrixScale>((uint32_t)getOpselB()), /*fmtScaleB=*/ROCDL::WMMAMatrixScaleFormat::e8, scaleB,
/*reuseA=*/getReuseA(), /*reuseB=*/getReuseB())
.getResult();
}
Expand All @@ -264,17 +264,17 @@ FailureOr<Value> 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<ROCDL::MatrixFormat>(*aFmt), a, /*fmtB=*/static_cast<ROCDL::MatrixFormat>(*bFmt), b,
/*modC=*/static_cast<ROCDL::WMMACModifier>((uint16_t)getModC()), c,
/*scaleAType=*/static_cast<ROCDL::WMMAMatrixScale>((uint32_t)getOpselA()), /*fmtScaleA=*/ROCDL::WMMAMatrixScaleFormat::e8, scaleA,
/*scaleBType=*/static_cast<ROCDL::WMMAMatrixScale>((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<ROCDL::MatrixFormat>(*aFmt), a, /*fmtB=*/static_cast<ROCDL::MatrixFormat>(*bFmt), b,
/*modC=*/static_cast<ROCDL::WMMACModifier>((uint16_t)getModC()), c,
/*scaleAType=*/static_cast<ROCDL::WMMAMatrixScale>((uint32_t)getOpselA()), /*fmtScaleA=*/ROCDL::WMMAMatrixScaleFormat::e8, scaleA,
/*scaleBType=*/static_cast<ROCDL::WMMAMatrixScale>((uint32_t)getOpselB()), /*fmtScaleB=*/ROCDL::WMMAMatrixScaleFormat::e8, scaleB,
/*reuseA=*/getReuseA(), /*reuseB=*/getReuseB())
.getResult();
}
Expand Down
2 changes: 1 addition & 1 deletion lib/Dialect/FlyROCDL/Ops.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@ using namespace mlir::fly_rocdl;
LogicalResult GetBufferRsrcOp::inferReturnTypes(MLIRContext *context,
std::optional<Location> location,
ValueRange operands, DictionaryAttr attributes,
OpaqueProperties properties, RegionRange regions,
PropertyRef properties, RegionRange regions,
SmallVectorImpl<Type> &inferredReturnTypes) {
auto ptrTy = dyn_cast<PointerType>(operands[0].getType());
if (!ptrTy)
Expand Down
3 changes: 3 additions & 0 deletions python/flydsl/compiler/backends/base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand Down
21 changes: 20 additions & 1 deletion python/flydsl/compiler/backends/rocm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<chip = "{chip}">']
Expand Down Expand Up @@ -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


Expand Down
30 changes: 24 additions & 6 deletions python/flydsl/compiler/jit_function.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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(
Expand Down Expand Up @@ -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(
Expand All @@ -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

Expand Down
8 changes: 4 additions & 4 deletions python/flydsl/expr/buffer_ops.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)
)
Loading