Add type-check to offload intrinisc calls

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.

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
+}

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
     };
 }

compiler/rustc_hir_typeck/src/typeck_root_ctxt.rs

@@ -57,6 +57,8 @@ pub(crate) struct TypeckRootCtxt<'tcx> {
 
     pub(super) deferred_transmute_checks: RefCell<Vec<(Ty<'tcx>, Ty<'tcx>, HirId)>>,
 
+    pub(super) deferred_offload_checks: RefCell<Vec<(Ty<'tcx>, Ty<'tcx>, Ty<'tcx>, HirId)>>,
+
     pub(super) deferred_asm_checks: RefCell<Vec<(&'tcx hir::InlineAsm<'tcx>, 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()),

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);

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);

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) }

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<Vec<(Ty<'tcx>, 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(),
         }
     }

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" {

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" {

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=<Device or Host=/absolute/path/to/device.bin>
 
 error: using the offload feature requires -C lto=fat
 

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=<Device or Host=/absolute/path/to/device.bin>
 //[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]) {}

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() {}

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`.

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) {}