From 7f90d7200671be243a65b18f9baa1ca5ad4f9475 Mon Sep 17 00:00:00 2001 From: Fabian Schuiki Date: Fri, 8 Nov 2024 14:20:24 -0800 Subject: [PATCH] [Module] Add afterModuleBuilt hook (#4479) Add the `afterModuleBuilt` hook to `RawModule`, similar to the existing `atModuleBodyEnd`. This hook is called after a module has been fully constructed and its resulting component definition has been added to the builder. Thus code inside the body of `afterModuleBuilt` can get the definition of the surrounding module and instantiate it, which is interesting for generating additional collateral like unit tests alongside a module. --- core/src/main/scala/chisel3/ModuleImpl.scala | 6 ++ core/src/main/scala/chisel3/RawModule.scala | 43 +++++++++++++++ .../scala/chiselTests/RawModuleSpec.scala | 55 ++++++++++++++++++- 3 files changed, 103 insertions(+), 1 deletion(-) diff --git a/core/src/main/scala/chisel3/ModuleImpl.scala b/core/src/main/scala/chisel3/ModuleImpl.scala index a04766b8c85..f75908d88c7 100644 --- a/core/src/main/scala/chisel3/ModuleImpl.scala +++ b/core/src/main/scala/chisel3/ModuleImpl.scala @@ -114,6 +114,7 @@ private[chisel3] trait ObjectModuleImpl { Builder.setPrefix(savePrefix) Builder.enabledLayers = saveEnabledLayers + module.moduleBuilt() module } @@ -644,6 +645,11 @@ package experimental { } } + /** Called once the module's definition has been fully built. At this point + * the module can be instantiated through its definition. + */ + protected[chisel3] def moduleBuilt(): Unit = {} + // // Chisel Internals // diff --git a/core/src/main/scala/chisel3/RawModule.scala b/core/src/main/scala/chisel3/RawModule.scala index 7ffb87a79bf..de74374d6f3 100644 --- a/core/src/main/scala/chisel3/RawModule.scala +++ b/core/src/main/scala/chisel3/RawModule.scala @@ -82,6 +82,46 @@ abstract class RawModule extends BaseModule { } private val _atModuleBodyEnd = new ArrayBuffer[() => Unit] + /** Hook to invoke hardware generators after a Module has been constructed and closed. + * + * This is useful for running hardware generators after a Module's constructor has run and its Definition is available, while still having access to arguments and definitions in the constructor. The Module itself can no longer be modified at this point. + * + * An interesting application of this is the generation of unit tests whenever a module is instantiated. For example: + * + * {{{ + * class Example(N: int) extends RawModule { + * private val someSecret: Int = ... + * + * afterModuleBuilt { + * // Executes once the surrounding module is closed. + * // We can get its definition at this point and pass it to another module. + * Definition(ExampleTest(this.definition, someSecret)) + * } + * } + * + * class ExampleTest(unitDef: Definition[Example], someSecret: Int) extends RawModule { + * // Instantiate the generated module and test it. + * val unit = Instance(unitDef) + * ... + * } + * + * class Parent extends RawModule { + * Instantiate(Example(42)) + * } + * + * // Resulting modules: + * // - Parent (top-level) + * // - instantiates Example + * // - ExampleTest (top-level) + * // - instantiates Example + * // - Example + * }}} + */ + protected def afterModuleBuilt(gen: => Unit): Unit = { + _afterModuleBuilt += { () => gen } + } + private val _afterModuleBuilt = new ArrayBuffer[() => Unit] + // // RTL construction internals // @@ -253,6 +293,9 @@ abstract class RawModule extends BaseModule { Builder.currentBlock.get.addSecretCommand(rhs) } + // Evaluate any afterModuleBuilt generators. + protected[chisel3] override def moduleBuilt(): Unit = _afterModuleBuilt.foreach(_()) + private[chisel3] def initializeInParent(): Unit = {} } diff --git a/src/test/scala/chiselTests/RawModuleSpec.scala b/src/test/scala/chiselTests/RawModuleSpec.scala index 09f73543eba..0c42de2d3a6 100644 --- a/src/test/scala/chiselTests/RawModuleSpec.scala +++ b/src/test/scala/chiselTests/RawModuleSpec.scala @@ -4,7 +4,7 @@ package chiselTests import chisel3._ import chisel3.aop.Select -import chisel3.experimental.hierarchy.Definition +import chisel3.experimental.hierarchy.{instantiable, public, Definition, Instance} import chisel3.reflect.DataMirror import chisel3.testers.BasicTester import circt.stage.ChiselStage @@ -170,4 +170,57 @@ class RawModuleSpec extends ChiselFlatSpec with Utils with MatchesAndOmits { } } } + + "RawModule with afterModuleBuilt" should "be able to create other modules" in { + val chirrtl = ChiselStage.emitCHIRRTL(new RawModule { + override def desiredName = "Foo" + val port0 = IO(Input(Bool())) + + afterModuleBuilt { + Module(new RawModule { + override def desiredName = "Bar" + val port1 = IO(Input(Bool())) + }) + } + }) + + matchesAndOmits(chirrtl)( + "module Foo", + "input port0 : UInt<1>", + "module Bar", + "input port1 : UInt<1>" + )() + } + + "RawModule with afterModuleBuilt" should "be able to instantiate the surrounding module" in { + @instantiable + class Foo extends RawModule { + override def desiredName = "Foo" + @public val port0 = IO(Input(Bool())) + + afterModuleBuilt { + val fooDef = this.toDefinition + Module(new RawModule { + override def desiredName = "Bar" + val port1 = IO(Input(Bool())) + val foo1 = Instance(fooDef) + val foo2 = Instance(fooDef) + foo1.port0 := port1 + foo2.port0 := port1 + }) + } + } + val chirrtl = ChiselStage.emitCHIRRTL(new Foo) + + matchesAndOmits(chirrtl)( + "module Foo :", + "input port0 : UInt<1>", + "module Bar", + "input port1 : UInt<1>", + "inst foo1 of Foo", + "inst foo2 of Foo", + "connect foo1.port0, port1", + "connect foo2.port0, port1" + )() + } }