From c266bb2f5307864efd38eea45343e0c650bc4470 Mon Sep 17 00:00:00 2001 From: Sa4dUs Date: Wed, 17 Jun 2026 16:19:15 +0200 Subject: [PATCH 1/3] get number of devices intrinisc --- compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs | 10 ++++++++++ compiler/rustc_codegen_llvm/src/intrinsic.rs | 9 ++++++++- compiler/rustc_codegen_ssa/src/mir/intrinsic.rs | 1 + compiler/rustc_hir_analysis/src/check/intrinsic.rs | 7 +++++++ compiler/rustc_span/src/symbol.rs | 1 + library/core/src/intrinsics/mod.rs | 5 +++++ 6 files changed, 32 insertions(+), 1 deletion(-) diff --git a/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs b/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs index 3d0bb6fcc48fd..5fc4f71f951ad 100644 --- a/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs +++ b/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs @@ -197,6 +197,16 @@ fn generate_launcher<'ll>(cx: &CodegenCx<'ll, '_>) -> (&'ll llvm::Value, &'ll ll (tgt_decl, tgt_fn_ty) } +pub(crate) fn generate_decl<'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 diff --git a/compiler/rustc_codegen_llvm/src/intrinsic.rs b/compiler/rustc_codegen_llvm/src/intrinsic.rs index ba11ef29fb536..da230a6da9ac2 100644 --- a/compiler/rustc_codegen_llvm/src/intrinsic.rs +++ b/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, + OffloadKernelDims, gen_call_handling, gen_define_handling, register_offload, generate_decl, }; 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) = generate_decl(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( diff --git a/compiler/rustc_codegen_ssa/src/mir/intrinsic.rs b/compiler/rustc_codegen_ssa/src/mir/intrinsic.rs index 77ab3bdce4689..23285a043af70 100644 --- a/compiler/rustc_codegen_ssa/src/mir/intrinsic.rs +++ b/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!( diff --git a/compiler/rustc_hir_analysis/src/check/intrinsic.rs b/compiler/rustc_hir_analysis/src/check/intrinsic.rs index 11d21744bd7ae..c96cd0a8c8a5f 100644 --- a/compiler/rustc_hir_analysis/src/check/intrinsic.rs +++ b/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,6 +385,12 @@ pub(crate) fn check_intrinsic_type( ], 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, diff --git a/compiler/rustc_span/src/symbol.rs b/compiler/rustc_span/src/symbol.rs index 45df107bf7469..b678f63495b82 100644 --- a/compiler/rustc_span/src/symbol.rs +++ b/compiler/rustc_span/src/symbol.rs @@ -1468,6 +1468,7 @@ symbols! { of, off, offload, + offload_get_num_devices, offload_kernel, offset, offset_of, diff --git a/library/core/src/intrinsics/mod.rs b/library/core/src/intrinsics/mod.rs index 4e67199bba8c9..eefb1b9ed7bd8 100644 --- a/library/core/src/intrinsics/mod.rs +++ b/library/core/src/intrinsics/mod.rs @@ -3814,6 +3814,11 @@ pub const fn offload( args: T, ) -> R; +// TODO(Sa4dUs): add docs +#[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)] From efbb356b50c0d7e8c743ad1fbb9f4227b050b6fb Mon Sep 17 00:00:00 2001 From: Marcelo Dominguez Date: Thu, 13 Aug 2026 12:35:46 +0300 Subject: [PATCH 2/3] Add device arg to offload intrinisc --- .../src/builder/gpu_offload.rs | 12 +++-------- compiler/rustc_codegen_llvm/src/intrinsic.rs | 9 ++++++-- .../rustc_hir_analysis/src/check/intrinsic.rs | 8 ++----- library/core/src/intrinsics/mod.rs | 12 +++++++++-- .../codegen-llvm/gpu_offload/control_flow.rs | 1 + tests/codegen-llvm/gpu_offload/gpu_host.rs | 16 +++++++------- tests/codegen-llvm/gpu_offload/scalar_host.rs | 9 +++++++- .../codegen-llvm/gpu_offload/slice_device.rs | 2 +- tests/codegen-llvm/gpu_offload/slice_host.rs | 9 +++++++- tests/ui/offload/check_config.rs | 2 +- tests/ui/offload/duplicate_kernel.rs | 2 +- tests/ui/offload/non_tuple_args.rs | 2 +- tests/ui/offload/non_tuple_args.stderr | 2 +- tests/ui/offload/type_mismatch.rs | 21 ++++++++++++------- tests/ui/offload/type_mismatch.stderr | 12 +++++------ 15 files changed, 73 insertions(+), 46 deletions(-) diff --git a/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs b/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs index 5fc4f71f951ad..4f8d918c72b51 100644 --- a/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs +++ b/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs @@ -601,6 +601,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 { @@ -785,15 +786,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) diff --git a/compiler/rustc_codegen_llvm/src/intrinsic.rs b/compiler/rustc_codegen_llvm/src/intrinsic.rs index da230a6da9ac2..fbd9a04ab2687 100644 --- a/compiler/rustc_codegen_llvm/src/intrinsic.rs +++ b/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, generate_decl, + OffloadKernelDims, gen_call_handling, gen_define_handling, generate_decl, register_offload, }; use crate::context::CodegenCx; use crate::declare::declare_raw_fn; @@ -1858,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(); @@ -1899,6 +1903,7 @@ fn codegen_offload<'ll, 'tcx>( offload_globals, &offload_dims, &dyn_cache, + &device_id, ); } diff --git a/compiler/rustc_hir_analysis/src/check/intrinsic.rs b/compiler/rustc_hir_analysis/src/check/intrinsic.rs index c96cd0a8c8a5f..39e68b7541d08 100644 --- a/compiler/rustc_hir_analysis/src/check/intrinsic.rs +++ b/compiler/rustc_hir_analysis/src/check/intrinsic.rs @@ -381,16 +381,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::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, diff --git a/library/core/src/intrinsics/mod.rs b/library/core/src/intrinsics/mod.rs index eefb1b9ed7bd8..c26c007a254e5 100644 --- a/library/core/src/intrinsics/mod.rs +++ b/library/core/src/intrinsics/mod.rs @@ -3781,13 +3781,15 @@ pub const fn autodiff(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")] @@ -3811,10 +3813,16 @@ pub const fn offload( workgroup_dim: [u32; 3], thread_dim: [u32; 3], dyn_cache: u32, + device_id: i32, args: T, ) -> R; -// TODO(Sa4dUs): add docs +/// 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; diff --git a/tests/codegen-llvm/gpu_offload/control_flow.rs b/tests/codegen-llvm/gpu_offload/control_flow.rs index da997de53a428..ef4881147e63b 100644 --- a/tests/codegen-llvm/gpu_offload/control_flow.rs +++ b/tests/codegen-llvm/gpu_offload/control_flow.rs @@ -33,6 +33,7 @@ unsafe fn main() { [256, 1, 1], [32, 1, 1], 0, + -1, (A.as_ptr() as *const [f32; 6],), ); } diff --git a/tests/codegen-llvm/gpu_offload/gpu_host.rs b/tests/codegen-llvm/gpu_offload/gpu_host.rs index 2bfaf89b45590..208417b98bcae 100644 --- a/tests/codegen-llvm/gpu_offload/gpu_host.rs +++ b/tests/codegen-llvm/gpu_offload/gpu_host.rs @@ -21,7 +21,7 @@ fn main() { } pub fn kernel_1(x: &mut [f32; 256], y: &[f32; 256]) { - core::intrinsics::offload(_kernel_1, [256, 1, 1], [32, 1, 1], 0, (x, y)) + core::intrinsics::offload(_kernel_1, [256, 1, 1], [32, 1, 1], 0, -1, (x, y)) } #[inline(never)] @@ -78,8 +78,10 @@ pub fn _kernel_1(x: &mut [f32; 256], y: &[f32; 256]) { // CHECK-NEXT: [[P32:%[^ ]+]] = getelementptr inbounds nuw i8, ptr %kernel_args, i64 32 // CHECK-NEXT: store ptr @.offload_maptypes.[[K]].kernel, ptr [[P32]], align 8 // CHECK-NEXT: [[P40:%[^ ]+]] = getelementptr inbounds nuw i8, ptr %kernel_args, i64 40 +// CHECK-NEXT: [[P64:%[^ ]+]] = getelementptr inbounds nuw i8, ptr %kernel_args, i64 64 +// CHECK-NEXT: call void @llvm.memset.p0.i64(ptr noundef nonnull align 8 dereferenceable(24) [[P40]], i8 0, i64 24, i1 false) +// CHECK-NEXT: store i64 64, ptr [[P64]], align 8 // CHECK-NEXT: [[P72:%[^ ]+]] = getelementptr inbounds nuw i8, ptr %kernel_args, i64 72 -// CHECK-NEXT: call void @llvm.memset.p0.i64(ptr noundef nonnull align 8 dereferenceable(32) [[P40]], i8 0, i64 32, i1 false) // CHECK-NEXT: store <4 x i32> , ptr [[P72]], align 8 // CHECK-NEXT: [[P88:%[^ ]+]] = getelementptr inbounds nuw i8, ptr %kernel_args, i64 88 // CHECK-NEXT: store i32 1, ptr [[P88]], align 8 @@ -95,17 +97,17 @@ pub fn _kernel_1(x: &mut [f32; 256], y: &[f32; 256]) { // CHECK: declare void @__tgt_register_lib(ptr) local_unnamed_addr // CHECK: declare void @__tgt_unregister_lib(ptr) local_unnamed_addr -// CHECK-LABEL: define internal void @.omp_offloading.descriptor_reg() section ".text.startup" { +// CHECK-LABEL: define internal void @.omp_offloading.descriptor_reg() section ".text.startup" !guid !{{[0-9]+}} { // CHECK-NEXT: entry: -// CHECK-NEXT: call void @__tgt_register_lib(ptr nonnull @.omp_offloading.descriptor) -// CHECK-NEXT: call void @__tgt_init_all_rtls() +// CHECK-NEXT: {{tail }}call void @__tgt_register_lib(ptr nonnull @.omp_offloading.descriptor) +// CHECK-NEXT: {{tail }}call void @__tgt_init_all_rtls() // CHECK-NEXT: %0 = {{tail }}call i32 @atexit(ptr nonnull @.omp_offloading.descriptor_unreg) // CHECK-NEXT: ret void // CHECK-NEXT: } -// CHECK-LABEL: define internal void @.omp_offloading.descriptor_unreg() section ".text.startup" { +// CHECK-LABEL: define internal void @.omp_offloading.descriptor_unreg() section ".text.startup" !guid !{{[0-9]+}} { // CHECK-NEXT: entry: -// CHECK-NEXT: call void @__tgt_unregister_lib(ptr nonnull @.omp_offloading.descriptor) +// CHECK-NEXT: {{tail }}call void @__tgt_unregister_lib(ptr nonnull @.omp_offloading.descriptor) // CHECK-NEXT: ret void // CHECK-NEXT: } diff --git a/tests/codegen-llvm/gpu_offload/scalar_host.rs b/tests/codegen-llvm/gpu_offload/scalar_host.rs index 66c910c439e46..950306c83a43f 100644 --- a/tests/codegen-llvm/gpu_offload/scalar_host.rs +++ b/tests/codegen-llvm/gpu_offload/scalar_host.rs @@ -28,7 +28,14 @@ fn main() { let mut x = 0.0f32; let k = core::hint::black_box(42.0f32); - core::intrinsics::offload::<_, _, ()>(foo, [1, 1, 1], [1, 1, 1], 0, (&mut x as *mut f32, k)); + core::intrinsics::offload::<_, _, ()>( + foo, + [1, 1, 1], + [1, 1, 1], + 0, + -1, + (&mut x as *mut f32, k), + ); } unsafe extern "C" { diff --git a/tests/codegen-llvm/gpu_offload/slice_device.rs b/tests/codegen-llvm/gpu_offload/slice_device.rs index 1abe04f8cc429..6e900c21ca7cb 100644 --- a/tests/codegen-llvm/gpu_offload/slice_device.rs +++ b/tests/codegen-llvm/gpu_offload/slice_device.rs @@ -15,7 +15,7 @@ extern crate minicore; // CHECK: ; Function Attrs // nvptx-NEXT: define ptx_kernel void @foo // amdgpu-NEXT: define amdgpu_kernel void @foo -// CHECK-SAME: ptr readnone captures(none) %dyn_ptr +// CHECK-SAME: ptr nofree readnone captures(none) %dyn_ptr // nvptx-SAME: [2 x i64] %0 // amdgpu-SAME: ptr noalias {{.*}} %0, i64 {{.*}} %1 // CHECK-NEXT: entry: diff --git a/tests/codegen-llvm/gpu_offload/slice_host.rs b/tests/codegen-llvm/gpu_offload/slice_host.rs index dfc7ec545630c..1b314d489612c 100644 --- a/tests/codegen-llvm/gpu_offload/slice_host.rs +++ b/tests/codegen-llvm/gpu_offload/slice_host.rs @@ -27,7 +27,14 @@ #[unsafe(no_mangle)] fn main() { let mut x = [0.0f32, 0.0, 0.0, 0.0]; - core::intrinsics::offload::<_, _, ()>(foo, [1, 1, 1], [1, 1, 1], 0, ((&mut x) as &mut [f32],)); + core::intrinsics::offload::<_, _, ()>( + foo, + [1, 1, 1], + [1, 1, 1], + 0, + -1, + ((&mut x) as &mut [f32],), + ); } unsafe extern "C" { diff --git a/tests/ui/offload/check_config.rs b/tests/ui/offload/check_config.rs index ff145f420e482..833ba39dffeca 100644 --- a/tests/ui/offload/check_config.rs +++ b/tests/ui/offload/check_config.rs @@ -17,7 +17,7 @@ fn main() { } fn kernel_1(x: &mut [f32; 256]) { - core::intrinsics::offload(_kernel_1, [1, 1, 1], [1, 1, 1], 0, (x,)) + core::intrinsics::offload(_kernel_1, [1, 1, 1], [1, 1, 1], 0, -1, (x,)) } fn _kernel_1(x: &mut [f32; 256]) {} diff --git a/tests/ui/offload/duplicate_kernel.rs b/tests/ui/offload/duplicate_kernel.rs index abde76137a37c..da667a0c0666a 100644 --- a/tests/ui/offload/duplicate_kernel.rs +++ b/tests/ui/offload/duplicate_kernel.rs @@ -18,5 +18,5 @@ fn kernel(_x: f32) {} fn main() { _RNvC19collision_kernels_a6kernel(0.0); - core::intrinsics::offload::<_, _, ()>(kernel, [1, 1, 1], [1, 1, 1], 0, (0.0f32,)); + core::intrinsics::offload::<_, _, ()>(kernel, [1, 1, 1], [1, 1, 1], 0, -1, (0.0f32,)); } diff --git a/tests/ui/offload/non_tuple_args.rs b/tests/ui/offload/non_tuple_args.rs index 0a07c99a26d34..14de21b2374a2 100644 --- a/tests/ui/offload/non_tuple_args.rs +++ b/tests/ui/offload/non_tuple_args.rs @@ -4,7 +4,7 @@ fn main() { // args_ty is not a tuple - core::intrinsics::offload::<_, _, ()>(kernel_0, [1, 1, 1], [1, 1, 1], 0, 42); + core::intrinsics::offload::<_, _, ()>(kernel_0, [1, 1, 1], [1, 1, 1], 0, -1, 42); //~^ ERROR `{integer}` is not a tuple } diff --git a/tests/ui/offload/non_tuple_args.stderr b/tests/ui/offload/non_tuple_args.stderr index 8b59d6828c6f2..90b0f16bec53e 100644 --- a/tests/ui/offload/non_tuple_args.stderr +++ b/tests/ui/offload/non_tuple_args.stderr @@ -1,7 +1,7 @@ error[E0277]: `{integer}` is not a tuple --> $DIR/non_tuple_args.rs:7:36 | -LL | core::intrinsics::offload::<_, _, ()>(kernel_0, [1, 1, 1], [1, 1, 1], 0, 42); +LL | core::intrinsics::offload::<_, _, ()>(kernel_0, [1, 1, 1], [1, 1, 1], 0, -1, 42); | ^ the nightly-only, unstable trait `std::marker::Tuple` is not implemented for `{integer}` | note: required by a bound in `offload` diff --git a/tests/ui/offload/type_mismatch.rs b/tests/ui/offload/type_mismatch.rs index 4079444a0aff1..a75f8358b7359 100644 --- a/tests/ui/offload/type_mismatch.rs +++ b/tests/ui/offload/type_mismatch.rs @@ -5,25 +5,32 @@ fn main() { // kernel_ty is not a function item let not_fn = 42; - core::intrinsics::offload::<_, _, ()>(not_fn, [1, 1, 1], [1, 1, 1], 0, ()); + core::intrinsics::offload::<_, _, ()>(not_fn, [1, 1, 1], [1, 1, 1], 0, -1, ()); //~^ ERROR expected a function item for the offload kernel, found `i32` // argument count mismatch - core::intrinsics::offload::<_, _, ()>(kernel_1, [1, 1, 1], [1, 1, 1], 0, ()); + core::intrinsics::offload::<_, _, ()>(kernel_1, [1, 1, 1], [1, 1, 1], 0, -1, ()); //~^ ERROR offload kernel expects 1 arguments, but 0 arguments were provided // argument type mismatch - core::intrinsics::offload::<_, _, ()>(kernel_1, [1, 1, 1], [1, 1, 1], 0, (42.0f64,)); + core::intrinsics::offload::<_, _, ()>(kernel_1, [1, 1, 1], [1, 1, 1], 0, -1, (42.0f64,)); //~^ ERROR type mismatch in offload kernel argument 0: expected `f32`, found `f64` // return type mismatch - let _: f64 = core::intrinsics::offload::<_, _, f64>(kernel_0, [1, 1, 1], [1, 1, 1], 0, ()); + let _: f64 = core::intrinsics::offload::<_, _, f64>(kernel_0, [1, 1, 1], [1, 1, 1], 0, -1, ()); //~^ ERROR offload kernel return type mismatch: kernel returns `()`, but offload call expects `f64` // multiple argument type mismatch - core::intrinsics::offload::<_, _, ()>(kernel_2, [1, 1, 1], [1, 1, 1], 0, (42.0f64, 42.0f64)); - //~^ ERROR type mismatch in offload kernel argument 0: expected `f32`, found `f64` - //~| ERROR type mismatch in offload kernel argument 1: expected `f32`, found `f64` + core::intrinsics::offload::<_, _, ()>( + //~^ ERROR type mismatch in offload kernel argument 0: expected `f32`, found `f64` + //~| ERROR type mismatch in offload kernel argument 1: expected `f32`, found `f64` + kernel_2, + [1, 1, 1], + [1, 1, 1], + 0, + -1, + (42.0f64, 42.0f64), + ); } fn kernel_0() {} diff --git a/tests/ui/offload/type_mismatch.stderr b/tests/ui/offload/type_mismatch.stderr index 8cf160ca09486..808768e7cd4f3 100644 --- a/tests/ui/offload/type_mismatch.stderr +++ b/tests/ui/offload/type_mismatch.stderr @@ -1,37 +1,37 @@ error: expected a function item for the offload kernel, found `i32` --> $DIR/type_mismatch.rs:8:5 | -LL | core::intrinsics::offload::<_, _, ()>(not_fn, [1, 1, 1], [1, 1, 1], 0, ()); +LL | core::intrinsics::offload::<_, _, ()>(not_fn, [1, 1, 1], [1, 1, 1], 0, -1, ()); | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ error: offload kernel expects 1 arguments, but 0 arguments were provided --> $DIR/type_mismatch.rs:12:5 | -LL | core::intrinsics::offload::<_, _, ()>(kernel_1, [1, 1, 1], [1, 1, 1], 0, ()); +LL | core::intrinsics::offload::<_, _, ()>(kernel_1, [1, 1, 1], [1, 1, 1], 0, -1, ()); | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ error: type mismatch in offload kernel argument 0: expected `f32`, found `f64` --> $DIR/type_mismatch.rs:16:5 | -LL | core::intrinsics::offload::<_, _, ()>(kernel_1, [1, 1, 1], [1, 1, 1], 0, (42.0f64,)); +LL | core::intrinsics::offload::<_, _, ()>(kernel_1, [1, 1, 1], [1, 1, 1], 0, -1, (42.0f64,)); | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ error: offload kernel return type mismatch: kernel returns `()`, but offload call expects `f64` --> $DIR/type_mismatch.rs:20:18 | -LL | let _: f64 = core::intrinsics::offload::<_, _, f64>(kernel_0, [1, 1, 1], [1, 1, 1], 0, ()); +LL | let _: f64 = core::intrinsics::offload::<_, _, f64>(kernel_0, [1, 1, 1], [1, 1, 1], 0, -1, ()); | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ error: type mismatch in offload kernel argument 0: expected `f32`, found `f64` --> $DIR/type_mismatch.rs:24:5 | -LL | core::intrinsics::offload::<_, _, ()>(kernel_2, [1, 1, 1], [1, 1, 1], 0, (42.0f64, 42.0f64)); +LL | core::intrinsics::offload::<_, _, ()>( | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ error: type mismatch in offload kernel argument 1: expected `f32`, found `f64` --> $DIR/type_mismatch.rs:24:5 | -LL | core::intrinsics::offload::<_, _, ()>(kernel_2, [1, 1, 1], [1, 1, 1], 0, (42.0f64, 42.0f64)); +LL | core::intrinsics::offload::<_, _, ()>( | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ error: aborting due to 6 previous errors From 3777642175b402d20d3957ba98ce49b31ed424fb Mon Sep 17 00:00:00 2001 From: Marcelo Dominguez Date: Fri, 14 Aug 2026 18:04:13 +0300 Subject: [PATCH 3/3] Update with macro and some tweaks --- .../src/builder/gpu_offload.rs | 6 +- compiler/rustc_codegen_llvm/src/intrinsic.rs | 9 +-- library/core/src/offload.md | 3 +- library/core/src/offload/mod.rs | 62 +++++++++++++------ .../codegen-llvm/gpu_offload/control_flow.rs | 16 +++-- .../codegen-llvm/gpu_offload/device_check.rs | 30 +++++++++ tests/codegen-llvm/gpu_offload/gpu_host.rs | 13 ++-- tests/codegen-llvm/gpu_offload/scalar_host.rs | 14 ++--- tests/codegen-llvm/gpu_offload/slice_host.rs | 14 ++--- .../offload-generic-manifest/generic.rs | 12 +++- tests/ui/offload/check_config.rs | 7 ++- tests/ui/offload/offload_macro.rs | 3 + tests/ui/offload/offload_macro.stderr | 10 ++- tests/ui/offload/offload_negative_device.rs | 12 ++++ .../ui/offload/offload_negative_device.stderr | 19 ++++++ 15 files changed, 168 insertions(+), 62 deletions(-) create mode 100644 tests/codegen-llvm/gpu_offload/device_check.rs create mode 100644 tests/ui/offload/offload_negative_device.rs create mode 100644 tests/ui/offload/offload_negative_device.stderr diff --git a/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs b/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs index 4f8d918c72b51..d20a73e8e6825 100644 --- a/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs +++ b/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs @@ -197,7 +197,11 @@ fn generate_launcher<'ll>(cx: &CodegenCx<'ll, '_>) -> (&'ll llvm::Value, &'ll ll (tgt_decl, tgt_fn_ty) } -pub(crate) fn generate_decl<'ll>(cx: &CodegenCx<'ll, '_>) -> (&'ll llvm::Value, &'ll llvm::Type) { +/// 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"; diff --git a/compiler/rustc_codegen_llvm/src/intrinsic.rs b/compiler/rustc_codegen_llvm/src/intrinsic.rs index fbd9a04ab2687..a5bb595d9b2c6 100644 --- a/compiler/rustc_codegen_llvm/src/intrinsic.rs +++ b/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, generate_decl, register_offload, + self, OffloadKernelDims, declare_omp_get_num_devices, register_offload, }; use crate::context::CodegenCx; use crate::declare::declare_raw_fn; @@ -242,7 +242,7 @@ impl<'ll, 'tcx> IntrinsicCallBuilderMethods<'tcx> for Builder<'_, 'll, 'tcx> { return IntrinsicResult::WroteIntoPlace; } sym::offload_get_num_devices => { - let (fn_decl, fn_ty) = generate_decl(self.cx); + let (fn_decl, fn_ty) = declare_omp_get_num_devices(self.cx); let llval = self.call(fn_ty, None, None, fn_decl, &[], None, None); @@ -1893,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, diff --git a/library/core/src/offload.md b/library/core/src/offload.md index 985a93a4294fa..726d0c7af1928 100644 --- a/library/core/src/offload.md +++ b/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]; diff --git a/library/core/src/offload/mod.rs b/library/core/src/offload/mod.rs index 17ff74f0bbfbb..3d85621361209 100644 --- a/library/core/src/offload/mod.rs +++ b/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 + } }; } diff --git a/tests/codegen-llvm/gpu_offload/control_flow.rs b/tests/codegen-llvm/gpu_offload/control_flow.rs index ef4881147e63b..8cafeda3395ce 100644 --- a/tests/codegen-llvm/gpu_offload/control_flow.rs +++ b/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,14 +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, - -1, - (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],), + } } } diff --git a/tests/codegen-llvm/gpu_offload/device_check.rs b/tests/codegen-llvm/gpu_offload/device_check.rs new file mode 100644 index 0000000000000..4eaa324a662a4 --- /dev/null +++ b/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 diff --git a/tests/codegen-llvm/gpu_offload/gpu_host.rs b/tests/codegen-llvm/gpu_offload/gpu_host.rs index 208417b98bcae..45fcf6cf5c3ce 100644 --- a/tests/codegen-llvm/gpu_offload/gpu_host.rs +++ b/tests/codegen-llvm/gpu_offload/gpu_host.rs @@ -7,8 +7,8 @@ // Better documentation to what each global or variable means is available in the gpu offload code, // or the LLVM offload documentation. +#![feature(gpu_offload)] #![feature(rustc_attrs)] -#![feature(core_intrinsics)] #![no_main] #[unsafe(no_mangle)] @@ -21,7 +21,12 @@ fn main() { } pub fn kernel_1(x: &mut [f32; 256], y: &[f32; 256]) { - core::intrinsics::offload(_kernel_1, [256, 1, 1], [32, 1, 1], 0, -1, (x, y)) + core::offload::offload! { + kernel = _kernel_1, + workgroup_dim = [256, 1, 1], + thread_dim = [32, 1, 1], + args = (x, y), + } } #[inline(never)] @@ -97,7 +102,7 @@ pub fn _kernel_1(x: &mut [f32; 256], y: &[f32; 256]) { // CHECK: declare void @__tgt_register_lib(ptr) local_unnamed_addr // CHECK: declare void @__tgt_unregister_lib(ptr) local_unnamed_addr -// CHECK-LABEL: define internal void @.omp_offloading.descriptor_reg() section ".text.startup" !guid !{{[0-9]+}} { +// CHECK-LABEL: define internal void @.omp_offloading.descriptor_reg() section ".text.startup" // CHECK-NEXT: entry: // CHECK-NEXT: {{tail }}call void @__tgt_register_lib(ptr nonnull @.omp_offloading.descriptor) // CHECK-NEXT: {{tail }}call void @__tgt_init_all_rtls() @@ -105,7 +110,7 @@ pub fn _kernel_1(x: &mut [f32; 256], y: &[f32; 256]) { // CHECK-NEXT: ret void // CHECK-NEXT: } -// CHECK-LABEL: define internal void @.omp_offloading.descriptor_unreg() section ".text.startup" !guid !{{[0-9]+}} { +// CHECK-LABEL: define internal void @.omp_offloading.descriptor_unreg() section ".text.startup" // CHECK-NEXT: entry: // CHECK-NEXT: {{tail }}call void @__tgt_unregister_lib(ptr nonnull @.omp_offloading.descriptor) // CHECK-NEXT: ret void diff --git a/tests/codegen-llvm/gpu_offload/scalar_host.rs b/tests/codegen-llvm/gpu_offload/scalar_host.rs index 950306c83a43f..807d08ddf1893 100644 --- a/tests/codegen-llvm/gpu_offload/scalar_host.rs +++ b/tests/codegen-llvm/gpu_offload/scalar_host.rs @@ -6,8 +6,8 @@ // the kernel as i64 #![feature(abi_gpu_kernel)] +#![feature(gpu_offload)] #![feature(rustc_attrs)] -#![feature(core_intrinsics)] #![no_main] // CHECK: define{{( dso_local)?}} void @main() @@ -28,14 +28,10 @@ fn main() { let mut x = 0.0f32; let k = core::hint::black_box(42.0f32); - core::intrinsics::offload::<_, _, ()>( - foo, - [1, 1, 1], - [1, 1, 1], - 0, - -1, - (&mut x as *mut f32, k), - ); + core::offload::offload! { + kernel = foo, + args = (&mut x as *mut f32, k), + } } unsafe extern "C" { diff --git a/tests/codegen-llvm/gpu_offload/slice_host.rs b/tests/codegen-llvm/gpu_offload/slice_host.rs index 1b314d489612c..ad47d2e76360a 100644 --- a/tests/codegen-llvm/gpu_offload/slice_host.rs +++ b/tests/codegen-llvm/gpu_offload/slice_host.rs @@ -5,8 +5,8 @@ // This test verifies that offload is properly handling slices passing them properly to the device #![feature(abi_gpu_kernel)] +#![feature(gpu_offload)] #![feature(rustc_attrs)] -#![feature(core_intrinsics)] #![no_main] // CHECK: @anon.[[ID:.*]].0 = private unnamed_addr constant [23 x i8] c";unknown;unknown;0;0;;\00", align 1 @@ -27,14 +27,10 @@ #[unsafe(no_mangle)] fn main() { let mut x = [0.0f32, 0.0, 0.0, 0.0]; - core::intrinsics::offload::<_, _, ()>( - foo, - [1, 1, 1], - [1, 1, 1], - 0, - -1, - ((&mut x) as &mut [f32],), - ); + core::offload::offload! { + kernel = foo, + args = ((&mut x) as &mut [f32],), + } } unsafe extern "C" { diff --git a/tests/run-make/offload-generic-manifest/generic.rs b/tests/run-make/offload-generic-manifest/generic.rs index eb356ad05c574..a6b8a7368858d 100644 --- a/tests/run-make/offload-generic-manifest/generic.rs +++ b/tests/run-make/offload-generic-manifest/generic.rs @@ -1,4 +1,4 @@ -#![feature(core_intrinsics, rustc_attrs)] +#![feature(gpu_offload, rustc_attrs)] #![allow(internal_features)] #![cfg_attr(device, no_main)] @@ -7,6 +7,12 @@ fn kernel(x: T) {} #[cfg(not(device))] fn main() { - core::intrinsics::offload::<_, _, ()>(kernel::, [1, 1, 1], [1, 1, 1], 0, (0.0f32,)); - core::intrinsics::offload::<_, _, ()>(kernel::, [1, 1, 1], [1, 1, 1], 0, (0i32,)); + core::offload::offload! { + kernel = kernel::, + args = (0.0f32,), + } + core::offload::offload! { + kernel = kernel::, + args = (0i32,), + } } diff --git a/tests/ui/offload/check_config.rs b/tests/ui/offload/check_config.rs index 833ba39dffeca..63388ce69ba62 100644 --- a/tests/ui/offload/check_config.rs +++ b/tests/ui/offload/check_config.rs @@ -9,7 +9,7 @@ //[fail]~? ERROR: using the offload feature requires -Z offload= //[fail]~? ERROR: using the offload feature requires -C lto=fat -#![feature(core_intrinsics)] +#![feature(gpu_offload)] fn main() { let mut x = [3.0; 256]; @@ -17,7 +17,10 @@ fn main() { } fn kernel_1(x: &mut [f32; 256]) { - core::intrinsics::offload(_kernel_1, [1, 1, 1], [1, 1, 1], 0, -1, (x,)) + core::offload::offload! { + kernel = _kernel_1, + args = (x,), + } } fn _kernel_1(x: &mut [f32; 256]) {} diff --git a/tests/ui/offload/offload_macro.rs b/tests/ui/offload/offload_macro.rs index 468820f08c291..4480f7dd9c80b 100644 --- a/tests/ui/offload/offload_macro.rs +++ b/tests/ui/offload/offload_macro.rs @@ -26,4 +26,7 @@ fn main() { core::offload::offload! { kernel = kernel, args = (), dyn_cache = 0, dyn_cache = 8 } //~^ ERROR duplicate field `dyn_cache` + + core::offload::offload! { kernel = kernel, args = (), device = 0, device = 1 } + //~^ ERROR duplicate field `device` } diff --git a/tests/ui/offload/offload_macro.stderr b/tests/ui/offload/offload_macro.stderr index cd85afeab373e..e2517a2f8cca1 100644 --- a/tests/ui/offload/offload_macro.stderr +++ b/tests/ui/offload/offload_macro.stderr @@ -62,5 +62,13 @@ LL | core::offload::offload! { kernel = kernel, args = (), dyn_cache = 0, dy | = note: this error originates in the macro `$crate::offload` which comes from the expansion of the macro `core::offload::offload` (in Nightly builds, run with -Z macro-backtrace for more info) -error: aborting due to 8 previous errors +error: duplicate field `device` + --> $DIR/offload_macro.rs:30:5 + | +LL | core::offload::offload! { kernel = kernel, args = (), device = 0, device = 1 } + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + | + = note: this error originates in the macro `$crate::offload` which comes from the expansion of the macro `core::offload::offload` (in Nightly builds, run with -Z macro-backtrace for more info) + +error: aborting due to 9 previous errors diff --git a/tests/ui/offload/offload_negative_device.rs b/tests/ui/offload/offload_negative_device.rs new file mode 100644 index 0000000000000..b47dc5f085315 --- /dev/null +++ b/tests/ui/offload/offload_negative_device.rs @@ -0,0 +1,12 @@ +//@ compile-flags: -Zunstable-options -Zoffload=Test -Clto=fat --emit=llvm-ir -Zdeduplicate-diagnostics=yes +//@ no-prefer-dynamic +//@ needs-offload + +#![feature(gpu_offload)] + +fn kernel() {} + +fn main() { + core::offload::offload! { kernel = kernel, args = (), device = -1 } + //~^ ERROR evaluation panicked: offload device must be non-negative; omit `device` to use the default device +} diff --git a/tests/ui/offload/offload_negative_device.stderr b/tests/ui/offload/offload_negative_device.stderr new file mode 100644 index 0000000000000..4be386dd436da --- /dev/null +++ b/tests/ui/offload/offload_negative_device.stderr @@ -0,0 +1,19 @@ +error[E0080]: evaluation panicked: offload device must be non-negative; omit `device` to use the default device + --> $DIR/offload_negative_device.rs:10:5 + | +LL | core::offload::offload! { kernel = kernel, args = (), device = -1 } + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ evaluation of `main::{constant#0}` failed here + | + = note: this error originates in the macro `$crate::panic::panic_2021` which comes from the expansion of the macro `core::offload::offload` (in Nightly builds, run with -Z macro-backtrace for more info) + +note: erroneous constant encountered + --> $DIR/offload_negative_device.rs:10:5 + | +LL | core::offload::offload! { kernel = kernel, args = (), device = -1 } + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + | + = note: this note originates in the macro `$crate::offload` which comes from the expansion of the macro `core::offload::offload` (in Nightly builds, run with -Z macro-backtrace for more info) + +error: aborting due to 1 previous error + +For more information about this error, try `rustc --explain E0080`.