Auto merge of #159064 - sgasho:offload-dist, r=<try>

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


try-job: dist-x86_64-linux

Author: bors

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

@@ -2278,15 +2278,14 @@ impl Step for Assemble {
             if let Some(_llvm_config) = builder.llvm_config(builder.config.host_target) {
                 let target_libdir =
                     builder.sysroot_target_libdir(target_compiler, target_compiler.host);
-                for p in offload_install.offload_paths() {
-                    let libname = p.file_name().unwrap();
-                    let dst_lib = target_libdir.join(libname);
+                let offload_libdir = builder.offload_out(build_compiler.host).join("lib");
+
+                for p in offload_install.artifact_paths_with_symlink_targets() {
+                    let relative = t!(p.strip_prefix(&offload_libdir));
+                    let dst_lib = target_libdir.join(relative);
+                    t!(fs::create_dir_all(dst_lib.parent().unwrap()));
                     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

@@ -2811,6 +2811,59 @@ impl Step for Enzyme {
     }
 }
 
+#[derive(Debug, Clone, Hash, PartialEq, Eq)]
+pub struct Offload {
+    pub target: TargetSelection,
+}
+
+impl Step 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 offload = builder.ensure(llvm::OmpOffload { target });
+
+        if builder.config.dry_run() {
+            return None;
+        }
+
+        let target_libdir = 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 offload_libdir = builder.offload_out(target).join("lib");
+
+        for path in offload.artifact_paths_with_symlink_targets() {
+            let relative = t!(path.strip_prefix(&offload_libdir));
+            let destdir = Path::new(&target_libdir).join(relative.parent().unwrap());
+
+            tarball.add_file(path, destdir, 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

@@ -941,8 +941,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
     }
 }
 
@@ -997,6 +1014,30 @@ impl Step 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();
@@ -1063,7 +1104,15 @@ impl Step 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.
@@ -1097,6 +1146,7 @@ impl Step 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

@@ -1017,6 +1017,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++
 
@@ -90,6 +90,7 @@ ENV RUST_CONFIGURE_ARGS="--enable-full-tools \
       --set llvm.thin-lto=true \
       --set llvm.ninja=false \
       --set llvm.libzstd=true \
+      --set llvm.offload-clang-dir="/rustroot/lib/cmake/clang" \
       --set rust.override-allocator=jemalloc \
       --set rust.bootstrap-override-lld=true \
       --set rust.lto=thin \

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

src/ci/docker/scripts/build-gcc.sh

@@ -1,4 +1,5 @@
 #!/usr/bin/env bash
+# ignore-tidy-file-linelength
 set -eux
 
 source shared.sh
@@ -29,7 +30,7 @@ cd gcc-$GCC
 # the expression will need to be updated. That new URL is available at:
 # https://github.com/gcc-mirror/gcc/blob/6e6e3f144a33ae504149dc992453b4f6dea12fdb/contrib/download_prerequisites#L35
 #
-sed -i'' 's|ftp://gcc\.gnu\.org/|https://gcc.gnu.org/|g' ./contrib/download_prerequisites
+sed -i'' 's|ftp://gcc\.gnu\.org/pub/gcc/infrastructure|https://ci-mirrors.rust-lang.org/rustc/gcc|g' ./contrib/download_prerequisites
 
 ./contrib/download_prerequisites
 mkdir ../gcc-build