Mark globals as used + some minor fixes

Author: Sa4dUs

compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs

@@ -14,7 +14,7 @@ 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
@@ -259,15 +259,14 @@ pub(crate) fn add_global<'ll>(
 // This function returns a memtransfer value which encodes how arguments to this kernel shall be
 // mapped to/from the gpu. It also returns a region_id with the name of this kernel, to be
 // concatenated into the list of region_ids.
-pub(crate) fn gen_define_handling<'ll, 'tcx>(
+pub(crate) fn gen_define_handling<'ll>(
     cx: &SimpleCx<'ll>,
-    tcx: TyCtxt<'tcx>,
-    kernel: &'ll llvm::Value,
+    llfn: &'ll llvm::Value,
     offload_entry_ty: &'ll llvm::Type,
-    metadata: Vec<OffloadMetadata>,
+    metadata: &Vec<OffloadMetadata>,
     symbol: &str,
 ) -> (&'ll llvm::Value, &'ll llvm::Value) {
-    let types = cx.func_params_types(cx.get_type_of_global(kernel));
+    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
@@ -277,7 +276,7 @@ pub(crate) fn gen_define_handling<'ll, 'tcx>(
             rustc_codegen_ssa::common::TypeKind::Pointer => Some(meta),
             _ => None,
         })
-        .collect::<Vec<OffloadMetadata>>();
+        .collect::<Vec<_>>();
 
     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<_>>();
@@ -286,7 +285,7 @@ pub(crate) fn gen_define_handling<'ll, 'tcx>(
     // 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.{symbol}"), &ptr_sizes);
+    let offload_sizes = 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
@@ -326,6 +325,8 @@ pub(crate) fn gen_define_handling<'ll, 'tcx>(
     llvm::set_alignment(llglobal, Align::EIGHT);
     let c_section_name = CString::new("llvm_offload_entries").unwrap();
     llvm::set_section(llglobal, &c_section_name);
+
+    add_to_llvm_used(cx, &[offload_sizes, memtransfer_types, region_id, llglobal]);
     (memtransfer_types, region_id)
 }
 
@@ -367,11 +368,10 @@ fn declare_offload_fn<'ll>(
 pub(crate) fn gen_call_handling<'ll>(
     cx: &SimpleCx<'ll>,
     bb: &BasicBlock,
-    kernel: &'ll llvm::Value,
     memtransfer_types: &[&'ll llvm::Value],
     region_ids: &[&'ll llvm::Value],
     llfn: &'ll Value,
-    metadata: Vec<OffloadMetadata>,
+    metadata: &Vec<OffloadMetadata>,
 ) {
     let (tgt_decl, tgt_target_kernel_ty) = generate_launcher(&cx);
     // %struct.__tgt_bin_desc = type { i32, ptr, ptr, ptr }
@@ -386,7 +386,7 @@ 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(kernel));
+    let types = cx.func_params_types(cx.get_type_of_global(llfn));
     let num_args = types.len() as u64;
 
     // Step 0)
@@ -442,7 +442,7 @@ pub(crate) fn gen_call_handling<'ll>(
         // As mentioned above, we don't use Rust type information yet. So for now we will just
         // assume that we have 1024 bytes, 256 f32 values.
         // FIXME(offload): write an offload frontend and handle arbitrary types.
-        builder.store(cx.get_const_i64(metadata[i].payload_size), gep3, Align::EIGHT);
+        builder.store(cx.get_const_i64(metadata[i as usize].payload_size), gep3, Align::EIGHT);
     }
 
     // For now we have a very simplistic indexing scheme into our
@@ -517,3 +517,41 @@ pub(crate) fn gen_call_handling<'ll>(
 
     drop(builder);
 }
+
+// TODO(Sa4dUs): check if there's a better way of doing this, also move to a proper location
+fn add_to_llvm_used<'ll>(cx: &'ll SimpleCx<'_>, globals: &[&'ll Value]) {
+    let ptr_ty = cx.type_ptr();
+    let arr_ty = cx.type_array(ptr_ty, globals.len() as u64);
+    let arr_val = cx.const_array(ptr_ty, globals);
+
+    let name = CString::new("llvm.used").unwrap();
+
+    let used_global_opt = unsafe { llvm::LLVMGetNamedGlobal(cx.llmod, name.as_ptr()) };
+
+    if used_global_opt.is_none() {
+        let new_global = unsafe { llvm::LLVMAddGlobal(cx.llmod, arr_ty, name.as_ptr()) };
+        unsafe { llvm::LLVMSetLinkage(new_global, llvm::Linkage::AppendingLinkage) };
+        unsafe {
+            llvm::LLVMSetSection(new_global, CString::new("llvm.metadata").unwrap().as_ptr())
+        };
+        unsafe { llvm::LLVMSetInitializer(new_global, arr_val) };
+        llvm::LLVMSetGlobalConstant(new_global, llvm::TRUE);
+        return;
+    }
+
+    let used_global = used_global_opt.expect("expected @llvm.used");
+    let mut combined: Vec<&'ll Value> = Vec::new();
+
+    if let Some(existing_init) = llvm::LLVMGetInitializer(used_global) {
+        let num_elems = unsafe { llvm::LLVMGetNumOperands(existing_init) };
+        for i in 0..num_elems {
+            if let Some(elem) = unsafe { llvm::LLVMGetOperand(existing_init, i) } {
+                combined.push(elem);
+            }
+        }
+    }
+
+    combined.extend_from_slice(globals);
+    let new_arr = cx.const_array(ptr_ty, &combined);
+    unsafe { llvm::LLVMSetInitializer(used_global, new_arr) };
+}

compiler/rustc_codegen_llvm/src/intrinsic.rs

@@ -1262,9 +1262,6 @@ fn codegen_offload<'ll, 'tcx>(
     };
 
     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 = TgtOffloadEntry::new_decl(&cx);
 
@@ -1273,29 +1270,26 @@ fn codegen_offload<'ll, 'tcx>(
     let inputs = sig.inputs();
 
     let metadata = inputs.iter().map(|ty| OffloadMetadata::from_ty(tcx, *ty)).collect::<Vec<_>>();
+    let llfn = bx.llfn();
 
     // 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,
-        tcx,
-        kernel,
+        llfn,
         offload_entry_ty,
-        metadata,
+        &metadata,
         &target_symbol,
     );
 
-    let llfn = bx.llfn();
-
     // TODO(Sa4dUs): this is just to a void lifetime's issues
     let bb = unsafe { llvm::LLVMGetInsertBlock(bx.llbuilder) };
     crate::builder::gpu_offload::gen_call_handling(
         cx,
         bb,
-        kernel,
         &[memtransfer_type],
         &[region_id],
         llfn,
-        metadata,
+        &metadata,
     );
 }
 

compiler/rustc_codegen_llvm/src/llvm/ffi.rs

@@ -1164,6 +1164,7 @@ unsafe extern "C" {
     pub(crate) fn LLVMGetOperand(Val: &Value, Index: c_uint) -> Option<&Value>;
     pub(crate) fn LLVMGetNextInstruction(Val: &Value) -> Option<&Value>;
     pub(crate) fn LLVMInstructionEraseFromParent(Val: &Value);
+    pub(crate) fn LLVMGetNumOperands(Val: &Value) -> c_uint;
 
     // Operations on call sites
     pub(crate) fn LLVMSetInstructionCallConv(Instr: &Value, CC: c_uint);

tests/codegen-llvm/gpu_offload/gpu_host.rs

@@ -11,12 +11,13 @@
 // when inside of a function called main. This, too, is a temporary workaround for not having a
 // frontend.
 
+#![feature(core_intrinsics)]
 #![no_main]
 
 #[unsafe(no_mangle)]
 fn main() {
     let mut x = [3.0; 256];
-    kernel_1(&mut x);
+    kernel(&mut x);
     core::hint::black_box(&x);
 }
 
@@ -25,13 +26,14 @@ fn main() {
 // CHECK: %struct.__tgt_bin_desc = type { i32, ptr, ptr, ptr }
 // CHECK: %struct.__tgt_kernel_arguments = type { i32, i32, ptr, ptr, ptr, ptr, ptr, ptr, i64, i64, [3 x i32], [3 x i32], i32 }
 
-// CHECK: @.offload_sizes.1 = private unnamed_addr constant [1 x i64] [i64 1024]
-// CHECK: @.offload_maptypes.1 = private unnamed_addr constant [1 x i64] [i64 35]
-// CHECK: @.kernel_1.region_id = weak unnamed_addr constant i8 0
-// CHECK: @.offloading.entry_name.1 = internal unnamed_addr constant [9 x i8] c"kernel_1\00", section ".llvm.rodata.offloading", align 1
-// CHECK: @.offloading.entry.kernel_1 = weak constant %struct.__tgt_offload_entry { i64 0, i16 1, i16 1, i32 0, ptr @.kernel_1.region_id, ptr @.offloading.entry_name.1, i64 0, i64 0, ptr null }, section "llvm_offload_entries", align 8
-// CHECK: @0 = private unnamed_addr constant [23 x i8] c";unknown;unknown;0;0;;\00", align 1
-// CHECK: @1 = private unnamed_addr constant %struct.ident_t { i32 0, i32 2, i32 0, i32 22, ptr @0 }, align 8
+// CHECK: @.offload_sizes._kernel = private unnamed_addr constant [1 x i64] [i64 1024]
+// CHECK: @.offload_maptypes._kernel = private unnamed_addr constant [1 x i64] [i64 35]
+// CHECK: @._kernel.region_id = weak unnamed_addr constant i8 0
+// CHECK: @.offloading.entry_name._kernel = internal unnamed_addr constant [8 x i8] c"_kernel\00", section ".llvm.rodata.offloading", align 1
+// CHECK: @.offloading.entry._kernel = weak constant %struct.__tgt_offload_entry { i64 0, i16 1, i16 1, i32 0, ptr @._kernel.region_id, ptr @.offloading.entry_name._kernel, i64 0, i64 0, ptr null }, section "llvm_offload_entries", align 8
+
+// CHECK: @anon.{{.*}}.0 = private unnamed_addr constant [23 x i8] c";unknown;unknown;0;0;;\00", align 1
+// CHECK: @anon.{{.*}}.1 = private unnamed_addr constant %struct.ident_t { i32 0, i32 2, i32 0, i32 22, ptr @anon.{{.*}}.0 }, align 8
 
 // CHECK:  Function Attrs:
 // CHECK-NEXT: define{{( dso_local)?}} void @main()
@@ -99,7 +101,13 @@ fn main() {
 
 #[unsafe(no_mangle)]
 #[inline(never)]
-pub fn kernel_1(x: &mut [f32; 256]) {
+pub fn kernel(x: &mut [f32; 256]) {
+    core::intrinsics::offload(_kernel)
+}
+
+#[unsafe(no_mangle)]
+#[inline(never)]
+pub fn _kernel(x: &mut [f32; 256]) {
     for i in 0..256 {
         x[i] = 21.0;
     }