Offload intrinsic

[Sa4dUs]

This PR implements the minimal mechanisms required to run a small subset of arbitrary offload kernels without relying on hardcoded names or metadata.

• offload(kernel, (..args)): an intrinsic that generates the necessary host-side LLVM-IR code.
• rustc_offload_kernel: a builtin attribute that marks device kernels to be handled appropriately.

Example usage (pseudocode):

fn kernel(x: *mut [f64; 128]) {
    core::intrinsics::offload(kernel_1, (x,))
}

#[cfg(target_os = "linux")]
extern "C" {
    pub fn kernel_1(array_b: *mut [f64; 128]);
}

#[cfg(not(target_os = "linux"))]
#[rustc_offload_kernel]
extern "gpu-kernel" fn kernel_1(x: *mut [f64; 128]) {
    unsafe { (*x)[0] = 21.0 };
}

[ZuseZ4]

If you already have the code here, I would add a small check for & or byVal (implies Mode ToGPU), vs &mut (implies Both).

In the future we would hope to analyze the & or byval case more, if we never read from it (before writing) then we could use a new mode 4, which allocates directly on the gpu.

[bors]

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

[bors]

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

[ZuseZ4]

can you iterate directly over args now?

[ZuseZ4]

I'd rather have that as a TODO, thus fixed before we land it. You can just copy the check I had above (and potentially error if you find an intrinsic, but not the flag set). The current main already requires it, and we shouldn't change nightly behaviour because of an experimental feature without a user setting a feature flag.

[ZuseZ4]

I'd expect all of these to not be handled correctly by offload without further work, so I'd just error.

The same goes for

            | ty::Bound(_, _)
            | ty::Placeholder(_)
            | ty::Infer(_)
            | ty::Error(_)

[rustbot]

Some changes occurred to the intrinsics. Make sure the CTFE / Miri interpreter
gets adapted for the changes, if necessary.

cc @rust-lang/miri, @RalfJung, @oli-obk, @lcnr

[RalfJung]

We have other backends besides LLVM, so intrinsics typically should be described in terms of what they do, not implementation details. Is that possible here?

[Sa4dUs]

this intrinsic only makes sense in LLVM right now because it relies directly on LLVM's offload feature. that's why i wanted to specify the backend

[Sa4dUs]

if there's a better way to proceed, please let me know

[RalfJung]

Hm, seems like we did something similar for the autodiff intrinsic. I would assume that the concept of offloading is independent of LLVM, but maybe we don't have to figure out that full story at this point.

Are there docs for the LLVM offload feature you could link to?

[Sa4dUs]

i'd say https://clang.llvm.org/docs/OffloadingDesign.html contains all the relevant details

ping @ZuseZ4 in case he has something better

[ZuseZ4]

Yes, your link is probably the best overview. Offload grew out of OpenMP, which is also supported by other compilers like GCC. LLVM just put in some effort to split Offloading and OpenMP, so that the former is easier to use independently. https://gcc.gnu.org/projects/gomp/

ping @antoyo just for awareness.

With respect to a high-level explanation of this intrinsic:
We use a single-source, two-pass compilation approach. We compile all functions that should be offloaded for the device (e.g nvptx64, amdgcn-amd-amdhsa, intel in the future) and which are marked by our intrinsic. We then compile the code for the host (e.g. x86-64), where most of the offloading logic happens. On the host side, we generate calls to the openmp offload runtime, to inform it about the layout of the types (a simplified version of the autodiff TypeTrees). We also use the type system to figure out whether kernel arguments have to be moved only to the device (e.g. &[f32;1024]), from the device, or both (e.g. &mut [f64]). We then launched the kernel, after which we inform the runtime to end this environment and move data back (as far as needed).

There are obviously a lot of features and optimizations which we want to add in the future. The Rust frontend currently also mostly uses the OpenMP API, since it was more stable back when I started working on it. We intend to move over to the newer offload API, which is slightly lower level.