Pass optional name in ImportDefinitionAnno (#2592) (#2594)

Used for separate elaboration of Definition and Instance

(cherry picked from commit 48d57cc8db6f38fdf0e23b7dce36caa404c871b8)

Co-authored-by: Girish Pai <girish.pai@sifive.com>

4 files changed, 182 insertions, 14 deletions

diff --git a/core/src/main/scala/chisel3/experimental/hierarchy/Definition.scala b/core/src/main/scala/chisel3/experimental/hierarchy/Definition.scala
index c79c26dc..36bf6f87 100644
--- a/core/src/main/scala/chisel3/experimental/hierarchy/Definition.scala
+++ b/core/src/main/scala/chisel3/experimental/hierarchy/Definition.scala
@@ -120,5 +120,7 @@ object Definition extends SourceInfoDoc {
 /** Stores a [[Definition]] that is imported so that its Instances can be
   * compiled separately.
   */
-case class ImportDefinitionAnnotation[T <: BaseModule with IsInstantiable](definition: Definition[T])
+case class ImportDefinitionAnnotation[T <: BaseModule with IsInstantiable](
+  definition:      Definition[T],
+  overrideDefName: Option[String] = None)
     extends NoTargetAnnotation
diff --git a/core/src/main/scala/chisel3/experimental/hierarchy/Instance.scala b/core/src/main/scala/chisel3/experimental/hierarchy/Instance.scala
index 9f96dff1..861023a1 100644
--- a/core/src/main/scala/chisel3/experimental/hierarchy/Instance.scala
+++ b/core/src/main/scala/chisel3/experimental/hierarchy/Instance.scala
@@ -10,6 +10,7 @@ import chisel3.internal.sourceinfo.{InstanceTransform, SourceInfo}
 import chisel3.experimental.{BaseModule, ExtModule}
 import chisel3.internal.firrtl.{Component, DefBlackBox, DefModule, Port}
 import firrtl.annotations.IsModule
+import chisel3.internal.throwException
 
 /** User-facing Instance type.
   * Represents a unique instance of type [[A]] which are marked as @instantiable
@@ -118,8 +119,14 @@ object Instance extends SourceInfoDoc {
     if (existingMod.isEmpty) {
       // Add a Definition that will get emitted as an ExtModule so that FIRRTL
       // does not complain about a missing element
+      val extModName = Builder.importDefinitionMap.getOrElse(
+        definition.proto.name,
+        throwException(
+          "Imported Definition information not found - possibly forgot to add ImportDefinition annotation?"
+        )
+      )
       class EmptyExtModule extends ExtModule {
-        override def desiredName: String = definition.proto.name
+        override def desiredName: String = extModName
         override def generateComponent(): Option[Component] = {
           require(!_closed, s"Can't generate $desiredName module more than once")
           _closed = true
diff --git a/core/src/main/scala/chisel3/internal/Builder.scala b/core/src/main/scala/chisel3/internal/Builder.scala
index 6fd9bdd5..35dd01ab 100644
--- a/core/src/main/scala/chisel3/internal/Builder.scala
+++ b/core/src/main/scala/chisel3/internal/Builder.scala
@@ -367,11 +367,30 @@ private[chisel3] class DynamicContext(
   val warnReflectiveNaming: Boolean) {
   val importDefinitionAnnos = annotationSeq.collect { case a: ImportDefinitionAnnotation[_] => a }
 
-  // Ensure there are no repeated names for imported Definitions
-  val importDefinitionNames = importDefinitionAnnos.map { a => a.definition.proto.name }
-  if (importDefinitionNames.distinct.length < importDefinitionNames.length) {
-    val duplicates = importDefinitionNames.diff(importDefinitionNames.distinct).mkString(", ")
-    throwException(s"Expected distinct imported Definition names but found duplicates for: $duplicates")
+  // Map holding the actual names of extModules
+  // Pick the definition name by default in case not passed through annotation.
+  val importDefinitionMap = importDefinitionAnnos
+    .map(a => a.definition.proto.name -> a.overrideDefName.getOrElse(a.definition.proto.name))
+    .toMap
+
+  // Helper function which does 2 things
+  // 1. Ensure there are no repeated names for imported Definitions - both Proto Names as well as ExtMod Names
+  // 2. Return the distinct definition / extMod names
+  private def checkAndGeDistinctProtoExtModNames() = {
+    val importAllDefinitionProtoNames = importDefinitionAnnos.map { a => a.definition.proto.name }
+    val importDistinctDefinitionProtoNames = importDefinitionMap.keys.toSeq
+    val importAllDefinitionExtModNames = importDefinitionMap.toSeq.map(_._2)
+    val importDistinctDefinitionExtModNames = importAllDefinitionExtModNames.distinct
+
+    if (importDistinctDefinitionProtoNames.length < importAllDefinitionProtoNames.length) {
+      val duplicates = importAllDefinitionProtoNames.diff(importDistinctDefinitionProtoNames).mkString(", ")
+      throwException(s"Expected distinct imported Definition names but found duplicates for: $duplicates")
+    }
+    if (importDistinctDefinitionExtModNames.length < importAllDefinitionExtModNames.length) {
+      val duplicates = importAllDefinitionExtModNames.diff(importDistinctDefinitionExtModNames).mkString(", ")
+      throwException(s"Expected distinct overrideDef names but found duplicates for: $duplicates")
+    }
+    (importAllDefinitionProtoNames ++ importAllDefinitionExtModNames).distinct
   }
 
   val globalNamespace = Namespace.empty
@@ -379,8 +398,8 @@ private[chisel3] class DynamicContext(
   // Ensure imported Definitions emit as ExtModules with the correct name so
   // that instantiations will also use the correct name and prevent any name
   // conflicts with Modules/Definitions in this elaboration
-  importDefinitionNames.foreach { importDefName =>
-    globalNamespace.name(importDefName)
+  checkAndGeDistinctProtoExtModNames().foreach {
+    globalNamespace.name(_)
   }
 
   val components = ArrayBuffer[Component]()
@@ -447,11 +466,12 @@ private[chisel3] object Builder extends LazyLogging {
 
   def idGen: IdGen = chiselContext.get.idGen
 
-  def globalNamespace: Namespace = dynamicContext.globalNamespace
-  def components:      ArrayBuffer[Component] = dynamicContext.components
-  def annotations:     ArrayBuffer[ChiselAnnotation] = dynamicContext.annotations
-  def annotationSeq:   AnnotationSeq = dynamicContext.annotationSeq
-  def namingStack:     NamingStack = dynamicContext.namingStack
+  def globalNamespace:     Namespace = dynamicContext.globalNamespace
+  def components:          ArrayBuffer[Component] = dynamicContext.components
+  def annotations:         ArrayBuffer[ChiselAnnotation] = dynamicContext.annotations
+  def annotationSeq:       AnnotationSeq = dynamicContext.annotationSeq
+  def namingStack:         NamingStack = dynamicContext.namingStack
+  def importDefinitionMap: Map[String, String] = dynamicContext.importDefinitionMap
 
   def unnamedViews:  ArrayBuffer[Data] = dynamicContext.unnamedViews
   def viewNamespace: Namespace = chiselContext.get.viewNamespace
diff --git a/src/test/scala/chiselTests/experimental/hierarchy/SeparateElaborationSpec.scala b/src/test/scala/chiselTests/experimental/hierarchy/SeparateElaborationSpec.scala
index 7555a1c4..25bbc474 100644
--- a/src/test/scala/chiselTests/experimental/hierarchy/SeparateElaborationSpec.scala
+++ b/src/test/scala/chiselTests/experimental/hierarchy/SeparateElaborationSpec.scala
@@ -353,4 +353,143 @@ class SeparateElaborationSpec extends ChiselFunSpec with Utils {
     )
   }
 
+  describe("(4): With ExtMod Names") {
+    it("(4.a): should pick correct ExtMod names when passed") {
+      val testDir = createTestDirectory(this.getClass.getSimpleName).toString
+
+      val dutDef = getAddOneDefinition(testDir)
+
+      class Testbench(defn: Definition[AddOne]) extends Module {
+        val mod = Module(new AddOne)
+        val inst = Instance(defn)
+
+        // Tie inputs to a value so ChiselStage does not complain
+        mod.in := 0.U
+        inst.in := 0.U
+        dontTouch(mod.out)
+      }
+
+      (new ChiselStage).run(
+        Seq(
+          ChiselGeneratorAnnotation(() => new Testbench(dutDef)),
+          TargetDirAnnotation(testDir),
+          ImportDefinitionAnnotation(dutDef, Some("CustomPrefix_AddOne_CustomSuffix"))
+        )
+      )
+
+      val tb_rtl = Source.fromFile(s"$testDir/Testbench.v").getLines.mkString
+
+      tb_rtl should include("module AddOne_1(")
+      tb_rtl should include("AddOne_1 mod (")
+      (tb_rtl should not).include("module AddOne(")
+      tb_rtl should include("CustomPrefix_AddOne_CustomSuffix inst (")
+    }
+  }
+
+  it(
+    "(4.b): should work if a list of imported Definitions is passed between Stages with ExtModName."
+  ) {
+    val testDir = createTestDirectory(this.getClass.getSimpleName).toString
+
+    val dutAnnos0 = (new ChiselStage).run(
+      Seq(
+        ChiselGeneratorAnnotation(() => new AddOneParameterized(4)),
+        TargetDirAnnotation(s"$testDir/dutDef0")
+      )
+    )
+    val dutDef0 = getDesignAnnotation(dutAnnos0).design.asInstanceOf[AddOneParameterized].toDefinition
+
+    val dutAnnos1 = (new ChiselStage).run(
+      Seq(
+        ChiselGeneratorAnnotation(() => new AddOneParameterized(8)),
+        TargetDirAnnotation(s"$testDir/dutDef1"),
+        // pass in previously elaborated Definitions
+        ImportDefinitionAnnotation(dutDef0)
+      )
+    )
+    val dutDef1 = getDesignAnnotation(dutAnnos1).design.asInstanceOf[AddOneParameterized].toDefinition
+
+    class Testbench(defn0: Definition[AddOneParameterized], defn1: Definition[AddOneParameterized]) extends Module {
+      val inst0 = Instance(defn0)
+      val inst1 = Instance(defn1)
+
+      // Tie inputs to a value so ChiselStage does not complain
+      inst0.in := 0.U
+      inst1.in := 0.U
+    }
+
+    (new ChiselStage).run(
+      Seq(
+        ChiselGeneratorAnnotation(() => new Testbench(dutDef0, dutDef1)),
+        TargetDirAnnotation(testDir),
+        ImportDefinitionAnnotation(dutDef0, Some("Inst1_Prefix_AddOnePramaterized_Inst1_Suffix")),
+        ImportDefinitionAnnotation(dutDef1, Some("Inst2_Prefix_AddOnePrameterized_1_Inst2_Suffix"))
+      )
+    )
+
+    val dutDef0_rtl = Source.fromFile(s"$testDir/dutDef0/AddOneParameterized.v").getLines.mkString
+    dutDef0_rtl should include("module AddOneParameterized(")
+    val dutDef1_rtl = Source.fromFile(s"$testDir/dutDef1/AddOneParameterized_1.v").getLines.mkString
+    dutDef1_rtl should include("module AddOneParameterized_1(")
+
+    val tb_rtl = Source.fromFile(s"$testDir/Testbench.v").getLines.mkString
+    tb_rtl should include("Inst1_Prefix_AddOnePramaterized_Inst1_Suffix inst0 (")
+    tb_rtl should include("Inst2_Prefix_AddOnePrameterized_1_Inst2_Suffix inst1 (")
+    (tb_rtl should not).include("module AddOneParameterized(")
+    (tb_rtl should not).include("module AddOneParameterized_1(")
+  }
+
+  it(
+    "(4.c): should throw an exception  if a list of imported Definitions is passed between Stages with same ExtModName."
+  ) {
+    val testDir = createTestDirectory(this.getClass.getSimpleName).toString
+
+    val dutAnnos0 = (new ChiselStage).run(
+      Seq(
+        ChiselGeneratorAnnotation(() => new AddOneParameterized(4)),
+        TargetDirAnnotation(s"$testDir/dutDef0")
+      )
+    )
+    val importDefinitionAnnos0 = allModulesToImportedDefs(dutAnnos0)
+    val dutDef0 = getDesignAnnotation(dutAnnos0).design.asInstanceOf[AddOneParameterized].toDefinition
+
+    val dutAnnos1 = (new ChiselStage).run(
+      Seq(
+        ChiselGeneratorAnnotation(() => new AddOneParameterized(8)),
+        TargetDirAnnotation(s"$testDir/dutDef1"),
+        // pass in previously elaborated Definitions
+        ImportDefinitionAnnotation(dutDef0)
+      )
+    )
+    val importDefinitionAnnos1 = allModulesToImportedDefs(dutAnnos1)
+    val dutDef1 = getDesignAnnotation(dutAnnos1).design.asInstanceOf[AddOneParameterized].toDefinition
+
+    class Testbench(defn0: Definition[AddOneParameterized], defn1: Definition[AddOneParameterized]) extends Module {
+      val inst0 = Instance(defn0)
+      val inst1 = Instance(defn1)
+
+      // Tie inputs to a value so ChiselStage does not complain
+      inst0.in := 0.U
+      inst1.in := 0.U
+    }
+
+    val dutDef0_rtl = Source.fromFile(s"$testDir/dutDef0/AddOneParameterized.v").getLines.mkString
+    dutDef0_rtl should include("module AddOneParameterized(")
+    val dutDef1_rtl = Source.fromFile(s"$testDir/dutDef1/AddOneParameterized_1.v").getLines.mkString
+    dutDef1_rtl should include("module AddOneParameterized_1(")
+
+    val errMsg = intercept[ChiselException] {
+      (new ChiselStage).run(
+        Seq(
+          ChiselGeneratorAnnotation(() => new Testbench(dutDef0, dutDef1)),
+          TargetDirAnnotation(testDir),
+          ImportDefinitionAnnotation(dutDef0, Some("Inst1_Prefix_AddOnePrameterized_Inst1_Suffix")),
+          ImportDefinitionAnnotation(dutDef1, Some("Inst1_Prefix_AddOnePrameterized_Inst1_Suffix"))
+        )
+      )
+    }
+    errMsg.getMessage should include(
+      "Expected distinct overrideDef names but found duplicates for: Inst1_Prefix_AddOnePrameterized_Inst1_Suffix"
+    )
+  }
 }