Offload safe mutable args with Region and PartitioningStrategy

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()

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<Ty<'tcx>> {
+    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() {

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,

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<Self::View<'a, T>>;
+
+    /// 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<Self::ViewMut<'a, T>>;
+}
+
+/// 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<D>(data: D) -> Self
+    where
+        D: Into<RawRegion<'a, T>>,
+    {
+        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<S::View<'_, T>> {
+        // 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<S::ViewMut<'_, T>> {
+        // 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 }
+    }
+}

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<Self::View<'a, T>> {
+        None
+    }
+
+    unsafe fn get_mut<'a, T>(_ptr: *mut T, _len: usize) -> Option<Self::ViewMut<'a, T>> {
+        None
+    }
+}

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::<f32, Dummy>::new(&mut x as &mut [f32]),),
+    };
+}
+
+fn foo(region: Region<'_, f32, Dummy>) {
+    unreachable!();
+}

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<Self::View<'a, T>> {
+        None
+    }
+
+    unsafe fn get_mut<'a, T>(_ptr: *mut T, _len: usize) -> Option<Self::ViewMut<'a, T>> {
+        None
+    }
+}

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::<f32, Dummy>::new(&mut x[..]);
+
+    x[0] = 1.0;
+    //~^ ERROR cannot assign to `x[_]` because it is borrowed
+
+    let _view = region.get();
+}

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::<f32, Dummy>::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`.

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::<f32, Dummy>::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,
+        (&region,),
+    );
+
+    let mut y = [0.0f32; 4];
+    let mut region = Region::<f32, Dummy>::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,),
+    );
+}

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
+