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 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>(
     tcx: TyCtxt<'tcx>,
     typing_env: ty::TypingEnv<'tcx>,
@@ -206,7 +211,23 @@ 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 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} contains a `Region` nested inside another \
+                        type. 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,25 @@ 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);
+            // `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 {
             BackendRepr::ScalarPair { a: _, b: _, b_offset: _ } => (0..2)
                 .map(|i| {
@@ -87,6 +107,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

@@ -1489,6 +1489,7 @@ symbols! {
         offload,
         offload_get_num_devices,
         offload_kernel,
+        offload_region,
         offset,
         offset_of,
         offset_of_enum,

tests/codegen-llvm/gpu_offload/auxiliary/offload_strategies.rs

@@ -0,0 +1,59 @@
+//@ edition: 2024
+
+#![feature(gpu_offload)]
+#![feature(offload)]
+
+use core::offload::{LaunchError, PartitioningStrategy};
+use core::ptr::NonNull;
+
+#[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: NonNull<T>, _len: usize) -> Option<Self::View<'a, T>> {
+        None
+    }
+
+    unsafe fn get_mut<'a, T>(_ptr: NonNull<T>, _len: usize) -> Option<Self::ViewMut<'a, T>> {
+        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<T>, _len: usize) -> Option<Self::View<'a, T>> {
+        None
+    }
+
+    unsafe fn get_mut<'a, T>(_ptr: NonNull<T>, _len: usize) -> Option<Self::ViewMut<'a, T>> {
+        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(""))
+        }
+    }
+}

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,59 @@
+//@ edition: 2024
+
+#![feature(gpu_offload)]
+#![feature(offload)]
+
+use core::offload::{LaunchError, PartitioningStrategy};
+use core::ptr::NonNull;
+
+#[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: NonNull<T>, _len: usize) -> Option<Self::View<'a, T>> {
+        None
+    }
+
+    unsafe fn get_mut<'a, T>(_ptr: NonNull<T>, _len: usize) -> Option<Self::ViewMut<'a, T>> {
+        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<T>, _len: usize) -> Option<Self::View<'a, T>> {
+        None
+    }
+
+    unsafe fn get_mut<'a, T>(_ptr: NonNull<T>, _len: usize) -> Option<Self::ViewMut<'a, T>> {
+        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(""))
+        }
+    }
+}

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,57 @@
+//@ compile-flags: -Zunstable-options -Zoffload=Device -Clto=fat
+//@ edition: 2024
+//@ aux-crate: offload_strategies=offload_strategies.rs
+
+// 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)]
+#![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 kernel_nested(_arg: (u32, 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 contains a `Region` nested inside another type
+        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 contains a `Region` nested inside another type
+        kernel_mut,
+        [1, 1, 1],
+        [1, 1, 1],
+        0,
+        -1,
+        (&mut region,),
+    );
+
+    let mut z = [0.0f32; 4];
+    let region = Region::<f32, Dummy>::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),),
+    );
+}

tests/ui/offload/region_by_ref.stderr

@@ -0,0 +1,20 @@
+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 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: 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
+