feat: use non-local annotations to mark memories with extmodule

This fixes #8994. We use NLA here to
pass information from LowerMemory to LayerSink to avoid marking memories
with an extmodule as not effectful.

Signed-off-by: Tianrui Wei <[email protected]>

Author: tianrui-wei

lib/Dialect/FIRRTL/Transforms/LayerSink.cpp

@@ -13,6 +13,7 @@
 #include "circt/Dialect/FIRRTL/FIRRTLInstanceGraph.h"
 #include "circt/Dialect/FIRRTL/FIRRTLOps.h"
 #include "circt/Dialect/FIRRTL/Passes.h"
+#include "circt/Dialect/HW/HWOps.h"
 #include "circt/Support/Debug.h"
 #include "mlir/IR/Dominance.h"
 #include "mlir/IR/Threading.h"
@@ -34,6 +35,9 @@ using namespace circt;
 using namespace firrtl;
 using namespace mlir;
 
+static constexpr llvm::StringLiteral kExtMemoryEffectNLAAttrName(
+    "circt.layerSink.externalMemoryEffect");
+
 //===----------------------------------------------------------------------===//
 // Helpers
 //===----------------------------------------------------------------------===//
@@ -87,7 +91,10 @@ static bool cloneable(Operation *op) {
 namespace {
 class EffectInfo {
 public:
-  EffectInfo(CircuitOp circuit, InstanceGraph &instanceGraph) {
+  EffectInfo(CircuitOp circuit, InstanceGraph &instanceGraph,
+             DenseSet<StringAttr> explicitEffectfulModules)
+      : explicitEffectfulModules(std::move(explicitEffectfulModules)) {
+    DenseSet<InstanceGraphNode *> visited;
     instanceGraph.walkPostOrder(
         [&](auto &node) { update(node.getModule().getOperation()); });
   }
@@ -111,6 +118,7 @@ class EffectInfo {
   }
 
 private:
+  DenseSet<StringAttr> explicitEffectfulModules;
   /// Record whether the module contains any effectful ops.
   void update(FModuleOp moduleOp) {
     moduleOp.getBodyBlock()->walk([&](Operation *op) {
@@ -130,6 +138,14 @@ class EffectInfo {
       if (!annos.empty())
         return markEffectful(moduleOp);
 
+    if (explicitEffectfulModules.contains(moduleOp.getModuleNameAttr()))
+      return markEffectful(moduleOp);
+
+    // Treat generated memory modules as effectful since they model stateful
+    // hardware even though they have no body.
+    if (isa<FMemModuleOp>(moduleOp.getOperation()))
+      return markEffectful(moduleOp);
+
     for (auto annos : moduleOp.getPortAnnotations())
       if (!cast<ArrayAttr>(annos).empty())
         return markEffectful(moduleOp);
@@ -486,7 +502,19 @@ void LayerSinkPass::runOnOperation() {
                  << "\n"
                  << "Circuit: '" << circuit.getName() << "'\n";);
   auto &instanceGraph = getAnalysis<InstanceGraph>();
-  EffectInfo effectInfo(circuit, instanceGraph);
+
+  DenseSet<StringAttr> explicitEffectModules;
+  SmallVector<hw::HierPathOp, 4> effectNLAs;
+  for (auto nla : circuit.getOps<hw::HierPathOp>()) {
+    if (!nla->hasAttr(kExtMemoryEffectNLAAttrName))
+      continue;
+    if (auto leaf = nla.leafMod())
+      explicitEffectModules.insert(leaf);
+    effectNLAs.push_back(nla);
+  }
+
+  EffectInfo effectInfo(circuit, instanceGraph,
+                        std::move(explicitEffectModules));
 
   std::atomic<bool> changed(false);
   parallelForEach(&getContext(), circuit.getOps<FModuleOp>(),
@@ -495,6 +523,9 @@ void LayerSinkPass::runOnOperation() {
                       changed = true;
                   });
 
+  for (auto nla : effectNLAs)
+    nla.erase();
+
   if (!changed)
     markAllAnalysesPreserved();
   else

lib/Dialect/FIRRTL/Transforms/LowerMemory.cpp

@@ -24,6 +24,7 @@
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/STLFunctionalExtras.h"
 #include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Parallel.h"
 #include <optional>
 #include <set>
@@ -189,9 +190,11 @@ LowerMemoryPass::getMemoryModulePorts(const FirMemory &mem) {
   return ports;
 }
 
-FMemModuleOp
-LowerMemoryPass::emitMemoryModule(MemOp op, const FirMemory &mem,
-                                  const SmallVectorImpl<PortInfo> &ports) {
+static constexpr llvm::StringLiteral kExtMemoryEffectNLAAttrName(
+    "circt.layerSink.externalMemoryEffect");
+
+FMemModuleOp LowerMemoryPass::emitMemoryModule(
+    MemOp op, const FirMemory &mem, const SmallVectorImpl<PortInfo> &ports) {
   // Get a non-colliding name for the memory module, and update the summary.
   StringRef prefix = "";
   if (mem.prefix)
@@ -208,6 +211,22 @@ LowerMemoryPass::emitMemoryModule(MemOp op, const FirMemory &mem,
       mem.numReadWritePorts, mem.dataWidth, mem.maskBits, mem.readLatency,
       mem.writeLatency, mem.depth,
       *symbolizeRUWBehavior(static_cast<uint32_t>(mem.readUnderWrite)));
+  // Mark the generated external memory as effectful so downstream passes do not
+  // treat it as pure and drop the wrapper wiring around it.  We create a
+  // temporary HierPathOp, which LayerSink will consume and erase.
+  auto circuit = op->getParentOfType<CircuitOp>();
+  OpBuilder nlaBuilder(circuit);
+  nlaBuilder.setInsertionPointToEnd(circuit.getBodyBlock());
+  auto pathAttr = nlaBuilder.getArrayAttr(
+      SmallVector<Attribute, 1>{
+          FlatSymbolRefAttr::get(moduleOp.getModuleNameAttr())});
+  auto nlaName = circuitNamespace.newName(
+      (moduleOp.getModuleNameAttr().getValue() + "_effect_nla").str());
+  auto nla = nlaBuilder.create<hw::HierPathOp>(
+      mem.loc, StringAttr::get(&getContext(), nlaName), pathAttr);
+  nla->setAttr(kExtMemoryEffectNLAAttrName,
+               FlatSymbolRefAttr::get(moduleOp.getModuleNameAttr()));
+  SymbolTable::setSymbolVisibility(nla, SymbolTable::Visibility::Private);
   SymbolTable::setSymbolVisibility(moduleOp, SymbolTable::Visibility::Private);
   return moduleOp;
 }