Adding a new, workign offload_args intrinsic, which only maps arguments, but calls some host code with device ptrs

Author: ZuseZ4

compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs

@@ -8,7 +8,7 @@ use rustc_middle::bug;
 use rustc_middle::ty::offload_meta::OffloadMetadata;
 
 use crate::builder::Builder;
-use crate::common::CodegenCx;
+use crate::common::{AsCCharPtr, CodegenCx};
 use crate::llvm::AttributePlace::Function;
 use crate::llvm::{self, Linkage, Type, Value};
 use crate::{SimpleCx, attributes};
@@ -288,7 +288,7 @@ pub(crate) struct OffloadKernelGlobals<'ll> {
     pub offload_sizes: &'ll llvm::Value,
     pub memtransfer_types: &'ll llvm::Value,
     pub region_id: &'ll llvm::Value,
-    pub offload_entry: &'ll llvm::Value,
+    pub offload_entry: Option<&'ll llvm::Value>,
 }
 
 fn gen_tgt_data_mappers<'ll>(
@@ -359,16 +359,19 @@ pub(crate) fn gen_define_handling<'ll>(
     types: &[&'ll Type],
     symbol: String,
     offload_globals: &OffloadGlobals<'ll>,
+    host: bool,
 ) -> OffloadKernelGlobals<'ll> {
     if let Some(entry) = cx.offload_kernel_cache.borrow().get(&symbol) {
         return *entry;
     }
 
     let offload_entry_ty = offload_globals.offload_entry_ty;
 
+    let mut arg_iter = types.iter().zip(metadata);
+
     // It seems like non-pointer values are automatically mapped. So here, we focus on pointer (or
     // reference) types.
-    let ptr_meta = types.iter().zip(metadata).filter_map(|(&x, meta)| match cx.type_kind(x) {
+    let ptr_meta = arg_iter.filter_map(|(&x, meta)| match cx.type_kind(x) {
         rustc_codegen_ssa::common::TypeKind::Pointer => Some(meta),
         _ => None,
     });
@@ -393,29 +396,34 @@ pub(crate) fn gen_define_handling<'ll>(
     let initializer = cx.get_const_i8(0);
     let region_id = add_global(&cx, &name, initializer, WeakAnyLinkage);
 
-    let c_entry_name = CString::new(symbol.clone()).unwrap();
-    let c_val = c_entry_name.as_bytes_with_nul();
-    let offload_entry_name = format!(".offloading.entry_name.{symbol}");
-
-    let initializer = crate::common::bytes_in_context(cx.llcx, c_val);
-    let llglobal = add_unnamed_global(&cx, &offload_entry_name, initializer, InternalLinkage);
-    llvm::set_alignment(llglobal, Align::ONE);
-    llvm::set_section(llglobal, c".llvm.rodata.offloading");
-
-    let name = format!(".offloading.entry.{symbol}");
-
-    // See the __tgt_offload_entry documentation above.
-    let elems = TgtOffloadEntry::new(&cx, region_id, llglobal);
-
-    let initializer = crate::common::named_struct(offload_entry_ty, &elems);
-    let c_name = CString::new(name).unwrap();
-    let offload_entry = llvm::add_global(cx.llmod, offload_entry_ty, &c_name);
-    llvm::set_global_constant(offload_entry, true);
-    llvm::set_linkage(offload_entry, WeakAnyLinkage);
-    llvm::set_initializer(offload_entry, initializer);
-    llvm::set_alignment(offload_entry, Align::EIGHT);
-    let c_section_name = CString::new("llvm_offload_entries").unwrap();
-    llvm::set_section(offload_entry, &c_section_name);
+    let offload_entry = if !host {
+        let c_entry_name = CString::new(symbol.clone()).unwrap();
+        let c_val = c_entry_name.as_bytes_with_nul();
+        let offload_entry_name = format!(".offloading.entry_name.{symbol}");
+
+        let initializer = crate::common::bytes_in_context(cx.llcx, c_val);
+        let llglobal = add_unnamed_global(&cx, &offload_entry_name, initializer, InternalLinkage);
+        llvm::set_alignment(llglobal, Align::ONE);
+        llvm::set_section(llglobal, c".llvm.rodata.offloading");
+
+        let name = format!(".offloading.entry.{symbol}");
+
+        // See the __tgt_offload_entry documentation above.
+        let elems = TgtOffloadEntry::new(&cx, region_id, llglobal);
+
+        let initializer = crate::common::named_struct(offload_entry_ty, &elems);
+        let c_name = CString::new(name).unwrap();
+        let offload_entry = llvm::add_global(cx.llmod, offload_entry_ty, &c_name);
+        llvm::set_global_constant(offload_entry, true);
+        llvm::set_linkage(offload_entry, WeakAnyLinkage);
+        llvm::set_initializer(offload_entry, initializer);
+        llvm::set_alignment(offload_entry, Align::EIGHT);
+        let c_section_name = CString::new("llvm_offload_entries").unwrap();
+        llvm::set_section(offload_entry, &c_section_name);
+        Some(offload_entry)
+    } else {
+        None
+    };
 
     let result =
         OffloadKernelGlobals { offload_sizes, memtransfer_types, region_id, offload_entry };
@@ -467,7 +475,12 @@ pub(crate) fn gen_call_handling<'ll, 'tcx>(
     metadata: &[OffloadMetadata],
     offload_globals: &OffloadGlobals<'ll>,
     offload_dims: &OffloadKernelDims<'ll>,
+    host: bool,
+    host_fnc_name: String,
+    host_llfn: &'ll Value,
+    host_llty: &'ll Type,
 ) {
+    dbg!(&host_llfn);
     let cx = builder.cx;
     let OffloadKernelGlobals { offload_sizes, offload_entry, memtransfer_types, region_id } =
         offload_data;
@@ -490,7 +503,8 @@ pub(crate) fn gen_call_handling<'ll, 'tcx>(
 
     // FIXME(Sa4dUs): dummy loads are a temp workaround, we should find a proper way to prevent these
     // variables from being optimized away
-    for val in [offload_sizes, offload_entry] {
+    for val in [offload_sizes] {
+        //for val in [offload_sizes, offload_entry] {
         unsafe {
             let dummy = llvm::LLVMBuildLoad2(
                 &builder.llbuilder,
@@ -623,14 +637,59 @@ pub(crate) fn gen_call_handling<'ll, 'tcx>(
     let args = vec![
         s_ident_t,
         // FIXME(offload) give users a way to select which GPU to use.
+        //cx.get_const_i64(0), // MAX == -1.
         cx.get_const_i64(u64::MAX), // MAX == -1.
         num_workgroups,
         threads_per_block,
         region_id,
         a5,
     ];
-    builder.call(tgt_target_kernel_ty, None, None, tgt_decl, &args, None, None);
-    // %41 = call i32 @__tgt_target_kernel(ptr @1, i64 -1, i32 2097152, i32 256, ptr @.kernel_1.region_id, ptr %kernel_args)
+    if host {
+        //let new_args = vec![];
+        //dbg!(&args);// were overwritten above, now hold omp stuff (shadowed)
+        dbg!(&vals);
+        dbg!(&types);
+        //dbg!(&fn_ty);
+        //let host_llty = cx.llvm_ty(host_llfn);
+        dbg!(&host_fnc_name);
+        dbg!(cx.get_defined_value(&host_fnc_name));
+        //let user_fn = unsafe {
+        //    llvm::LLVMRustGetOrInsertFunction(
+        //        builder.llmod,
+        //        host_fnc_name.as_c_char_ptr(),
+        //        host_fnc_name.len(),
+        //        fn_ty,
+        //    )
+        //};
+        //dbg!(&user_fn);
+        // void *omp_get_mapped_ptr(void *ptr, int device_num);
+
+        let fn_name = "omp_get_mapped_ptr";
+        let ty2: &'ll Type = cx.type_func(&[cx.type_ptr(), cx.type_i32()], cx.type_ptr());
+        let mapper_fn = unsafe {
+            llvm::LLVMRustGetOrInsertFunction(
+                builder.llmod,
+                fn_name.as_c_char_ptr(),
+                fn_name.len(),
+                ty2,
+            )
+        };
+
+        let mut device_vals = Vec::with_capacity(vals.len());
+        let device_num = cx.get_const_i32(0);
+        for arg in vals {
+            dbg!(&mapper_fn);
+            dbg!(&ty2);
+            let device_arg =
+                builder.call(ty2, None, None, mapper_fn, &[arg, device_num], None, None);
+            device_vals.push(device_arg);
+        }
+        builder.call(host_llty, None, None, host_llfn, &device_vals, None, None);
+        //builder.call(fn_ty, None, None, host_llfn, &vals, None, None);
+    } else {
+        builder.call(tgt_target_kernel_ty, None, None, tgt_decl, &args, None, None);
+        // %41 = call i32 @__tgt_target_kernel(ptr @1, i64 -1, i32 2097152, i32 256, ptr @.kernel_1.region_id, ptr %kernel_args)
+    }
 
     // Step 4)
     let geps = get_geps(builder, ty, ty2, a1, a2, a4);

compiler/rustc_codegen_llvm/src/intrinsic.rs

@@ -215,7 +215,19 @@ impl<'ll, 'tcx> IntrinsicCallBuilderMethods<'tcx> for Builder<'_, 'll, 'tcx> {
                     let _ = tcx.dcx().emit_almost_fatal(OffloadWithoutFatLTO);
                 }
 
-                codegen_offload(self, tcx, instance, args);
+                codegen_offload(self, tcx, instance, args, false);
+                return Ok(());
+            }
+            sym::offload_args => {
+                if tcx.sess.opts.unstable_opts.offload.is_empty() {
+                    let _ = tcx.dcx().emit_almost_fatal(OffloadWithoutEnable);
+                }
+
+                if tcx.sess.lto() != rustc_session::config::Lto::Fat {
+                    let _ = tcx.dcx().emit_almost_fatal(OffloadWithoutFatLTO);
+                }
+
+                codegen_offload(self, tcx, instance, args, true);
                 return Ok(());
             }
             sym::is_val_statically_known => {
@@ -1362,8 +1374,11 @@ fn codegen_offload<'ll, 'tcx>(
     tcx: TyCtxt<'tcx>,
     instance: ty::Instance<'tcx>,
     args: &[OperandRef<'tcx, &'ll Value>],
+    host: bool,
 ) {
+    dbg!(&args[0]);
     let cx = bx.cx;
+    dbg!(&instance);
     let fn_args = instance.args;
 
     let (target_id, target_args) = match fn_args.into_type_list(tcx)[0].kind() {
@@ -1383,13 +1398,34 @@ fn codegen_offload<'ll, 'tcx>(
             return;
         }
     };
+    let llfn = cx.get_fn(fn_target);
+    dbg!(&llfn);
+    dbg!(&fn_target);
 
     let offload_dims = OffloadKernelDims::from_operands(bx, &args[1], &args[2]);
     let args = get_args_from_tuple(bx, args[3], fn_target);
+    dbg!(&args);
     let target_symbol = symbol_name_for_instance_in_crate(tcx, fn_target, LOCAL_CRATE);
+    dbg!(&target_symbol);
 
     let sig = tcx.fn_sig(fn_target.def_id()).skip_binder().skip_binder();
+    dbg!(&sig);
     let inputs = sig.inputs();
+    // [compiler/rustc_codegen_llvm/src/intrinsic.rs:1405:5] &args = [
+    //    (ptr:  %36 = load ptr, ptr %_14, align 8),
+    //    (ptr:  %38 = load ptr, ptr %37, align 8),
+    //    (ptr:  %40 = load ptr, ptr %39, align 8),
+    //    (ptr:  %42 = load ptr, ptr %41, align 8),
+    //    (ptr:  %44 = load ptr, ptr %43, align 8),
+    //    (ptr:  %46 = load ptr, ptr %45, align 8),
+    //    (ptr:  %48 = load ptr, ptr %47, align 8),
+    //    (ptr:  %50 = load ptr, ptr %49, align 8),
+    //    (ptr:  %52 = load ptr, ptr %51, align 8),
+    //    (ptr:  %54 = load ptr, ptr %53, align 8),
+    //]
+    //[compiler/rustc_codegen_llvm/src/intrinsic.rs:1407:5] &target_symbol = "rocblas_sgemv_wrapper"
+    //[compiler/rustc_codegen_llvm/src/intrinsic.rs:1410:5] &sig = fn(*const i32, *const i32, *const f32, *const f32, *const i32, *const f32, *const i32, *const f32, *mut f32, *const i32)
+    //)
 
     let metadata = inputs.iter().map(|ty| OffloadMetadata::from_ty(tcx, *ty)).collect::<Vec<_>>();
 
@@ -1403,8 +1439,26 @@ fn codegen_offload<'ll, 'tcx>(
             return;
         }
     };
-    let offload_data = gen_define_handling(&cx, &metadata, &types, target_symbol, offload_globals);
-    gen_call_handling(bx, &offload_data, &args, &types, &metadata, offload_globals, &offload_dims);
+    let instance = rustc_middle::ty::Instance::mono(tcx, fn_target.def_id());
+    let fn_abi = cx.fn_abi_of_instance(instance, tcx.mk_type_list(&[]));
+    let host_fn_ty = fn_abi.llvm_type(cx);
+    dbg!(&host_fn_ty);
+
+    let offload_data =
+        gen_define_handling(&cx, &metadata, &types, target_symbol.clone(), offload_globals, host);
+    gen_call_handling(
+        bx,
+        &offload_data,
+        &args,
+        &types,
+        &metadata,
+        offload_globals,
+        &offload_dims,
+        host,
+        target_symbol,
+        llfn,
+        host_fn_ty,
+    );
 }
 
 fn get_args_from_tuple<'ll, 'tcx>(

compiler/rustc_hir_analysis/src/check/intrinsic.rs

@@ -164,6 +164,7 @@ fn intrinsic_operation_unsafety(tcx: TyCtxt<'_>, intrinsic_id: LocalDefId) -> hi
         | sym::mul_with_overflow
         | sym::needs_drop
         | sym::offload
+        | sym::offload_args
         | sym::offset_of
         | sym::overflow_checks
         | sym::powf16
@@ -326,6 +327,17 @@ pub(crate) fn check_intrinsic_type(
             ],
             param(2),
         ),
+        sym::offload_args => (
+            3,
+            0,
+            vec![
+                param(0),
+                Ty::new_array_with_const_len(tcx, tcx.types.u32, Const::from_target_usize(tcx, 3)),
+                Ty::new_array_with_const_len(tcx, tcx.types.u32, Const::from_target_usize(tcx, 3)),
+                param(1),
+            ],
+            param(2),
+        ),
         sym::offset => (2, 0, vec![param(0), param(1)], param(0)),
         sym::arith_offset => (
             1,

compiler/rustc_span/src/symbol.rs

@@ -1596,6 +1596,7 @@ symbols! {
         of,
         off,
         offload,
+        offload_args,
         offset,
         offset_of,
         offset_of_enum,

library/core/src/intrinsics/mod.rs

@@ -3421,6 +3421,15 @@ pub const fn offload<F, T: crate::marker::Tuple, R>(
     args: T,
 ) -> R;
 
+#[rustc_nounwind]
+#[rustc_intrinsic]
+pub const fn offload_args<F, T: crate::marker::Tuple, R>(
+    f: F,
+    workgroup_dim: [u32; 3],
+    thread_dim: [u32; 3],
+    args: T,
+) -> R;
+
 /// Inform Miri that a given pointer definitely has a certain alignment.
 #[cfg(miri)]
 #[rustc_allow_const_fn_unstable(const_eval_select)]