Add hook to invoke hardware generators at the end of elaboration. (ch…

…ipsalliance#3446)
rahiCICD · Aug 4, 2023 · 1880b4d · 1880b4d
1 parent cede284
commit 1880b4d
Show file tree

Hide file tree

Showing 3 changed files with 158 additions and 5 deletions.
diff --git a/core/src/main/scala/chisel3/Module.scala b/core/src/main/scala/chisel3/Module.scala
@@ -63,17 +63,19 @@ object Module extends SourceInfoDoc {
           sourceInfo.makeMessage(" See " + _)
       )
     }
-    Builder.currentModule = parent // Back to parent!
-    Builder.whenStack = parentWhenStack
-    Builder.currentClock = saveClock // Back to clock and reset scope
-    Builder.currentReset = saveReset
 
     // Only add the component if the module generates one
     val componentOpt = module.generateComponent()
     for (component <- componentOpt) {
       Builder.components += component
     }
 
+    // Reset Builder state *after* generating the component, so any atModuleBodyEnd generators are still within the
+    // scope of the current Module.
+    Builder.currentModule = parent // Back to parent!
+    Builder.whenStack = parentWhenStack
+    Builder.currentClock = saveClock // Back to clock and reset scope
+    Builder.currentReset = saveReset
     Builder.setPrefix(savePrefix)
 
     // Handle connections at enclosing scope

diff --git a/core/src/main/scala/chisel3/RawModule.scala b/core/src/main/scala/chisel3/RawModule.scala
@@ -12,13 +12,75 @@ import chisel3.internal.Builder._
 import chisel3.internal.firrtl._
 import _root_.firrtl.annotations.{IsModule, ModuleTarget}
 import scala.collection.immutable.VectorBuilder
+import scala.collection.mutable.ArrayBuffer
 
 /** Abstract base class for Modules that contain Chisel RTL.
   * This abstract base class is a user-defined module which does not include implicit clock and reset and supports
   * multiple IO() declarations.
   */
 @nowarn("msg=class Port") // delete when Port becomes private
 abstract class RawModule extends BaseModule {
+
+  /** Hook to invoke hardware generators after the rest of the Module is constructed.
+    *
+    * This is a power-user API, and should not normally be needed.
+    *
+    * In rare cases, it is necessary to run hardware generators at a late stage, but still within the scope of the
+    * Module. In these situations, atModuleBodyEnd may be used to register such generators. For example:
+    *
+    *  {{{
+    *    class Example extends RawModule {
+    *      atModuleBodyEnd {
+    *        val extraPort0 = IO(Output(Bool()))
+    *        extraPort0 := 0.B
+    *      }
+    *    }
+    *  }}}
+    *
+    * Any generators registered with atModuleBodyEnd are the last code to execute when the Module is constructed. The
+    * execution order is:
+    *
+    *   - The constructors of any super classes or traits the Module extends
+    *   - The constructor of the Module itself
+    *   - The atModuleBodyEnd generators
+    *
+    * The atModuleBodyEnd generators execute in the lexical order they appear in the Module constructor.
+    *
+    * For example:
+    *
+    *  {{{
+    *    trait Parent {
+    *      // Executes first.
+    *      val foo = ...
+    *    }
+    *
+    *    class Example extends Parent {
+    *      // Executes second.
+    *      val bar = ...
+    *
+    *      atModuleBodyEnd {
+    *        // Executes fourth.
+    *        val qux = ...
+    *      }
+    *
+    *      atModuleBodyEnd {
+    *        // Executes fifth.
+    *        val quux = ...
+    *      }
+    *
+    *      // Executes third..
+    *      val baz = ...
+    *    }
+    *  }}}
+    *
+    * If atModuleBodyEnd is used in a Definition, any generated hardware will be included in the Definition. However, it
+    * is currently not possible to annotate any val within atModuleBodyEnd as @public.
+    */
+  protected def atModuleBodyEnd(gen: => Unit): Unit = {
+    _atModuleBodyEnd += { () => gen }
+  }
+  private val _atModuleBodyEnd = new ArrayBuffer[() => Unit]
+
   //
   // RTL construction internals
   //
@@ -54,6 +116,12 @@ abstract class RawModule extends BaseModule {
 
   private[chisel3] override def generateComponent(): Option[Component] = {
     require(!_closed, "Can't generate module more than once")
+
+    // Evaluate any atModuleBodyEnd generators.
+    _atModuleBodyEnd.foreach { gen =>
+      gen()
+    }
+
     _closed = true
 
     // Check to make sure that all ports can be named

diff --git a/src/test/scala/chiselTests/RawModuleSpec.scala b/src/test/scala/chiselTests/RawModuleSpec.scala
@@ -3,6 +3,9 @@
 package chiselTests
 
 import chisel3._
+import chisel3.aop.Select
+import chisel3.experimental.hierarchy.Definition
+import chisel3.reflect.DataMirror
 import chisel3.testers.BasicTester
 import circt.stage.ChiselStage
 
@@ -59,7 +62,7 @@ class ImplicitModuleDirectlyInRawModuleTester extends BasicTester {
   stop()
 }
 
-class RawModuleSpec extends ChiselFlatSpec with Utils {
+class RawModuleSpec extends ChiselFlatSpec with Utils with MatchesAndOmits {
   "RawModule" should "elaborate" in {
     ChiselStage.emitCHIRRTL { new RawModuleWithImplicitModule }
   }
@@ -68,6 +71,86 @@ class RawModuleSpec extends ChiselFlatSpec with Utils {
     assertTesterPasses({ new RawModuleTester })
   }
 
+  "RawModule with atModuleBodyEnd" should "support late stage generators" in {
+    val chirrtl = ChiselStage.emitCHIRRTL(new RawModule {
+      atModuleBodyEnd {
+        val extraPort0 = IO(Output(Bool()))
+        extraPort0 := 0.B
+      }
+
+      atModuleBodyEnd {
+        val extraPort1 = IO(Output(Bool()))
+        extraPort1 := 1.B
+      }
+    })
+
+    matchesAndOmits(chirrtl)(
+      "output extraPort0 : UInt<1>",
+      "output extraPort1 : UInt<1>",
+      "connect extraPort0, UInt<1>(0h0)",
+      "connect extraPort1, UInt<1>(0h1)"
+    )()
+  }
+
+  "RawModule with atModuleBodyEnd" should "support multiple connects" in {
+    val chirrtl = ChiselStage.emitCHIRRTL(new RawModule {
+      val port = IO(Output(UInt(2.W)))
+
+      atModuleBodyEnd {
+        port := 2.U
+      }
+
+      atModuleBodyEnd {
+        port := 3.U
+      }
+
+      port := 1.U
+    })
+
+    matchesAndOmits(chirrtl)(
+      "output port : UInt<2>",
+      "connect port, UInt<1>(0h1)",
+      "connect port, UInt<2>(0h2)",
+      "connect port, UInt<2>(0h3)"
+    )()
+  }
+
+  "RawModule with atModuleBodyEnd" should "support the added hardware in DataMirror" in {
+    ChiselStage.emitCHIRRTL(new RawModule {
+      val module = Module(new RawModule {
+        val port0 = IO(Output(Bool()))
+        port0 := 0.B
+
+        atModuleBodyEnd {
+          val port1 = IO(Output(Bool()))
+          port1 := 0.B
+        }
+      })
+
+      val mirroredPorts = DataMirror.modulePorts(module)
+
+      mirroredPorts should have size 2
+    })
+  }
+
+  "RawModule with atModuleBodyEnd" should "support the added hardware in Definition" in {
+    ChiselStage.emitCHIRRTL(new RawModule {
+      val definition = Definition(new RawModule {
+        val port0 = IO(Output(Bool()))
+        port0 := 0.B
+
+        atModuleBodyEnd {
+          val port1 = IO(Output(Bool()))
+          port1 := 0.B
+        }
+      })
+
+      val definitionPorts = Select.ios(definition)
+
+      definitionPorts should have size 2
+    })
+  }
+
   "ImplicitModule in a withClock block in a RawModule" should "work" in {
     assertTesterPasses({ new ImplicitModuleInRawModuleTester })
   }