Rollup merge of #158032 - Sa4dUs:offload-device-selection, r=ZuseZ4

Offload expose device selection

get amount of available devices with `omp_get_num_devices` and add arg to the intrinsic for device selection

r? @ZuseZ4

Author: jhpratt

compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs

@@ -197,6 +197,20 @@ fn generate_launcher<'ll>(cx: &CodegenCx<'ll, '_>) -> (&'ll llvm::Value, &'ll ll
     (tgt_decl, tgt_fn_ty)
 }
 
+/// Declares the `omp_get_num_devices` runtime function and returns the
+/// declaration together with its type.
+pub(crate) fn declare_omp_get_num_devices<'ll>(
+    cx: &CodegenCx<'ll, '_>,
+) -> (&'ll llvm::Value, &'ll llvm::Type) {
+    let ti32 = cx.type_i32();
+    let tgt_fn_ty = cx.type_func(&[], ti32);
+    let name = "omp_get_num_devices";
+    let tgt_decl = declare_offload_fn(&cx, name, tgt_fn_ty);
+    let nounwind = llvm::AttributeKind::NoUnwind.create_attr(cx.llcx);
+    attributes::apply_to_llfn(tgt_decl, Function, &[nounwind]);
+    (tgt_decl, tgt_fn_ty)
+}
+
 // What is our @1 here? A magic global, used in our data_{begin/update/end}_mapper:
 // @0 = private unnamed_addr constant [23 x i8] c";unknown;unknown;0;0;;\00", align 1
 // @1 = private unnamed_addr constant %struct.ident_t { i32 0, i32 2, i32 0, i32 22, ptr @0 }, align 8
@@ -591,6 +605,7 @@ pub(crate) fn gen_call_handling<'ll, 'tcx>(
     offload_globals: &OffloadGlobals<'ll>,
     offload_dims: &OffloadKernelDims<'ll>,
     dyn_cache: &'ll Value,
+    device_id: &'ll Value,
 ) {
     let cx = builder.cx;
     let OffloadKernelGlobals {
@@ -775,15 +790,8 @@ pub(crate) fn gen_call_handling<'ll, 'tcx>(
         builder.store(value.2, ptr, value.0);
     }
 
-    let args = vec![
-        s_ident_t,
-        // FIXME(offload) give users a way to select which GPU to use.
-        cx.get_const_i64(u64::MAX), // MAX == -1.
-        num_workgroups,
-        threads_per_block,
-        region_id,
-        a5,
-    ];
+    let device_id = builder.sext(device_id, cx.type_i64());
+    let args = vec![s_ident_t, device_id, num_workgroups, threads_per_block, region_id, a5];
     builder.call(tgt_target_kernel_ty, None, None, tgt_decl, &args, None, None);
     // %41 = call i32 @__tgt_target_kernel(ptr @1, i64 -1, i32 2097152, i32 256, ptr @.kernel_1.region_id, ptr %kernel_args)
 

compiler/rustc_codegen_llvm/src/intrinsic.rs

@@ -37,7 +37,7 @@ use crate::abi::FnAbiLlvmExt;
 use crate::builder::Builder;
 use crate::builder::autodiff::{adjust_activity_to_abi, generate_enzyme_call};
 use crate::builder::gpu_offload::{
-    OffloadKernelDims, gen_call_handling, gen_define_handling, register_offload,
+    self, OffloadKernelDims, declare_omp_get_num_devices, register_offload,
 };
 use crate::context::CodegenCx;
 use crate::declare::declare_raw_fn;
@@ -241,6 +241,13 @@ impl<'ll, 'tcx> IntrinsicCallBuilderMethods<'tcx> for Builder<'_, 'll, 'tcx> {
                 // offload *has* a return type, but somehow works without mentioning the place
                 return IntrinsicResult::WroteIntoPlace;
             }
+            sym::offload_get_num_devices => {
+                let (fn_decl, fn_ty) = declare_omp_get_num_devices(self.cx);
+
+                let llval = self.call(fn_ty, None, None, fn_decl, &[], None, None);
+
+                return IntrinsicResult::Operand(OperandValue::Immediate(llval));
+            },
             sym::is_val_statically_known => {
                 if let OperandValue::Immediate(imm) = args[0].val {
                     self.call_intrinsic(
@@ -1851,7 +1858,11 @@ fn codegen_offload<'ll, 'tcx>(
         OperandValue::Immediate(val) => val,
         _ => panic!("unparsable"),
     };
-    let args = get_args_from_tuple(bx, args[4], fn_target);
+    let device_id = match args[4].val {
+        OperandValue::Immediate(val) => val,
+        _ => panic!("unparsable"),
+    };
+    let args = get_args_from_tuple(bx, args[5], fn_target);
     let target_symbol = mangle_offload_export(tcx, fn_target);
 
     let sig = tcx.fn_sig(fn_target.def_id()).instantiate(tcx, fn_target.args).skip_norm_wip();
@@ -1882,8 +1893,9 @@ fn codegen_offload<'ll, 'tcx>(
         }
     };
     register_offload(cx);
-    let offload_data = gen_define_handling(&cx, &metadata, target_symbol, offload_globals);
-    gen_call_handling(
+    let offload_data =
+        gpu_offload::gen_define_handling(&cx, &metadata, target_symbol, offload_globals);
+    gpu_offload::gen_call_handling(
         bx,
         &offload_data,
         &args,
@@ -1892,6 +1904,7 @@ fn codegen_offload<'ll, 'tcx>(
         offload_globals,
         &offload_dims,
         &dyn_cache,
+        &device_id,
     );
 }
 

compiler/rustc_codegen_ssa/src/mir/intrinsic.rs

@@ -135,6 +135,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                 | sym::atomic_fence
                 | sym::atomic_singlethreadfence
                 | sym::caller_location
+                | sym::offload_get_num_devices
                 | sym::return_address => {}
                 _ => {
                     span_bug!(

compiler/rustc_hir_analysis/src/check/intrinsic.rs

@@ -168,6 +168,7 @@ fn intrinsic_operation_unsafety(tcx: TyCtxt<'_>, intrinsic_id: LocalDefId) -> hi
         | sym::needs_drop
         | sym::non_exhaustive
         | sym::offload
+        | sym::offload_get_num_devices
         | sym::offset_of
         | sym::overflow_checks
         | sym::powf16
@@ -384,10 +385,12 @@ pub(crate) fn check_intrinsic_type(
                 Ty::new_array_with_const_len(tcx, tcx.types.u32, Const::from_target_usize(tcx, 3)),
                 Ty::new_array_with_const_len(tcx, tcx.types.u32, Const::from_target_usize(tcx, 3)),
                 tcx.types.u32,
+                tcx.types.i32,
                 param(1),
             ],
             param(2),
         ),
+        sym::offload_get_num_devices => (0, 0, vec![], tcx.types.i32),
         sym::offset => (2, 0, vec![param(0), param(1)], param(0)),
         sym::arith_offset => (
             1,

compiler/rustc_span/src/symbol.rs

@@ -1471,6 +1471,7 @@ symbols! {
         of,
         off,
         offload,
+        offload_get_num_devices,
         offload_kernel,
         offset,
         offset_of,

library/core/src/intrinsics/mod.rs

@@ -3808,13 +3808,15 @@ pub const fn autodiff<F, G, T: crate::marker::Tuple, R>(f: F, df: G, args: T) ->
 /// - `f`: The kernel function to offload.
 /// - `workgroup_dim`: A 3D size specifying the number of workgroups to launch.
 /// - `thread_dim`: A 3D size specifying the number of threads per workgroup.
+/// - `dyn_cache`: The amount of dynamic shared memory to request for the kernel.
+/// - `device_id`: The device to offload to. Use `-1` to select the default device.
 /// - `args`: A tuple of arguments forwarded to `f`.
 ///
 /// Example usage (pseudocode):
 ///
 /// ```rust,ignore (pseudocode)
 /// fn kernel(x: *mut [f64; 128]) {
-///     core::intrinsics::offload(kernel_1, [256, 1, 1], [32, 1, 1], (x,))
+///     core::intrinsics::offload(kernel_1, [256, 1, 1], [32, 1, 1], 0, -1, (x,))
 /// }
 ///
 /// #[cfg(target_os = "linux")]
@@ -3838,9 +3840,20 @@ pub const fn offload<F, T: crate::marker::Tuple, R>(
     workgroup_dim: [u32; 3],
     thread_dim: [u32; 3],
     dyn_cache: u32,
+    device_id: i32,
     args: T,
 ) -> R;
 
+/// Returns the number of offload devices available on the system.
+///
+/// Use this to discover which `device_id` values are valid to pass to
+/// [`offload`]. Devices are numbered from `0` to the returned value minus one.
+///
+/// Returns `0` if no offloading devices are present.
+#[rustc_nounwind]
+#[rustc_intrinsic]
+pub const fn offload_get_num_devices() -> i32;
+
 /// Inform Miri that a given pointer definitely has a certain alignment.
 #[cfg(miri)]
 #[rustc_allow_const_fn_unstable(const_eval_select)]

library/core/src/offload.md

@@ -21,7 +21,8 @@ fn kernel(x: *mut [f64; 256]) {
 ```
 
 To launch an offloaded kernel, use the `offload!` macro. It lets you specify the kernel, the
-workgroup and thread dimensions, and the arguments to forward to the device.
+workgroup and thread dimensions, the device to offload to, and the arguments to forward to the
+device.
 
 ```rust,ignore (optional component)
 let mut x = [0.0f64; 256];

library/core/src/offload/mod.rs

@@ -19,6 +19,9 @@ pub use crate::offload;
 ///   Defaults to `[1, 1, 1]`.
 /// - `dyn_cache`: The amount of dynamic shared memory, in bytes, to allocate for the kernel.
 ///   Defaults to `0`.
+/// - `device`: The index of the device to offload to. Must be `>= 0`. If omitted, the
+///   default device is used. Use [`crate::intrinsics::offload_get_num_devices`] to discover
+///   which device ids are valid.
 ///
 /// Each argument may only be specified once.
 ///
@@ -43,61 +46,82 @@ macro_rules! offload {
             workgroup_dim = ([1, 1, 1]);
             thread_dim = ([1, 1, 1]);
             dyn_cache = (0);
+            device = NONE;
             args = NONE
         )
     };
 
-    (@munch [kernel = $val:expr $(, $rest_f:ident = $rest_v:expr)*]; kernel = NONE; workgroup_dim = $w:tt; thread_dim = $t:tt; dyn_cache = $d:tt; args = $a:tt) => {
-        $crate::offload!(@munch [$($rest_f = $rest_v),*]; kernel = (SOME $val); workgroup_dim = $w; thread_dim = $t; dyn_cache = $d; args = $a)
+    (@munch [kernel = $val:expr $(, $rest_f:ident = $rest_v:expr)*]; kernel = NONE; workgroup_dim = $w:tt; thread_dim = $t:tt; dyn_cache = $d:tt; device = $device:tt; args = $a:tt) => {
+        $crate::offload!(@munch [$($rest_f = $rest_v),*]; kernel = (SOME $val); workgroup_dim = $w; thread_dim = $t; dyn_cache = $d; device = $device; args = $a)
     };
-    (@munch [kernel = $val:expr $(, $rest_f:ident = $rest_v:expr)*]; kernel = (SOME $old:expr); workgroup_dim = $w:tt; thread_dim = $t:tt; dyn_cache = $d:tt; args = $a:tt) => {
+    (@munch [kernel = $val:expr $(, $rest_f:ident = $rest_v:expr)*]; kernel = (SOME $old:expr); workgroup_dim = $w:tt; thread_dim = $t:tt; dyn_cache = $d:tt; device = $device:tt; args = $a:tt) => {
         compile_error!("duplicate field `kernel`")
     };
-    (@munch [workgroup_dim = $val:expr $(, $rest_f:ident = $rest_v:expr)*]; kernel = $k:tt; workgroup_dim = ([1, 1, 1]); thread_dim = $t:tt; dyn_cache = $d:tt; args = $a:tt) => {
-        $crate::offload!(@munch [$($rest_f = $rest_v),*]; kernel = $k; workgroup_dim = (SOME $val); thread_dim = $t; dyn_cache = $d; args = $a)
+    (@munch [workgroup_dim = $val:expr $(, $rest_f:ident = $rest_v:expr)*]; kernel = $k:tt; workgroup_dim = ([1, 1, 1]); thread_dim = $t:tt; dyn_cache = $d:tt; device = $device:tt; args = $a:tt) => {
+        $crate::offload!(@munch [$($rest_f = $rest_v),*]; kernel = $k; workgroup_dim = (SOME $val); thread_dim = $t; dyn_cache = $d; device = $device; args = $a)
     };
-    (@munch [workgroup_dim = $val:expr $(, $rest_f:ident = $rest_v:expr)*]; kernel = $k:tt; workgroup_dim = (SOME $old:expr); thread_dim = $t:tt; dyn_cache = $d:tt; args = $a:tt) => {
+    (@munch [workgroup_dim = $val:expr $(, $rest_f:ident = $rest_v:expr)*]; kernel = $k:tt; workgroup_dim = (SOME $old:expr); thread_dim = $t:tt; dyn_cache = $d:tt; device = $device:tt; args = $a:tt) => {
         compile_error!("duplicate field `workgroup_dim`")
     };
-    (@munch [thread_dim = $val:expr $(, $rest_f:ident = $rest_v:expr)*]; kernel = $k:tt; workgroup_dim = $w:tt; thread_dim = ([1, 1, 1]); dyn_cache = $d:tt; args = $a:tt) => {
-        $crate::offload!(@munch [$($rest_f = $rest_v),*]; kernel = $k; workgroup_dim = $w; thread_dim = (SOME $val); dyn_cache = $d; args = $a)
+    (@munch [thread_dim = $val:expr $(, $rest_f:ident = $rest_v:expr)*]; kernel = $k:tt; workgroup_dim = $w:tt; thread_dim = ([1, 1, 1]); dyn_cache = $d:tt; device = $device:tt; args = $a:tt) => {
+        $crate::offload!(@munch [$($rest_f = $rest_v),*]; kernel = $k; workgroup_dim = $w; thread_dim = (SOME $val); dyn_cache = $d; device = $device; args = $a)
     };
-    (@munch [thread_dim = $val:expr $(, $rest_f:ident = $rest_v:expr)*]; kernel = $k:tt; workgroup_dim = $w:tt; thread_dim = (SOME $old:expr); dyn_cache = $d:tt; args = $a:tt) => {
+    (@munch [thread_dim = $val:expr $(, $rest_f:ident = $rest_v:expr)*]; kernel = $k:tt; workgroup_dim = $w:tt; thread_dim = (SOME $old:expr); dyn_cache = $d:tt; device = $device:tt; args = $a:tt) => {
         compile_error!("duplicate field `thread_dim`")
     };
-    (@munch [dyn_cache = $val:expr $(, $rest_f:ident = $rest_v:expr)*]; kernel = $k:tt; workgroup_dim = $w:tt; thread_dim = $t:tt; dyn_cache = (0); args = $a:tt) => {
-        $crate::offload!(@munch [$($rest_f = $rest_v),*]; kernel = $k; workgroup_dim = $w; thread_dim = $t; dyn_cache = (SOME $val); args = $a)
+    (@munch [dyn_cache = $val:expr $(, $rest_f:ident = $rest_v:expr)*]; kernel = $k:tt; workgroup_dim = $w:tt; thread_dim = $t:tt; dyn_cache = (0); device = $device:tt; args = $a:tt) => {
+        $crate::offload!(@munch [$($rest_f = $rest_v),*]; kernel = $k; workgroup_dim = $w; thread_dim = $t; dyn_cache = (SOME $val); device = $device; args = $a)
     };
-    (@munch [dyn_cache = $val:expr $(, $rest_f:ident = $rest_v:expr)*]; kernel = $k:tt; workgroup_dim = $w:tt; thread_dim = $t:tt; dyn_cache = (SOME $old:expr); args = $a:tt) => {
+    (@munch [dyn_cache = $val:expr $(, $rest_f:ident = $rest_v:expr)*]; kernel = $k:tt; workgroup_dim = $w:tt; thread_dim = $t:tt; dyn_cache = (SOME $old:expr); device = $device:tt; args = $a:tt) => {
         compile_error!("duplicate field `dyn_cache`")
     };
-    (@munch [args = $val:expr $(, $rest_f:ident = $rest_v:expr)*]; kernel = $k:tt; workgroup_dim = $w:tt; thread_dim = $t:tt; dyn_cache = $d:tt; args = NONE) => {
-        $crate::offload!(@munch [$($rest_f = $rest_v),*]; kernel = $k; workgroup_dim = $w; thread_dim = $t; dyn_cache = $d; args = (SOME $val))
+    (@munch [device = $val:expr $(, $rest_f:ident = $rest_v:expr)*]; kernel = $k:tt; workgroup_dim = $w:tt; thread_dim = $t:tt; dyn_cache = $d:tt; device = NONE; args = $a:tt) => {
+        $crate::offload!(@munch [$($rest_f = $rest_v),*]; kernel = $k; workgroup_dim = $w; thread_dim = $t; dyn_cache = $d; device = (SOME $val); args = $a)
     };
-    (@munch [args = $val:expr $(, $rest_f:ident = $rest_v:expr)*]; kernel = $k:tt; workgroup_dim = $w:tt; thread_dim = $t:tt; dyn_cache = $d:tt; args = (SOME $old:expr)) => {
+    (@munch [device = $val:expr $(, $rest_f:ident = $rest_v:expr)*]; kernel = $k:tt; workgroup_dim = $w:tt; thread_dim = $t:tt; dyn_cache = $d:tt; device = (SOME $old:expr); args = $a:tt) => {
+        compile_error!("duplicate field `device`")
+    };
+    (@munch [args = $val:expr $(, $rest_f:ident = $rest_v:expr)*]; kernel = $k:tt; workgroup_dim = $w:tt; thread_dim = $t:tt; dyn_cache = $d:tt; device = $device:tt; args = NONE) => {
+        $crate::offload!(@munch [$($rest_f = $rest_v),*]; kernel = $k; workgroup_dim = $w; thread_dim = $t; dyn_cache = $d; device = $device; args = (SOME $val))
+    };
+    (@munch [args = $val:expr $(, $rest_f:ident = $rest_v:expr)*]; kernel = $k:tt; workgroup_dim = $w:tt; thread_dim = $t:tt; dyn_cache = $d:tt; device = $device:tt; args = (SOME $old:expr)) => {
         compile_error!("duplicate field `args`")
     };
 
-    (@munch [$invalid:ident = $val:expr $(, $rest_f:ident = $rest_v:expr)*]; kernel = $k:tt; workgroup_dim = $w:tt; thread_dim = $t:tt; dyn_cache = $d:tt; args = $a:tt) => {
+    (@munch [$invalid:ident = $val:expr $(, $rest_f:ident = $rest_v:expr)*]; kernel = $k:tt; workgroup_dim = $w:tt; thread_dim = $t:tt; dyn_cache = $d:tt; device = $device:tt; args = $a:tt) => {
         compile_error!(concat!("unknown field `", stringify!($invalid), "`"))
     };
 
-    (@munch []; kernel = NONE; workgroup_dim = $w:tt; thread_dim = $t:tt; dyn_cache = $d:tt; args = $a:tt) => {
+    (@munch []; kernel = NONE; workgroup_dim = $w:tt; thread_dim = $t:tt; dyn_cache = $d:tt; device = $device:tt; args = $a:tt) => {
         compile_error!("missing `kernel`")
     };
-    (@munch []; kernel = $k:tt; workgroup_dim = $w:tt; thread_dim = $t:tt; dyn_cache = $d:tt; args = NONE) => {
+    (@munch []; kernel = $k:tt; workgroup_dim = $w:tt; thread_dim = $t:tt; dyn_cache = $d:tt; device = $device:tt; args = NONE) => {
         compile_error!("missing `args`")
     };
-    (@munch []; kernel = (SOME $kernel:expr); workgroup_dim = $w:tt; thread_dim = $t:tt; dyn_cache = $d:tt; args = (SOME $args:expr)) => {
+    (@munch []; kernel = (SOME $kernel:expr); workgroup_dim = $w:tt; thread_dim = $t:tt; dyn_cache = $d:tt; device = $device:tt; args = (SOME $args:expr)) => {
         $crate::intrinsics::offload::<_, _, ()>(
             $kernel,
             $crate::offload!(@value $w),
             $crate::offload!(@value $t),
             $crate::offload!(@value $d),
+            $crate::offload!(@device $device),
             $args,
         )
     };
 
     (@value (SOME $val:expr)) => { $val };
     (@value ($val:expr)) => { $val };
+
+    // if `device` is omitted (`NONE), we use the OpenMP default device (`-1`)
+    (@device NONE) => { -1 };
+    (@device (SOME $val:expr)) => { {
+        const { $crate::assert!($val >= 0, "offload device must be non-negative; omit `device` to use the default device") };
+        let device: i32 = $val;
+        $crate::assert!(
+            device < $crate::intrinsics::offload_get_num_devices(),
+            "offload device {} is not available",
+            device,
+        );
+        device
+    } };
 }

tests/codegen-llvm/gpu_offload/control_flow.rs

@@ -6,8 +6,8 @@
 // contains control flow.
 
 #![feature(abi_gpu_kernel)]
+#![feature(gpu_offload)]
 #![feature(rustc_attrs)]
-#![feature(core_intrinsics)]
 #![no_main]
 
 // CHECK: @.offload_sizes.[[K:[^ ]*foo]] = private unnamed_addr constant
@@ -28,13 +28,12 @@ unsafe fn main() {
     let A = [1.0, 1.0, 1.0, 1.0, 1.0, 1.0];
 
     for i in 0..100 {
-        core::intrinsics::offload::<_, _, ()>(
-            foo,
-            [256, 1, 1],
-            [32, 1, 1],
-            0,
-            (A.as_ptr() as *const [f32; 6],),
-        );
+        core::offload::offload! {
+            kernel = foo,
+            workgroup_dim = [256, 1, 1],
+            thread_dim = [32, 1, 1],
+            args = (A.as_ptr() as *const [f32; 6],),
+        }
     }
 }
 

tests/codegen-llvm/gpu_offload/device_check.rs

@@ -0,0 +1,30 @@
+//@ compile-flags: -Zoffload=Test -Zunstable-options -C opt-level=0 -Clto=fat
+//@ no-prefer-dynamic
+//@ needs-offload
+
+// This test verifies that selecting an unavailable `device` in the `offload` macro panics.
+
+#![feature(gpu_offload)]
+#![no_main]
+
+#[unsafe(no_mangle)]
+fn main() {
+    core::offload::offload! {
+        kernel = kernel,
+        device = 99,
+        args = (),
+    }
+}
+
+#[unsafe(no_mangle)]
+fn kernel() {}
+
+// CHECK-LABEL: define{{( dso_local)?}} void @main()
+// CHECK: store i32 99, ptr %device, align 4
+// CHECK-NEXT: %{{[0-9_]+}} = call i32 @omp_get_num_devices()
+// CHECK-NEXT: %{{[0-9_]+}} = load i32, ptr %device, align 4
+// CHECK-NEXT: %{{[0-9_]+}} = icmp slt i32 %{{[0-9_]+}}, %{{[0-9_]+}}
+// CHECK-NEXT: br i1 %{{[0-9_]+}}, label %bb{{[0-9]+}}, label %bb{{[0-9]+}}
+// CHECK: call void @{{.*}}panic_fmt
+// CHECK: unreachable
+// CHECK: call i32 @__tgt_target_kernel