rustc_target/callconv/
nvptx64.rs

1use rustc_abi::{HasDataLayout, Reg, Size, TyAbiInterface};
2
3use super::CastTarget;
4use crate::callconv::{ArgAbi, FnAbi, Uniform};
5
6fn classify_ret<Ty>(ret: &mut ArgAbi<'_, Ty>) {
7    if ret.layout.is_aggregate() && ret.layout.is_sized() {
8        classify_aggregate(ret)
9    } else if ret.layout.size.bits() < 32 && ret.layout.is_sized() {
10        ret.extend_integer_width_to(32);
11    }
12}
13
14fn classify_arg<'a, Ty, C>(cx: &C, arg: &mut ArgAbi<'a, Ty>)
15where
16    Ty: TyAbiInterface<'a, C> + Copy,
17{
18    if arg.layout.pass_indirectly_in_non_rustic_abis(cx) {
19        arg.make_indirect();
20        return;
21    }
22    if arg.layout.is_aggregate() && arg.layout.is_sized() {
23        classify_aggregate(arg)
24    } else if arg.layout.size.bits() < 32 && arg.layout.is_sized() {
25        arg.extend_integer_width_to(32);
26    }
27}
28
29/// the pass mode used for aggregates in arg and ret position
30fn classify_aggregate<Ty>(arg: &mut ArgAbi<'_, Ty>) {
31    let align_bytes = arg.layout.align.bytes();
32    let size = arg.layout.size;
33
34    let reg = match align_bytes {
35        1 => Reg::i8(),
36        2 => Reg::i16(),
37        4 => Reg::i32(),
38        8 => Reg::i64(),
39        16 => Reg::i128(),
40        _ => unreachable!("Align is given as power of 2 no larger than 16 bytes"),
41    };
42
43    if align_bytes == size.bytes() {
44        arg.cast_to(CastTarget::prefixed(
45            [Some(reg), None, None, None, None, None, None, None],
46            Uniform::new(Reg::i8(), Size::ZERO),
47        ));
48    } else {
49        arg.cast_to(Uniform::new(reg, size));
50    }
51}
52
53fn classify_arg_kernel<'a, Ty, C>(_cx: &C, arg: &mut ArgAbi<'a, Ty>)
54where
55    Ty: TyAbiInterface<'a, C> + Copy,
56    C: HasDataLayout,
57{
58    match arg.mode {
59        super::PassMode::Ignore | super::PassMode::Direct(_) => return,
60        super::PassMode::Pair(_, _) => {}
61        super::PassMode::Cast { .. } => unreachable!(),
62        super::PassMode::Indirect { .. } => {}
63    }
64
65    // FIXME only allow structs and wide pointers here
66    // panic!(
67    //     "`extern \"ptx-kernel\"` doesn't allow passing types other than primitives and structs"
68    // );
69
70    let align_bytes = arg.layout.align.bytes();
71
72    let unit = match align_bytes {
73        1 => Reg::i8(),
74        2 => Reg::i16(),
75        4 => Reg::i32(),
76        8 => Reg::i64(),
77        16 => Reg::i128(),
78        _ => unreachable!("Align is given as power of 2 no larger than 16 bytes"),
79    };
80    if arg.layout.size.bytes() / align_bytes == 1 {
81        // Make sure we pass the struct as array at the LLVM IR level and not as a single integer.
82        arg.cast_to(CastTarget::prefixed(
83            [Some(unit), None, None, None, None, None, None, None],
84            Uniform::new(unit, Size::ZERO),
85        ));
86    } else {
87        arg.cast_to(Uniform::new(unit, arg.layout.size));
88    }
89}
90
91pub(crate) fn compute_abi_info<'a, Ty, C>(cx: &C, fn_abi: &mut FnAbi<'a, Ty>)
92where
93    Ty: TyAbiInterface<'a, C> + Copy,
94{
95    if !fn_abi.ret.is_ignore() {
96        classify_ret(&mut fn_abi.ret);
97    }
98
99    for arg in fn_abi.args.iter_mut() {
100        if arg.is_ignore() {
101            continue;
102        }
103        classify_arg(cx, arg);
104    }
105}
106
107pub(crate) fn compute_ptx_kernel_abi_info<'a, Ty, C>(cx: &C, fn_abi: &mut FnAbi<'a, Ty>)
108where
109    Ty: TyAbiInterface<'a, C> + Copy,
110    C: HasDataLayout,
111{
112    if !fn_abi.ret.layout.is_unit() && !fn_abi.ret.layout.is_never() {
113        panic!("Kernels should not return anything other than () or !");
114    }
115
116    for arg in fn_abi.args.iter_mut() {
117        if arg.is_ignore() {
118            continue;
119        }
120        classify_arg_kernel(cx, arg);
121    }
122}