Add unstable -Zllvm-tpde flag to rustc_codegen_llvm

compiler/rustc_codegen_llvm/src/back/llvm_backend.rs

@@ -22,12 +22,12 @@ use rustc_errors::{DiagCtxt, DiagCtxtHandle};
 use rustc_metadata::EncodedMetadata;
 use rustc_middle::dep_graph::{WorkProduct, WorkProductMap};
 use rustc_middle::ty::TyCtxt;
-use rustc_session::config::{OptLevel, OutputFilenames, PrintKind, PrintRequest};
+use rustc_session::config::{LtoCli, OptLevel, OutputFilenames, PrintKind, PrintRequest};
 use rustc_session::{
     BorrowedIncrCompSession, CodegenBackendInit, EarlySession, IncrCompSession, Session,
 };
 use rustc_span::{Symbol, sym};
-use rustc_target::spec::{RelocModel, TlsModel};
+use rustc_target::spec::{RelocModel, SplitDebuginfo, TlsModel};
 
 use crate::back::owned_target_machine::OwnedTargetMachine;
 use crate::back::write::{create_informational_target_machine, create_target_machine};
@@ -213,6 +213,43 @@ impl CodegenBackend for LlvmCodegenBackend {
             }
         }
 
+        // -Zllvm-tpde requires the TPDE pass plugin, which may not be available. If TPDE is used but not available,
+        // we emit an early error here and abort compilation.
+        {
+            match llvm::TpdeWrapper::get_or_init(
+                &sess.opts.sysroot,
+                &sess.opts.unstable_opts.llvm_tpde,
+            ) {
+                Ok(_) => {}
+                Err(llvm::TpdeLibraryError::NotFound { err }) => {
+                    sess.dcx().emit_fatal(crate::diagnostics::TpdeComponentMissing { err });
+                }
+                Err(llvm::TpdeLibraryError::LoadFailed { err }) => {
+                    sess.dcx().emit_fatal(crate::diagnostics::TpdeComponentUnavailable { err });
+                }
+            }
+
+            if sess.opts.unstable_opts.llvm_tpde.is_some() {
+                if !matches!(sess.opts.optimize, OptLevel::No | OptLevel::Less) {
+                    sess.dcx().emit_warn(crate::diagnostics::TpdeHighOptLevel);
+                }
+
+                if sess.target.requires_lto
+                    || !matches!(sess.early_lto(), LtoCli::No | LtoCli::Unspecified)
+                {
+                    sess.dcx().emit_fatal(crate::diagnostics::TpdeLtoUsed);
+                }
+
+                // JUSTIFICATION: the session doesn't exist at this point.
+                #[allow(rustc::bad_opt_access)]
+                if sess.opts.cg.split_debuginfo.unwrap_or(sess.target.split_debuginfo)
+                    != SplitDebuginfo::Off
+                {
+                    sess.dcx().emit_fatal(crate::diagnostics::TpdeSplitDebugInfoUsed);
+                }
+            }
+        }
+
         // Intrinsics whose fallback body will not be used by the LLVM backend.
         let replaced_intrinsics = {
             #[rustfmt::skip]

compiler/rustc_codegen_llvm/src/back/write.rs

@@ -40,7 +40,7 @@ use crate::diagnostics::{
     UnsupportedCompression, WithLlvmError, WriteBytecode,
 };
 use crate::llvm::diagnostic::OptimizationDiagnosticKind::*;
-use crate::llvm::{self, DiagnosticInfo};
+use crate::llvm::{self, DiagnosticInfo, TpdeWrapper};
 use crate::type_::llvm_type_ptr;
 use crate::{LlvmCodegenBackend, ModuleLlvm, attributes, base, common, llvm_util};
 
@@ -71,6 +71,8 @@ fn write_output_file<'ll>(
     } else {
         std::ptr::null()
     };
+
+    let tpde = TpdeWrapper::get_instance();
     let result = unsafe {
         llvm::LLVMRustWriteOutputFile(
             target,
@@ -80,6 +82,7 @@ fn write_output_file<'ll>(
             file_type,
             verify_llvm_ir,
             no_builtins,
+            tpde.options().clone(),
         )
     };
 

compiler/rustc_codegen_llvm/src/diagnostics.rs

@@ -107,6 +107,33 @@ pub(crate) struct LtoBitcodeFromRlib {
     pub err: String,
 }
 
+#[derive(Diagnostic)]
+#[diag("failed to load the TPDE backend: {$err}")]
+pub(crate) struct TpdeComponentUnavailable {
+    pub err: String,
+}
+
+#[derive(Diagnostic)]
+#[diag("TPDE backend not found in the sysroot: {$err}")]
+pub(crate) struct TpdeComponentMissing {
+    pub err: String,
+}
+
+#[derive(Diagnostic)]
+#[diag("using -Zllvm-tpde with optimization levels higher than 1 is not recommended")]
+#[note(
+    "compilation with TPDE at higher opt levels is more likely to fail, and to have worse runtime performance than LLVM even if it succeeds"
+)]
+pub(crate) struct TpdeHighOptLevel;
+
+#[derive(Diagnostic)]
+#[diag("using -Zllvm-tpde with LTO enabled is not supported")]
+pub(crate) struct TpdeLtoUsed;
+
+#[derive(Diagnostic)]
+#[diag("using -Zllvm-tpde with split debuginfo is not supported")]
+pub(crate) struct TpdeSplitDebugInfoUsed;
+
 #[derive(Diagnostic)]
 pub(crate) enum LlvmError<'a> {
     #[diag("could not write output to {$path}")]

compiler/rustc_codegen_llvm/src/llvm/ffi.rs

@@ -25,8 +25,7 @@ use super::debuginfo::{
     DIArray, DIBuilder, DIDerivedType, DIDescriptor, DIFile, DIFlags, DILocation, DISPFlags,
     DIScope, DISubprogram, DITemplateTypeParameter, DIType, DebugEmissionKind, DebugNameTableKind,
 };
-use crate::llvm;
-use crate::llvm::MetadataKindId;
+use crate::llvm::{self, LLVMRustTpdeOptions, MetadataKindId};
 use crate::macros::TryFromU32;
 
 /// In the LLVM-C API, boolean values are passed as `typedef int LLVMBool`,
@@ -2439,6 +2438,7 @@ unsafe extern "C" {
         FileType: FileType,
         VerifyIR: bool,
         DisableSimplifyLibCalls: bool,
+        TpdeOptions: LLVMRustTpdeOptions,
     ) -> LLVMRustResult;
     pub(crate) fn LLVMRustOptimize<'a>(
         M: &'a Module,

compiler/rustc_codegen_llvm/src/llvm/mod.rs

@@ -21,9 +21,11 @@ pub(crate) mod enzyme_ffi;
 mod ffi;
 mod metadata_kind;
 pub(crate) mod offload_ffi;
+pub(crate) mod tpde_ffi;
 
 pub(crate) use self::enzyme_ffi::*;
 pub(crate) use self::offload_ffi::*;
+pub(crate) use self::tpde_ffi::*;
 
 impl LLVMRustResult {
     pub(crate) fn into_result(self) -> Result<(), ()> {

compiler/rustc_codegen_llvm/src/llvm/tpde_ffi.rs

@@ -0,0 +1,125 @@
+use std::path::PathBuf;
+use std::sync::OnceLock;
+
+use rustc_session::config::{self, Sysroot, host_tuple};
+use rustc_session::filesearch;
+
+unsafe extern "C" {
+    type LLVMPassPluginLibraryInfo;
+}
+
+type LLVMGetPassPluginInfoFn = unsafe extern "C" fn() -> LLVMPassPluginLibraryInfo;
+
+/// LLVMRustTpdeMode
+#[repr(C)]
+#[derive(Clone)]
+pub(crate) enum LLVMRustTpdeMode {
+    None,
+    Only,
+    Try,
+}
+
+/// LLVMRustTpdeOptions
+#[repr(C)]
+#[derive(Clone)]
+pub(crate) struct LLVMRustTpdeOptions {
+    mode: LLVMRustTpdeMode,
+    // `None` if mode is `LLVMRustTpdeMode::None`
+    plugin: Option<LLVMGetPassPluginInfoFn>,
+}
+
+pub(crate) struct TpdeWrapper {
+    options: LLVMRustTpdeOptions,
+    // Keep the dynamic library loaded while the function pointers are used.
+    _lib: Option<libloading::Library>,
+}
+
+static TPDE_INSTANCE: OnceLock<TpdeWrapper> = OnceLock::new();
+
+#[derive(Debug)]
+pub(crate) enum TpdeLibraryError {
+    NotFound { err: String },
+    LoadFailed { err: String },
+}
+
+impl From<libloading::Error> for TpdeLibraryError {
+    fn from(err: libloading::Error) -> Self {
+        Self::LoadFailed { err: format!("{err:?}") }
+    }
+}
+
+impl TpdeWrapper {
+    /// Initialize TpdeWrapper with the given sysroot and mode if not already initialized.
+    /// Safe to call multiple times - subsequent calls are no-ops due to OnceLock.
+    pub(crate) fn get_or_init(
+        sysroot: &rustc_session::config::Sysroot,
+        mode: &Option<config::LlvmTpde>,
+    ) -> Result<&'static Self, TpdeLibraryError> {
+        let wrapper = TPDE_INSTANCE.get_or_try_init(|| {
+            let w = match mode {
+                Some(mode) => Self::call_dynamic(sysroot, mode)?,
+                None => Self {
+                    options: LLVMRustTpdeOptions { mode: LLVMRustTpdeMode::None, plugin: None },
+                    _lib: None,
+                },
+            };
+            Ok::<_, TpdeLibraryError>(w)
+        })?;
+
+        Ok(wrapper)
+    }
+
+    /// Get the TpdeWrapper instance. Panics if not initialized.
+    pub(crate) fn get_instance() -> &'static Self {
+        TPDE_INSTANCE
+            .get()
+            .expect("TpdeWrapper not initialized. Call get_or_init with sysroot first.")
+    }
+
+    pub(crate) fn options(&self) -> &LLVMRustTpdeOptions {
+        &self.options
+    }
+
+    fn call_dynamic(
+        sysroot: &rustc_session::config::Sysroot,
+        mode: &config::LlvmTpde,
+    ) -> Result<Self, TpdeLibraryError> {
+        let mode = match mode {
+            config::LlvmTpde::Only => LLVMRustTpdeMode::Only,
+            config::LlvmTpde::Try => LLVMRustTpdeMode::Try,
+        };
+        let tpde_path = Self::get_tpde_path(sysroot)?;
+        let lib = unsafe { libloading::Library::new(tpde_path)? };
+        let llvm_get_pass_plugin_info =
+            *unsafe { lib.get::<LLVMGetPassPluginInfoFn>(b"llvmGetPassPluginInfo\0")? };
+        let options = LLVMRustTpdeOptions { mode, plugin: Some(llvm_get_pass_plugin_info) };
+        Ok(Self { options, _lib: Some(lib) })
+    }
+
+    fn get_tpde_path(sysroot: &Sysroot) -> Result<PathBuf, TpdeLibraryError> {
+        let path_buf = sysroot
+            .all_paths()
+            .map(|sysroot_path| {
+                filesearch::make_target_lib_path(sysroot_path, host_tuple())
+                    .join("lib")
+                    .with_file_name("tpde-plugin")
+                    .with_extension(std::env::consts::DLL_EXTENSION)
+            })
+            .find(|f| f.exists())
+            .ok_or_else(|| {
+                let candidates = sysroot
+                    .all_paths()
+                    .map(|p| p.join("lib").display().to_string())
+                    .collect::<Vec<String>>()
+                    .join("\n* ");
+                TpdeLibraryError::NotFound {
+                    err: format!(
+                        "failed to find the `tpde-plugin` shared library \
+                    in the sysroot candidates:\n* {candidates}"
+                    ),
+                }
+            })?;
+
+        Ok(path_buf)
+    }
+}

compiler/rustc_interface/src/tests.rs

@@ -14,10 +14,10 @@ use rustc_session::config::{
     CoverageOptions, DebugInfo, DumpMonoStatsFormat, ErrorOutputType, ExternEntry, ExternLocation,
     Externs, FmtDebug, FunctionReturn, IncrementalStateAssertion, InliningThreshold, Input,
     InstrumentCoverage, InstrumentMcount, InstrumentMcountOpts, InstrumentXRay, LinkSelfContained,
-    LinkerPluginLto, LocationDetail, LtoCli, MirIncludeSpans, NextSolverConfig, Offload, Options,
-    OutFileName, OutputType, OutputTypes, PAuthKey, PacRet, Passes, PatchableFunctionEntry,
-    Polonius, ProcMacroExecutionStrategy, Strip, SwitchWithOptPath, SymbolManglingVersion,
-    WasiExecModel, build_session_options, rustc_optgroups,
+    LinkerPluginLto, LlvmTpde, LocationDetail, LtoCli, MirIncludeSpans, NextSolverConfig, Offload,
+    Options, OutFileName, OutputType, OutputTypes, PAuthKey, PacRet, Passes,
+    PatchableFunctionEntry, Polonius, ProcMacroExecutionStrategy, Strip, SwitchWithOptPath,
+    SymbolManglingVersion, WasiExecModel, build_session_options, rustc_optgroups,
 };
 use rustc_session::search_paths::SearchPath;
 use rustc_session::utils::{CanonicalizedPath, NativeLib};
@@ -846,6 +846,8 @@ fn test_unstable_options_tracking_hash() {
     tracked!(lint_llvm_ir, true);
     tracked!(llvm_module_flag, vec![("bar".to_string(), 123, "max".to_string())]);
     tracked!(llvm_plugins, vec![String::from("plugin_name")]);
+    tracked!(llvm_tpde, Some(LlvmTpde::Only));
+    tracked!(llvm_tpde, Some(LlvmTpde::Try));
     tracked!(llvm_writable, true);
     tracked!(location_detail, LocationDetail { file: true, line: false, column: false });
     tracked!(maximal_hir_to_mir_coverage, true);

compiler/rustc_llvm/llvm-wrapper/PassWrapper.cpp

@@ -465,11 +465,23 @@ static CodeGenFileType fromRust(LLVMRustFileType Type) {
   }
 }
 
-extern "C" LLVMRustResult
-LLVMRustWriteOutputFile(LLVMTargetMachineRef Target, LLVMModuleRef M,
-                        const char *Path, const char *DwoPath,
-                        LLVMRustFileType RustFileType, bool VerifyIR,
-                        bool DisableSimplifyLibCalls) {
+extern "C" typedef llvm::PassPluginLibraryInfo (*LLVMGetPassPluginInfo)();
+
+enum class LLVMRustTpdeMode {
+  None,
+  Only,
+  Try,
+};
+
+struct LLVMRustTpdeOptions {
+  LLVMRustTpdeMode mode;
+  LLVMGetPassPluginInfo plugin;
+};
+
+extern "C" LLVMRustResult LLVMRustWriteOutputFile(
+    LLVMTargetMachineRef Target, LLVMModuleRef M, const char *Path,
+    const char *DwoPath, LLVMRustFileType RustFileType, bool VerifyIR,
+    bool DisableSimplifyLibCalls, LLVMRustTpdeOptions TpdeOptions) {
   std::unique_ptr<llvm::legacy::PassManager> PM =
       std::make_unique<llvm::legacy::PassManager>();
 
@@ -507,9 +519,25 @@ LLVMRustWriteOutputFile(LLVMTargetMachineRef Target, LLVMModuleRef M,
   }
 
   // TargetMachine::addPassesToEmitFile stores pointers to the output streams
-  // in a couple of places inside of the object. Explicitly delete the PM after
-  // we call run() to avoid dangling references.
+  // in a couple of places inside of the object. Explicitly delete the PM
+  // after we call run() to avoid dangling references.
   auto BOS = buffer_ostream(OS);
+
+  if (TpdeOptions.mode != LLVMRustTpdeMode::None) {
+    assert(
+        TpdeOptions.plugin &&
+        "rustc_codegen_llvm must set TpdeOptions.plugin if using -Zllvm-tpde");
+    llvm::PassPluginLibraryInfo plugin = TpdeOptions.plugin();
+    if (plugin.PreCodeGenCallback(*unwrap(M), *unwrap(Target), FileType, BOS)) {
+      return LLVMRustResult::Success;
+    } else if (TpdeOptions.mode == LLVMRustTpdeMode::Only) {
+      std::string error = "Failed to compile module " +
+                          unwrap(M)->getModuleIdentifier() + " with TPDE";
+      LLVMRustSetLastError(error.c_str());
+      return LLVMRustResult::Failure;
+    }
+  }
+
   if (DwoPath) {
     auto DOS = raw_fd_ostream(DwoPath, EC, sys::fs::OF_None);
     EC.clear();

compiler/rustc_macros/src/diagnostics/message.rs

@@ -138,6 +138,7 @@ const ALLOWED_CAPITALIZED_WORDS: &[&str] = &[
     "NaNs",
     "OK",
     "Rust",
+    "TPDE",
     "ThinLTO",
     "Unicode",
     "VS",

compiler/rustc_session/src/config.rs

@@ -3373,6 +3373,14 @@ impl PpMode {
     }
 }
 
+#[derive(Clone, Hash, PartialEq, Eq, Debug)]
+pub enum LlvmTpde {
+    /// Attempt to compile with TPDE-LLVM, but fall back to LLVM if unsupported LLVM IR is encountered
+    Try,
+    /// Attempt to compile with TPDE-LLVM, and bail if unsupported LLVM IR is encountered
+    Only,
+}
+
 #[derive(Clone, Hash, PartialEq, Eq, Debug)]
 pub enum WasiExecModel {
     Command,
@@ -3423,7 +3431,7 @@ pub(crate) mod dep_tracking {
         AnnotateMoves, AutoDiff, BranchProtection, CFGuard, CFProtection, CodegenRetagOptions,
         CoverageOptions, CrateType, DebugInfo, DebugInfoCompression, ErrorOutputType, FmtDebug,
         FunctionReturn, InliningThreshold, InstrumentCoverage, InstrumentMcount,
-        InstrumentMcountOpts, InstrumentXRay, LinkerPluginLto, LocationDetail, LtoCli,
+        InstrumentMcountOpts, InstrumentXRay, LinkerPluginLto, LlvmTpde, LocationDetail, LtoCli,
         MirStripDebugInfo, NextSolverConfig, Offload, OptLevel, OutFileName, OutputType,
         OutputTypes, PatchableFunctionEntry, PointerAuthOption, Polonius, ResolveDocLinks,
         SourceFileHashAlgorithm, SplitDwarfKind, SwitchWithOptPath, SymbolManglingVersion,
@@ -3496,6 +3504,7 @@ pub(crate) mod dep_tracking {
         InstrumentXRay,
         LanguageIdentifier,
         LinkerPluginLto,
+        LlvmTpde,
         LocationDetail,
         LtoCli,
         MergeFunctions,