Skip to content

Commit 8553efd

Browse files
matthias-springermatthias-springer
authored
[mlir][LLVM] Add OpBuilder & to lookupOrCreateFn functions (#136421)
These functions are called from lowering patterns. All IR modifications in a pattern must be performed through the provided rewriter, but these functions used to instantiate a new `OpBuilder`, bypassing the provided rewriter.
1 parent 71037ee commit 8553efd

File tree

8 files changed

+120
-98
lines changed

8 files changed

+120
-98
lines changed

‎mlir/include/mlir/Dialect/LLVMIR/FunctionCallUtils.h‎

Lines changed: 35 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -33,40 +33,53 @@ class LLVMFuncOp;
3333
/// implemented separately (e.g. as part of a support runtime library or as part
3434
/// of the libc).
3535
/// Failure if an unexpected version of function is found.
36-
FailureOr<LLVM::LLVMFuncOp> lookupOrCreatePrintI64Fn(Operation *moduleOp);
37-
FailureOr<LLVM::LLVMFuncOp> lookupOrCreatePrintU64Fn(Operation *moduleOp);
38-
FailureOr<LLVM::LLVMFuncOp> lookupOrCreatePrintF16Fn(Operation *moduleOp);
39-
FailureOr<LLVM::LLVMFuncOp> lookupOrCreatePrintBF16Fn(Operation *moduleOp);
40-
FailureOr<LLVM::LLVMFuncOp> lookupOrCreatePrintF32Fn(Operation *moduleOp);
41-
FailureOr<LLVM::LLVMFuncOp> lookupOrCreatePrintF64Fn(Operation *moduleOp);
36+
FailureOr<LLVM::LLVMFuncOp> lookupOrCreatePrintI64Fn(OpBuilder &b,
37+
Operation *moduleOp);
38+
FailureOr<LLVM::LLVMFuncOp> lookupOrCreatePrintU64Fn(OpBuilder &b,
39+
Operation *moduleOp);
40+
FailureOr<LLVM::LLVMFuncOp> lookupOrCreatePrintF16Fn(OpBuilder &b,
41+
Operation *moduleOp);
42+
FailureOr<LLVM::LLVMFuncOp> lookupOrCreatePrintBF16Fn(OpBuilder &b,
43+
Operation *moduleOp);
44+
FailureOr<LLVM::LLVMFuncOp> lookupOrCreatePrintF32Fn(OpBuilder &b,
45+
Operation *moduleOp);
46+
FailureOr<LLVM::LLVMFuncOp> lookupOrCreatePrintF64Fn(OpBuilder &b,
47+
Operation *moduleOp);
4248
/// Declares a function to print a C-string.
4349
/// If a custom runtime function is defined via `runtimeFunctionName`, it must
4450
/// have the signature void(char const*). The default function is `printString`.
4551
FailureOr<LLVM::LLVMFuncOp>
46-
lookupOrCreatePrintStringFn(Operation *moduleOp,
52+
lookupOrCreatePrintStringFn(OpBuilder &b, Operation *moduleOp,
4753
std::optional<StringRef> runtimeFunctionName = {});
48-
FailureOr<LLVM::LLVMFuncOp> lookupOrCreatePrintOpenFn(Operation *moduleOp);
49-
FailureOr<LLVM::LLVMFuncOp> lookupOrCreatePrintCloseFn(Operation *moduleOp);
50-
FailureOr<LLVM::LLVMFuncOp> lookupOrCreatePrintCommaFn(Operation *moduleOp);
51-
FailureOr<LLVM::LLVMFuncOp> lookupOrCreatePrintNewlineFn(Operation *moduleOp);
52-
FailureOr<LLVM::LLVMFuncOp> lookupOrCreateMallocFn(Operation *moduleOp,
53-
Type indexType);
54-
FailureOr<LLVM::LLVMFuncOp> lookupOrCreateAlignedAllocFn(Operation *moduleOp,
55-
Type indexType);
56-
FailureOr<LLVM::LLVMFuncOp> lookupOrCreateFreeFn(Operation *moduleOp);
57-
FailureOr<LLVM::LLVMFuncOp> lookupOrCreateGenericAllocFn(Operation *moduleOp,
58-
Type indexType);
54+
FailureOr<LLVM::LLVMFuncOp> lookupOrCreatePrintOpenFn(OpBuilder &b,
55+
Operation *moduleOp);
56+
FailureOr<LLVM::LLVMFuncOp> lookupOrCreatePrintCloseFn(OpBuilder &b,
57+
Operation *moduleOp);
58+
FailureOr<LLVM::LLVMFuncOp> lookupOrCreatePrintCommaFn(OpBuilder &b,
59+
Operation *moduleOp);
60+
FailureOr<LLVM::LLVMFuncOp> lookupOrCreatePrintNewlineFn(OpBuilder &b,
61+
Operation *moduleOp);
5962
FailureOr<LLVM::LLVMFuncOp>
60-
lookupOrCreateGenericAlignedAllocFn(Operation *moduleOp, Type indexType);
61-
FailureOr<LLVM::LLVMFuncOp> lookupOrCreateGenericFreeFn(Operation *moduleOp);
63+
lookupOrCreateMallocFn(OpBuilder &b, Operation *moduleOp, Type indexType);
6264
FailureOr<LLVM::LLVMFuncOp>
63-
lookupOrCreateMemRefCopyFn(Operation *moduleOp, Type indexType,
65+
lookupOrCreateAlignedAllocFn(OpBuilder &b, Operation *moduleOp, Type indexType);
66+
FailureOr<LLVM::LLVMFuncOp> lookupOrCreateFreeFn(OpBuilder &b,
67+
Operation *moduleOp);
68+
FailureOr<LLVM::LLVMFuncOp>
69+
lookupOrCreateGenericAllocFn(OpBuilder &b, Operation *moduleOp, Type indexType);
70+
FailureOr<LLVM::LLVMFuncOp>
71+
lookupOrCreateGenericAlignedAllocFn(OpBuilder &b, Operation *moduleOp,
72+
Type indexType);
73+
FailureOr<LLVM::LLVMFuncOp> lookupOrCreateGenericFreeFn(OpBuilder &b,
74+
Operation *moduleOp);
75+
FailureOr<LLVM::LLVMFuncOp>
76+
lookupOrCreateMemRefCopyFn(OpBuilder &b, Operation *moduleOp, Type indexType,
6477
Type unrankedDescriptorType);
6578

6679
/// Create a FuncOp with signature `resultType`(`paramTypes`)` and name `name`.
6780
/// Return a failure if the FuncOp found has unexpected signature.
6881
FailureOr<LLVM::LLVMFuncOp>
69-
lookupOrCreateFn(Operation *moduleOp, StringRef name,
82+
lookupOrCreateFn(OpBuilder &b, Operation *moduleOp, StringRef name,
7083
ArrayRef<Type> paramTypes = {}, Type resultType = {},
7184
bool isVarArg = false, bool isReserved = false);
7285

‎mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -395,7 +395,7 @@ class CoroBeginOpConversion : public AsyncOpConversionPattern<CoroBeginOp> {
395395

396396
// Allocate memory for the coroutine frame.
397397
auto allocFuncOp = LLVM::lookupOrCreateAlignedAllocFn(
398-
op->getParentOfType<ModuleOp>(), rewriter.getI64Type());
398+
rewriter, op->getParentOfType<ModuleOp>(), rewriter.getI64Type());
399399
if (failed(allocFuncOp))
400400
return failure();
401401
auto coroAlloc = rewriter.create<LLVM::CallOp>(
@@ -432,7 +432,7 @@ class CoroFreeOpConversion : public AsyncOpConversionPattern<CoroFreeOp> {
432432

433433
// Free the memory.
434434
auto freeFuncOp =
435-
LLVM::lookupOrCreateFreeFn(op->getParentOfType<ModuleOp>());
435+
LLVM::lookupOrCreateFreeFn(rewriter, op->getParentOfType<ModuleOp>());
436436
if (failed(freeFuncOp))
437437
return failure();
438438
rewriter.replaceOpWithNewOp<LLVM::CallOp>(op, freeFuncOp.value(),

‎mlir/lib/Conversion/LLVMCommon/Pattern.cpp‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -278,12 +278,12 @@ LogicalResult ConvertToLLVMPattern::copyUnrankedDescriptors(
278278
auto module = builder.getInsertionPoint()->getParentOfType<ModuleOp>();
279279
FailureOr<LLVM::LLVMFuncOp> freeFunc, mallocFunc;
280280
if (toDynamic) {
281-
mallocFunc = LLVM::lookupOrCreateMallocFn(module, indexType);
281+
mallocFunc = LLVM::lookupOrCreateMallocFn(builder, module, indexType);
282282
if (failed(mallocFunc))
283283
return failure();
284284
}
285285
if (!toDynamic) {
286-
freeFunc = LLVM::lookupOrCreateFreeFn(module);
286+
freeFunc = LLVM::lookupOrCreateFreeFn(builder, module);
287287
if (failed(freeFunc))
288288
return failure();
289289
}

‎mlir/lib/Conversion/LLVMCommon/PrintCallHelper.cpp‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,7 @@ LogicalResult mlir::LLVM::createPrintStrCall(
6060
Value gep =
6161
builder.create<LLVM::GEPOp>(loc, ptrTy, arrayTy, msgAddr, indices);
6262
FailureOr<LLVM::LLVMFuncOp> printer =
63-
LLVM::lookupOrCreatePrintStringFn(moduleOp, runtimeFunctionName);
63+
LLVM::lookupOrCreatePrintStringFn(builder, moduleOp, runtimeFunctionName);
6464
if (failed(printer))
6565
return failure();
6666
builder.create<LLVM::CallOp>(loc, TypeRange(),

‎mlir/lib/Conversion/MemRefToLLVM/AllocLikeConversion.cpp‎

Lines changed: 12 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -15,24 +15,24 @@
1515
using namespace mlir;
1616

1717
static FailureOr<LLVM::LLVMFuncOp>
18-
getNotalignedAllocFn(const LLVMTypeConverter *typeConverter, Operation *module,
19-
Type indexType) {
18+
getNotalignedAllocFn(OpBuilder &b, const LLVMTypeConverter *typeConverter,
19+
Operation *module, Type indexType) {
2020
bool useGenericFn = typeConverter->getOptions().useGenericFunctions;
2121
if (useGenericFn)
22-
return LLVM::lookupOrCreateGenericAllocFn(module, indexType);
22+
return LLVM::lookupOrCreateGenericAllocFn(b, module, indexType);
2323

24-
return LLVM::lookupOrCreateMallocFn(module, indexType);
24+
return LLVM::lookupOrCreateMallocFn(b, module, indexType);
2525
}
2626

2727
static FailureOr<LLVM::LLVMFuncOp>
28-
getAlignedAllocFn(const LLVMTypeConverter *typeConverter, Operation *module,
29-
Type indexType) {
28+
getAlignedAllocFn(OpBuilder &b, const LLVMTypeConverter *typeConverter,
29+
Operation *module, Type indexType) {
3030
bool useGenericFn = typeConverter->getOptions().useGenericFunctions;
3131

3232
if (useGenericFn)
33-
return LLVM::lookupOrCreateGenericAlignedAllocFn(module, indexType);
33+
return LLVM::lookupOrCreateGenericAlignedAllocFn(b, module, indexType);
3434

35-
return LLVM::lookupOrCreateAlignedAllocFn(module, indexType);
35+
return LLVM::lookupOrCreateAlignedAllocFn(b, module, indexType);
3636
}
3737

3838
Value AllocationOpLLVMLowering::createAligned(
@@ -75,8 +75,8 @@ std::tuple<Value, Value> AllocationOpLLVMLowering::allocateBufferManuallyAlign(
7575
Type elementPtrType = this->getElementPtrType(memRefType);
7676
assert(elementPtrType && "could not compute element ptr type");
7777
FailureOr<LLVM::LLVMFuncOp> allocFuncOp = getNotalignedAllocFn(
78-
getTypeConverter(), op->getParentWithTrait<OpTrait::SymbolTable>(),
79-
getIndexType());
78+
rewriter, getTypeConverter(),
79+
op->getParentWithTrait<OpTrait::SymbolTable>(), getIndexType());
8080
if (failed(allocFuncOp))
8181
return std::make_tuple(Value(), Value());
8282
auto results =
@@ -144,8 +144,8 @@ Value AllocationOpLLVMLowering::allocateBufferAutoAlign(
144144

145145
Type elementPtrType = this->getElementPtrType(memRefType);
146146
FailureOr<LLVM::LLVMFuncOp> allocFuncOp = getAlignedAllocFn(
147-
getTypeConverter(), op->getParentWithTrait<OpTrait::SymbolTable>(),
148-
getIndexType());
147+
rewriter, getTypeConverter(),
148+
op->getParentWithTrait<OpTrait::SymbolTable>(), getIndexType());
149149
if (failed(allocFuncOp))
150150
return Value();
151151
auto results = rewriter.create<LLVM::CallOp>(

‎mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp‎

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -43,13 +43,14 @@ static bool isStaticStrideOrOffset(int64_t strideOrOffset) {
4343
}
4444

4545
static FailureOr<LLVM::LLVMFuncOp>
46-
getFreeFn(const LLVMTypeConverter *typeConverter, ModuleOp module) {
46+
getFreeFn(OpBuilder &b, const LLVMTypeConverter *typeConverter,
47+
ModuleOp module) {
4748
bool useGenericFn = typeConverter->getOptions().useGenericFunctions;
4849

4950
if (useGenericFn)
50-
return LLVM::lookupOrCreateGenericFreeFn(module);
51+
return LLVM::lookupOrCreateGenericFreeFn(b, module);
5152

52-
return LLVM::lookupOrCreateFreeFn(module);
53+
return LLVM::lookupOrCreateFreeFn(b, module);
5354
}
5455

5556
struct AllocOpLowering : public AllocLikeOpLLVMLowering {
@@ -223,8 +224,8 @@ struct DeallocOpLowering : public ConvertOpToLLVMPattern<memref::DeallocOp> {
223224
matchAndRewrite(memref::DeallocOp op, OpAdaptor adaptor,
224225
ConversionPatternRewriter &rewriter) const override {
225226
// Insert the `free` declaration if it is not already present.
226-
FailureOr<LLVM::LLVMFuncOp> freeFunc =
227-
getFreeFn(getTypeConverter(), op->getParentOfType<ModuleOp>());
227+
FailureOr<LLVM::LLVMFuncOp> freeFunc = getFreeFn(
228+
rewriter, getTypeConverter(), op->getParentOfType<ModuleOp>());
228229
if (failed(freeFunc))
229230
return failure();
230231
Value allocatedPtr;
@@ -834,7 +835,8 @@ struct MemRefCopyOpLowering : public ConvertOpToLLVMPattern<memref::CopyOp> {
834835
// potential alignment
835836
auto elemSize = getSizeInBytes(loc, srcType.getElementType(), rewriter);
836837
auto copyFn = LLVM::lookupOrCreateMemRefCopyFn(
837-
op->getParentOfType<ModuleOp>(), getIndexType(), sourcePtr.getType());
838+
rewriter, op->getParentOfType<ModuleOp>(), getIndexType(),
839+
sourcePtr.getType());
838840
if (failed(copyFn))
839841
return failure();
840842
rewriter.create<LLVM::CallOp>(loc, copyFn.value(),

‎mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp‎

Lines changed: 11 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1570,13 +1570,13 @@ class VectorPrintOpConversion : public ConvertOpToLLVMPattern<vector::PrintOp> {
15701570
FailureOr<LLVM::LLVMFuncOp> op = [&]() {
15711571
switch (punct) {
15721572
case PrintPunctuation::Close:
1573-
return LLVM::lookupOrCreatePrintCloseFn(parent);
1573+
return LLVM::lookupOrCreatePrintCloseFn(rewriter, parent);
15741574
case PrintPunctuation::Open:
1575-
return LLVM::lookupOrCreatePrintOpenFn(parent);
1575+
return LLVM::lookupOrCreatePrintOpenFn(rewriter, parent);
15761576
case PrintPunctuation::Comma:
1577-
return LLVM::lookupOrCreatePrintCommaFn(parent);
1577+
return LLVM::lookupOrCreatePrintCommaFn(rewriter, parent);
15781578
case PrintPunctuation::NewLine:
1579-
return LLVM::lookupOrCreatePrintNewlineFn(parent);
1579+
return LLVM::lookupOrCreatePrintNewlineFn(rewriter, parent);
15801580
default:
15811581
llvm_unreachable("unexpected punctuation");
15821582
}
@@ -1610,17 +1610,17 @@ class VectorPrintOpConversion : public ConvertOpToLLVMPattern<vector::PrintOp> {
16101610
PrintConversion conversion = PrintConversion::None;
16111611
FailureOr<Operation *> printer;
16121612
if (printType.isF32()) {
1613-
printer = LLVM::lookupOrCreatePrintF32Fn(parent);
1613+
printer = LLVM::lookupOrCreatePrintF32Fn(rewriter, parent);
16141614
} else if (printType.isF64()) {
1615-
printer = LLVM::lookupOrCreatePrintF64Fn(parent);
1615+
printer = LLVM::lookupOrCreatePrintF64Fn(rewriter, parent);
16161616
} else if (printType.isF16()) {
16171617
conversion = PrintConversion::Bitcast16; // bits!
1618-
printer = LLVM::lookupOrCreatePrintF16Fn(parent);
1618+
printer = LLVM::lookupOrCreatePrintF16Fn(rewriter, parent);
16191619
} else if (printType.isBF16()) {
16201620
conversion = PrintConversion::Bitcast16; // bits!
1621-
printer = LLVM::lookupOrCreatePrintBF16Fn(parent);
1621+
printer = LLVM::lookupOrCreatePrintBF16Fn(rewriter, parent);
16221622
} else if (printType.isIndex()) {
1623-
printer = LLVM::lookupOrCreatePrintU64Fn(parent);
1623+
printer = LLVM::lookupOrCreatePrintU64Fn(rewriter, parent);
16241624
} else if (auto intTy = dyn_cast<IntegerType>(printType)) {
16251625
// Integers need a zero or sign extension on the operand
16261626
// (depending on the source type) as well as a signed or
@@ -1630,7 +1630,7 @@ class VectorPrintOpConversion : public ConvertOpToLLVMPattern<vector::PrintOp> {
16301630
if (width <= 64) {
16311631
if (width < 64)
16321632
conversion = PrintConversion::ZeroExt64;
1633-
printer = LLVM::lookupOrCreatePrintU64Fn(parent);
1633+
printer = LLVM::lookupOrCreatePrintU64Fn(rewriter, parent);
16341634
} else {
16351635
return failure();
16361636
}
@@ -1643,7 +1643,7 @@ class VectorPrintOpConversion : public ConvertOpToLLVMPattern<vector::PrintOp> {
16431643
conversion = PrintConversion::ZeroExt64;
16441644
else if (width < 64)
16451645
conversion = PrintConversion::SignExt64;
1646-
printer = LLVM::lookupOrCreatePrintI64Fn(parent);
1646+
printer = LLVM::lookupOrCreatePrintI64Fn(rewriter, parent);
16471647
} else {
16481648
return failure();
16491649
}

0 commit comments

Comments
 (0)