Set maptypes using offload metadata

Author: Sa4dUs

compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs

@@ -270,19 +270,17 @@ pub(crate) fn gen_define_handling<'ll, 'tcx>(
     let types = cx.func_params_types(cx.get_type_of_global(kernel));
     // It seems like non-pointer values are automatically mapped. So here, we focus on pointer (or
     // reference) types.
-    let num_ptr_types = types
-        .iter()
-        .filter(|&x| matches!(cx.type_kind(x), rustc_codegen_ssa::common::TypeKind::Pointer))
-        .count();
-
-    let ptr_sizes = types
+    let ptr_meta = types
         .iter()
         .zip(metadata)
         .filter_map(|(&x, meta)| match cx.type_kind(x) {
-            rustc_codegen_ssa::common::TypeKind::Pointer => Some(meta.payload_size),
+            rustc_codegen_ssa::common::TypeKind::Pointer => Some(meta),
             _ => None,
         })
-        .collect::<Vec<u64>>();
+        .collect::<Vec<OffloadMetadata>>();
+
+    let ptr_sizes = ptr_meta.iter().map(|m| m.payload_size).collect::<Vec<_>>();
+    let ptr_transfer = ptr_meta.iter().map(|m| m.mode as u64 | 0x20).collect::<Vec<_>>();
 
     // We do not know their size anymore at this level, so hardcode a placeholder.
     // A follow-up pr will track these from the frontend, where we still have Rust types.
@@ -294,11 +292,8 @@ pub(crate) fn gen_define_handling<'ll, 'tcx>(
     // A non-mutable reference or pointer will be 1, an array that's not read, but fully overwritten
     // will be 2. For now, everything is 3, until we have our frontend set up.
     // 1+2+32: 1 (MapTo), 2 (MapFrom), 32 (Add one extra input ptr per function, to be used later).
-    let memtransfer_types = add_priv_unnamed_arr(
-        &cx,
-        &format!(".offload_maptypes.{symbol}"),
-        &vec![1 + 2 + 32; num_ptr_types],
-    );
+    let memtransfer_types =
+        add_priv_unnamed_arr(&cx, &format!(".offload_maptypes.{symbol}"), &ptr_transfer);
 
     // Next: For each function, generate these three entries. A weak constant,
     // the llvm.rodata entry name, and  the llvm_offload_entries value

compiler/rustc_middle/src/ty/offload_meta.rs

@@ -6,10 +6,13 @@ pub struct OffloadMetadata {
     pub mode: TransferKind,
 }
 
+// TODO(Sa4dUs): add `OMP_MAP_TARGET_PARAM = 0x20` flag only when needed
+#[repr(u64)]
+#[derive(Debug, Copy, Clone)]
 pub enum TransferKind {
     FromGpu = 1,
     ToGpu = 2,
-    Both = 3,
+    Both = 1 + 2,
 }
 
 impl OffloadMetadata {
@@ -18,7 +21,10 @@ impl OffloadMetadata {
     }
 
     pub fn from_ty<'tcx>(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>) -> Self {
-        OffloadMetadata { payload_size: get_payload_size(tcx, ty), mode: TransferKind::Both }
+        OffloadMetadata {
+            payload_size: get_payload_size(tcx, ty),
+            mode: TransferKind::from_ty(tcx, ty),
+        }
     }
 }
 
@@ -68,3 +74,49 @@ fn get_payload_size<'tcx>(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>) -> u64 {
             .bytes(),
     }
 }
+
+impl TransferKind {
+    pub fn from_ty<'tcx>(_tcx: TyCtxt<'tcx>, ty: Ty<'tcx>) -> Self {
+        // TODO(Sa4dUs): this logic is probs not fully correct, but it works for now
+        match ty.kind() {
+            rustc_type_ir::TyKind::Bool
+            | rustc_type_ir::TyKind::Char
+            | rustc_type_ir::TyKind::Int(_)
+            | rustc_type_ir::TyKind::Uint(_)
+            | rustc_type_ir::TyKind::Float(_) => TransferKind::ToGpu,
+
+            rustc_type_ir::TyKind::Adt(_, _)
+            | rustc_type_ir::TyKind::Tuple(_)
+            | rustc_type_ir::TyKind::Array(_, _) => TransferKind::ToGpu,
+
+            rustc_type_ir::TyKind::RawPtr(_, rustc_ast::Mutability::Not)
+            | rustc_type_ir::TyKind::Ref(_, _, rustc_ast::Mutability::Not) => TransferKind::ToGpu,
+
+            rustc_type_ir::TyKind::RawPtr(_, rustc_ast::Mutability::Mut)
+            | rustc_type_ir::TyKind::Ref(_, _, rustc_ast::Mutability::Mut) => TransferKind::Both,
+
+            rustc_type_ir::TyKind::Slice(_)
+            | rustc_type_ir::TyKind::Str
+            | rustc_type_ir::TyKind::Dynamic(_, _) => TransferKind::Both,
+
+            rustc_type_ir::TyKind::FnDef(_, _)
+            | rustc_type_ir::TyKind::FnPtr(_, _)
+            | rustc_type_ir::TyKind::Closure(_, _)
+            | rustc_type_ir::TyKind::CoroutineClosure(_, _)
+            | rustc_type_ir::TyKind::Coroutine(_, _)
+            | rustc_type_ir::TyKind::CoroutineWitness(_, _) => TransferKind::ToGpu,
+
+            rustc_type_ir::TyKind::Alias(_, _)
+            | rustc_type_ir::TyKind::Param(_)
+            | rustc_type_ir::TyKind::Bound(_, _)
+            | rustc_type_ir::TyKind::Placeholder(_)
+            | rustc_type_ir::TyKind::Infer(_)
+            | rustc_type_ir::TyKind::Error(_) => TransferKind::ToGpu,
+
+            rustc_type_ir::TyKind::Never => TransferKind::ToGpu,
+            rustc_type_ir::TyKind::Foreign(_) => TransferKind::Both,
+            rustc_type_ir::TyKind::Pat(_, _) => TransferKind::Both,
+            rustc_type_ir::TyKind::UnsafeBinder(_) => TransferKind::Both,
+        }
+    }
+}

tests/codegen-llvm/gpu_offload/offload_intrinsic.rs

@@ -1,4 +1,4 @@
-//@ compile-flags: -Zoffload=Enable -Zunstable-options -C opt-level=0  -Clto=fat
+//@ compile-flags: -Zoffload=Enable -Zunstable-options -C opt-level=3  -Clto=fat
 //@ no-prefer-dynamic
 //@ needs-enzyme