Improve performance of RenameMap in LowerTypes (#2233)

LowerTypes creates a lot of mappings for the RenameMap. The built-in
.distinct of renames becomes a performance program for designs with
deeply nested Aggregates. Because LowerTypes does not create duplicate
renames, it can safely eschew the safety of using .distinct via a
private internal API.

2 files changed, 20 insertions, 3 deletions

diff --git a/src/main/scala/firrtl/RenameMap.scala b/src/main/scala/firrtl/RenameMap.scala
index 82c00ca5..d39d8106 100644
--- a/src/main/scala/firrtl/RenameMap.scala
+++ b/src/main/scala/firrtl/RenameMap.scala
@@ -78,6 +78,11 @@ object RenameMap {
   /** Initialize a new RenameMap */
   def apply(): RenameMap = new RenameMap
 
+  // This is a private internal API for transforms where the .distinct operation is very expensive
+  // (eg. LowerTypes). The onus is on the user of this API to be very careful and not inject
+  // duplicates. This is a bad, hacky API that no one should use
+  private[firrtl] def noDistinct(): RenameMap = new RenameMap(doDistinct = false)
+
   abstract class RenameTargetException(reason: String) extends Exception(reason)
   case class IllegalRenameException(reason: String) extends RenameTargetException(reason)
   case class CircularRenameException(reason: String) extends RenameTargetException(reason)
@@ -94,7 +99,11 @@ object RenameMap {
 final class RenameMap private (
   val underlying: mutable.HashMap[CompleteTarget, Seq[CompleteTarget]] =
     mutable.HashMap[CompleteTarget, Seq[CompleteTarget]](),
-  val chained: Option[RenameMap] = None) {
+  val chained: Option[RenameMap] = None,
+  // This is a private internal API for transforms where the .distinct operation is very expensive
+  // (eg. LowerTypes). The onus is on the user of this API to be very careful and not inject
+  // duplicates. This is a bad, hacky API that no one should use
+  doDistinct: Boolean = true) {
 
   /** Chain a [[RenameMap]] with this [[RenameMap]]
     * @param next the map to chain with this map
@@ -662,7 +671,10 @@ final class RenameMap private (
   private def completeRename(from: CompleteTarget, tos: Seq[CompleteTarget]): Unit = {
     tos.foreach { recordSensitivity(from, _) }
     val existing = underlying.getOrElse(from, Vector.empty)
-    val updated = (existing ++ tos).distinct
+    val updated = {
+      val all = (existing ++ tos)
+      if (doDistinct) all.distinct else all
+    }
     underlying(from) = updated
     traverseTokensCache.clear()
     traverseHierarchyCache.clear()
diff --git a/src/main/scala/firrtl/passes/LowerTypes.scala b/src/main/scala/firrtl/passes/LowerTypes.scala
index 0bd44a8c..7ba320d0 100644
--- a/src/main/scala/firrtl/passes/LowerTypes.scala
+++ b/src/main/scala/firrtl/passes/LowerTypes.scala
@@ -75,7 +75,12 @@ object LowerTypes extends Transform with DependencyAPIMigration {
     val memInitByModule = memInitAnnos.map(_.asInstanceOf[MemoryInitAnnotation]).groupBy(_.target.encapsulatingModule)
 
     val c = CircuitTarget(state.circuit.main)
-    val refRenameMap = RenameMap()
+    // By default, the RenameMap enforces a .distinct invariant for renames. This helps transform
+    // writers not mess up because violating that invariant can cause problems for transform
+    // writers. Unfortunately, when you have lots of renames, this is very expensive
+    // performance-wise. We use a private internal API that does not run .distinct to improve
+    // performance, but we must be careful to not insert any duplicates.
+    val refRenameMap = RenameMap.noDistinct()
     val resultAndRenames =
       state.circuit.modules.map(m => onModule(c, m, memInitByModule.getOrElse(m.name, Seq()), refRenameMap))
     val result = state.circuit.copy(modules = resultAndRenames.map(_._1))