Multi protobuf module emission and consumption

src/main/scala/firrtl/Emitter.scala

@@ -99,7 +99,7 @@ object EmitCircuitAnnotation extends HasShellOptions {
           case "low-opt" =>
             Seq(
               RunFirrtlTransformAnnotation(new ProtoEmitter.OptLow),
-              EmitCircuitAnnotation(classOf[ProtoEmitter.Low])
+              EmitCircuitAnnotation(classOf[ProtoEmitter.OptLow])
             )
           case _ => throw new PhaseException(s"Unknown emitter '$a'! (Did you misspell it?)")
         },
@@ -151,6 +151,43 @@ object EmitAllModulesAnnotation extends HasShellOptions {
       helpText = "Run the specified module emitter (one file per module)",
       shortOption = Some("e"),
       helpValueName = Some("<chirrtl|high|middle|low|verilog|mverilog|sverilog>")
+    ),
+    new ShellOption[String](
+      longOption = "emit-modules-protobuf",
+      toAnnotationSeq = (a: String) =>
+        a match {
+          case "chirrtl" =>
+            Seq(
+              RunFirrtlTransformAnnotation(new ProtoEmitter.Chirrtl),
+              EmitAllModulesAnnotation(classOf[ProtoEmitter.Chirrtl])
+            )
+          case "mhigh" =>
+            Seq(
+              RunFirrtlTransformAnnotation(new ProtoEmitter.MHigh),
+              EmitAllModulesAnnotation(classOf[ProtoEmitter.MHigh])
+            )
+          case "high" =>
+            Seq(
+              RunFirrtlTransformAnnotation(new ProtoEmitter.High),
+              EmitAllModulesAnnotation(classOf[ProtoEmitter.High])
+            )
+          case "middle" =>
+            Seq(
+              RunFirrtlTransformAnnotation(new ProtoEmitter.Middle),
+              EmitAllModulesAnnotation(classOf[ProtoEmitter.Middle])
+            )
+          case "low" =>
+            Seq(RunFirrtlTransformAnnotation(new ProtoEmitter.Low), EmitAllModulesAnnotation(classOf[ProtoEmitter.Low]))
+          case "low-opt" =>
+            Seq(
+              RunFirrtlTransformAnnotation(new ProtoEmitter.OptLow),
+              EmitAllModulesAnnotation(classOf[ProtoEmitter.OptLow])
+            )
+          case _ => throw new PhaseException(s"Unknown emitter '$a'! (Did you misspell it?)")
+        },
+      helpText = "Run the specified module emitter (one protobuf per module)",
+      shortOption = Some("p"),
+      helpValueName = Some("<chirrtl|mhigh|high|middle|low|low-opt>")
     )
   )
 

src/main/scala/firrtl/Utils.scala

@@ -978,6 +978,45 @@ object Utils extends LazyLogging {
     map.view.map({ case (k, vs) => k -> vs.toList }).toList
   }
 
+  /** Combines several separate circuit modules (typically emitted by -e or -p compiler options) into a single circuit */
+  def combine(circuits: Seq[Circuit]): Circuit = {
+    def dedup(modules: Seq[DefModule]): Seq[Either[Module, DefModule]] = {
+      // Left means module with no ExtModules, Right means child modules or lone ExtModules
+      val module: Option[Module] = {
+        val found: Seq[Module] = modules.collect { case m: Module => m }
+        assert(found.size <= 1)
+        found.headOption
+      }
+      val extModules: Seq[ExtModule] = modules.collect { case e: ExtModule => e }.distinct
+
+      // If the module is a lone module (no extmodule references in any other file)
+      if (extModules.isEmpty && !module.isEmpty)
+        Seq(Left(module.get))
+      // If a module has extmodules, but no other file contains the implementation
+      else if (!extModules.isEmpty && module.isEmpty)
+        extModules.map(Right(_))
+      // Otherwise there is a module implementation with extmodule references
+      else
+        Seq(Right(module.get))
+    }
+
+    // 1. Combine modules
+    val grouped: Seq[(String, Seq[DefModule])] = groupByIntoSeq(circuits.flatMap(_.modules))({
+      case mod: Module    => mod.name
+      case ext: ExtModule => ext.defname
+    })
+    val deduped: Iterable[Either[Module, DefModule]] = grouped.flatMap { case (_, insts) => dedup(insts) }
+
+    // 2. Determine top
+    val top = {
+      val found = deduped.collect { case Left(m) => m }
+      assert(found.size == 1)
+      found.head
+    }
+    val res = deduped.collect { case Right(m) => m }
+    ir.Circuit(NoInfo, top +: res.toSeq, top.name)
+  }
+
   object True {
     private val _True = UIntLiteral(1, IntWidth(1))
 

src/main/scala/firrtl/backends/proto/ProtoBufEmitter.scala

@@ -1,15 +1,18 @@
 // SPDX-License-Identifier: Apache-2.0
 package firrtl.backends.proto
 
-import firrtl.{AnnotationSeq, CircuitState, DependencyAPIMigration, Transform}
-import firrtl.ir
+import firrtl._
+import firrtl.ir._
 import firrtl.annotations.NoTargetAnnotation
 import firrtl.options.CustomFileEmission
 import firrtl.options.Viewer.view
 import firrtl.proto.ToProto
 import firrtl.stage.{FirrtlOptions, Forms}
 import firrtl.stage.TransformManager.TransformDependency
+import firrtl.traversals.Foreachers._
 import java.io.{ByteArrayOutputStream, Writer}
+import scala.collection.mutable.ArrayBuffer
+import Utils.throwInternalError
 
 /** This object defines Annotations that are used by Protocol Buffer emission.
   */
@@ -59,10 +62,47 @@ sealed abstract class ProtoBufEmitter(prereqs: Seq[TransformDependency])
   override def optionalPrerequisiteOf = Seq.empty
   override def invalidates(a: Transform) = false
 
-  override def execute(state: CircuitState) =
-    state.copy(annotations = state.annotations :+ Annotation.ProtoBufSerialization(state.circuit, Some(outputSuffix)))
+  private def emitAllModules(circuit: Circuit): Seq[Annotation.ProtoBufSerialization] = {
+    // For a given module, returns a Seq of all modules instantited inside of it
+    def collectInstantiatedModules(mod: Module, map: Map[String, DefModule]): Seq[DefModule] = {
+      // Use list instead of set to maintain order
+      val modules = ArrayBuffer.empty[DefModule]
+      def onStmt(stmt: Statement): Unit = stmt match {
+        case DefInstance(_, _, name, _) => modules += map(name)
+        case _: WDefInstanceConnector => throwInternalError(s"unrecognized statement: $stmt")
+        case other => other.foreach(onStmt)
+      }
+      onStmt(mod.body)
+      modules.distinct.toSeq
+    }
+    val modMap = circuit.modules.map(m => m.name -> m).toMap
+    // Turn each module into it's own circuit with it as the top and all instantied modules as ExtModules
+    circuit.modules.collect {
+      case m: Module =>
+        val instModules = collectInstantiatedModules(m, modMap)
+        val extModules = instModules.map {
+          case Module(info, name, ports, _) => ExtModule(info, name, ports, name, Seq.empty)
+          case ext: ExtModule => ext
+        }
+        val newCircuit = Circuit(m.info, extModules :+ m, m.name)
+        Annotation.ProtoBufSerialization(newCircuit, Some(outputSuffix))
+    }
+  }
+
+  override def execute(state: CircuitState) = {
+    val newAnnos = state.annotations.flatMap {
+      case EmitCircuitAnnotation(a) if this.getClass == a =>
+        Seq(
+          Annotation.ProtoBufSerialization(state.circuit, Some(outputSuffix))
+        )
+      case EmitAllModulesAnnotation(a) if this.getClass == a =>
+        emitAllModules(state.circuit)
+      case _ => Seq()
+    }
+    state.copy(annotations = newAnnos ++ state.annotations)
+  }
 
-  override def emit(state: CircuitState, writer: Writer): Unit = {
+  def emit(state: CircuitState, writer: Writer): Unit = {
     val ostream = new java.io.ByteArrayOutputStream
     ToProto.writeToStream(ostream, state.circuit)
     writer.write(ostream.toString())

src/main/scala/firrtl/proto/FromProto.scala

@@ -9,6 +9,11 @@ import collection.JavaConverters._
 import FirrtlProtos._
 import com.google.protobuf.CodedInputStream
 import Firrtl.Statement.{Formal, ReadUnderWrite}
+import firrtl.ir.DefModule
+import Utils.combine
+import java.io.FileNotFoundException
+import firrtl.options.OptionsException
+import java.nio.file.NotDirectoryException
 
 object FromProto {
 
@@ -33,6 +38,26 @@ object FromProto {
     proto.FromProto.convert(pb)
   }
 
+  /** Deserialize all the ProtoBuf representations of [[ir.Circuit]] in @dir
+    *
+    * @param dir directory containing ProtoBuf representation(s)
+    * @return Deserialized FIRRTL Circuit
+    * @throws java.io.FileNotFoundException if dir does not exist
+    * @throws java.nio.file.NotDirectoryException if dir exists but is not a directory
+    */
+  def fromDirectory(dir: String): ir.Circuit = {
+    val d = new File(dir)
+    if (!d.exists) {
+      throw new FileNotFoundException
+    }
+    if (!d.isDirectory) {
+      throw new NotDirectoryException("Not a directory")
+    }
+
+    val fileList = d.listFiles.filter(_.isFile).toList
+    combine(fileList.map(f => fromInputStream(new FileInputStream(f))))
+  }
+
   // Convert from ProtoBuf message repeated Statements to FIRRRTL Block
   private def compressStmts(stmts: scala.collection.Seq[ir.Statement]): ir.Statement = stmts match {
     case scala.collection.Seq()     => ir.EmptyStmt

src/main/scala/firrtl/stage/FirrtlAnnotations.scala

@@ -6,8 +6,8 @@ import firrtl._
 import firrtl.ir.Circuit
 import firrtl.annotations.{Annotation, NoTargetAnnotation}
 import firrtl.options.{Dependency, HasShellOptions, OptionsException, ShellOption, Unserializable}
-import java.io.FileNotFoundException
-import java.nio.file.NoSuchFileException
+import java.io.{File, FileNotFoundException}
+import java.nio.file.{NoSuchFileException, NotDirectoryException}
 
 import firrtl.stage.TransformManager.TransformDependency
 
@@ -64,6 +64,41 @@ object FirrtlFileAnnotation extends HasShellOptions {
 
 }
 
+/** Read a directory of FIRRTL files or ProtoBufs
+  *  - set with `-I/--input-directory`
+  * @param file input filename
+  */
+case class FirrtlDirectoryAnnotation(dir: String) extends NoTargetAnnotation with CircuitOption {
+
+  def toCircuit(info: Parser.InfoMode): FirrtlCircuitAnnotation = {
+    val circuit =
+      try {
+        proto.FromProto.fromDirectory(dir)
+      } catch {
+        case a @ (_: FileNotFoundException | _: NoSuchFileException) =>
+          throw new OptionsException(s"Directory '$dir' not found! (Did you misspell it?)", a)
+        case _: NotDirectoryException =>
+          throw new OptionsException(s"Directory '$dir' is not a directory")
+      }
+    FirrtlCircuitAnnotation(circuit)
+  }
+
+}
+
+object FirrtlDirectoryAnnotation extends HasShellOptions {
+
+  val options = Seq(
+    new ShellOption[String](
+      longOption = "input-directory",
+      toAnnotationSeq = a => Seq(FirrtlDirectoryAnnotation(a)),
+      helpText = "A directory of FIRRTL files",
+      shortOption = Some("I"),
+      helpValueName = Some("<directory>")
+    )
+  )
+
+}
+
 /** An explicit output file the emitter will write to
   *   - set with `-o/--output-file`
   *  @param file output filename

src/main/scala/firrtl/stage/FirrtlCli.scala

@@ -13,6 +13,7 @@ trait FirrtlCli { this: Shell =>
   parser.note("FIRRTL Compiler Options")
   Seq(
     FirrtlFileAnnotation,
+    FirrtlDirectoryAnnotation,
     OutputFileAnnotation,
     InfoModeAnnotation,
     FirrtlSourceAnnotation,

src/main/scala/firrtl/stage/phases/Checks.scala

@@ -31,34 +31,39 @@ class Checks extends Phase {
     * @throws firrtl.options.OptionsException if any checks fail
     */
   def transform(annos: AnnotationSeq): AnnotationSeq = {
-    val inF, inS, eam, ec, outF, emitter, im, inC = collection.mutable.ListBuffer[Annotation]()
+    val inF, inS, inD, eam, ec, outF, emitter, im, inC = collection.mutable.ListBuffer[Annotation]()
     annos.foreach(_ match {
-      case a: FirrtlFileAnnotation     => a +=: inF
-      case a: FirrtlSourceAnnotation   => a +=: inS
-      case a: EmitAllModulesAnnotation => a +=: eam
-      case a: EmitCircuitAnnotation    => a +=: ec
-      case a: OutputFileAnnotation     => a +=: outF
-      case a: InfoModeAnnotation       => a +=: im
-      case a: FirrtlCircuitAnnotation  => a +=: inC
+      case a: FirrtlFileAnnotation      => a +=: inF
+      case a: FirrtlSourceAnnotation    => a +=: inS
+      case a: FirrtlDirectoryAnnotation => a +=: inD
+      case a: EmitAllModulesAnnotation  => a +=: eam
+      case a: EmitCircuitAnnotation     => a +=: ec
+      case a: OutputFileAnnotation      => a +=: outF
+      case a: InfoModeAnnotation        => a +=: im
+      case a: FirrtlCircuitAnnotation   => a +=: inC
       case a @ RunFirrtlTransformAnnotation(_: firrtl.Emitter) => a +=: emitter
       case _ =>
     })
 
     /* At this point, only a FIRRTL Circuit should exist */
-    if (inF.isEmpty && inS.isEmpty && inC.isEmpty) {
-      throw new OptionsException(s"""|Unable to determine FIRRTL source to read. None of the following were found:
-                                     |    - an input file:  -i, --input-file,    FirrtlFileAnnotation
-                                     |    - FIRRTL source:      --firrtl-source, FirrtlSourceAnnotation
-                                     |    - FIRRTL circuit:                      FirrtlCircuitAnnotation""".stripMargin)
+    if (inF.isEmpty && inS.isEmpty && inD.isEmpty && inC.isEmpty) {
+      throw new OptionsException(
+        s"""|Unable to determine FIRRTL source to read. None of the following were found:
+            |    - an input file:  -i, --input-file,      FirrtlFileAnnotation
+            |    - an input dir:   -I, --input-directory, FirrtlDirectoryAnnotation
+            |    - FIRRTL source:      --firrtl-source,   FirrtlSourceAnnotation
+            |    - FIRRTL circuit:                        FirrtlCircuitAnnotation""".stripMargin
+      )
     }
 
     /* Only one FIRRTL input can exist */
     if (inF.size + inS.size + inC.size > 1) {
       throw new OptionsException(
         s"""|Multiply defined input FIRRTL sources. More than one of the following was found:
-            |    - an input file (${inF.size} times):  -i, --input-file,    FirrtlFileAnnotation
-            |    - FIRRTL source (${inS.size} times):      --firrtl-source, FirrtlSourceAnnotation
-            |    - FIRRTL circuit (${inC.size} times):                      FirrtlCircuitAnnotation""".stripMargin
+            |    - an input file (${inF.size} times):  -i, --input-file,      FirrtlFileAnnotation
+            |    - an input dir (${inD.size} times):   -I, --input-directory, FirrtlDirectoryAnnotation
+            |    - FIRRTL source (${inS.size} times):      --firrtl-source,   FirrtlSourceAnnotation
+            |    - FIRRTL circuit (${inC.size} times):                        FirrtlCircuitAnnotation""".stripMargin
       )
     }
 

src/test/scala/firrtlTests/stage/FirrtlMainSpec.scala

@@ -246,6 +246,16 @@ class FirrtlMainSpec
         args = Array("-X", "sverilog", "-e", "sverilog"),
         files = Seq("Top.sv", "Child.sv")
       ),
+      /* Test all one protobuf per module emitters */
+      FirrtlMainTest(
+        args = Array("-X", "none", "--emit-modules-protobuf", "chirrtl"),
+        files = Seq("Top.pb", "Child.pb")
+      ),
+      FirrtlMainTest(args = Array("-X", "none", "-p", "mhigh"), files = Seq("Top.mhi.pb", "Child.mhi.pb")),
+      FirrtlMainTest(args = Array("-X", "none", "-p", "high"), files = Seq("Top.hi.pb", "Child.hi.pb")),
+      FirrtlMainTest(args = Array("-X", "none", "-p", "middle"), files = Seq("Top.mid.pb", "Child.mid.pb")),
+      FirrtlMainTest(args = Array("-X", "none", "-p", "low"), files = Seq("Top.lo.pb", "Child.lo.pb")),
+      FirrtlMainTest(args = Array("-X", "none", "-p", "low-opt"), files = Seq("Top.lo.pb", "Child.lo.pb")),
       /* Test mixing of -E with -e */
       FirrtlMainTest(
         args = Array("-X", "middle", "-E", "high", "-e", "middle"),
@@ -324,6 +334,41 @@ class FirrtlMainSpec
       new File(td.buildDir + "/Foo.hi.fir") should (exist)
     }
 
+    Scenario("User compiles to multiple Protocol Buffers") {
+      val f = new FirrtlMainFixture
+      val td = new TargetDirectoryFixture("multi-protobuf")
+      val c = new SimpleFirrtlCircuitFixture
+      val protobufs = Seq("Top.hi.pb", "Child.hi.pb")
+
+      And("some input multi-module FIRRTL IR")
+      val inputFile: Array[String] = {
+        val in = new File(td.dir, c.main)
+        val pw = new PrintWriter(in)
+        pw.write(c.input)
+        pw.close()
+        Array("-i", in.toString)
+      }
+
+      When("the user tries to compile to multiple Protocol Buffers in the target directory")
+      f.stage.main(
+        inputFile ++ Array("-X", "none", "-p", "high", "-td", td.buildDir.toString)
+      )
+
+      protobufs.foreach { f =>
+        Then(s"file '$f' should be emitted")
+        val out = new File(td.buildDir + s"/$f")
+        out should (exist)
+      }
+
+      When("the user compiles the Protobufs to a single High FIRRTL IR")
+      f.stage.main(
+        Array("-I", td.buildDir.toString, "-X", "none", "-E", "high", "-td", td.buildDir.toString, "-o", "Foo")
+      )
+
+      Then("one single High FIRRTL file should be emitted")
+      new File(td.buildDir + "/Foo.hi.fir") should (exist)
+    }
+
   }
 
   info("As a FIRRTL command line user")
@@ -381,6 +426,12 @@ class FirrtlMainSpec
         circuit = None,
         stdout = Some("Unknown compiler name 'Verilog'! (Did you misspell it?)"),
         result = 1
+      ),
+      FirrtlMainTest(
+        args = Array("-I", "test_run_dir/I-DO-NOT-EXIST"),
+        circuit = None,
+        stdout = Some("Directory 'test_run_dir/I-DO-NOT-EXIST' not found!"),
+        result = 1
       )
     )
       .foreach(runStageExpectFiles)