ptx_parser/type/
function.rs

1use super::common::CodeLinkage;
2use super::instruction::Instruction;
3use super::variable::VariableDirective;
4
5/// All directives that describe kernel/function entities.
6#[derive(Debug, Clone, PartialEq)]
7pub enum FunctionKernelDirective {
8    Entry(EntryFunction),
9    Func(FuncFunction),
10    Alias(FunctionAlias),
11}
12
13/// Alias directive relating one function symbol to another.
14#[derive(Debug, Clone, PartialEq)]
15pub struct FunctionAlias {
16    pub alias: String,
17    pub target: String,
18    pub raw: String,
19}
20
21/// A PTX kernel declared with the `.entry` directive.
22#[derive(Debug, Clone, PartialEq)]
23pub struct EntryFunction {
24    pub name: String,
25    pub directives: Vec<FunctionHeaderDirective>,
26    pub params: Vec<VariableDirective>,
27    pub body: FunctionBody,
28}
29
30/// A PTX device function declared with the `.func` directive.
31#[derive(Debug, Clone, PartialEq)]
32pub struct FuncFunction {
33    pub name: String,
34    pub directives: Vec<FunctionHeaderDirective>,
35    pub return_param: Option<VariableDirective>,
36    pub params: Vec<VariableDirective>,
37    pub body: FunctionBody,
38}
39
40/// Statements contained within a PTX function body.
41#[derive(Debug, Clone, PartialEq, Default)]
42pub struct FunctionBody {
43    pub entry_directives: Vec<FunctionEntryDirective>,
44    pub statements: Vec<FunctionStatement>,
45}
46
47/// Directive tokens that may decorate a PTX function header.
48#[derive(Debug, Clone, PartialEq)]
49pub enum FunctionHeaderDirective {
50    Linkage(CodeLinkage),
51    NoReturn,
52    AbiPreserve(u32),
53    AbiPreserveControl(u32),
54    MaxClusterRank(u32),
55    BlocksAreClusters,
56    ExplicitCluster(FunctionDim3),
57    ReqNctaPerCluster(FunctionDim3),
58    MaxNReg(u32),
59    MaxNTid(FunctionDim3),
60    MinNCtaPerSm(u32),
61    ReqNTid(FunctionDim3),
62    MaxNCtaPerSm(u32),
63    Pragma(Vec<String>),
64}
65
66/// Dimension triplet used by several function header directives.
67#[derive(Debug, Clone, PartialEq)]
68pub struct FunctionDim3 {
69    pub x: u32,
70    pub y: Option<u32>,
71    pub z: Option<u32>,
72}
73
74/// Entry directives that appear before executable statements in a function body.
75#[derive(Debug, Clone, PartialEq)]
76pub enum FunctionEntryDirective {
77    Reg(RegisterDirective),
78    Local(VariableDirective),
79    Param(VariableDirective),
80    Shared(VariableDirective),
81    Pragma(PragmaDirective),
82    Loc(LocationDirective),
83    Dwarf(DwarfDirective),
84}
85
86/// Nested statement block enclosed in braces.
87/// Executable items that appear within a function body.
88#[derive(Debug, Clone, PartialEq)]
89pub enum FunctionStatement {
90    Label(String),
91    Directive(StatementDirective),
92    Instruction(Instruction),
93    ExternCallBlock(ExternCallBlock),
94}
95
96#[derive(Debug, Clone, PartialEq)]
97pub struct ExternCallBlock {
98    pub declarations: Vec<FunctionEntryDirective>,
99    pub setup: Vec<ExternCallSetup>,
100    pub call: Instruction,
101    pub post_call: Vec<Instruction>,
102}
103
104#[derive(Debug, Clone, PartialEq)]
105pub enum ExternCallSetup {
106    Param(VariableDirective),
107    Store(Instruction),
108}
109
110/// .reg .ty name<range>
111#[derive(Debug, Clone, PartialEq)]
112pub struct RegisterDirective {
113    pub name: String,
114    pub ty: Option<String>,
115    pub range: Option<u32>,
116    pub comment: Option<String>,
117    pub raw: String,
118}
119
120/// Directive that applies to individual statements.
121#[derive(Debug, Clone, PartialEq)]
122pub enum StatementDirective {
123    Dwarf(DwarfDirective),
124    Loc(LocationDirective),
125    Pragma(PragmaDirective),
126    Section(StatementSectionDirective),
127}
128
129/// Raw dwarf directive emitted by the compiler (e.g. @@dwarf).
130#[derive(Debug, Clone, PartialEq)]
131pub struct DwarfDirective {
132    pub keyword: String,
133    pub arguments: Vec<String>,
134    pub comment: Option<String>,
135    pub raw: String,
136}
137
138/// Structured representation of a `.loc` directive inside a PTX function.
139#[derive(Debug, Clone, PartialEq)]
140pub struct LocationDirective {
141    pub file_index: u32,
142    pub line: u32,
143    pub column: u32,
144    pub options: Vec<String>,
145    pub comment: Option<String>,
146    pub raw: String,
147}
148
149/// Structured representation of a `.pragma` directive.
150#[derive(Debug, Clone, PartialEq)]
151pub struct PragmaDirective {
152    pub arguments: Vec<String>,
153    pub comment: Option<String>,
154    pub raw: String,
155}
156
157/// Structured representation of a `.section` directive inside a function body.
158#[derive(Debug, Clone, PartialEq)]
159pub struct StatementSectionDirective {
160    pub name: String,
161    pub arguments: Vec<String>,
162    pub comment: Option<String>,
163    pub raw: String,
164}