Add checks for gpu-kernel calling conv

compiler/rustc_abi/src/extern_abi.rs

@@ -67,7 +67,6 @@ pub enum ExternAbi {
 
     /* gpu */
     /// An entry-point function called by the GPU's host
-    // FIXME: should not be callable from Rust on GPU targets, is for host's use only
     GpuKernel,
     /// An entry-point function called by the GPU's host
     // FIXME: why do we have two of these?

compiler/rustc_ast_passes/src/ast_validation.rs

@@ -401,9 +401,16 @@ impl<'a> AstValidator<'a> {
                     | CanonAbi::Rust
                     | CanonAbi::RustCold
                     | CanonAbi::Arm(_)
-                    | CanonAbi::GpuKernel
                     | CanonAbi::X86(_) => { /* nothing to check */ }
 
+                    CanonAbi::GpuKernel => {
+                        // An `extern "gpu-kernel"` function cannot be `async` and/or `gen`.
+                        self.reject_coroutine(abi, sig);
+
+                        // An `extern "gpu-kernel"` function cannot return a value.
+                        self.reject_return(abi, sig);
+                    }
+
                     CanonAbi::Custom => {
                         // An `extern "custom"` function must be unsafe.
                         self.reject_safe_fn(abi, ctxt, sig);
@@ -433,18 +440,7 @@ impl<'a> AstValidator<'a> {
                                 self.dcx().emit_err(errors::AbiX86Interrupt { spans, param_count });
                             }
 
-                            if let FnRetTy::Ty(ref ret_ty) = sig.decl.output
-                                && match &ret_ty.kind {
-                                    TyKind::Never => false,
-                                    TyKind::Tup(tup) if tup.is_empty() => false,
-                                    _ => true,
-                                }
-                            {
-                                self.dcx().emit_err(errors::AbiMustNotHaveReturnType {
-                                    span: ret_ty.span,
-                                    abi,
-                                });
-                            }
+                            self.reject_return(abi, sig);
                         } else {
                             // An `extern "interrupt"` function must have type `fn()`.
                             self.reject_params_or_return(abi, ident, sig);
@@ -496,6 +492,18 @@ impl<'a> AstValidator<'a> {
         }
     }
 
+    fn reject_return(&self, abi: ExternAbi, sig: &FnSig) {
+        if let FnRetTy::Ty(ref ret_ty) = sig.decl.output
+            && match &ret_ty.kind {
+                TyKind::Never => false,
+                TyKind::Tup(tup) if tup.is_empty() => false,
+                _ => true,
+            }
+        {
+            self.dcx().emit_err(errors::AbiMustNotHaveReturnType { span: ret_ty.span, abi });
+        }
+    }
+
     fn reject_params_or_return(&self, abi: ExternAbi, ident: &Ident, sig: &FnSig) {
         let mut spans: Vec<_> = sig.decl.inputs.iter().map(|p| p.span).collect();
         if let FnRetTy::Ty(ref ret_ty) = sig.decl.output

compiler/rustc_hir_typeck/messages.ftl

@@ -133,6 +133,10 @@ hir_typeck_fru_suggestion =
 hir_typeck_functional_record_update_on_non_struct =
     functional record update syntax requires a struct
 
+hir_typeck_gpu_kernel_abi_cannot_be_called =
+    functions with the "gpu-kernel" ABI cannot be called
+    .note = an `extern "gpu-kernel"` function can only be launched on the GPU through an API
+
 hir_typeck_help_set_edition_cargo = set `edition = "{$edition}"` in `Cargo.toml`
 hir_typeck_help_set_edition_standalone = pass `--edition {$edition}` to `rustc`
 

compiler/rustc_hir_typeck/src/callee.rs

@@ -169,27 +169,27 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
             }
         };
 
-        let valid = match canon_abi {
+        match canon_abi {
             // Rust doesn't know how to call functions with this ABI.
-            CanonAbi::Custom => false,
-
-            // These is an entry point for the host, and cannot be called on the GPU.
-            CanonAbi::GpuKernel => false,
-
+            CanonAbi::Custom
             // The interrupt ABIs should only be called by the CPU. They have complex
             // pre- and postconditions, and can use non-standard instructions like `iret` on x86.
-            CanonAbi::Interrupt(_) => false,
+            | CanonAbi::Interrupt(_) => {
+                let err = crate::errors::AbiCannotBeCalled { span, abi };
+                self.tcx.dcx().emit_err(err);
+            }
+
+            // This is an entry point for the host, and cannot be called directly.
+            CanonAbi::GpuKernel => {
+                let err = crate::errors::GpuKernelAbiCannotBeCalled { span };
+                self.tcx.dcx().emit_err(err);
+            }
 
             CanonAbi::C
             | CanonAbi::Rust
             | CanonAbi::RustCold
             | CanonAbi::Arm(_)
-            | CanonAbi::X86(_) => true,
-        };
-
-        if !valid {
-            let err = crate::errors::AbiCannotBeCalled { span, abi };
-            self.tcx.dcx().emit_err(err);
+            | CanonAbi::X86(_) => {}
         }
     }
 

compiler/rustc_hir_typeck/src/errors.rs

@@ -1198,6 +1198,14 @@ pub(crate) struct AbiCannotBeCalled {
     pub abi: ExternAbi,
 }
 
+#[derive(Diagnostic)]
+#[diag(hir_typeck_gpu_kernel_abi_cannot_be_called)]
+pub(crate) struct GpuKernelAbiCannotBeCalled {
+    #[primary_span]
+    #[note]
+    pub span: Span,
+}
+
 #[derive(Diagnostic)]
 #[diag(hir_typeck_const_continue_bad_label)]
 pub(crate) struct ConstContinueBadLabel {

compiler/rustc_lint/messages.ftl

@@ -470,6 +470,9 @@ lint_improper_ctypes_union_non_exhaustive = this union is non-exhaustive
 
 lint_improper_ctypes_unsafe_binder = unsafe binders are incompatible with foreign function interfaces
 
+lint_improper_gpu_kernel_arg = passing type `{$ty}` to a function with "gpu-kernel" ABI may have unexpected behavior
+    .help = use primitive types and raw pointers to get reliable behavior
+
 lint_int_to_ptr_transmutes = transmuting an integer to a pointer creates a pointer without provenance
     .note = this is dangerous because dereferencing the resulting pointer is undefined behavior
     .note_exposed_provenance = exposed provenance semantics can be used to create a pointer based on some previously exposed provenance
@@ -597,6 +600,10 @@ lint_mismatched_lifetime_syntaxes_suggestion_mixed =
 lint_mismatched_lifetime_syntaxes_suggestion_mixed_only_paths =
     use `'_` for type paths
 
+lint_missing_gpu_kernel_export_name = function with the "gpu-kernel" ABI has a mangled name
+    .note = mangled names make it hard to find the kernel, this is usually not intended
+    .help = use `unsafe(no_mangle)` or `unsafe(export_name = "<name>")`
+
 lint_mixed_script_confusables =
     the usage of Script Group `{$set}` in this crate consists solely of mixed script confusables
     .includes_note = the usage includes {$includes}

compiler/rustc_lint/src/gpukernel_abi.rs

@@ -0,0 +1,177 @@
+use std::iter;
+
+use rustc_abi::ExternAbi;
+use rustc_hir::attrs::AttributeKind;
+use rustc_hir::{self as hir, find_attr};
+use rustc_middle::ty;
+use rustc_session::{declare_lint, declare_lint_pass};
+use rustc_span::Span;
+use rustc_span::def_id::LocalDefId;
+
+use crate::lints::{ImproperGpuKernelArg, MissingGpuKernelExportName};
+use crate::{LateContext, LateLintPass, LintContext};
+
+declare_lint! {
+    /// The `improper_gpu_kernel_arg` lint detects incorrect use of types in `gpu-kernel`
+    /// arguments.
+    ///
+    /// ### Example
+    ///
+    /// ```rust,ignore (fails on non-GPU targets)
+    /// #[unsafe(no_mangle)]
+    /// extern "gpu-kernel" fn kernel(_: [i32; 10]) {}
+    /// ```
+    ///
+    /// This will produce:
+    ///
+    /// ```text
+    /// warning: passing type `[i32; 10]` to a function with "gpu-kernel" ABI may have unexpected behavior
+    ///  --> t.rs:2:34
+    ///   |
+    /// 2 | extern "gpu-kernel" fn kernel(_: [i32; 10]) {}
+    ///   |                                  ^^^^^^^^^
+    ///   |
+    ///   = help: use primitive types and raw pointers to get reliable behavior
+    ///   = note: `#[warn(improper_gpu_kernel_arg)]` on by default
+    /// ```
+    ///
+    /// ### Explanation
+    ///
+    /// The compiler has several checks to verify that types used as arguments in `gpu-kernel`
+    /// functions follow certain rules to ensure proper compatibility with the foreign interfaces.
+    /// This lint is issued when it detects a probable mistake in a signature.
+    IMPROPER_GPU_KERNEL_ARG,
+    Warn,
+    "simple arguments of gpu-kernel functions"
+}
+
+declare_lint! {
+    /// The `missing_gpu_kernel_export_name` lint detects `gpu-kernel` functions that have a mangled name.
+    ///
+    /// ### Example
+    ///
+    /// ```rust,ignore (fails on non-GPU targets)
+    /// extern "gpu-kernel" fn kernel() { }
+    /// ```
+    ///
+    /// This will produce:
+    ///
+    /// ```text
+    /// warning: function with the "gpu-kernel" ABI has a mangled name
+    ///  --> t.rs:1:1
+    ///   |
+    /// 1 | extern "gpu-kernel" fn kernel() {}
+    ///   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+    ///   |
+    ///   = help: use `unsafe(no_mangle)` or `unsafe(export_name = "<name>")`
+    ///   = note: mangled names make it hard to find the kernel, this is usually not intended
+    ///   = note: `#[warn(missing_gpu_kernel_export_name)]` on by default
+    /// ```
+    ///
+    /// ### Explanation
+    ///
+    /// `gpu-kernel` functions are usually searched by name in the compiled file.
+    /// A mangled name is usually unintentional as it would need to be searched by the mangled name.
+    ///
+    /// To use an unmangled name for the kernel, either `no_mangle` or `export_name` can be used.
+    /// ```rust,ignore (fails on non-GPU targets)
+    /// // Can be found by the name "kernel"
+    /// #[unsafe(no_mangle)]
+    /// extern "gpu-kernel" fn kernel() { }
+    ///
+    /// // Can be found by the name "new_name"
+    /// #[unsafe(export_name = "new_name")]
+    /// extern "gpu-kernel" fn other_kernel() { }
+    /// ```
+    MISSING_GPU_KERNEL_EXPORT_NAME,
+    Warn,
+    "mangled gpu-kernel function"
+}
+
+declare_lint_pass!(ImproperGpuKernelLint => [
+    IMPROPER_GPU_KERNEL_ARG,
+    MISSING_GPU_KERNEL_EXPORT_NAME,
+]);
+
+/// `ImproperGpuKernelLint` checks `gpu-kernel` function definitions:
+///
+/// - `extern "gpu-kernel" fn` arguments should be simple.
+/// - `extern "gpu-kernel" fn` should have an unmangled name.
+impl<'tcx> LateLintPass<'tcx> for ImproperGpuKernelLint {
+    fn check_fn(
+        &mut self,
+        cx: &LateContext<'tcx>,
+        kind: hir::intravisit::FnKind<'tcx>,
+        decl: &'tcx hir::FnDecl<'_>,
+        _: &'tcx hir::Body<'_>,
+        span: Span,
+        id: LocalDefId,
+    ) {
+        use hir::intravisit::FnKind;
+
+        let abi = match kind {
+            FnKind::ItemFn(_, _, header, ..) => header.abi,
+            FnKind::Method(_, sig, ..) => sig.header.abi,
+            _ => return,
+        };
+
+        if abi != ExternAbi::GpuKernel {
+            return;
+        }
+
+        let sig = cx.tcx.fn_sig(id).instantiate_identity();
+        let sig = cx.tcx.instantiate_bound_regions_with_erased(sig);
+
+        for (input_ty, input_hir) in iter::zip(sig.inputs(), decl.inputs) {
+            let is_valid_arg = match input_ty.kind() {
+                ty::Bool
+                | ty::Char
+                | ty::Int(_)
+                | ty::Uint(_)
+                | ty::Float(_)
+                | ty::RawPtr(_, _) => true,
+
+                ty::Adt(_, _)
+                | ty::Alias(_, _)
+                | ty::Array(_, _)
+                | ty::Bound(_, _)
+                | ty::Closure(_, _)
+                | ty::Coroutine(_, _)
+                | ty::CoroutineClosure(_, _)
+                | ty::CoroutineWitness(..)
+                | ty::Dynamic(_, _)
+                | ty::Error(_)
+                | ty::FnDef(_, _)
+                | ty::FnPtr(..)
+                | ty::Foreign(_)
+                | ty::Infer(_)
+                | ty::Never
+                | ty::Param(_)
+                | ty::Pat(_, _)
+                | ty::Placeholder(_)
+                | ty::Ref(_, _, _)
+                | ty::Slice(_)
+                | ty::Str
+                | ty::Tuple(_)
+                | ty::UnsafeBinder(_) => false,
+            };
+
+            if !is_valid_arg {
+                cx.tcx.emit_node_span_lint(
+                    IMPROPER_GPU_KERNEL_ARG,
+                    input_hir.hir_id,
+                    input_hir.span,
+                    ImproperGpuKernelArg { ty: *input_ty },
+                );
+            }
+        }
+
+        // Check for no_mangle/export_name, so the kernel can be found when querying the compiled object for the kernel function by name
+        if !find_attr!(
+            cx.tcx.get_all_attrs(id),
+            AttributeKind::NoMangle(..) | AttributeKind::ExportName { .. }
+        ) {
+            cx.emit_span_lint(MISSING_GPU_KERNEL_EXPORT_NAME, span, MissingGpuKernelExportName);
+        }
+    }
+}

compiler/rustc_lint/src/lib.rs

@@ -47,6 +47,7 @@ mod expect;
 mod for_loops_over_fallibles;
 mod foreign_modules;
 mod function_cast_as_integer;
+mod gpukernel_abi;
 mod if_let_rescope;
 mod impl_trait_overcaptures;
 mod interior_mutable_consts;
@@ -93,6 +94,7 @@ use drop_forget_useless::*;
 use enum_intrinsics_non_enums::EnumIntrinsicsNonEnums;
 use for_loops_over_fallibles::*;
 use function_cast_as_integer::*;
+use gpukernel_abi::*;
 use if_let_rescope::IfLetRescope;
 use impl_trait_overcaptures::ImplTraitOvercaptures;
 use interior_mutable_consts::*;
@@ -197,6 +199,7 @@ late_lint_methods!(
             DerefIntoDynSupertrait: DerefIntoDynSupertrait,
             DropForgetUseless: DropForgetUseless,
             ImproperCTypesLint: ImproperCTypesLint,
+            ImproperGpuKernelLint: ImproperGpuKernelLint,
             InvalidFromUtf8: InvalidFromUtf8,
             VariantSizeDifferences: VariantSizeDifferences,
             PathStatements: PathStatements,

compiler/rustc_lint/src/lints.rs

@@ -2008,6 +2008,19 @@ impl<'a> LintDiagnostic<'a, ()> for ImproperCTypes<'_> {
     }
 }
 
+#[derive(LintDiagnostic)]
+#[diag(lint_improper_gpu_kernel_arg)]
+#[help]
+pub(crate) struct ImproperGpuKernelArg<'a> {
+    pub ty: Ty<'a>,
+}
+
+#[derive(LintDiagnostic)]
+#[diag(lint_missing_gpu_kernel_export_name)]
+#[help]
+#[note]
+pub(crate) struct MissingGpuKernelExportName;
+
 #[derive(LintDiagnostic)]
 #[diag(lint_variant_size_differences)]
 pub(crate) struct VariantSizeDifferencesDiag {