Use symbol tables more efficiently (#59)

* [mlir][debuginfo] Add flags to attribute creation where necessary

* [MLIR][mlir-link] Cache symbol tables and use SymbolUserMap for replacement lookups

* [MLIR][mlir-link] Use ModuleOp as key for symbolUserMap holding map

Author: Jezurko

mlir/include/mlir/Linker/LinkerInterface.h

@@ -37,7 +37,8 @@ enum LinkerFlags {
 class LinkState {
 public:
   LinkState(ModuleOp dst)
-      : mapping(std::make_shared<IRMapping>()), builder(dst.getBodyRegion()) {}
+      : mapping(std::make_shared<IRMapping>()), builder(dst.getBodyRegion()),
+        symbolTableCollection(), moduleMaps() {}
 
   Operation *clone(Operation *src);
   Operation *cloneWithoutRegions(Operation *src);
@@ -49,14 +50,21 @@ class LinkState {
   LinkState nest(ModuleOp submod) const;
 
   IRMapping &getMapping();
+  SymbolTableCollection &getSymbolTableCollection() {
+    return symbolTableCollection;
+  }
+  SymbolUserMap &getSymbolUserMap(ModuleOp mod);
 
 private:
   // Private constructor used by nest()
   LinkState(ModuleOp dst, std::shared_ptr<IRMapping> mapping)
-      : mapping(std::move(mapping)), builder(dst.getBodyRegion()) {}
+      : mapping(std::move(mapping)), builder(dst.getBodyRegion()),
+        symbolTableCollection(), moduleMaps() {}
 
   std::shared_ptr<IRMapping> mapping;
   OpBuilder builder;
+  SymbolTableCollection symbolTableCollection;
+  DenseMap<ModuleOp, SymbolUserMap> moduleMaps;
 };
 
 struct Conflict {

mlir/lib/Linker/LinkerInterface.cpp

@@ -53,7 +53,10 @@ LinkState LinkState::nest(ModuleOp submod) const {
   return LinkState(submod, mapping);
 }
 
-IRMapping &LinkState::getMapping() { return *mapping.get(); }
+IRMapping &LinkState::getMapping() { return *mapping; }
+SymbolUserMap &LinkState::getSymbolUserMap(ModuleOp mod) {
+  return moduleMaps.try_emplace(mod, symbolTableCollection, mod).first->second;
+}
 
 //===----------------------------------------------------------------------===//
 // SymbolAttrLinkerInterface
@@ -83,21 +86,17 @@ static void renameSymbolRefIn(Operation *op, StringAttr newName) {
 static LogicalResult renameRemappedUsersOf(Operation *op, StringAttr newName,
                                            LinkState &state) {
   ModuleOp module = op->getParentOfType<ModuleOp>();
+  SymbolUserMap &userMap = state.getSymbolUserMap(module);
   // TODO: use something like SymbolTableAnalysis
-  SymbolTable src(module);
-  if (auto uses = src.getSymbolUses(op, module)) {
-    for (SymbolTable::SymbolUse use : *uses) {
-      // TODO: add test where user is not remapped
-      Operation *dstUser = state.remapped(use.getUser());
-      if (!dstUser)
-        continue;
-      renameSymbolRefIn(dstUser, newName);
-    }
-
-    return success();
+  auto users = userMap.getUsers(op);
+  for (Operation *user : users) {
+    // TODO: add test where user is not remapped
+    Operation *dstUser = state.remapped(user);
+    if (!dstUser)
+      continue;
+    renameSymbolRefIn(dstUser, newName);
   }
-
-  return op->emitError("failed to rename symbol to a unique name");
+  return success();
 }
 
 StringRef SymbolAttrLinkerInterface::getSymbol(Operation *op) const {
@@ -119,7 +118,8 @@ void SymbolAttrLinkerInterface::registerForLink(Operation *op) {
 }
 
 LogicalResult SymbolAttrLinkerInterface::link(LinkState &state) const {
-  SymbolTable st(state.getDestinationOp());
+  SymbolTable &st =
+      state.getSymbolTableCollection().getSymbolTable(state.getDestinationOp());
 
   auto materializeError = [&](Operation *op) {
     return op->emitError("failed to materialize symbol for linking");
@@ -145,7 +145,7 @@ LogicalResult SymbolAttrLinkerInterface::link(LinkState &state) const {
     if (st.lookup(name)) {
       StringAttr newName = getUniqueNameIn(st, name);
       st.setSymbolName(materialized, newName);
-      toRenameUsers.push_back({op, newName});
+      toRenameUsers.emplace_back(op, newName);
     }
 
     st.insert(materialized);

mlir/test/CAPI/llvm.c

@@ -303,12 +303,12 @@ static void testDebugInfoAttributes(MlirContext ctx) {
   // CHECK: #llvm.di_derived_type<{{.*}}>
   // CHECK-NOT: dwarfAddressSpace
   mlirAttributeDump(mlirLLVMDIDerivedTypeAttrGet(
-      ctx, 0, bar, di_type, 64, 8, 0, MLIR_CAPI_DWARF_ADDRESS_SPACE_NULL,
+      ctx, 0, bar, di_type, 0, 64, 8, 0, MLIR_CAPI_DWARF_ADDRESS_SPACE_NULL,
       di_type));
 
   // CHECK: #llvm.di_derived_type<{{.*}} dwarfAddressSpace = 3{{.*}}>
   mlirAttributeDump(
-      mlirLLVMDIDerivedTypeAttrGet(ctx, 0, bar, di_type, 64, 8, 0, 3, di_type));
+      mlirLLVMDIDerivedTypeAttrGet(ctx, 0, bar, di_type, 0, 64, 8, 0, 3, di_type));
 
   MlirAttribute subroutine_type =
       mlirLLVMDISubroutineTypeAttrGet(ctx, 0x0, 1, &di_type);