From ad8352b497bf1c2e4bf421bf69a74bb0b9f5a079 Mon Sep 17 00:00:00 2001 From: Sa4dUs Date: Thu, 18 Jun 2026 12:19:34 +0200 Subject: [PATCH 1/4] Add `Region` and `PartitioningStrategy` to `core::offload` module --- library/core/src/offload/mod.rs | 90 +++++++++++++++++++++++++++++++++ 1 file changed, 90 insertions(+) diff --git a/library/core/src/offload/mod.rs b/library/core/src/offload/mod.rs index 3d85621361209..9a85d325ffd7d 100644 --- a/library/core/src/offload/mod.rs +++ b/library/core/src/offload/mod.rs @@ -1,6 +1,7 @@ // offload module #[unstable(feature = "gpu_offload", issue = "131513")] pub use crate::macros::builtin::offload_kernel; +use crate::marker::PhantomData; #[unstable(feature = "gpu_offload", issue = "131513")] pub use crate::offload; @@ -125,3 +126,92 @@ macro_rules! offload { device } }; } + +// Region & Partitioning Strategy + +/// Defines how execution units access memory regions. +/// +/// # Safety +/// +/// Implementations must guarantee that generated views are disjoint. +#[unstable(feature = "offload", issue = "124509")] +pub unsafe trait PartitioningStrategy { + /// Read-only view type for the partitioned memory region. + type View<'a, T: 'a>; + + /// Mutable view type for the partitioned memory region. + type ViewMut<'a, T: 'a>; + + /// Returns the execution index of the current unit. + fn index() -> usize; + + /// Returns a read-only view of the region for the current execution context. + /// + /// # Safety + /// + /// `ptr` must point to `len` valid, initialized elements of type `T`. + /// The memory must stay valid for lifetime `'a`. + unsafe fn get<'a, T>(ptr: *const T, len: usize) -> Option>; + + /// Returns a mutable view of the region for the current execution context. + /// + /// # Safety + /// + /// `ptr` must point to `len` valid, initialized elements of type `T`. + /// The memory must stay valid for lifetime `'a`. + /// The returned view must be disjoint from all other active views. + unsafe fn get_mut<'a, T>(ptr: *mut T, len: usize) -> Option>; +} + +/// A memory region bound to a partitioning strategy. +#[derive(Copy, Clone)] +#[unstable(feature = "offload", issue = "124509")] +pub struct Region<'a, T, S: PartitioningStrategy> { + ptr: *mut T, + len: usize, + _marker: core::marker::PhantomData<(&'a mut [T], S)>, +} + +/// Raw representation used to build a [`Region`] from common aggregate types. +struct RawRegion<'a, T> { + pub ptr: *mut T, + pub len: usize, + _marker: core::marker::PhantomData<&'a mut [T]>, +} + +impl<'a, T> From<&'a mut [T]> for RawRegion<'a, T> { + fn from(data: &'a mut [T]) -> Self { + Self { ptr: data.as_mut_ptr(), len: data.len(), _marker: core::marker::PhantomData } + } +} + +impl<'a, T, const N: usize> From<&'a mut [T; N]> for RawRegion<'a, T> { + fn from(data: &'a mut [T; N]) -> Self { + Self { ptr: data.as_mut_ptr(), len: N, _marker: core::marker::PhantomData } + } +} + +#[unstable(feature = "offload", issue = "124509")] +impl<'a, T, S: PartitioningStrategy> Region<'a, T, S> { + /// Creates a new partitioned region from data convertible into a [`RawRegion`]. + pub fn new(data: D) -> Self + where + D: Into>, + { + let raw = data.into(); + Self { ptr: raw.ptr, len: raw.len, _marker: core::marker::PhantomData } + } + + /// Returns a read-only view for the current execution context. + pub fn get(&self) -> Option> { + // SAFETY: `self.ptr` points to `self.len` valid elements for lifetime `'a`. + unsafe { S::get(self.ptr as *const T, self.len) } + } + + /// Returns a mutable view for the current execution context. + pub fn get_mut(&mut self) -> Option> { + // SAFETY: `self.ptr` points to `self.len` valid elements for lifetime `'a`. + // The strategy guarantees that the returned view is disjoint. + unsafe { S::get_mut(self.ptr, self.len) } + } +} From 37dbb16f79a500d4680c3b7fa39374270f8fc3e1 Mon Sep 17 00:00:00 2001 From: Marcelo Dominguez Date: Thu, 13 Aug 2026 17:53:16 +0300 Subject: [PATCH 2/4] Map `Region` the same way as slices --- compiler/rustc_middle/src/ty/offload_meta.rs | 17 +++++ compiler/rustc_span/src/symbol.rs | 1 + library/core/src/offload/mod.rs | 12 ++-- tests/codegen-llvm/gpu_offload/region_host.rs | 65 +++++++++++++++++++ 4 files changed, 90 insertions(+), 5 deletions(-) create mode 100644 tests/codegen-llvm/gpu_offload/region_host.rs diff --git a/compiler/rustc_middle/src/ty/offload_meta.rs b/compiler/rustc_middle/src/ty/offload_meta.rs index a58e517e05f61..c689c67397b2b 100644 --- a/compiler/rustc_middle/src/ty/offload_meta.rs +++ b/compiler/rustc_middle/src/ty/offload_meta.rs @@ -1,5 +1,6 @@ use bitflags::bitflags; use rustc_abi::{BackendRepr, TyAbiInterface}; +use rustc_span::sym; use rustc_target::callconv::ArgAbi; use crate::ty::{self, PseudoCanonicalInput, Ty, TyCtxt, TypingEnv}; @@ -75,6 +76,12 @@ impl OffloadMetadata { where Ty<'tcx>: TyAbiInterface<'tcx, C>, { + if let Some(elem_ty) = region_element_ty(tcx, ty) { + let ptr = OffloadMetadata::from_ty(tcx, Ty::new_slice(tcx, elem_ty)); + let len = OffloadMetadata::from_ty(tcx, tcx.types.usize); + return vec![(ptr, Ty::new_mut_ptr(tcx, elem_ty)), (len, tcx.types.usize)]; + } + match arg_abi.layout.backend_repr { BackendRepr::ScalarPair { a: _, b: _, b_offset: _ } => (0..2) .map(|i| { @@ -87,6 +94,16 @@ impl OffloadMetadata { } } +fn region_element_ty<'tcx>(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>) -> Option> { + if let ty::Adt(def, args) = ty.kind() + && Some(def.did()) == tcx.get_diagnostic_item(sym::offload_region) + { + Some(args.type_at(1)) + } else { + None + } +} + // FIXME(Sa4dUs): implement a solid logic to determine the payload size fn get_payload_size<'tcx>(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>) -> OffloadSize { match ty.kind() { diff --git a/compiler/rustc_span/src/symbol.rs b/compiler/rustc_span/src/symbol.rs index 4c1b9c78b7963..54e9a3bdbab2a 100644 --- a/compiler/rustc_span/src/symbol.rs +++ b/compiler/rustc_span/src/symbol.rs @@ -1489,6 +1489,7 @@ symbols! { offload, offload_get_num_devices, offload_kernel, + offload_region, offset, offset_of, offset_of_enum, diff --git a/library/core/src/offload/mod.rs b/library/core/src/offload/mod.rs index 9a85d325ffd7d..d0854e6722967 100644 --- a/library/core/src/offload/mod.rs +++ b/library/core/src/offload/mod.rs @@ -1,7 +1,6 @@ // offload module #[unstable(feature = "gpu_offload", issue = "131513")] pub use crate::macros::builtin::offload_kernel; -use crate::marker::PhantomData; #[unstable(feature = "gpu_offload", issue = "131513")] pub use crate::offload; @@ -164,8 +163,9 @@ pub unsafe trait PartitioningStrategy { } /// A memory region bound to a partitioning strategy. -#[derive(Copy, Clone)] +#[derive(Copy, Clone, Debug)] #[unstable(feature = "offload", issue = "124509")] +#[rustc_diagnostic_item = "offload_region"] pub struct Region<'a, T, S: PartitioningStrategy> { ptr: *mut T, len: usize, @@ -173,9 +173,11 @@ pub struct Region<'a, T, S: PartitioningStrategy> { } /// Raw representation used to build a [`Region`] from common aggregate types. -struct RawRegion<'a, T> { - pub ptr: *mut T, - pub len: usize, +#[derive(Debug)] +#[unstable(feature = "offload", issue = "124509")] +pub struct RawRegion<'a, T> { + ptr: *mut T, + len: usize, _marker: core::marker::PhantomData<&'a mut [T]>, } diff --git a/tests/codegen-llvm/gpu_offload/region_host.rs b/tests/codegen-llvm/gpu_offload/region_host.rs new file mode 100644 index 0000000000000..602dc247f6540 --- /dev/null +++ b/tests/codegen-llvm/gpu_offload/region_host.rs @@ -0,0 +1,65 @@ +//@ compile-flags: -Zoffload=Test -Zunstable-options -C opt-level=1 -Clto=fat +//@ no-prefer-dynamic +//@ needs-offload + +// This test verifies that a `Region` kernel argument is mapped like a slice. +#![feature(abi_gpu_kernel)] +#![feature(core_intrinsics)] +#![feature(gpu_offload)] +#![feature(offload)] +#![feature(rustc_attrs)] +#![no_main] + +extern crate core; + +use core::offload::{PartitioningStrategy, Region}; + +struct Dummy; + +unsafe impl PartitioningStrategy for Dummy { + type View<'a, T: 'a> = &'a T; + type ViewMut<'a, T: 'a> = &'a mut T; + + fn index() -> usize { + 0 + } + + unsafe fn get<'a, T>(_ptr: *const T, _len: usize) -> Option> { + None + } + + unsafe fn get_mut<'a, T>(_ptr: *mut T, _len: usize) -> Option> { + None + } +} + +// CHECK: @anon.[[ID:.*]].0 = private unnamed_addr constant [23 x i8] c";unknown;unknown;0;0;;\00", align 1 + +// CHECK-DAG: @.offload_sizes.[[K:[^ ]*foo]] = private unnamed_addr constant [2 x i64] [i64 0, i64 8] +// CHECK-DAG: @.offload_maptypes.[[K]].begin = private unnamed_addr constant [2 x i64] [i64 1, i64 768] +// CHECK-DAG: @.offload_maptypes.[[K]].kernel = private unnamed_addr constant [2 x i64] [i64 32, i64 800] +// CHECK-DAG: @.offload_maptypes.[[K]].end = private unnamed_addr constant [2 x i64] [i64 2, i64 0] + +// CHECK: define{{( dso_local)?}} void @main() +// CHECK: %.offload_sizes = alloca [2 x i64], align 8 +// CHECK: call void @llvm.memcpy.p0.p0.i64(ptr {{.*}} %.offload_sizes, ptr {{.*}} @.offload_sizes.[[K]], i64 16, i1 false) +// CHECK: store i64 16, ptr %.offload_sizes, align 8 +// CHECK: call void @__tgt_target_data_begin_mapper(ptr nonnull @anon.[[ID]].1, i64 -1, i32 2, ptr nonnull %.offload_baseptrs, ptr nonnull %.offload_ptrs, ptr nonnull %.offload_sizes, ptr nonnull @.offload_maptypes.[[K]].begin, ptr null, ptr null) +// CHECK: call i32 @__tgt_target_kernel(ptr nonnull @anon.[[ID]].1, i64 -1, i32 1, i32 1, ptr nonnull @.[[K]].region_id, ptr nonnull %kernel_args) +// CHECK-NEXT: call void @__tgt_target_data_end_mapper(ptr nonnull @anon.[[ID]].1, i64 -1, i32 2, ptr nonnull %.offload_baseptrs, ptr nonnull %.offload_ptrs, ptr nonnull %.offload_sizes, ptr nonnull @.offload_maptypes.[[K]].end, ptr null, ptr null) + +#[unsafe(no_mangle)] +fn main() { + let mut x = [0.0f32; 4]; + core::intrinsics::offload::<_, _, ()>( + foo, + [1, 1, 1], + [1, 1, 1], + 0, + (Region::::new(&mut x as &mut [f32]),), + ); +} + +fn foo(region: Region<'_, f32, Dummy>) { + unreachable!(); +} From db165341db2fea5b505635a1397e7cc90ecfbb9f Mon Sep 17 00:00:00 2001 From: Marcelo Dominguez Date: Sat, 5 Sep 2026 13:38:07 +0300 Subject: [PATCH 3/4] Region byval and improve testing --- compiler/rustc_hir_typeck/src/intrinsicck.rs | 21 +++++++++- compiler/rustc_middle/src/ty/offload_meta.rs | 12 ++++-- library/core/src/offload/mod.rs | 16 ++++--- .../auxiliary/offload_strategies.rs | 26 ++++++++++++ tests/codegen-llvm/gpu_offload/region_host.rs | 34 ++++----------- .../offload/auxiliary/offload_strategies.rs | 26 ++++++++++++ tests/ui/offload/region_borrow.rs | 22 ++++++++++ tests/ui/offload/region_borrow.stderr | 15 +++++++ tests/ui/offload/region_by_ref.rs | 42 +++++++++++++++++++ tests/ui/offload/region_by_ref.stderr | 14 +++++++ tests/ui/offload/region_not_copy.rs | 19 +++++++++ tests/ui/offload/region_not_copy.stderr | 13 ++++++ 12 files changed, 225 insertions(+), 35 deletions(-) create mode 100644 tests/codegen-llvm/gpu_offload/auxiliary/offload_strategies.rs create mode 100644 tests/ui/offload/auxiliary/offload_strategies.rs create mode 100644 tests/ui/offload/region_borrow.rs create mode 100644 tests/ui/offload/region_borrow.stderr create mode 100644 tests/ui/offload/region_by_ref.rs create mode 100644 tests/ui/offload/region_by_ref.stderr create mode 100644 tests/ui/offload/region_not_copy.rs create mode 100644 tests/ui/offload/region_not_copy.stderr diff --git a/compiler/rustc_hir_typeck/src/intrinsicck.rs b/compiler/rustc_hir_typeck/src/intrinsicck.rs index d63bffab88221..4b7d38ea1162a 100644 --- a/compiler/rustc_hir_typeck/src/intrinsicck.rs +++ b/compiler/rustc_hir_typeck/src/intrinsicck.rs @@ -7,6 +7,7 @@ use rustc_hir as hir; use rustc_index::Idx; use rustc_middle::bug; use rustc_middle::ty::layout::{LayoutError, SizeSkeleton}; +use rustc_middle::ty::offload_meta::is_region_ty; use rustc_middle::ty::{self, Ty, TyCtxt, Unnormalized}; use rustc_span::ErrorGuaranteed; use rustc_span::def_id::LocalDefId; @@ -135,6 +136,10 @@ fn check_transmute<'tcx>( } } +fn is_offload_region_ref<'tcx>(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>) -> bool { + matches!(ty.kind(), ty::Ref(_, inner, _) if is_region_ty(tcx, *inner)) +} + fn check_offload<'tcx>( tcx: TyCtxt<'tcx>, typing_env: ty::TypingEnv<'tcx>, @@ -206,7 +211,21 @@ fn check_offload<'tcx>( { let norm_input_ty = normalize(input_ty); let norm_arg_ty = normalize(arg_ty); - if norm_input_ty != norm_arg_ty { + + if is_offload_region_ref(tcx, norm_input_ty) || is_offload_region_ref(tcx, norm_arg_ty) { + let err = tcx + .sess + .dcx() + .struct_span_err( + span, + format!( + "offload kernel argument {i} is a reference to a `Region`. Pass the \ + `Region` by value so it can be mapped like a slice" + ), + ) + .emit(); + result = Err(err); + } else if norm_input_ty != norm_arg_ty { let err = tcx .sess .dcx() diff --git a/compiler/rustc_middle/src/ty/offload_meta.rs b/compiler/rustc_middle/src/ty/offload_meta.rs index c689c67397b2b..c51220912eb26 100644 --- a/compiler/rustc_middle/src/ty/offload_meta.rs +++ b/compiler/rustc_middle/src/ty/offload_meta.rs @@ -94,10 +94,16 @@ impl OffloadMetadata { } } +pub fn is_region_ty<'tcx>(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>) -> bool { + matches!( + ty.kind(), + ty::Adt(def, _) if Some(def.did()) == tcx.get_diagnostic_item(sym::offload_region) + ) +} + fn region_element_ty<'tcx>(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>) -> Option> { - if let ty::Adt(def, args) = ty.kind() - && Some(def.did()) == tcx.get_diagnostic_item(sym::offload_region) - { + if is_region_ty(tcx, ty) { + let ty::Adt(_, args) = ty.kind() else { unreachable!() }; Some(args.type_at(1)) } else { None diff --git a/library/core/src/offload/mod.rs b/library/core/src/offload/mod.rs index d0854e6722967..33d69ceb37d60 100644 --- a/library/core/src/offload/mod.rs +++ b/library/core/src/offload/mod.rs @@ -133,7 +133,7 @@ macro_rules! offload { /// # Safety /// /// Implementations must guarantee that generated views are disjoint. -#[unstable(feature = "offload", issue = "124509")] +#[unstable(feature = "offload", issue = "131513")] pub unsafe trait PartitioningStrategy { /// Read-only view type for the partitioned memory region. type View<'a, T: 'a>; @@ -163,8 +163,8 @@ pub unsafe trait PartitioningStrategy { } /// A memory region bound to a partitioning strategy. -#[derive(Copy, Clone, Debug)] -#[unstable(feature = "offload", issue = "124509")] +#[derive(Debug)] +#[unstable(feature = "offload", issue = "131513")] #[rustc_diagnostic_item = "offload_region"] pub struct Region<'a, T, S: PartitioningStrategy> { ptr: *mut T, @@ -174,7 +174,7 @@ pub struct Region<'a, T, S: PartitioningStrategy> { /// Raw representation used to build a [`Region`] from common aggregate types. #[derive(Debug)] -#[unstable(feature = "offload", issue = "124509")] +#[unstable(feature = "offload", issue = "131513")] pub struct RawRegion<'a, T> { ptr: *mut T, len: usize, @@ -193,7 +193,7 @@ impl<'a, T, const N: usize> From<&'a mut [T; N]> for RawRegion<'a, T> { } } -#[unstable(feature = "offload", issue = "124509")] +#[unstable(feature = "offload", issue = "131513")] impl<'a, T, S: PartitioningStrategy> Region<'a, T, S> { /// Creates a new partitioned region from data convertible into a [`RawRegion`]. pub fn new(data: D) -> Self @@ -216,4 +216,10 @@ impl<'a, T, S: PartitioningStrategy> Region<'a, T, S> { // The strategy guarantees that the returned view is disjoint. unsafe { S::get_mut(self.ptr, self.len) } } + + /// Reborrows the region, producing a new region that aliases the same memory with + /// the lifetime of the borrow. + pub fn reborrow(&mut self) -> Region<'_, T, S> { + Region { ptr: self.ptr, len: self.len, _marker: core::marker::PhantomData } + } } diff --git a/tests/codegen-llvm/gpu_offload/auxiliary/offload_strategies.rs b/tests/codegen-llvm/gpu_offload/auxiliary/offload_strategies.rs new file mode 100644 index 0000000000000..1e3271e5a8547 --- /dev/null +++ b/tests/codegen-llvm/gpu_offload/auxiliary/offload_strategies.rs @@ -0,0 +1,26 @@ +//@ edition: 2024 + +#![feature(gpu_offload)] +#![feature(offload)] + +use core::offload::PartitioningStrategy; + +#[derive(Debug, Clone, Copy)] +pub struct Dummy; + +unsafe impl PartitioningStrategy for Dummy { + type View<'a, T: 'a> = &'a T; + type ViewMut<'a, T: 'a> = &'a mut T; + + fn index() -> usize { + 0 + } + + unsafe fn get<'a, T>(_ptr: *const T, _len: usize) -> Option> { + None + } + + unsafe fn get_mut<'a, T>(_ptr: *mut T, _len: usize) -> Option> { + None + } +} diff --git a/tests/codegen-llvm/gpu_offload/region_host.rs b/tests/codegen-llvm/gpu_offload/region_host.rs index 602dc247f6540..975096c7fb5d4 100644 --- a/tests/codegen-llvm/gpu_offload/region_host.rs +++ b/tests/codegen-llvm/gpu_offload/region_host.rs @@ -1,6 +1,8 @@ //@ compile-flags: -Zoffload=Test -Zunstable-options -C opt-level=1 -Clto=fat //@ no-prefer-dynamic //@ needs-offload +//@ edition: 2024 +//@ aux-crate: offload_strategies=offload_strategies.rs // This test verifies that a `Region` kernel argument is mapped like a slice. #![feature(abi_gpu_kernel)] @@ -12,26 +14,9 @@ extern crate core; -use core::offload::{PartitioningStrategy, Region}; +use core::offload::Region; -struct Dummy; - -unsafe impl PartitioningStrategy for Dummy { - type View<'a, T: 'a> = &'a T; - type ViewMut<'a, T: 'a> = &'a mut T; - - fn index() -> usize { - 0 - } - - unsafe fn get<'a, T>(_ptr: *const T, _len: usize) -> Option> { - None - } - - unsafe fn get_mut<'a, T>(_ptr: *mut T, _len: usize) -> Option> { - None - } -} +use offload_strategies::Dummy; // CHECK: @anon.[[ID:.*]].0 = private unnamed_addr constant [23 x i8] c";unknown;unknown;0;0;;\00", align 1 @@ -51,13 +36,10 @@ unsafe impl PartitioningStrategy for Dummy { #[unsafe(no_mangle)] fn main() { let mut x = [0.0f32; 4]; - core::intrinsics::offload::<_, _, ()>( - foo, - [1, 1, 1], - [1, 1, 1], - 0, - (Region::::new(&mut x as &mut [f32]),), - ); + core::offload::offload! { + kernel = foo, + args = (Region::::new(&mut x as &mut [f32]),), + }; } fn foo(region: Region<'_, f32, Dummy>) { diff --git a/tests/ui/offload/auxiliary/offload_strategies.rs b/tests/ui/offload/auxiliary/offload_strategies.rs new file mode 100644 index 0000000000000..1e3271e5a8547 --- /dev/null +++ b/tests/ui/offload/auxiliary/offload_strategies.rs @@ -0,0 +1,26 @@ +//@ edition: 2024 + +#![feature(gpu_offload)] +#![feature(offload)] + +use core::offload::PartitioningStrategy; + +#[derive(Debug, Clone, Copy)] +pub struct Dummy; + +unsafe impl PartitioningStrategy for Dummy { + type View<'a, T: 'a> = &'a T; + type ViewMut<'a, T: 'a> = &'a mut T; + + fn index() -> usize { + 0 + } + + unsafe fn get<'a, T>(_ptr: *const T, _len: usize) -> Option> { + None + } + + unsafe fn get_mut<'a, T>(_ptr: *mut T, _len: usize) -> Option> { + None + } +} diff --git a/tests/ui/offload/region_borrow.rs b/tests/ui/offload/region_borrow.rs new file mode 100644 index 0000000000000..d1424d1cde35c --- /dev/null +++ b/tests/ui/offload/region_borrow.rs @@ -0,0 +1,22 @@ +//@ edition: 2024 +//@ aux-crate: offload_strategies=offload_strategies.rs + +// This test checks that the borrow checker errors when writing into the data a +// `Region` was created from while the `Region` is still alive. + +#![feature(gpu_offload)] +#![feature(offload)] +#![allow(unused_assignments)] + +use core::offload::Region; +use offload_strategies::Dummy; + +fn main() { + let mut x = [0.0f32; 4]; + let region = Region::::new(&mut x[..]); + + x[0] = 1.0; + //~^ ERROR cannot assign to `x[_]` because it is borrowed + + let _view = region.get(); +} diff --git a/tests/ui/offload/region_borrow.stderr b/tests/ui/offload/region_borrow.stderr new file mode 100644 index 0000000000000..4f7a153bf30fc --- /dev/null +++ b/tests/ui/offload/region_borrow.stderr @@ -0,0 +1,15 @@ +error[E0506]: cannot assign to `x[_]` because it is borrowed + --> $DIR/region_borrow.rs:18:5 + | +LL | let region = Region::::new(&mut x[..]); + | - `x[_]` is borrowed here +LL | +LL | x[0] = 1.0; + | ^^^^^^^^^^ `x[_]` is assigned to here but it was already borrowed +... +LL | let _view = region.get(); + | ------ borrow later used here + +error: aborting due to 1 previous error + +For more information about this error, try `rustc --explain E0506`. diff --git a/tests/ui/offload/region_by_ref.rs b/tests/ui/offload/region_by_ref.rs new file mode 100644 index 0000000000000..8bdc92815525f --- /dev/null +++ b/tests/ui/offload/region_by_ref.rs @@ -0,0 +1,42 @@ +//@ compile-flags: -Zunstable-options -Zoffload=Device -Clto=fat +//@ edition: 2024 +//@ aux-crate: offload_strategies=offload_strategies.rs + +// This tests ensures an error is emmited with passing a `&Region<'_, _, _>` args to offload. + +#![feature(core_intrinsics)] +#![feature(gpu_offload)] +#![feature(offload)] + +use core::offload::Region; +use offload_strategies::Dummy; + +fn kernel_shared(_region: &Region<'_, f32, Dummy>) {} + +fn kernel_mut(_region: &mut Region<'_, f32, Dummy>) {} + +fn main() { + let mut x = [0.0f32; 4]; + let region = Region::::new(&mut x[..]); + core::intrinsics::offload::<_, _, ()>( + //~^ ERROR offload kernel argument 0 is a reference to a `Region` + kernel_shared, + [1, 1, 1], + [1, 1, 1], + 0, + -1, + (®ion,), + ); + + let mut y = [0.0f32; 4]; + let mut region = Region::::new(&mut y[..]); + core::intrinsics::offload::<_, _, ()>( + //~^ ERROR offload kernel argument 0 is a reference to a `Region` + kernel_mut, + [1, 1, 1], + [1, 1, 1], + 0, + -1, + (&mut region,), + ); +} diff --git a/tests/ui/offload/region_by_ref.stderr b/tests/ui/offload/region_by_ref.stderr new file mode 100644 index 0000000000000..25c278983ea38 --- /dev/null +++ b/tests/ui/offload/region_by_ref.stderr @@ -0,0 +1,14 @@ +error: offload kernel argument 0 is a reference to a `Region`. Pass the `Region` by value so it can be mapped like a slice + --> $DIR/region_by_ref.rs:21:5 + | +LL | core::intrinsics::offload::<_, _, ()>( + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +error: offload kernel argument 0 is a reference to a `Region`. Pass the `Region` by value so it can be mapped like a slice + --> $DIR/region_by_ref.rs:33:5 + | +LL | core::intrinsics::offload::<_, _, ()>( + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +error: aborting due to 2 previous errors + diff --git a/tests/ui/offload/region_not_copy.rs b/tests/ui/offload/region_not_copy.rs new file mode 100644 index 0000000000000..2ab8253b89f73 --- /dev/null +++ b/tests/ui/offload/region_not_copy.rs @@ -0,0 +1,19 @@ +//@ edition: 2024 +//@ aux-crate: offload_strategies=offload_strategies.rs + +// This tests checks `Region` doesn't implement Copy. + +#![feature(gpu_offload)] +#![feature(offload)] + +use core::offload::Region; +use offload_strategies::Dummy; + +fn main() { + let mut x = [0.0f32; 4]; + let mut a = Region::::new(&mut x[..]); + let mut b = a; + if let (Some(_), Some(_)) = (a.get_mut(), b.get_mut()) { + //~^ ERROR borrow of moved value: `a` + } +} diff --git a/tests/ui/offload/region_not_copy.stderr b/tests/ui/offload/region_not_copy.stderr new file mode 100644 index 0000000000000..9af4bc064585d --- /dev/null +++ b/tests/ui/offload/region_not_copy.stderr @@ -0,0 +1,13 @@ +error[E0382]: borrow of moved value: `a` + --> $DIR/region_not_copy.rs:16:34 + | +LL | let mut a = Region::::new(&mut x[..]); + | ----- move occurs because `a` has type `Region<'_, f32, Dummy>`, which does not implement the `Copy` trait +LL | let mut b = a; + | - value moved here +LL | if let (Some(_), Some(_)) = (a.get_mut(), b.get_mut()) { + | ^ value borrowed here after move + +error: aborting due to 1 previous error + +For more information about this error, try `rustc --explain E0382`. From 3b4f613c61c4f60fe9ebc5ee41363f7650d3ac11 Mon Sep 17 00:00:00 2001 From: Marcelo Dominguez Date: Sun, 13 Sep 2026 14:39:53 +0200 Subject: [PATCH 4/4] Fixes and launch checks --- compiler/rustc_hir_typeck/src/intrinsicck.rs | 12 +- compiler/rustc_middle/src/ty/offload_meta.rs | 15 +- library/core/src/offload/mod.rs | 182 ++++++++++++++---- .../auxiliary/offload_strategies.rs | 39 +++- .../offload/auxiliary/offload_strategies.rs | 39 +++- tests/ui/offload/region_by_ref.rs | 21 +- tests/ui/offload/region_by_ref.stderr | 16 +- tests/ui/offload/region_launch_bounds.rs | 19 ++ 8 files changed, 284 insertions(+), 59 deletions(-) create mode 100644 tests/ui/offload/region_launch_bounds.rs diff --git a/compiler/rustc_hir_typeck/src/intrinsicck.rs b/compiler/rustc_hir_typeck/src/intrinsicck.rs index 4b7d38ea1162a..0ed4085a1a2d6 100644 --- a/compiler/rustc_hir_typeck/src/intrinsicck.rs +++ b/compiler/rustc_hir_typeck/src/intrinsicck.rs @@ -136,8 +136,8 @@ fn check_transmute<'tcx>( } } -fn is_offload_region_ref<'tcx>(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>) -> bool { - matches!(ty.kind(), ty::Ref(_, inner, _) if is_region_ty(tcx, *inner)) +fn contains_nested_offload_region<'tcx>(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>) -> bool { + ty.walk().skip(1).any(|arg| arg.as_type().is_some_and(|ty| is_region_ty(tcx, ty))) } fn check_offload<'tcx>( @@ -212,15 +212,17 @@ fn check_offload<'tcx>( let norm_input_ty = normalize(input_ty); let norm_arg_ty = normalize(arg_ty); - if is_offload_region_ref(tcx, norm_input_ty) || is_offload_region_ref(tcx, norm_arg_ty) { + if contains_nested_offload_region(tcx, norm_input_ty) + || contains_nested_offload_region(tcx, norm_arg_ty) + { let err = tcx .sess .dcx() .struct_span_err( span, format!( - "offload kernel argument {i} is a reference to a `Region`. Pass the \ - `Region` by value so it can be mapped like a slice" + "offload kernel argument {i} contains a `Region` nested inside another \ + type. Pass the `Region` by value so it can be mapped like a slice" ), ) .emit(); diff --git a/compiler/rustc_middle/src/ty/offload_meta.rs b/compiler/rustc_middle/src/ty/offload_meta.rs index c51220912eb26..6cf8abc482260 100644 --- a/compiler/rustc_middle/src/ty/offload_meta.rs +++ b/compiler/rustc_middle/src/ty/offload_meta.rs @@ -79,7 +79,20 @@ impl OffloadMetadata { if let Some(elem_ty) = region_element_ty(tcx, ty) { let ptr = OffloadMetadata::from_ty(tcx, Ty::new_slice(tcx, elem_ty)); let len = OffloadMetadata::from_ty(tcx, tcx.types.usize); - return vec![(ptr, Ty::new_mut_ptr(tcx, elem_ty)), (len, tcx.types.usize)]; + // `Region` is a `{ ptr, len }` pair, but field order is not guaranteed. + return arg_abi + .layout + .fields + .index_by_increasing_offset() + .filter(|&i| arg_abi.layout.field(cx, i).size.bytes() != 0) + .map(|i| { + if arg_abi.layout.field(cx, i).ty == tcx.types.usize { + (len, tcx.types.usize) + } else { + (ptr, Ty::new_mut_ptr(tcx, elem_ty)) + } + }) + .collect(); } match arg_abi.layout.backend_repr { diff --git a/library/core/src/offload/mod.rs b/library/core/src/offload/mod.rs index 33d69ceb37d60..b3cd2b8b5cd55 100644 --- a/library/core/src/offload/mod.rs +++ b/library/core/src/offload/mod.rs @@ -1,4 +1,6 @@ // offload module +use core::ptr::NonNull; + #[unstable(feature = "gpu_offload", issue = "131513")] pub use crate::macros::builtin::offload_kernel; #[unstable(feature = "gpu_offload", issue = "131513")] @@ -37,58 +39,52 @@ pub use crate::offload; /// ``` #[macro_export] #[unstable(feature = "gpu_offload", issue = "131513")] -#[allow_internal_unstable(core_intrinsics)] +#[allow_internal_unstable(core_intrinsics, offload)] macro_rules! offload { - ( $($field:ident = $val:expr),* $(,)? ) => { - $crate::offload!(@munch - [ $($field = $val),* ]; - kernel = NONE; - 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; 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:tt)*)?]; kernel = NONE; workgroup_dim = $w:tt; thread_dim = $t:tt; dyn_cache = $d:tt; device = $device:tt; args = $a:tt) => { + $crate::offload!(@munch [ $($($rest)*)? ]; 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; device = $device:tt; args = $a:tt) => { + (@munch [kernel = $val:expr $(, $($rest:tt)*)?]; 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; 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:tt)*)?]; 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)*)? ]; 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; device = $device:tt; args = $a:tt) => { + (@munch [workgroup_dim = $val:expr $(, $($rest:tt)*)?]; 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; 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:tt)*)?]; 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)*)? ]; 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; device = $device:tt; args = $a:tt) => { + (@munch [thread_dim = $val:expr $(, $($rest:tt)*)?]; 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); 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:tt)*)?]; kernel = $k:tt; workgroup_dim = $w:tt; thread_dim = $t:tt; dyn_cache = (0); device = $device:tt; args = $a:tt) => { + $crate::offload!(@munch [ $($($rest)*)? ]; 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); device = $device:tt; args = $a:tt) => { + (@munch [dyn_cache = $val:expr $(, $($rest:tt)*)?]; 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 [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 [device = $val:expr $(, $($rest:tt)*)?]; kernel = $k:tt; workgroup_dim = $w:tt; thread_dim = $t:tt; dyn_cache = $d:tt; device = NONE; args = $a:tt) => { + $crate::offload!(@munch [ $($($rest)*)? ]; kernel = $k; workgroup_dim = $w; thread_dim = $t; dyn_cache = $d; device = (SOME $val); args = $a) }; - (@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) => { + (@munch [device = $val:expr $(, $($rest:tt)*)?]; 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 = ($($arg:expr),* $(,)?) $(, $($rest:tt)*)?]; kernel = $k:tt; workgroup_dim = $w:tt; thread_dim = $t:tt; dyn_cache = $d:tt; device = $device:tt; args = NONE) => { + $crate::offload!(@munch [ $($($rest)*)? ]; kernel = $k; workgroup_dim = $w; thread_dim = $t; dyn_cache = $d; device = $device; args = (SOME_TUPLE ($($arg),*))) + }; + (@munch [args = $val:expr $(, $($rest:tt)*)?]; kernel = $k:tt; workgroup_dim = $w:tt; thread_dim = $t:tt; dyn_cache = $d:tt; device = $device:tt; args = NONE) => { + $crate::offload!(@munch [ $($($rest)*)? ]; 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)) => { + (@munch [args = $val:expr $(, $($rest:tt)*)?]; 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 [args = $val:expr $(, $($rest:tt)*)?]; kernel = $k:tt; workgroup_dim = $w:tt; thread_dim = $t:tt; dyn_cache = $d:tt; device = $device:tt; args = (SOME_TUPLE $old:tt)) => { 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; device = $device:tt; args = $a:tt) => { + (@munch [$invalid:ident = $val:expr $(, $($rest:tt)*)?]; 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), "`")) }; @@ -98,6 +94,26 @@ macro_rules! offload { (@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; device = $device:tt; args = (SOME_TUPLE ($($arg:expr),* $(,)?))) => {{ + #[allow(unused_imports)] + use $crate::offload::{DefaultLaunchCheck as _, RegionLaunchCheck as _}; + + let __offload_grid = $crate::offload!(@value $w); + let __offload_block = $crate::offload!(@value $t); + + $crate::intrinsics::offload::<_, _, ()>( + $kernel, + __offload_grid, + __offload_block, + $crate::offload!(@value $d), + $crate::offload!(@device $device), + ($({ + let __offload_arg = $arg; + (&__offload_arg).__offload_check_launch(__offload_grid, __offload_block); + __offload_arg + },)*), + ) + }}; (@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, @@ -124,10 +140,45 @@ macro_rules! offload { ); device } }; + + ( $($tt:tt)* ) => { + $crate::offload!(@munch + [ $($tt)* ]; + kernel = NONE; + workgroup_dim = ([1, 1, 1]); + thread_dim = ([1, 1, 1]); + dyn_cache = (0); + device = NONE; + args = NONE + ) + }; } // Region & Partitioning Strategy +/// Error returned by [`PartitioningStrategy::check_launch`] when a launch +/// configuration is not compatible with a partitioning strategy. +#[derive(Debug)] +#[unstable(feature = "offload", issue = "131513")] +pub struct LaunchError { + message: &'static str, +} + +impl LaunchError { + /// Creates a new [`LaunchError`] with the given message. + #[unstable(feature = "offload", issue = "131513")] + pub const fn new(message: &'static str) -> Self { + Self { message } + } +} + +#[unstable(feature = "offload", issue = "131513")] +impl core::fmt::Display for LaunchError { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + f.write_str(self.message) + } +} + /// Defines how execution units access memory regions. /// /// # Safety @@ -150,7 +201,7 @@ pub unsafe trait PartitioningStrategy { /// /// `ptr` must point to `len` valid, initialized elements of type `T`. /// The memory must stay valid for lifetime `'a`. - unsafe fn get<'a, T>(ptr: *const T, len: usize) -> Option>; + unsafe fn get<'a, T>(ptr: NonNull, len: usize) -> Option>; /// Returns a mutable view of the region for the current execution context. /// @@ -159,7 +210,13 @@ pub unsafe trait PartitioningStrategy { /// `ptr` must point to `len` valid, initialized elements of type `T`. /// The memory must stay valid for lifetime `'a`. /// The returned view must be disjoint from all other active views. - unsafe fn get_mut<'a, T>(ptr: *mut T, len: usize) -> Option>; + unsafe fn get_mut<'a, T>(ptr: NonNull, len: usize) -> Option>; + + /// Checks that a kernel using this strategy can be launched with `len` + /// elements on a `grid` by `block` launch configuration. + /// + /// This is called automatically before the kernel is launched. + fn check_launch(len: usize, grid: [u32; 3], block: [u32; 3]) -> Result<(), LaunchError>; } /// A memory region bound to a partitioning strategy. @@ -167,7 +224,7 @@ pub unsafe trait PartitioningStrategy { #[unstable(feature = "offload", issue = "131513")] #[rustc_diagnostic_item = "offload_region"] pub struct Region<'a, T, S: PartitioningStrategy> { - ptr: *mut T, + ptr: NonNull, len: usize, _marker: core::marker::PhantomData<(&'a mut [T], S)>, } @@ -176,20 +233,24 @@ pub struct Region<'a, T, S: PartitioningStrategy> { #[derive(Debug)] #[unstable(feature = "offload", issue = "131513")] pub struct RawRegion<'a, T> { - ptr: *mut T, + ptr: NonNull, len: usize, _marker: core::marker::PhantomData<&'a mut [T]>, } impl<'a, T> From<&'a mut [T]> for RawRegion<'a, T> { fn from(data: &'a mut [T]) -> Self { - Self { ptr: data.as_mut_ptr(), len: data.len(), _marker: core::marker::PhantomData } + // SAFETY: `data.as_mut_ptr()` is non-null, because it is derived from a reference. + let ptr = unsafe { NonNull::new_unchecked(data.as_mut_ptr()) }; + Self { ptr, len: data.len(), _marker: core::marker::PhantomData } } } impl<'a, T, const N: usize> From<&'a mut [T; N]> for RawRegion<'a, T> { fn from(data: &'a mut [T; N]) -> Self { - Self { ptr: data.as_mut_ptr(), len: N, _marker: core::marker::PhantomData } + // SAFETY: `data.as_mut_ptr()` is non-null, because it is derived from a reference. + let ptr = unsafe { NonNull::new_unchecked(data.as_mut_ptr()) }; + Self { ptr, len: N, _marker: core::marker::PhantomData } } } @@ -207,7 +268,7 @@ impl<'a, T, S: PartitioningStrategy> Region<'a, T, S> { /// Returns a read-only view for the current execution context. pub fn get(&self) -> Option> { // SAFETY: `self.ptr` points to `self.len` valid elements for lifetime `'a`. - unsafe { S::get(self.ptr as *const T, self.len) } + unsafe { S::get(self.ptr, self.len) } } /// Returns a mutable view for the current execution context. @@ -217,9 +278,52 @@ impl<'a, T, S: PartitioningStrategy> Region<'a, T, S> { unsafe { S::get_mut(self.ptr, self.len) } } + /// Checks that this region can be launched with the given configuration. + /// + /// Called automatically by the [`offload!`] macro before launching a kernel. + pub fn check_launch(&self, grid: [u32; 3], block: [u32; 3]) -> Result<(), LaunchError> { + S::check_launch(self.len, grid, block) + } + /// Reborrows the region, producing a new region that aliases the same memory with /// the lifetime of the borrow. pub fn reborrow(&mut self) -> Region<'_, T, S> { Region { ptr: self.ptr, len: self.len, _marker: core::marker::PhantomData } } } + +// Launch validation helpers. + +/// Implementation detail of [`offload!`]: validates a [`Region`] argument against +/// the launch configuration of an offload call. +#[doc(hidden)] +#[unstable(feature = "offload", issue = "131513")] +pub trait RegionLaunchCheck { + /// Panics if the region's partitioning strategy rejects the launch. + fn __offload_check_launch(&self, grid: [u32; 3], block: [u32; 3]); +} + +#[unstable(feature = "offload", issue = "131513")] +impl<'a, T, S: PartitioningStrategy> RegionLaunchCheck for Region<'a, T, S> { + fn __offload_check_launch(&self, grid: [u32; 3], block: [u32; 3]) { + if let Err(err) = self.check_launch(grid, block) { + panic!( + "offload launch is not supported by the region's partitioning strategy: {}", + err, + ); + } + } +} + +/// no-op launch check for kernel arguments that are not [`Region`]s. +#[doc(hidden)] +#[unstable(feature = "offload", issue = "131513")] +pub trait DefaultLaunchCheck { + /// Does nothing. + fn __offload_check_launch(&self, grid: [u32; 3], block: [u32; 3]); +} + +#[unstable(feature = "offload", issue = "131513")] +impl DefaultLaunchCheck for &T { + fn __offload_check_launch(&self, _grid: [u32; 3], _block: [u32; 3]) {} +} diff --git a/tests/codegen-llvm/gpu_offload/auxiliary/offload_strategies.rs b/tests/codegen-llvm/gpu_offload/auxiliary/offload_strategies.rs index 1e3271e5a8547..68cfe40fec6c6 100644 --- a/tests/codegen-llvm/gpu_offload/auxiliary/offload_strategies.rs +++ b/tests/codegen-llvm/gpu_offload/auxiliary/offload_strategies.rs @@ -3,7 +3,8 @@ #![feature(gpu_offload)] #![feature(offload)] -use core::offload::PartitioningStrategy; +use core::offload::{LaunchError, PartitioningStrategy}; +use core::ptr::NonNull; #[derive(Debug, Clone, Copy)] pub struct Dummy; @@ -16,11 +17,43 @@ unsafe impl PartitioningStrategy for Dummy { 0 } - unsafe fn get<'a, T>(_ptr: *const T, _len: usize) -> Option> { + unsafe fn get<'a, T>(_ptr: NonNull, _len: usize) -> Option> { None } - unsafe fn get_mut<'a, T>(_ptr: *mut T, _len: usize) -> Option> { + unsafe fn get_mut<'a, T>(_ptr: NonNull, _len: usize) -> Option> { None } + + fn check_launch(_len: usize, _grid: [u32; 3], _block: [u32; 3]) -> Result<(), LaunchError> { + Ok(()) + } +} + +#[derive(Debug, Clone, Copy)] +pub struct Linear1D; + +unsafe impl PartitioningStrategy for Linear1D { + type View<'a, T: 'a> = &'a T; + type ViewMut<'a, T: 'a> = &'a mut T; + + fn index() -> usize { + 0 + } + + unsafe fn get<'a, T>(_ptr: NonNull, _len: usize) -> Option> { + None + } + + unsafe fn get_mut<'a, T>(_ptr: NonNull, _len: usize) -> Option> { + None + } + + fn check_launch(_len: usize, grid: [u32; 3], block: [u32; 3]) -> Result<(), LaunchError> { + if grid[1] == 1 && grid[2] == 1 && block[1] == 1 && block[2] == 1 { + Ok(()) + } else { + Err(LaunchError::new("")) + } + } } diff --git a/tests/ui/offload/auxiliary/offload_strategies.rs b/tests/ui/offload/auxiliary/offload_strategies.rs index 1e3271e5a8547..68cfe40fec6c6 100644 --- a/tests/ui/offload/auxiliary/offload_strategies.rs +++ b/tests/ui/offload/auxiliary/offload_strategies.rs @@ -3,7 +3,8 @@ #![feature(gpu_offload)] #![feature(offload)] -use core::offload::PartitioningStrategy; +use core::offload::{LaunchError, PartitioningStrategy}; +use core::ptr::NonNull; #[derive(Debug, Clone, Copy)] pub struct Dummy; @@ -16,11 +17,43 @@ unsafe impl PartitioningStrategy for Dummy { 0 } - unsafe fn get<'a, T>(_ptr: *const T, _len: usize) -> Option> { + unsafe fn get<'a, T>(_ptr: NonNull, _len: usize) -> Option> { None } - unsafe fn get_mut<'a, T>(_ptr: *mut T, _len: usize) -> Option> { + unsafe fn get_mut<'a, T>(_ptr: NonNull, _len: usize) -> Option> { None } + + fn check_launch(_len: usize, _grid: [u32; 3], _block: [u32; 3]) -> Result<(), LaunchError> { + Ok(()) + } +} + +#[derive(Debug, Clone, Copy)] +pub struct Linear1D; + +unsafe impl PartitioningStrategy for Linear1D { + type View<'a, T: 'a> = &'a T; + type ViewMut<'a, T: 'a> = &'a mut T; + + fn index() -> usize { + 0 + } + + unsafe fn get<'a, T>(_ptr: NonNull, _len: usize) -> Option> { + None + } + + unsafe fn get_mut<'a, T>(_ptr: NonNull, _len: usize) -> Option> { + None + } + + fn check_launch(_len: usize, grid: [u32; 3], block: [u32; 3]) -> Result<(), LaunchError> { + if grid[1] == 1 && grid[2] == 1 && block[1] == 1 && block[2] == 1 { + Ok(()) + } else { + Err(LaunchError::new("")) + } + } } diff --git a/tests/ui/offload/region_by_ref.rs b/tests/ui/offload/region_by_ref.rs index 8bdc92815525f..3328cc9493c89 100644 --- a/tests/ui/offload/region_by_ref.rs +++ b/tests/ui/offload/region_by_ref.rs @@ -2,7 +2,8 @@ //@ edition: 2024 //@ aux-crate: offload_strategies=offload_strategies.rs -// This tests ensures an error is emmited with passing a `&Region<'_, _, _>` args to offload. +// This test ensures an error is emitted when a `Region` is nested inside another +// type instead of being passed by value. #![feature(core_intrinsics)] #![feature(gpu_offload)] @@ -15,11 +16,13 @@ fn kernel_shared(_region: &Region<'_, f32, Dummy>) {} fn kernel_mut(_region: &mut Region<'_, f32, Dummy>) {} +fn kernel_nested(_arg: (u32, Region<'_, f32, Dummy>)) {} + fn main() { let mut x = [0.0f32; 4]; let region = Region::::new(&mut x[..]); core::intrinsics::offload::<_, _, ()>( - //~^ ERROR offload kernel argument 0 is a reference to a `Region` + //~^ ERROR offload kernel argument 0 contains a `Region` nested inside another type kernel_shared, [1, 1, 1], [1, 1, 1], @@ -31,7 +34,7 @@ fn main() { let mut y = [0.0f32; 4]; let mut region = Region::::new(&mut y[..]); core::intrinsics::offload::<_, _, ()>( - //~^ ERROR offload kernel argument 0 is a reference to a `Region` + //~^ ERROR offload kernel argument 0 contains a `Region` nested inside another type kernel_mut, [1, 1, 1], [1, 1, 1], @@ -39,4 +42,16 @@ fn main() { -1, (&mut region,), ); + + let mut z = [0.0f32; 4]; + let region = Region::::new(&mut z[..]); + core::intrinsics::offload::<_, _, ()>( + //~^ ERROR offload kernel argument 0 contains a `Region` nested inside another type + kernel_nested, + [1, 1, 1], + [1, 1, 1], + 0, + -1, + ((0u32, region),), + ); } diff --git a/tests/ui/offload/region_by_ref.stderr b/tests/ui/offload/region_by_ref.stderr index 25c278983ea38..8653d0d55b156 100644 --- a/tests/ui/offload/region_by_ref.stderr +++ b/tests/ui/offload/region_by_ref.stderr @@ -1,14 +1,20 @@ -error: offload kernel argument 0 is a reference to a `Region`. Pass the `Region` by value so it can be mapped like a slice - --> $DIR/region_by_ref.rs:21:5 +error: offload kernel argument 0 contains a `Region` nested inside another type. Pass the `Region` by value so it can be mapped like a slice + --> $DIR/region_by_ref.rs:24:5 | LL | core::intrinsics::offload::<_, _, ()>( | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -error: offload kernel argument 0 is a reference to a `Region`. Pass the `Region` by value so it can be mapped like a slice - --> $DIR/region_by_ref.rs:33:5 +error: offload kernel argument 0 contains a `Region` nested inside another type. Pass the `Region` by value so it can be mapped like a slice + --> $DIR/region_by_ref.rs:36:5 | LL | core::intrinsics::offload::<_, _, ()>( | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -error: aborting due to 2 previous errors +error: offload kernel argument 0 contains a `Region` nested inside another type. Pass the `Region` by value so it can be mapped like a slice + --> $DIR/region_by_ref.rs:48:5 + | +LL | core::intrinsics::offload::<_, _, ()>( + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +error: aborting due to 3 previous errors diff --git a/tests/ui/offload/region_launch_bounds.rs b/tests/ui/offload/region_launch_bounds.rs new file mode 100644 index 0000000000000..8b9f1d70ca7aa --- /dev/null +++ b/tests/ui/offload/region_launch_bounds.rs @@ -0,0 +1,19 @@ +//@ run-fail +//@ edition: 2024 +//@ aux-crate: offload_strategies=offload_strategies.rs +//@ error-pattern: offload launch is not supported by the region's partitioning strategy + +#![feature(gpu_offload)] +#![feature(offload)] + +use core::offload::{Region, RegionLaunchCheck}; +use offload_strategies::Linear1D; + +fn main() { + let mut x = [0.0f32; 4]; + let region = Region::::new(&mut x[..]); + + // `Linear1D` only supports single-dimensional launches, so `[4, 2, 1]` is + // rejected. + region.__offload_check_launch([4, 2, 1], [1, 1, 1]); +}