mirror of
https://github.com/rust-lang/rust.git
synced 2026-04-26 13:01:27 +03:00
Flakebi
13ec3de673
Workgroup memory is a memory region that is shared between all threads in a workgroup on GPUs. Workgroup memory can be allocated statically or after compilation, when launching a gpu-kernel. The intrinsic added here returns the pointer to the memory that is allocated at launch-time. # Interface With this change, workgroup memory can be accessed in Rust by calling the new `gpu_launch_sized_workgroup_mem<T>() -> *mut T` intrinsic. It returns the pointer to workgroup memory guaranteeing that it is aligned to at least the alignment of `T`. The pointer is dereferencable for the size specified when launching the current gpu-kernel (which may be the size of `T` but can also be larger or smaller or zero). All calls to this intrinsic return a pointer to the same address. See the intrinsic documentation for more details. ## Alternative Interfaces It was also considered to expose dynamic workgroup memory as extern static variables in Rust, like they are represented in LLVM IR. However, due to the pointer not being guaranteed to be dereferencable (that depends on the allocated size at runtime), such a global must be zero-sized, which makes global variables a bad fit. # Implementation Details Workgroup memory in amdgpu and nvptx lives in address space 3. Workgroup memory from a launch is implemented by creating an external global variable in address space 3. The global is declared with size 0, as the actual size is only known at runtime. It is defined behavior in LLVM to access an external global outside the defined size. There is no similar way to get the allocated size of launch-sized workgroup memory on amdgpu an nvptx, so users have to pass this out-of-band or rely on target specific ways for now.
42 lines
2.0 KiB
Rust
42 lines
2.0 KiB
Rust
// Checks that the GPU intrinsic to get launch-sized workgroup memory works
|
|
// and correctly aligns the `external addrspace(...) global`s over multiple calls.
|
|
|
|
//@ revisions: amdgpu nvptx-pre-llvm-23 nvptx-post-llvm-23
|
|
//@ compile-flags: --crate-type=rlib -Copt-level=1
|
|
//
|
|
//@ [amdgpu] compile-flags: --target amdgcn-amd-amdhsa -Ctarget-cpu=gfx900
|
|
//@ [amdgpu] needs-llvm-components: amdgpu
|
|
|
|
//@ [nvptx-pre-llvm-23] compile-flags: --target nvptx64-nvidia-cuda
|
|
//@ [nvptx-pre-llvm-23] needs-llvm-components: nvptx
|
|
//@ [nvptx-pre-llvm-23] max-llvm-major-version: 22
|
|
//@ [nvptx-post-llvm-23] compile-flags: --target nvptx64-nvidia-cuda
|
|
//@ [nvptx-post-llvm-23] needs-llvm-components: nvptx
|
|
//@ [nvptx-post-llvm-23] min-llvm-version: 23
|
|
//@ add-minicore
|
|
#![feature(intrinsics, no_core, rustc_attrs)]
|
|
#![no_core]
|
|
|
|
extern crate minicore;
|
|
|
|
#[rustc_intrinsic]
|
|
#[rustc_nounwind]
|
|
fn gpu_launch_sized_workgroup_mem<T>() -> *mut T;
|
|
|
|
// amdgpu-DAG: @[[SMALL:[^ ]+]] = external addrspace(3) global [0 x i8], align 4
|
|
// amdgpu-DAG: @[[BIG:[^ ]+]] = external addrspace(3) global [0 x i8], align 8
|
|
// amdgpu: ret { ptr, ptr } { ptr addrspacecast (ptr addrspace(3) @[[SMALL]] to ptr), ptr addrspacecast (ptr addrspace(3) @[[BIG]] to ptr) }
|
|
|
|
// nvptx-pre-llvm-23: @[[BIG:[^ ]+]] = external addrspace(3) global [0 x i8], align 8
|
|
// nvptx-pre-llvm-23: ret { ptr, ptr } { ptr addrspacecast (ptr addrspace(3) @[[BIG]] to ptr), ptr addrspacecast (ptr addrspace(3) @[[BIG]] to ptr) }
|
|
|
|
// nvptx-post-llvm-23-DAG: @[[SMALL:[^ ]+]] = external addrspace(3) global [0 x i8], align 4
|
|
// nvptx-post-llvm-23-DAG: @[[BIG:[^ ]+]] = external addrspace(3) global [0 x i8], align 8
|
|
// nvptx-post-llvm-23: ret { ptr, ptr } { ptr addrspacecast (ptr addrspace(3) @[[SMALL]] to ptr), ptr addrspacecast (ptr addrspace(3) @[[BIG]] to ptr) }
|
|
#[unsafe(no_mangle)]
|
|
pub fn fun() -> (*mut i32, *mut f64) {
|
|
let small = gpu_launch_sized_workgroup_mem::<i32>();
|
|
let big = gpu_launch_sized_workgroup_mem::<f64>(); // Increase alignment to 8
|
|
(small, big)
|
|
}
|