@@ -18,21 +18,41 @@ pub(crate) fn handle_gpu_code<'ll>(
1818 // The offload memory transfer type for each kernel
1919 let mut o_types = vec ! [ ] ;
2020 let mut kernels = vec ! [ ] ;
21+ let mut region_ids = vec ! [ ] ;
2122 let offload_entry_ty = add_tgt_offload_entry ( & cx) ;
2223 for num in 0 ..9 {
2324 let kernel = cx. get_function ( & format ! ( "kernel_{num}" ) ) ;
2425 if let Some ( kernel) = kernel {
25- o_types. push ( gen_define_handling ( & cx, kernel, offload_entry_ty, num) ) ;
26+ let ( o, k) = gen_define_handling ( & cx, kernel, offload_entry_ty, num) ;
27+ o_types. push ( o) ;
28+ region_ids. push ( k) ;
2629 kernels. push ( kernel) ;
2730 }
2831 }
2932
30- gen_call_handling ( & cx, & kernels, & o_types) ;
33+ gen_call_handling ( & cx, & kernels, & o_types, & region_ids) ;
34+ }
35+
36+ // ; Function Attrs: nounwind
37+ // declare i32 @__tgt_target_kernel(ptr, i64, i32, i32, ptr, ptr) #2
38+ fn generate_launcher < ' ll > ( cx : & ' ll SimpleCx < ' _ > ) -> ( & ' ll llvm:: Value , & ' ll llvm:: Type ) {
39+ let tptr = cx. type_ptr ( ) ;
40+ let ti64 = cx. type_i64 ( ) ;
41+ let ti32 = cx. type_i32 ( ) ;
42+ let args = vec ! [ tptr, ti64, ti32, ti32, tptr, tptr] ;
43+ let tgt_fn_ty = cx. type_func ( & args, ti32) ;
44+ let name = "__tgt_target_kernel" ;
45+ let tgt_decl = declare_offload_fn ( & cx, name, tgt_fn_ty) ;
46+ let nounwind = llvm:: AttributeKind :: NoUnwind . create_attr ( cx. llcx ) ;
47+ attributes:: apply_to_llfn ( tgt_decl, Function , & [ nounwind] ) ;
48+ ( tgt_decl, tgt_fn_ty)
3149}
3250
3351// What is our @1 here? A magic global, used in our data_{begin/update/end}_mapper:
3452// @0 = private unnamed_addr constant [23 x i8] c";unknown;unknown;0;0;;\00", align 1
3553// @1 = private unnamed_addr constant %struct.ident_t { i32 0, i32 2, i32 0, i32 22, ptr @0 }, align 8
54+ // FIXME(offload): @0 should include the file name (e.g. lib.rs) in which the function to be
55+ // offloaded was defined.
3656fn generate_at_one < ' ll > ( cx : & ' ll SimpleCx < ' _ > ) -> & ' ll llvm:: Value {
3757 // @0 = private unnamed_addr constant [23 x i8] c";unknown;unknown;0;0;;\00", align 1
3858 let unknown_txt = ";unknown;unknown;0;0;;" ;
@@ -83,7 +103,7 @@ pub(crate) fn add_tgt_offload_entry<'ll>(cx: &'ll SimpleCx<'_>) -> &'ll llvm::Ty
83103 offload_entry_ty
84104}
85105
86- fn gen_tgt_kernel_global < ' ll > ( cx : & ' ll SimpleCx < ' _ > ) {
106+ fn gen_tgt_kernel_global < ' ll > ( cx : & ' ll SimpleCx < ' _ > ) -> & ' ll llvm :: Type {
87107 let kernel_arguments_ty = cx. type_named_struct ( "struct.__tgt_kernel_arguments" ) ;
88108 let tptr = cx. type_ptr ( ) ;
89109 let ti64 = cx. type_i64 ( ) ;
@@ -107,7 +127,7 @@ fn gen_tgt_kernel_global<'ll>(cx: &'ll SimpleCx<'_>) {
107127 // uint64_t NoWait : 1; // Was this kernel spawned with a `nowait` clause.
108128 // uint64_t IsCUDA : 1; // Was this kernel spawned via CUDA.
109129 // uint64_t Unused : 62;
110- // } Flags = {0, 0, 0};
130+ // } Flags = {0, 0, 0}; // totals to 64 Bit, 8 Byte
111131 // // The number of teams (for x,y,z dimension).
112132 // uint32_t NumTeams[3] = {0, 0, 0};
113133 // // The number of threads (for x,y,z dimension).
@@ -118,9 +138,7 @@ fn gen_tgt_kernel_global<'ll>(cx: &'ll SimpleCx<'_>) {
118138 vec ! [ ti32, ti32, tptr, tptr, tptr, tptr, tptr, tptr, ti64, ti64, tarr, tarr, ti32] ;
119139
120140 cx. set_struct_body ( kernel_arguments_ty, & kernel_elements, false ) ;
121- // For now we don't handle kernels, so for now we just add a global dummy
122- // to make sure that the __tgt_offload_entry is defined and handled correctly.
123- cx. declare_global ( "my_struct_global2" , kernel_arguments_ty) ;
141+ kernel_arguments_ty
124142}
125143
126144fn gen_tgt_data_mappers < ' ll > (
@@ -187,7 +205,7 @@ fn gen_define_handling<'ll>(
187205 kernel : & ' ll llvm:: Value ,
188206 offload_entry_ty : & ' ll llvm:: Type ,
189207 num : i64 ,
190- ) -> & ' ll llvm:: Value {
208+ ) -> ( & ' ll llvm:: Value , & ' ll llvm :: Value ) {
191209 let types = cx. func_params_types ( cx. get_type_of_global ( kernel) ) ;
192210 // It seems like non-pointer values are automatically mapped. So here, we focus on pointer (or
193211 // reference) types.
@@ -205,10 +223,14 @@ fn gen_define_handling<'ll>(
205223 // or both to and from the gpu (=3). Other values shouldn't affect us for now.
206224 // A non-mutable reference or pointer will be 1, an array that's not read, but fully overwritten
207225 // will be 2. For now, everything is 3, until we have our frontend set up.
208- let o_types =
209- add_priv_unnamed_arr ( & cx, & format ! ( ".offload_maptypes.{num}" ) , & vec ! [ 3 ; num_ptr_types] ) ;
226+ // 1+2+32: 1 (MapTo), 2 (MapFrom), 32 (Add one extra input ptr per function, to be used later).
227+ let o_types = add_priv_unnamed_arr (
228+ & cx,
229+ & format ! ( ".offload_maptypes.{num}" ) ,
230+ & vec ! [ 1 + 2 + 32 ; num_ptr_types] ,
231+ ) ;
210232 // Next: For each function, generate these three entries. A weak constant,
211- // the llvm.rodata entry name, and the omp_offloading_entries value
233+ // the llvm.rodata entry name, and the llvm_offload_entries value
212234
213235 let name = format ! ( ".kernel_{num}.region_id" ) ;
214236 let initializer = cx. get_const_i8 ( 0 ) ;
@@ -242,13 +264,13 @@ fn gen_define_handling<'ll>(
242264 llvm:: set_global_constant ( llglobal, true ) ;
243265 llvm:: set_linkage ( llglobal, WeakAnyLinkage ) ;
244266 llvm:: set_initializer ( llglobal, initializer) ;
245- llvm:: set_alignment ( llglobal, Align :: ONE ) ;
246- let c_section_name = CString :: new ( ".omp_offloading_entries " ) . unwrap ( ) ;
267+ llvm:: set_alignment ( llglobal, Align :: EIGHT ) ;
268+ let c_section_name = CString :: new ( "llvm_offload_entries " ) . unwrap ( ) ;
247269 llvm:: set_section ( llglobal, & c_section_name) ;
248- o_types
270+ ( o_types, region_id )
249271}
250272
251- fn declare_offload_fn < ' ll > (
273+ pub ( crate ) fn declare_offload_fn < ' ll > (
252274 cx : & ' ll SimpleCx < ' _ > ,
253275 name : & str ,
254276 ty : & ' ll llvm:: Type ,
@@ -287,15 +309,17 @@ fn gen_call_handling<'ll>(
287309 cx : & ' ll SimpleCx < ' _ > ,
288310 _kernels : & [ & ' ll llvm:: Value ] ,
289311 o_types : & [ & ' ll llvm:: Value ] ,
312+ region_ids : & [ & ' ll llvm:: Value ] ,
290313) {
314+ let ( tgt_decl, tgt_target_kernel_ty) = generate_launcher ( & cx) ;
291315 // %struct.__tgt_bin_desc = type { i32, ptr, ptr, ptr }
292316 let tptr = cx. type_ptr ( ) ;
293317 let ti32 = cx. type_i32 ( ) ;
294318 let tgt_bin_desc_ty = vec ! [ ti32, tptr, tptr, tptr] ;
295319 let tgt_bin_desc = cx. type_named_struct ( "struct.__tgt_bin_desc" ) ;
296320 cx. set_struct_body ( tgt_bin_desc, & tgt_bin_desc_ty, false ) ;
297321
298- gen_tgt_kernel_global ( & cx) ;
322+ let tgt_kernel_decl = gen_tgt_kernel_global ( & cx) ;
299323 let ( begin_mapper_decl, _, end_mapper_decl, fn_ty) = gen_tgt_data_mappers ( & cx) ;
300324
301325 let main_fn = cx. get_function ( "main" ) ;
@@ -329,35 +353,32 @@ fn gen_call_handling<'ll>(
329353 // These represent the sizes in bytes, e.g. the entry for `&[f64; 16]` will be 8*16.
330354 let ty2 = cx. type_array ( cx. type_i64 ( ) , num_args) ;
331355 let a4 = builder. direct_alloca ( ty2, Align :: EIGHT , ".offload_sizes" ) ;
356+
357+ //%kernel_args = alloca %struct.__tgt_kernel_arguments, align 8
358+ let a5 = builder. direct_alloca ( tgt_kernel_decl, Align :: EIGHT , "kernel_args" ) ;
359+
360+ // Step 1)
361+ unsafe { llvm:: LLVMRustPositionBefore ( builder. llbuilder , kernel_call) } ;
362+ builder. memset ( tgt_bin_desc_alloca, cx. get_const_i8 ( 0 ) , cx. get_const_i64 ( 32 ) , Align :: EIGHT ) ;
363+
332364 // Now we allocate once per function param, a copy to be passed to one of our maps.
333365 let mut vals = vec ! [ ] ;
334366 let mut geps = vec ! [ ] ;
335367 let i32_0 = cx. get_const_i32 ( 0 ) ;
336- for ( index, in_ty) in types. iter ( ) . enumerate ( ) {
337- // get function arg, store it into the alloca, and read it.
338- let p = llvm:: get_param ( called, index as u32 ) ;
339- let name = llvm:: get_value_name ( p) ;
340- let name = str:: from_utf8 ( & name) . unwrap ( ) ;
341- let arg_name = format ! ( "{name}.addr" ) ;
342- let alloca = builder. direct_alloca ( in_ty, Align :: EIGHT , & arg_name) ;
343-
344- builder. store ( p, alloca, Align :: EIGHT ) ;
345- let val = builder. load ( in_ty, alloca, Align :: EIGHT ) ;
346- let gep = builder. inbounds_gep ( cx. type_f32 ( ) , val, & [ i32_0] ) ;
347- vals. push ( val) ;
368+ for index in 0 ..types. len ( ) {
369+ let v = unsafe { llvm:: LLVMGetOperand ( kernel_call, index as u32 ) . unwrap ( ) } ;
370+ let gep = builder. inbounds_gep ( cx. type_f32 ( ) , v, & [ i32_0] ) ;
371+ vals. push ( v) ;
348372 geps. push ( gep) ;
349373 }
350374
351- // Step 1)
352- unsafe { llvm:: LLVMRustPositionBefore ( builder. llbuilder , kernel_call) } ;
353- builder. memset ( tgt_bin_desc_alloca, cx. get_const_i8 ( 0 ) , cx. get_const_i64 ( 32 ) , Align :: EIGHT ) ;
354-
355375 let mapper_fn_ty = cx. type_func ( & [ cx. type_ptr ( ) ] , cx. type_void ( ) ) ;
356376 let register_lib_decl = declare_offload_fn ( & cx, "__tgt_register_lib" , mapper_fn_ty) ;
357377 let unregister_lib_decl = declare_offload_fn ( & cx, "__tgt_unregister_lib" , mapper_fn_ty) ;
358378 let init_ty = cx. type_func ( & [ ] , cx. type_void ( ) ) ;
359379 let init_rtls_decl = declare_offload_fn ( cx, "__tgt_init_all_rtls" , init_ty) ;
360380
381+ // FIXME(offload): Later we want to add them to the wrapper code, rather than our main function.
361382 // call void @__tgt_register_lib(ptr noundef %6)
362383 builder. call ( mapper_fn_ty, register_lib_decl, & [ tgt_bin_desc_alloca] , None ) ;
363384 // call void @__tgt_init_all_rtls()
@@ -420,17 +441,58 @@ fn gen_call_handling<'ll>(
420441 generate_mapper_call ( & mut builder, & cx, geps, o, begin_mapper_decl, fn_ty, num_args, s_ident_t) ;
421442
422443 // Step 3)
423- // Here we will add code for the actual kernel launches in a follow-up PR.
424- // FIXME(offload): launch kernels
444+ let mut values = vec ! [ ] ;
445+ let offload_version = cx. get_const_i32 ( 3 ) ;
446+ values. push ( ( 4 , offload_version) ) ;
447+ values. push ( ( 4 , cx. get_const_i32 ( num_args) ) ) ;
448+ values. push ( ( 8 , geps. 0 ) ) ;
449+ values. push ( ( 8 , geps. 1 ) ) ;
450+ values. push ( ( 8 , geps. 2 ) ) ;
451+ values. push ( ( 8 , o_types[ 0 ] ) ) ;
452+ // The next two are debug infos. FIXME(offload) set them
453+ values. push ( ( 8 , cx. const_null ( cx. type_ptr ( ) ) ) ) ;
454+ values. push ( ( 8 , cx. const_null ( cx. type_ptr ( ) ) ) ) ;
455+ values. push ( ( 8 , cx. get_const_i64 ( 0 ) ) ) ;
456+ values. push ( ( 8 , cx. get_const_i64 ( 0 ) ) ) ;
457+ let ti32 = cx. type_i32 ( ) ;
458+ let ci32_0 = cx. get_const_i32 ( 0 ) ;
459+ values. push ( ( 4 , cx. const_array ( ti32, & vec ! [ cx. get_const_i32( 2097152 ) , ci32_0, ci32_0] ) ) ) ;
460+ values. push ( ( 4 , cx. const_array ( ti32, & vec ! [ cx. get_const_i32( 256 ) , ci32_0, ci32_0] ) ) ) ;
461+ values. push ( ( 4 , cx. get_const_i32 ( 0 ) ) ) ;
462+
463+ for ( i, value) in values. iter ( ) . enumerate ( ) {
464+ let ptr = builder. inbounds_gep ( tgt_kernel_decl, a5, & [ i32_0, cx. get_const_i32 ( i as u64 ) ] ) ;
465+ builder. store ( value. 1 , ptr, Align :: from_bytes ( value. 0 ) . unwrap ( ) ) ;
466+ }
467+
468+ let args = vec ! [
469+ s_ident_t,
470+ // MAX == -1
471+ cx. get_const_i64( u64 :: MAX ) ,
472+ cx. get_const_i32( 2097152 ) ,
473+ cx. get_const_i32( 256 ) ,
474+ region_ids[ 0 ] ,
475+ a5,
476+ ] ;
477+ let offload_success = builder. call ( tgt_target_kernel_ty, tgt_decl, & args, None ) ;
478+ // %41 = call i32 @__tgt_target_kernel(ptr @1, i64 -1, i32 2097152, i32 256, ptr @.kernel_1.region_id, ptr %kernel_args)
479+ unsafe {
480+ let next = llvm:: LLVMGetNextInstruction ( offload_success) . unwrap ( ) ;
481+ llvm:: LLVMRustPositionAfter ( builder. llbuilder , next) ;
482+ llvm:: LLVMInstructionEraseFromParent ( next) ;
483+ }
425484
426485 // Step 4)
427- unsafe { llvm:: LLVMRustPositionAfter ( builder. llbuilder , kernel_call) } ;
486+ // unsafe { llvm::LLVMRustPositionAfter(builder.llbuilder, kernel_call) };
428487
429488 let geps = get_geps ( & mut builder, & cx, ty, ty2, a1, a2, a4) ;
430489 generate_mapper_call ( & mut builder, & cx, geps, o, end_mapper_decl, fn_ty, num_args, s_ident_t) ;
431490
432491 builder. call ( mapper_fn_ty, unregister_lib_decl, & [ tgt_bin_desc_alloca] , None ) ;
433492
493+ drop ( builder) ;
494+ unsafe { llvm:: LLVMDeleteFunction ( called) } ;
495+
434496 // With this we generated the following begin and end mappers. We could easily generate the
435497 // update mapper in an update.
436498 // 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)
0 commit comments