ptx_parser/unparser/
function.rs

1use crate::unparser::common::push_register;
2use crate::{
3    lexer::PtxToken,
4    r#type::{function::*, variable::VariableDirective},
5    unparser::*,
6};
7
8fn push_register_components(tokens: &mut Vec<PtxToken>, name: &str) {
9    if let Some(stripped) = name.strip_prefix('%') {
10        let mut parts = stripped.split('.');
11        if let Some(first) = parts.next() {
12            let register_name = format!("%{first}");
13            push_register(tokens, &register_name);
14        }
15        for part in parts {
16            if part.is_empty() {
17                continue;
18            }
19            push_directive(tokens, part);
20        }
21    } else {
22        push_identifier(tokens, name);
23    }
24}
25
26fn unparse_register_directive(tokens: &mut Vec<PtxToken>, directive: &RegisterDirective) {
27    push_directive(tokens, "reg");
28    if let Some(ty) = &directive.ty {
29        push_directive(tokens, ty);
30    }
31    push_register_components(tokens, &directive.name);
32    if let Some(range) = directive.range {
33        tokens.push(PtxToken::LAngle);
34        push_decimal(tokens, range);
35        tokens.push(PtxToken::RAngle);
36    }
37    tokens.push(PtxToken::Semicolon);
38}
39
40fn unparse_entry_directive(tokens: &mut Vec<PtxToken>, directive: &FunctionEntryDirective) {
41    match directive {
42        FunctionEntryDirective::Reg(register) => unparse_register_directive(tokens, register),
43        FunctionEntryDirective::Local(variable) => variable.unparse_tokens(tokens),
44        FunctionEntryDirective::Param(variable) => variable.unparse_tokens(tokens),
45        FunctionEntryDirective::Shared(variable) => variable.unparse_tokens(tokens),
46        FunctionEntryDirective::Pragma(_) => {
47            panic!("unimplemented: unparsing .pragma function entry directives");
48        }
49        FunctionEntryDirective::Loc(_) => {
50            panic!("unimplemented: unparsing .loc function entry directives");
51        }
52        FunctionEntryDirective::Dwarf(_) => {
53            panic!("unimplemented: unparsing dwarf function entry directives");
54        }
55    }
56}
57
58fn unparse_extern_call_setup(tokens: &mut Vec<PtxToken>, setup: &ExternCallSetup) {
59    match setup {
60        ExternCallSetup::Param(variable) => variable.unparse_tokens(tokens),
61        ExternCallSetup::Store(instruction) => instruction.unparse_tokens(tokens),
62    }
63}
64
65fn unparse_extern_call_block(tokens: &mut Vec<PtxToken>, block: &ExternCallBlock) {
66    tokens.push(PtxToken::LBrace);
67    for directive in &block.declarations {
68        unparse_entry_directive(tokens, directive);
69    }
70    for entry in &block.setup {
71        unparse_extern_call_setup(tokens, entry);
72    }
73    block.call.unparse_tokens(tokens);
74    for instruction in &block.post_call {
75        instruction.unparse_tokens(tokens);
76    }
77    tokens.push(PtxToken::RBrace);
78}
79
80fn unparse_function_statement(tokens: &mut Vec<PtxToken>, statement: &FunctionStatement) {
81    match statement {
82        FunctionStatement::Instruction(instruction) => instruction.unparse_tokens(tokens),
83        FunctionStatement::ExternCallBlock(block) => unparse_extern_call_block(tokens, block),
84        FunctionStatement::Directive(_) => {
85            panic!("unimplemented: unparsing function statement directives");
86        }
87    }
88}
89
90fn unparse_function_dim(tokens: &mut Vec<PtxToken>, dim: &FunctionDim3) {
91    push_decimal(tokens, dim.x);
92    if let Some(y) = dim.y {
93        tokens.push(PtxToken::Comma);
94        push_decimal(tokens, y);
95    }
96    if let Some(z) = dim.z {
97        tokens.push(PtxToken::Comma);
98        push_decimal(tokens, z);
99    }
100}
101
102fn unparse_param(tokens: &mut Vec<PtxToken>, param: &VariableDirective) {
103    let mut param_tokens = param.to_tokens();
104    if matches!(param_tokens.last(), Some(PtxToken::Semicolon)) {
105        param_tokens.pop();
106    }
107    tokens.extend(param_tokens);
108}
109
110fn unparse_param_list(tokens: &mut Vec<PtxToken>, params: &[VariableDirective]) {
111    for (idx, param) in params.iter().enumerate() {
112        if idx > 0 {
113            tokens.push(PtxToken::Comma);
114        }
115        unparse_param(tokens, param);
116    }
117}
118
119fn unparse_function_header_directive(
120    tokens: &mut Vec<PtxToken>,
121    directive: &FunctionHeaderDirective,
122) {
123    match directive {
124        FunctionHeaderDirective::Linkage(linkage) => linkage.unparse_tokens(tokens),
125        FunctionHeaderDirective::NoReturn => push_directive(tokens, "noreturn"),
126        FunctionHeaderDirective::AbiPreserve(value) => {
127            push_directive(tokens, "abipreserve");
128            push_decimal(tokens, *value);
129        }
130        FunctionHeaderDirective::AbiPreserveControl(value) => {
131            push_directive(tokens, "abipreserve_control");
132            push_decimal(tokens, *value);
133        }
134        FunctionHeaderDirective::MaxClusterRank(value) => {
135            push_directive(tokens, "maxclusterrank");
136            push_decimal(tokens, *value);
137        }
138        FunctionHeaderDirective::BlocksAreClusters => push_directive(tokens, "blocksareclusters"),
139        FunctionHeaderDirective::ExplicitCluster(dim) => {
140            push_directive(tokens, "explicitcluster");
141            unparse_function_dim(tokens, dim);
142        }
143        FunctionHeaderDirective::ReqNctaPerCluster(dim) => {
144            push_directive(tokens, "reqnctapercluster");
145            unparse_function_dim(tokens, dim);
146        }
147        FunctionHeaderDirective::MaxNReg(value) => {
148            push_directive(tokens, "maxnreg");
149            push_decimal(tokens, *value);
150        }
151        FunctionHeaderDirective::MaxNTid(dim) => {
152            push_directive(tokens, "maxntid");
153            unparse_function_dim(tokens, dim);
154        }
155        FunctionHeaderDirective::MinNCtaPerSm(value) => {
156            push_directive(tokens, "minnctapersm");
157            push_decimal(tokens, *value);
158        }
159        FunctionHeaderDirective::ReqNTid(dim) => {
160            push_directive(tokens, "reqntid");
161            unparse_function_dim(tokens, dim);
162        }
163        FunctionHeaderDirective::MaxNCtaPerSm(value) => {
164            push_directive(tokens, "maxnctapersm");
165            push_decimal(tokens, *value);
166        }
167        FunctionHeaderDirective::Pragma(arguments) => {
168            push_directive(tokens, "pragma");
169            for argument in arguments {
170                tokens.push(PtxToken::Identifier(argument.clone()));
171            }
172        }
173    }
174}
175
176fn unparse_function_headers(tokens: &mut Vec<PtxToken>, directives: &[FunctionHeaderDirective]) {
177    for directive in directives {
178        unparse_function_header_directive(tokens, directive);
179    }
180}
181
182fn unparse_function_body(tokens: &mut Vec<PtxToken>, body: &FunctionBody, prefer_braces: bool) {
183    if body.entry_directives.is_empty() && body.statements.is_empty() {
184        if prefer_braces {
185            tokens.push(PtxToken::LBrace);
186            tokens.push(PtxToken::RBrace);
187        } else {
188            tokens.push(PtxToken::Semicolon);
189        }
190        return;
191    }
192
193    tokens.push(PtxToken::LBrace);
194    for directive in &body.entry_directives {
195        unparse_entry_directive(tokens, directive);
196    }
197    for statement in &body.statements {
198        unparse_function_statement(tokens, statement);
199    }
200    tokens.push(PtxToken::RBrace);
201}
202
203impl PtxUnparser for FunctionAlias {
204    fn unparse_tokens(&self, tokens: &mut Vec<PtxToken>) {
205        push_directive(tokens, "alias");
206        push_identifier(tokens, &self.alias);
207        tokens.push(PtxToken::Comma);
208        push_identifier(tokens, &self.target);
209        tokens.push(PtxToken::Semicolon);
210    }
211}
212
213impl PtxUnparser for EntryFunction {
214    fn unparse_tokens(&self, tokens: &mut Vec<PtxToken>) {
215        unparse_function_headers(tokens, &self.directives);
216        push_directive(tokens, "entry");
217        push_identifier(tokens, &self.name);
218        tokens.push(PtxToken::LParen);
219        unparse_param_list(tokens, &self.params);
220        tokens.push(PtxToken::RParen);
221        unparse_function_body(tokens, &self.body, true);
222    }
223}
224
225impl PtxUnparser for FuncFunction {
226    fn unparse_tokens(&self, tokens: &mut Vec<PtxToken>) {
227        unparse_function_headers(tokens, &self.directives);
228        push_directive(tokens, "func");
229        if let Some(ret) = &self.return_param {
230            tokens.push(PtxToken::LParen);
231            unparse_param(tokens, ret);
232            tokens.push(PtxToken::RParen);
233        }
234        push_identifier(tokens, &self.name);
235        tokens.push(PtxToken::LParen);
236        unparse_param_list(tokens, &self.params);
237        tokens.push(PtxToken::RParen);
238        unparse_function_body(tokens, &self.body, false);
239    }
240}
241
242impl PtxUnparser for FunctionKernelDirective {
243    fn unparse_tokens(&self, tokens: &mut Vec<PtxToken>) {
244        match self {
245            FunctionKernelDirective::Entry(entry) => entry.unparse_tokens(tokens),
246            FunctionKernelDirective::Func(func) => func.unparse_tokens(tokens),
247            FunctionKernelDirective::Alias(alias) => alias.unparse_tokens(tokens),
248        }
249    }
250}
ptx_parser/unparser/function.rs

ptx_parser/unparser/
function.rs
