diff --git a/compiler/rustc_codegen_llvm/src/back/write.rs b/compiler/rustc_codegen_llvm/src/back/write.rs index 4426f6ebb3c17..93ed54fd69fba 100644 --- a/compiler/rustc_codegen_llvm/src/back/write.rs +++ b/compiler/rustc_codegen_llvm/src/back/write.rs @@ -32,7 +32,7 @@ use crate::back::profiling::{ LlvmSelfProfiler, selfprofile_after_pass_callback, selfprofile_before_pass_callback, }; use crate::builder::SBuilder; -use crate::builder::gpu_offload::scalar_width; +use crate::builder::gpu_helper::scalar_width; use crate::common::AsCCharPtr; use crate::errors::{ CopyBitcode, FromLlvmDiag, FromLlvmOptimizationDiag, LlvmError, ParseTargetMachineConfig, diff --git a/compiler/rustc_codegen_llvm/src/builder.rs b/compiler/rustc_codegen_llvm/src/builder.rs index 134bc5006dd00..e3078b2abbe25 100644 --- a/compiler/rustc_codegen_llvm/src/builder.rs +++ b/compiler/rustc_codegen_llvm/src/builder.rs @@ -4,6 +4,7 @@ use std::ops::Deref; use rustc_ast::expand::typetree::FncTree; pub(crate) mod autodiff; +pub(crate) mod gpu_helper; pub(crate) mod gpu_offload; use libc::{c_char, c_uint}; diff --git a/compiler/rustc_codegen_llvm/src/builder/gpu_helper.rs b/compiler/rustc_codegen_llvm/src/builder/gpu_helper.rs new file mode 100644 index 0000000000000..825f76c22b806 --- /dev/null +++ b/compiler/rustc_codegen_llvm/src/builder/gpu_helper.rs @@ -0,0 +1,240 @@ +use crate::SimpleCx; +use crate::builder::Builder; +use crate::builder::gpu_offload::{OffloadGlobals, get_or_create_async_info}; +use crate::intrinsic::TransferType; +use crate::llvm; +use crate::llvm::{Type, Value}; +use rustc_abi::Align; +use rustc_codegen_ssa::MemFlags; +use rustc_codegen_ssa::common::TypeKind; +use rustc_codegen_ssa::traits::{BaseTypeCodegenMethods, BuilderMethods}; +use rustc_middle::bug; +use rustc_middle::ty::offload_meta::{OffloadMetadata, OffloadSize}; + +pub(crate) fn scalar_width<'ll>(cx: &'ll SimpleCx<'_>, ty: &'ll Type) -> u64 { + match cx.type_kind(ty) { + TypeKind::Half + | TypeKind::Float + | TypeKind::Double + | TypeKind::X86_FP80 + | TypeKind::FP128 + | TypeKind::PPC_FP128 => cx.float_width(ty) as u64, + TypeKind::Integer => cx.int_width(ty), + other => bug!("scalar_width was called on a non scalar type {other:?}"), + } +} + +fn get_runtime_size<'ll, 'tcx>( + builder: &mut Builder<'_, 'll, 'tcx>, + args: &[&'ll Value], + index: usize, + meta: &OffloadMetadata, +) -> &'ll Value { + match meta.payload_size { + OffloadSize::Slice { element_size } => { + let length_idx = index + 1; + let length = args[length_idx]; + let length_i64 = builder.intcast(length, builder.cx.type_i64(), false); + builder.mul(length_i64, builder.cx.get_const_i64(element_size)) + } + _ => bug!("unexpected offload size {:?}", meta.payload_size), + } +} + +// For now we have a very simplistic indexing scheme into our +// offload_{baseptrs,ptrs,sizes}. We will probably improve this along with our gpu frontend pr. +pub(crate) fn get_geps<'ll, 'tcx>( + builder: &mut Builder<'_, 'll, 'tcx>, + ty: &'ll Type, + ty2: &'ll Type, + a1: &'ll Value, + a2: &'ll Value, + a4: &'ll Value, + is_dynamic: bool, +) -> [&'ll Value; 3] { + let cx = builder.cx; + let i32_0 = cx.get_const_i32(0); + + let gep1 = builder.inbounds_gep(ty, a1, &[i32_0, i32_0]); + let gep2 = builder.inbounds_gep(ty, a2, &[i32_0, i32_0]); + let gep3 = if is_dynamic { builder.inbounds_gep(ty2, a4, &[i32_0, i32_0]) } else { a4 }; + [gep1, gep2, gep3] +} + +pub(crate) fn synchronize_async_info<'ll, 'tcx>( + builder: &mut Builder<'_, 'll, 'tcx>, + offload_globals: &OffloadGlobals<'ll>, +) { + let cx = builder.cx; + let ptr_ty = cx.type_ptr(); + let null = cx.const_null(ptr_ty); + + let async_info = builder.load(ptr_ty, offload_globals.async_info_global, Align::EIGHT); + + let is_null = builder.icmp(rustc_codegen_ssa::common::IntPredicate::IntEQ, async_info, null); + + let sync_bb = Builder::append_block(cx, builder.llfn(), "offload.async.sync"); + let done_bb = Builder::append_block(cx, builder.llfn(), "offload.async.sync.done"); + + builder.cond_br(is_null, done_bb, sync_bb); + + unsafe { + llvm::LLVMPositionBuilderAtEnd(&builder.llbuilder, sync_bb); + } + + builder.call( + offload_globals.async_info_synchronize_ty, + None, + None, + offload_globals.async_info_synchronize, + &[async_info], + None, + None, + ); + + builder.br(done_bb); + + unsafe { + llvm::LLVMPositionBuilderAtEnd(&builder.llbuilder, done_bb); + } +} + +pub(crate) fn generate_mapper_call<'ll, 'tcx>( + builder: &mut Builder<'_, 'll, 'tcx>, + geps: [&'ll Value; 3], + offload_globals: &OffloadGlobals<'ll>, + o_type: &'ll Value, + fn_to_call: &'ll Value, + fn_ty: &'ll Type, + num_args: u64, + s_ident_t: &'ll Value, + transfer: TransferType, +) { + let cx = builder.cx; + let nullptr = cx.const_null(cx.type_ptr()); + let i64_max = cx.get_const_i64(u64::MAX); + let num_args = cx.get_const_i32(num_args); + let mut args = + vec![s_ident_t, i64_max, num_args, geps[0], geps[1], geps[2], o_type, nullptr, nullptr]; + if matches!(transfer, TransferType::NowaitBegin) { + let i32_0 = cx.get_const_i32(0); + args.append(&mut vec![i32_0, nullptr, i32_0, nullptr]); + } + if matches!(transfer, TransferType::End) { + synchronize_async_info(builder, offload_globals); + //let a = offload_globals.taskwait; + //let b = offload_globals.taskwait_ty; + //let c = offload_globals.threadnum; + //let d = offload_globals.threadnum_ty; + //dbg!(&c); + //dbg!(&d); + //dbg!("first"); + //let tid = builder.call(d, None, None, c, &vec![s_ident_t], None, None); + //let args2 = vec![s_ident_t, tid]; + //dbg!(&a); + //dbg!(&b); + //dbg!("second"); + //builder.call(b, None, None, a, &args2, None, None); + } + builder.call(fn_ty, None, None, fn_to_call, &args, None, None); +} + +pub(crate) fn preper_datatransfers<'ll, 'tcx>( + builder: &mut Builder<'_, 'll, 'tcx>, + args: &[&'ll Value], + types: &[&Type], + offload_sizes: &'ll Value, + metadata: &[OffloadMetadata], + has_dynamic: bool, +) -> (&'ll Type, &'ll Type, &'ll Value, &'ll Value, &'ll Value) { + let cx = builder.cx; + let num_args = types.len() as u64; + let bb = builder.llbb(); + + // Step 0) + unsafe { + llvm::LLVMRustPositionBuilderPastAllocas(&builder.llbuilder, builder.llfn()); + } + + let ty = cx.type_array(cx.type_ptr(), num_args); + // Baseptr are just the input pointer to the kernel, stored in a local alloca + let a1 = builder.direct_alloca(ty, Align::EIGHT, ".offload_baseptrs"); + // Ptrs are the result of a gep into the baseptr, at least for our trivial types. + let a2 = builder.direct_alloca(ty, Align::EIGHT, ".offload_ptrs"); + // 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 = if has_dynamic { + let alloc = builder.direct_alloca(ty2, Align::EIGHT, ".offload_sizes"); + + builder.memcpy( + alloc, + Align::EIGHT, + offload_sizes, + Align::EIGHT, + cx.get_const_i64(8 * args.len() as u64), + MemFlags::empty(), + None, + ); + + alloc + } else { + offload_sizes + }; + + // Step 1) + unsafe { + llvm::LLVMPositionBuilderAtEnd(&builder.llbuilder, bb); + } + + // 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 &v in args { + let ty = cx.val_ty(v); + let ty_kind = cx.type_kind(ty); + let (base_val, gep_base) = match ty_kind { + TypeKind::Pointer => (v, v), + TypeKind::Half | TypeKind::Float | TypeKind::Double | TypeKind::Integer => { + // FIXME(Sa4dUs): check for `f128` support, latest NVIDIA cards support it + let num_bits = scalar_width(cx, ty); + + let bb = builder.llbb(); + unsafe { + llvm::LLVMRustPositionBuilderPastAllocas(builder.llbuilder, builder.llfn()); + } + let addr = builder.direct_alloca(cx.type_i64(), Align::EIGHT, "addr"); + unsafe { + llvm::LLVMPositionBuilderAtEnd(builder.llbuilder, bb); + } + + let cast = builder.bitcast(v, cx.type_ix(num_bits)); + let value = builder.zext(cast, cx.type_i64()); + builder.store(value, addr, Align::EIGHT); + (value, addr) + } + other => bug!("offload does not support {other:?}"), + }; + + let gep = builder.inbounds_gep(cx.type_f32(), gep_base, &[i32_0]); + + vals.push(base_val); + geps.push(gep); + } + + for i in 0..num_args { + let idx = cx.get_const_i32(i); + let gep1 = builder.inbounds_gep(ty, a1, &[i32_0, idx]); + builder.store(vals[i as usize], gep1, Align::EIGHT); + let gep2 = builder.inbounds_gep(ty, a2, &[i32_0, idx]); + builder.store(geps[i as usize], gep2, Align::EIGHT); + + if !matches!(metadata[i as usize].payload_size, OffloadSize::Static(_)) { + let gep3 = builder.inbounds_gep(ty2, a4, &[i32_0, idx]); + let size_val = get_runtime_size(builder, args, i as usize, &metadata[i as usize]); + builder.store(size_val, gep3, Align::EIGHT); + } + } + (ty, ty2, a1, a2, a4) +} diff --git a/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs b/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs index 0b009321802cf..84b1f42dee396 100644 --- a/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs +++ b/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs @@ -3,19 +3,23 @@ use std::ffi::CString; use bitflags::Flags; use llvm::Linkage::*; use rustc_abi::Align; -use rustc_codegen_ssa::MemFlags; -use rustc_codegen_ssa::common::TypeKind; use rustc_codegen_ssa::mir::operand::{OperandRef, OperandValue}; use rustc_codegen_ssa::traits::{BaseTypeCodegenMethods, BuilderMethods}; use rustc_middle::bug; use rustc_middle::ty::offload_meta::{MappingFlags, OffloadMetadata, OffloadSize}; use crate::builder::Builder; +use crate::builder::gpu_helper::*; use crate::common::CodegenCx; +use crate::intrinsic::TransferType; use crate::llvm::AttributePlace::Function; use crate::llvm::{self, Linkage, Type, Value}; use crate::{SimpleCx, attributes}; +use rustc_target::spec::HasTargetSpec; +//int32 __kmpc_omp_taskwait(ident_t *loc_ref, kmp_int32 gtid); +//int32 __kmpc_global_thread_num(ident_t *); + // LLVM kernel-independent globals required for offloading pub(crate) struct OffloadGlobals<'ll> { pub launcher_fn: &'ll llvm::Value, @@ -29,7 +33,27 @@ pub(crate) struct OffloadGlobals<'ll> { pub end_mapper: &'ll llvm::Value, pub mapper_fn_ty: &'ll llvm::Type, + pub nowait_begin_mapper: &'ll llvm::Value, + pub nowait_mapper_fn_ty: &'ll llvm::Type, + pub ident_t_global: &'ll llvm::Value, + + //pub taskwait: &'ll llvm::Value, + //pub taskwait_ty: &'ll llvm::Type, + //pub threadnum: &'ll llvm::Value, + //pub threadnum_ty: &'ll llvm::Type, + pub async_info_create: &'ll Value, + pub async_info_create_ty: &'ll Type, + + pub async_info_synchronize: &'ll Value, + pub async_info_synchronize_ty: &'ll Type, + + pub async_info_destroy: &'ll Value, + pub async_info_destroy_ty: &'ll Type, + + pub async_kernel_launcher: &'ll Value, + pub async_kernel_launcher_ty: &'ll Type, + pub async_info_global: &'ll llvm::Value, } impl<'ll> OffloadGlobals<'ll> { @@ -38,25 +62,150 @@ impl<'ll> OffloadGlobals<'ll> { let kernel_args_ty = KernelArgsTy::new_decl(cx); let offload_entry_ty = TgtOffloadEntry::new_decl(cx); let (begin_mapper, _, end_mapper, mapper_fn_ty) = gen_tgt_data_mappers(cx); + let (nowait_begin_mapper, nowait_mapper_fn_ty) = gen_tgt_data_nowait_mappers(cx); let ident_t_global = generate_at_one(cx); - + //let (taskwait, taskwait_ty, threadnum, threadnum_ty) = generate_sync(cx); + // ptr __tgt_async_info_create(i64) + let async_info_create_ty = cx.type_func(&[cx.type_i64()], cx.type_ptr()); + + // i32 __tgt_async_info_synchronize(ptr) + let async_info_synchronize_ty = cx.type_func(&[cx.type_ptr()], cx.type_i32()); + + // void __tgt_async_info_destroy(ptr) + let async_info_destroy_ty = cx.type_func(&[cx.type_ptr()], cx.type_void()); + + // i32 __tgt_target_kernel_async( + // ptr ident, + // i64 device, + // i32 num_teams, + // i32 thread_limit, + // ptr host_ptr, + // ptr kernel_args, + // ptr async_info) + let async_kernel_launcher_ty = cx.type_func( + &[ + cx.type_ptr(), + cx.type_i64(), + cx.type_i32(), + cx.type_i32(), + cx.type_ptr(), + cx.type_ptr(), + cx.type_ptr(), + ], + cx.type_i32(), + ); // We want LLVM's openmp-opt pass to pick up and optimize this module, since it covers both // openmp and offload optimizations. llvm::add_module_flag_u32(cx.llmod(), llvm::ModuleFlagMergeBehavior::Max, "openmp", 51); + let async_info_create = + declare_offload_fn(cx, "__tgt_async_info_create", async_info_create_ty); + + let async_info_synchronize = + declare_offload_fn(cx, "__tgt_async_info_synchronize", async_info_synchronize_ty); + + let async_info_destroy = + declare_offload_fn(cx, "__tgt_async_info_destroy", async_info_destroy_ty); + + let async_kernel_launcher = + declare_offload_fn(cx, "__tgt_target_kernel_async", async_kernel_launcher_ty); + let async_info_global = get_or_create_async_info_global(cx); OffloadGlobals { launcher_fn, launcher_ty, kernel_args_ty, offload_entry_ty, begin_mapper, + nowait_begin_mapper, end_mapper, mapper_fn_ty, + nowait_mapper_fn_ty, ident_t_global, + //taskwait, + //taskwait_ty, + //threadnum, + //threadnum_ty, + async_info_create, + async_info_create_ty, + + async_info_synchronize, + async_info_synchronize_ty, + + async_info_destroy, + async_info_destroy_ty, + + async_kernel_launcher, + async_kernel_launcher_ty, + async_info_global, } } } +fn get_or_create_async_info_global<'ll, 'tcx>(cx: &CodegenCx<'ll, 'tcx>) -> &'ll llvm::Value { + let name = c"__rust_offload_async_info"; + + if let Some(global) = unsafe { llvm::LLVMGetNamedGlobal(cx.llmod, name.as_ptr()) } { + return global; + } + + let global = unsafe { llvm::LLVMAddGlobal(cx.llmod, cx.type_ptr(), name.as_ptr()) }; + + unsafe { + llvm::LLVMSetInitializer(global, cx.const_null(cx.type_ptr())); + llvm::set_thread_local_mode(global, llvm::ThreadLocalMode::GeneralDynamic); + llvm::LLVMSetLinkage(global, llvm::Linkage::LinkOnceODRLinkage); + } + + global +} + +use rustc_codegen_ssa::common::IntPredicate; + +pub(crate) fn get_or_create_async_info<'ll, 'tcx>( + builder: &mut Builder<'_, 'll, 'tcx>, + offload_globals: &OffloadGlobals<'ll>, +) -> &'ll Value { + let cx = builder.cx; + let ptr_ty = cx.type_ptr(); + let null = cx.const_null(ptr_ty); + + let current = builder.load(ptr_ty, offload_globals.async_info_global, Align::EIGHT); + + let is_null = builder.icmp(IntPredicate::IntEQ, current, null); + + let create_bb = Builder::append_block(cx, builder.llfn(), "offload.async.create"); + let ready_bb = Builder::append_block(cx, builder.llfn(), "offload.async.ready"); + + builder.cond_br(is_null, create_bb, ready_bb); + + unsafe { + llvm::LLVMPositionBuilderAtEnd(&builder.llbuilder, create_bb); + } + + let device_id = cx.get_const_i64(u64::MAX); // -1/default device + + let created = builder.call( + offload_globals.async_info_create_ty, + None, + None, + offload_globals.async_info_create, + &[device_id], + None, + None, + ); + + builder.store(created, offload_globals.async_info_global, Align::EIGHT); + + builder.br(ready_bb); + + unsafe { + llvm::LLVMPositionBuilderAtEnd(&builder.llbuilder, ready_bb); + } + + // Reload instead of needing a phi. + builder.load(ptr_ty, offload_globals.async_info_global, Align::EIGHT) +} + // We need to register offload before using it. We also should unregister it once we are done, for // good measures. Previously we have done so before and after each individual offload intrinsic // call, but that comes at a performance cost. The repeated (un)register calls might also confuse @@ -197,6 +346,28 @@ fn generate_launcher<'ll>(cx: &CodegenCx<'ll, '_>) -> (&'ll llvm::Value, &'ll ll (tgt_decl, tgt_fn_ty) } +//fn generate_sync<'ll>( +// cx: &CodegenCx<'ll, '_>, +//) -> (&'ll llvm::Value, &'ll llvm::Type, &'ll llvm::Value, &'ll llvm::Type) { +// let tptr = cx.type_ptr(); +// let ti32 = cx.type_i32(); +// let args1 = vec![tptr, ti32]; +// let args2 = vec![tptr]; +// let fn_ty1 = cx.type_func(&args1, ti32); +// let fn_ty2 = cx.type_func(&args2, ti32); +// let name1 = "__kmpc_omp_taskwait"; +// let name2 = "__kmpc_global_thread_num"; +// let decl1 = declare_offload_fn(&cx, name1, fn_ty1); +// let decl2 = declare_offload_fn(&cx, name2, fn_ty2); +// +// let nounwind = llvm::AttributeKind::NoUnwind.create_attr(cx.llcx); +// attributes::apply_to_llfn(decl1, Function, &[nounwind]); +// attributes::apply_to_llfn(decl2, Function, &[nounwind]); +// //int32 __kmpc_omp_taskwait(ident_t *loc_ref, kmp_int32 gtid); +// //int32 __kmpc_global_thread_num(ident_t *); +// (decl1, fn_ty1, decl2, fn_ty2) +//} + // 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 @@ -348,11 +519,29 @@ impl KernelArgsTy { pub(crate) struct OffloadKernelGlobals<'ll> { pub offload_sizes: &'ll llvm::Value, pub memtransfer_begin: &'ll llvm::Value, - pub memtransfer_kernel: &'ll llvm::Value, + pub memtransfer_kernel: Option<&'ll llvm::Value>, pub memtransfer_end: &'ll llvm::Value, pub region_id: &'ll llvm::Value, } +fn gen_tgt_data_nowait_mappers<'ll>( + cx: &CodegenCx<'ll, '_>, +) -> (&'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, tptr, tptr, tptr, tptr, tptr, tptr, ti32, tptr, ti32, tptr]; + let mapper_fn_ty = cx.type_func(&args, cx.type_void()); + let mapper_begin = "__tgt_target_data_begin_nowait_mapper"; + let begin_mapper_decl = declare_offload_fn(&cx, mapper_begin, mapper_fn_ty); + + let nounwind = llvm::AttributeKind::NoUnwind.create_attr(cx.llcx); + attributes::apply_to_llfn(begin_mapper_decl, Function, &[nounwind]); + + (begin_mapper_decl, mapper_fn_ty) +} + fn gen_tgt_data_mappers<'ll>( cx: &CodegenCx<'ll, '_>, ) -> (&'ll llvm::Value, &'ll llvm::Value, &'ll llvm::Value, &'ll llvm::Type) { @@ -420,12 +609,14 @@ pub(crate) fn gen_define_handling<'ll>( metadata: &[OffloadMetadata], symbol: String, offload_globals: &OffloadGlobals<'ll>, + transfer: TransferType, ) -> OffloadKernelGlobals<'ll> { if let Some(entry) = cx.offload_kernel_cache.borrow().get(&symbol) { return *entry; } let offload_entry_ty = offload_globals.offload_entry_ty; + let gen_kernel = matches!(transfer, TransferType::Kernel); let (sizes, transfer): (Vec<_>, Vec<_>) = metadata.iter().map(|m| (m.payload_size, m.mode)).unzip(); @@ -447,6 +638,7 @@ pub(crate) fn gen_define_handling<'ll>( let valid_begin_mappings = MappingFlags::TO | MappingFlags::LITERAL | MappingFlags::IMPLICIT; let transfer_to: Vec = transfer.iter().map(|m| m.intersection(valid_begin_mappings).bits()).collect(); + //dbg!(&transfer); let transfer_from: Vec = transfer.iter().map(|m| m.intersection(MappingFlags::FROM).bits()).collect(); let valid_kernel_mappings = MappingFlags::LITERAL | MappingFlags::IMPLICIT; @@ -468,8 +660,16 @@ pub(crate) fn gen_define_handling<'ll>( add_priv_unnamed_arr(&cx, &format!(".offload_sizes.{symbol}"), &actual_sizes); let memtransfer_begin = add_priv_unnamed_arr(&cx, &format!(".offload_maptypes.{symbol}.begin"), &transfer_to); - let memtransfer_kernel = - add_priv_unnamed_arr(&cx, &format!(".offload_maptypes.{symbol}.kernel"), &transfer_kernel); + + let memtransfer_kernel = if gen_kernel { + Some(add_priv_unnamed_arr( + &cx, + &format!(".offload_maptypes.{symbol}.kernel"), + &transfer_kernel, + )) + } else { + None + }; let memtransfer_end = add_priv_unnamed_arr(&cx, &format!(".offload_maptypes.{symbol}.end"), &transfer_from); @@ -480,31 +680,33 @@ pub(crate) fn gen_define_handling<'ll>( let initializer = cx.get_const_i8(0); let region_id = add_global(&cx, &name, initializer, WeakAnyLinkage); - let c_entry_name = CString::new(symbol.clone()).unwrap(); - let c_val = c_entry_name.as_bytes_with_nul(); - let offload_entry_name = format!(".offloading.entry_name.{symbol}"); + if gen_kernel { + let c_entry_name = CString::new(symbol.clone()).unwrap(); + let c_val = c_entry_name.as_bytes_with_nul(); + let offload_entry_name = format!(".offloading.entry_name.{symbol}"); - let initializer = crate::common::bytes_in_context(cx.llcx, c_val); - let llglobal = add_unnamed_global(&cx, &offload_entry_name, initializer, InternalLinkage); - llvm::set_alignment(llglobal, Align::ONE); - llvm::set_section(llglobal, c".llvm.rodata.offloading"); + let initializer = crate::common::bytes_in_context(cx.llcx, c_val); + let llglobal = add_unnamed_global(&cx, &offload_entry_name, initializer, InternalLinkage); + llvm::set_alignment(llglobal, Align::ONE); + llvm::set_section(llglobal, c".llvm.rodata.offloading"); - let name = format!(".offloading.entry.{symbol}"); + let name = format!(".offloading.entry.{symbol}"); - // See the __tgt_offload_entry documentation above. - let elems = TgtOffloadEntry::new(&cx, region_id, llglobal); + // See the __tgt_offload_entry documentation above. + let elems = TgtOffloadEntry::new(&cx, region_id, llglobal); - let initializer = crate::common::named_struct(offload_entry_ty, &elems); - let c_name = CString::new(name).unwrap(); - let offload_entry = llvm::add_global(cx.llmod, offload_entry_ty, &c_name); - llvm::set_global_constant(offload_entry, true); - llvm::set_linkage(offload_entry, WeakAnyLinkage); - llvm::set_initializer(offload_entry, initializer); - llvm::set_alignment(offload_entry, Align::EIGHT); - let c_section_name = CString::new("llvm_offload_entries").unwrap(); - llvm::set_section(offload_entry, &c_section_name); + let initializer = crate::common::named_struct(offload_entry_ty, &elems); + let c_name = CString::new(name).unwrap(); + let offload_entry = llvm::add_global(cx.llmod, offload_entry_ty, &c_name); + llvm::set_global_constant(offload_entry, true); + llvm::set_linkage(offload_entry, WeakAnyLinkage); + llvm::set_initializer(offload_entry, initializer); + llvm::set_alignment(offload_entry, Align::EIGHT); + let c_section_name = CString::new("llvm_offload_entries").unwrap(); + llvm::set_section(offload_entry, &c_section_name); - cx.add_compiler_used_global(offload_entry); + cx.add_compiler_used_global(offload_entry); + } let result = OffloadKernelGlobals { offload_sizes, @@ -534,36 +736,6 @@ fn declare_offload_fn<'ll>( ) } -pub(crate) fn scalar_width<'ll>(cx: &'ll SimpleCx<'_>, ty: &'ll Type) -> u64 { - match cx.type_kind(ty) { - TypeKind::Half - | TypeKind::Float - | TypeKind::Double - | TypeKind::X86_FP80 - | TypeKind::FP128 - | TypeKind::PPC_FP128 => cx.float_width(ty) as u64, - TypeKind::Integer => cx.int_width(ty), - other => bug!("scalar_width was called on a non scalar type {other:?}"), - } -} - -fn get_runtime_size<'ll, 'tcx>( - builder: &mut Builder<'_, 'll, 'tcx>, - args: &[&'ll Value], - index: usize, - meta: &OffloadMetadata, -) -> &'ll Value { - match meta.payload_size { - OffloadSize::Slice { element_size } => { - let length_idx = index + 1; - let length = args[length_idx]; - let length_i64 = builder.intcast(length, builder.cx.type_i64(), false); - builder.mul(length_i64, builder.cx.get_const_i64(element_size)) - } - _ => bug!("unexpected offload size {:?}", meta.payload_size), - } -} - // For each kernel *call*, we now use some of our previous declared globals to move data to and from // the gpu. For now, we only handle the data transfer part of it. // If two consecutive kernels use the same memory, we still move it to the host and back to the gpu. @@ -600,161 +772,50 @@ pub(crate) fn gen_call_handling<'ll, 'tcx>( memtransfer_end, region_id, } = offload_data; + let memtransfer_kernel = memtransfer_kernel.unwrap(); let OffloadKernelDims { num_workgroups, threads_per_block, workgroup_dims, thread_dims } = offload_dims; let has_dynamic = metadata.iter().any(|m| !matches!(m.payload_size, OffloadSize::Static(_))); - let tgt_decl = offload_globals.launcher_fn; - let tgt_target_kernel_ty = offload_globals.launcher_ty; + let tgt_decl = offload_globals.async_kernel_launcher; //launcher_fn + //let tgt_target_kernel_ty = offload_globals.launcher_ty; + let tgt_target_kernel_ty = offload_globals.async_kernel_launcher_ty; let tgt_kernel_decl = offload_globals.kernel_args_ty; let begin_mapper_decl = offload_globals.begin_mapper; let end_mapper_decl = offload_globals.end_mapper; let fn_ty = offload_globals.mapper_fn_ty; + let (ty, ty2, a1, a2, a4) = + preper_datatransfers(builder, args, types, offload_sizes, metadata, has_dynamic); let num_args = types.len() as u64; - let bb = builder.llbb(); + assert_eq!(num_args as usize, args.len()); - // Step 0) + let bb = builder.llbb(); unsafe { llvm::LLVMRustPositionBuilderPastAllocas(&builder.llbuilder, builder.llfn()); } - - let ty = cx.type_array(cx.type_ptr(), num_args); - // Baseptr are just the input pointer to the kernel, stored in a local alloca - let a1 = builder.direct_alloca(ty, Align::EIGHT, ".offload_baseptrs"); - // Ptrs are the result of a gep into the baseptr, at least for our trivial types. - let a2 = builder.direct_alloca(ty, Align::EIGHT, ".offload_ptrs"); - // 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 = if has_dynamic { - let alloc = builder.direct_alloca(ty2, Align::EIGHT, ".offload_sizes"); - - builder.memcpy( - alloc, - Align::EIGHT, - offload_sizes, - Align::EIGHT, - cx.get_const_i64(8 * args.len() as u64), - MemFlags::empty(), - None, - ); - - alloc - } else { - 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::LLVMPositionBuilderAtEnd(&builder.llbuilder, bb); } - // 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 &v in args { - let ty = cx.val_ty(v); - let ty_kind = cx.type_kind(ty); - let (base_val, gep_base) = match ty_kind { - TypeKind::Pointer => (v, v), - TypeKind::Half | TypeKind::Float | TypeKind::Double | TypeKind::Integer => { - // FIXME(Sa4dUs): check for `f128` support, latest NVIDIA cards support it - let num_bits = scalar_width(cx, ty); - - let bb = builder.llbb(); - unsafe { - llvm::LLVMRustPositionBuilderPastAllocas(builder.llbuilder, builder.llfn()); - } - let addr = builder.direct_alloca(cx.type_i64(), Align::EIGHT, "addr"); - unsafe { - llvm::LLVMPositionBuilderAtEnd(builder.llbuilder, bb); - } - - let cast = builder.bitcast(v, cx.type_ix(num_bits)); - let value = builder.zext(cast, cx.type_i64()); - builder.store(value, addr, Align::EIGHT); - (value, addr) - } - other => bug!("offload does not support {other:?}"), - }; - - let gep = builder.inbounds_gep(cx.type_f32(), gep_base, &[i32_0]); - - vals.push(base_val); - geps.push(gep); - } - - for i in 0..num_args { - let idx = cx.get_const_i32(i); - let gep1 = builder.inbounds_gep(ty, a1, &[i32_0, idx]); - builder.store(vals[i as usize], gep1, Align::EIGHT); - let gep2 = builder.inbounds_gep(ty, a2, &[i32_0, idx]); - builder.store(geps[i as usize], gep2, Align::EIGHT); - - if !matches!(metadata[i as usize].payload_size, OffloadSize::Static(_)) { - let gep3 = builder.inbounds_gep(ty2, a4, &[i32_0, idx]); - let size_val = get_runtime_size(builder, args, i as usize, &metadata[i as usize]); - builder.store(size_val, gep3, Align::EIGHT); - } - } - - // For now we have a very simplistic indexing scheme into our - // offload_{baseptrs,ptrs,sizes}. We will probably improve this along with our gpu frontend pr. - fn get_geps<'ll, 'tcx>( - builder: &mut Builder<'_, 'll, 'tcx>, - ty: &'ll Type, - ty2: &'ll Type, - a1: &'ll Value, - a2: &'ll Value, - a4: &'ll Value, - is_dynamic: bool, - ) -> [&'ll Value; 3] { - let cx = builder.cx; - let i32_0 = cx.get_const_i32(0); - - let gep1 = builder.inbounds_gep(ty, a1, &[i32_0, i32_0]); - let gep2 = builder.inbounds_gep(ty, a2, &[i32_0, i32_0]); - let gep3 = if is_dynamic { builder.inbounds_gep(ty2, a4, &[i32_0, i32_0]) } else { a4 }; - [gep1, gep2, gep3] - } - - fn generate_mapper_call<'ll, 'tcx>( - builder: &mut Builder<'_, 'll, 'tcx>, - geps: [&'ll Value; 3], - o_type: &'ll Value, - fn_to_call: &'ll Value, - fn_ty: &'ll Type, - num_args: u64, - s_ident_t: &'ll Value, - ) { - let cx = builder.cx; - let nullptr = cx.const_null(cx.type_ptr()); - let i64_max = cx.get_const_i64(u64::MAX); - let num_args = cx.get_const_i32(num_args); - let args = - vec![s_ident_t, i64_max, num_args, geps[0], geps[1], geps[2], o_type, nullptr, nullptr]; - builder.call(fn_ty, None, None, fn_to_call, &args, None, None); - } - // Step 2) let s_ident_t = offload_globals.ident_t_global; let geps = get_geps(builder, ty, ty2, a1, a2, a4, has_dynamic); - generate_mapper_call( - builder, - geps, - memtransfer_begin, - begin_mapper_decl, - fn_ty, - num_args, - s_ident_t, - ); + //generate_mapper_call( + // builder, + // geps, + // offload_globals, + // memtransfer_begin, + // begin_mapper_decl, + // fn_ty, + // num_args, + // s_ident_t, + // TransferType::Kernel, + //); let values = KernelArgsTy::new( &cx, num_args, @@ -767,6 +828,7 @@ pub(crate) fn gen_call_handling<'ll, 'tcx>( // Step 3) // Here we fill the KernelArgsTy, see the documentation above + let i32_0 = 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)]); let name = std::ffi::CString::new(value.1).unwrap(); @@ -774,7 +836,7 @@ pub(crate) fn gen_call_handling<'ll, 'tcx>( builder.store(value.2, ptr, value.0); } - + let async_info = get_or_create_async_info(builder, offload_globals); let args = vec![ s_ident_t, // FIXME(offload) give users a way to select which GPU to use. @@ -783,19 +845,22 @@ pub(crate) fn gen_call_handling<'ll, 'tcx>( threads_per_block, region_id, a5, + async_info, ]; builder.call(tgt_target_kernel_ty, None, None, tgt_decl, &args, None, None); // %41 = call i32 @__tgt_target_kernel(ptr @1, i64 -1, i32 2097152, i32 256, ptr @.kernel_1.region_id, ptr %kernel_args) // Step 4) let geps = get_geps(builder, ty, ty2, a1, a2, a4, has_dynamic); - generate_mapper_call( - builder, - geps, - memtransfer_end, - end_mapper_decl, - fn_ty, - num_args, - s_ident_t, - ); + //generate_mapper_call( + // builder, + // geps, + // offload_globals, + // memtransfer_end, + // end_mapper_decl, + // fn_ty, + // num_args, + // s_ident_t, + // TransferType::Kernel, + //); } diff --git a/compiler/rustc_codegen_llvm/src/intrinsic.rs b/compiler/rustc_codegen_llvm/src/intrinsic.rs index 1c7b415fd04c7..781c941d5c10e 100644 --- a/compiler/rustc_codegen_llvm/src/intrinsic.rs +++ b/compiler/rustc_codegen_llvm/src/intrinsic.rs @@ -19,7 +19,7 @@ use rustc_hir::def_id::LOCAL_CRATE; use rustc_hir::find_attr; use rustc_middle::mir::BinOp; use rustc_middle::ty::layout::{FnAbiOf, HasTyCtxt, HasTypingEnv, LayoutOf}; -use rustc_middle::ty::offload_meta::OffloadMetadata; +use rustc_middle::ty::offload_meta::{MappingFlags, OffloadMetadata}; use rustc_middle::ty::{self, GenericArgsRef, Instance, SimdAlign, Ty, TyCtxt, TypingEnv}; use rustc_middle::{bug, span_bug}; use rustc_session::config::CrateType; @@ -34,15 +34,14 @@ use tracing::debug; use crate::abi::FnAbiLlvmExt; use crate::builder::Builder; use crate::builder::autodiff::{adjust_activity_to_abi, generate_enzyme_call}; -use crate::builder::gpu_offload::{ - OffloadKernelDims, gen_call_handling, gen_define_handling, register_offload, -}; +use crate::builder::gpu_offload::*; use crate::context::CodegenCx; use crate::declare::declare_raw_fn; use crate::errors::{ AutoDiffWithoutEnable, AutoDiffWithoutLto, IntrinsicSignatureMismatch, IntrinsicWrongArch, OffloadWithoutEnable, OffloadWithoutFatLTO, UnknownIntrinsic, }; +use crate::intrinsic::ty::offload_meta::OffloadSize; use crate::llvm::{self, Type, Value}; use crate::type_of::LayoutLlvmExt; use crate::va_arg::emit_va_arg; @@ -230,14 +229,28 @@ impl<'ll, 'tcx> IntrinsicCallBuilderMethods<'tcx> for Builder<'_, 'll, 'tcx> { codegen_autodiff(self, tcx, instance, args, result); return IntrinsicResult::WroteIntoPlace; } + sym::offload_sync => { + codegen_offload_sync(self); + return IntrinsicResult::WroteIntoPlace; + } + + sym::offload_preload => { + codegen_offload_preload(self, tcx, instance, args); + return IntrinsicResult::WroteIntoPlace; + } + + sym::offload_preload_end => { + codegen_offload_preload_drop(self, tcx, instance, args); + return IntrinsicResult::WroteIntoPlace; + } sym::offload => { if tcx.sess.opts.unstable_opts.offload.is_empty() { let _ = tcx.dcx().emit_almost_fatal(OffloadWithoutEnable); } - if tcx.sess.lto() != rustc_session::config::Lto::Fat { - let _ = tcx.dcx().emit_almost_fatal(OffloadWithoutFatLTO); - } + //if tcx.sess.lto() != rustc_session::config::Lto::Fat { + // let _ = tcx.dcx().emit_almost_fatal(OffloadWithoutFatLTO); + //} codegen_offload(self, tcx, instance, args); // offload *has* a return type, but somehow works without mentioning the place @@ -1894,6 +1907,205 @@ fn codegen_autodiff<'ll, 'tcx>( ); } +fn offload_bool_arg<'ll, 'tcx>(args: &[OperandRef<'tcx, &'ll llvm::Value>], idx: usize) -> bool { + let arg = &args[idx]; + + if !arg.layout.ty.is_bool() { + bug!("expected bool argument at index {idx}, got {:?}", arg.layout.ty); + } + + let OperandValue::Immediate(v) = arg.val else { + bug!("expected immediate bool argument at index {idx}"); + }; + + let Some(ci) = (unsafe { llvm::LLVMIsAConstantInt(v) }) else { + bug!("expected constant bool argument at index {idx}"); + }; + + let mut raw = 0u64; + let ok = unsafe { llvm::LLVMRustConstIntGetZExtValue(ci, &mut raw) }; + + if !ok { + bug!("failed to extract constant bool argument at index {idx}"); + } + + raw != 0 +} +fn codegen_offload_sync<'ll, 'tcx>(bx: &mut Builder<'_, 'll, 'tcx>) { + let cx = bx.cx; + + //register_offload(cx); + + let offload_globals_ref = cx.offload_globals.borrow(); + let offload_globals = match offload_globals_ref.as_ref() { + Some(globals) => globals, + None => bug!("offload globals were not initialized"), + }; + + crate::builder::gpu_helper::synchronize_async_info(bx, offload_globals); +} +fn codegen_offload_preload_drop<'ll, 'tcx>( + bx: &mut Builder<'_, 'll, 'tcx>, + tcx: TyCtxt<'tcx>, + _instance: ty::Instance<'tcx>, + args: &[OperandRef<'tcx, &'ll llvm::Value>], +) { + let cx = bx.cx; + dbg!("Starting the PreloadMut drop handling!"); + // PreloadMut<'a, T> -> extract T. + let ptr_arg = &args[0]; + let is_mut: bool = offload_bool_arg(args, 1); + + let pointee_ty = match *ptr_arg.layout.ty.kind() { + ty::RawPtr(pointee_ty, _) => pointee_ty, + _ => bug!("expected raw pointer argument"), + }; + + let ptr = match ptr_arg.val { + OperandValue::Immediate(ptr) => ptr, + _ => bug!("not handled"), + }; + + let args = vec![ptr]; + + let mut meta = OffloadMetadata::from_ty(tcx, pointee_ty); + // We end a mut Mapper. Unless the user never mutated a mut variable passed in a mutable way, we + // must return it from the device to update the host version. If they never mutated it, they + // surely got a clippy or rustc warning, so it's up to them for wasting time. + if is_mut { + meta.mode |= MappingFlags::FROM; + } else { + // We still want the refcounter to go down, so the runtime nows when it can free the data. + meta.mode |= MappingFlags::NONE; + } + dbg!(&meta); + let metadata: &[OffloadMetadata; 1] = &[meta]; + + let types: &Type = cx.layout_of(pointee_ty).llvm_type(cx); + + let offload_globals_ref = cx.offload_globals.borrow(); + let offload_globals = match offload_globals_ref.as_ref() { + Some(globals) => globals, + None => { + dbg!("Have to initialize offload? This is a bug!"); + return; + } + }; + + let target_symbol = cx.generate_local_symbol_name(""); + dbg!("done for now"); + let offload_data = + gen_define_handling(&cx, metadata, target_symbol, offload_globals, TransferType::End); + let has_dynamic = metadata.iter().any(|m| !matches!(m.payload_size, OffloadSize::Static(_))); + let (ty, ty2, a1, a2, a4) = crate::builder::gpu_helper::preper_datatransfers( + bx, + &args, + &[types], + offload_data.offload_sizes, + metadata, + has_dynamic, + ); + let geps = crate::builder::gpu_helper::get_geps(bx, ty, ty2, a1, a2, a4, has_dynamic); + + crate::builder::gpu_helper::generate_mapper_call( + bx, + geps, + offload_globals, + offload_data.memtransfer_end, + offload_globals.end_mapper, + offload_globals.mapper_fn_ty, + 1, + offload_globals.ident_t_global, + TransferType::End, + ); +} + +// For each PreLoad *call*, we now use some of our previous declared globals to move data to the gpu. +// For now, we only handle the data transfer part of it. Consecutive calls become a no-op on the +// LLVM side. +// +// Current steps: +// 0. Alloca some variables for the following steps +// 1. set insert point before PreLoad call. +// 2. generate all the GEPS and stores, to be used in 3) +// 3. generate __tgt_target_data_begin calls to move data to the GPU +// +// unchanged: keep kernel call. Later move the kernel to the GPU +// +// 4. set insert point after kernel call. +// 5. generate all the GEPS and stores, to be used in 6) +// 6. generate __tgt_target_data_end calls to move data from the GPU +fn codegen_offload_preload<'ll, 'tcx>( + bx: &mut Builder<'_, 'll, 'tcx>, + tcx: TyCtxt<'tcx>, + _instance: ty::Instance<'tcx>, + args: &[OperandRef<'tcx, &'ll Value>], +) { + dbg!("Starting the preload handling!"); + let cx = bx.cx; + register_offload(cx); + + let arg: &OperandRef<'_, &'ll Value> = &args[0]; + let args = match arg.val { + OperandValue::Immediate(val) => vec![val], + _ => bug!("not yet handled"), + }; + + let arg_ty = arg.layout.ty; + + let pointee_ty: Ty<'tcx> = match *arg_ty.kind() { + ty::RawPtr(pointee_ty, _) => pointee_ty, + _ => bug!("expected preload argument to be a raw pointer, got {arg_ty:?}"), + }; + + let meta = OffloadMetadata::from_ty(tcx, pointee_ty); + let metadata = &[meta]; + let types = cx.layout_of(pointee_ty).llvm_type(cx); + + let offload_globals_ref = cx.offload_globals.borrow(); + let offload_globals = match offload_globals_ref.as_ref() { + Some(globals) => globals, + None => { + dbg!("Have to initialize offload? This is a bug!"); + // Offload is not initialized, cannot continue + return; + } + }; + dbg!("asdf"); + let target_symbol = cx.generate_local_symbol_name(""); + let offload_data = + gen_define_handling(&cx, metadata, target_symbol, offload_globals, TransferType::Begin); + let has_dynamic = metadata.iter().any(|m| !matches!(m.payload_size, OffloadSize::Static(_))); + let (ty, ty2, a1, a2, a4) = crate::builder::gpu_helper::preper_datatransfers( + bx, + &args, + &[types], + offload_data.offload_sizes, + metadata, + has_dynamic, + ); + let geps = crate::builder::gpu_helper::get_geps(bx, ty, ty2, a1, a2, a4, has_dynamic); + + crate::builder::gpu_helper::generate_mapper_call( + bx, + geps, + offload_globals, + offload_data.memtransfer_begin, + offload_globals.begin_mapper, + offload_globals.mapper_fn_ty, + 1, + offload_globals.ident_t_global, + TransferType::Begin, + ); +} + +pub(crate) enum TransferType { + NowaitBegin, + Begin, + Kernel, + End, +} + // 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 @@ -1905,6 +2117,7 @@ fn codegen_offload<'ll, 'tcx>( args: &[OperandRef<'tcx, &'ll Value>], ) { let cx = bx.cx; + register_offload(cx); let fn_args = instance.args; let (target_id, target_args) = match fn_args.into_type_list(tcx)[0].kind() { @@ -1960,8 +2173,8 @@ fn codegen_offload<'ll, 'tcx>( return; } }; - register_offload(cx); - let offload_data = gen_define_handling(&cx, &metadata, target_symbol, offload_globals); + let offload_data = + gen_define_handling(&cx, &metadata, target_symbol, offload_globals, TransferType::Kernel); gen_call_handling( bx, &offload_data, diff --git a/compiler/rustc_codegen_ssa/src/mir/intrinsic.rs b/compiler/rustc_codegen_ssa/src/mir/intrinsic.rs index ac6c6a0a52efa..a6bdedc7877e7 100644 --- a/compiler/rustc_codegen_ssa/src/mir/intrinsic.rs +++ b/compiler/rustc_codegen_ssa/src/mir/intrinsic.rs @@ -124,6 +124,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> { | sym::atomic_fence | sym::atomic_singlethreadfence | sym::caller_location + | sym::offload_sync | sym::return_address => {} _ => { span_bug!( diff --git a/compiler/rustc_hir/src/lang_items.rs b/compiler/rustc_hir/src/lang_items.rs index 4a3615e5421fe..6780fd4e77e39 100644 --- a/compiler/rustc_hir/src/lang_items.rs +++ b/compiler/rustc_hir/src/lang_items.rs @@ -325,6 +325,12 @@ language_item_table! { DropGlue, sym::drop_glue, drop_glue_fn, Target::Fn, GenericRequirement::Exact(1); AllocLayout, sym::alloc_layout, alloc_layout, Target::Struct, GenericRequirement::None; + // Compiler-generated mapper functions for gpu offloading + PreloadStruct, sym::preload_type, preload_type, Target::Struct, GenericRequirement::None; + PreloadMutStruct, sym::preload_mut_type, preload_mut_type, Target::Struct, GenericRequirement::None; + PreloadFn, sym::preload, preload_fn, Target::Fn, GenericRequirement::None; + PreloadMutFn, sym::preload_mut, preload_mut_fn, Target::Fn, GenericRequirement::None; + /// For all binary crates without `#![no_main]`, Rust will generate a "main" function. /// The exact name and signature are target-dependent. The "main" function will invoke /// this lang item, passing it the `argc` and `argv` (or null, if those don't exist diff --git a/compiler/rustc_hir_analysis/src/check/intrinsic.rs b/compiler/rustc_hir_analysis/src/check/intrinsic.rs index a4d6056e47bdc..ee2b872e3ad0c 100644 --- a/compiler/rustc_hir_analysis/src/check/intrinsic.rs +++ b/compiler/rustc_hir_analysis/src/check/intrinsic.rs @@ -164,6 +164,9 @@ fn intrinsic_operation_unsafety(tcx: TyCtxt<'_>, intrinsic_id: LocalDefId) -> hi | sym::mul_with_overflow | sym::needs_drop | sym::offload + | sym::offload_sync + | sym::offload_preload + | sym::offload_preload_end | sym::offset_of | sym::overflow_checks | sym::powf16 @@ -361,6 +364,10 @@ pub(crate) fn check_intrinsic_type( ], param(2), ), + sym::offload_preload | sym::offload_preload_end => { + (1, 0, vec![Ty::new_imm_ptr(tcx, param(0)), tcx.types.bool], tcx.types.unit) + } + sym::offload_sync => (0, 0, vec![], tcx.types.unit), sym::offset => (2, 0, vec![param(0), param(1)], param(0)), sym::arith_offset => ( 1, diff --git a/compiler/rustc_span/src/symbol.rs b/compiler/rustc_span/src/symbol.rs index 7263680c302f1..73d61097eaa81 100644 --- a/compiler/rustc_span/src/symbol.rs +++ b/compiler/rustc_span/src/symbol.rs @@ -1425,6 +1425,9 @@ symbols! { of, off, offload, + offload_sync, + offload_preload, + offload_preload_end, offload_kernel, offset, offset_of, @@ -1563,6 +1566,10 @@ symbols! { prefetch_write_instruction, prefix_nops, preg, + preload, + preload_mut, + preload_type, + preload_mut_type, prelude, prelude_import, preserves_flags, diff --git a/library/core/src/intrinsics/mod.rs b/library/core/src/intrinsics/mod.rs index 78d7314c58110..8de3ffa459679 100644 --- a/library/core/src/intrinsics/mod.rs +++ b/library/core/src/intrinsics/mod.rs @@ -3573,6 +3573,19 @@ pub const fn offload( args: T, ) -> R; +#[rustc_intrinsic] +#[rustc_nounwind] +pub fn offload_preload(ptr: *const T, is_mut: bool); + +/// Waits for previously submitted offload operations to complete. +#[rustc_intrinsic] +#[rustc_nounwind] +pub fn offload_sync(); + +#[rustc_intrinsic] +#[rustc_nounwind] +pub fn offload_preload_end(ptr: *const T, is_mut: bool); + /// Inform Miri that a given pointer definitely has a certain alignment. #[cfg(miri)] #[rustc_allow_const_fn_unstable(const_eval_select)] diff --git a/library/core/src/offload/mod.rs b/library/core/src/offload/mod.rs index 164bbb9f0047d..5f49e298acb75 100644 --- a/library/core/src/offload/mod.rs +++ b/library/core/src/offload/mod.rs @@ -3,3 +3,78 @@ pub use crate::macros::builtin::offload_kernel; #[unstable(feature = "gpu_offload", issue = "131513")] pub use crate::offload; + +use crate::marker::PhantomData; + +// We store a raw pointer instead of a reference, since the real location of the data will be on a +// GPU, at a different address. We only use the CPU pointer as a key to our runtime cpu-gpu pointer +// map. In the future we might even directly store the gpu ptr here, which would make it even +// clearer why we are using raw pointers instead of references. +// We still use a lifetime marker to prevent writes into the original cpu version of the object, +// while the data is on the gpu. Dropping this struct will inform the runtime that this pointer can +// no longer be used to access the gpu copy of the data. If the reference counter reaches zero, the +// runtime might delete the gpu copy of the preloaded value. +#[lang = "preload_type"] +#[unstable(feature = "offload", issue = "124509")] +pub struct Preload<'a, T: ?Sized> { + pub cpu_ptr: *const T, + _marker: PhantomData<&'a T>, +} + +/// Waits until all previously submitted offload kernels on the current +/// host thread have completed. +/// +/// This does not copy preloaded values back to the host and does not release +/// their device mappings. Those operations still occur when the corresponding +/// [`Preload`] or [`PreloadMut`] guard is dropped. +#[inline(always)] +#[unstable(feature = "offload", issue = "124509")] +pub fn offload_sync() { + crate::intrinsics::offload_sync(); +} + +// We store a raw pointer instead of a reference, since the real location of the data will be on a +// GPU, at a different address. We only use the CPU pointer as a key to our runtime cpu-gpu pointer +// map. In the future we might even directly store the gpu ptr here, which would make it even +// clearer why we are using raw pointers instead of references. +// We still use a lifetime marker to prevent writes into the original cpu version of the object, +// while the data is on the gpu. Dropping this struct will force a copy of the data back from the +// gpu to the cpu, after which we can again safely use the original mutable reference. +#[lang = "preload_mut_type"] +#[unstable(feature = "offload", issue = "124509")] +pub struct PreloadMut<'a, T: ?Sized> { + pub cpu_ptr: *mut T, + _marker: PhantomData<&'a mut T>, +} + +#[lang = "preload"] +#[unstable(feature = "offload", issue = "124509")] +pub fn preload<'a, T: ?Sized>(x: &'a T) -> Preload<'a, T> { + let p = Preload { cpu_ptr: x as *const T, _marker: PhantomData }; + + core::intrinsics::offload_preload(p.cpu_ptr, false); + + p +} + +#[lang = "preload_mut"] +#[unstable(feature = "offload", issue = "124509")] +pub fn preload_mut<'a, T: ?Sized>(x: &'a mut T) -> PreloadMut<'a, T> { + let p = PreloadMut { cpu_ptr: x as *mut T, _marker: PhantomData }; + + core::intrinsics::offload_preload(p.cpu_ptr, true); + + p +} + +impl Drop for PreloadMut<'_, T> { + fn drop(&mut self) { + core::intrinsics::offload_preload_end(self.cpu_ptr, true); + } +} + +impl Drop for Preload<'_, T> { + fn drop(&mut self) { + core::intrinsics::offload_preload_end(self.cpu_ptr, false); + } +} diff --git a/src/llvm-project b/src/llvm-project index 08c84e69a84d9..35ab97bae2da4 160000 --- a/src/llvm-project +++ b/src/llvm-project @@ -1 +1 @@ -Subproject commit 08c84e69a84d95936296dfcab0e38b34100725d5 +Subproject commit 35ab97bae2da45a7a773f9589dfbc66369ee6060 diff --git a/tests/codegen-llvm/gpu_offload/explicit_memtransfer.rs b/tests/codegen-llvm/gpu_offload/explicit_memtransfer.rs new file mode 100644 index 0000000000000..dd2a85f9e621e --- /dev/null +++ b/tests/codegen-llvm/gpu_offload/explicit_memtransfer.rs @@ -0,0 +1,28 @@ +#![feature(abi_gpu_kernel, gpu_offload, offload)] +#![no_std] + +use core::offload::offload::*; + +#[cfg(target_os = "linux")] +#[unsafe(no_mangle)] +fn main() { + //println!("Hello, world!"); + let mut x = [1234.0f64; 256]; + let p: PreloadMut<[f64; 256]> = preload_mut(&mut x); + // The next line does not compile + //let p2: PreloadMut<[f64; 256]> = preload_mut(&mut x); + core::hint::black_box(p); + let y = [1234.0f64; 128]; + let q: Preload<[f64; 128]> = preload(&y); + let r: Preload<[f64; 128]> = preload(&y); + core::hint::black_box(&q); + core::hint::black_box(&r); + core::hint::black_box(&q); +} + +use core::offload::offload_kernel; + +//#[offload_kernel] +//fn foo(a: &[f32], b: &[f32], c: *mut f32) { +// unsafe { *c = a[0] + b[0] }; +//}