ptx_parser/unparser/
common.rs

1use crate::{
2    lexer::{PtxToken, tokenize},
3    r#type::{
4        common::{
5            AddressBase, AddressOffset, AddressOperand, AttributeDirective, Axis, CodeLinkage,
6            DataLinkage, DataType, FunctionSymbol, GeneralOperand, Immediate, Label, Operand,
7            PredicateRegister, RegisterOperand, Sign, SpecialRegister, TexHandler2, TexHandler3,
8            TexHandler3Optional, VariableSymbol, VectorOperand,
9        },
10        function::{DwarfDirective, DwarfDirectiveKind},
11        variable::ParamStateSpace,
12    },
13    unparser::{PtxUnparser, push_newline, push_space},
14};
15
16fn push_tokenized(tokens: &mut Vec<PtxToken>, text: &str) {
17    if text.trim().is_empty() {
18        return;
19    }
20    let lexemes =
21        tokenize(text).unwrap_or_else(|_| panic!("failed to tokenize literal {:?}", text));
22    tokens.extend(lexemes.into_iter().map(|(token, _)| token));
23}
24
25pub(crate) fn push_directive(tokens: &mut Vec<PtxToken>, name: &str) {
26    let raw = if name.starts_with('.') {
27        name.to_string()
28    } else {
29        format!(".{}", name)
30    };
31    push_tokenized(tokens, &raw);
32}
33
34pub(crate) fn push_token_from_str(tokens: &mut Vec<PtxToken>, value: &str) {
35    push_tokenized(tokens, value);
36}
37
38pub(crate) fn push_identifier(tokens: &mut Vec<PtxToken>, name: &str) {
39    tokens.push(PtxToken::Identifier(name.to_string()));
40}
41
42pub(crate) fn push_register(tokens: &mut Vec<PtxToken>, name: &str) {
43    tokens.push(PtxToken::Register(name.to_string()));
44}
45
46pub(crate) fn push_decimal<T: ToString>(tokens: &mut Vec<PtxToken>, value: T) {
47    tokens.push(PtxToken::DecimalInteger(value.to_string()));
48}
49
50fn push_hex_literal(tokens: &mut Vec<PtxToken>, value: u64) {
51    tokens.push(PtxToken::HexInteger(format!("0x{:x}", value)));
52}
53
54pub(crate) fn push_opcode(tokens: &mut Vec<PtxToken>, opcode: &str) {
55    push_identifier(tokens, opcode);
56}
57
58fn push_register_with_axis(tokens: &mut Vec<PtxToken>, base: &str, axis: &Axis) {
59    push_register(tokens, base);
60    match axis {
61        Axis::None { .. } => {}
62        Axis::X { .. } => push_directive(tokens, "x"),
63        Axis::Y { .. } => push_directive(tokens, "y"),
64        Axis::Z { .. } => push_directive(tokens, "z"),
65    };
66}
67
68fn numeric_token(literal: &str) -> PtxToken {
69    if literal.starts_with("0f") || literal.starts_with("0F") {
70        PtxToken::HexFloatSingle(literal.to_string())
71    } else if literal.starts_with("0d") || literal.starts_with("0D") {
72        PtxToken::HexFloatDouble(literal.to_string())
73    } else if literal.starts_with("0x") || literal.starts_with("0X") {
74        PtxToken::HexInteger(literal.to_string())
75    } else if literal.starts_with("0b") || literal.starts_with("0B") {
76        PtxToken::BinaryInteger(literal.to_string())
77    } else if literal.len() > 1
78        && literal.starts_with('0')
79        && literal.chars().all(|c| c >= '0' && c <= '7')
80    {
81        PtxToken::OctalInteger(literal.to_string())
82    } else if literal.contains('e') || literal.contains('E') {
83        PtxToken::FloatExponent(literal.to_string())
84    } else if literal.contains('.') {
85        PtxToken::Float(literal.to_string())
86    } else {
87        PtxToken::DecimalInteger(literal.to_string())
88    }
89}
90
91fn push_numeric(tokens: &mut Vec<PtxToken>, literal: &str) {
92    tokens.push(numeric_token(literal));
93}
94
95fn push_dwarf_values<I>(tokens: &mut Vec<PtxToken>, iter: I, spaced: bool)
96where
97    I: IntoIterator<Item = u64>,
98{
99    for (idx, value) in iter.into_iter().enumerate() {
100        if idx > 0 {
101            tokens.push(PtxToken::Comma);
102            push_space(tokens, spaced);
103        }
104        push_space(tokens, spaced);
105        push_hex_literal(tokens, value);
106    }
107}
108
109impl PtxUnparser for DwarfDirective {
110    fn unparse_tokens(&self, tokens: &mut Vec<PtxToken>) {
111        self.unparse_tokens_mode(tokens, false);
112    }
113
114    fn unparse_tokens_mode(&self, tokens: &mut Vec<PtxToken>, spaced: bool) {
115        push_directive(tokens, "dwarf");
116        push_space(tokens, spaced);
117        match &self.kind {
118            DwarfDirectiveKind::ByteValues(values) => {
119                push_directive(tokens, "byte");
120                push_space(tokens, spaced);
121                push_dwarf_values(tokens, values.iter().map(|v| u64::from(*v)), spaced);
122            }
123            DwarfDirectiveKind::FourByteValues(values) => {
124                push_directive(tokens, "4byte");
125                push_space(tokens, spaced);
126                push_dwarf_values(tokens, values.iter().map(|v| u64::from(*v)), spaced);
127            }
128            DwarfDirectiveKind::QuadValues(values) => {
129                push_directive(tokens, "quad");
130                push_space(tokens, spaced);
131                push_dwarf_values(tokens, values.iter().copied(), spaced);
132            }
133            DwarfDirectiveKind::FourByteLabel(label) => {
134                push_directive(tokens, "4byte");
135                push_space(tokens, spaced);
136                push_identifier(tokens, &label.val);
137            }
138            DwarfDirectiveKind::QuadLabel(label) => {
139                push_directive(tokens, "quad");
140                push_space(tokens, spaced);
141                push_identifier(tokens, &label.val);
142            }
143        }
144        tokens.push(PtxToken::Semicolon);
145        push_newline(tokens, spaced);
146    }
147}
148
149impl PtxUnparser for CodeLinkage {
150    fn unparse_tokens(&self, tokens: &mut Vec<PtxToken>) {
151        match self {
152            CodeLinkage::Visible { .. } => push_directive(tokens, "visible"),
153            CodeLinkage::Extern { .. } => push_directive(tokens, "extern"),
154            CodeLinkage::Weak { .. } => push_directive(tokens, "weak"),
155        }
156    }
157}
158
159impl PtxUnparser for DataLinkage {
160    fn unparse_tokens(&self, tokens: &mut Vec<PtxToken>) {
161        match self {
162            DataLinkage::Visible { .. } => push_directive(tokens, "visible"),
163            DataLinkage::Extern { .. } => push_directive(tokens, "extern"),
164            DataLinkage::Weak { .. } => push_directive(tokens, "weak"),
165            DataLinkage::Common { .. } => push_directive(tokens, "common"),
166        }
167    }
168}
169
170impl PtxUnparser for ParamStateSpace {
171    fn unparse_tokens(&self, tokens: &mut Vec<PtxToken>) {
172        match self {
173            ParamStateSpace::Const { .. } => push_directive(tokens, "const"),
174            ParamStateSpace::Global { .. } => push_directive(tokens, "global"),
175            ParamStateSpace::Local { .. } => push_directive(tokens, "local"),
176            ParamStateSpace::Shared { .. } => push_directive(tokens, "shared"),
177        }
178    }
179}
180
181impl PtxUnparser for AttributeDirective {
182    fn unparse_tokens(&self, tokens: &mut Vec<PtxToken>) {
183        match self {
184            AttributeDirective::Unified { uuid1, uuid2, .. } => {
185                push_directive(tokens, "unified");
186                tokens.push(PtxToken::LParen);
187                let first = uuid1.to_string();
188                push_numeric(tokens, &first);
189                tokens.push(PtxToken::Comma);
190                let second = uuid2.to_string();
191                push_numeric(tokens, &second);
192                tokens.push(PtxToken::RParen);
193            }
194            AttributeDirective::Managed { .. } => push_directive(tokens, "managed"),
195        }
196    }
197}
198
199impl PtxUnparser for DataType {
200    fn unparse_tokens(&self, tokens: &mut Vec<PtxToken>) {
201        let directive = match self {
202            DataType::U8 { .. } => "u8",
203            DataType::U16 { .. } => "u16",
204            DataType::U32 { .. } => "u32",
205            DataType::U64 { .. } => "u64",
206            DataType::S8 { .. } => "s8",
207            DataType::S16 { .. } => "s16",
208            DataType::S32 { .. } => "s32",
209            DataType::S64 { .. } => "s64",
210            DataType::F16 { .. } => "f16",
211            DataType::F16x2 { .. } => "f16x2",
212            DataType::F32 { .. } => "f32",
213            DataType::F64 { .. } => "f64",
214            DataType::B8 { .. } => "b8",
215            DataType::B16 { .. } => "b16",
216            DataType::B32 { .. } => "b32",
217            DataType::B64 { .. } => "b64",
218            DataType::B128 { .. } => "b128",
219            DataType::Pred { .. } => "pred",
220            // Texture types (merged from TexType)
221            DataType::TexRef { .. } => "texref",
222            DataType::SamplerRef { .. } => "samplerref",
223            DataType::SurfRef { .. } => "surfref",
224        };
225        push_directive(tokens, directive);
226    }
227}
228
229impl PtxUnparser for Sign {
230    fn unparse_tokens(&self, tokens: &mut Vec<PtxToken>) {
231        match self {
232            Sign::Negative { .. } => tokens.push(PtxToken::Minus),
233            Sign::Positive { .. } => tokens.push(PtxToken::Plus),
234        }
235    }
236}
237
238impl PtxUnparser for Immediate {
239    fn unparse_tokens(&self, tokens: &mut Vec<PtxToken>) {
240        let literal = self.value.as_str();
241        if let Some(rest) = literal.strip_prefix('-') {
242            tokens.push(PtxToken::Minus);
243            push_numeric(tokens, rest);
244        } else if let Some(rest) = literal.strip_prefix('+') {
245            tokens.push(PtxToken::Plus);
246            push_numeric(tokens, rest);
247        } else {
248            push_numeric(tokens, literal);
249        }
250    }
251}
252
253impl PtxUnparser for RegisterOperand {
254    fn unparse_tokens(&self, tokens: &mut Vec<PtxToken>) {
255        let mut repr = self.name.clone();
256        if let Some(component) = &self.component {
257            repr.push('.');
258            repr.push_str(component);
259        }
260        push_register(tokens, &repr);
261    }
262}
263
264impl PtxUnparser for PredicateRegister {
265    fn unparse_tokens(&self, tokens: &mut Vec<PtxToken>) {
266        push_register(tokens, &self.name);
267    }
268}
269
270impl PtxUnparser for Label {
271    fn unparse_tokens(&self, tokens: &mut Vec<PtxToken>) {
272        push_identifier(tokens, &self.val);
273    }
274}
275
276impl PtxUnparser for SpecialRegister {
277    fn unparse_tokens(&self, tokens: &mut Vec<PtxToken>) {
278        let name = match self {
279            SpecialRegister::AggrSmemSize { .. } => "%aggr_smem_size".to_string(),
280            SpecialRegister::DynamicSmemSize { .. } => "%dynamic_smem_size".to_string(),
281            SpecialRegister::LanemaskGt { .. } => "%lanemask_gt".to_string(),
282            SpecialRegister::ReservedSmemOffsetBegin { .. } => {
283                "%reserved_smem_offset_begin".to_string()
284            }
285            SpecialRegister::Clock { .. } => "%clock".to_string(),
286            SpecialRegister::Envreg { index, .. } => format!("%envreg{}", index),
287            SpecialRegister::LanemaskLe { .. } => "%lanemask_le".to_string(),
288            SpecialRegister::ReservedSmemOffsetCap { .. } => {
289                "%reserved_smem_offset_cap".to_string()
290            }
291            SpecialRegister::Clock64 { .. } => "%clock64".to_string(),
292            SpecialRegister::Globaltimer { .. } => "%globaltimer".to_string(),
293            SpecialRegister::LanemaskLt { .. } => "%lanemask_lt".to_string(),
294            SpecialRegister::ReservedSmemOffsetEnd { .. } => {
295                "%reserved_smem_offset_end".to_string()
296            }
297            SpecialRegister::ClusterCtaid { axis, .. } => {
298                push_register_with_axis(tokens, "%cluster_ctaid", axis);
299                return;
300            }
301            SpecialRegister::GlobaltimerHi { .. } => "%globaltimer_hi".to_string(),
302            SpecialRegister::Nclusterid { .. } => "%nclusterid".to_string(),
303            SpecialRegister::Smid { .. } => "%smid".to_string(),
304            SpecialRegister::ClusterCtarank { axis, .. } => {
305                push_register_with_axis(tokens, "%cluster_ctarank", axis);
306                return;
307            }
308            SpecialRegister::GlobaltimerLo { .. } => "%globaltimer_lo".to_string(),
309            SpecialRegister::Nctaid { axis, .. } => {
310                push_register_with_axis(tokens, "%nctaid", axis);
311                return;
312            }
313            SpecialRegister::Tid { axis, .. } => {
314                push_register_with_axis(tokens, "%tid", axis);
315                return;
316            }
317            SpecialRegister::ClusterNctaid { axis, .. } => {
318                push_register_with_axis(tokens, "%cluster_nctaid", axis);
319                return;
320            }
321            SpecialRegister::Gridid { .. } => "%gridid".to_string(),
322            SpecialRegister::Nsmid { .. } => "%nsmid".to_string(),
323            SpecialRegister::TotalSmemSize { .. } => "%total_smem_size".to_string(),
324            SpecialRegister::ClusterNctarank { axis, .. } => {
325                push_register_with_axis(tokens, "%cluster_nctarank", axis);
326                return;
327            }
328            SpecialRegister::IsExplicitCluster { .. } => "%is_explicit_cluster".to_string(),
329            SpecialRegister::Ntid { axis, .. } => {
330                push_register_with_axis(tokens, "%ntid", axis);
331                return;
332            }
333            SpecialRegister::Warpid { .. } => "%warpid".to_string(),
334            SpecialRegister::Clusterid { .. } => "%clusterid".to_string(),
335            SpecialRegister::Laneid { .. } => "%laneid".to_string(),
336            SpecialRegister::Nwarpid { .. } => "%nwarpid".to_string(),
337            SpecialRegister::WARPSZ { .. } => "%WARPSZ".to_string(),
338            SpecialRegister::Ctaid { axis, .. } => {
339                push_register_with_axis(tokens, "%ctaid", axis);
340                return;
341            }
342            SpecialRegister::LanemaskEq { .. } => "%lanemask_eq".to_string(),
343            SpecialRegister::Pm { index, .. } => format!("%pm{}", index),
344            SpecialRegister::Pm64 { index, .. } => format!("%pm{}_64", index),
345            SpecialRegister::CurrentGraphExec { .. } => "%current_graph_exec".to_string(),
346            SpecialRegister::LanemaskGe { .. } => "%lanemask_ge".to_string(),
347            SpecialRegister::ReservedSmemOffset { index, .. } => {
348                format!("%reserved_smem_offset_{}", index)
349            }
350        };
351        push_register(tokens, &name);
352    }
353}
354
355impl PtxUnparser for Operand {
356    fn unparse_tokens(&self, tokens: &mut Vec<PtxToken>) {
357        match self {
358            Operand::Register {
359                operand: register, ..
360            } => register.unparse_tokens(tokens),
361            Operand::Immediate {
362                operand: immediate, ..
363            } => immediate.unparse_tokens(tokens),
364            Operand::Symbol { name: symbol, .. } => push_identifier(tokens, symbol),
365            Operand::SymbolOffset { symbol, offset, .. } => {
366                push_identifier(tokens, symbol);
367                tokens.push(PtxToken::Plus);
368                offset.unparse_tokens(tokens);
369            }
370        }
371    }
372}
373
374impl PtxUnparser for VectorOperand {
375    fn unparse_tokens(&self, tokens: &mut Vec<PtxToken>) {
376        tokens.push(PtxToken::LBrace);
377        match self {
378            VectorOperand::Vector1 { operand: item, .. } => item.unparse_tokens(tokens),
379            VectorOperand::Vector2 {
380                operands: items, ..
381            } => {
382                for (idx, item) in items.iter().enumerate() {
383                    if idx > 0 {
384                        tokens.push(PtxToken::Comma);
385                    }
386                    item.unparse_tokens(tokens);
387                }
388            }
389            VectorOperand::Vector3 {
390                operands: items, ..
391            } => {
392                for (idx, item) in items.iter().enumerate() {
393                    if idx > 0 {
394                        tokens.push(PtxToken::Comma);
395                    }
396                    item.unparse_tokens(tokens);
397                }
398            }
399            VectorOperand::Vector4 {
400                operands: items, ..
401            } => {
402                for (idx, item) in items.iter().enumerate() {
403                    if idx > 0 {
404                        tokens.push(PtxToken::Comma);
405                    }
406                    item.unparse_tokens(tokens);
407                }
408            }
409            VectorOperand::Vector8 {
410                operands: items, ..
411            } => {
412                for (idx, item) in items.iter().enumerate() {
413                    if idx > 0 {
414                        tokens.push(PtxToken::Comma);
415                    }
416                    item.unparse_tokens(tokens);
417                }
418            }
419        }
420        tokens.push(PtxToken::RBrace);
421    }
422}
423
424impl PtxUnparser for GeneralOperand {
425    fn unparse_tokens(&self, tokens: &mut Vec<PtxToken>) {
426        match self {
427            GeneralOperand::Vec {
428                operand: vector, ..
429            } => vector.unparse_tokens(tokens),
430            GeneralOperand::Single { operand, .. } => operand.unparse_tokens(tokens),
431        }
432    }
433}
434
435impl PtxUnparser for TexHandler2 {
436    fn unparse_tokens(&self, tokens: &mut Vec<PtxToken>) {
437        tokens.push(PtxToken::LBracket);
438        for (idx, item) in self.operands.iter().enumerate() {
439            if idx > 0 {
440                tokens.push(PtxToken::Comma);
441            }
442            item.unparse_tokens(tokens);
443        }
444        tokens.push(PtxToken::RBracket);
445    }
446}
447
448impl PtxUnparser for TexHandler3 {
449    fn unparse_tokens(&self, tokens: &mut Vec<PtxToken>) {
450        tokens.push(PtxToken::LBracket);
451        self.handle.unparse_tokens(tokens);
452        tokens.push(PtxToken::Comma);
453        self.sampler.unparse_tokens(tokens);
454        tokens.push(PtxToken::Comma);
455        self.coords.unparse_tokens(tokens);
456        tokens.push(PtxToken::RBracket);
457    }
458}
459
460impl PtxUnparser for TexHandler3Optional {
461    fn unparse_tokens(&self, tokens: &mut Vec<PtxToken>) {
462        tokens.push(PtxToken::LBracket);
463        self.handle.unparse_tokens(tokens);
464        tokens.push(PtxToken::Comma);
465        if let Some(sampler) = &self.sampler {
466            sampler.unparse_tokens(tokens);
467            tokens.push(PtxToken::Comma);
468        }
469        self.coords.unparse_tokens(tokens);
470        tokens.push(PtxToken::RBracket);
471    }
472}
473
474impl PtxUnparser for AddressBase {
475    fn unparse_tokens(&self, tokens: &mut Vec<PtxToken>) {
476        match self {
477            AddressBase::Register {
478                operand: register, ..
479            } => register.unparse_tokens(tokens),
480            AddressBase::Variable { symbol, .. } => symbol.unparse_tokens(tokens),
481        }
482    }
483}
484
485impl PtxUnparser for AddressOffset {
486    fn unparse_tokens(&self, tokens: &mut Vec<PtxToken>) {
487        match self {
488            AddressOffset::Register {
489                operand: register, ..
490            } => {
491                tokens.push(PtxToken::Plus);
492                register.unparse_tokens(tokens);
493            }
494            AddressOffset::Immediate {
495                sign,
496                value: immediate,
497                ..
498            } => {
499                sign.unparse_tokens(tokens);
500                immediate.unparse_tokens(tokens);
501            }
502        }
503    }
504}
505
506impl PtxUnparser for AddressOperand {
507    fn unparse_tokens(&self, tokens: &mut Vec<PtxToken>) {
508        match self {
509            AddressOperand::Array { base, index, .. } => {
510                base.unparse_tokens(tokens);
511                tokens.push(PtxToken::LBracket);
512                index.unparse_tokens(tokens);
513                tokens.push(PtxToken::RBracket);
514            }
515            AddressOperand::ImmediateAddress { addr, .. } => {
516                tokens.push(PtxToken::LBracket);
517                addr.unparse_tokens(tokens);
518                tokens.push(PtxToken::RBracket);
519            }
520            AddressOperand::Offset { base, offset, .. } => {
521                tokens.push(PtxToken::LBracket);
522                base.unparse_tokens(tokens);
523                if let Some(offset) = offset {
524                    offset.unparse_tokens(tokens);
525                }
526                tokens.push(PtxToken::RBracket);
527            }
528        }
529    }
530}
531
532impl PtxUnparser for FunctionSymbol {
533    fn unparse_tokens(&self, tokens: &mut Vec<PtxToken>) {
534        push_identifier(tokens, &self.val);
535    }
536}
537
538impl PtxUnparser for VariableSymbol {
539    fn unparse_tokens(&self, tokens: &mut Vec<PtxToken>) {
540        push_identifier(tokens, &self.val);
541    }
542}
543
544impl PtxUnparser for crate::r#type::common::Instruction {
545    fn unparse_tokens(&self, tokens: &mut Vec<PtxToken>) {
546        self.unparse_tokens_mode(tokens, false);
547    }
548
549    fn unparse_tokens_mode(&self, tokens: &mut Vec<PtxToken>, spaced: bool) {
550        // Emit predicate if present
551        if let Some(predicate) = &self.predicate {
552            tokens.push(PtxToken::At);
553            if predicate.negated {
554                tokens.push(PtxToken::Exclaim);
555            }
556            predicate.operand.unparse_tokens_mode(tokens, spaced);
557            push_space(tokens, spaced);
558        }
559
560        // Emit the instruction
561        self.inst.unparse_tokens_mode(tokens, spaced);
562    }
563}
ptx_parser/unparser/common.rs

ptx_parser/unparser/
common.rs
