[Mlir-commits] [mlir] [MLIR][Conversion] Fix assert and print lowering in gpu.module (PR #214997)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Aug 11 00:23:23 PDT 2026
https://github.com/RohithPariki updated https://github.com/llvm/llvm-project/pull/214997
>From 96365456c30a12c45551e8d6e6f8354c0bebe6d8 Mon Sep 17 00:00:00 2001
From: Rohith Pariki <rohithpariki at gmail.com>
Date: Sat, 8 Aug 2026 19:12:11 +0530
Subject: [PATCH 1/3] [MLIR][Conversion] Fix assert and print lowering in
gpu.module
This fixes an issue where cf.assert or vector.print inside a gpu.module creates the global string or abort function in the top-level ModuleOp instead of the nearest SymbolTable (e.g. gpu.module), causing undefined references when the GPU module is compiled.
---
mlir/include/mlir/Conversion/LLVMCommon/PrintCallHelper.h | 2 +-
.../Conversion/ControlFlowToLLVM/ControlFlowToLLVM.cpp | 6 +++---
mlir/lib/Conversion/LLVMCommon/PrintCallHelper.cpp | 8 ++++----
mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp | 4 ++--
4 files changed, 10 insertions(+), 10 deletions(-)
diff --git a/mlir/include/mlir/Conversion/LLVMCommon/PrintCallHelper.h b/mlir/include/mlir/Conversion/LLVMCommon/PrintCallHelper.h
index d7de40555bb6a..193b50f71dfe2 100644
--- a/mlir/include/mlir/Conversion/LLVMCommon/PrintCallHelper.h
+++ b/mlir/include/mlir/Conversion/LLVMCommon/PrintCallHelper.h
@@ -25,7 +25,7 @@ namespace LLVM {
/// If a custom runtime function is defined via `runtimeFunctionName`, it must
/// have the signature void(char const*). The default function is `printString`.
LogicalResult createPrintStrCall(
- OpBuilder &builder, Location loc, ModuleOp moduleOp, StringRef symbolName,
+ OpBuilder &builder, Location loc, Operation *moduleOp, StringRef symbolName,
StringRef string, const LLVMTypeConverter &typeConverter,
bool addNewline = true, std::optional<StringRef> runtimeFunctionName = {},
SymbolTableCollection *symbolTables = nullptr);
diff --git a/mlir/lib/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.cpp b/mlir/lib/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.cpp
index fef78a46d69fc..6506e819151a1 100644
--- a/mlir/lib/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.cpp
+++ b/mlir/lib/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.cpp
@@ -50,7 +50,7 @@ struct AssertOpLowering : public ConvertOpToLLVMPattern<cf::AssertOp> {
matchAndRewrite(cf::AssertOp op, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
auto loc = op.getLoc();
- auto module = op->getParentOfType<ModuleOp>();
+ auto module = SymbolTable::getNearestSymbolTable(op);
// Split block at `assert` operation.
Block *opBlock = rewriter.getInsertionBlock();
@@ -68,10 +68,10 @@ struct AssertOpLowering : public ConvertOpToLLVMPattern<cf::AssertOp> {
if (abortOnFailedAssert) {
// Insert the `abort` declaration if necessary.
- auto abortFunc = module.lookupSymbol<LLVM::LLVMFuncOp>("abort");
+ auto abortFunc = SymbolTable::lookupSymbolIn<LLVM::LLVMFuncOp>(module, "abort");
if (!abortFunc) {
OpBuilder::InsertionGuard guard(rewriter);
- rewriter.setInsertionPointToStart(module.getBody());
+ rewriter.setInsertionPointToStart(&module->getRegion(0).front());
auto abortFuncTy = LLVM::LLVMFunctionType::get(getVoidType(), {});
abortFunc = LLVM::LLVMFuncOp::create(rewriter, rewriter.getUnknownLoc(),
"abort", abortFuncTy);
diff --git a/mlir/lib/Conversion/LLVMCommon/PrintCallHelper.cpp b/mlir/lib/Conversion/LLVMCommon/PrintCallHelper.cpp
index da4443dc86053..dde6a1fa1f21e 100644
--- a/mlir/lib/Conversion/LLVMCommon/PrintCallHelper.cpp
+++ b/mlir/lib/Conversion/LLVMCommon/PrintCallHelper.cpp
@@ -22,7 +22,7 @@ using namespace llvm;
/// returned. Otherwise, a unique and yet unused identifier is computed starting
/// from the requested one.
static std::string
-ensureSymbolNameIsUnique(ModuleOp moduleOp, StringRef symbolName,
+ensureSymbolNameIsUnique(Operation *moduleOp, StringRef symbolName,
SymbolTableCollection *symbolTables = nullptr) {
if (symbolTables) {
SymbolTable &symbolTable = symbolTables->getSymbolTable(moduleOp);
@@ -39,19 +39,19 @@ ensureSymbolNameIsUnique(ModuleOp moduleOp, StringRef symbolName,
static int counter = 0;
std::string uniqueName = std::string(symbolName);
- while (moduleOp.lookupSymbol(uniqueName)) {
+ while (SymbolTable::lookupSymbolIn(moduleOp, uniqueName)) {
uniqueName = std::string(symbolName) + "_" + std::to_string(counter++);
}
return uniqueName;
}
LogicalResult mlir::LLVM::createPrintStrCall(
- OpBuilder &builder, Location loc, ModuleOp moduleOp, StringRef symbolName,
+ OpBuilder &builder, Location loc, Operation *moduleOp, StringRef symbolName,
StringRef string, const LLVMTypeConverter &typeConverter, bool addNewline,
std::optional<StringRef> runtimeFunctionName,
SymbolTableCollection *symbolTables) {
auto ip = builder.saveInsertionPoint();
- builder.setInsertionPointToStart(moduleOp.getBody());
+ builder.setInsertionPointToStart(&moduleOp->getRegion(0).front());
MLIRContext *ctx = builder.getContext();
// Create a zero-terminated byte representation and allocate global symbol.
diff --git a/mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp b/mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp
index 54c117bae476b..cb6400d121fa8 100644
--- a/mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp
+++ b/mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp
@@ -1589,7 +1589,7 @@ class VectorPrintOpConversion : public ConvertOpToLLVMPattern<vector::PrintOp> {
LogicalResult
matchAndRewrite(vector::PrintOp printOp, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
- auto parent = printOp->getParentOfType<ModuleOp>();
+ auto parent = SymbolTable::getNearestSymbolTable(printOp);
if (!parent)
return failure();
@@ -1653,7 +1653,7 @@ class VectorPrintOpConversion : public ConvertOpToLLVMPattern<vector::PrintOp> {
};
LogicalResult emitScalarPrint(ConversionPatternRewriter &rewriter,
- ModuleOp parent, Location loc, Type printType,
+ Operation *parent, Location loc, Type printType,
Value value) const {
if (typeConverter->convertType(printType) == nullptr)
return failure();
>From b52ccd7b23897b66e08d9f7174ca43aa159eb115 Mon Sep 17 00:00:00 2001
From: Rohith Pariki <rohithpariki at gmail.com>
Date: Tue, 11 Aug 2026 01:48:48 +0530
Subject: [PATCH 2/3] test(mlir): add unit tests for cf.assert and vector.print
lowering inside gpu.module
---
.../ControlFlowToLLVM/ControlFlowToLLVM.cpp | 3 ++-
.../Conversion/ControlFlowToLLVM/assert.mlir | 16 ++++++++++++++++
.../VectorToLLVM/vector-to-llvm-interface.mlir | 14 ++++++++++++++
3 files changed, 32 insertions(+), 1 deletion(-)
diff --git a/mlir/lib/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.cpp b/mlir/lib/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.cpp
index 6506e819151a1..ac916564dea37 100644
--- a/mlir/lib/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.cpp
+++ b/mlir/lib/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.cpp
@@ -68,7 +68,8 @@ struct AssertOpLowering : public ConvertOpToLLVMPattern<cf::AssertOp> {
if (abortOnFailedAssert) {
// Insert the `abort` declaration if necessary.
- auto abortFunc = SymbolTable::lookupSymbolIn<LLVM::LLVMFuncOp>(module, "abort");
+ auto abortFunc =
+ SymbolTable::lookupSymbolIn<LLVM::LLVMFuncOp>(module, "abort");
if (!abortFunc) {
OpBuilder::InsertionGuard guard(rewriter);
rewriter.setInsertionPointToStart(&module->getRegion(0).front());
diff --git a/mlir/test/Conversion/ControlFlowToLLVM/assert.mlir b/mlir/test/Conversion/ControlFlowToLLVM/assert.mlir
index 18d0526ecf1a9..913671b6350c7 100644
--- a/mlir/test/Conversion/ControlFlowToLLVM/assert.mlir
+++ b/mlir/test/Conversion/ControlFlowToLLVM/assert.mlir
@@ -20,3 +20,19 @@ func.func @main() {
// CHECK: %[[ADDRESS_OF:.*]] = llvm.mlir.addressof @{{.*}} : !llvm.ptr{{$}}
// CHECK: %[[GEP:.*]] = llvm.getelementptr %[[ADDRESS_OF]][0] : (!llvm.ptr) -> !llvm.ptr, !llvm.array<{{[0-9]+}} x i8>
// CHECK: llvm.call @puts(%[[GEP]]) : (!llvm.ptr) -> ()
+
+// -----
+
+gpu.module @gpu_module {
+ gpu.func @kernel() {
+ %cond = arith.constant 0 : i1
+ cf.assert %cond, "gpu assertion"
+ gpu.return
+ }
+}
+
+// CHECK-LABEL: gpu.module @gpu_module
+// CHECK: llvm.func @puts(!llvm.ptr)
+// CHECK-LABEL: gpu.func @kernel
+// CHECK: %[[ADDRESS_OF:.*]] = llvm.mlir.addressof @assert_msg : !llvm.ptr
+
diff --git a/mlir/test/Conversion/VectorToLLVM/vector-to-llvm-interface.mlir b/mlir/test/Conversion/VectorToLLVM/vector-to-llvm-interface.mlir
index 4cbea04e8076e..e1cafab798879 100644
--- a/mlir/test/Conversion/VectorToLLVM/vector-to-llvm-interface.mlir
+++ b/mlir/test/Conversion/VectorToLLVM/vector-to-llvm-interface.mlir
@@ -878,6 +878,20 @@ func.func @print_string() {
// -----
+gpu.module @gpu_module {
+ gpu.func @kernel() {
+ vector.print str "GPU Print"
+ gpu.return
+ }
+}
+// CHECK-LABEL: gpu.module @gpu_module
+// CHECK: llvm.mlir.global private constant @[[GPU_STR:.*]]("GPU Print\00")
+// CHECK-LABEL: gpu.func @kernel
+// CHECK: %[[GLOBAL_ADDR:.*]] = llvm.mlir.addressof @[[GPU_STR]] : !llvm.ptr
+
+// -----
+
+
//===----------------------------------------------------------------------===//
// vector.reduction
//===----------------------------------------------------------------------===//
>From 8e16ff1dbd3d2b423bc010333c7b1946a3ff0d5f Mon Sep 17 00:00:00 2001
From: Rohith Pariki <rohithpariki at gmail.com>
Date: Tue, 11 Aug 2026 12:52:54 +0530
Subject: [PATCH 3/3] fix(mlir): fix SymbolTable lookup for abort in
ControlFlowToLLVM
---
mlir/lib/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.cpp | 5 ++++-
1 file changed, 4 insertions(+), 1 deletion(-)
diff --git a/mlir/lib/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.cpp b/mlir/lib/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.cpp
index ac916564dea37..c5a61f78dd61f 100644
--- a/mlir/lib/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.cpp
+++ b/mlir/lib/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.cpp
@@ -69,7 +69,10 @@ struct AssertOpLowering : public ConvertOpToLLVMPattern<cf::AssertOp> {
if (abortOnFailedAssert) {
// Insert the `abort` declaration if necessary.
auto abortFunc =
- SymbolTable::lookupSymbolIn<LLVM::LLVMFuncOp>(module, "abort");
+ symbolTables
+ ? symbolTables->lookupSymbolIn<LLVM::LLVMFuncOp>(module, "abort")
+ : dyn_cast_or_null<LLVM::LLVMFuncOp>(
+ SymbolTable::lookupSymbolIn(module, "abort"));
if (!abortFunc) {
OpBuilder::InsertionGuard guard(rewriter);
rewriter.setInsertionPointToStart(&module->getRegion(0).front());
More information about the Mlir-commits
mailing list