LLVM Bitcode Linker: Added crate

11 files changed, 338 insertions, 0 deletions

diff --git a/Cargo.lock b/Cargo.lock
index 3290741f128..53bd088b557 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2266,6 +2266,17 @@ name = "lld-wrapper"
 version = "0.1.0"
 
 [[package]]
+name = "llvm-bitcode-linker"
+version = "0.0.1"
+dependencies = [
+ "anyhow",
+ "clap",
+ "thiserror",
+ "tracing",
+ "tracing-subscriber",
+]
+
+[[package]]
 name = "lock_api"
 version = "0.4.11"
 source = "registry+https://github.com/rust-lang/crates.io-index"
diff --git a/Cargo.toml b/Cargo.toml
index 5847a817e76..5dd315ef2f7 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -34,6 +34,7 @@ members = [
   "src/tools/expand-yaml-anchors",
   "src/tools/jsondocck",
   "src/tools/jsondoclint",
+  "src/tools/llvm-bitcode-linker",
   "src/tools/html-checker",
   "src/tools/bump-stage0",
   "src/tools/replace-version-placeholder",
diff --git a/src/bootstrap/src/core/build_steps/tool.rs b/src/bootstrap/src/core/build_steps/tool.rs
index 889876f461d..53dc1cff0ae 100644
--- a/src/bootstrap/src/core/build_steps/tool.rs
+++ b/src/bootstrap/src/core/build_steps/tool.rs
@@ -795,6 +795,7 @@ tool_extended!((self, builder),
     Rls, "src/tools/rls", "rls", stable=true, tool_std=true;
     RustDemangler, "src/tools/rust-demangler", "rust-demangler", stable=false, tool_std=true;
     Rustfmt, "src/tools/rustfmt", "rustfmt", stable=true, add_bins_to_sysroot = ["rustfmt", "cargo-fmt"];
+    LlvmBitcodeLinker, "src/tools/llvm-bitcode-linker", "llvm-bitcode-linker", stable=false, add_bins_to_sysroot = ["llvm-bitcode-linker"];
 );
 
 impl<'a> Builder<'a> {
diff --git a/src/bootstrap/src/core/builder.rs b/src/bootstrap/src/core/builder.rs
index 4927c837608..5e5d6d024ee 100644
--- a/src/bootstrap/src/core/builder.rs
+++ b/src/bootstrap/src/core/builder.rs
@@ -763,6 +763,7 @@ impl<'a> Builder<'a> {
                 tool::RustdocGUITest,
                 tool::OptimizedDist,
                 tool::CoverageDump,
+                tool::LlvmBitcodeLinker
             ),
             Kind::Check | Kind::Clippy | Kind::Fix => describe!(
                 check::Std,
diff --git a/src/tools/llvm-bitcode-linker/Cargo.toml b/src/tools/llvm-bitcode-linker/Cargo.toml
new file mode 100644
index 00000000000..a9210b562f3
--- /dev/null
+++ b/src/tools/llvm-bitcode-linker/Cargo.toml
@@ -0,0 +1,14 @@
+[package]
+name = "llvm-bitcode-linker"
+version = "0.0.1"
+description = "A self-contained linker for llvm bitcode"
+license = "MIT OR Apache-2.0"
+edition = "2021"
+publish = false
+
+[dependencies]
+anyhow = "1.0"
+tracing = "0.1"
+tracing-subscriber = {version = "0.3.0", features = ["std"] }
+clap = { version = "4.3", features = ["derive"] }
+thiserror = "1.0.24"
diff --git a/src/tools/llvm-bitcode-linker/README.md b/src/tools/llvm-bitcode-linker/README.md
new file mode 100644
index 00000000000..9b8719e3a77
--- /dev/null
+++ b/src/tools/llvm-bitcode-linker/README.md
@@ -0,0 +1,5 @@
+# LLVM Bitcode Linker
+The LLVM bitcode linker can be used to link targets without any dependency on system libraries.
+The code will be linked in llvm-bc before compiling to native code. For some of these targets
+(e.g. ptx) there does not exist a sensible way to link the native format at all. A bitcode linker
+is required to link code compiled for such targets.
diff --git a/src/tools/llvm-bitcode-linker/src/bin/llvm-bitcode-linker.rs b/src/tools/llvm-bitcode-linker/src/bin/llvm-bitcode-linker.rs
new file mode 100644
index 00000000000..a92af71a36e
--- /dev/null
+++ b/src/tools/llvm-bitcode-linker/src/bin/llvm-bitcode-linker.rs
@@ -0,0 +1,62 @@
+use std::path::PathBuf;
+
+use clap::Parser;
+
+use llvm_bitcode_linker::{Optimization, Session, Target};
+
+#[derive(Debug, Parser)]
+/// Linker for embedded code without any system dependencies
+pub struct Args {
+    /// Input files - objects, archives and static libraries.
+    ///
+    /// An archive can be, but not required to be, a Rust rlib.
+    files: Vec<PathBuf>,
+
+    /// A symbol that should be exported
+    #[arg(long)]
+    export_symbol: Vec<String>,
+
+    /// Input files directory
+    #[arg(short = 'L')]
+    input_dir: Vec<PathBuf>,
+
+    /// Target triple for which the code is compiled
+    #[arg(long)]
+    target: Target,
+
+    /// The target cpu
+    #[arg(long)]
+    target_cpu: Option<String>,
+
+    /// Write output to the filename
+    #[arg(short, long)]
+    output: PathBuf,
+
+    // Enable link time optimization
+    #[arg(long)]
+    lto: bool,
+
+    /// Emit debug information
+    #[arg(long)]
+    debug: bool,
+
+    /// The optimization level
+    #[arg(short = 'O', value_enum, default_value = "0")]
+    optimization: Optimization,
+}
+
+fn main() -> anyhow::Result<()> {
+    tracing_subscriber::FmtSubscriber::builder().with_max_level(tracing::Level::DEBUG).init();
+
+    let args = Args::parse();
+
+    let mut linker = Session::new(args.target, args.target_cpu, args.output);
+
+    linker.add_exported_symbols(args.export_symbol);
+
+    for rlib in args.files {
+        linker.add_file(rlib);
+    }
+
+    linker.lto(args.optimization, args.debug)
+}
diff --git a/src/tools/llvm-bitcode-linker/src/lib.rs b/src/tools/llvm-bitcode-linker/src/lib.rs
new file mode 100644
index 00000000000..48f918f631c
--- /dev/null
+++ b/src/tools/llvm-bitcode-linker/src/lib.rs
@@ -0,0 +1,7 @@
+mod linker;
+mod opt;
+mod target;
+
+pub use linker::Session;
+pub use opt::Optimization;
+pub use target::Target;
diff --git a/src/tools/llvm-bitcode-linker/src/linker.rs b/src/tools/llvm-bitcode-linker/src/linker.rs
new file mode 100644
index 00000000000..aa6b6443e4d
--- /dev/null
+++ b/src/tools/llvm-bitcode-linker/src/linker.rs
@@ -0,0 +1,163 @@
+use std::path::PathBuf;
+
+use anyhow::Context;
+
+use crate::Optimization;
+use crate::Target;
+
+#[derive(Debug)]
+pub struct Session {
+    target: Target,
+    cpu: Option<String>,
+    symbols: Vec<String>,
+
+    /// A file that `llvm-link` supports, like a bitcode file or an archive.
+    files: Vec<PathBuf>,
+
+    // Output files
+    link_path: PathBuf,
+    opt_path: PathBuf,
+    sym_path: PathBuf,
+    out_path: PathBuf,
+}
+
+impl Session {
+    pub fn new(target: crate::Target, cpu: Option<String>, out_path: PathBuf) -> Self {
+        let link_path = out_path.with_extension("o");
+        let opt_path = out_path.with_extension("optimized.o");
+        let sym_path = out_path.with_extension("symbols.txt");
+
+        Session {
+            target,
+            cpu,
+            symbols: Vec::new(),
+            files: Vec::new(),
+            link_path,
+            opt_path,
+            sym_path,
+            out_path,
+        }
+    }
+
+    /// Add a file, like an rlib or bitcode file that should be linked
+    pub fn add_file(&mut self, path: PathBuf) {
+        self.files.push(path);
+    }
+
+    /// Add a Vec of symbols to the list of exported symbols
+    pub fn add_exported_symbols(&mut self, symbols: Vec<String>) {
+        self.symbols.extend(symbols);
+    }
+
+    /// Reads every file that was added to the session and link them without optimization.
+    ///
+    /// The resulting artifact will be written to a file that can later be read to perform
+    /// optimizations and/or compilation from bitcode to the final artifact.
+    fn link(&mut self) -> anyhow::Result<()> {
+        tracing::info!("Linking {} files using llvm-link", self.files.len());
+
+        let llvm_link_output = std::process::Command::new("llvm-link")
+            .arg("--ignore-non-bitcode")
+            .args(&self.files)
+            .arg("-o")
+            .arg(&self.link_path)
+            .output()
+            .unwrap();
+
+        if !llvm_link_output.status.success() {
+            tracing::error!(
+                "llvm-link returned with Exit status: {}\n stdout: {}\n stderr: {}",
+                llvm_link_output.status,
+                String::from_utf8(llvm_link_output.stdout).unwrap(),
+                String::from_utf8(llvm_link_output.stderr).unwrap(),
+            );
+            anyhow::bail!("llvm-link failed to link files {:?}", self.files);
+        }
+
+        Ok(())
+    }
+
+    /// Optimize and compile to native format using `opt` and `llc`
+    ///
+    /// Before this can be called `link` needs to be called
+    fn optimize(&mut self, optimization: Optimization, mut debug: bool) -> anyhow::Result<()> {
+        let mut passes = format!("default<{}>", optimization);
+
+        // FIXME(@kjetilkjeka) Debug symbol generation is broken for nvptx64 so we must remove them even in debug mode
+        if debug && self.target == crate::Target::Nvptx64NvidiaCuda {
+            tracing::warn!("nvptx64 target detected - stripping debug symbols");
+            debug = false;
+        }
+
+        // We add an internalize pass as the rust compiler as we require exported symbols to be explicitly marked
+        passes.push_str(",internalize,globaldce");
+        let symbol_file_content = self.symbols.iter().fold(String::new(), |s, x| s + &x + "\n");
+        std::fs::write(&self.sym_path, symbol_file_content)
+            .context(format!("Failed to write symbol file: {}", self.sym_path.display()))?;
+
+        tracing::info!("optimizing bitcode with passes: {}", passes);
+        let mut opt_cmd = std::process::Command::new("opt");
+        opt_cmd
+            .arg(&self.link_path)
+            .arg("-o")
+            .arg(&self.opt_path)
+            .arg(format!("--internalize-public-api-file={}", self.sym_path.display()))
+            .arg(format!("--passes={}", passes));
+
+        if !debug {
+            opt_cmd.arg("--strip-debug");
+        }
+
+        let opt_output = opt_cmd.output().unwrap();
+
+        if !opt_output.status.success() {
+            tracing::error!(
+                "opt returned with Exit status: {}\n stdout: {}\n stderr: {}",
+                opt_output.status,
+                String::from_utf8(opt_output.stdout).unwrap(),
+                String::from_utf8(opt_output.stderr).unwrap(),
+            );
+            anyhow::bail!("opt failed optimize bitcode: {}", self.link_path.display());
+        };
+
+        Ok(())
+    }
+
+    /// Compile the optimized bitcode file to native format using `llc`
+    ///
+    /// Before this can be called `optimize` needs to be called
+    fn compile(&mut self) -> anyhow::Result<()> {
+        let mut lcc_command = std::process::Command::new("llc");
+
+        if let Some(mcpu) = &self.cpu {
+            lcc_command.arg("--mcpu").arg(mcpu);
+        }
+
+        let lcc_output =
+            lcc_command.arg(&self.opt_path).arg("-o").arg(&self.out_path).output().unwrap();
+
+        if !lcc_output.status.success() {
+            tracing::error!(
+                "llc returned with Exit status: {}\n stdout: {}\n stderr: {}",
+                lcc_output.status,
+                String::from_utf8(lcc_output.stdout).unwrap(),
+                String::from_utf8(lcc_output.stderr).unwrap(),
+            );
+
+            anyhow::bail!(
+                "llc failed to compile {} into {}",
+                self.opt_path.display(),
+                self.out_path.display()
+            );
+        }
+
+        Ok(())
+    }
+
+    /// Links, optimizes and compiles to the native format
+    pub fn lto(&mut self, optimization: crate::Optimization, debug: bool) -> anyhow::Result<()> {
+        self.link()?;
+        self.optimize(optimization, debug)?;
+        self.compile()
+    }
+}
diff --git a/src/tools/llvm-bitcode-linker/src/opt.rs b/src/tools/llvm-bitcode-linker/src/opt.rs
new file mode 100644
index 00000000000..93f0ec7e2d6
--- /dev/null
+++ b/src/tools/llvm-bitcode-linker/src/opt.rs
@@ -0,0 +1,53 @@
+use std::fmt::Display;
+use std::fmt::Formatter;
+
+#[derive(Debug, Clone, Copy, Default, Hash, Eq, PartialEq, clap::ValueEnum)]
+pub enum Optimization {
+    #[default]
+    #[value(name = "0")]
+    O0,
+    #[value(name = "1")]
+    O1,
+    #[value(name = "2")]
+    O2,
+    #[value(name = "3")]
+    O3,
+    #[value(name = "s")]
+    Os,
+    #[value(name = "z")]
+    Oz,
+}
+
+#[derive(Debug, Clone, Copy, thiserror::Error)]
+/// An invalid optimization level
+#[error("invalid optimization level")]
+pub struct InvalidOptimizationLevel;
+
+impl std::str::FromStr for Optimization {
+    type Err = InvalidOptimizationLevel;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        match s {
+            "0" | "O0" => Ok(Optimization::O0),
+            "1" | "O1" => Ok(Optimization::O1),
+            "2" | "O2" => Ok(Optimization::O2),
+            "3" | "O3" => Ok(Optimization::O3),
+            "s" | "Os" => Ok(Optimization::Os),
+            "z" | "Oz" => Ok(Optimization::Oz),
+            _ => Err(InvalidOptimizationLevel),
+        }
+    }
+}
+
+impl Display for Optimization {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        match *self {
+            Optimization::O0 => write!(f, "O0"),
+            Optimization::O1 => write!(f, "O1"),
+            Optimization::O2 => write!(f, "O2"),
+            Optimization::O3 => write!(f, "O3"),
+            Optimization::Os => write!(f, "Os"),
+            Optimization::Oz => write!(f, "Oz"),
+        }
+    }
+}
diff --git a/src/tools/llvm-bitcode-linker/src/target.rs b/src/tools/llvm-bitcode-linker/src/target.rs
new file mode 100644
index 00000000000..d9f8ff3852b
--- /dev/null
+++ b/src/tools/llvm-bitcode-linker/src/target.rs
@@ -0,0 +1,20 @@
+#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash, clap::ValueEnum)]
+pub enum Target {
+    Nvptx64NvidiaCuda,
+}
+
+#[derive(Debug, Clone, Copy, thiserror::Error)]
+/// The target is not supported by this linker
+#[error("unsupported target")]
+pub struct UnsupportedTarget;
+
+impl std::str::FromStr for Target {
+    type Err = UnsupportedTarget;
+
+    fn from_str(s: &str) -> Result<Target, UnsupportedTarget> {
+        match s {
+            "nvptx64-nvidia-cuda" => Ok(Target::Nvptx64NvidiaCuda),
+            _ => Err(UnsupportedTarget),
+        }
+    }
+}