Add support for outputting multiple modules, one per entry point

Cargo.toml

@@ -7,6 +7,7 @@ members = [
     "examples/shaders/simplest-shader",
     "examples/shaders/compute-shader",
     "examples/shaders/mouse-shader",
+    "examples/multibuilder",
 
     "crates/rustc_codegen_spirv",
     "crates/spirv-builder",

crates/rustc_codegen_spirv/Cargo.toml

@@ -31,6 +31,7 @@ bimap = "0.6"
 indexmap = "1.6.0"
 rspirv = { git = "https://github.com/gfx-rs/rspirv.git", rev = "ee1e913" }
 rustc-demangle = "0.1.18"
+sanitize-filename = "0.3"
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
 smallvec = "1.6.1"

crates/rustc_codegen_spirv/src/codegen_cx/mod.rs

@@ -33,6 +33,7 @@ use std::cell::{Cell, RefCell};
 use std::collections::HashMap;
 use std::iter::once;
 use std::rc::Rc;
+use std::str::FromStr;
 
 pub struct CodegenCx<'tcx> {
     pub tcx: TyCtxt<'tcx>,
@@ -69,6 +70,8 @@ pub struct CodegenCx<'tcx> {
     /// Some runtimes (e.g. intel-compute-runtime) disallow atomics on i8 and i16, even though it's allowed by the spec.
     /// This enables/disables them.
     pub i8_i16_atomics_allowed: bool,
+
+    pub codegen_args: CodegenArgs,
 }
 
 impl<'tcx> CodegenCx<'tcx> {
@@ -102,6 +105,7 @@ impl<'tcx> CodegenCx<'tcx> {
                 tcx.sess.err(&format!("Unknown feature {}", feature));
             }
         }
+        let codegen_args = CodegenArgs::from_session(tcx.sess);
         Self {
             tcx,
             codegen_unit,
@@ -120,6 +124,7 @@ impl<'tcx> CodegenCx<'tcx> {
             panic_fn_id: Default::default(),
             panic_bounds_check_fn_id: Default::default(),
             i8_i16_atomics_allowed: false,
+            codegen_args,
         }
     }
 
@@ -256,6 +261,51 @@ impl<'tcx> CodegenCx<'tcx> {
     }
 }
 
+pub struct CodegenArgs {
+    pub module_output_type: ModuleOutputType,
+}
+
+impl CodegenArgs {
+    pub fn from_session(sess: &Session) -> Self {
+        match CodegenArgs::parse(&sess.opts.cg.llvm_args) {
+            Ok(ok) => ok,
+            Err(err) => sess.fatal(&format!("Unable to parse llvm-args: {}", err)),
+        }
+    }
+
+    pub fn parse(args: &[String]) -> Result<Self, rustc_session::getopts::Fail> {
+        use rustc_session::getopts;
+        let mut opts = getopts::Options::new();
+        opts.optopt(
+            "",
+            "module-output",
+            "single output or multiple output",
+            "[single|multiple]",
+        );
+        let matches = opts.parse(args)?;
+        let module_output_type =
+            matches.opt_get_default("module-output", ModuleOutputType::Single)?;
+        Ok(Self { module_output_type })
+    }
+}
+
+#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
+pub enum ModuleOutputType {
+    Single,
+    Multiple,
+}
+
+impl FromStr for ModuleOutputType {
+    type Err = rustc_session::getopts::Fail;
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        match s {
+            "single" => Ok(Self::Single),
+            "multiple" => Ok(Self::Multiple),
+            v => Err(Self::Err::UnrecognizedOption(v.to_string())),
+        }
+    }
+}
+
 impl<'tcx> BackendTypes for CodegenCx<'tcx> {
     type Value = SpirvValue;
     type Function = SpirvValue;

crates/rustc_codegen_spirv/src/lib.rs

@@ -118,7 +118,7 @@ mod spirv_type_constraints;
 mod symbols;
 
 use builder::Builder;
-use codegen_cx::CodegenCx;
+use codegen_cx::{CodegenArgs, CodegenCx, ModuleOutputType};
 pub use rspirv;
 use rspirv::binary::Assemble;
 use rustc_ast::expand::allocator::AllocatorKind;
@@ -380,6 +380,8 @@ impl CodegenBackend for SpirvCodegenBackend {
     ) -> Result<(), ErrorReported> {
         // TODO: Can we merge this sym with the one in symbols.rs?
         let legalize = !sess.target_features.contains(&Symbol::intern("kernel"));
+        let codegen_args = CodegenArgs::from_session(sess);
+        let emit_multiple_modules = codegen_args.module_output_type == ModuleOutputType::Multiple;
 
         let timer = sess.timer("link_crate");
         link::link(
@@ -388,6 +390,7 @@ impl CodegenBackend for SpirvCodegenBackend {
             outputs,
             &codegen_results.crate_name.as_str(),
             legalize,
+            emit_multiple_modules,
         );
         drop(timer);