Repl seq mem renaming (#1286)

* Consume NoDedupMemAnnotations in ResolveMemoryReference

The ComponentName being pointed to by the annotation no longer exists
after ReplaceSeqMems so we should consume the annotations

* Support renaming in ReplaceMemMacros

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

3 files changed, 112 insertions, 13 deletions

diff --git a/src/main/scala/firrtl/passes/memlib/ReplaceMemMacros.scala b/src/main/scala/firrtl/passes/memlib/ReplaceMemMacros.scala
index f81dc71b..7b3608a3 100644
--- a/src/main/scala/firrtl/passes/memlib/ReplaceMemMacros.scala
+++ b/src/main/scala/firrtl/passes/memlib/ReplaceMemMacros.scala
@@ -12,12 +12,25 @@ import MemTransformUtils._
 import firrtl.annotations._
 import wiring._
 
+import scala.collection.mutable
+
 
 /** Annotates the name of the pins to add for WiringTransform */
 case class PinAnnotation(pins: Seq[String]) extends NoTargetAnnotation
 
 object ReplaceMemMacros {
   class UnsupportedBlackboxMemoryException(msg: String) extends PassException(msg)
+
+  /** Mapping from (module, memory name) pairs to wrapper Module names */
+  private type NameMap = mutable.HashMap[(String, String), String]
+
+  /** Mutable datastructure representing mapping of smems to extracted blackboxes and their wrappers
+    *
+    * - nameMap: mapping from every (module, memory name) to mem blackbox wrapper Module name
+    * - bbMap: mapping from wrapper Module name to (blackbox instance, blackbox module)
+    * - For bbMap: instance and module name match in the code, but that could be changed
+    */
+  class MemMapping(val nameMap: NameMap, val bbMap: mutable.HashMap[String, (String, String)])
 }
 
 /** Replace DefAnnotatedMemory with memory blackbox + wrapper + conf file.
@@ -28,6 +41,8 @@ class ReplaceMemMacros(writer: ConfWriter) extends Transform {
   def inputForm = MidForm
   def outputForm = MidForm
 
+  import ReplaceMemMacros._
+
   /** Return true if mask granularity is per bit, false if per byte or unspecified
     */
   private def getFillWMask(mem: DefAnnotatedMemory) = mem.maskGran match {
@@ -197,8 +212,6 @@ class ReplaceMemMacros(writer: ConfWriter) extends Transform {
     }
   }
 
-  /** Mapping from (module, memory name) pairs to blackbox names */
-  private type NameMap = collection.mutable.HashMap[(String, String), String]
   /** Construct NameMap by assigning unique names for each memory blackbox */
   def constructNameMap(namespace: Namespace, nameMap: NameMap, mname: String)(s: Statement): Statement = {
     s match {
@@ -211,12 +224,23 @@ class ReplaceMemMacros(writer: ConfWriter) extends Transform {
     s map constructNameMap(namespace, nameMap, mname)
   }
 
+  // For 1.2.x backwards compatibility
   def updateMemStmts(namespace: Namespace,
                      nameMap: NameMap,
                      mname: String,
                      memPortMap: MemPortMap,
-                     memMods: Modules)
-                     (s: Statement): Statement = s match {
+                     memMods: Modules
+                    )(s: Statement): Statement =
+    updateMemStmts(namespace, nameMap, mname, memPortMap, memMods, None)(s)
+
+  // memMapping is only Option for backwards compatibility
+  def updateMemStmts(namespace: Namespace,
+                     nameMap: NameMap,
+                     mname: String,
+                     memPortMap: MemPortMap,
+                     memMods: Modules,
+                     memMapping: Option[MemMapping]
+                    )(s: Statement): Statement = s match {
     case m: DefAnnotatedMemory =>
       if (m.maskGran.isEmpty) {
         m.writers foreach { w => memPortMap(s"${m.name}.$w.mask") = EmptyExpression }
@@ -228,18 +252,39 @@ class ReplaceMemMacros(writer: ConfWriter) extends Transform {
           val newWrapperName = nameMap(mname -> m.name)
           val newMemBBName = namespace newName s"${newWrapperName}_ext"
           val newMem = m copy (name = newMemBBName)
+          // Record for annotation renaming
+          memMapping.foreach { mapping =>
+            mapping.nameMap += ((mname, m.name) -> newWrapperName)
+            mapping.bbMap += newWrapperName -> (newMemBBName, newMemBBName)
+          }
           memMods ++= createMemModule(newMem, newWrapperName)
           WDefInstance(m.info, m.name, newWrapperName, UnknownType)
         case Some((module, mem)) =>
-          WDefInstance(m.info, m.name, nameMap(module -> mem), UnknownType)
+          val wrapperName = nameMap(module -> mem)
+          // Record for annotation renaming
+          memMapping.foreach(_.nameMap += ((mname, m.name) -> wrapperName))
+          WDefInstance(m.info, m.name, wrapperName, UnknownType)
       }
-    case sx => sx map updateMemStmts(namespace, nameMap, mname, memPortMap, memMods)
+    case sx => sx map updateMemStmts(namespace, nameMap, mname, memPortMap, memMods, memMapping)
   }
 
-  def updateMemMods(namespace: Namespace, nameMap: NameMap, memMods: Modules)(m: DefModule) = {
+
+  // For 1.2.x backwards compatibility
+  def updateMemMods(namespace: Namespace,
+                    nameMap: NameMap,
+                    memMods: Modules
+                   )(m: DefModule): DefModule =
+    updateMemMods(namespace, nameMap, memMods, None)(m)
+
+  // memMapping is only Option for backwards compatibility
+  def updateMemMods(namespace: Namespace,
+                    nameMap: NameMap,
+                    memMods: Modules,
+                    memMapping: Option[MemMapping]
+                   )(m: DefModule): DefModule = {
     val memPortMap = new MemPortMap
 
-    (m map updateMemStmts(namespace, nameMap, m.name, memPortMap, memMods)
+    (m map updateMemStmts(namespace, nameMap, m.name, memPortMap, memMods, memMapping)
        map updateStmtRefs(memPortMap))
   }
 
@@ -249,7 +294,18 @@ class ReplaceMemMacros(writer: ConfWriter) extends Transform {
     val memMods = new Modules
     val nameMap = new NameMap
     c.modules map (m => m map constructNameMap(namespace, nameMap, m.name))
-    val modules = c.modules map updateMemMods(namespace, nameMap, memMods)
+    val memMapping = new MemMapping(new NameMap, new mutable.HashMap)
+    val modules = c.modules map updateMemMods(namespace, nameMap, memMods, Some(memMapping))
+
+    // Rename replaced memories with new blackbox
+    val renames = RenameMap.create {
+      val top = CircuitTarget(c.main)
+      memMapping.nameMap.map { case ((mod, mem), wrap) =>
+        val (bbInst, bbMod) = memMapping.bbMap(wrap)
+        top.module(mod).ref(mem) -> Seq(top.module(mod).instOf(mem, wrap).instOf(bbInst, bbMod))
+      }
+    }
+
     // print conf
     writer.serialize()
     val pannos = state.annotations.collect { case a: PinAnnotation => a }
@@ -263,6 +319,6 @@ class ReplaceMemMacros(writer: ConfWriter) extends Transform {
         case m: ExtModule => SinkAnnotation(ModuleName(m.name, CircuitName(c.main)), pin)
       }
     } ++ state.annotations
-    CircuitState(c.copy(modules = modules ++ memMods), inputForm, annos)
+    CircuitState(c.copy(modules = modules ++ memMods), inputForm, annos, renames = Some(renames))
   }
 }
diff --git a/src/main/scala/firrtl/passes/memlib/ResolveMemoryReference.scala b/src/main/scala/firrtl/passes/memlib/ResolveMemoryReference.scala
index b0d3731f..094c2929 100644
--- a/src/main/scala/firrtl/passes/memlib/ResolveMemoryReference.scala
+++ b/src/main/scala/firrtl/passes/memlib/ResolveMemoryReference.scala
@@ -68,7 +68,8 @@ class ResolveMemoryReference extends Transform {
     val noDedups = state.annotations.collect {
       case NoDedupMemAnnotation(ComponentName(cn, ModuleName(mn, _))) => mn -> cn
     }
+    val annos = state.annotations.filterNot(_.isInstanceOf[NoDedupMemAnnotation])
     val noDedupMap: Map[String, Set[String]] = noDedups.groupBy(_._1).mapValues(_.map(_._2).toSet)
-    state.copy(circuit = run(state.circuit, noDedupMap))
+    state.copy(circuit = run(state.circuit, noDedupMap), annotations = annos)
   }
 }
diff --git a/src/test/scala/firrtlTests/ReplSeqMemTests.scala b/src/test/scala/firrtlTests/ReplSeqMemTests.scala
index 72171d43..6dd52aa5 100644
--- a/src/test/scala/firrtlTests/ReplSeqMemTests.scala
+++ b/src/test/scala/firrtlTests/ReplSeqMemTests.scala
@@ -11,6 +11,11 @@ import firrtl.FileUtils
 import annotations._
 import FirrtlCheckers._
 
+// It's not clear if this should be IsComponent or IsMember
+case class MemAnnotation(target: IsComponent) extends SingleTargetAnnotation[IsComponent] {
+  def duplicate(n: IsComponent): MemAnnotation = this.copy(n)
+}
+
 class ReplSeqMemSpec extends SimpleTransformSpec {
   def emitter = new LowFirrtlEmitter
   def transforms = Seq(
@@ -269,10 +274,15 @@ circuit CustomMemory :
       MemConf("mem_0_ext", 7, 16, Map(WritePort -> 1, ReadPort -> 1), None),
       MemConf("mem_1_ext", 7, 16, Map(WritePort -> 1, ReadPort -> 1), None)
     )
+    val mod = CircuitTarget("CustomMemory").module("CustomMemory")
     val confLoc = "ReplSeqMemTests.confTEMP"
     val annos = Seq(
       ReplSeqMemAnnotation.parse("-c:CustomMemory:-o:"+confLoc),
-      NoDedupMemAnnotation(ComponentName("mem_1", ModuleName("CustomMemory",CircuitName("CustomMemory")))))
+      NoDedupMemAnnotation(ComponentName("mem_1", ModuleName("CustomMemory",CircuitName("CustomMemory")))),
+      MemAnnotation(mod.ref("mem_0")),
+      MemAnnotation(mod.ref("mem_1")),
+      MemAnnotation(mod.ref("mem_2"))
+    )
     val res = compileAndEmit(CircuitState(parse(input), ChirrtlForm, annos))
     // Check correctness of firrtl
     val circuit = parse(res.getEmittedCircuit.value)
@@ -283,6 +293,14 @@ circuit CustomMemory :
     numExtMods should be (2)
     // Check the emitted conf
     checkMemConf(confLoc, mems)
+    // Check annotation renaming
+    val expectedTargets = Seq(
+      mod.instOf("mem_0", "mem_0").instOf("mem_0_ext", "mem_0_ext"),
+      mod.instOf("mem_1", "mem_1").instOf("mem_1_ext", "mem_1_ext"),
+      mod.instOf("mem_2", "mem_0").instOf("mem_0_ext", "mem_0_ext")
+    )
+    res.annotations.collect { case MemAnnotation(t) => t } should equal (expectedTargets)
+
     (new java.io.File(confLoc)).delete()
   }
 
@@ -522,5 +540,29 @@ circuit Top :
     }
   }
 
+  "ReplSeqMem" should "rename annotations" in {
+    val input =
+      """|circuit CustomMemory :
+         |  module CustomMemory :
+         |    input clock : Clock
+         |    output io : { flip en : UInt<1>, out : UInt<8>[2], flip raddr : UInt<10>, flip waddr : UInt<10>, flip wdata : UInt<8>[2] }
+         |
+         |    smem mem : UInt<8>[2][1024]
+         |    read mport r = mem[io.raddr], clock
+         |    io.out <= r
+         |
+         |    when io.en :
+         |      write mport w = mem[io.waddr], clock
+         |      w <= io.wdata
+         |""".stripMargin
+
+    val mod = CircuitTarget("CustomMemory").module("CustomMemory")
+    val anno = MemAnnotation(mod.ref("mem"))
+    val confLoc = "ReplSeqMemTests.confTEMP"
+    val annos = Seq(ReplSeqMemAnnotation.parse("-c:CustomMemory:-o:"+confLoc), anno)
+    val res = compileAndEmit(CircuitState(parse(input), ChirrtlForm, annos))
+    val expectedTarget = mod.instOf("mem", "mem").instOf("mem_ext", "mem_ext")
+    res.annotations.collect { case MemAnnotation(t) => t } should equal (Seq(expectedTarget))
+    (new java.io.File(confLoc)).delete()
+  }
 }
-