Skip to content

Commit 10b32d6

Browse files
committed
Add checks for gpu-kernel calling conv
The `gpu-kernel` calling convention has several restrictions that were not enforced by the compiler until now. Add the following restrictions: 1. Cannot be async 2. Cannot be called 3. Cannot return values, return type must be `()` or `!` 4. Arguments should be simple, i.e. passed by value. More complicated types can work when you know what you are doing, but it is rather unintuitive, one needs to know ABI/compiler internals. 5. Export name should be unmangled, either through `no_mangle` or `export_name`. Kernels are searched by name on the CPU side, having a mangled name makes it hard to find and probably almost always unintentional.
1 parent 08de25c commit 10b32d6

35 files changed

+981
-152
lines changed

compiler/rustc_abi/src/extern_abi.rs

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -67,7 +67,6 @@ pub enum ExternAbi {
6767

6868
/* gpu */
6969
/// An entry-point function called by the GPU's host
70-
// FIXME: should not be callable from Rust on GPU targets, is for host's use only
7170
GpuKernel,
7271
/// An entry-point function called by the GPU's host
7372
// FIXME: why do we have two of these?

compiler/rustc_ast_passes/src/ast_validation.rs

Lines changed: 21 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -401,9 +401,16 @@ impl<'a> AstValidator<'a> {
401401
| CanonAbi::Rust
402402
| CanonAbi::RustCold
403403
| CanonAbi::Arm(_)
404-
| CanonAbi::GpuKernel
405404
| CanonAbi::X86(_) => { /* nothing to check */ }
406405

406+
CanonAbi::GpuKernel => {
407+
// An `extern "gpu-kernel"` function cannot be `async` and/or `gen`.
408+
self.reject_coroutine(abi, sig);
409+
410+
// An `extern "gpu-kernel"` function cannot return a value.
411+
self.reject_return(abi, sig);
412+
}
413+
407414
CanonAbi::Custom => {
408415
// An `extern "custom"` function must be unsafe.
409416
self.reject_safe_fn(abi, ctxt, sig);
@@ -433,18 +440,7 @@ impl<'a> AstValidator<'a> {
433440
self.dcx().emit_err(errors::AbiX86Interrupt { spans, param_count });
434441
}
435442

436-
if let FnRetTy::Ty(ref ret_ty) = sig.decl.output
437-
&& match &ret_ty.kind {
438-
TyKind::Never => false,
439-
TyKind::Tup(tup) if tup.is_empty() => false,
440-
_ => true,
441-
}
442-
{
443-
self.dcx().emit_err(errors::AbiMustNotHaveReturnType {
444-
span: ret_ty.span,
445-
abi,
446-
});
447-
}
443+
self.reject_return(abi, sig);
448444
} else {
449445
// An `extern "interrupt"` function must have type `fn()`.
450446
self.reject_params_or_return(abi, ident, sig);
@@ -496,6 +492,18 @@ impl<'a> AstValidator<'a> {
496492
}
497493
}
498494

495+
fn reject_return(&self, abi: ExternAbi, sig: &FnSig) {
496+
if let FnRetTy::Ty(ref ret_ty) = sig.decl.output
497+
&& match &ret_ty.kind {
498+
TyKind::Never => false,
499+
TyKind::Tup(tup) if tup.is_empty() => false,
500+
_ => true,
501+
}
502+
{
503+
self.dcx().emit_err(errors::AbiMustNotHaveReturnType { span: ret_ty.span, abi });
504+
}
505+
}
506+
499507
fn reject_params_or_return(&self, abi: ExternAbi, ident: &Ident, sig: &FnSig) {
500508
let mut spans: Vec<_> = sig.decl.inputs.iter().map(|p| p.span).collect();
501509
if let FnRetTy::Ty(ref ret_ty) = sig.decl.output

compiler/rustc_hir_typeck/messages.ftl

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -133,6 +133,10 @@ hir_typeck_fru_suggestion =
133133
hir_typeck_functional_record_update_on_non_struct =
134134
functional record update syntax requires a struct
135135
136+
hir_typeck_gpu_kernel_abi_cannot_be_called =
137+
functions with the "gpu-kernel" ABI cannot be called
138+
.note = an `extern "gpu-kernel"` function can only be launched on the GPU through an API
139+
136140
hir_typeck_help_set_edition_cargo = set `edition = "{$edition}"` in `Cargo.toml`
137141
hir_typeck_help_set_edition_standalone = pass `--edition {$edition}` to `rustc`
138142

compiler/rustc_hir_typeck/src/callee.rs

Lines changed: 13 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -169,27 +169,27 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
169169
}
170170
};
171171

172-
let valid = match canon_abi {
172+
match canon_abi {
173173
// Rust doesn't know how to call functions with this ABI.
174-
CanonAbi::Custom => false,
175-
176-
// These is an entry point for the host, and cannot be called on the GPU.
177-
CanonAbi::GpuKernel => false,
178-
174+
CanonAbi::Custom
179175
// The interrupt ABIs should only be called by the CPU. They have complex
180176
// pre- and postconditions, and can use non-standard instructions like `iret` on x86.
181-
CanonAbi::Interrupt(_) => false,
177+
| CanonAbi::Interrupt(_) => {
178+
let err = crate::errors::AbiCannotBeCalled { span, abi };
179+
self.tcx.dcx().emit_err(err);
180+
}
181+
182+
// This is an entry point for the host, and cannot be called directly.
183+
CanonAbi::GpuKernel => {
184+
let err = crate::errors::GpuKernelAbiCannotBeCalled { span };
185+
self.tcx.dcx().emit_err(err);
186+
}
182187

183188
CanonAbi::C
184189
| CanonAbi::Rust
185190
| CanonAbi::RustCold
186191
| CanonAbi::Arm(_)
187-
| CanonAbi::X86(_) => true,
188-
};
189-
190-
if !valid {
191-
let err = crate::errors::AbiCannotBeCalled { span, abi };
192-
self.tcx.dcx().emit_err(err);
192+
| CanonAbi::X86(_) => {}
193193
}
194194
}
195195

compiler/rustc_hir_typeck/src/errors.rs

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1198,6 +1198,14 @@ pub(crate) struct AbiCannotBeCalled {
11981198
pub abi: ExternAbi,
11991199
}
12001200

1201+
#[derive(Diagnostic)]
1202+
#[diag(hir_typeck_gpu_kernel_abi_cannot_be_called)]
1203+
pub(crate) struct GpuKernelAbiCannotBeCalled {
1204+
#[primary_span]
1205+
#[note]
1206+
pub span: Span,
1207+
}
1208+
12011209
#[derive(Diagnostic)]
12021210
#[diag(hir_typeck_const_continue_bad_label)]
12031211
pub(crate) struct ConstContinueBadLabel {

compiler/rustc_lint/messages.ftl

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -470,6 +470,9 @@ lint_improper_ctypes_union_non_exhaustive = this union is non-exhaustive
470470
471471
lint_improper_ctypes_unsafe_binder = unsafe binders are incompatible with foreign function interfaces
472472
473+
lint_improper_gpu_kernel_arg = passing type `{$ty}` to a function with "gpu-kernel" ABI may have unexpected behavior
474+
.help = use primitive types and raw pointers to get reliable behavior
475+
473476
lint_int_to_ptr_transmutes = transmuting an integer to a pointer creates a pointer without provenance
474477
.note = this is dangerous because dereferencing the resulting pointer is undefined behavior
475478
.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 =
597600
lint_mismatched_lifetime_syntaxes_suggestion_mixed_only_paths =
598601
use `'_` for type paths
599602
603+
lint_missing_gpu_kernel_export_name = function with the "gpu-kernel" ABI has a mangled name
604+
.note = mangled names make it hard to find the kernel, this is usually not intended
605+
.help = use `unsafe(no_mangle)` or `unsafe(export_name = "<name>")`
606+
600607
lint_mixed_script_confusables =
601608
the usage of Script Group `{$set}` in this crate consists solely of mixed script confusables
602609
.includes_note = the usage includes {$includes}
Lines changed: 177 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,177 @@
1+
use std::iter;
2+
3+
use rustc_abi::ExternAbi;
4+
use rustc_hir::attrs::AttributeKind;
5+
use rustc_hir::{self as hir, find_attr};
6+
use rustc_middle::ty;
7+
use rustc_session::{declare_lint, declare_lint_pass};
8+
use rustc_span::Span;
9+
use rustc_span::def_id::LocalDefId;
10+
11+
use crate::lints::{ImproperGpuKernelArg, MissingGpuKernelExportName};
12+
use crate::{LateContext, LateLintPass, LintContext};
13+
14+
declare_lint! {
15+
/// The `improper_gpu_kernel_arg` lint detects incorrect use of types in `gpu-kernel`
16+
/// arguments.
17+
///
18+
/// ### Example
19+
///
20+
/// ```rust,ignore (fails on non-GPU targets)
21+
/// #[unsafe(no_mangle)]
22+
/// extern "gpu-kernel" fn kernel(_: [i32; 10]) {}
23+
/// ```
24+
///
25+
/// This will produce:
26+
///
27+
/// ```text
28+
/// warning: passing type `[i32; 10]` to a function with "gpu-kernel" ABI may have unexpected behavior
29+
/// --> t.rs:2:34
30+
/// |
31+
/// 2 | extern "gpu-kernel" fn kernel(_: [i32; 10]) {}
32+
/// | ^^^^^^^^^
33+
/// |
34+
/// = help: use primitive types and raw pointers to get reliable behavior
35+
/// = note: `#[warn(improper_gpu_kernel_arg)]` on by default
36+
/// ```
37+
///
38+
/// ### Explanation
39+
///
40+
/// The compiler has several checks to verify that types used as arguments in `gpu-kernel`
41+
/// functions follow certain rules to ensure proper compatibility with the foreign interfaces.
42+
/// This lint is issued when it detects a probable mistake in a signature.
43+
IMPROPER_GPU_KERNEL_ARG,
44+
Warn,
45+
"simple arguments of gpu-kernel functions"
46+
}
47+
48+
declare_lint! {
49+
/// The `missing_gpu_kernel_export_name` lint detects `gpu-kernel` functions that have a mangled name.
50+
///
51+
/// ### Example
52+
///
53+
/// ```rust,ignore (fails on non-GPU targets)
54+
/// extern "gpu-kernel" fn kernel() { }
55+
/// ```
56+
///
57+
/// This will produce:
58+
///
59+
/// ```text
60+
/// warning: function with the "gpu-kernel" ABI has a mangled name
61+
/// --> t.rs:1:1
62+
/// |
63+
/// 1 | extern "gpu-kernel" fn kernel() {}
64+
/// | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
65+
/// |
66+
/// = help: use `unsafe(no_mangle)` or `unsafe(export_name = "<name>")`
67+
/// = note: mangled names make it hard to find the kernel, this is usually not intended
68+
/// = note: `#[warn(missing_gpu_kernel_export_name)]` on by default
69+
/// ```
70+
///
71+
/// ### Explanation
72+
///
73+
/// `gpu-kernel` functions are usually searched by name in the compiled file.
74+
/// A mangled name is usually unintentional as it would need to be searched by the mangled name.
75+
///
76+
/// To use an unmangled name for the kernel, either `no_mangle` or `export_name` can be used.
77+
/// ```rust,ignore (fails on non-GPU targets)
78+
/// // Can be found by the name "kernel"
79+
/// #[unsafe(no_mangle)]
80+
/// extern "gpu-kernel" fn kernel() { }
81+
///
82+
/// // Can be found by the name "new_name"
83+
/// #[unsafe(export_name = "new_name")]
84+
/// extern "gpu-kernel" fn other_kernel() { }
85+
/// ```
86+
MISSING_GPU_KERNEL_EXPORT_NAME,
87+
Warn,
88+
"mangled gpu-kernel function"
89+
}
90+
91+
declare_lint_pass!(ImproperGpuKernelLint => [
92+
IMPROPER_GPU_KERNEL_ARG,
93+
MISSING_GPU_KERNEL_EXPORT_NAME,
94+
]);
95+
96+
/// `ImproperGpuKernelLint` checks `gpu-kernel` function definitions:
97+
///
98+
/// - `extern "gpu-kernel" fn` arguments should be simple.
99+
/// - `extern "gpu-kernel" fn` should have an unmangled name.
100+
impl<'tcx> LateLintPass<'tcx> for ImproperGpuKernelLint {
101+
fn check_fn(
102+
&mut self,
103+
cx: &LateContext<'tcx>,
104+
kind: hir::intravisit::FnKind<'tcx>,
105+
decl: &'tcx hir::FnDecl<'_>,
106+
_: &'tcx hir::Body<'_>,
107+
span: Span,
108+
id: LocalDefId,
109+
) {
110+
use hir::intravisit::FnKind;
111+
112+
let abi = match kind {
113+
FnKind::ItemFn(_, _, header, ..) => header.abi,
114+
FnKind::Method(_, sig, ..) => sig.header.abi,
115+
_ => return,
116+
};
117+
118+
if abi != ExternAbi::GpuKernel {
119+
return;
120+
}
121+
122+
let sig = cx.tcx.fn_sig(id).instantiate_identity();
123+
let sig = cx.tcx.instantiate_bound_regions_with_erased(sig);
124+
125+
for (input_ty, input_hir) in iter::zip(sig.inputs(), decl.inputs) {
126+
let is_valid_arg = match input_ty.kind() {
127+
ty::Bool
128+
| ty::Char
129+
| ty::Int(_)
130+
| ty::Uint(_)
131+
| ty::Float(_)
132+
| ty::RawPtr(_, _) => true,
133+
134+
ty::Adt(_, _)
135+
| ty::Alias(_, _)
136+
| ty::Array(_, _)
137+
| ty::Bound(_, _)
138+
| ty::Closure(_, _)
139+
| ty::Coroutine(_, _)
140+
| ty::CoroutineClosure(_, _)
141+
| ty::CoroutineWitness(..)
142+
| ty::Dynamic(_, _)
143+
| ty::Error(_)
144+
| ty::FnDef(_, _)
145+
| ty::FnPtr(..)
146+
| ty::Foreign(_)
147+
| ty::Infer(_)
148+
| ty::Never
149+
| ty::Param(_)
150+
| ty::Pat(_, _)
151+
| ty::Placeholder(_)
152+
| ty::Ref(_, _, _)
153+
| ty::Slice(_)
154+
| ty::Str
155+
| ty::Tuple(_)
156+
| ty::UnsafeBinder(_) => false,
157+
};
158+
159+
if !is_valid_arg {
160+
cx.tcx.emit_node_span_lint(
161+
IMPROPER_GPU_KERNEL_ARG,
162+
input_hir.hir_id,
163+
input_hir.span,
164+
ImproperGpuKernelArg { ty: *input_ty },
165+
);
166+
}
167+
}
168+
169+
// Check for no_mangle/export_name, so the kernel can be found when querying the compiled object for the kernel function by name
170+
if !find_attr!(
171+
cx.tcx.get_all_attrs(id),
172+
AttributeKind::NoMangle(..) | AttributeKind::ExportName { .. }
173+
) {
174+
cx.emit_span_lint(MISSING_GPU_KERNEL_EXPORT_NAME, span, MissingGpuKernelExportName);
175+
}
176+
}
177+
}

compiler/rustc_lint/src/lib.rs

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,7 @@ mod expect;
4747
mod for_loops_over_fallibles;
4848
mod foreign_modules;
4949
mod function_cast_as_integer;
50+
mod gpukernel_abi;
5051
mod if_let_rescope;
5152
mod impl_trait_overcaptures;
5253
mod interior_mutable_consts;
@@ -93,6 +94,7 @@ use drop_forget_useless::*;
9394
use enum_intrinsics_non_enums::EnumIntrinsicsNonEnums;
9495
use for_loops_over_fallibles::*;
9596
use function_cast_as_integer::*;
97+
use gpukernel_abi::*;
9698
use if_let_rescope::IfLetRescope;
9799
use impl_trait_overcaptures::ImplTraitOvercaptures;
98100
use interior_mutable_consts::*;
@@ -197,6 +199,7 @@ late_lint_methods!(
197199
DerefIntoDynSupertrait: DerefIntoDynSupertrait,
198200
DropForgetUseless: DropForgetUseless,
199201
ImproperCTypesLint: ImproperCTypesLint,
202+
ImproperGpuKernelLint: ImproperGpuKernelLint,
200203
InvalidFromUtf8: InvalidFromUtf8,
201204
VariantSizeDifferences: VariantSizeDifferences,
202205
PathStatements: PathStatements,

compiler/rustc_lint/src/lints.rs

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2008,6 +2008,19 @@ impl<'a> LintDiagnostic<'a, ()> for ImproperCTypes<'_> {
20082008
}
20092009
}
20102010

2011+
#[derive(LintDiagnostic)]
2012+
#[diag(lint_improper_gpu_kernel_arg)]
2013+
#[help]
2014+
pub(crate) struct ImproperGpuKernelArg<'a> {
2015+
pub ty: Ty<'a>,
2016+
}
2017+
2018+
#[derive(LintDiagnostic)]
2019+
#[diag(lint_missing_gpu_kernel_export_name)]
2020+
#[help]
2021+
#[note]
2022+
pub(crate) struct MissingGpuKernelExportName;
2023+
20112024
#[derive(LintDiagnostic)]
20122025
#[diag(lint_variant_size_differences)]
20132026
pub(crate) struct VariantSizeDifferencesDiag {

0 commit comments

Comments
 (0)