From 7aee9c82a25042350baab6a45e87ebc5f5b21056 Mon Sep 17 00:00:00 2001
From: Robert Young <rwy0717@gmail.com>
Date: Fri, 7 Mar 2025 15:41:25 -0800
Subject: [PATCH] [LowerLayers][FIRRTL] Add the ability to sink ops by cloning

cloning
- constant op
- special constant op
- subfield op
- subindex op

(Not cloning subaccess ops.)
---
 .../FIRRTL/Transforms/AdvancedLayerSink.cpp   | 62 +++++++++++++++----
 1 file changed, 51 insertions(+), 11 deletions(-)

diff --git a/lib/Dialect/FIRRTL/Transforms/AdvancedLayerSink.cpp b/lib/Dialect/FIRRTL/Transforms/AdvancedLayerSink.cpp
index af3f16de78e1..803b608a758e 100644
--- a/lib/Dialect/FIRRTL/Transforms/AdvancedLayerSink.cpp
+++ b/lib/Dialect/FIRRTL/Transforms/AdvancedLayerSink.cpp
@@ -63,6 +63,14 @@ static bool isAncestor(Block *block, Block *other) {
   return block->getParent()->isProperAncestor(other->getParent());
 }
 
+//===----------------------------------------------------------------------===//
+// Clone-ability.
+//===----------------------------------------------------------------------===//
+
+static bool cloneable(Operation *op) {
+  return isa<ConstantOp, SpecialConstantOp, SubfieldOp, SubindexOp>(op);
+}
+
 //===----------------------------------------------------------------------===//
 // Effectfulness Analysis.
 //===----------------------------------------------------------------------===//
@@ -357,6 +365,10 @@ class ModuleLayerSink {
   void moveLayersToBack(Operation *op);
   void moveLayersToBack(Block *block);
 
+  void erase(Operation *op);
+  void sinkOpByCloning(Operation *op);
+  void sinkOpByMoving(Operation *op, Block *dest);
+
   FModuleOp moduleOp;
   const EffectInfo &effectInfo;
   bool changed = false;
@@ -403,22 +415,50 @@ void ModuleLayerSink::moveLayersToBack(Block *block) {
   }
 }
 
+void ModuleLayerSink::erase(Operation *op) {
+  op->erase();
+  changed = true;
+}
+
+void ModuleLayerSink::sinkOpByCloning(Operation *op) {
+  // Copy the op down to each block where there is a user.
+  // Use a cache to ensure we don't create more than one copy per block.
+  DenseMap<Block *, Operation *> cache;
+  for (unsigned i = 0, e = op->getNumResults(); i < e; ++i) {
+    for (auto &use : llvm::make_early_inc_range(op->getResult(i).getUses())) {
+      auto *dest = use.getOwner()->getBlock();
+      if (dest != op->getBlock()) {
+        auto &clone = cache[dest];
+        if (!clone)
+          clone = OpBuilder::atBlockBegin(dest).clone(*op);
+        use.set(clone->getResult(i));
+        changed = true;
+      }
+    }
+  }
+
+  // If there are no remaining uses of the original op, erase it.
+  if (op->use_empty())
+    erase(op);
+}
+
+void ModuleLayerSink::sinkOpByMoving(Operation *op, Block *dest) {
+  if (dest != op->getBlock()) {
+    op->moveBefore(dest, dest->begin());
+    changed = true;
+  }
+}
+
 bool ModuleLayerSink::operator()() {
   moveLayersToBack(moduleOp.getBodyBlock());
   DemandInfo demandInfo(effectInfo, moduleOp);
   walkBwd(moduleOp.getBodyBlock(), [&](Operation *op) {
     auto demand = demandInfo.getDemandFor(op);
-    if (!demand) {
-      op->erase();
-      changed = true;
-      return;
-    }
-
-    if (demand == op->getBlock())
-      return;
-
-    op->moveBefore(demand.block, demand.block->begin());
-    changed = true;
+    if (!demand)
+      return erase(op);
+    if (cloneable(op))
+      return sinkOpByCloning(op);
+    sinkOpByMoving(op, demand.block);
   });
 
   return changed;