Offload safe mutable args with Region and PartitioningStrategy

compiler/rustc_hir_typeck/src/intrinsicck.rs

@@ -1,12 +1,14 @@
 use hir::HirId;
 use rustc_abi::Primitive::Pointer;
 use rustc_abi::VariantIdx;
+use rustc_data_structures::fx::FxHashSet;
 use rustc_errors::codes::*;
 use rustc_errors::struct_span_code_err;
 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 +137,28 @@ fn check_transmute<'tcx>(
     }
 }
 
+fn contains_nested_offload_region<'tcx>(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>) -> bool {
+    let mut visited = FxHashSet::default();
+    let mut stack = vec![ty];
+    let mut is_root = true;
+
+    while let Some(ty) = stack.pop() {
+        if !is_root && is_region_ty(tcx, ty) {
+            return true;
+        }
+        is_root = false;
+
+        if let ty::Adt(def, args) = *ty.kind()
+            && visited.insert(def.did())
+        {
+            stack.extend(def.all_fields().map(|field| field.ty(tcx, args).skip_norm_wip()));
+        }
+        stack.extend(ty.walk().skip(1).filter_map(|arg| arg.as_type()));
+    }
+
+    false
+}
+
 fn check_offload<'tcx>(
     tcx: TyCtxt<'tcx>,
     typing_env: ty::TypingEnv<'tcx>,
@@ -206,7 +230,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,