Fix memory annotation deduplication (#2286)

* Add transform to deduplicate memory annotations
* Add annotation deduplication to Dedup stage
* ResolveAnnotationPaths and EliminateTargetPaths now invalidate the dedup annotations transform
* Verilog emitter now throws exception when memory annotations fail to dedup

Co-authored-by: Jack Koenig <koenig@sifive.com>

9 files changed, 173 insertions, 3 deletions

diff --git a/src/main/scala/firrtl/Compiler.scala b/src/main/scala/firrtl/Compiler.scala
index 7354b8ee..2017633d 100644
--- a/src/main/scala/firrtl/Compiler.scala
+++ b/src/main/scala/firrtl/Compiler.scala
@@ -14,6 +14,7 @@ import firrtl.Utils.throwInternalError
 import firrtl.annotations.transforms.{EliminateTargetPaths, ResolvePaths}
 import firrtl.options.{Dependency, DependencyAPI, StageUtils, TransformLike}
 import firrtl.stage.Forms
+import firrtl.transforms.DedupAnnotationsTransform
 
 /** Container of all annotations for a Firrtl compiler */
 class AnnotationSeq private (private[firrtl] val underlying: List[Annotation]) {
@@ -396,6 +397,9 @@ trait ResolvedAnnotationPaths {
   override def prepare(state: CircuitState): CircuitState = {
     state.resolvePathsOf(annotationClasses.toSeq: _*)
   }
+
+  // Any transform with this trait invalidates DedupAnnotationsTransform
+  override def invalidates(a: Transform) = a.isInstanceOf[DedupAnnotationsTransform]
 }
 
 /** Defines old API for Emission. Deprecated */
diff --git a/src/main/scala/firrtl/Utils.scala b/src/main/scala/firrtl/Utils.scala
index a58b6997..e29c1a3b 100644
--- a/src/main/scala/firrtl/Utils.scala
+++ b/src/main/scala/firrtl/Utils.scala
@@ -355,6 +355,7 @@ object Utils extends LazyLogging {
     onExp(expression)
     ReferenceTarget(main, module, Nil, ref, tokens.toSeq)
   }
+
   @deprecated("get_flip is fundamentally slow, use to_flip(flow(expr))", "FIRRTL 1.2")
   def get_flip(t: Type, i: Int, f: Orientation): Orientation = {
     if (i >= get_size(t)) throwInternalError(s"get_flip: shouldn't be here - $i >= get_size($t)")
diff --git a/src/main/scala/firrtl/annotations/Annotation.scala b/src/main/scala/firrtl/annotations/Annotation.scala
index 08555a84..f4d6ee55 100644
--- a/src/main/scala/firrtl/annotations/Annotation.scala
+++ b/src/main/scala/firrtl/annotations/Annotation.scala
@@ -38,6 +38,21 @@ trait Annotation extends Product {
     * @return
     */
   def getTargets: Seq[Target] = extractComponents(productIterator.toIterable).toSeq
+
+  /** Returns a deduplicable representation of this [[Annotation]]: a 3-tuple of the
+    * deduplicated annotation's "dedup key", the deduplicated [[Annotation]], and the
+    * [[firrtl.annotations.ReferenceTarget ReferenceTarget]](s) to the annotated objects.
+    *
+    * If two absolute instances of this [[Annotation]] would deduplicate to the same
+    * local form, both of their "dedup key"s must be equivalent.
+    *
+    * A deduplication key is typically taken to be a 2-tuple of the pathless target and
+    * the annotation's value.
+    *
+    * Returning None signifies this annotation will not deduplicate.
+    * @return
+    */
+  private[firrtl] def dedup: Option[(Any, Annotation, ReferenceTarget)] = None
 }
 
 /** If an Annotation does not target any [[Named]] thing in the circuit, then all updates just
diff --git a/src/main/scala/firrtl/annotations/MemoryInitAnnotation.scala b/src/main/scala/firrtl/annotations/MemoryInitAnnotation.scala
index 1e81301d..e87066fb 100644
--- a/src/main/scala/firrtl/annotations/MemoryInitAnnotation.scala
+++ b/src/main/scala/firrtl/annotations/MemoryInitAnnotation.scala
@@ -8,7 +8,8 @@ import firrtl.{
   MemoryFileInlineInit,
   MemoryInitValue,
   MemoryRandomInit,
-  MemoryScalarInit
+  MemoryScalarInit,
+  Utils
 }
 
 /**
@@ -25,6 +26,9 @@ case class MemoryRandomInitAnnotation(target: ReferenceTarget) extends MemoryIni
   override def duplicate(n: ReferenceTarget): Annotation = copy(n)
   override def initValue:    MemoryInitValue = MemoryRandomInit
   override def isRandomInit: Boolean = true
+  override private[firrtl] def dedup: Option[(Any, Annotation, ReferenceTarget)] = Some(
+    ((target.pathlessTarget, Nil), copy(target = target.pathlessTarget), target)
+  )
 }
 
 /** Initialize all entries of the `target` memory with the scalar `value`. */
@@ -32,6 +36,9 @@ case class MemoryScalarInitAnnotation(target: ReferenceTarget, value: BigInt) ex
   override def duplicate(n: ReferenceTarget): Annotation = copy(n)
   override def initValue:    MemoryInitValue = MemoryScalarInit(value)
   override def isRandomInit: Boolean = false
+  override private[firrtl] def dedup: Option[(Any, Annotation, ReferenceTarget)] = Some(
+    ((target.pathlessTarget, value), copy(target = target.pathlessTarget), target)
+  )
 }
 
 /** Initialize the `target` memory with the array of `values` which must be the same size as the memory depth. */
@@ -39,6 +46,9 @@ case class MemoryArrayInitAnnotation(target: ReferenceTarget, values: Seq[BigInt
   override def duplicate(n: ReferenceTarget): Annotation = copy(n)
   override def initValue:    MemoryInitValue = MemoryArrayInit(values)
   override def isRandomInit: Boolean = false
+  override private[firrtl] def dedup: Option[(Any, Annotation, ReferenceTarget)] = Some(
+    ((target.pathlessTarget, values), copy(target = target.pathlessTarget), target)
+  )
 }
 
 /** Initialize the `target` memory with inline readmem[hb] statement. */
@@ -51,6 +61,9 @@ case class MemoryFileInlineAnnotation(
   override def duplicate(n: ReferenceTarget): Annotation = copy(n)
   override def initValue:    MemoryInitValue = MemoryFileInlineInit(filename, hexOrBinary)
   override def isRandomInit: Boolean = false
+  override private[firrtl] def dedup: Option[(Any, Annotation, ReferenceTarget)] = Some(
+    ((target.pathlessTarget, filename), copy(target = target.pathlessTarget), target)
+  )
 }
 
 /** Initializes the memory inside the `ifndef SYNTHESIS` block (default) */
diff --git a/src/main/scala/firrtl/annotations/Target.scala b/src/main/scala/firrtl/annotations/Target.scala
index 92339946..02ec42b8 100644
--- a/src/main/scala/firrtl/annotations/Target.scala
+++ b/src/main/scala/firrtl/annotations/Target.scala
@@ -695,6 +695,15 @@ case class ReferenceTarget(
     }
   }
 
+  /** Returns the local form of this [[ReferenceTarget]]
+    *
+    * For example, given `~Top|Top/foo:Foo/bar:Bar>x`,
+    *
+    * `.pathlessTarget` returns `~Top|Bar>x`
+    *
+    * This is useful for cases in which annotations must point to the module itself rather than
+    *   an absolute *instance* of the module (e.g. deduplication).
+    */
   override def pathlessTarget: ReferenceTarget = ReferenceTarget(circuit, encapsulatingModule, Nil, ref, component)
 
   override def setPathTarget(newPath: IsModule): ReferenceTarget =
@@ -789,6 +798,15 @@ case class InstanceTarget(
 
   override def asPath: Seq[(Instance, OfModule)] = path :+ ((Instance(instance), OfModule(ofModule)))
 
+  /** Returns the local form of this [[InstanceTarget]]
+    *
+    * For example, given `~Top|Top/foo:Foo/bar:Bar`,
+    *
+    * `.pathlessTarget` returns `~Top|Foo/bar:Bar`
+    *
+    * This is useful for cases in which annotations must point to the module itself rather than
+    *   an absolute *instance* of the module (e.g. deduplication).
+    */
   override def pathlessTarget: InstanceTarget = InstanceTarget(circuit, encapsulatingModule, Nil, instance, ofModule)
 
   override def notPath = Seq(Instance(instance), OfModule(ofModule))
diff --git a/src/main/scala/firrtl/annotations/transforms/EliminateTargetPaths.scala b/src/main/scala/firrtl/annotations/transforms/EliminateTargetPaths.scala
index b63dd13e..104aafc3 100644
--- a/src/main/scala/firrtl/annotations/transforms/EliminateTargetPaths.scala
+++ b/src/main/scala/firrtl/annotations/transforms/EliminateTargetPaths.scala
@@ -12,6 +12,7 @@ import firrtl.ir._
 import firrtl.{AnnotationSeq, CircuitState, DependencyAPIMigration, FirrtlInternalException, RenameMap, Transform}
 import firrtl.stage.Forms
 import firrtl.transforms.DedupedResult
+import firrtl.transforms.DedupAnnotationsTransform
 
 import scala.collection.mutable
 
@@ -105,7 +106,7 @@ class EliminateTargetPaths extends Transform with DependencyAPIMigration {
   override def prerequisites = Forms.MinimalHighForm
   override def optionalPrerequisites = Seq.empty
   override def optionalPrerequisiteOf = Seq.empty
-  override def invalidates(a: Transform) = false
+  override def invalidates(a: Transform) = a.isInstanceOf[DedupAnnotationsTransform]
 
   /** Replaces old ofModules with new ofModules by calling dupMap methods
     * Updates oldUsedOfModules, newUsedOfModules
diff --git a/src/main/scala/firrtl/backends/verilog/VerilogEmitter.scala b/src/main/scala/firrtl/backends/verilog/VerilogEmitter.scala
index 13f7b793..8b20d365 100644
--- a/src/main/scala/firrtl/backends/verilog/VerilogEmitter.scala
+++ b/src/main/scala/firrtl/backends/verilog/VerilogEmitter.scala
@@ -9,6 +9,7 @@ import firrtl.WrappedExpression._
 import firrtl.traversals.Foreachers._
 import firrtl.annotations.{
   CircuitTarget,
+  MemoryInitAnnotation,
   MemoryLoadFileType,
   MemoryNoSynthInit,
   MemorySynthInit,
@@ -507,6 +508,18 @@ class VerilogEmitter extends SeqTransform with Emitter {
     private val emissionAnnos = annotations.collect {
       case m: SingleTargetAnnotation[ReferenceTarget] @unchecked with EmissionOption => m
     }
+
+    // Check for non-local memory annotations (error if found)
+    emissionAnnos.foreach {
+      case a: MemoryInitAnnotation => {
+        if (!a.target.isLocal)
+          throw new FirrtlUserException(
+            "At least one memory annotation did not deduplicate: got non-local annotation $a from [[DedupAnnotationsTransform]]"
+          )
+      }
+      case _ =>
+    }
+
     // using multiple foreach instead of a single partial function as an Annotation can gather multiple EmissionOptions for simplicity
     emissionAnnos.foreach {
       case a: MemoryEmissionOption => memoryEmissionOption += ((a.target, a))
diff --git a/src/main/scala/firrtl/stage/Forms.scala b/src/main/scala/firrtl/stage/Forms.scala
index 4132f758..c7ae648a 100644
--- a/src/main/scala/firrtl/stage/Forms.scala
+++ b/src/main/scala/firrtl/stage/Forms.scala
@@ -47,7 +47,11 @@ object Forms {
       Dependency[firrtl.transforms.InferResets]
     )
 
-  val Deduped: Seq[TransformDependency] = Resolved :+ Dependency[firrtl.transforms.DedupModules]
+  val Deduped: Seq[TransformDependency] = Resolved ++
+    Seq(
+      Dependency[firrtl.transforms.DedupModules],
+      Dependency[firrtl.transforms.DedupAnnotationsTransform]
+    )
 
   val HighForm: Seq[TransformDependency] = ChirrtlForm ++
     MinimalHighForm ++
diff --git a/src/main/scala/firrtl/transforms/DedupAnnotations.scala b/src/main/scala/firrtl/transforms/DedupAnnotations.scala
new file mode 100644
index 00000000..9355b5c3
--- /dev/null
+++ b/src/main/scala/firrtl/transforms/DedupAnnotations.scala
@@ -0,0 +1,101 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package firrtl
+package transforms
+
+import firrtl.ir._
+import firrtl.Mappers._
+import firrtl.options.Dependency
+import firrtl.Utils.BoolType
+import firrtl.annotations.Annotation
+import scala.collection.mutable.Buffer
+import firrtl.annotations.MemoryFileInlineAnnotation
+import firrtl.passes.PassException
+import firrtl.annotations.ReferenceTarget
+import firrtl.annotations._
+import firrtl.analyses.InstanceKeyGraph
+
+import scala.collection.mutable.ArrayBuffer
+
+object DedupAnnotationsTransform {
+
+  final class DifferingModuleAnnotationsException private (msg: String) extends PassException(msg)
+  object DifferingModuleAnnotationsException {
+    def apply(left: ReferenceTarget, right: ReferenceTarget): DifferingModuleAnnotationsException = {
+      val msg = s"${left.serialize} and ${right.serialize} have differing module binaries"
+      new DifferingModuleAnnotationsException(msg)
+    }
+  }
+
+  private case class DedupableRepr(
+    dedupKey:       Any,
+    deduped:        Annotation,
+    original:       Annotation,
+    absoluteTarget: ReferenceTarget)
+  private object DedupableRepr {
+    def apply(annotation: Annotation): Option[DedupableRepr] = annotation.dedup match {
+      case Some((dedupKey, dedupedAnno, absoluteTarget)) =>
+        Some(new DedupableRepr(dedupKey, dedupedAnno, annotation, absoluteTarget))
+      case _ => None
+    }
+  }
+
+  private type InstancePath = Seq[(TargetToken.Instance, TargetToken.OfModule)]
+
+  private def checkInstanceGraph(
+    module:        String,
+    graph:         InstanceKeyGraph,
+    absolutePaths: Seq[InstancePath]
+  ): Boolean = graph.findInstancesInHierarchy(module).size == absolutePaths.size
+
+  def dedupAnnotations(annotations: Seq[Annotation], graph: InstanceKeyGraph): Seq[Annotation] = {
+    val canDedup = ArrayBuffer.empty[DedupableRepr]
+    val outAnnos = ArrayBuffer.empty[Annotation]
+
+    // Extract the annotations which can be deduplicated
+    annotations.foreach { anno =>
+      DedupableRepr(anno) match {
+        case Some(repr) => canDedup += repr
+        case None       => outAnnos += anno
+      }
+    }
+
+    // Partition the dedupable annotations into groups that *should* deduplicate into the same annotation
+    val shouldDedup: Map[Any, ArrayBuffer[DedupableRepr]] = canDedup.groupBy(_.dedupKey)
+    shouldDedup.foreach {
+      case ((target: ReferenceTarget, _), dedupableAnnos) =>
+        val originalAnnos = dedupableAnnos.map(_.original)
+        val uniqueDedupedAnnos = dedupableAnnos.map(_.deduped).distinct
+        // TODO: Extend this to support multi-target annotations
+        val instancePaths = dedupableAnnos.map(_.absoluteTarget.path).toSeq
+        // The annotation deduplication is only legal if it applies to *all* instances of a
+        // deduplicated module -- requires an instance graph check
+        if (uniqueDedupedAnnos.size == 1 && checkInstanceGraph(target.encapsulatingModule, graph, instancePaths))
+          outAnnos += uniqueDedupedAnnos.head
+        else
+          outAnnos ++= originalAnnos
+    }
+
+    outAnnos.toSeq
+  }
+}
+
+/** Deduplicates memory annotations
+  */
+class DedupAnnotationsTransform extends Transform with DependencyAPIMigration {
+
+  override def prerequisites = Nil
+
+  override def optionalPrerequisites = Nil
+
+  override def optionalPrerequisiteOf = Nil
+
+  override def invalidates(a: Transform) = false
+
+  def execute(state: CircuitState): CircuitState = CircuitState(
+    state.circuit,
+    state.form,
+    DedupAnnotationsTransform.dedupAnnotations(state.annotations.underlying, InstanceKeyGraph(state.circuit)),
+    state.renames
+  )
+}