[mlir] Use SymbolTableCollection to lookup referenced symbol in AddressOfOp

[ezhulenev]

Summary

Depends On D131285

[ezhulenev]

Revert accidental formatting changes

[Mogball]

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

1034

And below as well

[ezhulenev]

Add argument name to extra decls

[ezhulenev]

Rebased

[ezhulenev]

Rename args

[ezhulenev]

Rename args one more time

[Mogball]

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

1740

In the declaration, the argument is named symbolMap instead of symbolTable..

[Mogball]

Thanks lol

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

[{
             LLVM::LLVMPointerType::get(func.getFunctionType()), func.getName());
       $_state.addAttributes(attrs);
     }]>
   ];
 
   let extraClassDeclaration = [{
     /// Return the llvm.mlir.global operation that defined the value referenced
     /// here.
-    GlobalOp getGlobal();
+    GlobalOp getGlobal(SymbolTableCollection &symbolTable);
 
     /// Return the llvm.func operation that is referenced here.
-    LLVMFuncOp getFunction();
+    LLVMFuncOp getFunction(SymbolTableCollection &symbolTable);
   }];
 
   let assemblyFormat = "$global_name attr-dict `:` type($res)";
 }
 
 def LLVM_MetadataOp : LLVM_Op<"metadata", [
    NoRegionArguments, SymbolTable, Symbol
 ]> {

mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_TARGET_LLVMIR_MODULETRANSLATION_H
 #define MLIR_TARGET_LLVMIR_MODULETRANSLATION_H
 
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/Value.h"
+#include "mlir/IR/SymbolTable.h"
 #include "mlir/Target/LLVMIR/Export.h"
 #include "mlir/Target/LLVMIR/LLVMTranslationInterface.h"
 #include "mlir/Target/LLVMIR/TypeToLLVM.h"
 
 #include "llvm/ADT/SetVector.h"
 #include "llvm/Frontend/OpenMP/OMPIRBuilder.h"
 
 namespace llvm {
explicit SaveStack(ModuleTranslation &m, Args &&...args)
       moduleTranslation.stackPush<T>(std::forward<Args>(args)...);
     }
     ~SaveStack() { moduleTranslation.stackPop(); }
 
   private:
     ModuleTranslation &moduleTranslation;
   };
 
+  SymbolTableCollection& symbolTable() { return symbolTableCollection; }
+
 private:
   ModuleTranslation(Operation *module,
                     std::unique_ptr<llvm::Module> llvmModule);
   ~ModuleTranslation();
 
   /// Converts individual components.
   LogicalResult convertOperation(Operation &op, llvm::IRBuilderBase &builder);
   LogicalResult convertFunctionSignatures();
private:
 
   /// Mapping from an access scope metadata operation to its LLVM metadata.
   /// This map is populated on module entry.
   DenseMap<Operation *, llvm::MDNode *> aliasScopeMetadataMapping;
 
   /// Stack of user-specified state elements, useful when translating operations
   /// with regions.
   SmallVector<std::unique_ptr<StackFrame>> stack;
+
+  /// A cache for the symbol tables constructed during symbols lookup.
+  SymbolTableCollection symbolTableCollection;
 };
 
 namespace detail {
 /// For all blocks in the region that were converted to LLVM IR using the given
 /// ModuleTranslation, connect the PHI nodes of the corresponding LLVM IR blocks
 /// to the results of preceding blocks.
 void connectPHINodes(Region &region, const ModuleTranslation &state);
 

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

 static Operation *parentLLVMModule(Operation *op) {
   Operation *module = op->getParentOp();
   while (module && !satisfiesLLVMModule(module))
     module = module->getParentOp();
   assert(module && "unexpected operation outside of a module");
   return module;
 }
 
-GlobalOp AddressOfOp::getGlobal() {
+GlobalOp AddressOfOp::getGlobal(SymbolTableCollection &symbolTable) {
   return dyn_cast_or_null<GlobalOp>(
-      SymbolTable::lookupSymbolIn(parentLLVMModule(*this), getGlobalName()));
+      symbolTable.lookupSymbolIn(parentLLVMModule(*this), getGlobalNameAttr()));
 }
 
-LLVMFuncOp AddressOfOp::getFunction() {
+LLVMFuncOp AddressOfOp::getFunction(SymbolTableCollection &symbolTable) {
   return dyn_cast_or_null<LLVMFuncOp>(
-      SymbolTable::lookupSymbolIn(parentLLVMModule(*this), getGlobalName()));
+      symbolTable.lookupSymbolIn(parentLLVMModule(*this), getGlobalNameAttr()));
 }
 
 LogicalResult
 AddressOfOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
   Operation *symbol =
       symbolTable.lookupSymbolIn(parentLLVMModule(*this), getGlobalNameAttr());
 
   auto global = dyn_cast_or_null<GlobalOp>(symbol);

mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp

if (auto switchOp = dyn_cast<LLVM::SwitchOp>(opInst)) {
     moduleTranslation.mapBranch(&opInst, switchInst);
     return success();
   }
 
   // Emit addressof.  We need to look up the global value referenced by the
   // operation and store it in the MLIR-to-LLVM value mapping.  This does not
   // emit any LLVM instruction.
   if (auto addressOfOp = dyn_cast<LLVM::AddressOfOp>(opInst)) {
-    LLVM::GlobalOp global = addressOfOp.getGlobal();
-    LLVM::LLVMFuncOp function = addressOfOp.getFunction();
+    LLVM::GlobalOp global =
+        addressOfOp.getGlobal(moduleTranslation.symbolTable());
+    LLVM::LLVMFuncOp function =
+        addressOfOp.getFunction(moduleTranslation.symbolTable());
 
     // The verifier should not have allowed this.
     assert((global || function) &&
            "referencing an undefined global or function");
 
     moduleTranslation.mapValue(
         addressOfOp.getResult(),
         global ? moduleTranslation.lookupGlobal(global)

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

convertOmpThreadprivate(Operation &opInst, llvm::IRBuilderBase &builder,
   auto threadprivateOp = cast<omp::ThreadprivateOp>(opInst);
 
   Value symAddr = threadprivateOp.sym_addr();
   auto *symOp = symAddr.getDefiningOp();
   if (!isa<LLVM::AddressOfOp>(symOp))
     return opInst.emitError("Addressing symbol not found");
   LLVM::AddressOfOp addressOfOp = dyn_cast<LLVM::AddressOfOp>(symOp);
 
-  LLVM::GlobalOp global = addressOfOp.getGlobal();
+  LLVM::GlobalOp global =
+      addressOfOp.getGlobal(moduleTranslation.symbolTable());
   llvm::GlobalValue *globalValue = moduleTranslation.lookupGlobal(global);
   llvm::Value *data =
       builder.CreateBitCast(globalValue, builder.getInt8PtrTy());
   llvm::Type *type = globalValue->getValueType();
   llvm::TypeSize typeSize =
       builder.GetInsertBlock()->getModule()->getDataLayout().getTypeStoreSize(
           type);
   llvm::ConstantInt *size = builder.getInt64(typeSize.getFixedSize());