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

ptx_parser/type/
function.rs
