From fae2f07321fef096b77d841df654656f054aecd6 Mon Sep 17 00:00:00 2001 From: Manuel Drehwald Date: Sun, 31 May 2026 16:52:05 -0700 Subject: [PATCH 01/28] Enable explicit data transfers to gpu --- compiler/rustc_codegen_llvm/src/back/write.rs | 2 +- compiler/rustc_codegen_llvm/src/builder.rs | 1 + .../src/builder/gpu_helper.rs | 178 ++++++++++++++++++ .../src/builder/gpu_offload.rs | 157 +-------------- compiler/rustc_codegen_llvm/src/intrinsic.rs | 101 +++++++++- compiler/rustc_codegen_ssa/src/mir/block.rs | 18 ++ .../rustc_codegen_ssa/src/traits/intrinsic.rs | 7 + compiler/rustc_hir/src/lang_items.rs | 6 + compiler/rustc_span/src/symbol.rs | 4 + library/core/src/offload/mod.rs | 28 +++ 10 files changed, 346 insertions(+), 156 deletions(-) create mode 100644 compiler/rustc_codegen_llvm/src/builder/gpu_helper.rs 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..c18146a670019 --- /dev/null +++ b/compiler/rustc_codegen_llvm/src/builder/gpu_helper.rs @@ -0,0 +1,178 @@ +use crate::SimpleCx; +use crate::builder::Builder; +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 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); +} + +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..a89040efcf09f 100644 --- a/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs +++ b/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs @@ -3,14 +3,13 @@ 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::llvm::AttributePlace::Function; use crate::llvm::{self, Linkage, Type, Value}; @@ -534,36 +533,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. @@ -613,136 +582,21 @@ pub(crate) fn gen_call_handling<'ll, 'tcx>( 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); @@ -767,6 +621,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(); diff --git a/compiler/rustc_codegen_llvm/src/intrinsic.rs b/compiler/rustc_codegen_llvm/src/intrinsic.rs index 1c7b415fd04c7..7f405ba461ae6 100644 --- a/compiler/rustc_codegen_llvm/src/intrinsic.rs +++ b/compiler/rustc_codegen_llvm/src/intrinsic.rs @@ -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; @@ -171,6 +170,24 @@ fn call_simple_intrinsic<'ll, 'tcx>( } impl<'ll, 'tcx> IntrinsicCallBuilderMethods<'tcx> for Builder<'_, 'll, 'tcx> { + fn codegen_offload_preload_call( + &mut self, + instance: ty::Instance<'tcx>, + args: &[OperandRef<'tcx, &'ll llvm::Value>], + is_mut: bool, + ) { + let tcx = self.tcx; + 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_warn(OffloadWithoutFatLTO); + } + + codegen_offload_preload(self, tcx, instance, args); + } + fn codegen_intrinsic_call( &mut self, instance: ty::Instance<'tcx>, @@ -1894,6 +1911,82 @@ fn codegen_autodiff<'ll, 'tcx>( ); } +// 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 ty::Ref(_, pointee_ty, _) = *arg_ty.kind() else { + bug!("expected preload argument to be a reference, 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 = "asdf_I_ll_nameclash".to_owned(); + let target_symbol = cx.generate_local_symbol_name(""); + let offload_data = gen_define_handling(&cx, metadata, target_symbol, offload_globals); + 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_data.memtransfer_begin, + offload_globals.begin_mapper, + offload_globals.mapper_fn_ty, + 1, + offload_globals.ident_t_global, + ); +} + // 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 +1998,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,7 +2054,6 @@ fn codegen_offload<'ll, 'tcx>( return; } }; - register_offload(cx); let offload_data = gen_define_handling(&cx, &metadata, target_symbol, offload_globals); gen_call_handling( bx, diff --git a/compiler/rustc_codegen_ssa/src/mir/block.rs b/compiler/rustc_codegen_ssa/src/mir/block.rs index b6b95c5f12aae..8320bbb3efdd1 100644 --- a/compiler/rustc_codegen_ssa/src/mir/block.rs +++ b/compiler/rustc_codegen_ssa/src/mir/block.rs @@ -919,6 +919,24 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> { fn_span, ); + if Some(instance.def_id()) == bx.tcx().lang_items().preload_fn() { + let cg_args: Vec<_> = + args.iter().map(|arg| self.codegen_operand(bx, &arg.node)).collect(); + + bx.codegen_offload_preload_call( + instance, &cg_args, false, // immutable preload + ); + } + + if Some(instance.def_id()) == bx.tcx().lang_items().preload_mut_fn() { + let cg_args: Vec<_> = + args.iter().map(|arg| self.codegen_operand(bx, &arg.node)).collect(); + + bx.codegen_offload_preload_call( + instance, &cg_args, true, // mutable preload + ); + } + match instance.def { // We don't need AsyncDropGlueCtorShim here because it is not `noop func`, // it is `func returning noop future` diff --git a/compiler/rustc_codegen_ssa/src/traits/intrinsic.rs b/compiler/rustc_codegen_ssa/src/traits/intrinsic.rs index dcd4e722a27a8..d6d5f43ca952a 100644 --- a/compiler/rustc_codegen_ssa/src/traits/intrinsic.rs +++ b/compiler/rustc_codegen_ssa/src/traits/intrinsic.rs @@ -31,6 +31,13 @@ pub trait IntrinsicCallBuilderMethods<'tcx>: BackendTypes { span: Span, ) -> IntrinsicResult<'tcx, Self::Value>; + fn codegen_offload_preload_call( + &mut self, + instance: ty::Instance<'tcx>, + args: &[OperandRef<'tcx, Self::Value>], + is_mut: bool, + ); + fn codegen_llvm_intrinsic_call( &mut self, instance: ty::Instance<'tcx>, 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_span/src/symbol.rs b/compiler/rustc_span/src/symbol.rs index 7263680c302f1..cdcac84708ec5 100644 --- a/compiler/rustc_span/src/symbol.rs +++ b/compiler/rustc_span/src/symbol.rs @@ -1563,6 +1563,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/offload/mod.rs b/library/core/src/offload/mod.rs index 164bbb9f0047d..fa2835a5a6c25 100644 --- a/library/core/src/offload/mod.rs +++ b/library/core/src/offload/mod.rs @@ -3,3 +3,31 @@ pub use crate::macros::builtin::offload_kernel; #[unstable(feature = "gpu_offload", issue = "131513")] pub use crate::offload; + +use crate::marker::PhantomData; + +#[lang = "preload_type"] +#[unstable(feature = "offload", issue = "124509")] +pub struct Preload<'a, T: ?Sized> { + cpu_ptr: *const T, + _marker: PhantomData<&'a T>, +} + +#[lang = "preload_mut_type"] +#[unstable(feature = "offload", issue = "124509")] +pub struct PreloadMut<'a, T: ?Sized> { + 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> { + Preload { cpu_ptr: x as *const T, _marker: PhantomData } +} + +#[lang = "preload_mut"] +#[unstable(feature = "offload", issue = "124509")] +pub fn preload_mut<'a, T: ?Sized>(x: &'a mut T) -> PreloadMut<'a, T> { + PreloadMut { cpu_ptr: x as *mut T, _marker: PhantomData } +} From 6c8bec9c2fe72fa4e83c76b6e26b1354e9bf1e14 Mon Sep 17 00:00:00 2001 From: Manuel Drehwald Date: Sun, 31 May 2026 18:36:24 -0700 Subject: [PATCH 02/28] impl drop for PreloadMut, which must returned the mutated value --- .../src/builder/gpu_offload.rs | 2 + compiler/rustc_codegen_llvm/src/intrinsic.rs | 90 ++++++++++++++++++- compiler/rustc_codegen_ssa/src/mir/block.rs | 16 +++- .../rustc_codegen_ssa/src/traits/intrinsic.rs | 8 +- library/core/src/offload/mod.rs | 10 +++ 5 files changed, 123 insertions(+), 3 deletions(-) diff --git a/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs b/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs index a89040efcf09f..d3412b64cdc99 100644 --- a/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs +++ b/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs @@ -446,6 +446,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; @@ -469,6 +470,7 @@ pub(crate) fn gen_define_handling<'ll>( 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); + dbg!(&transfer_from); let memtransfer_end = add_priv_unnamed_arr(&cx, &format!(".offload_maptypes.{symbol}.end"), &transfer_from); diff --git a/compiler/rustc_codegen_llvm/src/intrinsic.rs b/compiler/rustc_codegen_llvm/src/intrinsic.rs index 7f405ba461ae6..d6c6d90eafb7e 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; @@ -188,6 +188,17 @@ impl<'ll, 'tcx> IntrinsicCallBuilderMethods<'tcx> for Builder<'_, 'll, 'tcx> { codegen_offload_preload(self, tcx, instance, args); } + fn codegen_offload_preload_mut_drop( + &mut self, + preload_ty: Ty<'tcx>, + place: PlaceRef<'tcx, &'ll llvm::Value>, + ) { + let tcx = self.tcx; + dbg!("Dropping PreloadMut; emit offload end mapper"); + + codegen_offload_preload_mut_drop(self, tcx, preload_ty, place); + } + fn codegen_intrinsic_call( &mut self, instance: ty::Instance<'tcx>, @@ -1911,6 +1922,83 @@ fn codegen_autodiff<'ll, 'tcx>( ); } +fn codegen_offload_preload_mut_drop<'ll, 'tcx>( + bx: &mut Builder<'_, 'll, 'tcx>, + tcx: TyCtxt<'tcx>, + preload_ty: Ty<'tcx>, + place: PlaceRef<'tcx, &'ll llvm::Value>, +) { + let cx = bx.cx; + dbg!("Starting the PreloadMut drop handling!"); + // PreloadMut<'a, T> -> extract T. + let ty::Adt(_adt_def, generic_args) = preload_ty.kind() else { + bug!("expected PreloadMut ADT, got {preload_ty:?}"); + }; + + // This should be the `T` parameter of PreloadMut<'a, T>. + // If this indexes the lifetime in your tree, use the correct type arg index + // or `generic_args.types().next().unwrap()`. + let pointee_ty: Ty<'tcx> = + generic_args.types().next().unwrap_or_else(|| bug!("PreloadMut without type parameter")); + + // Load field 0: `cpu_ptr: *mut T`. + let cpu_ptr_place = place.project_field(bx, 0); + dbg!(&cpu_ptr_place); + let cpu_ptr_operand = bx.load_operand(cpu_ptr_place); + dbg!(&cpu_ptr_operand); + + let args: Vec<&'ll Value> = match cpu_ptr_operand.val { + OperandValue::Immediate(ptr) => vec![ptr], + OperandValue::Pair(_data, _meta) => { + bug!("unsized PreloadMut drop not handled yet") + } + _ => bug!("unexpected PreloadMut cpu_ptr operand"), + }; + + 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. + meta.mode |= MappingFlags::FROM; + 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); + 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_data.memtransfer_end, + offload_globals.end_mapper, + offload_globals.mapper_fn_ty, + 1, + offload_globals.ident_t_global, + ); +} + // 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. diff --git a/compiler/rustc_codegen_ssa/src/mir/block.rs b/compiler/rustc_codegen_ssa/src/mir/block.rs index 8320bbb3efdd1..9204c87d17466 100644 --- a/compiler/rustc_codegen_ssa/src/mir/block.rs +++ b/compiler/rustc_codegen_ssa/src/mir/block.rs @@ -9,7 +9,7 @@ use rustc_lint_defs::builtin::TAIL_CALL_TRACK_CALLER; use rustc_middle::mir::{self, AssertKind, InlineAsmMacro, SwitchTargets, UnwindTerminateReason}; use rustc_middle::ty::layout::{HasTyCtxt, LayoutOf, ValidityRequirement}; use rustc_middle::ty::print::{with_no_trimmed_paths, with_no_visible_paths}; -use rustc_middle::ty::{self, Instance, Ty, TypeVisitableExt}; +use rustc_middle::ty::{self, Instance, Ty, TyCtxt, TypeVisitableExt}; use rustc_middle::{bug, span_bug}; use rustc_session::config::OptLevel; use rustc_span::{Span, Spanned}; @@ -27,6 +27,14 @@ use crate::mir::IntrinsicResult; use crate::traits::*; use crate::{MemFlags, meth}; +fn is_preload_mut_type<'tcx>(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>) -> bool { + let ty::Adt(adt_def, _) = ty.kind() else { + return false; + }; + + Some(adt_def.did()) == tcx.lang_items().preload_mut_type() +} + // Indicates if we are in the middle of merging a BB's successor into it. This // can happen when BB jumps directly to its successor and the successor has no // other predecessors. @@ -604,6 +612,12 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> { ) -> MergingSucc { let ty = location.ty(self.mir, bx.tcx()).ty; let ty = self.monomorphize(ty); + + if is_preload_mut_type(bx.tcx(), ty) { + let place = self.codegen_place(bx, location.as_ref()); + + bx.codegen_offload_preload_mut_drop(ty, place); + } let drop_fn = Instance::resolve_drop_glue(bx.tcx(), ty); if let ty::InstanceKind::DropGlue(_, None) = drop_fn.def { diff --git a/compiler/rustc_codegen_ssa/src/traits/intrinsic.rs b/compiler/rustc_codegen_ssa/src/traits/intrinsic.rs index d6d5f43ca952a..c6e24cae0b000 100644 --- a/compiler/rustc_codegen_ssa/src/traits/intrinsic.rs +++ b/compiler/rustc_codegen_ssa/src/traits/intrinsic.rs @@ -5,7 +5,7 @@ use super::BackendTypes; use crate::RetagInfo; use crate::mir::IntrinsicResult; use crate::mir::operand::OperandRef; -use crate::mir::place::PlaceValue; +use crate::mir::place::{PlaceRef, PlaceValue}; pub trait IntrinsicCallBuilderMethods<'tcx>: BackendTypes { /// Higher-level interface to emitting calls to intrinsics @@ -38,6 +38,12 @@ pub trait IntrinsicCallBuilderMethods<'tcx>: BackendTypes { is_mut: bool, ); + fn codegen_offload_preload_mut_drop( + &mut self, + preload_ty: ty::Ty<'tcx>, + place: PlaceRef<'tcx, Self::Value>, + ); + fn codegen_llvm_intrinsic_call( &mut self, instance: ty::Instance<'tcx>, diff --git a/library/core/src/offload/mod.rs b/library/core/src/offload/mod.rs index fa2835a5a6c25..5fc5bc13adc3d 100644 --- a/library/core/src/offload/mod.rs +++ b/library/core/src/offload/mod.rs @@ -31,3 +31,13 @@ pub fn preload<'a, T: ?Sized>(x: &'a T) -> Preload<'a, T> { pub fn preload_mut<'a, T: ?Sized>(x: &'a mut T) -> PreloadMut<'a, T> { PreloadMut { cpu_ptr: x as *mut T, _marker: PhantomData } } + +impl Drop for PreloadMut<'_, T> { + fn drop(&mut self) { + // Intentionally empty. + // + // This exists so MIR creates Drop terminators for PreloadMut. + // rustc codegen intercepts those terminators and emits the + // offload return mapper. + } +} From 198c71eeea364e19e7a6d070e8738692da09c1ef Mon Sep 17 00:00:00 2001 From: Manuel Drehwald Date: Sun, 31 May 2026 18:56:21 -0700 Subject: [PATCH 03/28] Also handle drop for Preload --- compiler/rustc_codegen_llvm/src/intrinsic.rs | 15 +++++++++++---- compiler/rustc_codegen_ssa/src/mir/block.rs | 15 ++++++++++++++- .../rustc_codegen_ssa/src/traits/intrinsic.rs | 3 ++- library/core/src/offload/mod.rs | 10 ++++++++++ 4 files changed, 37 insertions(+), 6 deletions(-) diff --git a/compiler/rustc_codegen_llvm/src/intrinsic.rs b/compiler/rustc_codegen_llvm/src/intrinsic.rs index d6c6d90eafb7e..b44b3a5fd78d3 100644 --- a/compiler/rustc_codegen_llvm/src/intrinsic.rs +++ b/compiler/rustc_codegen_llvm/src/intrinsic.rs @@ -188,15 +188,16 @@ impl<'ll, 'tcx> IntrinsicCallBuilderMethods<'tcx> for Builder<'_, 'll, 'tcx> { codegen_offload_preload(self, tcx, instance, args); } - fn codegen_offload_preload_mut_drop( + fn codegen_offload_preload_drop( &mut self, preload_ty: Ty<'tcx>, place: PlaceRef<'tcx, &'ll llvm::Value>, + is_mut: bool, ) { let tcx = self.tcx; dbg!("Dropping PreloadMut; emit offload end mapper"); - codegen_offload_preload_mut_drop(self, tcx, preload_ty, place); + codegen_offload_preload_drop(self, tcx, preload_ty, place, is_mut); } fn codegen_intrinsic_call( @@ -1922,11 +1923,12 @@ fn codegen_autodiff<'ll, 'tcx>( ); } -fn codegen_offload_preload_mut_drop<'ll, 'tcx>( +fn codegen_offload_preload_drop<'ll, 'tcx>( bx: &mut Builder<'_, 'll, 'tcx>, tcx: TyCtxt<'tcx>, preload_ty: Ty<'tcx>, place: PlaceRef<'tcx, &'ll llvm::Value>, + is_mut: bool, ) { let cx = bx.cx; dbg!("Starting the PreloadMut drop handling!"); @@ -1959,7 +1961,12 @@ fn codegen_offload_preload_mut_drop<'ll, 'tcx>( // 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. - meta.mode |= MappingFlags::FROM; + 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]; diff --git a/compiler/rustc_codegen_ssa/src/mir/block.rs b/compiler/rustc_codegen_ssa/src/mir/block.rs index 9204c87d17466..61ad6bfaf3067 100644 --- a/compiler/rustc_codegen_ssa/src/mir/block.rs +++ b/compiler/rustc_codegen_ssa/src/mir/block.rs @@ -35,6 +35,14 @@ fn is_preload_mut_type<'tcx>(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>) -> bool { Some(adt_def.did()) == tcx.lang_items().preload_mut_type() } +fn is_preload_type<'tcx>(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>) -> bool { + let ty::Adt(adt_def, _) = ty.kind() else { + return false; + }; + + Some(adt_def.did()) == tcx.lang_items().preload_type() +} + // Indicates if we are in the middle of merging a BB's successor into it. This // can happen when BB jumps directly to its successor and the successor has no // other predecessors. @@ -616,7 +624,12 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> { if is_preload_mut_type(bx.tcx(), ty) { let place = self.codegen_place(bx, location.as_ref()); - bx.codegen_offload_preload_mut_drop(ty, place); + bx.codegen_offload_preload_drop(ty, place, true); + } + if is_preload_type(bx.tcx(), ty) { + let place = self.codegen_place(bx, location.as_ref()); + + bx.codegen_offload_preload_drop(ty, place, false); } let drop_fn = Instance::resolve_drop_glue(bx.tcx(), ty); diff --git a/compiler/rustc_codegen_ssa/src/traits/intrinsic.rs b/compiler/rustc_codegen_ssa/src/traits/intrinsic.rs index c6e24cae0b000..f6e6bb0a49f80 100644 --- a/compiler/rustc_codegen_ssa/src/traits/intrinsic.rs +++ b/compiler/rustc_codegen_ssa/src/traits/intrinsic.rs @@ -38,10 +38,11 @@ pub trait IntrinsicCallBuilderMethods<'tcx>: BackendTypes { is_mut: bool, ); - fn codegen_offload_preload_mut_drop( + fn codegen_offload_preload_drop( &mut self, preload_ty: ty::Ty<'tcx>, place: PlaceRef<'tcx, Self::Value>, + mut_drop: bool, ); fn codegen_llvm_intrinsic_call( diff --git a/library/core/src/offload/mod.rs b/library/core/src/offload/mod.rs index 5fc5bc13adc3d..942bc677c14c3 100644 --- a/library/core/src/offload/mod.rs +++ b/library/core/src/offload/mod.rs @@ -41,3 +41,13 @@ impl Drop for PreloadMut<'_, T> { // offload return mapper. } } + +impl Drop for Preload<'_, T> { + fn drop(&mut self) { + // Intentionally empty. + // + // This exists so MIR creates Drop terminators for Preload. + // rustc codegen intercepts those terminators and emits the + // offload return mapper. + } +} From 2d429ff43c93a509d2ace12a3e73a81c077a72be Mon Sep 17 00:00:00 2001 From: Manuel Drehwald Date: Sun, 31 May 2026 18:59:02 -0700 Subject: [PATCH 04/28] wip test --- .../gpu_offload/explicit_memtransfer.rs | 23 +++++++++++++++++++ 1 file changed, 23 insertions(+) create mode 100644 tests/codegen-llvm/gpu_offload/explicit_memtransfer.rs 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..d9ef23851fb97 --- /dev/null +++ b/tests/codegen-llvm/gpu_offload/explicit_memtransfer.rs @@ -0,0 +1,23 @@ +#![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); + core::hint::black_box(p); + let y = [1234.0f64; 128]; + let q: Preload<[f64; 128]> = preload(&y); + 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] }; +//} From 555b339761b7bfe99b67ce6d852f919b46f90a1d Mon Sep 17 00:00:00 2001 From: Manuel Drehwald Date: Sun, 31 May 2026 19:00:55 -0700 Subject: [PATCH 05/28] wip test --- tests/codegen-llvm/gpu_offload/explicit_memtransfer.rs | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tests/codegen-llvm/gpu_offload/explicit_memtransfer.rs b/tests/codegen-llvm/gpu_offload/explicit_memtransfer.rs index d9ef23851fb97..dd2a85f9e621e 100644 --- a/tests/codegen-llvm/gpu_offload/explicit_memtransfer.rs +++ b/tests/codegen-llvm/gpu_offload/explicit_memtransfer.rs @@ -9,10 +9,15 @@ 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); - core::hint::black_box(q); + 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; From d204beb6ed940ba2ad861f4dfc3c8004ab1e3c20 Mon Sep 17 00:00:00 2001 From: Manuel Drehwald Date: Fri, 5 Jun 2026 18:16:06 -0700 Subject: [PATCH 06/28] wip, moving to intrinsics --- compiler/rustc_codegen_llvm/src/intrinsic.rs | 54 ++++++-------------- compiler/rustc_span/src/symbol.rs | 2 + library/core/src/intrinsics/mod.rs | 8 +++ library/core/src/offload/mod.rs | 47 ++++++++++++----- 4 files changed, 60 insertions(+), 51 deletions(-) diff --git a/compiler/rustc_codegen_llvm/src/intrinsic.rs b/compiler/rustc_codegen_llvm/src/intrinsic.rs index b44b3a5fd78d3..e69262f1b0f54 100644 --- a/compiler/rustc_codegen_llvm/src/intrinsic.rs +++ b/compiler/rustc_codegen_llvm/src/intrinsic.rs @@ -170,36 +170,6 @@ fn call_simple_intrinsic<'ll, 'tcx>( } impl<'ll, 'tcx> IntrinsicCallBuilderMethods<'tcx> for Builder<'_, 'll, 'tcx> { - fn codegen_offload_preload_call( - &mut self, - instance: ty::Instance<'tcx>, - args: &[OperandRef<'tcx, &'ll llvm::Value>], - is_mut: bool, - ) { - let tcx = self.tcx; - 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_warn(OffloadWithoutFatLTO); - } - - codegen_offload_preload(self, tcx, instance, args); - } - - fn codegen_offload_preload_drop( - &mut self, - preload_ty: Ty<'tcx>, - place: PlaceRef<'tcx, &'ll llvm::Value>, - is_mut: bool, - ) { - let tcx = self.tcx; - dbg!("Dropping PreloadMut; emit offload end mapper"); - - codegen_offload_preload_drop(self, tcx, preload_ty, place, is_mut); - } - fn codegen_intrinsic_call( &mut self, instance: ty::Instance<'tcx>, @@ -259,6 +229,15 @@ impl<'ll, 'tcx> IntrinsicCallBuilderMethods<'tcx> for Builder<'_, 'll, 'tcx> { codegen_autodiff(self, tcx, instance, args, result); return IntrinsicResult::WroteIntoPlace; } + sym::offload_preload => { + codegen_offload_preload(self, tcx, instance, args, false); + 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); @@ -1933,16 +1912,13 @@ fn codegen_offload_preload_drop<'ll, 'tcx>( let cx = bx.cx; dbg!("Starting the PreloadMut drop handling!"); // PreloadMut<'a, T> -> extract T. - let ty::Adt(_adt_def, generic_args) = preload_ty.kind() else { - bug!("expected PreloadMut ADT, got {preload_ty:?}"); - }; - - // This should be the `T` parameter of PreloadMut<'a, T>. - // If this indexes the lifetime in your tree, use the correct type arg index - // or `generic_args.types().next().unwrap()`. - let pointee_ty: Ty<'tcx> = - generic_args.types().next().unwrap_or_else(|| bug!("PreloadMut without type parameter")); + let arg = &args[0]; + let arg_ty = arg.layout.ty; + let pointee_ty = match *arg_ty.kind() { + ty::RawPtr(pointee_ty, _) => pointee_ty, + _ => bug!("expected raw pointer argument, got {arg_ty:?}"), + }; // Load field 0: `cpu_ptr: *mut T`. let cpu_ptr_place = place.project_field(bx, 0); dbg!(&cpu_ptr_place); diff --git a/compiler/rustc_span/src/symbol.rs b/compiler/rustc_span/src/symbol.rs index cdcac84708ec5..8a98da74da688 100644 --- a/compiler/rustc_span/src/symbol.rs +++ b/compiler/rustc_span/src/symbol.rs @@ -1425,6 +1425,8 @@ symbols! { of, off, offload, + offload_preload, + offload_preload_end, offload_kernel, offset, offset_of, diff --git a/library/core/src/intrinsics/mod.rs b/library/core/src/intrinsics/mod.rs index 78d7314c58110..c3cbed7cb1bb7 100644 --- a/library/core/src/intrinsics/mod.rs +++ b/library/core/src/intrinsics/mod.rs @@ -3573,6 +3573,14 @@ pub const fn offload( args: T, ) -> R; +#[rustc_intrinsic] +#[rustc_nounwind] +pub unsafe fn offload_preload(ptr: *mut T, is_mut: bool); + +#[rustc_intrinsic] +#[rustc_nounwind] +pub unsafe fn offload_preload_end(ptr: *mut 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 942bc677c14c3..58c90be6119bf 100644 --- a/library/core/src/offload/mod.rs +++ b/library/core/src/offload/mod.rs @@ -6,6 +6,14 @@ 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> { @@ -13,6 +21,13 @@ pub struct Preload<'a, T: ?Sized> { _marker: PhantomData<&'a T>, } +// 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> { @@ -23,31 +38,39 @@ pub struct PreloadMut<'a, T: ?Sized> { #[lang = "preload"] #[unstable(feature = "offload", issue = "124509")] pub fn preload<'a, T: ?Sized>(x: &'a T) -> Preload<'a, T> { - Preload { cpu_ptr: x as *const T, _marker: PhantomData } + let p = Preload { cpu_ptr: x as *const T, _marker: PhantomData }; + + unsafe { + 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> { - PreloadMut { cpu_ptr: x as *mut T, _marker: PhantomData } + let p = PreloadMut { cpu_ptr: x as *mut T, _marker: PhantomData }; + + unsafe { + core::intrinsics::offload_preload(p.cpu_ptr, true); + } + + p } impl Drop for PreloadMut<'_, T> { fn drop(&mut self) { - // Intentionally empty. - // - // This exists so MIR creates Drop terminators for PreloadMut. - // rustc codegen intercepts those terminators and emits the - // offload return mapper. + unsafe { + core::intrinsics::offload_preload_mut_end(self.cpu_ptr, true); + } } } impl Drop for Preload<'_, T> { fn drop(&mut self) { - // Intentionally empty. - // - // This exists so MIR creates Drop terminators for Preload. - // rustc codegen intercepts those terminators and emits the - // offload return mapper. + unsafe { + core::intrinsics::offload_preload_mut_end(self.cpu_ptr, false); + } } } From f8bda403c91d3a647aa4e02be62874eeac0d9a97 Mon Sep 17 00:00:00 2001 From: Manuel Drehwald Date: Fri, 5 Jun 2026 19:50:45 -0700 Subject: [PATCH 07/28] partial cleanup --- compiler/rustc_codegen_llvm/src/intrinsic.rs | 8 ++-- compiler/rustc_codegen_ssa/src/mir/block.rs | 44 ------------------- .../rustc_codegen_ssa/src/traits/intrinsic.rs | 14 ------ library/core/src/offload/mod.rs | 4 +- 4 files changed, 6 insertions(+), 64 deletions(-) diff --git a/compiler/rustc_codegen_llvm/src/intrinsic.rs b/compiler/rustc_codegen_llvm/src/intrinsic.rs index e69262f1b0f54..749704934ec35 100644 --- a/compiler/rustc_codegen_llvm/src/intrinsic.rs +++ b/compiler/rustc_codegen_llvm/src/intrinsic.rs @@ -235,7 +235,7 @@ impl<'ll, 'tcx> IntrinsicCallBuilderMethods<'tcx> for Builder<'_, 'll, 'tcx> { } sym::offload_preload_end => { - codegen_offload_preload_drop(self, tcx, instance, args); + codegen_offload_preload_drop(self, tcx, instance, args, false); return IntrinsicResult::WroteIntoPlace; } sym::offload => { @@ -1901,12 +1901,11 @@ fn codegen_autodiff<'ll, 'tcx>( fnc_tree, ); } - fn codegen_offload_preload_drop<'ll, 'tcx>( bx: &mut Builder<'_, 'll, 'tcx>, tcx: TyCtxt<'tcx>, - preload_ty: Ty<'tcx>, - place: PlaceRef<'tcx, &'ll llvm::Value>, + _instance: ty::Instance<'tcx>, + args: &[OperandRef<'tcx, &'ll llvm::Value>], is_mut: bool, ) { let cx = bx.cx; @@ -2002,6 +2001,7 @@ fn codegen_offload_preload<'ll, 'tcx>( tcx: TyCtxt<'tcx>, _instance: ty::Instance<'tcx>, args: &[OperandRef<'tcx, &'ll Value>], + is_mut: bool, ) { dbg!("Starting the preload handling!"); let cx = bx.cx; diff --git a/compiler/rustc_codegen_ssa/src/mir/block.rs b/compiler/rustc_codegen_ssa/src/mir/block.rs index 61ad6bfaf3067..003a01df973d8 100644 --- a/compiler/rustc_codegen_ssa/src/mir/block.rs +++ b/compiler/rustc_codegen_ssa/src/mir/block.rs @@ -27,22 +27,6 @@ use crate::mir::IntrinsicResult; use crate::traits::*; use crate::{MemFlags, meth}; -fn is_preload_mut_type<'tcx>(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>) -> bool { - let ty::Adt(adt_def, _) = ty.kind() else { - return false; - }; - - Some(adt_def.did()) == tcx.lang_items().preload_mut_type() -} - -fn is_preload_type<'tcx>(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>) -> bool { - let ty::Adt(adt_def, _) = ty.kind() else { - return false; - }; - - Some(adt_def.did()) == tcx.lang_items().preload_type() -} - // Indicates if we are in the middle of merging a BB's successor into it. This // can happen when BB jumps directly to its successor and the successor has no // other predecessors. @@ -621,16 +605,6 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> { let ty = location.ty(self.mir, bx.tcx()).ty; let ty = self.monomorphize(ty); - if is_preload_mut_type(bx.tcx(), ty) { - let place = self.codegen_place(bx, location.as_ref()); - - bx.codegen_offload_preload_drop(ty, place, true); - } - if is_preload_type(bx.tcx(), ty) { - let place = self.codegen_place(bx, location.as_ref()); - - bx.codegen_offload_preload_drop(ty, place, false); - } let drop_fn = Instance::resolve_drop_glue(bx.tcx(), ty); if let ty::InstanceKind::DropGlue(_, None) = drop_fn.def { @@ -946,24 +920,6 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> { fn_span, ); - if Some(instance.def_id()) == bx.tcx().lang_items().preload_fn() { - let cg_args: Vec<_> = - args.iter().map(|arg| self.codegen_operand(bx, &arg.node)).collect(); - - bx.codegen_offload_preload_call( - instance, &cg_args, false, // immutable preload - ); - } - - if Some(instance.def_id()) == bx.tcx().lang_items().preload_mut_fn() { - let cg_args: Vec<_> = - args.iter().map(|arg| self.codegen_operand(bx, &arg.node)).collect(); - - bx.codegen_offload_preload_call( - instance, &cg_args, true, // mutable preload - ); - } - match instance.def { // We don't need AsyncDropGlueCtorShim here because it is not `noop func`, // it is `func returning noop future` diff --git a/compiler/rustc_codegen_ssa/src/traits/intrinsic.rs b/compiler/rustc_codegen_ssa/src/traits/intrinsic.rs index f6e6bb0a49f80..410e885576743 100644 --- a/compiler/rustc_codegen_ssa/src/traits/intrinsic.rs +++ b/compiler/rustc_codegen_ssa/src/traits/intrinsic.rs @@ -31,20 +31,6 @@ pub trait IntrinsicCallBuilderMethods<'tcx>: BackendTypes { span: Span, ) -> IntrinsicResult<'tcx, Self::Value>; - fn codegen_offload_preload_call( - &mut self, - instance: ty::Instance<'tcx>, - args: &[OperandRef<'tcx, Self::Value>], - is_mut: bool, - ); - - fn codegen_offload_preload_drop( - &mut self, - preload_ty: ty::Ty<'tcx>, - place: PlaceRef<'tcx, Self::Value>, - mut_drop: bool, - ); - fn codegen_llvm_intrinsic_call( &mut self, instance: ty::Instance<'tcx>, diff --git a/library/core/src/offload/mod.rs b/library/core/src/offload/mod.rs index 58c90be6119bf..2d5150b06e82b 100644 --- a/library/core/src/offload/mod.rs +++ b/library/core/src/offload/mod.rs @@ -62,7 +62,7 @@ pub fn preload_mut<'a, T: ?Sized>(x: &'a mut T) -> PreloadMut<'a, T> { impl Drop for PreloadMut<'_, T> { fn drop(&mut self) { unsafe { - core::intrinsics::offload_preload_mut_end(self.cpu_ptr, true); + core::intrinsics::offload_preload_end(self.cpu_ptr, true); } } } @@ -70,7 +70,7 @@ impl Drop for PreloadMut<'_, T> { impl Drop for Preload<'_, T> { fn drop(&mut self) { unsafe { - core::intrinsics::offload_preload_mut_end(self.cpu_ptr, false); + core::intrinsics::offload_preload_end(self.cpu_ptr, false); } } } From 1e3127fbeab03eafcb468e147c654fd37813656e Mon Sep 17 00:00:00 2001 From: Manuel Drehwald Date: Fri, 5 Jun 2026 20:25:38 -0700 Subject: [PATCH 08/28] bootstrap compiles, test fails. wip --- compiler/rustc_codegen_llvm/src/intrinsic.rs | 25 +++++++------------ .../rustc_hir_analysis/src/check/intrinsic.rs | 5 ++++ library/core/src/intrinsics/mod.rs | 4 +-- 3 files changed, 16 insertions(+), 18 deletions(-) diff --git a/compiler/rustc_codegen_llvm/src/intrinsic.rs b/compiler/rustc_codegen_llvm/src/intrinsic.rs index 749704934ec35..ee72834a1029d 100644 --- a/compiler/rustc_codegen_llvm/src/intrinsic.rs +++ b/compiler/rustc_codegen_llvm/src/intrinsic.rs @@ -1911,27 +1911,20 @@ fn codegen_offload_preload_drop<'ll, 'tcx>( let cx = bx.cx; dbg!("Starting the PreloadMut drop handling!"); // PreloadMut<'a, T> -> extract T. - let arg = &args[0]; - let arg_ty = arg.layout.ty; + let ptr_arg = &args[0]; - let pointee_ty = match *arg_ty.kind() { + let pointee_ty = match *ptr_arg.layout.ty.kind() { ty::RawPtr(pointee_ty, _) => pointee_ty, - _ => bug!("expected raw pointer argument, got {arg_ty:?}"), + _ => bug!("expected raw pointer argument"), }; - // Load field 0: `cpu_ptr: *mut T`. - let cpu_ptr_place = place.project_field(bx, 0); - dbg!(&cpu_ptr_place); - let cpu_ptr_operand = bx.load_operand(cpu_ptr_place); - dbg!(&cpu_ptr_operand); - - let args: Vec<&'ll Value> = match cpu_ptr_operand.val { - OperandValue::Immediate(ptr) => vec![ptr], - OperandValue::Pair(_data, _meta) => { - bug!("unsized PreloadMut drop not handled yet") - } - _ => bug!("unexpected PreloadMut cpu_ptr operand"), + + 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 diff --git a/compiler/rustc_hir_analysis/src/check/intrinsic.rs b/compiler/rustc_hir_analysis/src/check/intrinsic.rs index a4d6056e47bdc..ce3e18d41fcbe 100644 --- a/compiler/rustc_hir_analysis/src/check/intrinsic.rs +++ b/compiler/rustc_hir_analysis/src/check/intrinsic.rs @@ -164,6 +164,8 @@ fn intrinsic_operation_unsafety(tcx: TyCtxt<'_>, intrinsic_id: LocalDefId) -> hi | sym::mul_with_overflow | sym::needs_drop | sym::offload + | sym::offload_preload + | sym::offload_preload_end | sym::offset_of | sym::overflow_checks | sym::powf16 @@ -361,6 +363,9 @@ 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::offset => (2, 0, vec![param(0), param(1)], param(0)), sym::arith_offset => ( 1, diff --git a/library/core/src/intrinsics/mod.rs b/library/core/src/intrinsics/mod.rs index c3cbed7cb1bb7..7bf77bd22803a 100644 --- a/library/core/src/intrinsics/mod.rs +++ b/library/core/src/intrinsics/mod.rs @@ -3575,11 +3575,11 @@ pub const fn offload( #[rustc_intrinsic] #[rustc_nounwind] -pub unsafe fn offload_preload(ptr: *mut T, is_mut: bool); +pub fn offload_preload(ptr: *const T, is_mut: bool); #[rustc_intrinsic] #[rustc_nounwind] -pub unsafe fn offload_preload_end(ptr: *mut T, is_mut: bool); +pub fn offload_preload_end(ptr: *const T, is_mut: bool); /// Inform Miri that a given pointer definitely has a certain alignment. #[cfg(miri)] From 2da8e12a9266be3b8d2741496f7598933aeeb5a7 Mon Sep 17 00:00:00 2001 From: Manuel Drehwald Date: Fri, 5 Jun 2026 20:59:33 -0700 Subject: [PATCH 09/28] passing test --- compiler/rustc_codegen_llvm/src/intrinsic.rs | 37 ++++++++++++++++---- 1 file changed, 31 insertions(+), 6 deletions(-) diff --git a/compiler/rustc_codegen_llvm/src/intrinsic.rs b/compiler/rustc_codegen_llvm/src/intrinsic.rs index ee72834a1029d..6f651a0396579 100644 --- a/compiler/rustc_codegen_llvm/src/intrinsic.rs +++ b/compiler/rustc_codegen_llvm/src/intrinsic.rs @@ -230,12 +230,12 @@ impl<'ll, 'tcx> IntrinsicCallBuilderMethods<'tcx> for Builder<'_, 'll, 'tcx> { return IntrinsicResult::WroteIntoPlace; } sym::offload_preload => { - codegen_offload_preload(self, tcx, instance, args, false); + codegen_offload_preload(self, tcx, instance, args); return IntrinsicResult::WroteIntoPlace; } sym::offload_preload_end => { - codegen_offload_preload_drop(self, tcx, instance, args, false); + codegen_offload_preload_drop(self, tcx, instance, args); return IntrinsicResult::WroteIntoPlace; } sym::offload => { @@ -1901,17 +1901,42 @@ fn codegen_autodiff<'ll, 'tcx>( fnc_tree, ); } + +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_preload_drop<'ll, 'tcx>( bx: &mut Builder<'_, 'll, 'tcx>, tcx: TyCtxt<'tcx>, _instance: ty::Instance<'tcx>, args: &[OperandRef<'tcx, &'ll llvm::Value>], - is_mut: bool, ) { 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, @@ -1994,7 +2019,6 @@ fn codegen_offload_preload<'ll, 'tcx>( tcx: TyCtxt<'tcx>, _instance: ty::Instance<'tcx>, args: &[OperandRef<'tcx, &'ll Value>], - is_mut: bool, ) { dbg!("Starting the preload handling!"); let cx = bx.cx; @@ -2008,8 +2032,9 @@ fn codegen_offload_preload<'ll, 'tcx>( let arg_ty = arg.layout.ty; - let ty::Ref(_, pointee_ty, _) = *arg_ty.kind() else { - bug!("expected preload argument to be a reference, got {arg_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); From bbb2b7dca827d28f61c9bf66c49254dd6af034f2 Mon Sep 17 00:00:00 2001 From: Manuel Drehwald Date: Fri, 5 Jun 2026 21:06:00 -0700 Subject: [PATCH 10/28] cleanup --- compiler/rustc_codegen_ssa/src/mir/block.rs | 3 +-- compiler/rustc_codegen_ssa/src/traits/intrinsic.rs | 2 +- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/compiler/rustc_codegen_ssa/src/mir/block.rs b/compiler/rustc_codegen_ssa/src/mir/block.rs index 003a01df973d8..b6b95c5f12aae 100644 --- a/compiler/rustc_codegen_ssa/src/mir/block.rs +++ b/compiler/rustc_codegen_ssa/src/mir/block.rs @@ -9,7 +9,7 @@ use rustc_lint_defs::builtin::TAIL_CALL_TRACK_CALLER; use rustc_middle::mir::{self, AssertKind, InlineAsmMacro, SwitchTargets, UnwindTerminateReason}; use rustc_middle::ty::layout::{HasTyCtxt, LayoutOf, ValidityRequirement}; use rustc_middle::ty::print::{with_no_trimmed_paths, with_no_visible_paths}; -use rustc_middle::ty::{self, Instance, Ty, TyCtxt, TypeVisitableExt}; +use rustc_middle::ty::{self, Instance, Ty, TypeVisitableExt}; use rustc_middle::{bug, span_bug}; use rustc_session::config::OptLevel; use rustc_span::{Span, Spanned}; @@ -604,7 +604,6 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> { ) -> MergingSucc { let ty = location.ty(self.mir, bx.tcx()).ty; let ty = self.monomorphize(ty); - let drop_fn = Instance::resolve_drop_glue(bx.tcx(), ty); if let ty::InstanceKind::DropGlue(_, None) = drop_fn.def { diff --git a/compiler/rustc_codegen_ssa/src/traits/intrinsic.rs b/compiler/rustc_codegen_ssa/src/traits/intrinsic.rs index 410e885576743..dcd4e722a27a8 100644 --- a/compiler/rustc_codegen_ssa/src/traits/intrinsic.rs +++ b/compiler/rustc_codegen_ssa/src/traits/intrinsic.rs @@ -5,7 +5,7 @@ use super::BackendTypes; use crate::RetagInfo; use crate::mir::IntrinsicResult; use crate::mir::operand::OperandRef; -use crate::mir::place::{PlaceRef, PlaceValue}; +use crate::mir::place::PlaceValue; pub trait IntrinsicCallBuilderMethods<'tcx>: BackendTypes { /// Higher-level interface to emitting calls to intrinsics From 445f1cb37de1cc357b3c106576252209a0ff5ae2 Mon Sep 17 00:00:00 2001 From: Manuel Drehwald Date: Fri, 5 Jun 2026 22:23:01 -0700 Subject: [PATCH 11/28] cleanup --- library/core/src/offload/mod.rs | 16 ++++------------ 1 file changed, 4 insertions(+), 12 deletions(-) diff --git a/library/core/src/offload/mod.rs b/library/core/src/offload/mod.rs index 2d5150b06e82b..603dfa4135819 100644 --- a/library/core/src/offload/mod.rs +++ b/library/core/src/offload/mod.rs @@ -40,9 +40,7 @@ pub struct PreloadMut<'a, T: ?Sized> { pub fn preload<'a, T: ?Sized>(x: &'a T) -> Preload<'a, T> { let p = Preload { cpu_ptr: x as *const T, _marker: PhantomData }; - unsafe { - core::intrinsics::offload_preload(p.cpu_ptr, false); - } + core::intrinsics::offload_preload(p.cpu_ptr, false); p } @@ -52,25 +50,19 @@ pub fn preload<'a, T: ?Sized>(x: &'a T) -> Preload<'a, T> { 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 }; - unsafe { - core::intrinsics::offload_preload(p.cpu_ptr, true); - } + core::intrinsics::offload_preload(p.cpu_ptr, true); p } impl Drop for PreloadMut<'_, T> { fn drop(&mut self) { - unsafe { - core::intrinsics::offload_preload_end(self.cpu_ptr, true); - } + core::intrinsics::offload_preload_end(self.cpu_ptr, true); } } impl Drop for Preload<'_, T> { fn drop(&mut self) { - unsafe { - core::intrinsics::offload_preload_end(self.cpu_ptr, false); - } + core::intrinsics::offload_preload_end(self.cpu_ptr, false); } } From fc71123605e4e37c0576c63485a068eecbd4448d Mon Sep 17 00:00:00 2001 From: Manuel Drehwald Date: Sun, 7 Jun 2026 02:02:18 -0700 Subject: [PATCH 12/28] can compile and run all rust_perf benchmarks --- .../src/builder/gpu_offload.rs | 20 ++++++++++++++----- compiler/rustc_codegen_llvm/src/intrinsic.rs | 12 +++++------ library/core/src/offload/mod.rs | 4 ++-- 3 files changed, 23 insertions(+), 13 deletions(-) diff --git a/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs b/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs index d3412b64cdc99..e6df35c864d55 100644 --- a/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs +++ b/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs @@ -347,7 +347,7 @@ 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, } @@ -419,6 +419,7 @@ pub(crate) fn gen_define_handling<'ll>( metadata: &[OffloadMetadata], symbol: String, offload_globals: &OffloadGlobals<'ll>, + gen_kernel: bool, ) -> OffloadKernelGlobals<'ll> { if let Some(entry) = cx.offload_kernel_cache.borrow().get(&symbol) { return *entry; @@ -446,7 +447,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); + //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,9 +469,17 @@ 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); - dbg!(&transfer_from); + + let memtransfer_kernel = if gen_kernel { + Some(add_priv_unnamed_arr( + &cx, + &format!(".offload_maptypes.{symbol}.kernel"), + &transfer_kernel, + )) + } else { + None + }; + //dbg!(&transfer_from); let memtransfer_end = add_priv_unnamed_arr(&cx, &format!(".offload_maptypes.{symbol}.end"), &transfer_from); @@ -571,6 +580,7 @@ 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; diff --git a/compiler/rustc_codegen_llvm/src/intrinsic.rs b/compiler/rustc_codegen_llvm/src/intrinsic.rs index 6f651a0396579..7f4956e821f24 100644 --- a/compiler/rustc_codegen_llvm/src/intrinsic.rs +++ b/compiler/rustc_codegen_llvm/src/intrinsic.rs @@ -243,9 +243,9 @@ impl<'ll, 'tcx> IntrinsicCallBuilderMethods<'tcx> for Builder<'_, 'll, 'tcx> { 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 @@ -1976,7 +1976,7 @@ fn codegen_offload_preload_drop<'ll, 'tcx>( let target_symbol = cx.generate_local_symbol_name(""); dbg!("done for now"); - let offload_data = gen_define_handling(&cx, metadata, target_symbol, offload_globals); + let offload_data = gen_define_handling(&cx, metadata, target_symbol, offload_globals, false); 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, @@ -2053,7 +2053,7 @@ fn codegen_offload_preload<'ll, 'tcx>( dbg!("asdf"); //let target_symbol = "asdf_I_ll_nameclash".to_owned(); let target_symbol = cx.generate_local_symbol_name(""); - let offload_data = gen_define_handling(&cx, metadata, target_symbol, offload_globals); + let offload_data = gen_define_handling(&cx, metadata, target_symbol, offload_globals, false); 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, @@ -2143,7 +2143,7 @@ fn codegen_offload<'ll, 'tcx>( return; } }; - let offload_data = gen_define_handling(&cx, &metadata, target_symbol, offload_globals); + let offload_data = gen_define_handling(&cx, &metadata, target_symbol, offload_globals, true); gen_call_handling( bx, &offload_data, diff --git a/library/core/src/offload/mod.rs b/library/core/src/offload/mod.rs index 603dfa4135819..acd24dc99ecc2 100644 --- a/library/core/src/offload/mod.rs +++ b/library/core/src/offload/mod.rs @@ -17,7 +17,7 @@ use crate::marker::PhantomData; #[lang = "preload_type"] #[unstable(feature = "offload", issue = "124509")] pub struct Preload<'a, T: ?Sized> { - cpu_ptr: *const T, + pub cpu_ptr: *const T, _marker: PhantomData<&'a T>, } @@ -31,7 +31,7 @@ pub struct Preload<'a, T: ?Sized> { #[lang = "preload_mut_type"] #[unstable(feature = "offload", issue = "124509")] pub struct PreloadMut<'a, T: ?Sized> { - cpu_ptr: *mut T, + pub cpu_ptr: *mut T, _marker: PhantomData<&'a mut T>, } From 432af008d2f0bfe181220b13d5e5ebd2819d7acf Mon Sep 17 00:00:00 2001 From: Manuel Drehwald Date: Sun, 7 Jun 2026 02:12:15 -0700 Subject: [PATCH 13/28] only generate kernel globals for kernels, not for preload/drop memtransfers --- .../src/builder/gpu_offload.rs | 52 ++++++++++--------- 1 file changed, 27 insertions(+), 25 deletions(-) diff --git a/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs b/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs index e6df35c864d55..096488a9f98ff 100644 --- a/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs +++ b/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs @@ -490,31 +490,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}"); - - 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}"); - - // 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); - - cx.add_compiler_used_global(offload_entry); + 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 name = format!(".offloading.entry.{symbol}"); + + // 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); + + cx.add_compiler_used_global(offload_entry); + } let result = OffloadKernelGlobals { offload_sizes, From 23cd070821f8358813c08245e2ece3831474b43f Mon Sep 17 00:00:00 2001 From: Manuel Drehwald Date: Wed, 10 Jun 2026 15:37:23 -0700 Subject: [PATCH 14/28] make preload nowait, todo: await on kernel launch --- .../src/builder/gpu_helper.rs | 8 ++++- .../src/builder/gpu_offload.rs | 31 +++++++++++++++++-- compiler/rustc_codegen_llvm/src/intrinsic.rs | 27 ++++++++++++---- 3 files changed, 57 insertions(+), 9 deletions(-) diff --git a/compiler/rustc_codegen_llvm/src/builder/gpu_helper.rs b/compiler/rustc_codegen_llvm/src/builder/gpu_helper.rs index c18146a670019..fafb84ac81508 100644 --- a/compiler/rustc_codegen_llvm/src/builder/gpu_helper.rs +++ b/compiler/rustc_codegen_llvm/src/builder/gpu_helper.rs @@ -1,5 +1,6 @@ use crate::SimpleCx; use crate::builder::Builder; +use crate::intrinsic::TransferType; use crate::llvm; use crate::llvm::{Type, Value}; use rustc_abi::Align; @@ -67,13 +68,18 @@ pub(crate) fn generate_mapper_call<'ll, 'tcx>( 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 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]); + } builder.call(fn_ty, None, None, fn_to_call, &args, None, None); } diff --git a/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs b/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs index 096488a9f98ff..d2d2cde2019e0 100644 --- a/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs +++ b/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs @@ -11,6 +11,7 @@ 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}; @@ -28,6 +29,9 @@ 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, } @@ -37,6 +41,7 @@ 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); // We want LLVM's openmp-opt pass to pick up and optimize this module, since it covers both @@ -49,8 +54,10 @@ impl<'ll> OffloadGlobals<'ll> { kernel_args_ty, offload_entry_ty, begin_mapper, + nowait_begin_mapper, end_mapper, mapper_fn_ty, + nowait_mapper_fn_ty, ident_t_global, } } @@ -352,6 +359,24 @@ pub(crate) struct OffloadKernelGlobals<'ll> { 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) { @@ -419,13 +444,14 @@ pub(crate) fn gen_define_handling<'ll>( metadata: &[OffloadMetadata], symbol: String, offload_globals: &OffloadGlobals<'ll>, - gen_kernel: bool, + 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(); @@ -479,7 +505,6 @@ pub(crate) fn gen_define_handling<'ll>( } else { None }; - //dbg!(&transfer_from); let memtransfer_end = add_priv_unnamed_arr(&cx, &format!(".offload_maptypes.{symbol}.end"), &transfer_from); @@ -622,6 +647,7 @@ pub(crate) fn gen_call_handling<'ll, 'tcx>( fn_ty, num_args, s_ident_t, + TransferType::Kernel, ); let values = KernelArgsTy::new( &cx, @@ -666,5 +692,6 @@ pub(crate) fn gen_call_handling<'ll, 'tcx>( 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 7f4956e821f24..74cf676a85cc6 100644 --- a/compiler/rustc_codegen_llvm/src/intrinsic.rs +++ b/compiler/rustc_codegen_llvm/src/intrinsic.rs @@ -1976,7 +1976,8 @@ fn codegen_offload_preload_drop<'ll, 'tcx>( let target_symbol = cx.generate_local_symbol_name(""); dbg!("done for now"); - let offload_data = gen_define_handling(&cx, metadata, target_symbol, offload_globals, false); + 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, @@ -1996,6 +1997,7 @@ fn codegen_offload_preload_drop<'ll, 'tcx>( offload_globals.mapper_fn_ty, 1, offload_globals.ident_t_global, + TransferType::End, ); } @@ -2051,9 +2053,14 @@ fn codegen_offload_preload<'ll, 'tcx>( } }; dbg!("asdf"); - //let target_symbol = "asdf_I_ll_nameclash".to_owned(); let target_symbol = cx.generate_local_symbol_name(""); - let offload_data = gen_define_handling(&cx, metadata, target_symbol, offload_globals, false); + let offload_data = gen_define_handling( + &cx, + metadata, + target_symbol, + offload_globals, + TransferType::NowaitBegin, + ); 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, @@ -2069,13 +2076,20 @@ fn codegen_offload_preload<'ll, 'tcx>( bx, geps, offload_data.memtransfer_begin, - offload_globals.begin_mapper, - offload_globals.mapper_fn_ty, + offload_globals.nowait_begin_mapper, + offload_globals.nowait_mapper_fn_ty, 1, offload_globals.ident_t_global, + TransferType::NowaitBegin, ); } +pub(crate) enum TransferType { + NowaitBegin, + 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 @@ -2143,7 +2157,8 @@ fn codegen_offload<'ll, 'tcx>( return; } }; - let offload_data = gen_define_handling(&cx, &metadata, target_symbol, offload_globals, true); + let offload_data = + gen_define_handling(&cx, &metadata, target_symbol, offload_globals, TransferType::Kernel); gen_call_handling( bx, &offload_data, From f0e7169c7fe6a6127961ef478d886775febcf9d5 Mon Sep 17 00:00:00 2001 From: Manuel Drehwald Date: Thu, 2 Jul 2026 06:17:35 -0700 Subject: [PATCH 15/28] temporarily make transfer sync --- compiler/rustc_codegen_llvm/src/intrinsic.rs | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/compiler/rustc_codegen_llvm/src/intrinsic.rs b/compiler/rustc_codegen_llvm/src/intrinsic.rs index 74cf676a85cc6..84f3c6ddedb60 100644 --- a/compiler/rustc_codegen_llvm/src/intrinsic.rs +++ b/compiler/rustc_codegen_llvm/src/intrinsic.rs @@ -2054,13 +2054,8 @@ fn codegen_offload_preload<'ll, 'tcx>( }; dbg!("asdf"); let target_symbol = cx.generate_local_symbol_name(""); - let offload_data = gen_define_handling( - &cx, - metadata, - target_symbol, - offload_globals, - TransferType::NowaitBegin, - ); + 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, @@ -2080,12 +2075,13 @@ fn codegen_offload_preload<'ll, 'tcx>( offload_globals.nowait_mapper_fn_ty, 1, offload_globals.ident_t_global, - TransferType::NowaitBegin, + TransferType::Begin, ); } pub(crate) enum TransferType { NowaitBegin, + Begin, Kernel, End, } From eef39d91f11ff58d230975bd6f3ae625d29492c4 Mon Sep 17 00:00:00 2001 From: Manuel Drehwald Date: Tue, 7 Jul 2026 08:11:33 -0700 Subject: [PATCH 16/28] wip --- .../src/builder/gpu_offload.rs | 35 +++++++++++++++++++ 1 file changed, 35 insertions(+) diff --git a/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs b/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs index d2d2cde2019e0..1ce8660895109 100644 --- a/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs +++ b/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs @@ -16,6 +16,9 @@ use crate::llvm::AttributePlace::Function; use crate::llvm::{self, Linkage, Type, Value}; use crate::{SimpleCx, attributes}; +//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, @@ -33,6 +36,11 @@ pub(crate) struct OffloadGlobals<'ll> { 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, } impl<'ll> OffloadGlobals<'ll> { @@ -43,6 +51,7 @@ impl<'ll> OffloadGlobals<'ll> { 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); // We want LLVM's openmp-opt pass to pick up and optimize this module, since it covers both // openmp and offload optimizations. @@ -59,6 +68,10 @@ impl<'ll> OffloadGlobals<'ll> { mapper_fn_ty, nowait_mapper_fn_ty, ident_t_global, + taskwait, + taskwait_ty, + threadnum, + threadnum_ty, } } } @@ -203,6 +216,28 @@ fn generate_launcher<'ll>(cx: &CodegenCx<'ll, '_>) -> (&'ll llvm::Value, &'ll ll (tgt_decl, tgt_fn_ty) } +fn generate_sync( + 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 From 9b87e43c85f4a65490b60665abc1ac50a57f1e13 Mon Sep 17 00:00:00 2001 From: Manuel Drehwald Date: Tue, 7 Jul 2026 08:48:06 -0700 Subject: [PATCH 17/28] compiles-rustc --- compiler/rustc_codegen_llvm/src/builder/gpu_helper.rs | 11 +++++++++++ .../rustc_codegen_llvm/src/builder/gpu_offload.rs | 4 +++- compiler/rustc_codegen_llvm/src/intrinsic.rs | 2 ++ 3 files changed, 16 insertions(+), 1 deletion(-) diff --git a/compiler/rustc_codegen_llvm/src/builder/gpu_helper.rs b/compiler/rustc_codegen_llvm/src/builder/gpu_helper.rs index fafb84ac81508..2ca22f67b415c 100644 --- a/compiler/rustc_codegen_llvm/src/builder/gpu_helper.rs +++ b/compiler/rustc_codegen_llvm/src/builder/gpu_helper.rs @@ -1,5 +1,6 @@ use crate::SimpleCx; use crate::builder::Builder; +use crate::builder::gpu_offload::OffloadGlobals; use crate::intrinsic::TransferType; use crate::llvm; use crate::llvm::{Type, Value}; @@ -63,6 +64,7 @@ pub(crate) fn get_geps<'ll, 'tcx>( 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, @@ -80,6 +82,15 @@ pub(crate) fn generate_mapper_call<'ll, 'tcx>( let i32_0 = cx.get_const_i32(0); args.append(&mut vec![i32_0, nullptr, i32_0, nullptr]); } + if matches!(transfer, TransferType::End) { + let a = offload_globals.taskwait; + let b = offload_globals.taskwait_ty; + let c = offload_globals.threadnum; + let d = offload_globals.threadnum_ty; + let tid = builder.call(d, None, None, c, &vec![s_ident_t], None, None); + let args2 = vec![s_ident_t, tid]; + builder.call(b, None, None, a, &args2, None, None); + } builder.call(fn_ty, None, None, fn_to_call, &args, None, None); } diff --git a/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs b/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs index 1ce8660895109..6f8459e8b9dbd 100644 --- a/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs +++ b/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs @@ -216,7 +216,7 @@ fn generate_launcher<'ll>(cx: &CodegenCx<'ll, '_>) -> (&'ll llvm::Value, &'ll ll (tgt_decl, tgt_fn_ty) } -fn generate_sync( +fn generate_sync<'ll>( cx: &CodegenCx<'ll, '_>, ) -> (&'ll llvm::Value, &'ll llvm::Type, &'ll llvm::Value, &'ll llvm::Type) { let tptr = cx.type_ptr(); @@ -677,6 +677,7 @@ pub(crate) fn gen_call_handling<'ll, 'tcx>( generate_mapper_call( builder, geps, + offload_globals, memtransfer_begin, begin_mapper_decl, fn_ty, @@ -722,6 +723,7 @@ pub(crate) fn gen_call_handling<'ll, 'tcx>( generate_mapper_call( builder, geps, + offload_globals, memtransfer_end, end_mapper_decl, fn_ty, diff --git a/compiler/rustc_codegen_llvm/src/intrinsic.rs b/compiler/rustc_codegen_llvm/src/intrinsic.rs index 84f3c6ddedb60..a0e3e6446aa2a 100644 --- a/compiler/rustc_codegen_llvm/src/intrinsic.rs +++ b/compiler/rustc_codegen_llvm/src/intrinsic.rs @@ -1992,6 +1992,7 @@ fn codegen_offload_preload_drop<'ll, 'tcx>( crate::builder::gpu_helper::generate_mapper_call( bx, geps, + offload_globals, offload_data.memtransfer_end, offload_globals.end_mapper, offload_globals.mapper_fn_ty, @@ -2070,6 +2071,7 @@ fn codegen_offload_preload<'ll, 'tcx>( crate::builder::gpu_helper::generate_mapper_call( bx, geps, + offload_globals, offload_data.memtransfer_begin, offload_globals.nowait_begin_mapper, offload_globals.nowait_mapper_fn_ty, From 1928e81ff5559170ad0e3c47b5d852ff4d21189e Mon Sep 17 00:00:00 2001 From: Manuel Drehwald Date: Tue, 7 Jul 2026 09:07:15 -0700 Subject: [PATCH 18/28] compiles rustperf --- compiler/rustc_codegen_llvm/src/builder/gpu_helper.rs | 6 ++++++ compiler/rustc_codegen_llvm/src/intrinsic.rs | 4 ++-- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/compiler/rustc_codegen_llvm/src/builder/gpu_helper.rs b/compiler/rustc_codegen_llvm/src/builder/gpu_helper.rs index 2ca22f67b415c..45d6a5b212a9c 100644 --- a/compiler/rustc_codegen_llvm/src/builder/gpu_helper.rs +++ b/compiler/rustc_codegen_llvm/src/builder/gpu_helper.rs @@ -87,8 +87,14 @@ pub(crate) fn generate_mapper_call<'ll, 'tcx>( 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); diff --git a/compiler/rustc_codegen_llvm/src/intrinsic.rs b/compiler/rustc_codegen_llvm/src/intrinsic.rs index a0e3e6446aa2a..f6e7dd1908f7a 100644 --- a/compiler/rustc_codegen_llvm/src/intrinsic.rs +++ b/compiler/rustc_codegen_llvm/src/intrinsic.rs @@ -2073,8 +2073,8 @@ fn codegen_offload_preload<'ll, 'tcx>( geps, offload_globals, offload_data.memtransfer_begin, - offload_globals.nowait_begin_mapper, - offload_globals.nowait_mapper_fn_ty, + offload_globals.begin_mapper, + offload_globals.mapper_fn_ty, 1, offload_globals.ident_t_global, TransferType::Begin, From b991caf36e0b40e8e57a8733c9d7ab7e3cc7d82e Mon Sep 17 00:00:00 2001 From: Manuel Drehwald Date: Wed, 15 Jul 2026 23:51:58 -0700 Subject: [PATCH 19/28] prepare-async-llvm-offload --- src/llvm-project | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/llvm-project b/src/llvm-project index 08c84e69a84d9..2c1d94c5cf004 160000 --- a/src/llvm-project +++ b/src/llvm-project @@ -1 +1 @@ -Subproject commit 08c84e69a84d95936296dfcab0e38b34100725d5 +Subproject commit 2c1d94c5cf004d5878257f388c7e3f5fd33d5442 From 03b4df1305b04ecc4d3dea36473b9df3f75b6671 Mon Sep 17 00:00:00 2001 From: Manuel Drehwald Date: Mon, 20 Jul 2026 19:34:41 -0700 Subject: [PATCH 20/28] add rust declarations for async to ir --- .../src/builder/gpu_offload.rs | 64 ++++++++++++++++++- 1 file changed, 63 insertions(+), 1 deletion(-) diff --git a/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs b/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs index 6f8459e8b9dbd..f135f8b9a91b8 100644 --- a/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs +++ b/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs @@ -41,6 +41,18 @@ pub(crate) struct OffloadGlobals<'ll> { 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, } impl<'ll> OffloadGlobals<'ll> { @@ -52,11 +64,50 @@ impl<'ll> OffloadGlobals<'ll> { 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); OffloadGlobals { launcher_fn, launcher_ty, @@ -72,6 +123,17 @@ impl<'ll> OffloadGlobals<'ll> { 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, } } } From 7844895b8bd56f4aa30472cb0e11235acce74d9f Mon Sep 17 00:00:00 2001 From: Manuel Drehwald Date: Mon, 20 Jul 2026 22:04:42 -0700 Subject: [PATCH 21/28] generate __rust_offload_async_info --- .../src/builder/gpu_offload.rs | 22 +++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs b/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs index f135f8b9a91b8..be3d6472a566d 100644 --- a/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs +++ b/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs @@ -16,6 +16,7 @@ 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 *); @@ -53,6 +54,7 @@ pub(crate) struct OffloadGlobals<'ll> { pub async_kernel_launcher: &'ll Value, pub async_kernel_launcher_ty: &'ll Type, + pub async_info_global: &'ll llvm::Value, } impl<'ll> OffloadGlobals<'ll> { @@ -108,6 +110,7 @@ impl<'ll> OffloadGlobals<'ll> { 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, @@ -134,10 +137,29 @@ impl<'ll> OffloadGlobals<'ll> { 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 +} + // 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 From 4b278c6271650fb4b1e0bca41d45d80eb31bd853 Mon Sep 17 00:00:00 2001 From: Manuel Drehwald Date: Mon, 20 Jul 2026 22:27:55 -0700 Subject: [PATCH 22/28] wip --- .../src/builder/gpu_helper.rs | 20 ++++++- .../src/builder/gpu_offload.rs | 53 ++++++++++++++++++- 2 files changed, 70 insertions(+), 3 deletions(-) diff --git a/compiler/rustc_codegen_llvm/src/builder/gpu_helper.rs b/compiler/rustc_codegen_llvm/src/builder/gpu_helper.rs index 45d6a5b212a9c..41787943adc98 100644 --- a/compiler/rustc_codegen_llvm/src/builder/gpu_helper.rs +++ b/compiler/rustc_codegen_llvm/src/builder/gpu_helper.rs @@ -1,6 +1,6 @@ use crate::SimpleCx; use crate::builder::Builder; -use crate::builder::gpu_offload::OffloadGlobals; +use crate::builder::gpu_offload::{OffloadGlobals, get_or_create_async_info}; use crate::intrinsic::TransferType; use crate::llvm; use crate::llvm::{Type, Value}; @@ -61,6 +61,23 @@ pub(crate) fn get_geps<'ll, 'tcx>( [gep1, gep2, gep3] } +fn synchronize_async_info<'ll, 'tcx>( + builder: &mut Builder<'_, 'll, 'tcx>, + offload_globals: &OffloadGlobals<'ll>, +) { + let async_info = get_or_create_async_info(builder, offload_globals); + + builder.call( + offload_globals.async_info_synchronize_ty, + None, + None, + offload_globals.async_info_synchronize, + &[async_info], + None, + None, + ); +} + pub(crate) fn generate_mapper_call<'ll, 'tcx>( builder: &mut Builder<'_, 'll, 'tcx>, geps: [&'ll Value; 3], @@ -83,6 +100,7 @@ pub(crate) fn generate_mapper_call<'ll, 'tcx>( 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; diff --git a/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs b/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs index be3d6472a566d..d58b819cf3082 100644 --- a/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs +++ b/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs @@ -160,6 +160,53 @@ fn get_or_create_async_info_global<'ll, 'tcx>(cx: &CodegenCx<'ll, 'tcx>) -> &'ll 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 @@ -733,7 +780,8 @@ pub(crate) fn gen_call_handling<'ll, 'tcx>( 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_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; @@ -789,7 +837,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. @@ -798,6 +846,7 @@ 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) From 4f86b1bce2fe0965226e9486859ad603cbe22992 Mon Sep 17 00:00:00 2001 From: Manuel Drehwald Date: Mon, 20 Jul 2026 23:55:28 -0700 Subject: [PATCH 23/28] try-llvm-async --- src/llvm-project | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/llvm-project b/src/llvm-project index 2c1d94c5cf004..35ab97bae2da4 160000 --- a/src/llvm-project +++ b/src/llvm-project @@ -1 +1 @@ -Subproject commit 2c1d94c5cf004d5878257f388c7e3f5fd33d5442 +Subproject commit 35ab97bae2da45a7a773f9589dfbc66369ee6060 From 238c16f4631e0c2e3a00058f2b3aeb015e3f65f3 Mon Sep 17 00:00:00 2001 From: Manuel Drehwald Date: Tue, 21 Jul 2026 00:31:10 -0700 Subject: [PATCH 24/28] move to async launch, someitmes ~2x slwoer, sometimes 2x faster --- compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs b/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs index d58b819cf3082..92ad1add1b97c 100644 --- a/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs +++ b/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs @@ -779,7 +779,7 @@ pub(crate) fn gen_call_handling<'ll, 'tcx>( let has_dynamic = metadata.iter().any(|m| !matches!(m.payload_size, OffloadSize::Static(_))); - let tgt_decl = offload_globals.launcher_fn; + 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; From 0762f57fccfbd59e7ad9a2fdd282e2474e7a371b Mon Sep 17 00:00:00 2001 From: Manuel Drehwald Date: Tue, 21 Jul 2026 01:31:11 -0700 Subject: [PATCH 25/28] add sync() call to synchronize before drop --- .../src/builder/gpu_helper.rs | 4 +- .../src/builder/gpu_offload.rs | 44 +++++++++---------- compiler/rustc_codegen_llvm/src/intrinsic.rs | 18 ++++++++ .../rustc_hir_analysis/src/check/intrinsic.rs | 2 + compiler/rustc_span/src/symbol.rs | 1 + library/core/src/intrinsics/mod.rs | 5 +++ library/core/src/offload/mod.rs | 12 +++++ 7 files changed, 62 insertions(+), 24 deletions(-) diff --git a/compiler/rustc_codegen_llvm/src/builder/gpu_helper.rs b/compiler/rustc_codegen_llvm/src/builder/gpu_helper.rs index 41787943adc98..5d50025bebab5 100644 --- a/compiler/rustc_codegen_llvm/src/builder/gpu_helper.rs +++ b/compiler/rustc_codegen_llvm/src/builder/gpu_helper.rs @@ -61,7 +61,7 @@ pub(crate) fn get_geps<'ll, 'tcx>( [gep1, gep2, gep3] } -fn synchronize_async_info<'ll, 'tcx>( +pub(crate) fn synchronize_async_info<'ll, 'tcx>( builder: &mut Builder<'_, 'll, 'tcx>, offload_globals: &OffloadGlobals<'ll>, ) { @@ -100,7 +100,7 @@ pub(crate) fn generate_mapper_call<'ll, 'tcx>( args.append(&mut vec![i32_0, nullptr, i32_0, nullptr]); } if matches!(transfer, TransferType::End) { - synchronize_async_info(builder, offload_globals); + //synchronize_async_info(builder, offload_globals); let a = offload_globals.taskwait; let b = offload_globals.taskwait_ty; let c = offload_globals.threadnum; diff --git a/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs b/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs index 92ad1add1b97c..44ca8b22e5385 100644 --- a/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs +++ b/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs @@ -806,17 +806,17 @@ pub(crate) fn gen_call_handling<'ll, 'tcx>( // 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, - offload_globals, - memtransfer_begin, - begin_mapper_decl, - fn_ty, - num_args, - s_ident_t, - TransferType::Kernel, - ); + //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, @@ -853,15 +853,15 @@ pub(crate) fn gen_call_handling<'ll, 'tcx>( // Step 4) let geps = get_geps(builder, ty, ty2, a1, a2, a4, has_dynamic); - generate_mapper_call( - builder, - geps, - offload_globals, - memtransfer_end, - end_mapper_decl, - fn_ty, - num_args, - s_ident_t, - TransferType::Kernel, - ); + //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 f6e7dd1908f7a..e12b4c839e9e7 100644 --- a/compiler/rustc_codegen_llvm/src/intrinsic.rs +++ b/compiler/rustc_codegen_llvm/src/intrinsic.rs @@ -229,6 +229,11 @@ 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; @@ -1926,6 +1931,19 @@ fn offload_bool_arg<'ll, 'tcx>(args: &[OperandRef<'tcx, &'ll llvm::Value>], 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>, diff --git a/compiler/rustc_hir_analysis/src/check/intrinsic.rs b/compiler/rustc_hir_analysis/src/check/intrinsic.rs index ce3e18d41fcbe..ee2b872e3ad0c 100644 --- a/compiler/rustc_hir_analysis/src/check/intrinsic.rs +++ b/compiler/rustc_hir_analysis/src/check/intrinsic.rs @@ -164,6 +164,7 @@ 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 @@ -366,6 +367,7 @@ pub(crate) fn check_intrinsic_type( 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 8a98da74da688..73d61097eaa81 100644 --- a/compiler/rustc_span/src/symbol.rs +++ b/compiler/rustc_span/src/symbol.rs @@ -1425,6 +1425,7 @@ symbols! { of, off, offload, + offload_sync, offload_preload, offload_preload_end, offload_kernel, diff --git a/library/core/src/intrinsics/mod.rs b/library/core/src/intrinsics/mod.rs index 7bf77bd22803a..8de3ffa459679 100644 --- a/library/core/src/intrinsics/mod.rs +++ b/library/core/src/intrinsics/mod.rs @@ -3577,6 +3577,11 @@ pub const fn offload( #[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); diff --git a/library/core/src/offload/mod.rs b/library/core/src/offload/mod.rs index acd24dc99ecc2..5f49e298acb75 100644 --- a/library/core/src/offload/mod.rs +++ b/library/core/src/offload/mod.rs @@ -21,6 +21,18 @@ pub struct Preload<'a, T: ?Sized> { _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 From bf5ff81772f5163b0fd2840fff5004b830bf56b0 Mon Sep 17 00:00:00 2001 From: Manuel Drehwald Date: Tue, 21 Jul 2026 02:41:06 -0700 Subject: [PATCH 26/28] cleanup old task based async --- .../src/builder/gpu_helper.rs | 28 ++++----- .../src/builder/gpu_offload.rs | 61 +++++++++---------- 2 files changed, 44 insertions(+), 45 deletions(-) diff --git a/compiler/rustc_codegen_llvm/src/builder/gpu_helper.rs b/compiler/rustc_codegen_llvm/src/builder/gpu_helper.rs index 5d50025bebab5..0480da45dae1f 100644 --- a/compiler/rustc_codegen_llvm/src/builder/gpu_helper.rs +++ b/compiler/rustc_codegen_llvm/src/builder/gpu_helper.rs @@ -100,20 +100,20 @@ pub(crate) fn generate_mapper_call<'ll, 'tcx>( 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); + 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); } diff --git a/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs b/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs index 44ca8b22e5385..84b1f42dee396 100644 --- a/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs +++ b/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs @@ -38,11 +38,10 @@ pub(crate) struct OffloadGlobals<'ll> { 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 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, @@ -65,7 +64,7 @@ impl<'ll> OffloadGlobals<'ll> { 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); + //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()); @@ -122,10 +121,10 @@ impl<'ll> OffloadGlobals<'ll> { mapper_fn_ty, nowait_mapper_fn_ty, ident_t_global, - taskwait, - taskwait_ty, - threadnum, - threadnum_ty, + //taskwait, + //taskwait_ty, + //threadnum, + //threadnum_ty, async_info_create, async_info_create_ty, @@ -347,27 +346,27 @@ 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) -} +//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 From a054f75450a63b9e69991fe76f675ac91a01fff5 Mon Sep 17 00:00:00 2001 From: Manuel Drehwald Date: Tue, 21 Jul 2026 03:27:28 -0700 Subject: [PATCH 27/28] bugfix --- compiler/rustc_codegen_ssa/src/mir/intrinsic.rs | 1 + 1 file changed, 1 insertion(+) 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!( From c3f20f2f572ea4afd7cce72d3d5c73f837eab690 Mon Sep 17 00:00:00 2001 From: Manuel Drehwald Date: Tue, 21 Jul 2026 04:19:18 -0700 Subject: [PATCH 28/28] try --- .../src/builder/gpu_helper.rs | 23 ++++++++++++++++++- compiler/rustc_codegen_llvm/src/intrinsic.rs | 2 +- 2 files changed, 23 insertions(+), 2 deletions(-) diff --git a/compiler/rustc_codegen_llvm/src/builder/gpu_helper.rs b/compiler/rustc_codegen_llvm/src/builder/gpu_helper.rs index 0480da45dae1f..825f76c22b806 100644 --- a/compiler/rustc_codegen_llvm/src/builder/gpu_helper.rs +++ b/compiler/rustc_codegen_llvm/src/builder/gpu_helper.rs @@ -65,7 +65,22 @@ pub(crate) fn synchronize_async_info<'ll, 'tcx>( builder: &mut Builder<'_, 'll, 'tcx>, offload_globals: &OffloadGlobals<'ll>, ) { - let async_info = get_or_create_async_info(builder, offload_globals); + 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, @@ -76,6 +91,12 @@ pub(crate) fn synchronize_async_info<'ll, 'tcx>( None, None, ); + + builder.br(done_bb); + + unsafe { + llvm::LLVMPositionBuilderAtEnd(&builder.llbuilder, done_bb); + } } pub(crate) fn generate_mapper_call<'ll, 'tcx>( diff --git a/compiler/rustc_codegen_llvm/src/intrinsic.rs b/compiler/rustc_codegen_llvm/src/intrinsic.rs index e12b4c839e9e7..781c941d5c10e 100644 --- a/compiler/rustc_codegen_llvm/src/intrinsic.rs +++ b/compiler/rustc_codegen_llvm/src/intrinsic.rs @@ -1934,7 +1934,7 @@ fn offload_bool_arg<'ll, 'tcx>(args: &[OperandRef<'tcx, &'ll llvm::Value>], idx: fn codegen_offload_sync<'ll, 'tcx>(bx: &mut Builder<'_, 'll, 'tcx>) { let cx = bx.cx; - register_offload(cx); + //register_offload(cx); let offload_globals_ref = cx.offload_globals.borrow(); let offload_globals = match offload_globals_ref.as_ref() {