ptx_parser/type/function.rs
1use super::common::{Instruction, Label};
2use super::variable::{ParameterDirective, VariableDirective};
3use crate::Spanned;
4use crate::parser::Span;
5use crate::r#type::{AttributeDirective, DataType, FunctionSymbol, VariableSymbol};
6
7/// Alias directive relating one function symbol to another.
8///
9/// Syntax:
10/// .alias fAlias, fAliasee;
11///
12/// Example:
13/// .alias foo, bar;
14#[derive(Debug, Clone, PartialEq, Spanned)]
15pub struct AliasFunctionDirective {
16 pub alias: FunctionSymbol,
17 pub target: FunctionSymbol,
18 pub span: Span,
19}
20
21/// A PTX kernel declared with the `.entry` directive.
22#[derive(Debug, Clone, PartialEq, Spanned)]
23pub struct FuncFunctionDirective {
24 /// Example:
25 /// .func .attribute(.unified(0xAB, 0xCD)) bar() { ... }
26 pub attributes: Vec<AttributeDirective>,
27 /// Optional return param.
28 ///
29 /// Example:
30 /// .func (.param .u32 rval) bar(.param .u32 N, .param .align 4 .b8 numbers[]) { ... }
31 pub return_param: Option<ParameterDirective>,
32 /// Function name.
33 pub name: FunctionSymbol,
34 /// Function parameters.
35 ///
36 /// Example:
37 /// .func (.param .u32 rval) bar(.param .u32 N, .param .align 4 .b8 numbers[])
38 pub params: Vec<ParameterDirective>,
39 /// Optional directives.
40 ///
41 /// Example:
42 /// .func foo (.reg .b32 N, .reg .f64 dbl) .noreturn { ... }
43 pub directives: Vec<FuncFunctionHeaderDirective>,
44 /// Optional function body. Without body represents a function prototype.
45 pub body: Option<FunctionBody>,
46 pub span: Span,
47}
48
49/// A PTX device function declared with the `.func` directive.
50#[derive(Debug, Clone, PartialEq, Spanned)]
51pub struct EntryFunctionDirective {
52 /// Name of the entry function.
53 pub name: FunctionSymbol,
54 /// Function parameters.
55 pub params: Vec<ParameterDirective>,
56 /// Optional directives.
57 pub directives: Vec<EntryFunctionHeaderDirective>,
58 /// Optional function body. Without body represents a function prototype.
59 pub body: Option<FunctionBody>,
60 pub span: Span,
61}
62
63/// Directive tokens that may decorate a PTX function header.
64#[derive(Debug, Clone, PartialEq, Spanned)]
65pub enum FuncFunctionHeaderDirective {
66 /// Syntax:
67 /// .noreturn
68 ///
69 /// Example:
70 /// .func foo .noreturn { ... }
71 NoReturn { span: Span },
72 /// Syntax:
73 /// .pragma list-of-strings ;
74 ///
75 /// Example:
76 ///.entry foo .pragma "nounroll"; { ... } // disable unrolling for current kernel
77 Pragma { args: Vec<String>, span: Span },
78 /// Syntax:
79 /// .abi_preserve N
80 ///
81 /// Example:
82 /// .entry foo .abi_preserve 8 { ... }
83 AbiPreserve { value: u32, span: Span },
84 /// Syntax:
85 /// .abi_preserve_control N
86 ///
87 /// Example:
88 /// .entry foo .abi_preserve_control 16 { ... }
89 AbiPreserveControl { value: u32, span: Span },
90}
91
92/// Directive tokens that may decorate a PTX function header.
93#[derive(Debug, Clone, PartialEq, Spanned)]
94pub enum EntryFunctionHeaderDirective {
95 /// Syntax:
96 /// .maxnreg n
97 ///
98 /// Example:
99 /// .entry foo .maxnreg 16 { ... } // max regs per thread = 16
100 MaxNReg { value: u32, span: Span },
101 /// Syntax:
102 /// .maxntid nx
103 /// .maxntid nx, ny
104 /// .maxntid nx, ny, nz
105 ///
106 /// Example:
107 /// .entry foo .maxntid 256 { ... } // max threads = 256
108 /// .entry bar .maxntid 16,16,4 { ... } // max threads = 1024
109 MaxNTid { dim: FunctionDim, span: Span },
110 /// Syntax:
111 /// .reqntid nx
112 /// .reqntid nx, ny
113 /// .reqntid nx, ny, nz
114 ///
115 /// Example:
116 /// .entry foo .reqntid 256 { ... } // num threads = 256
117 /// .entry bar .reqntid 16,16,4 { ... } // num threads = 1024
118 ReqNTid { dim: FunctionDim, span: Span },
119 /// Syntax:
120 /// .minnctapersm ncta
121 ///
122 /// Example:
123 /// .entry foo .maxntid 256 .minnctapersm 4 { ... }
124 MinNCtaPerSm { value: u32, span: Span },
125 /// Syntax:
126 /// .maxnctapersm ncta
127 ///
128 /// Example:
129 /// .entry foo .maxntid 256 .maxnctapersm 4 { ... }
130 MaxNCtaPerSm { value: u32, span: Span },
131 /// Syntax:
132 /// .pragma list-of-strings ;
133 ///
134 /// Example:
135 ///.entry foo .pragma "nounroll"; { ... } // disable unrolling for current kernel
136 Pragma { args: Vec<String>, span: Span },
137 /// Syntax:
138 /// .reqnctapercluster nx
139 /// .reqnctapercluster nx, ny
140 /// .reqnctapercluster nx, ny, nz
141 ///
142 /// Example:
143 /// .entry foo .reqnctapercluster 2 { . . . }
144 /// .entry bar .reqnctapercluster 2, 2, 1 { . . . }
145 /// .entry ker .reqnctapercluster 3, 2 { . . . }
146 ReqNctaPerCluster { dim: FunctionDim, span: Span },
147 /// Syntax:
148 /// .explicitcluster
149 ///
150 /// Example:
151 /// .entry foo .explicitcluster { . . . }
152 ExplicitCluster { span: Span },
153 /// Syntax:
154 /// .maxclusterrank n
155 ///
156 /// Example:
157 /// .entry foo ..maxclusterrank 8 { . . . }
158 MaxClusterRank { value: u32, span: Span },
159 /// Syntax:
160 ///.blocksareclusters
161 ///
162 /// Example:
163 /// .entry foo .reqntid 32, 32, 1 .reqnctapercluster 32, 32, 1 .blocksareclusters { ... } // only allowed when with .reqnctapercluster and .reqntid
164 BlocksAreClusters { span: Span },
165}
166
167/// Statements contained within a PTX function body.
168#[derive(Debug, Clone, Default, PartialEq, Spanned)]
169pub struct FunctionBody {
170 pub statements: Vec<FunctionStatement>,
171 pub span: Span,
172}
173
174/// Nested statement block enclosed in braces.
175/// Executable items that appear within a function body.
176#[derive(Debug, Clone, PartialEq, Spanned)]
177pub enum FunctionStatement {
178 Label {
179 label: Label,
180 span: Span,
181 },
182 Directive {
183 directive: StatementDirective,
184 span: Span,
185 },
186 Instruction {
187 instruction: Instruction,
188 span: Span,
189 },
190 Block {
191 statements: Vec<FunctionStatement>,
192 span: Span,
193 },
194}
195
196/// Directive that declares a register variable inside a function body.
197///
198/// Syntax:
199/// .reg .ty name<range>
200#[derive(Debug, Clone, PartialEq, Spanned)]
201pub struct RegisterDirective {
202 pub ty: DataType,
203 pub registers: Vec<RegisterTarget>,
204 pub span: Span,
205}
206
207#[derive(Debug, Clone, PartialEq, Spanned)]
208pub struct RegisterTarget {
209 pub name: VariableSymbol,
210 pub range: Option<u32>,
211 pub span: Span,
212}
213
214/// Directive that applies to individual statements.
215#[derive(Debug, Clone, PartialEq, Spanned)]
216pub enum StatementDirective {
217 Loc {
218 directive: LocationDirective,
219 span: Span,
220 },
221 Pragma {
222 directive: PragmaDirective,
223 span: Span,
224 },
225 Section {
226 directive: SectionDirective,
227 span: Span,
228 },
229 Reg {
230 directive: RegisterDirective,
231 span: Span,
232 },
233 Local {
234 directive: VariableDirective,
235 span: Span,
236 },
237 Param {
238 directive: VariableDirective,
239 span: Span,
240 },
241 Shared {
242 directive: VariableDirective,
243 span: Span,
244 },
245 Dwarf {
246 directive: DwarfDirective,
247 span: Span,
248 },
249 BranchTargets {
250 directive: BranchTargetsDirective,
251 span: Span,
252 },
253 CallTargets {
254 directive: CallTargetsDirective,
255 span: Span,
256 },
257 CallPrototype {
258 directive: CallPrototypeDirective,
259 span: Span,
260 },
261}
262
263/// Raw dwarf directive emitted by the compiler (e.g. `@@dwarf`).
264///
265/// Syntax:
266/// ```text
267/// @@DWARF dwarf-string
268///
269/// dwarf-string may have one of the
270/// .byte byte-list // comma-separated hexadecimal byte values
271/// .4byte int32-list // comma-separated hexadecimal integers in range [0..2^32-1]
272/// .quad int64-list // comma-separated hexadecimal integers in range [0..2^64-1]
273/// .4byte label
274/// .quad label
275/// ```
276#[derive(Debug, Clone, PartialEq, Spanned)]
277pub struct DwarfDirective {
278 pub kind: DwarfDirectiveKind,
279 pub span: Span,
280}
281
282#[derive(Debug, Clone, PartialEq)]
283pub enum DwarfDirectiveKind {
284 ByteValues(Vec<u8>),
285 FourByteValues(Vec<u32>),
286 QuadValues(Vec<u64>),
287 FourByteLabel(Label),
288 QuadLabel(Label),
289}
290
291/// Structured representation of a `.section` directive inside a function body.
292///
293/// Syntax:
294/// ```text
295/// .section section_name { dwarf-lines }
296///
297/// dwarf-lines have the following formats:
298/// .b8 byte-list // integers in [-128..255]
299/// .b16 int16-list // integers in [-2^15..2^16-1]
300/// .b32 int32-list // integers in [-2^31..2^32-1]
301/// label: // define label inside the debug section
302/// .b64 int64-list // integers in [-2^63..2^64-1]
303/// .b32 label
304/// .b64 label
305/// .b32 label+imm // label plus constant integer byte offset (32-bit)
306/// .b64 label+imm // label plus constant integer byte offset (64-bit)
307/// .b32 label1-label2 // difference between labels in same section (32-bit)
308/// .b64 label3-label4 // difference between labels in same section (64-bit)
309/// ```
310///
311/// Example:
312/// ```text
313/// .section .debug_str {
314/// info_string0:
315/// .b8 95 // _
316/// .b8 90 // z
317/// .b8 51 // 3
318/// .b8 102 // f
319/// .b8 111 // o
320/// .b8 111 // o
321/// .b8 118 // v
322/// .b8 0
323/// info_string1:
324/// .b8 95 // _
325/// .b8 90 // z
326/// .b8 51 // 3
327/// .b8 98 // b
328/// .b8 97 // a
329/// .b8 114 // r
330/// .b8 118 // v
331/// .b8 0
332/// .b8 95 // _
333/// .b8 90 // z
334/// .b8 51 // 3
335/// .b8 99 // c
336/// .b8 97 // a
337/// .b8 114 // r
338/// .b8 118 // v
339/// .b8 0
340/// }
341/// ```
342#[derive(Debug, Clone, PartialEq, Spanned)]
343pub struct SectionDirective {
344 pub name: String,
345 pub entries: Vec<SectionEntry>,
346 pub span: Span,
347}
348
349#[derive(Debug, Clone, PartialEq)]
350pub enum SectionEntry {
351 Label { label: Label, span: Span },
352 Directive(StatementSectionDirectiveLine),
353}
354
355#[derive(Debug, Clone, PartialEq, Spanned)]
356pub enum StatementSectionDirectiveLine {
357 B8 { values: Vec<i16>, span: Span },
358 B16 { values: Vec<i32>, span: Span },
359 B32Immediate { values: Vec<i64>, span: Span },
360 B64Immediate { values: Vec<i128>, span: Span },
361 B32Label { labels: Label, span: Span },
362 B64Label { labels: Label, span: Span },
363 B32LabelPlusImm { entries: (Label, i32), span: Span },
364 B64LabelPlusImm { entries: (Label, i64), span: Span },
365 B32LabelDiff { entries: (Label, Label), span: Span },
366 B64LabelDiff { entries: (Label, Label), span: Span },
367}
368
369/// Structured representation of a `.loc` directive inside a PTX function.
370///
371/// Syntax:
372/// .loc file_index line_number column_position
373/// .loc file_index line_number column_position,function_name label {+ immediate }, inlined_at file_index2 line_number2 column_position2
374#[derive(Debug, Clone, PartialEq, Spanned)]
375pub struct LocationDirective {
376 pub file_index: u32,
377 pub line: u32,
378 pub column: u32,
379 pub inlined_at: Option<LocationInlinedAt>,
380 pub span: Span,
381}
382
383#[derive(Debug, Clone, PartialEq, Spanned)]
384pub struct LocationInlinedAt {
385 pub file_index: u32,
386 pub line: u32,
387 pub column: u32,
388 pub function_name: FunctionSymbol,
389 pub label: Label,
390 pub label_offset: Option<i64>,
391 pub span: Span,
392}
393
394/// Structured representation of a `.pragma` directive.
395///
396/// Syntax:
397/// .pragma "nounroll";
398/// .pragma "used_bytes_mask mask";
399/// .pragma "enable_smem_spilling";
400/// .pragma "frequency n";
401#[derive(Debug, Clone, PartialEq, Spanned)]
402pub struct PragmaDirective {
403 pub kind: PragmaDirectiveKind,
404 pub span: Span,
405}
406
407#[derive(Debug, Clone, PartialEq)]
408pub enum PragmaDirectiveKind {
409 Nounroll,
410 UsedBytesMask { mask: String },
411 EnableSmemSpilling,
412 Frequency { value: u32 },
413 Raw(String),
414}
415
416/// Structured representation of a `.branchtargets` directive.
417///
418/// Syntax:
419/// .branchtargets label1, label2, label3, ...;
420#[derive(Debug, Clone, PartialEq, Spanned)]
421pub struct BranchTargetsDirective {
422 pub labels: Vec<Label>,
423 pub span: Span,
424}
425
426/// Structured representation of a `.calltargets` directive.
427///
428/// Syntax:
429/// .calltargets func1, func2, func3, ...;
430#[derive(Debug, Clone, PartialEq, Spanned)]
431pub struct CallTargetsDirective {
432 pub targets: Vec<FunctionSymbol>,
433 pub span: Span,
434}
435
436/// Structured representation of a `.callprototype` directive.
437///
438/// Syntax:
439/// // no input or return parameters
440/// label: .callprototype _ .noreturn {.abi_preserve N} {.abi_preserve_control N};
441/// // input params, no return params
442/// label: .callprototype _ (param-list) .noreturn {.abi_preserve N} {.abi_preserve_control N};
443/// // no input params, // return params
444/// label: .callprototype (ret-param) _ {.abi_preserve N} {.abi_preserve_control N};
445/// // input, return parameters
446/// label: .callprototype (ret-param) _ (param-list) {.abi_preserve N} {.abi_preserve_control N};
447#[derive(Debug, Clone, PartialEq, Spanned)]
448pub struct CallPrototypeDirective {
449 pub return_param: Option<ParameterDirective>,
450 pub params: Vec<ParameterDirective>,
451 pub noreturn: bool,
452 pub abi_preserve: Option<u32>,
453 pub abi_preserve_control: Option<u32>,
454 pub span: Span,
455}
456
457/// Dimension triplet used by several function header directives.
458#[derive(Debug, Clone, PartialEq, Spanned)]
459pub enum FunctionDim {
460 X { x: u32, span: Span },
461 XY { x: u32, y: u32, span: Span },
462 XYZ { x: u32, y: u32, z: u32, span: Span },
463}