[Mlir-commits] [mlir] [MLIR][Conversion] Fix assert and print lowering in gpu.module (PR #214997)

[llvmlistbot at llvm.org]

https://github.com/RohithPariki created https://github.com/llvm/llvm-project/pull/214997

Fixes #214402.

### Description
This fixes an issue where operations like `cf.assert` or `vector.print` that are located inside a `gpu.module` would incorrectly create their global strings or `abort` function declarations in the top-level `ModuleOp` instead of the nearest `SymbolTable` (which is the `gpu.module` itself). 

This behavior caused "undefined reference" errors when compiling the GPU module, because the required globals or function declarations were placed in the host module rather than the device module. 

The fix updates `AssertOpLowering` and `PrintCallHelper` to use `SymbolTable::getNearestSymbolTable()` instead of strictly assuming a top-level `ModuleOp`. It replaces instances of `ModuleOp` with `Operation*` representing the nearest symbol table to accurately place these declarations within their correct local context.

**Note to reviewers:**
* AI was used to help identify the root cause and generate this patch. 


>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] [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();