ptx_parser/type/
function.rs

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