Only try to link against offload functions if llvm.enzyme is enabled

Author: ZuseZ4

compiler/rustc/Cargo.toml

@@ -31,6 +31,7 @@ check_only = ['rustc_driver_impl/check_only']
 jemalloc = ['dep:tikv-jemalloc-sys']
 llvm = ['rustc_driver_impl/llvm']
 llvm_enzyme = ['rustc_driver_impl/llvm_enzyme']
+llvm_offload = ['rustc_driver_impl/llvm_offload']
 max_level_info = ['rustc_driver_impl/max_level_info']
 rustc_randomized_layouts = ['rustc_driver_impl/rustc_randomized_layouts']
 # tidy-alphabetical-end

compiler/rustc_codegen_llvm/Cargo.toml

@@ -47,5 +47,6 @@ tracing = "0.1"
 # tidy-alphabetical-start
 check_only = ["rustc_llvm/check_only"]
 llvm_enzyme = []
+llvm_offload = []
 # tidy-alphabetical-end
 

compiler/rustc_codegen_llvm/src/llvm/ffi.rs

@@ -1641,9 +1641,6 @@ unsafe extern "C" {
         Name: *const c_char,
     ) -> &'a Value;
 
-    /// Processes the module and writes it in an offload compatible way into a "host.out" file.
-    pub(crate) fn LLVMRustBundleImages<'a>(M: &'a Module, TM: &'a TargetMachine) -> bool;
-
     /// Writes a module to the specified path. Returns 0 on success.
     pub(crate) fn LLVMWriteBitcodeToFile(M: &Module, Path: *const c_char) -> c_int;
 
@@ -1721,6 +1718,37 @@ unsafe extern "C" {
     ) -> &'a Value;
 }
 
+#[cfg(feature = "llvm_offload")]
+pub(crate) use self::Offload::*;
+
+#[cfg(feature = "llvm_offload")]
+mod Offload {
+    use super::*;
+    unsafe extern "C" {
+        /// Processes the module and writes it in an offload compatible way into a "host.out" file.
+        pub(crate) fn LLVMRustBundleImages<'a>(M: &'a Module, TM: &'a TargetMachine) -> bool;
+        pub(crate) fn LLVMRustOffloadMapper<'a>(OldFn: &'a Value, NewFn: &'a Value);
+    }
+}
+
+#[cfg(not(feature = "llvm_offload"))]
+pub(crate) use self::Offload_fallback::*;
+
+#[cfg(not(feature = "llvm_offload"))]
+mod Offload_fallback {
+    use super::*;
+    /// Processes the module and writes it in an offload compatible way into a "host.out" file.
+    /// Marked as unsafe to match the real offload wrapper which is unsafe due to FFI.
+    #[allow(unused_unsafe)]
+    pub(crate) unsafe fn LLVMRustBundleImages<'a>(_M: &'a Module, _TM: &'a TargetMachine) -> bool {
+        unimplemented!("This rustc version was not built with LLVM Offload support!");
+    }
+    #[allow(unused_unsafe)]
+    pub(crate) unsafe fn LLVMRustOffloadMapper<'a>(_OldFn: &'a Value, _NewFn: &'a Value) {
+        unimplemented!("This rustc version was not built with LLVM Offload support!");
+    }
+}
+
 // FFI bindings for `DIBuilder` functions in the LLVM-C API.
 // Try to keep these in the same order as in `llvm/include/llvm-c/DebugInfo.h`.
 //
@@ -2028,7 +2056,6 @@ unsafe extern "C" {
     ) -> &Attribute;
 
     // Operations on functions
-    pub(crate) fn LLVMRustOffloadMapper<'a>(Fn: &'a Value, Fn: &'a Value);
     pub(crate) fn LLVMRustGetOrInsertFunction<'a>(
         M: &'a Module,
         Name: *const c_char,

compiler/rustc_driver_impl/Cargo.toml

@@ -75,6 +75,7 @@ ctrlc = "3.4.4"
 check_only = ['rustc_interface/check_only']
 llvm = ['rustc_interface/llvm']
 llvm_enzyme = ['rustc_interface/llvm_enzyme']
+llvm_offload = ['rustc_interface/llvm_offload']
 max_level_info = ['rustc_log/max_level_info']
 rustc_randomized_layouts = [
     'rustc_index/rustc_randomized_layouts',

compiler/rustc_interface/Cargo.toml

@@ -59,4 +59,5 @@ rustc_abi = { path = "../rustc_abi" }
 check_only = ['rustc_codegen_llvm?/check_only']
 llvm = ['dep:rustc_codegen_llvm']
 llvm_enzyme = ['rustc_builtin_macros/llvm_enzyme', 'rustc_codegen_llvm/llvm_enzyme']
+llvm_offload = ['rustc_codegen_llvm/llvm_offload']
 # tidy-alphabetical-end

compiler/rustc_llvm/build.rs

@@ -214,6 +214,10 @@ fn main() {
         cfg.define("ENZYME", None);
     }
 
+    if tracked_env_var_os("LLVM_OFFLOAD").is_some() {
+        cfg.define("OFFLOAD", None);
+    }
+
     if tracked_env_var_os("LLVM_RUSTLLVM").is_some() {
         cfg.define("LLVM_RUSTLLVM", None);
     }

compiler/rustc_llvm/llvm-wrapper/RustWrapper.cpp

@@ -25,7 +25,6 @@
 #include "llvm/IR/Module.h"
 #include "llvm/IR/Value.h"
 #include "llvm/Object/COFFImportFile.h"
-#include "llvm/Object/OffloadBinary.h"
 #include "llvm/Remarks/RemarkFormat.h"
 #include "llvm/Remarks/RemarkSerializer.h"
 #include "llvm/Remarks/RemarkStreamer.h"
@@ -36,11 +35,18 @@
 #include "llvm/Support/Signals.h"
 #include "llvm/Support/Timer.h"
 #include "llvm/Support/ToolOutputFile.h"
-#include "llvm/Target/TargetMachine.h"
 #include "llvm/Transforms/Utils/Cloning.h"
 #include "llvm/Transforms/Utils/ValueMapper.h"
 #include <iostream>
 
+// Some of the functions below rely on LLVM modules that may not always be
+// available. As such, we only try to build it in the first place, if
+// llvm.offload is enabled.
+#ifdef OFFLOAD
+#include "llvm/Object/OffloadBinary.h"
+#include "llvm/Target/TargetMachine.h"
+#endif
+
 // for raw `write` in the bad-alloc handler
 #ifdef _MSC_VER
 #include <io.h>
@@ -146,6 +152,10 @@ extern "C" void LLVMRustPrintStatistics(RustStringRef OutBuf) {
   llvm::PrintStatistics(OS);
 }
 
+// Some of the functions here rely on LLVM modules that may not always be
+// available. As such, we only try to build it in the first place, if
+// llvm.offload is enabled.
+#ifdef OFFLOAD
 static Error writeFile(StringRef Filename, StringRef Data) {
   Expected<std::unique_ptr<FileOutputBuffer>> OutputOrErr =
       FileOutputBuffer::create(Filename, Data.size());
@@ -210,6 +220,7 @@ extern "C" void LLVMRustOffloadMapper(LLVMValueRef OldFn, LLVMValueRef NewFn) {
                           llvm::CloneFunctionChangeType::LocalChangesOnly,
                           returns);
 }
+#endif
 
 extern "C" LLVMValueRef LLVMRustGetNamedValue(LLVMModuleRef M, const char *Name,
                                               size_t NameLen) {

src/bootstrap/src/core/build_steps/compile.rs

@@ -1436,6 +1436,9 @@ fn rustc_llvm_env(builder: &Builder<'_>, cargo: &mut Cargo, target: TargetSelect
     if builder.config.llvm_enzyme {
         cargo.env("LLVM_ENZYME", "1");
     }
+    if builder.config.llvm_offload {
+        cargo.env("LLVM_OFFLOAD", "1");
+    }
     let llvm::LlvmResult { host_llvm_config, .. } = builder.ensure(llvm::Llvm { target });
     cargo.env("LLVM_CONFIG", &host_llvm_config);
 

src/bootstrap/src/lib.rs

@@ -873,6 +873,9 @@ impl Build {
         if self.config.llvm_enzyme {
             features.push("llvm_enzyme");
         }
+        if self.config.llvm_offload {
+            features.push("llvm_offload");
+        }
         // keep in sync with `bootstrap/compile.rs:rustc_cargo_env`
         if self.config.rust_randomize_layout && check("rustc_randomized_layouts") {
             features.push("rustc_randomized_layouts");