[mlir][memref] Do not access erased op in memref.global lowering (#148355)

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`.

Author: matthias-springer

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();
   }