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 a58e517e05f61..c51220912eb26 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,22 @@ 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 is_region_ty(tcx, ty) { + let ty::Adt(_, args) = ty.kind() else { unreachable!() }; + 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 3fcb59ea2d701..edd13176987ba 100644 --- a/compiler/rustc_span/src/symbol.rs +++ b/compiler/rustc_span/src/symbol.rs @@ -1462,6 +1462,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 3d85621361209..33d69ceb37d60 100644 --- a/library/core/src/offload/mod.rs +++ b/library/core/src/offload/mod.rs @@ -125,3 +125,101 @@ 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 = "131513")] +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(Debug)] +#[unstable(feature = "offload", issue = "131513")] +#[rustc_diagnostic_item = "offload_region"] +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. +#[derive(Debug)] +#[unstable(feature = "offload", issue = "131513")] +pub struct RawRegion<'a, T> { + ptr: *mut T, + 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 = "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 + 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) } + } + + /// 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 new file mode 100644 index 0000000000000..975096c7fb5d4 --- /dev/null +++ b/tests/codegen-llvm/gpu_offload/region_host.rs @@ -0,0 +1,47 @@ +//@ 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)] +#![feature(core_intrinsics)] +#![feature(gpu_offload)] +#![feature(offload)] +#![feature(rustc_attrs)] +#![no_main] + +extern crate core; + +use core::offload::Region; + +use offload_strategies::Dummy; + +// 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::offload::offload! { + kernel = foo, + args = (Region::::new(&mut x as &mut [f32]),), + }; +} + +fn foo(region: Region<'_, f32, Dummy>) { + unreachable!(); +} 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`.