also generate allocas in the caller of kernels

ZuseZ4 · ZuseZ4 · commit d1ef75584268 · 2025-06-09T14:35:24.000-07:00
diff --git a/compiler/rustc_codegen_llvm/src/back/lto.rs b/compiler/rustc_codegen_llvm/src/back/lto.rs
@@ -7,6 +7,7 @@ use std::sync::Arc;
 use std::{io, iter, slice};
 
 use object::read::archive::ArchiveFile;
+use rustc_abi::{Align, Size};
 use rustc_codegen_ssa::back::lto::{LtoModuleCodegen, SerializedModule, ThinModule, ThinShared};
 use rustc_codegen_ssa::back::symbol_export;
 use rustc_codegen_ssa::back::write::{CodegenContext, FatLtoInput};
@@ -648,7 +649,7 @@ fn gen_globals<'ll>(cx: &'ll SimpleCx<'_>) -> (&'ll llvm::Type, &'ll llvm::Value
     let c_val = c_entry_name.as_bytes_with_nul();
     let initializer = crate::common::bytes_in_context(cx.llcx, c_val);
     let at_zero = add_unnamed_global(&cx, &"", initializer, PrivateLinkage);
-    llvm::set_alignment(at_zero, rustc_abi::Align::ONE);
+    llvm::set_alignment(at_zero, Align::ONE);
 
     // @1 = private unnamed_addr constant %struct.ident_t { i32 0, i32 2, i32 0, i32 22, ptr @0 }, align 8
     let struct_ident_ty = cx.type_named_struct("struct.ident_t");
@@ -657,7 +658,7 @@ fn gen_globals<'ll>(cx: &'ll SimpleCx<'_>) -> (&'ll llvm::Type, &'ll llvm::Value
     let initializer = crate::common::named_struct(struct_ident_ty, &struct_elems);
     cx.set_struct_body(struct_ident_ty, &struct_elems_ty, false);
     let at_one = add_unnamed_global(&cx, &"", initializer, PrivateLinkage);
-    llvm::set_alignment(at_one, rustc_abi::Align::EIGHT);
+    llvm::set_alignment(at_one, Align::EIGHT);
 
     // coppied from LLVM
     // typedef struct {
@@ -711,7 +712,7 @@ fn gen_globals<'ll>(cx: &'ll SimpleCx<'_>) -> (&'ll llvm::Type, &'ll llvm::Value
     (offload_entry_ty, at_one, foo, bar, baz, mapper_fn_ty)
 }
 
-fn add_priv_unnamed_arr<'ll>(cx: &SimpleCx<'ll>, name: &str, vals: &[u64]) -> &'ll llvm::Value{
+fn add_priv_unnamed_arr<'ll>(cx: &SimpleCx<'ll>, name: &str, vals: &[u64]) -> &'ll llvm::Value {
     let ti64 = cx.type_i64();
     let size_ty = cx.type_array(ti64, vals.len() as u64);
     let mut size_val = Vec::with_capacity(vals.len());
@@ -739,7 +740,7 @@ fn add_global<'ll>(cx: &SimpleCx<'ll>, name: &str, initializer: &'ll llvm::Value
 
 
 
-fn gen_define_handling<'ll>(cx: &'ll SimpleCx<'_>, offload_entry_ty: &'ll llvm::Type, num: i64) {
+fn gen_define_handling<'ll>(cx: &'ll SimpleCx<'_>, offload_entry_ty: &'ll llvm::Type, num: i64) -> &'ll llvm::Value {
     // We add a pair of sizes and maptypes per offloadable function.
     // @.offload_maptypes = private unnamed_addr constant [4 x i64] [i64 800, i64 544, i64 547, i64 544]
     let o_sizes = add_priv_unnamed_arr(&cx, &format!(".offload_sizes.{num}"), &vec![8u64,0,16,0]);
@@ -763,7 +764,7 @@ fn gen_define_handling<'ll>(cx: &'ll SimpleCx<'_>, offload_entry_ty: &'ll llvm::
 
     let initializer = crate::common::bytes_in_context(cx.llcx, c_val);
     let llglobal = add_unnamed_global(&cx, &foo, initializer, InternalLinkage);
-    llvm::set_alignment(llglobal, rustc_abi::Align::ONE);
+    llvm::set_alignment(llglobal, Align::ONE);
     let c_section_name = CString::new(".llvm.rodata.offloading").unwrap();
     llvm::set_section(llglobal, &c_section_name);
 
@@ -780,7 +781,7 @@ fn gen_define_handling<'ll>(cx: &'ll SimpleCx<'_>, offload_entry_ty: &'ll llvm::
     llvm::set_global_constant(llglobal, true);
     llvm::set_linkage(llglobal, WeakAnyLinkage);
     llvm::set_initializer(llglobal, initializer);
-    llvm::set_alignment(llglobal, rustc_abi::Align::ONE);
+    llvm::set_alignment(llglobal, Align::ONE);
     let c_section_name = CString::new(".omp_offloading_entries").unwrap();
     llvm::set_section(llglobal, &c_section_name);
     // rustc
@@ -795,9 +796,10 @@ fn gen_define_handling<'ll>(cx: &'ll SimpleCx<'_>, offload_entry_ty: &'ll llvm::
     // 3. @.__omp_offloading_<hash>_fnc_name_<hash> = weak constant i8 0
     // 4. @.offloading.entry_name = internal unnamed_addr constant [66 x i8] c"__omp_offloading_86fafab6_c40006a1__Z3fooPSt7complexIdES1_S0_m_l7\00", section ".llvm.rodata.offloading", align 1
     // 5. @.offloading.entry.__omp_offloading_86fafab6_c40006a1__Z3fooPSt7complexIdES1_S0_m_l7 = weak constant %struct.__tgt_offload_entry { i64 0, i16 1, i16 1, i32 0, ptr @.__omp_offloading_86fafab6_c40006a1__Z3fooPSt7complexIdES1_S0_m_l7.region_id, ptr @.offloading.entry_name, i64 0, i64 0, ptr null }, section "omp_offloading_entries", align 1
+    o_types
 }
 
-fn gen_call_handling<'ll>(cx: &'ll SimpleCx<'_>, s_ident_t: &'ll llvm::Value, begin: &'ll llvm::Value, update: &'ll llvm::Value, end: &'ll llvm::Value, fn_ty: &'ll llvm::Type) {
+fn gen_call_handling<'ll>(cx: &'ll SimpleCx<'_>, s_ident_t: &'ll llvm::Value, begin: &'ll llvm::Value, update: &'ll llvm::Value, end: &'ll llvm::Value, fn_ty: &'ll llvm::Type, o_types: &[&'ll llvm::Value]) {
 
     let main_fn = cx.get_function("main");
     if let Some(main_fn) = main_fn {
@@ -811,20 +813,33 @@ fn gen_call_handling<'ll>(cx: &'ll SimpleCx<'_>, s_ident_t: &'ll llvm::Value, be
         };
         let kernel_call_bb = unsafe {llvm::LLVMGetInstructionParent(kernel_call)};
         let mut builder = SBuilder::build(cx, kernel_call_bb);
+
+        // First we generate a few variables used for the data mappers below.
+        // %.offload_baseptrs = alloca [3 x ptr], align 8
+        // %.offload_ptrs = alloca [3 x ptr], align 8
+        // %.offload_mappers = alloca [3 x ptr], align 8
+        // %.offload_sizes = alloca [3 x i64], align 8
+        unsafe{llvm::LLVMRustPositionBuilderPastAllocas(builder.llbuilder, main_fn)};
+        let ty = cx.type_array(cx.type_ptr(), 3);
+        builder.my_alloca2(ty, Align::EIGHT, ".offload_baseptrs");
+        builder.my_alloca2(ty, Align::EIGHT, ".offload_ptrs");
+        builder.my_alloca2(ty, Align::EIGHT, ".offload_mappers");
+        let ty = cx.type_array(cx.type_i64(), 3);
+        builder.my_alloca2(ty, Align::EIGHT, ".offload_sizes");
+
+
+        // Now we generate the __tgt_target_data calls
         unsafe {llvm::LLVMRustPositionBefore(builder.llbuilder, kernel_call)};
         dbg!("positioned builder, ready");
 
         let nullptr = cx.const_null(cx.type_ptr());
-        let args = vec![s_ident_t, cx.get_const_i64(u64::MAX), cx.get_const_i32(3), nullptr, nullptr, nullptr, nullptr, nullptr, nullptr];
+        let o_type = o_types[0];
+        let args = vec![s_ident_t, cx.get_const_i64(u64::MAX), cx.get_const_i32(3), nullptr, nullptr, nullptr, o_type, nullptr, nullptr];
         dbg!(&fn_ty);
         dbg!(&begin);
         dbg!(&args);
         builder.call(fn_ty, begin, &args, None);
         dbg!("called begin");
-        //llty: &'ll Type,
-        //llfn: &'ll Value,
-        //args: &[&'ll Value],
-        //funclet: Option<&Funclet<'ll>>,
 
         // 1. set insert point before kernel call.
         // 2. generate all the GEPS and stores.
@@ -895,14 +910,15 @@ pub(crate) fn run_pass_manager(
             let (offload_entry_ty, at_one, foo, bar, baz, fn_ty) = gen_globals(&cx);
 
             dbg!("created struct");
+            let mut o_types = vec![];
             for num in 0..9 {
                 if !cx.get_function(&format!("kernel_{num}")).is_some() {
                     continue;
                 }
                 // TODO: replace num by proper fn name
-                gen_define_handling(&cx, offload_entry_ty, num);
+                o_types.push(gen_define_handling(&cx, offload_entry_ty, num));
             }
-            gen_call_handling(&cx, at_one, foo, bar, baz, fn_ty);
+            gen_call_handling(&cx, at_one, foo, bar, baz, fn_ty, &o_types);
         } else {
             dbg!("no marker found");
         }
diff --git a/compiler/rustc_codegen_llvm/src/builder.rs b/compiler/rustc_codegen_llvm/src/builder.rs
@@ -88,6 +88,7 @@ impl<'a, 'll> SBuilder<'a, 'll> {
         };
         call
     }
+
 }
 
 impl<'a, 'll, CX: Borrow<SCx<'ll>>> GenericBuilder<'a, 'll, CX> {
@@ -118,6 +119,20 @@ impl<'a, 'll, CX: Borrow<SCx<'ll>>> GenericBuilder<'a, 'll, CX> {
         }
         bx
     }
+
+    pub(crate) fn my_alloca2(&mut self, ty: &'ll Type, align: Align, name: &str) -> &'ll Value {
+        let val = unsafe {
+            let alloca = llvm::LLVMBuildAlloca(self.llbuilder, ty, UNNAMED);
+            llvm::LLVMSetAlignment(alloca, align.bytes() as c_uint);
+            // Cast to default addrspace if necessary
+            llvm::LLVMBuildPointerCast(self.llbuilder, alloca, self.cx.type_ptr(), UNNAMED)
+        };
+        if name != "" {
+            let name = std::ffi::CString::new(name).unwrap();
+            unsafe {llvm::set_value_name(val, &name.as_bytes())};
+        }
+        val
+    }
 }
 
 /// Empty string, to be used where LLVM expects an instruction name, indicating
@@ -1261,7 +1276,7 @@ impl<'a, 'll, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> {
         unsafe {
             llvm::LLVMBuildCleanupRet(self.llbuilder, funclet.cleanuppad(), unwind)
                 .expect("LLVM does not have support for cleanupret");
-        }
+            }
     }
 
     fn catch_pad(&mut self, parent: &'ll Value, args: &[&'ll Value]) -> Funclet<'ll> {
@@ -1631,14 +1646,14 @@ impl<'a, 'll, CX: Borrow<SCx<'ll>>> GenericBuilder<'a, 'll, CX> {
                     debug!(
                         "type mismatch in function call of {:?}. \
                             Expected {:?} for param {}, got {:?}; injecting bitcast",
-                        llfn, expected_ty, i, actual_ty
+                            llfn, expected_ty, i, actual_ty
                     );
                     self.bitcast(actual_val, expected_ty)
                 } else {
                     actual_val
                 }
             })
-            .collect();
+        .collect();
 
         Cow::Owned(casted_args)
     }
@@ -1791,48 +1806,48 @@ impl<'a, 'll, 'tcx> Builder<'a, 'll, 'tcx> {
         let is_indirect_call = unsafe { llvm::LLVMRustIsNonGVFunctionPointerTy(llfn) };
         if self.tcx.sess.is_sanitizer_cfi_enabled()
             && let Some(fn_abi) = fn_abi
-            && is_indirect_call
-        {
-            if let Some(fn_attrs) = fn_attrs
-                && fn_attrs.no_sanitize.contains(SanitizerSet::CFI)
+                && is_indirect_call
             {
-                return;
-            }
+                if let Some(fn_attrs) = fn_attrs
+                    && fn_attrs.no_sanitize.contains(SanitizerSet::CFI)
+                    {
+                        return;
+                    }
 
-            let mut options = cfi::TypeIdOptions::empty();
-            if self.tcx.sess.is_sanitizer_cfi_generalize_pointers_enabled() {
-                options.insert(cfi::TypeIdOptions::GENERALIZE_POINTERS);
-            }
-            if self.tcx.sess.is_sanitizer_cfi_normalize_integers_enabled() {
-                options.insert(cfi::TypeIdOptions::NORMALIZE_INTEGERS);
-            }
+                let mut options = cfi::TypeIdOptions::empty();
+                if self.tcx.sess.is_sanitizer_cfi_generalize_pointers_enabled() {
+                    options.insert(cfi::TypeIdOptions::GENERALIZE_POINTERS);
+                }
+                if self.tcx.sess.is_sanitizer_cfi_normalize_integers_enabled() {
+                    options.insert(cfi::TypeIdOptions::NORMALIZE_INTEGERS);
+                }
 
-            let typeid = if let Some(instance) = instance {
-                cfi::typeid_for_instance(self.tcx, instance, options)
-            } else {
-                cfi::typeid_for_fnabi(self.tcx, fn_abi, options)
-            };
-            let typeid_metadata = self.cx.typeid_metadata(typeid).unwrap();
-            let dbg_loc = self.get_dbg_loc();
-
-            // Test whether the function pointer is associated with the type identifier.
-            let cond = self.type_test(llfn, typeid_metadata);
-            let bb_pass = self.append_sibling_block("type_test.pass");
-            let bb_fail = self.append_sibling_block("type_test.fail");
-            self.cond_br(cond, bb_pass, bb_fail);
-
-            self.switch_to_block(bb_fail);
-            if let Some(dbg_loc) = dbg_loc {
-                self.set_dbg_loc(dbg_loc);
-            }
-            self.abort();
-            self.unreachable();
+                let typeid = if let Some(instance) = instance {
+                    cfi::typeid_for_instance(self.tcx, instance, options)
+                } else {
+                    cfi::typeid_for_fnabi(self.tcx, fn_abi, options)
+                };
+                let typeid_metadata = self.cx.typeid_metadata(typeid).unwrap();
+                let dbg_loc = self.get_dbg_loc();
+
+                // Test whether the function pointer is associated with the type identifier.
+                let cond = self.type_test(llfn, typeid_metadata);
+                let bb_pass = self.append_sibling_block("type_test.pass");
+                let bb_fail = self.append_sibling_block("type_test.fail");
+                self.cond_br(cond, bb_pass, bb_fail);
+
+                self.switch_to_block(bb_fail);
+                if let Some(dbg_loc) = dbg_loc {
+                    self.set_dbg_loc(dbg_loc);
+                }
+                self.abort();
+                self.unreachable();
 
-            self.switch_to_block(bb_pass);
-            if let Some(dbg_loc) = dbg_loc {
-                self.set_dbg_loc(dbg_loc);
+                self.switch_to_block(bb_pass);
+                if let Some(dbg_loc) = dbg_loc {
+                    self.set_dbg_loc(dbg_loc);
+                }
             }
-        }
     }
 
     // Emits KCFI operand bundles.
@@ -1847,31 +1862,31 @@ impl<'a, 'll, 'tcx> Builder<'a, 'll, 'tcx> {
         let kcfi_bundle = if self.tcx.sess.is_sanitizer_kcfi_enabled()
             && let Some(fn_abi) = fn_abi
             && is_indirect_call
-        {
-            if let Some(fn_attrs) = fn_attrs
-                && fn_attrs.no_sanitize.contains(SanitizerSet::KCFI)
             {
-                return None;
-            }
+                if let Some(fn_attrs) = fn_attrs
+                    && fn_attrs.no_sanitize.contains(SanitizerSet::KCFI)
+                    {
+                        return None;
+                    }
 
-            let mut options = kcfi::TypeIdOptions::empty();
-            if self.tcx.sess.is_sanitizer_cfi_generalize_pointers_enabled() {
-                options.insert(kcfi::TypeIdOptions::GENERALIZE_POINTERS);
-            }
-            if self.tcx.sess.is_sanitizer_cfi_normalize_integers_enabled() {
-                options.insert(kcfi::TypeIdOptions::NORMALIZE_INTEGERS);
-            }
+                let mut options = kcfi::TypeIdOptions::empty();
+                if self.tcx.sess.is_sanitizer_cfi_generalize_pointers_enabled() {
+                    options.insert(kcfi::TypeIdOptions::GENERALIZE_POINTERS);
+                }
+                if self.tcx.sess.is_sanitizer_cfi_normalize_integers_enabled() {
+                    options.insert(kcfi::TypeIdOptions::NORMALIZE_INTEGERS);
+                }
 
-            let kcfi_typeid = if let Some(instance) = instance {
-                kcfi::typeid_for_instance(self.tcx, instance, options)
+                let kcfi_typeid = if let Some(instance) = instance {
+                    kcfi::typeid_for_instance(self.tcx, instance, options)
+                } else {
+                    kcfi::typeid_for_fnabi(self.tcx, fn_abi, options)
+                };
+
+                Some(llvm::OperandBundleBox::new("kcfi", &[self.const_u32(kcfi_typeid)]))
             } else {
-                kcfi::typeid_for_fnabi(self.tcx, fn_abi, options)
+                None
             };
-
-            Some(llvm::OperandBundleBox::new("kcfi", &[self.const_u32(kcfi_typeid)]))
-        } else {
-            None
-        };
         kcfi_bundle
     }
 
diff --git a/compiler/rustc_codegen_llvm/src/llvm/ffi.rs b/compiler/rustc_codegen_llvm/src/llvm/ffi.rs
@@ -2560,6 +2560,7 @@ unsafe extern "C" {
 
     pub(crate) fn LLVMRustSetDataLayoutFromTargetMachine<'a>(M: &'a Module, TM: &'a TargetMachine);
 
+    pub(crate) fn LLVMRustPositionBuilderPastAllocas<'a>(B: &Builder<'a>, Fn: &'a Value);
     pub(crate) fn LLVMRustPositionBuilderAtStart<'a>(B: &Builder<'a>, BB: &'a BasicBlock);
 
     pub(crate) fn LLVMRustSetModulePICLevel(M: &Module);
diff --git a/compiler/rustc_llvm/llvm-wrapper/RustWrapper.cpp b/compiler/rustc_llvm/llvm-wrapper/RustWrapper.cpp
@@ -1591,6 +1591,11 @@ extern "C" LLVMValueRef LLVMRustBuildMemSet(LLVMBuilderRef B, LLVMValueRef Dst,
                                       MaybeAlign(DstAlign), IsVolatile));
 }
 
+extern "C" void LLVMRustPositionBuilderPastAllocas(LLVMBuilderRef B,
+                                                   LLVMValueRef Fn) {
+  Function *F = unwrap<Function>(Fn);
+  unwrap(B)->SetInsertPointPastAllocas(F);
+}
 extern "C" void LLVMRustPositionBuilderAtStart(LLVMBuilderRef B,
                                                LLVMBasicBlockRef BB) {
   auto Point = unwrap(BB)->getFirstInsertionPt();
