impl drop for PreloadMut, which must returned the mutated value

Author: ZuseZ4

compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs

@@ -446,6 +446,7 @@ pub(crate) fn gen_define_handling<'ll>(
     let valid_begin_mappings = MappingFlags::TO | MappingFlags::LITERAL | MappingFlags::IMPLICIT;
     let transfer_to: Vec<u64> =
         transfer.iter().map(|m| m.intersection(valid_begin_mappings).bits()).collect();
+    dbg!(&transfer);
     let transfer_from: Vec<u64> =
         transfer.iter().map(|m| m.intersection(MappingFlags::FROM).bits()).collect();
     let valid_kernel_mappings = MappingFlags::LITERAL | MappingFlags::IMPLICIT;
@@ -469,6 +470,7 @@ pub(crate) fn gen_define_handling<'ll>(
         add_priv_unnamed_arr(&cx, &format!(".offload_maptypes.{symbol}.begin"), &transfer_to);
     let memtransfer_kernel =
         add_priv_unnamed_arr(&cx, &format!(".offload_maptypes.{symbol}.kernel"), &transfer_kernel);
+    dbg!(&transfer_from);
     let memtransfer_end =
         add_priv_unnamed_arr(&cx, &format!(".offload_maptypes.{symbol}.end"), &transfer_from);
 

compiler/rustc_codegen_llvm/src/intrinsic.rs

@@ -19,7 +19,7 @@ use rustc_hir::def_id::LOCAL_CRATE;
 use rustc_hir::find_attr;
 use rustc_middle::mir::BinOp;
 use rustc_middle::ty::layout::{FnAbiOf, HasTyCtxt, HasTypingEnv, LayoutOf};
-use rustc_middle::ty::offload_meta::OffloadMetadata;
+use rustc_middle::ty::offload_meta::{MappingFlags, OffloadMetadata};
 use rustc_middle::ty::{self, GenericArgsRef, Instance, SimdAlign, Ty, TyCtxt, TypingEnv};
 use rustc_middle::{bug, span_bug};
 use rustc_session::config::CrateType;
@@ -188,6 +188,17 @@ impl<'ll, 'tcx> IntrinsicCallBuilderMethods<'tcx> for Builder<'_, 'll, 'tcx> {
         codegen_offload_preload(self, tcx, instance, args);
     }
 
+    fn codegen_offload_preload_mut_drop(
+        &mut self,
+        preload_ty: Ty<'tcx>,
+        place: PlaceRef<'tcx, &'ll llvm::Value>,
+    ) {
+        let tcx = self.tcx;
+        dbg!("Dropping PreloadMut; emit offload end mapper");
+
+        codegen_offload_preload_mut_drop(self, tcx, preload_ty, place);
+    }
+
     fn codegen_intrinsic_call(
         &mut self,
         instance: ty::Instance<'tcx>,
@@ -1911,6 +1922,83 @@ fn codegen_autodiff<'ll, 'tcx>(
     );
 }
 
+fn codegen_offload_preload_mut_drop<'ll, 'tcx>(
+    bx: &mut Builder<'_, 'll, 'tcx>,
+    tcx: TyCtxt<'tcx>,
+    preload_ty: Ty<'tcx>,
+    place: PlaceRef<'tcx, &'ll llvm::Value>,
+) {
+    let cx = bx.cx;
+    dbg!("Starting the PreloadMut drop handling!");
+    // PreloadMut<'a, T> -> extract T.
+    let ty::Adt(_adt_def, generic_args) = preload_ty.kind() else {
+        bug!("expected PreloadMut ADT, got {preload_ty:?}");
+    };
+
+    // This should be the `T` parameter of PreloadMut<'a, T>.
+    // If this indexes the lifetime in your tree, use the correct type arg index
+    // or `generic_args.types().next().unwrap()`.
+    let pointee_ty: Ty<'tcx> =
+        generic_args.types().next().unwrap_or_else(|| bug!("PreloadMut without type parameter"));
+
+    // Load field 0: `cpu_ptr: *mut T`.
+    let cpu_ptr_place = place.project_field(bx, 0);
+    dbg!(&cpu_ptr_place);
+    let cpu_ptr_operand = bx.load_operand(cpu_ptr_place);
+    dbg!(&cpu_ptr_operand);
+
+    let args: Vec<&'ll Value> = match cpu_ptr_operand.val {
+        OperandValue::Immediate(ptr) => vec![ptr],
+        OperandValue::Pair(_data, _meta) => {
+            bug!("unsized PreloadMut drop not handled yet")
+        }
+        _ => bug!("unexpected PreloadMut cpu_ptr operand"),
+    };
+
+    let mut meta = OffloadMetadata::from_ty(tcx, pointee_ty);
+    // We end a mut Mapper. Unless the user never mutated a mut variable passed in a mutable way, we
+    // must return it from the device to update the host version. If they never mutated it, they
+    // surely got a clippy or rustc warning, so it's up to them for wasting time.
+    meta.mode |= MappingFlags::FROM;
+    dbg!(&meta);
+    let metadata: &[OffloadMetadata; 1] = &[meta];
+
+    let types: &Type = cx.layout_of(pointee_ty).llvm_type(cx);
+
+    let offload_globals_ref = cx.offload_globals.borrow();
+    let offload_globals = match offload_globals_ref.as_ref() {
+        Some(globals) => globals,
+        None => {
+            dbg!("Have to initialize offload? This is a bug!");
+            return;
+        }
+    };
+
+    let target_symbol = cx.generate_local_symbol_name("");
+    dbg!("done for now");
+    let offload_data = gen_define_handling(&cx, metadata, target_symbol, offload_globals);
+    let has_dynamic = metadata.iter().any(|m| !matches!(m.payload_size, OffloadSize::Static(_)));
+    let (ty, ty2, a1, a2, a4) = crate::builder::gpu_helper::preper_datatransfers(
+        bx,
+        &args,
+        &[types],
+        offload_data.offload_sizes,
+        metadata,
+        has_dynamic,
+    );
+    let geps = crate::builder::gpu_helper::get_geps(bx, ty, ty2, a1, a2, a4, has_dynamic);
+
+    crate::builder::gpu_helper::generate_mapper_call(
+        bx,
+        geps,
+        offload_data.memtransfer_end,
+        offload_globals.end_mapper,
+        offload_globals.mapper_fn_ty,
+        1,
+        offload_globals.ident_t_global,
+    );
+}
+
 // For each PreLoad *call*, we now use some of our previous declared globals to move data to the gpu.
 // For now, we only handle the data transfer part of it. Consecutive calls become a no-op on the
 // LLVM side.

compiler/rustc_codegen_ssa/src/mir/block.rs

@@ -9,7 +9,7 @@ use rustc_lint_defs::builtin::TAIL_CALL_TRACK_CALLER;
 use rustc_middle::mir::{self, AssertKind, InlineAsmMacro, SwitchTargets, UnwindTerminateReason};
 use rustc_middle::ty::layout::{HasTyCtxt, LayoutOf, ValidityRequirement};
 use rustc_middle::ty::print::{with_no_trimmed_paths, with_no_visible_paths};
-use rustc_middle::ty::{self, Instance, Ty, TypeVisitableExt};
+use rustc_middle::ty::{self, Instance, Ty, TyCtxt, TypeVisitableExt};
 use rustc_middle::{bug, span_bug};
 use rustc_session::config::OptLevel;
 use rustc_span::{Span, Spanned};
@@ -27,6 +27,14 @@ use crate::mir::IntrinsicResult;
 use crate::traits::*;
 use crate::{MemFlags, meth};
 
+fn is_preload_mut_type<'tcx>(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>) -> bool {
+    let ty::Adt(adt_def, _) = ty.kind() else {
+        return false;
+    };
+
+    Some(adt_def.did()) == tcx.lang_items().preload_mut_type()
+}
+
 // Indicates if we are in the middle of merging a BB's successor into it. This
 // can happen when BB jumps directly to its successor and the successor has no
 // other predecessors.
@@ -604,6 +612,12 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
     ) -> MergingSucc {
         let ty = location.ty(self.mir, bx.tcx()).ty;
         let ty = self.monomorphize(ty);
+
+        if is_preload_mut_type(bx.tcx(), ty) {
+            let place = self.codegen_place(bx, location.as_ref());
+
+            bx.codegen_offload_preload_mut_drop(ty, place);
+        }
         let drop_fn = Instance::resolve_drop_glue(bx.tcx(), ty);
 
         if let ty::InstanceKind::DropGlue(_, None) = drop_fn.def {

compiler/rustc_codegen_ssa/src/traits/intrinsic.rs

@@ -5,7 +5,7 @@ use super::BackendTypes;
 use crate::RetagInfo;
 use crate::mir::IntrinsicResult;
 use crate::mir::operand::OperandRef;
-use crate::mir::place::PlaceValue;
+use crate::mir::place::{PlaceRef, PlaceValue};
 
 pub trait IntrinsicCallBuilderMethods<'tcx>: BackendTypes {
     /// Higher-level interface to emitting calls to intrinsics
@@ -38,6 +38,12 @@ pub trait IntrinsicCallBuilderMethods<'tcx>: BackendTypes {
         is_mut: bool,
     );
 
+    fn codegen_offload_preload_mut_drop(
+        &mut self,
+        preload_ty: ty::Ty<'tcx>,
+        place: PlaceRef<'tcx, Self::Value>,
+    );
+
     fn codegen_llvm_intrinsic_call(
         &mut self,
         instance: ty::Instance<'tcx>,

library/core/src/offload/mod.rs

@@ -31,3 +31,13 @@ pub fn preload<'a, T: ?Sized>(x: &'a T) -> Preload<'a, T> {
 pub fn preload_mut<'a, T: ?Sized>(x: &'a mut T) -> PreloadMut<'a, T> {
     PreloadMut { cpu_ptr: x as *mut T, _marker: PhantomData }
 }
+
+impl<T: ?Sized> Drop for PreloadMut<'_, T> {
+    fn drop(&mut self) {
+        // Intentionally empty.
+        //
+        // This exists so MIR creates Drop terminators for PreloadMut.
+        // rustc codegen intercepts those terminators and emits the
+        // offload return mapper.
+    }
+}