From 2711dd5a958af5180c33d15a3c648e7d91fe3d95 Mon Sep 17 00:00:00 2001
From: Nandor Licker <n@ndor.email>
Date: Sat, 18 Nov 2023 04:10:11 -0800
Subject: [PATCH] [NFC][InstanceGraph] Remove the slow target getter helper

---
 .../Dialect/FIRRTL/FIRRTLDeclarations.td      |  3 +--
 .../circt/Support/InstanceGraphInterface.td   | 11 ----------
 .../HandshakeToHW/HandshakeToHW.cpp           |  4 +++-
 lib/Dialect/FIRRTL/FIRRTLOps.cpp              | 21 +++----------------
 lib/Dialect/HW/Transforms/FlattenIO.cpp       |  8 +++++--
 lib/Dialect/SystemC/SystemCOps.cpp            |  3 ++-
 6 files changed, 15 insertions(+), 35 deletions(-)

diff --git a/include/circt/Dialect/FIRRTL/FIRRTLDeclarations.td b/include/circt/Dialect/FIRRTL/FIRRTLDeclarations.td
index 9767ae2aca..f3ed5f2d45 100644
--- a/include/circt/Dialect/FIRRTL/FIRRTLDeclarations.td
+++ b/include/circt/Dialect/FIRRTL/FIRRTLDeclarations.td
@@ -44,8 +44,7 @@ def InstanceOp : HardwareDeclOp<"instance", [
     DeclareOpInterfaceMethods<FInstanceLike>,
     DeclareOpInterfaceMethods<SymbolUserOpInterface>,
     DeclareOpInterfaceMethods<InstanceGraphInstanceOpInterface, [
-      "getModuleName",
-      "getReferencedModuleSlow"
+      "getModuleName"
     ]>,
   ]> {
   let summary = "Instantiate an instance of a module";
diff --git a/include/circt/Support/InstanceGraphInterface.td b/include/circt/Support/InstanceGraphInterface.td
index 7b0635d12d..e1ccc94da0 100644
--- a/include/circt/Support/InstanceGraphInterface.td
+++ b/include/circt/Support/InstanceGraphInterface.td
@@ -39,17 +39,6 @@ def InstanceGraphInstanceOpInterface : OpInterface<"InstanceOpInterface"> {
     "::mlir::StringAttr", "getReferencedModuleNameAttr", (ins),
     /*methodBody=*/[{}],
     /*defaultImplementation=*/[{ return $_op.getModuleNameAttr().getAttr(); }]>,
-
-    InterfaceMethod<[{
-      Get the referenced module (slow, unsafe).  This function directly accesses
-      the parent operation to lookup a symbol, which is unsafe in many contexts.
-    }],
-    "::mlir::Operation *", "getReferencedModuleSlow", (ins),
-    /*methodBody=*/[{}],
-    /*defaultImplementation=*/[{
-        return SymbolTable::lookupNearestSymbolFrom(
-            $_op, getReferencedModuleNameAttr());
-    }]>
   ];
 }
 
diff --git a/lib/Conversion/HandshakeToHW/HandshakeToHW.cpp b/lib/Conversion/HandshakeToHW/HandshakeToHW.cpp
index 4dc67f9670..b1cd634055 100644
--- a/lib/Conversion/HandshakeToHW/HandshakeToHW.cpp
+++ b/lib/Conversion/HandshakeToHW/HandshakeToHW.cpp
@@ -322,7 +322,9 @@ static LogicalResult convertExtMemoryOps(HWModuleOp mod) {
     // Get the attached extmemory external module.
     auto extmemInstance = cast<hw::InstanceOp>(*arg.getUsers().begin());
     auto extmemMod =
-        cast<hw::HWModuleExternOp>(extmemInstance.getReferencedModuleSlow());
+        cast<hw::HWModuleExternOp>(SymbolTable::lookupNearestSymbolFrom(
+            extmemInstance, extmemInstance.getReferencedModuleNameAttr()));
+
     ModulePortInfo portInfo(extmemMod.getPortList());
 
     // The extmemory external module's interface is a direct wrapping of the
diff --git a/lib/Dialect/FIRRTL/FIRRTLOps.cpp b/lib/Dialect/FIRRTL/FIRRTLOps.cpp
index 2a2d443a93..fea64bcecd 100644
--- a/lib/Dialect/FIRRTL/FIRRTLOps.cpp
+++ b/lib/Dialect/FIRRTL/FIRRTLOps.cpp
@@ -1913,18 +1913,11 @@ bool ExtClassOp::canDiscardOnUseEmpty() {
 // Declarations
 //===----------------------------------------------------------------------===//
 
-/// Lookup the module or extmodule for the symbol.  This returns null on
-/// invalid IR.
-Operation *InstanceOp::getReferencedModuleSlow() {
+SmallVector<::circt::hw::PortInfo> InstanceOp::getPortList() {
   auto circuit = (*this)->getParentOfType<CircuitOp>();
   if (!circuit)
-    return nullptr;
-
-  return circuit.lookupSymbol<FModuleLike>(getModuleNameAttr());
-}
-
-SmallVector<::circt::hw::PortInfo> InstanceOp::getPortList() {
-  return cast<hw::PortList>(getReferencedModuleSlow()).getPortList();
+    llvm::report_fatal_error("instance op not in circuit");
+  return circuit.lookupSymbol<hw::PortList>(getModuleNameAttr()).getPortList();
 }
 
 void InstanceOp::build(OpBuilder &builder, OperationState &result,
@@ -2930,14 +2923,6 @@ Operation *ObjectOp::getReferencedModule(const SymbolTable &symtbl) {
   return getReferencedClass(symtbl);
 }
 
-Operation *ObjectOp::getReferencedModuleSlow() {
-  auto circuit = (*this)->getParentOfType<CircuitOp>();
-  if (!circuit)
-    return nullptr;
-
-  return circuit.lookupSymbol<ClassLike>(getClassNameAttr());
-}
-
 StringRef ObjectOp::getInstanceName() { return getName(); }
 
 StringAttr ObjectOp::getInstanceNameAttr() { return getNameAttr(); }
diff --git a/lib/Dialect/HW/Transforms/FlattenIO.cpp b/lib/Dialect/HW/Transforms/FlattenIO.cpp
index 020302c33b..faf2be0fe7 100644
--- a/lib/Dialect/HW/Transforms/FlattenIO.cpp
+++ b/lib/Dialect/HW/Transforms/FlattenIO.cpp
@@ -94,8 +94,10 @@ struct InstanceOpConversion : public OpConversionPattern<hw::InstanceOp> {
     }
 
     // Create the new instance...
+    Operation *targetModule = SymbolTable::lookupNearestSymbolFrom(
+        op, op.getReferencedModuleNameAttr());
     auto newInstance = rewriter.create<hw::InstanceOp>(
-        loc, op.getReferencedModuleSlow(), op.getInstanceName(), convOperands);
+        loc, targetModule, op.getInstanceName(), convOperands);
 
     // re-create any structs in the result.
     llvm::SmallVector<Value> convResults;
@@ -381,7 +383,9 @@ static LogicalResult flattenOpsOfType(ModuleOp module, bool recursive) {
 
     // And likewise with the converted instance ops.
     for (auto instanceOp : convertedInstances) {
-      Operation *targetModule = instanceOp.getReferencedModuleSlow();
+      Operation *targetModule = SymbolTable::lookupNearestSymbolFrom(
+          instanceOp, instanceOp.getReferencedModuleNameAttr());
+
       auto ioInfo = ioInfoMap[targetModule];
       instanceOp.setInputNames(ArrayAttr::get(
           instanceOp.getContext(),
diff --git a/lib/Dialect/SystemC/SystemCOps.cpp b/lib/Dialect/SystemC/SystemCOps.cpp
index 95a6cf000f..6800ff5bed 100644
--- a/lib/Dialect/SystemC/SystemCOps.cpp
+++ b/lib/Dialect/SystemC/SystemCOps.cpp
@@ -531,7 +531,8 @@ InstanceDeclOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
 }
 
 SmallVector<hw::PortInfo> InstanceDeclOp::getPortList() {
-  return cast<hw::PortList>(getReferencedModuleSlow()).getPortList();
+  return cast<hw::PortList>(getReferencedModuleCached(/*cache=*/nullptr))
+      .getPortList();
 }
 
 //===----------------------------------------------------------------------===//