make SerializableModuleGenerator work with D/I (#4003)

* implement the instantiable SerializableModule * add test
chipsalliance · Apr 18, 2024 · 03ef61f · 03ef61f
1 parent 3b1ab45
commit 03ef61f
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 3 deletions.
diff --git a/core/src/main/scala/chisel3/experimental/SerializableModuleGenerator.scala b/core/src/main/scala/chisel3/experimental/SerializableModuleGenerator.scala
@@ -1,5 +1,7 @@
 package chisel3.experimental
 
+import chisel3.experimental.hierarchy.{Definition, Instance}
+import chisel3.internal.{instantiable, Builder, BuilderContextCache}
 import upickle.default._
 
 import scala.reflect.runtime.universe
@@ -17,6 +19,7 @@ trait SerializableModuleParameter
   * }}}
   * 3. user should guarantee the module is reproducible on their own.
   */
+@instantiable
 trait SerializableModule[T <: SerializableModuleParameter] { this: BaseModule =>
   val parameter: T
 }
@@ -42,6 +45,13 @@ object SerializableModuleGenerator {
       )
     }
   )
+
+  /** cache instance of a generator. */
+  private case class CacheKey[P <: SerializableModuleParameter, M <: SerializableModule[P]](
+    parameter: P,
+    mTypeTag:  universe.TypeTag[M])
+      extends BuilderContextCache.Key[Definition[M with BaseModule]]
+
 }
 
 /** the serializable module generator:
@@ -54,7 +64,7 @@ case class SerializableModuleGenerator[M <: SerializableModule[P], P <: Serializ
 )(
   implicit val pTag: universe.TypeTag[P],
   implicit val mTag: universe.TypeTag[M]) {
-  private[chisel3] def construct: M = {
+  private[chisel3] def construct: M with BaseModule = {
     require(
       generator.getConstructors.length == 1,
       s"""only allow constructing SerializableModule from SerializableModuleParameter via class Module(val parameter: Parameter),
@@ -75,9 +85,22 @@ case class SerializableModuleGenerator[M <: SerializableModule[P], P <: Serializ
          |${generator.getConstructors.head.getParameterTypes.mkString("\n")}
          |""".stripMargin
     )
-    generator.getConstructors.head.newInstance(parameter).asInstanceOf[M]
+    generator.getConstructors.head.newInstance(parameter).asInstanceOf[M with BaseModule]
   }
 
   /** elaborate a module from this generator. */
-  def module(): M = construct
+  def module(): M with BaseModule = construct
+
+  /** get the definition from this generator. */
+  def definition(): Definition[M with BaseModule] = Builder.contextCache
+    .getOrElseUpdate(
+      SerializableModuleGenerator.CacheKey(
+        parameter,
+        implicitly[universe.TypeTag[M]]
+      ),
+      Definition.do_apply(construct)(UnlocatableSourceInfo)
+    )
+
+  /** get an instance of from this generator. */
+  def instance(): Instance[M with BaseModule] = Instance.do_apply(definition())(UnlocatableSourceInfo)
 }
diff --git a/src/test/scala/chiselTests/experimental/SerializableModuleGeneratorSpec.scala b/src/test/scala/chiselTests/experimental/SerializableModuleGeneratorSpec.scala
@@ -59,6 +59,22 @@ class SerializableModuleGeneratorSpec extends ChiselFlatSpec with Utils {
     circt.stage.ChiselStage.emitCHIRRTL(g.module())
   }
 
+  "SerializableModuleGenerator" should "be able to elaborate with D/I" in {
+    val cir = circt.stage.ChiselStage.emitCHIRRTL(
+      new Module {
+        val i32 =
+          SerializableModuleGenerator(classOf[GCDSerializableModule], GCDSerializableModuleParameter(32)).instance()
+        val i16 =
+          SerializableModuleGenerator(classOf[GCDSerializableModule], GCDSerializableModuleParameter(16)).instance()
+        val ii32 =
+          SerializableModuleGenerator(classOf[GCDSerializableModule], GCDSerializableModuleParameter(32)).instance()
+      }
+    )
+    cir should include("inst i32 of GCDSerializableModule")
+    cir should include("inst i16 of GCDSerializableModule_1")
+    cir should include("inst ii32 of GCDSerializableModule")
+  }
+
   case class FooParameter() extends SerializableModuleParameter
 
   class InnerFoo(val parameter: FooParameter) extends RawModule with SerializableModule[FooParameter]