new code

Author: ZuseZ4

compiler/rustc_codegen_llvm/src/back/write.rs

@@ -926,7 +926,7 @@ pub(crate) fn codegen(
             // binaries. So we must clone the module to produce the asm output
             // if we are also producing object code.
             let llmod = if let EmitObj::ObjectCode(_) = config.emit_obj {
-                llvm::LLVMCloneModule(llmod)
+                unsafe { llvm::LLVMCloneModule(llmod) }
             } else {
                 llmod
             };

compiler/rustc_codegen_llvm/src/builder/gpu_device.rs

@@ -14,7 +14,7 @@ fn add_unnamed_global_in_addrspace<'ll>(
     addrspace: u32,
 ) -> &'ll llvm::Value {
     let llglobal = add_global_in_addrspace(cx, name, initializer, l, addrspace);
-    unsafe { llvm::LLVMSetUnnamedAddress(llglobal, llvm::UnnamedAddr::Global) };
+    llvm::LLVMSetUnnamedAddress(llglobal, llvm::UnnamedAddr::Global);
     llglobal
 }
 

compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs

@@ -18,18 +18,36 @@ pub(crate) fn handle_gpu_code<'ll>(
     // The offload memory transfer type for each kernel
     let mut o_types = vec![];
     let mut kernels = vec![];
+    let mut region_ids = vec![];
     let offload_entry_ty = add_tgt_offload_entry(&cx);
     for num in 0..9 {
         let kernel = cx.get_function(&format!("kernel_{num}"));
         if let Some(kernel) = kernel {
-            o_types.push(gen_define_handling(&cx, kernel, offload_entry_ty, num));
+            let (o, k) = gen_define_handling(&cx, kernel, offload_entry_ty, num);
+            o_types.push(o);
+            region_ids.push(k);
             kernels.push(kernel);
         }
     }
-    gen_call_handling(&cx, &kernels, &o_types);
+    gen_call_handling(&cx, &kernels, &o_types, &region_ids);
     crate::builder::gpu_wrapper::gen_image_wrapper_module(&cgcx);
 }
 
+// ; Function Attrs: nounwind
+// declare i32 @__tgt_target_kernel(ptr, i64, i32, i32, ptr, ptr) #2
+fn generate_launcher<'ll>(cx: &'ll SimpleCx<'_>) -> (&'ll llvm::Value, &'ll llvm::Type) {
+    let tptr = cx.type_ptr();
+    let ti64 = cx.type_i64();
+    let ti32 = cx.type_i32();
+    let args = vec![tptr, ti64, ti32, ti32, tptr, tptr];
+    let tgt_fn_ty = cx.type_func(&args, ti32);
+    let name = "__tgt_target_kernel";
+    let tgt_decl = declare_offload_fn(&cx, name, tgt_fn_ty);
+    let nounwind = llvm::AttributeKind::NoUnwind.create_attr(cx.llcx);
+    attributes::apply_to_llfn(tgt_decl, Function, &[nounwind]);
+    (tgt_decl, tgt_fn_ty)
+}
+
 // What is our @1 here? A magic global, used in our data_{begin/update/end}_mapper:
 // @0 = private unnamed_addr constant [23 x i8] c";unknown;unknown;0;0;;\00", align 1
 // @1 = private unnamed_addr constant %struct.ident_t { i32 0, i32 2, i32 0, i32 22, ptr @0 }, align 8
@@ -83,7 +101,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<'_>) {
+fn gen_tgt_kernel_global<'ll>(cx: &'ll SimpleCx<'_>) -> (&'ll llvm::Type, Vec<&'ll llvm::Type>) {
     let kernel_arguments_ty = cx.type_named_struct("struct.__tgt_kernel_arguments");
     let tptr = cx.type_ptr();
     let ti64 = cx.type_i64();
@@ -118,9 +136,10 @@ fn gen_tgt_kernel_global<'ll>(cx: &'ll SimpleCx<'_>) {
         vec![ti32, ti32, tptr, tptr, tptr, tptr, tptr, tptr, ti64, ti64, tarr, tarr, ti32];
 
     cx.set_struct_body(kernel_arguments_ty, &kernel_elements, false);
+    (kernel_arguments_ty, kernel_elements)
     // For now we don't handle kernels, so for now we just add a global dummy
     // to make sure that the __tgt_offload_entry is defined and handled correctly.
-    cx.declare_global("my_struct_global2", kernel_arguments_ty);
+    //cx.declare_global("my_struct_global2", kernel_arguments_ty);
 }
 
 fn gen_tgt_data_mappers<'ll>(
@@ -187,7 +206,7 @@ fn gen_define_handling<'ll>(
     kernel: &'ll llvm::Value,
     offload_entry_ty: &'ll llvm::Type,
     num: i64,
-) -> &'ll llvm::Value {
+) -> (&'ll llvm::Value, &'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
     // reference) types.
@@ -205,10 +224,11 @@ fn gen_define_handling<'ll>(
     // 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.
+    // 1+2+32: 1 (MapTo), 2 (MapFrom), 32 (Add extra input ptr once, idk, figure out later)
     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.{num}"), &vec![1+2+32; 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
+    // the llvm.rodata entry name, and  the llvm_offload_entries value
 
     let name = format!(".kernel_{num}.region_id");
     let initializer = cx.get_const_i8(0);
@@ -243,12 +263,12 @@ fn gen_define_handling<'ll>(
     llvm::set_linkage(llglobal, WeakAnyLinkage);
     llvm::set_initializer(llglobal, initializer);
     llvm::set_alignment(llglobal, Align::ONE);
-    let c_section_name = CString::new(".omp_offloading_entries").unwrap();
+    let c_section_name = CString::new("llvm_offload_entries").unwrap();
     llvm::set_section(llglobal, &c_section_name);
-    o_types
+    (o_types, region_id)
 }
 
-fn declare_offload_fn<'ll>(
+pub(crate) fn declare_offload_fn<'ll>(
     cx: &'ll SimpleCx<'_>,
     name: &str,
     ty: &'ll llvm::Type,
@@ -287,15 +307,17 @@ fn gen_call_handling<'ll>(
     cx: &'ll SimpleCx<'_>,
     _kernels: &[&'ll llvm::Value],
     o_types: &[&'ll llvm::Value],
+    region_ids: &[&'ll llvm::Value],
 ) {
+    let (tgt_decl, tgt_target_kernel_ty) = generate_launcher(&cx);
     // %struct.__tgt_bin_desc = type { i32, ptr, ptr, ptr }
     let tptr = cx.type_ptr();
     let ti32 = cx.type_i32();
     let tgt_bin_desc_ty = vec![ti32, tptr, tptr, tptr];
     let tgt_bin_desc = cx.type_named_struct("struct.__tgt_bin_desc");
     cx.set_struct_body(tgt_bin_desc, &tgt_bin_desc_ty, false);
 
-    gen_tgt_kernel_global(&cx);
+    let (tgt_kernel_decl, tgt_kernel_types) = 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");
@@ -329,29 +351,33 @@ fn gen_call_handling<'ll>(
     // These represent the sizes in bytes, e.g. the entry for `&[f64; 16]` will be 8*16.
     let ty2 = cx.type_array(cx.type_i64(), num_args);
     let a4 = builder.direct_alloca(ty2, Align::EIGHT, ".offload_sizes");
+
+    //%kernel_args = alloca %struct.__tgt_kernel_arguments, align 8
+    let a5 = builder.direct_alloca(tgt_kernel_decl, Align::EIGHT, "kernel_args");
+
+    // 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);
+
     // Now we allocate once per function param, a copy to be passed to one of our maps.
     let mut vals = vec![];
     let mut geps = vec![];
     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 name = llvm::get_value_name(p);
-        let name = str::from_utf8(&name).unwrap();
-        let arg_name = format!("{name}.addr");
-        let alloca = builder.direct_alloca(in_ty, Align::EIGHT, &arg_name);
-
-        builder.store(p, alloca, Align::EIGHT);
-        let val = builder.load(in_ty, alloca, Align::EIGHT);
-        let gep = builder.inbounds_gep(cx.type_f32(), val, &[i32_0]);
-        vals.push(val);
+        //let p = llvm::get_param(called, index as u32);
+        //let name = llvm::get_value_name(p);
+        //let name = str::from_utf8(&name).unwrap();
+        //let arg_name = format!("{name}.addr");
+        //let alloca = builder.direct_alloca(in_ty, Align::EIGHT, &arg_name);
+
+        let v = unsafe { llvm::LLVMGetOperand(kernel_call, index as u32).unwrap() };
+        let gep = builder.inbounds_gep(cx.type_f32(), v, &[i32_0]);
+        vals.push(v);
+        //vals.push(val);
         geps.push(gep);
     }
 
-    // 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());
     let register_lib_decl = declare_offload_fn(&cx, "__tgt_register_lib", mapper_fn_ty);
     let unregister_lib_decl = declare_offload_fn(&cx, "__tgt_unregister_lib", mapper_fn_ty);
@@ -421,16 +447,87 @@ fn gen_call_handling<'ll>(
 
     // Step 3)
     // Here we will add code for the actual kernel launches in a follow-up PR.
+    //%28 = getelementptr inbounds nuw %struct.__tgt_kernel_arguments, ptr %kernel_args, i32 0, i32 0
+    //store i32 3, ptr %28, align 4
+    //%29 = getelementptr inbounds nuw %struct.__tgt_kernel_arguments, ptr %kernel_args, i32 0, i32 1
+    //store i32 3, ptr %29, align 4
+    //%30 = getelementptr inbounds nuw %struct.__tgt_kernel_arguments, ptr %kernel_args, i32 0, i32 2
+    //store ptr %26, ptr %30, align 8
+    //%31 = getelementptr inbounds nuw %struct.__tgt_kernel_arguments, ptr %kernel_args, i32 0, i32 3
+    //store ptr %27, ptr %31, align 8
+    //%32 = getelementptr inbounds nuw %struct.__tgt_kernel_arguments, ptr %kernel_args, i32 0, i32 4
+    //store ptr @.offload_sizes, ptr %32, align 8
+    //%33 = getelementptr inbounds nuw %struct.__tgt_kernel_arguments, ptr %kernel_args, i32 0, i32 5
+    //store ptr @.offload_maptypes, ptr %33, align 8
+    //%34 = getelementptr inbounds nuw %struct.__tgt_kernel_arguments, ptr %kernel_args, i32 0, i32 6
+    //store ptr null, ptr %34, align 8
+    //%35 = getelementptr inbounds nuw %struct.__tgt_kernel_arguments, ptr %kernel_args, i32 0, i32 7
+    //store ptr null, ptr %35, align 8
+    //%36 = getelementptr inbounds nuw %struct.__tgt_kernel_arguments, ptr %kernel_args, i32 0, i32 8
+    //store i64 0, ptr %36, align 8
+    //%37 = getelementptr inbounds nuw %struct.__tgt_kernel_arguments, ptr %kernel_args, i32 0, i32 9
+    //store i64 0, ptr %37, align 8
+    //%38 = getelementptr inbounds nuw %struct.__tgt_kernel_arguments, ptr %kernel_args, i32 0, i32 10
+    //store [3 x i32] [i32 2097152, i32 0, i32 0], ptr %38, align 4
+    //%39 = getelementptr inbounds nuw %struct.__tgt_kernel_arguments, ptr %kernel_args, i32 0, i32 11
+    //store [3 x i32] [i32 256, i32 0, i32 0], ptr %39, align 4
+    //%40 = getelementptr inbounds nuw %struct.__tgt_kernel_arguments, ptr %kernel_args, i32 0, i32 12
+    //store i32 0, ptr %40, align 4
     // FIXME(offload): launch kernels
+    let mut values = vec![];
+    values.push((4, cx.get_const_i32(3)));
+    values.push((4, cx.get_const_i32(num_args)));
+    values.push((8, geps.0));
+    values.push((8, geps.1));
+    values.push((8, geps.2));
+    values.push((8, o_types[0]));
+    values.push((8, cx.const_null(cx.type_ptr())));
+    values.push((8, cx.const_null(cx.type_ptr())));
+    values.push((8, cx.get_const_i64(0)));
+    values.push((8, cx.get_const_i64(0)));
+    let ti32 = cx.type_i32();
+    let ci32_0 = cx.get_const_i32(0);
+    values.push((8, cx.const_array(ti32, &vec![cx.get_const_i32(2097152), ci32_0, ci32_0])));
+    values.push((8, cx.const_array(ti32, &vec![cx.get_const_i32(256), ci32_0, ci32_0])));
+    values.push((4, cx.get_const_i32(0)));
+
+    for (i, value) in values.iter().enumerate() {
+        let ptr = builder.inbounds_gep(tgt_kernel_decl, a5, &[i32_0, cx.get_const_i32(i as u64)]);
+        builder.store(value.1, ptr, Align::from_bytes(value.0).unwrap());
+    }
+
+    let args = vec![
+        s_ident_t,
+        // MAX == -1
+        cx.get_const_i64(u64::MAX),
+        cx.get_const_i32(2097152),
+        cx.get_const_i32(256),
+        region_ids[0],
+        a5,
+    ];
+    let offload_success = builder.call(tgt_target_kernel_ty, tgt_decl, &args, None);
+    // %41 = call i32 @__tgt_target_kernel(ptr @1, i64 -1, i32 2097152, i32 256, ptr @.kernel_1.region_id, ptr %kernel_args)
+    unsafe {
+        let next = llvm::LLVMGetNextInstruction(offload_success).unwrap();
+        dbg!(&next);
+        llvm::LLVMRustPositionAfter(builder.llbuilder, next);
+        let called_kernel = llvm::LLVMGetCalledValue(next).unwrap();
+        llvm::LLVMInstructionEraseFromParent(next);
+        dbg!(&called_kernel);
+    }
 
     // Step 4)
-    unsafe { llvm::LLVMRustPositionAfter(builder.llbuilder, kernel_call) };
+    //unsafe { llvm::LLVMRustPositionAfter(builder.llbuilder, kernel_call) };
 
     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);
 
     builder.call(mapper_fn_ty, unregister_lib_decl, &[tgt_bin_desc_alloca], None);
 
+    drop(builder);
+    unsafe { llvm::LLVMDeleteFunction(called) };
+    dbg!("survived");
+
     // With this we generated the following begin and end mappers. We could easily generate the
     // update mapper in an update.
     // call void @__tgt_target_data_begin_mapper(ptr @1, i64 -1, i32 3, ptr %27, ptr %28, ptr %29, ptr @.offload_maptypes, ptr null, ptr null)

compiler/rustc_codegen_llvm/src/builder/gpu_wrapper.rs

@@ -1,12 +1,12 @@
 use std::ffi::CString;
 
 use llvm::Linkage::*;
-use rustc_abi::Align;
+use rustc_abi::{AddressSpace, Align};
 use rustc_codegen_ssa::back::write::CodegenContext;
 use rustc_codegen_ssa::traits::BaseTypeCodegenMethods;
 
 use crate::builder::gpu_offload::*;
-use crate::llvm::{self, Visibility};
+use crate::llvm::{self, Linkage, Type, Visibility};
 use crate::{LlvmCodegenBackend, ModuleLlvm, SimpleCx};
 
 pub(crate) fn create_struct_ty<'ll>(
@@ -22,6 +22,23 @@ pub(crate) fn create_struct_ty<'ll>(
     }
 }
 
+pub(crate) fn add_global_decl<'ll>(
+    cx: &SimpleCx<'ll>,
+    ty: &'ll Type,
+    name: &str,
+    l: Linkage,
+    hidden: bool,
+) -> &'ll llvm::Value {
+    let c_name = CString::new(name).unwrap();
+    let llglobal: &'ll llvm::Value = llvm::add_global(cx.llmod, ty, &c_name);
+    llvm::set_global_constant(llglobal, true);
+    llvm::set_linkage(llglobal, l);
+    if hidden {
+        llvm::set_visibility(llglobal, Visibility::Hidden);
+    }
+    llglobal
+}
+
 // We don't copy types from other functions because we generate a new module and context.
 // Bringing in types from other contexts would likely cause issues.
 pub(crate) fn gen_image_wrapper_module(cgcx: &CodegenContext<LlvmCodegenBackend>) {
@@ -32,6 +49,7 @@ pub(crate) fn gen_image_wrapper_module(cgcx: &CodegenContext<LlvmCodegenBackend>
         ModuleLlvm::new_simple(name, dl_cstr.into_raw(), target_cstr.into_raw(), &cgcx).unwrap();
     let cx = SimpleCx::new(m.llmod(), m.llcx, cgcx.pointer_size);
     let tptr = cx.type_ptr();
+    let tptr1 = cx.type_ptr_ext(AddressSpace(1));
     let ti64 = cx.type_i64();
     let ti32 = cx.type_i32();
     let ti16 = cx.type_i16();
@@ -44,28 +62,22 @@ pub(crate) fn gen_image_wrapper_module(cgcx: &CodegenContext<LlvmCodegenBackend>
     let offload_entry_ty = add_tgt_offload_entry(&cx);
     let offload_entry_arr = cx.type_array(offload_entry_ty, 0);
 
-    let c_name = CString::new("__start_omp_offloading_entries").unwrap();
-    let llglobal = llvm::add_global(cx.llmod, offload_entry_arr, &c_name);
-    llvm::set_global_constant(llglobal, true);
-    llvm::set_linkage(llglobal, ExternalLinkage);
-    llvm::set_visibility(llglobal, Visibility::Hidden);
-    let c_name = CString::new("__stop_omp_offloading_entries").unwrap();
-    let llglobal = llvm::add_global(cx.llmod, offload_entry_arr, &c_name);
-    llvm::set_global_constant(llglobal, true);
-    llvm::set_linkage(llglobal, ExternalLinkage);
-    llvm::set_visibility(llglobal, Visibility::Hidden);
+    let name = "__start_omp_offloading_entries";
+    add_global_decl(&cx, offload_entry_arr, name, ExternalLinkage, true);
+
+    let name = "__stop_omp_offloading_entries";
+    add_global_decl(&cx, offload_entry_arr, name, ExternalLinkage, true);
+
+    let name = "__dummy.omp_offloading_entries";
+    let llglobal = add_global_decl(&cx, offload_entry_arr, name, InternalLinkage, false);
 
-    let c_name = CString::new("__dummy.omp_offloading_entries").unwrap();
-    let llglobal = llvm::add_global(cx.llmod, offload_entry_arr, &c_name);
-    llvm::set_global_constant(llglobal, true);
-    llvm::set_linkage(llglobal, InternalLinkage);
     let c_section_name = CString::new("omp_offloading_entries").unwrap();
     llvm::set_section(llglobal, &c_section_name);
     let zeroinit = cx.const_null(offload_entry_arr);
     llvm::set_initializer(llglobal, zeroinit);
 
     CString::new("llvm.compiler.used").unwrap();
-    let arr_val = cx.const_array(tptr, &[llglobal]);
+    let arr_val = cx.const_array(tptr1, &[llglobal]);
     let c_section_name = CString::new("llvm.metadata").unwrap();
     let llglobal = add_global(&cx, "llvm.compiler.used", arr_val, AppendingLinkage);
     llvm::set_section(llglobal, &c_section_name);
@@ -74,30 +86,9 @@ pub(crate) fn gen_image_wrapper_module(cgcx: &CodegenContext<LlvmCodegenBackend>
     //@llvm.compiler.used = appending global [1 x ptr] [ptr @__dummy.omp_offloading_entries], section "llvm.metadata"
 
     let mapper_fn_ty = cx.type_func(&[tptr], cx.type_void());
-    crate::declare::declare_simple_fn(
-        &cx,
-        &"__tgt_unregister_lib",
-        llvm::CallConv::CCallConv,
-        llvm::UnnamedAddr::No,
-        llvm::Visibility::Default,
-        mapper_fn_ty,
-    );
-    crate::declare::declare_simple_fn(
-        &cx,
-        &"__tgt_register_lib",
-        llvm::CallConv::CCallConv,
-        llvm::UnnamedAddr::No,
-        llvm::Visibility::Default,
-        mapper_fn_ty,
-    );
-    crate::declare::declare_simple_fn(
-        &cx,
-        &"atexit",
-        llvm::CallConv::CCallConv,
-        llvm::UnnamedAddr::No,
-        llvm::Visibility::Default,
-        cx.type_func(&[tptr], ti32),
-    );
+    declare_offload_fn(&cx, &"__tgt_register_lib", mapper_fn_ty);
+    declare_offload_fn(&cx, &"__tgt_unregister_lib", mapper_fn_ty);
+    declare_offload_fn(&cx, &"atexit", cx.type_func(&[tptr], ti32));
 
     let unknown_txt = "11111111111111";
     let c_entry_name = CString::new(unknown_txt).unwrap();