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

ptx_parser/type/
function.rs
