Get types from fn_sig

Author: Sa4dUs

compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs

@@ -261,12 +261,11 @@ pub(crate) fn add_global<'ll>(
 // concatenated into the list of region_ids.
 pub(crate) fn gen_define_handling<'ll>(
     cx: &SimpleCx<'ll>,
-    llfn: &'ll llvm::Value,
     offload_entry_ty: &'ll llvm::Type,
     metadata: &Vec<OffloadMetadata>,
+    types: &Vec<&Type>,
     symbol: &str,
 ) -> (&'ll llvm::Value, &'ll llvm::Value) {
-    let types = cx.func_params_types(cx.get_type_of_global(llfn));
     // It seems like non-pointer values are automatically mapped. So here, we focus on pointer (or
     // reference) types.
     let ptr_meta = types
@@ -371,6 +370,7 @@ pub(crate) fn gen_call_handling<'ll>(
     memtransfer_types: &[&'ll llvm::Value],
     region_ids: &[&'ll llvm::Value],
     llfn: &'ll Value,
+    types: &Vec<&Type>,
     metadata: &Vec<OffloadMetadata>,
 ) {
     let (tgt_decl, tgt_target_kernel_ty) = generate_launcher(&cx);
@@ -386,7 +386,6 @@ pub(crate) fn gen_call_handling<'ll>(
 
     let mut builder = SBuilder::build(cx, bb);
 
-    let types = cx.func_params_types(cx.get_type_of_global(llfn));
     let num_args = types.len() as u64;
 
     // Step 0)

compiler/rustc_codegen_llvm/src/intrinsic.rs

@@ -23,7 +23,7 @@ use rustc_target::callconv::PassMode;
 use rustc_target::spec::Os;
 use tracing::debug;
 
-use crate::abi::FnAbiLlvmExt;
+use crate::abi::{FnAbiLlvmExt, LlvmType};
 use crate::builder::Builder;
 use crate::builder::autodiff::{adjust_activity_to_abi, generate_enzyme_call};
 use crate::builder::gpu_offload::TgtOffloadEntry;
@@ -1276,12 +1276,14 @@ fn codegen_offload<'ll, 'tcx>(
     let metadata = inputs.iter().map(|ty| OffloadMetadata::from_ty(tcx, *ty)).collect::<Vec<_>>();
     let llfn = bx.llfn();
 
+    let types = inputs.iter().map(|ty| cx.layout_of(*ty).llvm_type(cx)).collect::<Vec<_>>();
+
     // TODO(Sa4dUs): separate globals from call-independent headers and use typetrees to reserve the correct amount of memory
     let (memtransfer_type, region_id) = crate::builder::gpu_offload::gen_define_handling(
         cx,
-        llfn,
         offload_entry_ty,
         &metadata,
+        &types,
         &target_symbol,
     );
 
@@ -1293,6 +1295,7 @@ fn codegen_offload<'ll, 'tcx>(
         &[memtransfer_type],
         &[region_id],
         llfn,
+        &types,
         &metadata,
     );
 }