ptx_parser/type/
function.rs

1use super::common::{CodeLinkage, Instruction};
2use super::variable::VariableDirective;
3use crate::parser::Span;
4
5/// All directives that describe kernel/function entities.
6#[derive(Debug, Clone, PartialEq)]
7pub enum FunctionKernelDirective {
8    Entry { function: EntryFunction, span: Span },
9    Func { function: FuncFunction, span: Span },
10    Alias { alias: FunctionAlias, span: Span },
11}
12
13impl FunctionKernelDirective {
14    pub fn span(&self) -> Span {
15        match self {
16            FunctionKernelDirective::Entry { span, .. } => span.clone(),
17            FunctionKernelDirective::Func { span, .. } => span.clone(),
18            FunctionKernelDirective::Alias { span, .. } => span.clone(),
19        }
20    }
21}
22
23/// Alias directive relating one function symbol to another.
24#[derive(Debug, Clone, PartialEq)]
25pub struct FunctionAlias {
26    pub alias: String,
27    pub target: String,
28    pub span: Span,
29}
30
31impl FunctionAlias {
32    pub fn with_span(mut self, span: Span) -> Self {
33        self.span = span;
34        self
35    }
36}
37
38/// A PTX kernel declared with the `.entry` directive.
39#[derive(Debug, Clone, PartialEq)]
40pub struct EntryFunction {
41    pub name: String,
42    pub directives: Vec<FunctionHeaderDirective>,
43    pub params: Vec<VariableDirective>,
44    pub body: FunctionBody,
45    pub span: Span,
46}
47
48impl EntryFunction {
49    pub fn with_span(mut self, span: Span) -> Self {
50        self.span = span;
51        self
52    }
53}
54
55/// A PTX device function declared with the `.func` directive.
56#[derive(Debug, Clone, PartialEq)]
57pub struct FuncFunction {
58    pub name: String,
59    pub directives: Vec<FunctionHeaderDirective>,
60    pub return_param: Option<VariableDirective>,
61    pub params: Vec<VariableDirective>,
62    pub body: FunctionBody,
63    pub span: Span,
64}
65
66impl FuncFunction {
67    pub fn with_span(mut self, span: Span) -> Self {
68        self.span = span;
69        self
70    }
71}
72
73/// Statements contained within a PTX function body.
74#[derive(Debug, Clone, PartialEq)]
75pub struct FunctionBody {
76    pub statements: Vec<FunctionStatement>,
77    pub span: Span,
78}
79
80impl FunctionBody {
81    pub fn with_span(mut self, span: Span) -> Self {
82        self.span = span;
83        self
84    }
85}
86
87impl Default for FunctionBody {
88    fn default() -> Self {
89        Self {
90            statements: Vec::new(),
91            span: 0..0,
92        }
93    }
94}
95
96/// Directive tokens that may decorate a PTX function header.
97#[derive(Debug, Clone, PartialEq)]
98pub enum FunctionHeaderDirective {
99    Linkage { linkage: CodeLinkage, span: Span },
100    NoReturn { span: Span },
101    AbiPreserve { value: u32, span: Span },
102    AbiPreserveControl { value: u32, span: Span },
103    MaxClusterRank { value: u32, span: Span },
104    BlocksAreClusters { span: Span },
105    ExplicitCluster { dim: FunctionDim3, span: Span },
106    ReqNctaPerCluster { dim: FunctionDim3, span: Span },
107    MaxNReg { value: u32, span: Span },
108    MaxNTid { dim: FunctionDim3, span: Span },
109    MinNCtaPerSm { value: u32, span: Span },
110    ReqNTid { dim: FunctionDim3, span: Span },
111    MaxNCtaPerSm { value: u32, span: Span },
112    Pragma { args: Vec<String>, span: Span },
113}
114
115impl FunctionHeaderDirective {
116    pub fn span(&self) -> Span {
117        match self {
118            FunctionHeaderDirective::Linkage { span, .. } => span.clone(),
119            FunctionHeaderDirective::NoReturn { span } => span.clone(),
120            FunctionHeaderDirective::AbiPreserve { span, .. } => span.clone(),
121            FunctionHeaderDirective::AbiPreserveControl { span, .. } => span.clone(),
122            FunctionHeaderDirective::MaxClusterRank { span, .. } => span.clone(),
123            FunctionHeaderDirective::BlocksAreClusters { span } => span.clone(),
124            FunctionHeaderDirective::ExplicitCluster { span, .. } => span.clone(),
125            FunctionHeaderDirective::ReqNctaPerCluster { span, .. } => span.clone(),
126            FunctionHeaderDirective::MaxNReg { span, .. } => span.clone(),
127            FunctionHeaderDirective::MaxNTid { span, .. } => span.clone(),
128            FunctionHeaderDirective::MinNCtaPerSm { span, .. } => span.clone(),
129            FunctionHeaderDirective::ReqNTid { span, .. } => span.clone(),
130            FunctionHeaderDirective::MaxNCtaPerSm { span, .. } => span.clone(),
131            FunctionHeaderDirective::Pragma { span, .. } => span.clone(),
132        }
133    }
134}
135
136/// Dimension triplet used by several function header directives.
137#[derive(Debug, Clone, PartialEq)]
138pub struct FunctionDim3 {
139    pub x: u32,
140    pub y: Option<u32>,
141    pub z: Option<u32>,
142    pub span: Span,
143}
144
145impl FunctionDim3 {
146    pub fn with_span(mut self, span: Span) -> Self {
147        self.span = span;
148        self
149    }
150}
151
152/// Nested statement block enclosed in braces.
153/// Executable items that appear within a function body.
154#[derive(Debug, Clone, PartialEq)]
155pub enum FunctionStatement {
156    Label { name: String, span: Span },
157    Directive { directive: StatementDirective, span: Span },
158    Instruction { instruction: Instruction, span: Span },
159    Block { statements: Vec<FunctionStatement>, span: Span },
160}
161
162impl FunctionStatement {
163    pub fn span(&self) -> Span {
164        match self {
165            FunctionStatement::Label { span, .. } => span.clone(),
166            FunctionStatement::Directive { span, .. } => span.clone(),
167            FunctionStatement::Instruction { span, .. } => span.clone(),
168            FunctionStatement::Block { span, .. } => span.clone(),
169        }
170    }
171}
172
173/// .reg .ty name<range>
174#[derive(Debug, Clone, PartialEq)]
175pub struct RegisterDirective {
176    pub name: String,
177    pub ty: Option<String>,
178    pub range: Option<u32>,
179    pub comment: Option<String>,
180    pub span: Span,
181}
182
183impl RegisterDirective {
184    pub fn with_span(mut self, span: Span) -> Self {
185        self.span = span;
186        self
187    }
188}
189
190/// Directive that applies to individual statements.
191#[derive(Debug, Clone, PartialEq)]
192pub enum StatementDirective {
193    Loc { directive: LocationDirective, span: Span },
194    Pragma { directive: PragmaDirective, span: Span },
195    Section { directive: StatementSectionDirective, span: Span },
196    Reg { directive: RegisterDirective, span: Span },
197    Local { directive: VariableDirective, span: Span },
198    Param { directive: VariableDirective, span: Span },
199    Shared { directive: VariableDirective, span: Span },
200    Dwarf { directive: DwarfDirective, span: Span },
201}
202
203impl StatementDirective {
204    pub fn span(&self) -> Span {
205        match self {
206            StatementDirective::Loc { span, .. } => span.clone(),
207            StatementDirective::Pragma { span, .. } => span.clone(),
208            StatementDirective::Section { span, .. } => span.clone(),
209            StatementDirective::Reg { span, .. } => span.clone(),
210            StatementDirective::Local { span, .. } => span.clone(),
211            StatementDirective::Param { span, .. } => span.clone(),
212            StatementDirective::Shared { span, .. } => span.clone(),
213            StatementDirective::Dwarf { span, .. } => span.clone(),
214        }
215    }
216}
217
218/// Raw dwarf directive emitted by the compiler (e.g. @@dwarf).
219#[derive(Debug, Clone, PartialEq)]
220pub struct DwarfDirective {
221    pub keyword: String,
222    pub arguments: Vec<String>,
223    pub comment: Option<String>,
224    pub span: Span,
225}
226
227impl DwarfDirective {
228    pub fn with_span(mut self, span: Span) -> Self {
229        self.span = span;
230        self
231    }
232}
233
234/// Structured representation of a `.loc` directive inside a PTX function.
235#[derive(Debug, Clone, PartialEq)]
236pub struct LocationDirective {
237    pub file_index: u32,
238    pub line: u32,
239    pub column: u32,
240    pub options: Vec<String>,
241    pub comment: Option<String>,
242    pub span: Span,
243}
244
245impl LocationDirective {
246    pub fn with_span(mut self, span: Span) -> Self {
247        self.span = span;
248        self
249    }
250}
251
252/// Structured representation of a `.pragma` directive.
253#[derive(Debug, Clone, PartialEq)]
254pub struct PragmaDirective {
255    pub arguments: Vec<String>,
256    pub comment: Option<String>,
257    pub span: Span,
258}
259
260impl PragmaDirective {
261    pub fn with_span(mut self, span: Span) -> Self {
262        self.span = span;
263        self
264    }
265}
266
267/// Structured representation of a `.section` directive inside a function body.
268#[derive(Debug, Clone, PartialEq)]
269pub struct StatementSectionDirective {
270    pub name: String,
271    pub arguments: Vec<String>,
272    pub comment: Option<String>,
273    pub span: Span,
274}
275
276impl StatementSectionDirective {
277    pub fn with_span(mut self, span: Span) -> Self {
278        self.span = span;
279        self
280    }
281}
ptx_parser/type/function.rs

ptx_parser/type/
function.rs
