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}