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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: RohithPariki

<details>
<summary>Changes</summary>

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. 


---
Full diff: https://github.com/llvm/llvm-project/pull/214997.diff


4 Files Affected:

- (modified) mlir/include/mlir/Conversion/LLVMCommon/PrintCallHelper.h (+1-1) 
- (modified) mlir/lib/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.cpp (+3-3) 
- (modified) mlir/lib/Conversion/LLVMCommon/PrintCallHelper.cpp (+4-4) 
- (modified) mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp (+2-2) 


``````````diff
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();

``````````

</details>


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