[mlir][memref] Do not access erased op in memref.global lowering

[matthias-springer]

Do not access the erased memref.global operation in the lowering pattern. That won't work anymore in a One-Shot Dialect Conversion and triggers a use-after-free sanitizer error.

After the One-Shot Dialect Conversion refactoring, a ConversionPatternRewriter will behave more like a normal PatternRewriter.

[llvmbot]

@llvm/pr-subscribers-mlir

Author: Matthias Springer (matthias-springer)

Changes

Do not access the erased memref.global operation in the lowering pattern. That won't work anymore in a One-Shot Dialect Conversion and triggers a use-after-free sanitizer error.

After the One-Shot Dialect Conversion refactoring, a ConversionPatternRewriter will behave more like a normal PatternRewriter.

---

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

1 Files Affected:

• (modified) mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp (+14-18)

diff --git a/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp b/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
index 00bbdcb12e326..83681b2d5fd87 100644
--- a/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
+++ b/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
@@ -754,9 +754,11 @@ class GlobalMemrefOpLowering : public ConvertOpToLLVMPattern<memref::GlobalOp> {
 
     LLVM::Linkage linkage =
         global.isPublic() ? LLVM::Linkage::External : LLVM::Linkage::Private;
+    bool isExternal = global.isExternal();
+    bool isUninitialized = global.isUninitialized();
 
     Attribute initialValue = nullptr;
-    if (!global.isExternal() && !global.isUninitialized()) {
+    if (!isExternal && !isUninitialized) {
       auto elementsAttr = llvm::cast<ElementsAttr>(*global.getInitialValue());
       initialValue = elementsAttr;
 
@@ -773,35 +775,29 @@ class GlobalMemrefOpLowering : public ConvertOpToLLVMPattern<memref::GlobalOp> {
       return global.emitOpError(
           "memory space cannot be converted to an integer address space");
 
+    // Remove old operation from symbol table.
+    SymbolTable *symbolTable = nullptr;
     if (symbolTables) {
       Operation *symbolTableOp =
           global->getParentWithTrait<OpTrait::SymbolTable>();
-
-      if (symbolTableOp) {
-        SymbolTable &symbolTable = symbolTables->getSymbolTable(symbolTableOp);
-        symbolTable.remove(global);
-      }
+      symbolTable = &symbolTables->getSymbolTable(symbolTableOp);
+      symbolTable->remove(global);
     }
 
+    // Create new operation.
     auto newGlobal = rewriter.replaceOpWithNewOp<LLVM::GlobalOp>(
         global, arrayTy, global.getConstant(), linkage, global.getSymName(),
         initialValue, alignment, *addressSpace);
 
-    if (symbolTables) {
-      Operation *symbolTableOp =
-          global->getParentWithTrait<OpTrait::SymbolTable>();
-
-      if (symbolTableOp) {
-        SymbolTable &symbolTable = symbolTables->getSymbolTable(symbolTableOp);
-        symbolTable.insert(newGlobal, rewriter.getInsertionPoint());
-      }
-    }
+    // Insert new operation into symbol table.
+    if (symbolTable)
+      symbolTable->insert(newGlobal, rewriter.getInsertionPoint());
 
-    if (!global.isExternal() && global.isUninitialized()) {
+    if (!isExternal && isUninitialized) {
       rewriter.createBlock(&newGlobal.getInitializerRegion());
       Value undef[] = {
-          rewriter.create<LLVM::UndefOp>(global.getLoc(), arrayTy)};
-      rewriter.create<LLVM::ReturnOp>(global.getLoc(), undef);
+          rewriter.create<LLVM::UndefOp>(newGlobal.getLoc(), arrayTy)};
+      rewriter.create<LLVM::ReturnOp>(newGlobal.getLoc(), undef);
     }
     return success();
   }

[zero9178]

LGTM!