Add offload support to the dist-x86_64-linux CI job

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

@@ -2281,9 +2281,9 @@ impl CommandLineStep for Assemble {
 
         if builder.config.llvm_offload && !builder.config.dry_run() {
             debug!("`llvm_offload` requested");
-            let rust_offload = builder.ensure(llvm::RustOffload { target: build_compiler.host });
-            let offload_install = builder.ensure(llvm::OmpOffload { target: build_compiler.host });
             if let Some(_llvm_config) = builder.llvm_config(builder.config.host_target) {
+                let rust_offload =
+                    builder.ensure(llvm::RustOffload { target: build_compiler.host });
                 let target_libdir =
                     builder.sysroot_target_libdir(target_compiler, target_compiler.host);
                 let rust_offload_dst_lib = target_libdir.join(rust_offload.rust_offload_filename());
@@ -2293,15 +2293,12 @@ impl CommandLineStep for Assemble {
                     FileType::NativeLibrary,
                 );
 
-                for p in offload_install.offload_paths() {
+                let omp_offload = builder.ensure(llvm::OmpOffload { target: build_compiler.host });
+                for p in omp_offload.artifact_paths_with_symlink_targets() {
                     let libname = p.file_name().unwrap();
                     let dst_lib = target_libdir.join(libname);
                     builder.resolve_symlink_and_copy(&p, &dst_lib);
                 }
-                // FIXME(offload): Add amdgcn-amd-amdhsa and nvptx64-nvidia-cuda folder
-                // This one is slightly more tricky, since we have the same file twice, in two
-                // subfolders for amdgcn and nvptx64. We'll likely find two more in the future, once
-                // Intel and Spir-V support lands in offload.
             }
         }
 

src/bootstrap/src/core/build_steps/dist.rs

@@ -2813,6 +2813,62 @@ impl CommandLineStep for Enzyme {
     }
 }
 
+#[derive(Debug, Clone, Hash, PartialEq, Eq)]
+pub struct Offload {
+    pub target: TargetSelection,
+}
+
+impl CommandLineStep for Offload {
+    type Output = Option<GeneratedTarball>;
+    const IS_HOST: bool = true;
+
+    fn should_run(run: ShouldRun<'_>) -> ShouldRun<'_> {
+        run.alias("offload")
+    }
+
+    fn is_default_step(builder: &Builder<'_>) -> bool {
+        builder.config.llvm_offload
+    }
+
+    fn make_run(run: RunConfig<'_>) {
+        run.builder.ensure(Offload { target: run.target });
+    }
+
+    fn run(self, builder: &Builder<'_>) -> Self::Output {
+        if !builder.unstable_features() {
+            return None;
+        }
+
+        let target = self.target;
+
+        let omp_offload = builder.ensure(llvm::OmpOffload { target });
+        let rust_offload = builder.ensure(llvm::RustOffload { target });
+
+        if builder.config.dry_run() {
+            return None;
+        }
+
+        let target_libdir = PathBuf::from(format!("lib/rustlib/{}/lib", target.triple));
+
+        let mut tarball = Tarball::new(builder, "offload", &target.triple);
+        tarball.set_overlay(OverlayKind::Offload);
+        tarball.is_preview(true);
+
+        let omp_offload_libdir = builder.out.join(target).join("offload").join("lib");
+
+        for path in omp_offload.artifact_paths_with_symlink_targets() {
+            let relative = t!(path.strip_prefix(&omp_offload_libdir));
+            let destdir = target_libdir.join(relative.parent().unwrap());
+
+            tarball.add_file(path, destdir, FileType::NativeLibrary);
+        }
+
+        tarball.add_file(rust_offload.rust_offload_path(), target_libdir, FileType::NativeLibrary);
+
+        Some(tarball.generate())
+    }
+}
+
 /// Tarball intended for internal consumption to ease rustc/std development.
 ///
 /// Should not be considered stable by end users.

src/bootstrap/src/core/build_steps/llvm.rs

@@ -1046,8 +1046,25 @@ pub struct BuiltOmpOffload {
 }
 
 impl BuiltOmpOffload {
-    pub fn offload_paths(&self) -> Vec<PathBuf> {
-        self.offload.clone()
+    pub fn artifact_paths_with_symlink_targets(&self) -> Vec<PathBuf> {
+        let mut paths = self.offload.clone();
+
+        for path in &self.offload {
+            let mut current = path.clone();
+
+            while t!(fs::symlink_metadata(&current)).file_type().is_symlink() {
+                let target = t!(fs::read_link(&current));
+                current = current.parent().unwrap().join(target);
+
+                if paths.contains(&current) {
+                    break;
+                }
+
+                paths.push(current.clone());
+            }
+        }
+
+        paths
     }
 }
 
@@ -1101,6 +1118,30 @@ impl CommandLineStep for OmpOffload {
         files.push(out_dir.join("lib").join("libLLVMOffload").with_extension(lib_ext));
         files.push(out_dir.join("lib").join("libomp").with_extension(lib_ext));
         files.push(out_dir.join("lib").join("libomptarget").with_extension(lib_ext));
+        files.push(
+            out_dir.join("lib").join("amdgcn-amd-amdhsa").join("libompdevice").with_extension("a"),
+        );
+        files.push(
+            out_dir
+                .join("lib")
+                .join("amdgcn-amd-amdhsa")
+                .join("libomptarget-amdgpu")
+                .with_extension("bc"),
+        );
+        files.push(
+            out_dir
+                .join("lib")
+                .join("nvptx64-nvidia-cuda")
+                .join("libompdevice")
+                .with_extension("a"),
+        );
+        files.push(
+            out_dir
+                .join("lib")
+                .join("nvptx64-nvidia-cuda")
+                .join("libomptarget-nvptx")
+                .with_extension("bc"),
+        );
 
         // Offload/OpenMP are just subfolders of LLVM, so we can use the LLVM sha.
         static STAMP_HASH_MEMO: OnceLock<String> = OnceLock::new();
@@ -1167,7 +1208,15 @@ impl CommandLineStep for OmpOffload {
                 cflags.push_all(format!(" -I {inc_dir}"));
             }
 
-            configure_cmake(builder, target, &mut cfg, true, LdFlags::default(), cflags, &[]);
+            // Logic copied from `configure_llvm`
+            // ThinLTO is only available when building with LLVM, enabling LLD is required.
+            // Apple's linker ld64 supports ThinLTO out of the box though, so don't use LLD on Darwin.
+            let mut ldflags = LdFlags::default();
+            if builder.config.llvm_thin_lto && !target.contains("apple") {
+                ldflags.push_all("-fuse-ld=lld");
+            }
+
+            configure_cmake(builder, target, &mut cfg, true, ldflags, cflags, &[]);
 
             // Re-use the same flags as llvm to control the level of debug information
             // generated for offload.
@@ -1196,6 +1245,7 @@ impl CommandLineStep for OmpOffload {
                 cfg.define("LLVM_ENABLE_RUNTIMES", "openmp;offload");
             } else {
                 // OpenMP provides some device libraries, so we also compile it for all gpu targets.
+                cfg.define("OPENMP_INSTALL_LIBDIR", Path::new("lib").join(omp_target));
                 cfg.define("LLVM_USE_LINKER", "lld");
                 cfg.define("LLVM_ENABLE_RUNTIMES", "openmp");
                 cfg.define("CMAKE_C_COMPILER_TARGET", omp_target);

src/bootstrap/src/core/builder/mod.rs

@@ -1018,6 +1018,7 @@ impl<'a> Builder<'a> {
                 dist::LlvmBitcodeLinker,
                 dist::RustDev,
                 dist::Enzyme,
+                dist::Offload,
                 dist::Bootstrap,
                 dist::Extended,
                 // It seems that PlainSourceTarball somehow changes how some of the tools

src/bootstrap/src/utils/tarball.rs

@@ -29,6 +29,7 @@ pub(crate) enum OverlayKind {
     Gcc,
     LlvmBitcodeLinker,
     Enzyme,
+    Offload,
 }
 
 impl OverlayKind {
@@ -39,6 +40,9 @@ impl OverlayKind {
                 &["src/llvm-project/llvm/LICENSE.TXT", "src/llvm-project/llvm/README.txt"]
             }
             OverlayKind::Enzyme => &["src/tools/enzyme/LICENSE", "src/tools/enzyme/Readme.md"],
+            OverlayKind::Offload => {
+                &["src/llvm-project/openmp/LICENSE.TXT", "src/llvm-project/offload/README.md"]
+            }
             OverlayKind::Cargo => &[
                 "src/tools/cargo/README.md",
                 "src/tools/cargo/LICENSE-MIT",
@@ -114,6 +118,7 @@ impl OverlayKind {
             OverlayKind::LlvmBitcodeLinker => builder.rust_version(),
             OverlayKind::Gcc => builder.rust_version(),
             OverlayKind::Enzyme => builder.rust_version(),
+            OverlayKind::Offload => builder.rust_version(),
         }
     }
 }

src/ci/docker/host-x86_64/dist-x86_64-linux/Dockerfile

@@ -65,7 +65,7 @@ RUN ./cmake.sh
 # Now build LLVM+Clang, afterwards configuring further compilations to use the
 # clang/clang++ compilers.
 COPY scripts/build-clang.sh /tmp/
-ENV LLVM_BUILD_TARGETS=X86
+ENV LLVM_BUILD_TARGETS="X86;AMDGPU;NVPTX"
 RUN ./build-clang.sh
 ENV CC=clang CXX=clang++
 
@@ -91,6 +91,7 @@ ENV RUST_CONFIGURE_ARGS="--enable-full-tools \
       --set llvm.ninja=false \
       --set llvm.libzstd=true \
       --set build.allocator=jemalloc \
+      --set llvm.offload-clang-dir="/rustroot/lib/cmake/clang" \
       --set rust.bootstrap-override-lld=true \
       --set rust.lto=thin \
       --set rust.codegen-units=1"

src/ci/docker/host-x86_64/dist-x86_64-linux/dist.sh

@@ -10,6 +10,7 @@ python3 ../x.py build --set rust.debug=true opt-dist
     build-manifest \
     bootstrap \
     enzyme \
+    offload \
     rustc_codegen_gcc
 
 # Use GCC for building GCC components, as it seems to behave badly when built with Clang