gpu offload host code generation #142097

ZuseZ4 · 2025-06-05T20:01:04Z

r? ghost

This will generate most of the host side code to use llvm's offload feature.
The first PR will only handle automatic mem-transfers to and from the device.
So if a user calls a kernel, we will copy inputs back and forth, but we won't do the actual kernel launch.
Before merging, we will use LLVM's Info infrastructure to verify that the memcopies match what openmp offloa generates in C++. LIBOMPTARGET_INFO=-1 ./my_rust_binary should print that a memcpy to and later from the device is happening.

A follow-up PR will generate the actual device-side kernel which will then do computations on the GPU.
A third PR will implement manual host2device and device2host functionality, but the goal is to minimize cases where a user has to overwrite our default handling due to performance issues.

I'm trying to get a full MVP out first, so this just recognizes GPU functions based on magic names. The final frontend will obviously move this over to use proper macros, like I'm already doing it for the autodiff work.
This work will also be compatible with std::autodiff, so one can differentiate GPU kernels.

Tracking:

Tracking Issue for GPU-offload #131513

ZuseZ4 · 2025-06-10T03:09:23Z

@oli-obk Featurewise, I am almost done. I'll add a few more lines to describe the layout of Rust types to the offload library, but in this PR I only intend to support one type or two (maybe array's, raw pointer, or slices). I might even hardcode the length in the very first approach. In a follow-up PR I'll do some proper type parsing on a higher level, similar to what I did in the past with Rust TypeTrees. This work is much simpler and more reliable though, since offload doesn't care what type something has, just how many bytes it is large, and hence need to be moved to/from the GPU.

I was able to just move a few of the builder methods I needed to the generic builder.
However, there are also around 7 that I had to duplicate. I guess at some point I'll need to do the proper work of enabling the trait implementations for both builders :/
Once I have everything working, I'll clean it up and add some tests and docs.

compiler/rustc_codegen_llvm/src/back/lto.rs

compiler/rustc_codegen_llvm/src/context.rs

ZuseZ4 · 2025-06-12T14:32:20Z

Not fully ready yet, I apparently missed yet another global to initialize the offload runtime. But at least it compiles successfully to a binary if I emit the IR from Rust, and then use clang for the rest. I'll add the global today, then I should be done and will clean it up

bors · 2025-06-18T04:08:45Z

☔ The latest upstream changes (presumably #142644) made this pull request unmergeable. Please resolve the merge conflicts.

ZuseZ4 · 2025-06-18T21:57:35Z

Jay, turns out the only issue in my test binary was a bug in LLVM, which was already fixed upstream in llvm/llvm-project#143638.
Once rustc syncs the llvm submodule again (in a week or so), we should get the fix. This does not affect the llvm-ir we generate with rustc, it just affects how clang compiles the llvm-ir from rustc to a binary. Therefore we don't have to wait for it. I'll add an llvm-ir test to make sure we generate the right things and clean it up a bit more.

rust-log-analyzer · 2025-06-18T23:54:51Z

The job mingw-check-tidy failed! Check out the build log: (web) (plain)

Click to see the possible cause of the failure (guessed by this bot)

fmt: checked 6074 files
tidy check
Running eslint on rustdoc JS files
tidy: Skipping binary file check, read-only filesystem
##[error]tidy error: /checkout/compiler/rustc_codegen_llvm/src/back/lto.rs:672: `dbg!` macro is intended as a debugging tool. It should not be in version control.
##[error]tidy error: /checkout/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs:121: TODO is used for tasks that should be done before merging a PR; If you want to leave a message in the codebase use FIXME
##[error]tidy error: /checkout/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs:157: `dbg!` macro is intended as a debugging tool. It should not be in version control.
##[error]tidy error: /checkout/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs:167: `dbg!` macro is intended as a debugging tool. It should not be in version control.
##[error]tidy error: /checkout/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs:250: `dbg!` macro is intended as a debugging tool. It should not be in version control.
##[error]tidy error: /checkout/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs:457: `dbg!` macro is intended as a debugging tool. It should not be in version control.
##[error]tidy error: /checkout/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs:467: `dbg!` macro is intended as a debugging tool. It should not be in version control.
removing old virtual environment
creating virtual environment at '/checkout/obj/build/venv' using 'python3.10' and 'venv'
creating virtual environment at '/checkout/obj/build/venv' using 'python3.10' and 'virtualenv'
Requirement already satisfied: pip in ./build/venv/lib/python3.10/site-packages (25.1.1)
linting python files
All checks passed!
checking python file formatting
28 files already formatted
checking C++ file formatting
some tidy checks failed
Command has failed. Rerun with -v to see more details.
Build completed unsuccessfully in 0:01:25
  local time: Wed Jun 18 23:54:41 UTC 2025
  network time: Wed, 18 Jun 2025 23:54:41 GMT
##[error]Process completed with exit code 1.
Post job cleanup.

rustbot · 2025-06-19T00:39:01Z

Some changes occurred in compiler/rustc_codegen_ssa

cc @WaffleLapkin

ZuseZ4

I did the first round of reviews for myself, I'll address them tomorrow.
I'll also clean up the code in gpu_builder more, it has a lot of duplications and IR comments from when I was trying to figure out what to generate..

ZuseZ4 · 2025-06-19T00:40:07Z

compiler/rustc_codegen_llvm/src/builder.rs

@@ -117,6 +118,70 @@ impl<'a, 'll, CX: Borrow<SCx<'ll>>> GenericBuilder<'a, 'll, CX> {
        }
        bx
    }
+
+    pub(crate) fn my_alloca2(&mut self, ty: &'ll Type, align: Align, name: &str) -> &'ll Value {


I'll find a better name for it.

ZuseZ4 · 2025-06-19T00:40:23Z

compiler/rustc_codegen_llvm/src/back/lto.rs

@@ -667,6 +668,13 @@ pub(crate) fn run_pass_manager(
        write::llvm_optimize(cgcx, dcx, module, None, config, opt_level, opt_stage, stage)?;
    }

+    if cfg!(llvm_enzyme) && enable_gpu && !thin {
+        dbg!(&enable_gpu);


I'll remove the dbg statements.

ZuseZ4 · 2025-06-19T00:41:09Z

compiler/rustc_codegen_llvm/src/declare.rs

@@ -215,7 +215,9 @@ impl<'ll, 'tcx> CodegenCx<'ll, 'tcx> {

        llfn
    }
+}


Unfortunately I forgot this one in the first PR, I can make another one?

ZuseZ4 · 2025-06-19T00:41:59Z

compiler/rustc_codegen_llvm/src/llvm/ffi.rs

@@ -1004,13 +1004,22 @@ unsafe extern "C" {
        SLen: c_uint,
    ) -> MetadataKindId;

-    // Create modules.
+    // Create, print, and destroy modules.
+    pub(crate) fn LLVMPrintModuleToFile(


Can go into the follow-up PR (I guess I am not getting all unused function warnings for some reason).

ZuseZ4 · 2025-06-19T00:42:41Z

compiler/rustc_llvm/llvm-wrapper/RustWrapper.cpp

+        const llvm::Function *calledFunc = callInst->getCalledFunction();
+        if (calledFunc && calledFunc->getName() == targetName) {
+          // Found a call to the target function
+          llvm::errs() << "Found call: " << *callInst << "\n";


rustbot added A-LLVM Area: Code generation parts specific to LLVM. Both correctness bugs and optimization-related issues. T-compiler Relevant to the compiler team, which will review and decide on the PR/issue. labels Jun 5, 2025

ZuseZ4 added F-gpu_offload `#![feature(gpu_offload)]` and removed A-LLVM Area: Code generation parts specific to LLVM. Both correctness bugs and optimization-related issues. T-compiler Relevant to the compiler team, which will review and decide on the PR/issue. labels Jun 5, 2025

This comment has been minimized.

Sign in to view

rustbot added A-LLVM Area: Code generation parts specific to LLVM. Both correctness bugs and optimization-related issues. T-compiler Relevant to the compiler team, which will review and decide on the PR/issue. labels Jun 5, 2025

This comment has been minimized.

Sign in to view

rustbot added the F-autodiff `#![feature(autodiff)]` label Jun 9, 2025

This comment has been minimized.

Sign in to view

ZuseZ4 mentioned this pull request Mar 4, 2025

Tracking Issue for GPU-offload #131513

Open

5 tasks

oli-obk reviewed Jun 12, 2025

View reviewed changes

compiler/rustc_codegen_llvm/src/back/lto.rs Outdated Show resolved Hide resolved

oli-obk reviewed Jun 12, 2025

View reviewed changes

compiler/rustc_codegen_llvm/src/context.rs Outdated Show resolved Hide resolved

oli-obk reviewed Jun 12, 2025

View reviewed changes

compiler/rustc_codegen_llvm/src/context.rs Outdated Show resolved Hide resolved

ZuseZ4 mentioned this pull request Jun 15, 2025

Expose experimental LLVM features for GPU offloading rust-lang/rust-project-goals#109

Open

4 tasks

ZuseZ4 force-pushed the offload-host1 branch from c8d7349 to 1c1953d Compare June 17, 2025 04:02

rustbot added the T-bootstrap Relevant to the bootstrap subteam: Rust's build system (x.py and src/bootstrap) label Jun 17, 2025

This comment has been minimized.

Sign in to view

ZuseZ4 force-pushed the offload-host1 branch from 1c1953d to f185093 Compare June 17, 2025 21:35

This comment has been minimized.

Sign in to view

bors added the S-waiting-on-author Status: This is awaiting some action (such as code changes or more information) from the author. label Jun 18, 2025

ZuseZ4 force-pushed the offload-host1 branch 2 times, most recently from 100f9f3 to 0fb93f0 Compare June 18, 2025 21:55

This comment has been minimized.

Sign in to view

ZuseZ4 added 5 commits June 18, 2025 15:25

make more builder functions generic

3debe14

add -Zoffload=Enable flag, to enable gpu (host) code generation

22a499f

various wrappers

3dc0c3b

gpu host code generation

c1e5961

add gpu offload codegen host side test

5d3ce24

ZuseZ4 force-pushed the offload-host1 branch from 0fb93f0 to 5d3ce24 Compare June 18, 2025 23:47

ZuseZ4 marked this pull request as ready for review June 19, 2025 00:38

rustbot added the S-waiting-on-review Status: Awaiting review from the assignee but also interested parties. label Jun 19, 2025

ZuseZ4 commented Jun 19, 2025

View reviewed changes

@@ @@ -215,7 +215,9 @@ impl<'ll, 'tcx> CodegenCx<'ll, 'tcx> { @@
                       llfn
                   }
+              }

gpu offload host code generation #142097

Are you sure you want to change the base?

gpu offload host code generation #142097

Conversation

ZuseZ4 commented Jun 5, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

This comment has been minimized.

This comment has been minimized.

This comment has been minimized.

This comment has been minimized.

This comment has been minimized.

This comment has been minimized.

This comment has been minimized.

This comment has been minimized.

This comment has been minimized.

This comment has been minimized.

This comment has been minimized.

ZuseZ4 commented Jun 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ZuseZ4 commented Jun 12, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

This comment has been minimized.

This comment has been minimized.

bors commented Jun 18, 2025

Uh oh!

ZuseZ4 commented Jun 18, 2025

Uh oh!

This comment has been minimized.

rust-log-analyzer commented Jun 18, 2025

Uh oh!

rustbot commented Jun 19, 2025

Uh oh!

ZuseZ4 left a comment

Choose a reason for hiding this comment

Uh oh!

ZuseZ4 Jun 19, 2025

Choose a reason for hiding this comment

Uh oh!

ZuseZ4 Jun 19, 2025

Choose a reason for hiding this comment

Uh oh!

ZuseZ4 Jun 19, 2025

Choose a reason for hiding this comment

Uh oh!

ZuseZ4 Jun 19, 2025

Choose a reason for hiding this comment

Uh oh!

ZuseZ4 Jun 19, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

ZuseZ4 commented Jun 5, 2025 •

edited

Loading

ZuseZ4 commented Jun 10, 2025 •

edited

Loading

ZuseZ4 commented Jun 12, 2025 •

edited

Loading