diff --git a/compiler/rustc_hir_typeck/src/expr.rs b/compiler/rustc_hir_typeck/src/expr.rs index 3bcad2460e78f..385ecc30bf877 100644 --- a/compiler/rustc_hir_typeck/src/expr.rs +++ b/compiler/rustc_hir_typeck/src/expr.rs @@ -614,7 +614,7 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> { } }; - if let ty::FnDef(did, _) = *ty.kind() { + if let ty::FnDef(did, args) = *ty.kind() { let fn_sig = ty.fn_sig(tcx); if tcx.is_intrinsic(did, sym::transmute) { @@ -631,6 +631,14 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> { // be known if explicitly specified via turbofish). self.deferred_transmute_checks.borrow_mut().push((*from, to, expr.hir_id)); } + if tcx.is_intrinsic(did, sym::offload) { + let args = args.skip_binder(); + let f = args.type_at(0); + let t = args.type_at(1); + let r = args.type_at(2); + // Defer offload checks to check generics later once types are fully inferred. + self.deferred_offload_checks.borrow_mut().push((f, t, r, expr.hir_id)); + } if !tcx.features().unsized_fn_params() { // We want to remove some Sized bounds from std functions, // but don't want to expose the removal to stable Rust. diff --git a/compiler/rustc_hir_typeck/src/intrinsicck.rs b/compiler/rustc_hir_typeck/src/intrinsicck.rs index 430f4d828b91f..e009be599e36c 100644 --- a/compiler/rustc_hir_typeck/src/intrinsicck.rs +++ b/compiler/rustc_hir_typeck/src/intrinsicck.rs @@ -140,6 +140,109 @@ fn check_transmute<'tcx>( } } +fn check_offload<'tcx>( + tcx: TyCtxt<'tcx>, + typing_env: ty::TypingEnv<'tcx>, + kernel_ty: Ty<'tcx>, + args_ty: Ty<'tcx>, + ret_ty: Ty<'tcx>, + hir_id: HirId, +) -> Result<(), ErrorGuaranteed> { + let span = tcx.hir_span(hir_id); + let ty::FnDef(kernel_def_id, kernel_args) = *kernel_ty.kind() else { + let err = tcx + .sess + .dcx() + .struct_span_err( + span, + format!("expected a function item for the offload kernel, found `{}`", kernel_ty), + ) + .emit(); + return Err(err); + }; + + let kernel_sig = + tcx.fn_sig(kernel_def_id).instantiate(tcx, kernel_args.skip_binder()).skip_norm_wip(); + let kernel_sig = tcx.instantiate_bound_regions_with_erased(kernel_sig); + + let ty::Tuple(tuple_fields) = *args_ty.kind() else { + let err = tcx + .sess + .dcx() + .struct_span_err( + span, + format!("expected a tuple for the offload arguments, found `{}`", args_ty), + ) + .emit(); + return Err(err); + }; + + if kernel_sig.inputs().len() != tuple_fields.len() { + let err = tcx + .sess + .dcx() + .struct_span_err( + span, + format!( + "offload kernel expects {} arguments, but {} arguments were provided", + kernel_sig.inputs().len(), + tuple_fields.len() + ), + ) + .emit(); + return Err(err); + } + + let normalize = |ty| { + if let Ok(ty) = tcx.try_normalize_erasing_regions(typing_env, Unnormalized::new_wip(ty)) { + ty + } else { + Ty::new_error_with_message( + tcx, + span, + format!("tried to normalize non-wf type {ty:#?} in check_offload"), + ) + } + }; + + let mut result = Ok(()); + + for (i, (&input_ty, arg_ty)) in kernel_sig.inputs().iter().zip(tuple_fields.iter()).enumerate() + { + let norm_input_ty = normalize(input_ty); + let norm_arg_ty = normalize(arg_ty); + if norm_input_ty != norm_arg_ty { + let err = tcx + .sess + .dcx() + .struct_span_err( + span, + format!( + "type mismatch in offload kernel argument {}: expected `{}`, found `{}`", + i, norm_input_ty, norm_arg_ty + ), + ) + .emit(); + result = Err(err); + } + } + + let norm_kernel_ret = normalize(kernel_sig.output()); + let norm_offload_ret = normalize(ret_ty); + if norm_kernel_ret != norm_offload_ret { + let err = tcx.sess.dcx().struct_span_err( + span, + format!( + "offload kernel return type mismatch: kernel returns `{}`, but offload call expects `{}`", + norm_kernel_ret, norm_offload_ret + ) + ).emit(); + result = Err(err); + } + + result +} + pub(crate) fn check_transmutes(tcx: TyCtxt<'_>, owner: LocalDefId) -> Result<(), ErrorGuaranteed> { assert!(!tcx.is_typeck_child(owner.to_def_id())); let typeck_results = tcx.typeck(owner); @@ -154,3 +257,18 @@ pub(crate) fn check_transmutes(tcx: TyCtxt<'_>, owner: LocalDefId) -> Result<(), } result } + +pub(crate) fn check_offloads(tcx: TyCtxt<'_>, owner: LocalDefId) -> Result<(), ErrorGuaranteed> { + assert!(!tcx.is_typeck_child(owner.to_def_id())); + let typeck_results = tcx.typeck(owner); + if let Some(e) = typeck_results.tainted_by_errors { + return Err(e); + }; + + let typing_env = ty::TypingEnv::codegen(tcx, owner); + let mut result = Ok(()); + for &(kernel_ty, args_ty, ret_ty, hir_id) in &typeck_results.offloads_to_check { + result = result.and(check_offload(tcx, typing_env, kernel_ty, args_ty, ret_ty, hir_id)); + } + result +} diff --git a/compiler/rustc_hir_typeck/src/lib.rs b/compiler/rustc_hir_typeck/src/lib.rs index dce7f0bd67794..5c5bf77609ce2 100644 --- a/compiler/rustc_hir_typeck/src/lib.rs +++ b/compiler/rustc_hir_typeck/src/lib.rs @@ -721,6 +721,7 @@ pub fn provide(providers: &mut Providers) { typeck_root, used_trait_imports, check_transmutes: intrinsicck::check_transmutes, + check_offloads: intrinsicck::check_offloads, ..*providers }; } diff --git a/compiler/rustc_hir_typeck/src/typeck_root_ctxt.rs b/compiler/rustc_hir_typeck/src/typeck_root_ctxt.rs index e4dcead1f7954..a475d073f452d 100644 --- a/compiler/rustc_hir_typeck/src/typeck_root_ctxt.rs +++ b/compiler/rustc_hir_typeck/src/typeck_root_ctxt.rs @@ -57,6 +57,8 @@ pub(crate) struct TypeckRootCtxt<'tcx> { pub(super) deferred_transmute_checks: RefCell, Ty<'tcx>, HirId)>>, + pub(super) deferred_offload_checks: RefCell, Ty<'tcx>, Ty<'tcx>, HirId)>>, + pub(super) deferred_asm_checks: RefCell, HirId)>>, pub(super) deferred_repeat_expr_checks: @@ -97,6 +99,7 @@ impl<'tcx> TypeckRootCtxt<'tcx> { deferred_call_resolutions: RefCell::new(Default::default()), deferred_cast_checks: RefCell::new(Vec::new()), deferred_transmute_checks: RefCell::new(Vec::new()), + deferred_offload_checks: RefCell::new(Vec::new()), deferred_asm_checks: RefCell::new(Vec::new()), deferred_repeat_expr_checks: RefCell::new(Vec::new()), diverging_type_vars: RefCell::new(Default::default()), diff --git a/compiler/rustc_hir_typeck/src/writeback.rs b/compiler/rustc_hir_typeck/src/writeback.rs index 4161975d88ea9..7b1f38f882747 100644 --- a/compiler/rustc_hir_typeck/src/writeback.rs +++ b/compiler/rustc_hir_typeck/src/writeback.rs @@ -76,6 +76,7 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> { wbcx.visit_user_provided_sigs(); wbcx.visit_coroutine_interior(); wbcx.visit_transmutes(); + wbcx.visit_offloads(); wbcx.visit_offset_of_container_types(); wbcx.visit_potentially_region_dependent_goals(); @@ -544,6 +545,21 @@ impl<'cx, 'tcx> WritebackCx<'cx, 'tcx> { } } + fn visit_offloads(&mut self) { + let tcx = self.tcx(); + let fcx_typeck_results = self.fcx.typeck_results.borrow(); + assert_eq!(fcx_typeck_results.hir_owner, self.typeck_results.hir_owner); + for &(kernel_ty, args_ty, ret_ty, hir_id) in + self.fcx.deferred_offload_checks.borrow().iter() + { + let span = tcx.hir_span(hir_id); + let kernel_ty = self.resolve(kernel_ty, &span); + let args_ty = self.resolve(args_ty, &span); + let ret_ty = self.resolve(ret_ty, &span); + self.typeck_results.offloads_to_check.push((kernel_ty, args_ty, ret_ty, hir_id)); + } + } + fn visit_opaque_types_next(&mut self) { let mut fcx_typeck_results = self.fcx.typeck_results.borrow_mut(); assert_eq!(fcx_typeck_results.hir_owner, self.typeck_results.hir_owner); diff --git a/compiler/rustc_interface/src/passes.rs b/compiler/rustc_interface/src/passes.rs index 2f32a6b208b6c..a914fa997106b 100644 --- a/compiler/rustc_interface/src/passes.rs +++ b/compiler/rustc_interface/src/passes.rs @@ -1153,6 +1153,7 @@ fn run_required_analyses(tcx: TyCtxt<'_>) { if not_typeck_child { tcx.ensure_ok().mir_borrowck(def_id); tcx.ensure_ok().check_transmutes(def_id); + tcx.ensure_ok().check_offloads(def_id); } tcx.ensure_ok().has_ffi_unwind_calls(def_id); tcx.ensure_ok().check_liveness(def_id); diff --git a/compiler/rustc_middle/src/queries.rs b/compiler/rustc_middle/src/queries.rs index cbd54ec959c6a..50f59c2c95c5a 100644 --- a/compiler/rustc_middle/src/queries.rs +++ b/compiler/rustc_middle/src/queries.rs @@ -1123,6 +1123,11 @@ rustc_queries! { desc { "check transmute calls inside `{}`", tcx.def_path_str(key) } } + /// Unsafety-check this `LocalDefId`. + query check_offloads(key: LocalDefId) -> Result<(), ErrorGuaranteed> { + desc { "check offload calls inside `{}`", tcx.def_path_str(key) } + } + /// Unsafety-check this `LocalDefId`. query check_unsafety(key: LocalDefId) { desc { "unsafety-checking `{}`", tcx.def_path_str(key) } diff --git a/compiler/rustc_middle/src/ty/typeck_results.rs b/compiler/rustc_middle/src/ty/typeck_results.rs index e11bc6d38495b..7624d860aa1ac 100644 --- a/compiler/rustc_middle/src/ty/typeck_results.rs +++ b/compiler/rustc_middle/src/ty/typeck_results.rs @@ -223,6 +223,9 @@ pub struct TypeckResults<'tcx> { /// computation. pub transmutes_to_check: Vec<(Ty<'tcx>, Ty<'tcx>, HirId)>, + /// Stores the types involved in calls to `offload` intrinsic. + pub offloads_to_check: Vec<(Ty<'tcx>, Ty<'tcx>, Ty<'tcx>, HirId)>, + /// Container types and field indices of `offset_of!` expressions offset_of_data: ItemLocalMap, VariantIdx, FieldIdx)>>, } @@ -256,6 +259,7 @@ impl<'tcx> TypeckResults<'tcx> { potentially_region_dependent_goals: Default::default(), closure_size_eval: Default::default(), transmutes_to_check: Default::default(), + offloads_to_check: Default::default(), offset_of_data: Default::default(), } } diff --git a/tests/codegen-llvm/gpu_offload/scalar_host.rs b/tests/codegen-llvm/gpu_offload/scalar_host.rs index 3470761be06c8..66c910c439e46 100644 --- a/tests/codegen-llvm/gpu_offload/scalar_host.rs +++ b/tests/codegen-llvm/gpu_offload/scalar_host.rs @@ -13,21 +13,22 @@ // CHECK: define{{( dso_local)?}} void @main() // CHECK-NOT: define // CHECK: %addr = alloca i64, align 8 -// CHECK: store double 4.200000e+01, ptr [[TMP:%[^,]+]], align 8 -// CHECK: [[VAL:%[0-9]+]] = load double, ptr [[TMP]], align 8 -// CHECK: store double [[VAL]], ptr %addr, align 8 -// CHECK: %1 = getelementptr inbounds nuw i8, ptr %.offload_baseptrs, i64 8 -// CHECK-NEXT: store double [[VAL]], ptr %1, align 8 -// CHECK-NEXT: %2 = getelementptr inbounds nuw i8, ptr %.offload_ptrs, i64 8 -// CHECK-NEXT: store ptr %addr, ptr %2, align 8 +// CHECK: store float 4.200000e+01, ptr [[TMP:%[^,]+]], align 4 +// CHECK: [[VAL:%[0-9]+]] = load i32, ptr [[TMP]], align 4 +// CHECK: [[VAL_I64:%[0-9]+]] = zext i32 [[VAL]] to i64 +// CHECK: store i64 [[VAL_I64]], ptr %addr, align 8 +// CHECK: [[REG_GEP1:%[^,]+]] = getelementptr inbounds nuw i8, ptr %.offload_baseptrs, i64 8 +// CHECK-NEXT: store i64 [[VAL_I64]], ptr [[REG_GEP1]], align 8 +// CHECK-NEXT: [[REG_GEP2:%[^,]+]] = getelementptr inbounds nuw i8, ptr %.offload_ptrs, i64 8 +// CHECK-NEXT: store ptr %addr, ptr [[REG_GEP2]], align 8 // CHECK-NEXT: call void @__tgt_target_data_begin_mapper #[unsafe(no_mangle)] fn main() { - let mut x = 0.0; - let k = core::hint::black_box(42.0); + let mut x = 0.0f32; + let k = core::hint::black_box(42.0f32); - core::intrinsics::offload::<_, _, ()>(foo, [1, 1, 1], [1, 1, 1], 0, (&mut x, k)); + core::intrinsics::offload::<_, _, ()>(foo, [1, 1, 1], [1, 1, 1], 0, (&mut x as *mut f32, k)); } unsafe extern "C" { diff --git a/tests/codegen-llvm/gpu_offload/slice_host.rs b/tests/codegen-llvm/gpu_offload/slice_host.rs index d4157d24e03dd..0f27821ef765c 100644 --- a/tests/codegen-llvm/gpu_offload/slice_host.rs +++ b/tests/codegen-llvm/gpu_offload/slice_host.rs @@ -26,8 +26,8 @@ #[unsafe(no_mangle)] fn main() { - let mut x = [0.0, 0.0, 0.0, 0.0]; - core::intrinsics::offload::<_, _, ()>(foo, [1, 1, 1], [1, 1, 1], 0, ((&mut x) as &mut [f64],)); + let mut x = [0.0f32, 0.0, 0.0, 0.0]; + core::intrinsics::offload::<_, _, ()>(foo, [1, 1, 1], [1, 1, 1], 0, ((&mut x) as &mut [f32],)); } unsafe extern "C" { diff --git a/tests/ui/offload/check_config.fail.stderr b/tests/ui/offload/check_config.fail.stderr index a9162ed926cb0..f1c73687889fd 100644 --- a/tests/ui/offload/check_config.fail.stderr +++ b/tests/ui/offload/check_config.fail.stderr @@ -1,4 +1,4 @@ -error: using the offload feature requires -Z offload=Enable +error: using the offload feature requires -Z offload= error: using the offload feature requires -C lto=fat diff --git a/tests/ui/offload/check_config.rs b/tests/ui/offload/check_config.rs index 667c6d9788bae..69afe65a308b4 100644 --- a/tests/ui/offload/check_config.rs +++ b/tests/ui/offload/check_config.rs @@ -3,10 +3,10 @@ //@ needs-enzyme //@[pass] build-pass //@[fail] build-fail -//@[pass] compile-flags: -Zunstable-options -Zoffload=Enable -Clto=fat --emit=metadata +//@[pass] compile-flags: -Zunstable-options -Zoffload=Device -Clto=fat --emit=metadata //@[fail] compile-flags: -Clto=thin -//[fail]~? ERROR: using the offload feature requires -Z offload=Enable +//[fail]~? ERROR: using the offload feature requires -Z offload= //[fail]~? ERROR: using the offload feature requires -C lto=fat #![feature(core_intrinsics)] @@ -17,7 +17,7 @@ fn main() { } fn kernel_1(x: &mut [f32; 256]) { - core::intrinsics::offload(_kernel_1, (x,)) + core::intrinsics::offload(_kernel_1, [1, 1, 1], [1, 1, 1], 0, (x,)) } fn _kernel_1(x: &mut [f32; 256]) {} diff --git a/tests/ui/offload/non_tuple_args.rs b/tests/ui/offload/non_tuple_args.rs new file mode 100644 index 0000000000000..0a07c99a26d34 --- /dev/null +++ b/tests/ui/offload/non_tuple_args.rs @@ -0,0 +1,11 @@ +//@ compile-flags: -Zunstable-options -Zoffload=Device -Clto=fat + +#![feature(core_intrinsics)] + +fn main() { + // args_ty is not a tuple + core::intrinsics::offload::<_, _, ()>(kernel_0, [1, 1, 1], [1, 1, 1], 0, 42); + //~^ ERROR `{integer}` is not a tuple +} + +fn kernel_0() {} diff --git a/tests/ui/offload/non_tuple_args.stderr b/tests/ui/offload/non_tuple_args.stderr new file mode 100644 index 0000000000000..8b59d6828c6f2 --- /dev/null +++ b/tests/ui/offload/non_tuple_args.stderr @@ -0,0 +1,12 @@ +error[E0277]: `{integer}` is not a tuple + --> $DIR/non_tuple_args.rs:7:36 + | +LL | core::intrinsics::offload::<_, _, ()>(kernel_0, [1, 1, 1], [1, 1, 1], 0, 42); + | ^ the nightly-only, unstable trait `std::marker::Tuple` is not implemented for `{integer}` + | +note: required by a bound in `offload` + --> $SRC_DIR/core/src/intrinsics/mod.rs:LL:COL + +error: aborting due to 1 previous error + +For more information about this error, try `rustc --explain E0277`. diff --git a/tests/ui/offload/type_mismatch.rs b/tests/ui/offload/type_mismatch.rs new file mode 100644 index 0000000000000..4079444a0aff1 --- /dev/null +++ b/tests/ui/offload/type_mismatch.rs @@ -0,0 +1,31 @@ +//@ compile-flags: -Zunstable-options -Zoffload=Device -Clto=fat + +#![feature(core_intrinsics)] + +fn main() { + // kernel_ty is not a function item + let not_fn = 42; + core::intrinsics::offload::<_, _, ()>(not_fn, [1, 1, 1], [1, 1, 1], 0, ()); + //~^ ERROR expected a function item for the offload kernel, found `i32` + + // argument count mismatch + core::intrinsics::offload::<_, _, ()>(kernel_1, [1, 1, 1], [1, 1, 1], 0, ()); + //~^ ERROR offload kernel expects 1 arguments, but 0 arguments were provided + + // argument type mismatch + core::intrinsics::offload::<_, _, ()>(kernel_1, [1, 1, 1], [1, 1, 1], 0, (42.0f64,)); + //~^ ERROR type mismatch in offload kernel argument 0: expected `f32`, found `f64` + + // return type mismatch + let _: f64 = core::intrinsics::offload::<_, _, f64>(kernel_0, [1, 1, 1], [1, 1, 1], 0, ()); + //~^ ERROR offload kernel return type mismatch: kernel returns `()`, but offload call expects `f64` + + // multiple argument type mismatch + core::intrinsics::offload::<_, _, ()>(kernel_2, [1, 1, 1], [1, 1, 1], 0, (42.0f64, 42.0f64)); + //~^ ERROR type mismatch in offload kernel argument 0: expected `f32`, found `f64` + //~| ERROR type mismatch in offload kernel argument 1: expected `f32`, found `f64` +} + +fn kernel_0() {} +fn kernel_1(_x: f32) {} +fn kernel_2(_x: f32, _y: f32) {} diff --git a/tests/ui/offload/type_mismatch.stderr b/tests/ui/offload/type_mismatch.stderr new file mode 100644 index 0000000000000..8cf160ca09486 --- /dev/null +++ b/tests/ui/offload/type_mismatch.stderr @@ -0,0 +1,38 @@ +error: expected a function item for the offload kernel, found `i32` + --> $DIR/type_mismatch.rs:8:5 + | +LL | core::intrinsics::offload::<_, _, ()>(not_fn, [1, 1, 1], [1, 1, 1], 0, ()); + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +error: offload kernel expects 1 arguments, but 0 arguments were provided + --> $DIR/type_mismatch.rs:12:5 + | +LL | core::intrinsics::offload::<_, _, ()>(kernel_1, [1, 1, 1], [1, 1, 1], 0, ()); + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +error: type mismatch in offload kernel argument 0: expected `f32`, found `f64` + --> $DIR/type_mismatch.rs:16:5 + | +LL | core::intrinsics::offload::<_, _, ()>(kernel_1, [1, 1, 1], [1, 1, 1], 0, (42.0f64,)); + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +error: offload kernel return type mismatch: kernel returns `()`, but offload call expects `f64` + --> $DIR/type_mismatch.rs:20:18 + | +LL | let _: f64 = core::intrinsics::offload::<_, _, f64>(kernel_0, [1, 1, 1], [1, 1, 1], 0, ()); + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +error: type mismatch in offload kernel argument 0: expected `f32`, found `f64` + --> $DIR/type_mismatch.rs:24:5 + | +LL | core::intrinsics::offload::<_, _, ()>(kernel_2, [1, 1, 1], [1, 1, 1], 0, (42.0f64, 42.0f64)); + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +error: type mismatch in offload kernel argument 1: expected `f32`, found `f64` + --> $DIR/type_mismatch.rs:24:5 + | +LL | core::intrinsics::offload::<_, _, ()>(kernel_2, [1, 1, 1], [1, 1, 1], 0, (42.0f64, 42.0f64)); + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +error: aborting due to 6 previous errors +