Docs and minor fixes

Author: Sa4dUs

compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs

@@ -170,6 +170,14 @@ impl KernelArgsTy {
     }
 }
 
+// Contains LLVM values needed to manage offloading for a single kernel.
+pub(crate) struct OffloadKernelData<'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,
+}
+
 fn gen_tgt_data_mappers<'ll>(
     cx: &'ll SimpleCx<'_>,
 ) -> (&'ll llvm::Value, &'ll llvm::Value, &'ll llvm::Value, &'ll llvm::Type) {
@@ -238,7 +246,7 @@ pub(crate) fn gen_define_handling<'ll>(
     metadata: &[OffloadMetadata],
     types: &[&Type],
     symbol: &str,
-) -> (&'ll llvm::Value, &'ll llvm::Value, &'ll llvm::Value, &'ll llvm::Value) {
+) -> OffloadKernelData<'ll> {
     // 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) {
@@ -290,7 +298,7 @@ pub(crate) fn gen_define_handling<'ll>(
     let c_section_name = CString::new("llvm_offload_entries").unwrap();
     llvm::set_section(offload_entry, &c_section_name);
 
-    (offload_sizes, memtransfer_types, region_id, offload_entry)
+    OffloadKernelData { offload_sizes, memtransfer_types, region_id, offload_entry }
 }
 
 fn declare_offload_fn<'ll>(
@@ -330,14 +338,13 @@ fn declare_offload_fn<'ll>(
 pub(crate) fn gen_call_handling<'ll>(
     cx: &SimpleCx<'ll>,
     bb: &BasicBlock,
-    offload_sizes: &'ll llvm::Value,
-    offload_entry: &'ll llvm::Value,
-    memtransfer_types: &'ll llvm::Value,
-    region_id: &'ll llvm::Value,
+    offload_data: &OffloadKernelData<'ll>,
     args: &[&'ll Value],
     types: &[&Type],
     metadata: &[OffloadMetadata],
 ) {
+    let OffloadKernelData { offload_sizes, offload_entry, memtransfer_types, region_id } =
+        offload_data;
     let (tgt_decl, tgt_target_kernel_ty) = generate_launcher(&cx);
     // %struct.__tgt_bin_desc = type { i32, ptr, ptr, ptr }
     let tptr = cx.type_ptr();
@@ -351,7 +358,11 @@ pub(crate) fn gen_call_handling<'ll>(
 
     let mut builder = SBuilder::build(cx, bb);
 
-    // prevent these globals from being optimized away
+    let num_args = types.len() as u64;
+    let ip = unsafe { llvm::LLVMRustGetInsertPoint(&builder.llbuilder) };
+
+    // 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] {
         unsafe {
             let dummy = llvm::LLVMBuildLoad2(
@@ -364,11 +375,13 @@ pub(crate) fn gen_call_handling<'ll>(
         }
     }
 
-    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
+    let llfn = unsafe { llvm::LLVMGetBasicBlockParent(bb) };
+    unsafe {
+        llvm::LLVMRustPositionBuilderPastAllocas(&builder.llbuilder, llfn);
+    }
     let tgt_bin_desc_alloca = builder.direct_alloca(tgt_bin_desc, Align::EIGHT, "EmptyDesc");
 
     let ty = cx.type_array(cx.type_ptr(), num_args);
@@ -384,6 +397,9 @@ pub(crate) fn gen_call_handling<'ll>(
     let a5 = builder.direct_alloca(tgt_kernel_decl, Align::EIGHT, "kernel_args");
 
     // Step 1)
+    unsafe {
+        llvm::LLVMRustRestoreInsertPoint(&builder.llbuilder, ip);
+    }
     builder.memset(tgt_bin_desc_alloca, cx.get_const_i8(0), cx.get_const_i64(32), Align::EIGHT);
 
     // Now we allocate once per function param, a copy to be passed to one of our maps.
@@ -458,9 +474,17 @@ pub(crate) fn gen_call_handling<'ll>(
 
     // Step 2)
     let s_ident_t = generate_at_one(&cx);
-    let o = memtransfer_types;
     let geps = get_geps(&mut builder, &cx, ty, ty2, a1, a2, a4);
-    generate_mapper_call(&mut builder, &cx, geps, o, begin_mapper_decl, fn_ty, num_args, s_ident_t);
+    generate_mapper_call(
+        &mut builder,
+        &cx,
+        geps,
+        memtransfer_types,
+        begin_mapper_decl,
+        fn_ty,
+        num_args,
+        s_ident_t,
+    );
     let values = KernelArgsTy::new(&cx, num_args, memtransfer_types, geps);
 
     // Step 3)
@@ -485,7 +509,16 @@ pub(crate) fn gen_call_handling<'ll>(
 
     // 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);
+    generate_mapper_call(
+        &mut builder,
+        &cx,
+        geps,
+        memtransfer_types,
+        end_mapper_decl,
+        fn_ty,
+        num_args,
+        s_ident_t,
+    );
 
     builder.call(mapper_fn_ty, unregister_lib_decl, &[tgt_bin_desc_alloca], None);
 

compiler/rustc_codegen_llvm/src/intrinsic.rs

@@ -1243,6 +1243,10 @@ fn codegen_autodiff<'ll, 'tcx>(
     );
 }
 
+// Generates the LLVM code to offload a Rust function to a target device (e.g., GPU).
+// For each kernel call, it generates the necessary globals (including metadata such as
+// size and pass mode), manages memory mapping to and from the device, handles all
+// data transfers, and launches the kernel on the target device.
 fn codegen_offload<'ll, 'tcx>(
     bx: &mut Builder<'_, 'll, 'tcx>,
     tcx: TyCtxt<'tcx>,
@@ -1282,27 +1286,16 @@ fn codegen_offload<'ll, 'tcx>(
 
     let types = inputs.iter().map(|ty| cx.layout_of(*ty).llvm_type(cx)).collect::<Vec<_>>();
 
-    let (offload_sizes, memtransfer_types, region_id, offload_entry) =
-        crate::builder::gpu_offload::gen_define_handling(
-            cx,
-            offload_entry_ty,
-            &metadata,
-            &types,
-            &target_symbol,
-        );
-
-    let bb = unsafe { llvm::LLVMGetInsertBlock(bx.llbuilder) };
-    crate::builder::gpu_offload::gen_call_handling(
+    let offload_data = crate::builder::gpu_offload::gen_define_handling(
         cx,
-        bb,
-        offload_sizes,
-        offload_entry,
-        memtransfer_types,
-        region_id,
-        &args,
-        &types,
+        offload_entry_ty,
         &metadata,
+        &types,
+        &target_symbol,
     );
+
+    let bb = unsafe { llvm::LLVMGetInsertBlock(bx.llbuilder) };
+    crate::builder::gpu_offload::gen_call_handling(cx, bb, &offload_data, &args, &types, &metadata);
 }
 
 fn get_args_from_tuple<'ll, 'tcx>(

compiler/rustc_codegen_llvm/src/llvm/ffi.rs

@@ -2448,7 +2448,10 @@ 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 LLVMRustGetInsertPoint<'a>(B: &Builder<'a>) -> &'a Value;
+    pub(crate) fn LLVMRustRestoreInsertPoint<'a>(B: &Builder<'a>, IP: &'a Value);
 
     pub(crate) fn LLVMRustSetModulePICLevel(M: &Module);
     pub(crate) fn LLVMRustSetModulePIELevel(M: &Module);

compiler/rustc_llvm/llvm-wrapper/RustWrapper.cpp

@@ -1378,6 +1378,39 @@ extern "C" void LLVMRustPositionAfter(LLVMBuilderRef B, LLVMValueRef Instr) {
   }
 }
 
+extern "C" LLVMValueRef LLVMRustGetInsertPoint(LLVMBuilderRef B) {
+    llvm::IRBuilderBase &IRB = *unwrap(B);
+
+    llvm::IRBuilderBase::InsertPoint ip = IRB.saveIP();
+    llvm::BasicBlock *BB = ip.getBlock();
+
+    if (!BB)
+        return nullptr;
+
+    auto it = ip.getPoint();
+
+    if (it == BB->end())
+        return nullptr;
+
+    llvm::Instruction *I = &*it;
+    return wrap(I);
+}
+
+extern "C" void LLVMRustRestoreInsertPoint(LLVMBuilderRef B, LLVMValueRef Instr) {
+    llvm::IRBuilderBase &IRB = *unwrap(B);
+
+    if (!Instr) {
+        llvm::BasicBlock *BB = IRB.GetInsertBlock();
+        if (BB)
+            IRB.SetInsertPoint(BB);
+        return;
+    }
+
+    llvm::Instruction *I = unwrap<llvm::Instruction>(Instr);
+    IRB.SetInsertPoint(I);
+}
+
+
 extern "C" LLVMValueRef
 LLVMRustGetFunctionCall(LLVMValueRef Fn, const char *Name, size_t NameLen) {
   auto targetName = StringRef(Name, NameLen);

tests/codegen-llvm/gpu_offload/gpu_host.rs

@@ -51,13 +51,13 @@ fn main() {
 
 // CHECK:      define{{( dso_local)?}} void @kernel_1(ptr noalias noundef align 4 dereferenceable(1024) %x)
 // CHECK-NEXT: start:
-// CHECK-NEXT:   %dummy = load volatile ptr, ptr @.offload_sizes._kernel_1, align 8
-// CHECK-NEXT:   %dummy1 = load volatile ptr, ptr @.offloading.entry._kernel_1, align 8
 // CHECK-NEXT:   %EmptyDesc = alloca %struct.__tgt_bin_desc, align 8
 // CHECK-NEXT:   %.offload_baseptrs = alloca [1 x ptr], align 8
 // CHECK-NEXT:   %.offload_ptrs = alloca [1 x ptr], align 8
 // CHECK-NEXT:   %.offload_sizes = alloca [1 x i64], align 8
 // CHECK-NEXT:   %kernel_args = alloca %struct.__tgt_kernel_arguments, align 8
+// CHECK-NEXT:   %dummy = load volatile ptr, ptr @.offload_sizes._kernel_1, align 8
+// CHECK-NEXT:   %dummy1 = load volatile ptr, ptr @.offloading.entry._kernel_1, align 8
 // CHECK-NEXT:   call void @llvm.memset.p0.i64(ptr noundef nonnull align 8 dereferenceable(32) %EmptyDesc, i8 0, i64 32, i1 false)
 // CHECK-NEXT:   call void @__tgt_register_lib(ptr nonnull %EmptyDesc)
 // CHECK-NEXT:   call void @__tgt_init_all_rtls()