Offload safe mutable args with Region and PartitioningStrategy

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

compiler/rustc_span/src/symbol.rs

@@ -1469,6 +1469,7 @@ symbols! {
         off,
         offload,
         offload_kernel,
+        offload_region,
         offset,
         offset_of,
         offset_of_enum,

library/core/src/offload/mod.rs

@@ -101,3 +101,95 @@ macro_rules! offload {
     (@value (SOME $val:expr)) => { $val };
     (@value ($val:expr)) => { $val };
 }
+
+// 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<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(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,
+    _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 = "124509")]
+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 = "124509")]
+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) }
+    }
+}

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<Self::View<'a, T>> {
+        None
+    }
+
+    unsafe fn get_mut<'a, T>(_ptr: *mut T, _len: usize) -> Option<Self::ViewMut<'a, T>> {
+        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::<f32, Dummy>::new(&mut x as &mut [f32]),),
+    );
+}
+
+fn foo(region: Region<'_, f32, Dummy>) {
+    unreachable!();
+}