first definition of offload intrinsic (dirty code)

Author: Sa4dUs

compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs

@@ -4,17 +4,18 @@ use llvm::Linkage::*;
 use rustc_abi::Align;
 use rustc_codegen_ssa::back::write::CodegenContext;
 use rustc_codegen_ssa::traits::BaseTypeCodegenMethods;
+use rustc_middle::ty::{self, PseudoCanonicalInput, Ty, TyCtxt, TypingEnv};
 
 use crate::builder::SBuilder;
-use crate::common::AsCCharPtr;
 use crate::llvm::AttributePlace::Function;
-use crate::llvm::{self, Linkage, Type, Value};
+use crate::llvm::{self, BasicBlock, Linkage, Type, Value};
 use crate::{LlvmCodegenBackend, SimpleCx, attributes};
 
 pub(crate) fn handle_gpu_code<'ll>(
     _cgcx: &CodegenContext<LlvmCodegenBackend>,
-    cx: &'ll SimpleCx<'_>,
+    _cx: &'ll SimpleCx<'_>,
 ) {
+    /*
     // The offload memory transfer type for each kernel
     let mut o_types = vec![];
     let mut kernels = vec![];
@@ -28,6 +29,7 @@ pub(crate) fn handle_gpu_code<'ll>(
     }
 
     gen_call_handling(&cx, &kernels, &o_types);
+    */
 }
 
 // What is our @1 here? A magic global, used in our data_{begin/update/end}_mapper:
@@ -83,7 +85,7 @@ pub(crate) fn add_tgt_offload_entry<'ll>(cx: &'ll SimpleCx<'_>) -> &'ll llvm::Ty
     offload_entry_ty
 }
 
-fn gen_tgt_kernel_global<'ll>(cx: &'ll SimpleCx<'_>) {
+pub(crate) fn gen_tgt_kernel_global<'ll>(cx: &'ll SimpleCx<'_>) {
     let kernel_arguments_ty = cx.type_named_struct("struct.__tgt_kernel_arguments");
     let tptr = cx.type_ptr();
     let ti64 = cx.type_i64();
@@ -182,11 +184,14 @@ pub(crate) fn add_global<'ll>(
     llglobal
 }
 
-fn gen_define_handling<'ll>(
-    cx: &'ll SimpleCx<'_>,
+pub(crate) fn gen_define_handling<'ll, 'tcx>(
+    cx: &SimpleCx<'ll>,
+    tcx: TyCtxt<'tcx>,
     kernel: &'ll llvm::Value,
     offload_entry_ty: &'ll llvm::Type,
-    num: i64,
+    // TODO(Sa4dUs): Define a typetree once i have a better idea of what do we exactly need
+    tt: Vec<Ty<'tcx>>,
+    symbol: &str,
 ) -> &'ll llvm::Value {
     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
@@ -196,35 +201,47 @@ fn gen_define_handling<'ll>(
         .filter(|&x| matches!(cx.type_kind(x), rustc_codegen_ssa::common::TypeKind::Pointer))
         .count();
 
+    // TODO(Sa4dUs): Add typetrees here
+    let ptr_sizes = types
+        .iter()
+        .zip(tt)
+        .filter_map(|(&x, ty)| match cx.type_kind(x) {
+            rustc_codegen_ssa::common::TypeKind::Pointer => Some(get_payload_size(tcx, ty)),
+            _ => None,
+        })
+        .collect::<Vec<u64>>();
+
     // 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.
     // Then, we will be able to figure out that e.g. `&[f32;256]` will result in 4*256 bytes.
     // I decided that 1024 bytes is a great placeholder value for now.
-    add_priv_unnamed_arr(&cx, &format!(".offload_sizes.{num}"), &vec![1024; num_ptr_types]);
+    add_priv_unnamed_arr(&cx, &format!(".offload_sizes.{symbol}"), &ptr_sizes);
     // Here we figure out whether something needs to be copied to the gpu (=1), from the gpu (=2),
     // or both to and from the gpu (=3). Other values shouldn't affect us for now.
     // 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.
+
+    // TODO(Sa4dUs): Check the way to figure out this
     let o_types =
-        add_priv_unnamed_arr(&cx, &format!(".offload_maptypes.{num}"), &vec![3; num_ptr_types]);
+        add_priv_unnamed_arr(&cx, &format!(".offload_maptypes.{symbol}"), &vec![3; num_ptr_types]);
     // Next: For each function, generate these three entries. A weak constant,
     // the llvm.rodata entry name, and  the omp_offloading_entries value
 
-    let name = format!(".kernel_{num}.region_id");
+    let name = format!(".{symbol}.region_id");
     let initializer = cx.get_const_i8(0);
     let region_id = add_unnamed_global(&cx, &name, initializer, WeakAnyLinkage);
 
-    let c_entry_name = CString::new(format!("kernel_{num}")).unwrap();
+    let c_entry_name = CString::new(symbol).unwrap();
     let c_val = c_entry_name.as_bytes_with_nul();
-    let offload_entry_name = format!(".offloading.entry_name.{num}");
+    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");
 
     // Not actively used yet, for calling real kernels
-    let name = format!(".offloading.entry.kernel_{num}");
+    let name = format!(".offloading.entry.{symbol}");
 
     // See the __tgt_offload_entry documentation above.
     let reserved = cx.get_const_i64(0);
@@ -248,6 +265,56 @@ fn gen_define_handling<'ll>(
     o_types
 }
 
+// TODO(Sa4dUs): move this to a proper place
+fn get_payload_size<'tcx>(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>) -> u64 {
+    match ty.kind() {
+        /*
+        rustc_middle::infer::canonical::ir::TyKind::Bool => todo!(),
+        rustc_middle::infer::canonical::ir::TyKind::Char => todo!(),
+        rustc_middle::infer::canonical::ir::TyKind::Int(int_ty) => todo!(),
+        rustc_middle::infer::canonical::ir::TyKind::Uint(uint_ty) => todo!(),
+        rustc_middle::infer::canonical::ir::TyKind::Float(float_ty) => todo!(),
+        rustc_middle::infer::canonical::ir::TyKind::Adt(_, _) => todo!(),
+        rustc_middle::infer::canonical::ir::TyKind::Foreign(_) => todo!(),
+        rustc_middle::infer::canonical::ir::TyKind::Str => todo!(),
+        rustc_middle::infer::canonical::ir::TyKind::Array(_, _) => todo!(),
+        rustc_middle::infer::canonical::ir::TyKind::Pat(_, _) => todo!(),
+        rustc_middle::infer::canonical::ir::TyKind::Slice(_) => todo!(),
+        rustc_middle::infer::canonical::ir::TyKind::RawPtr(_, mutability) => todo!(),
+        */
+        ty::Ref(_, inner, _) => get_payload_size(tcx, *inner),
+        /*
+        rustc_middle::infer::canonical::ir::TyKind::FnDef(_, _) => todo!(),
+        rustc_middle::infer::canonical::ir::TyKind::FnPtr(binder, fn_header) => todo!(),
+        rustc_middle::infer::canonical::ir::TyKind::UnsafeBinder(unsafe_binder_inner) => todo!(),
+        rustc_middle::infer::canonical::ir::TyKind::Dynamic(_, _) => todo!(),
+        rustc_middle::infer::canonical::ir::TyKind::Closure(_, _) => todo!(),
+        rustc_middle::infer::canonical::ir::TyKind::CoroutineClosure(_, _) => todo!(),
+        rustc_middle::infer::canonical::ir::TyKind::Coroutine(_, _) => todo!(),
+        rustc_middle::infer::canonical::ir::TyKind::CoroutineWitness(_, _) => todo!(),
+        rustc_middle::infer::canonical::ir::TyKind::Never => todo!(),
+        rustc_middle::infer::canonical::ir::TyKind::Tuple(_) => todo!(),
+        rustc_middle::infer::canonical::ir::TyKind::Alias(alias_ty_kind, alias_ty) => todo!(),
+        rustc_middle::infer::canonical::ir::TyKind::Param(_) => todo!(),
+        rustc_middle::infer::canonical::ir::TyKind::Bound(bound_var_index_kind, _) => todo!(),
+        rustc_middle::infer::canonical::ir::TyKind::Placeholder(_) => todo!(),
+        rustc_middle::infer::canonical::ir::TyKind::Infer(infer_ty) => todo!(),
+        rustc_middle::infer::canonical::ir::TyKind::Error(_) => todo!(),
+        */
+        _ => {
+            tcx
+                // TODO(Sa4dUs): Maybe `.as_query_input()`?
+                .layout_of(PseudoCanonicalInput {
+                    typing_env: TypingEnv::fully_monomorphized(),
+                    value: ty,
+                })
+                .unwrap()
+                .size
+                .bytes()
+        }
+    }
+}
+
 fn declare_offload_fn<'ll>(
     cx: &'ll SimpleCx<'_>,
     name: &str,
@@ -283,10 +350,13 @@ fn declare_offload_fn<'ll>(
 // 4. set insert point after kernel call.
 // 5. generate all the GEPS and stores, to be used in 6)
 // 6. generate __tgt_target_data_end calls to move data from the GPU
-fn gen_call_handling<'ll>(
-    cx: &'ll SimpleCx<'_>,
-    _kernels: &[&'ll llvm::Value],
+pub(crate) fn gen_call_handling<'ll>(
+    cx: &SimpleCx<'ll>,
+    bb: &BasicBlock,
+    kernels: &[&'ll llvm::Value],
     o_types: &[&'ll llvm::Value],
+    llty: &'ll Type,
+    llfn: &'ll Value,
 ) {
     // %struct.__tgt_bin_desc = type { i32, ptr, ptr, ptr }
     let tptr = cx.type_ptr();
@@ -298,27 +368,14 @@ fn gen_call_handling<'ll>(
     gen_tgt_kernel_global(&cx);
     let (begin_mapper_decl, _, end_mapper_decl, fn_ty) = gen_tgt_data_mappers(&cx);
 
-    let main_fn = cx.get_function("main");
-    let Some(main_fn) = main_fn else { return };
-    let kernel_name = "kernel_1";
-    let call = unsafe {
-        llvm::LLVMRustGetFunctionCall(main_fn, kernel_name.as_c_char_ptr(), kernel_name.len())
-    };
-    let Some(kernel_call) = call else {
-        return;
-    };
-    let kernel_call_bb = unsafe { llvm::LLVMGetInstructionParent(kernel_call) };
-    let called = unsafe { llvm::LLVMGetCalledValue(kernel_call).unwrap() };
-    let mut builder = SBuilder::build(cx, kernel_call_bb);
-
-    let types = cx.func_params_types(cx.get_type_of_global(called));
+    let mut builder = SBuilder::build(cx, bb);
+
+    let types = cx.func_params_types(cx.get_type_of_global(kernels[0]));
     let num_args = types.len() as u64;
 
     // Step 0)
     // %struct.__tgt_bin_desc = type { i32, ptr, ptr, ptr }
     // %6 = alloca %struct.__tgt_bin_desc, align 8
-    unsafe { llvm::LLVMRustPositionBuilderPastAllocas(builder.llbuilder, main_fn) };
-
     let tgt_bin_desc_alloca = builder.direct_alloca(tgt_bin_desc, Align::EIGHT, "EmptyDesc");
 
     let ty = cx.type_array(cx.type_ptr(), num_args);
@@ -335,7 +392,7 @@ fn gen_call_handling<'ll>(
     let i32_0 = cx.get_const_i32(0);
     for (index, in_ty) in types.iter().enumerate() {
         // get function arg, store it into the alloca, and read it.
-        let p = llvm::get_param(called, index as u32);
+        let p = llvm::get_param(kernels[0], index as u32);
         let name = llvm::get_value_name(p);
         let name = str::from_utf8(&name).unwrap();
         let arg_name = format!("{name}.addr");
@@ -349,7 +406,6 @@ fn gen_call_handling<'ll>(
     }
 
     // Step 1)
-    unsafe { llvm::LLVMRustPositionBefore(builder.llbuilder, kernel_call) };
     builder.memset(tgt_bin_desc_alloca, cx.get_const_i8(0), cx.get_const_i64(32), Align::EIGHT);
 
     let mapper_fn_ty = cx.type_func(&[cx.type_ptr()], cx.type_void());
@@ -422,10 +478,16 @@ fn gen_call_handling<'ll>(
     // Step 3)
     // Here we will add code for the actual kernel launches in a follow-up PR.
     // FIXME(offload): launch kernels
+    let nparams = llvm::LLVMCountParams(llfn);
+    let mut args = Vec::with_capacity(nparams as usize);
+    for i in 0..nparams {
+        let param = unsafe { llvm::LLVMGetParam(llfn, i) };
+        args.push(param);
+    }
 
-    // Step 4)
-    unsafe { llvm::LLVMRustPositionAfter(builder.llbuilder, kernel_call) };
+    builder.call(llty, kernels[0], &args, None);
 
+    // Step 4)
     let geps = get_geps(&mut builder, &cx, ty, ty2, a1, a2, a4);
     generate_mapper_call(&mut builder, &cx, geps, o, end_mapper_decl, fn_ty, num_args, s_ident_t);
 

compiler/rustc_codegen_llvm/src/intrinsic.rs

@@ -195,6 +195,10 @@ impl<'ll, 'tcx> IntrinsicCallBuilderMethods<'tcx> for Builder<'_, 'll, 'tcx> {
                 codegen_autodiff(self, tcx, instance, args, result);
                 return Ok(());
             }
+            sym::offload => {
+                codegen_offload(self, tcx, instance, args, result);
+                return Ok(());
+            }
             sym::is_val_statically_known => {
                 if let OperandValue::Immediate(imm) = args[0].val {
                     self.call_intrinsic(
@@ -1227,6 +1231,73 @@ fn codegen_autodiff<'ll, 'tcx>(
     );
 }
 
+fn codegen_offload<'ll, 'tcx>(
+    bx: &mut Builder<'_, 'll, 'tcx>,
+    tcx: TyCtxt<'tcx>,
+    instance: ty::Instance<'tcx>,
+    _args: &[OperandRef<'tcx, &'ll Value>],
+    _result: PlaceRef<'tcx, &'ll Value>,
+) {
+    let cx = bx.cx;
+    let fn_args = instance.args;
+
+    let (target_id, target_args) = match fn_args.into_type_list(tcx)[0].kind() {
+        ty::FnDef(def_id, params) => (def_id, params),
+        _ => bug!("invalid offload intrinsic arg"),
+    };
+
+    let fn_target = match Instance::try_resolve(tcx, cx.typing_env(), *target_id, target_args) {
+        Ok(Some(instance)) => instance,
+        Ok(None) => bug!(
+            "could not resolve ({:?}, {:?}) to a specific offload instance",
+            target_id,
+            target_args
+        ),
+        Err(_) => {
+            // An error has already been emitted
+            return;
+        }
+    };
+
+    // TODO(Sa4dUs): Will need typetrees
+    let target_symbol = symbol_name_for_instance_in_crate(tcx, fn_target.clone(), LOCAL_CRATE);
+    let Some(kernel) = cx.get_function(&target_symbol) else {
+        bug!("could not find target function")
+    };
+
+    let offload_entry_ty = crate::builder::gpu_offload::add_tgt_offload_entry(cx);
+
+    // Build TypeTree (or something similar)
+    let sig = tcx.fn_sig(fn_target.def_id()).skip_binder().skip_binder();
+    let inputs = sig.inputs();
+
+    // TODO(Sa4dUs): separate globals from call-independent headers and use typetrees to reserve the correct amount of memory
+    let o = crate::builder::gpu_offload::gen_define_handling(
+        cx,
+        tcx,
+        kernel,
+        offload_entry_ty,
+        inputs.to_vec(),
+        &target_symbol,
+    );
+
+    let kernels = &[kernel];
+    let o_types = &[o];
+
+    let llvm_args = inputs.iter().map(|ty| bx.layout_of(*ty).llvm_type(cx)).collect::<Vec<_>>();
+    let ret_ty = match sig.output().kind() {
+        // TODO(Sa4dUs): dunno if there's a better way of doing this
+        ty::Tuple(tys) if tys.is_empty() => bx.type_void(),
+        _ => bx.layout_of(sig.output()).llvm_type(cx),
+    };
+    let llty = bx.type_func(&llvm_args, ret_ty);
+    let llfn = bx.llfn();
+
+    // TODO(Sa4dUs): this is a patch for delaying lifetime's issue fix
+    let bb = unsafe { llvm::LLVMGetInsertBlock(bx.llbuilder) };
+    crate::builder::gpu_offload::gen_call_handling(cx, bb, kernels, o_types, llty, llfn);
+}
+
 fn get_args_from_tuple<'ll, 'tcx>(
     bx: &mut Builder<'_, 'll, 'tcx>,
     tuple_op: OperandRef<'tcx, &'ll Value>,

compiler/rustc_codegen_llvm/src/lib.rs

@@ -4,6 +4,7 @@
 //!
 //! This API is completely unstable and subject to change.
 
+#![allow(unused)]
 // tidy-alphabetical-start
 #![allow(internal_features)]
 #![doc(html_root_url = "https://doc.rust-lang.org/nightly/nightly-rustc/")]

compiler/rustc_hir_analysis/src/check/intrinsic.rs

@@ -163,6 +163,7 @@ fn intrinsic_operation_unsafety(tcx: TyCtxt<'_>, intrinsic_id: LocalDefId) -> hi
         | sym::minnumf128
         | sym::mul_with_overflow
         | sym::needs_drop
+        | sym::offload
         | sym::powf16
         | sym::powf32
         | sym::powf64
@@ -310,6 +311,7 @@ pub(crate) fn check_intrinsic_type(
             let type_id = tcx.type_of(tcx.lang_items().type_id().unwrap()).instantiate_identity();
             (0, 0, vec![type_id, type_id], tcx.types.bool)
         }
+        sym::offload => (2, 0, vec![param(0)], param(1)),
         sym::offset => (2, 0, vec![param(0), param(1)], param(0)),
         sym::arith_offset => (
             1,

compiler/rustc_span/src/symbol.rs

@@ -1558,6 +1558,7 @@ symbols! {
         object_safe_for_dispatch,
         of,
         off,
+        offload,
         offset,
         offset_of,
         offset_of_enum,

library/core/src/intrinsics/mod.rs

@@ -3262,6 +3262,10 @@ pub const fn copysignf128(x: f128, y: f128) -> f128;
 #[rustc_intrinsic]
 pub const fn autodiff<F, G, T: crate::marker::Tuple, R>(f: F, df: G, args: T) -> R;
 
+#[rustc_nounwind]
+#[rustc_intrinsic]
+pub const fn offload<F, R>(f: F) -> R;
+
 /// Inform Miri that a given pointer definitely has a certain alignment.
 #[cfg(miri)]
 #[rustc_allow_const_fn_unstable(const_eval_select)]