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        push_directive(tokens, "attribute");
184        tokens.push(PtxToken::LParen);
185        match self {
186            AttributeDirective::Unified { uuid1, uuid2, .. } => {
187                push_directive(tokens, "unified");
188                tokens.push(PtxToken::LParen);
189                let first = uuid1.to_string();
190                push_numeric(tokens, &first);
191                tokens.push(PtxToken::Comma);
192                let second = uuid2.to_string();
193                push_numeric(tokens, &second);
194                tokens.push(PtxToken::RParen);
195            }
196            AttributeDirective::Managed { .. } => push_directive(tokens, "managed"),
197        }
198        tokens.push(PtxToken::RParen);
199    }
200}
201
202impl PtxUnparser for DataType {
203    fn unparse_tokens(&self, tokens: &mut Vec<PtxToken>) {
204        let directive = match self {
205            DataType::U8 { .. } => "u8",
206            DataType::U16 { .. } => "u16",
207            DataType::U32 { .. } => "u32",
208            DataType::U64 { .. } => "u64",
209            DataType::S8 { .. } => "s8",
210            DataType::S16 { .. } => "s16",
211            DataType::S32 { .. } => "s32",
212            DataType::S64 { .. } => "s64",
213            DataType::F16 { .. } => "f16",
214            DataType::F16x2 { .. } => "f16x2",
215            DataType::F32 { .. } => "f32",
216            DataType::F64 { .. } => "f64",
217            DataType::B8 { .. } => "b8",
218            DataType::B16 { .. } => "b16",
219            DataType::B32 { .. } => "b32",
220            DataType::B64 { .. } => "b64",
221            DataType::B128 { .. } => "b128",
222            DataType::Pred { .. } => "pred",
223            // Texture types (merged from TexType)
224            DataType::TexRef { .. } => "texref",
225            DataType::SamplerRef { .. } => "samplerref",
226            DataType::SurfRef { .. } => "surfref",
227        };
228        push_directive(tokens, directive);
229    }
230}
231
232impl PtxUnparser for Sign {
233    fn unparse_tokens(&self, tokens: &mut Vec<PtxToken>) {
234        match self {
235            Sign::Negative { .. } => tokens.push(PtxToken::Minus),
236            Sign::Positive { .. } => tokens.push(PtxToken::Plus),
237        }
238    }
239}
240
241impl PtxUnparser for Immediate {
242    fn unparse_tokens(&self, tokens: &mut Vec<PtxToken>) {
243        let literal = self.value.as_str();
244        if let Some(rest) = literal.strip_prefix('-') {
245            tokens.push(PtxToken::Minus);
246            push_numeric(tokens, rest);
247        } else if let Some(rest) = literal.strip_prefix('+') {
248            tokens.push(PtxToken::Plus);
249            push_numeric(tokens, rest);
250        } else {
251            push_numeric(tokens, literal);
252        }
253    }
254}
255
256impl PtxUnparser for RegisterOperand {
257    fn unparse_tokens(&self, tokens: &mut Vec<PtxToken>) {
258        let mut repr = self.name.clone();
259        if let Some(component) = &self.component {
260            repr.push('.');
261            repr.push_str(component);
262        }
263        push_register(tokens, &repr);
264    }
265}
266
267impl PtxUnparser for PredicateRegister {
268    fn unparse_tokens(&self, tokens: &mut Vec<PtxToken>) {
269        push_register(tokens, &self.name);
270    }
271}
272
273impl PtxUnparser for Label {
274    fn unparse_tokens(&self, tokens: &mut Vec<PtxToken>) {
275        // Design note: labels may be directive-shaped section names such as
276        // `.debug_str`; tokenizing preserves both those and plain identifiers.
277        push_token_from_str(tokens, &self.val);
278    }
279}
280
281impl PtxUnparser for SpecialRegister {
282    fn unparse_tokens(&self, tokens: &mut Vec<PtxToken>) {
283        let name = match self {
284            SpecialRegister::AggrSmemSize { .. } => "%aggr_smem_size".to_string(),
285            SpecialRegister::DynamicSmemSize { .. } => "%dynamic_smem_size".to_string(),
286            SpecialRegister::LanemaskGt { .. } => "%lanemask_gt".to_string(),
287            SpecialRegister::ReservedSmemOffsetBegin { .. } => {
288                "%reserved_smem_offset_begin".to_string()
289            }
290            SpecialRegister::Clock { .. } => "%clock".to_string(),
291            SpecialRegister::Envreg { index, .. } => format!("%envreg{}", index),
292            SpecialRegister::LanemaskLe { .. } => "%lanemask_le".to_string(),
293            SpecialRegister::ReservedSmemOffsetCap { .. } => {
294                "%reserved_smem_offset_cap".to_string()
295            }
296            SpecialRegister::Clock64 { .. } => "%clock64".to_string(),
297            SpecialRegister::Globaltimer { .. } => "%globaltimer".to_string(),
298            SpecialRegister::LanemaskLt { .. } => "%lanemask_lt".to_string(),
299            SpecialRegister::ReservedSmemOffsetEnd { .. } => {
300                "%reserved_smem_offset_end".to_string()
301            }
302            SpecialRegister::ClusterCtaid { axis, .. } => {
303                push_register_with_axis(tokens, "%cluster_ctaid", axis);
304                return;
305            }
306            SpecialRegister::GlobaltimerHi { .. } => "%globaltimer_hi".to_string(),
307            SpecialRegister::Nclusterid { .. } => "%nclusterid".to_string(),
308            SpecialRegister::Smid { .. } => "%smid".to_string(),
309            SpecialRegister::ClusterCtarank { axis, .. } => {
310                push_register_with_axis(tokens, "%cluster_ctarank", axis);
311                return;
312            }
313            SpecialRegister::GlobaltimerLo { .. } => "%globaltimer_lo".to_string(),
314            SpecialRegister::Nctaid { axis, .. } => {
315                push_register_with_axis(tokens, "%nctaid", axis);
316                return;
317            }
318            SpecialRegister::Tid { axis, .. } => {
319                push_register_with_axis(tokens, "%tid", axis);
320                return;
321            }
322            SpecialRegister::ClusterNctaid { axis, .. } => {
323                push_register_with_axis(tokens, "%cluster_nctaid", axis);
324                return;
325            }
326            SpecialRegister::Gridid { .. } => "%gridid".to_string(),
327            SpecialRegister::Nsmid { .. } => "%nsmid".to_string(),
328            SpecialRegister::TotalSmemSize { .. } => "%total_smem_size".to_string(),
329            SpecialRegister::ClusterNctarank { axis, .. } => {
330                push_register_with_axis(tokens, "%cluster_nctarank", axis);
331                return;
332            }
333            SpecialRegister::IsExplicitCluster { .. } => "%is_explicit_cluster".to_string(),
334            SpecialRegister::Ntid { axis, .. } => {
335                push_register_with_axis(tokens, "%ntid", axis);
336                return;
337            }
338            SpecialRegister::Warpid { .. } => "%warpid".to_string(),
339            SpecialRegister::Clusterid { .. } => "%clusterid".to_string(),
340            SpecialRegister::Laneid { .. } => "%laneid".to_string(),
341            SpecialRegister::Nwarpid { .. } => "%nwarpid".to_string(),
342            SpecialRegister::WARPSZ { .. } => "%WARPSZ".to_string(),
343            SpecialRegister::Ctaid { axis, .. } => {
344                push_register_with_axis(tokens, "%ctaid", axis);
345                return;
346            }
347            SpecialRegister::LanemaskEq { .. } => "%lanemask_eq".to_string(),
348            SpecialRegister::Pm { index, .. } => format!("%pm{}", index),
349            SpecialRegister::Pm64 { index, .. } => format!("%pm{}_64", index),
350            SpecialRegister::CurrentGraphExec { .. } => "%current_graph_exec".to_string(),
351            SpecialRegister::LanemaskGe { .. } => "%lanemask_ge".to_string(),
352            SpecialRegister::ReservedSmemOffset { index, .. } => {
353                format!("%reserved_smem_offset_{}", index)
354            }
355        };
356        push_register(tokens, &name);
357    }
358}
359
360impl PtxUnparser for Operand {
361    fn unparse_tokens(&self, tokens: &mut Vec<PtxToken>) {
362        match self {
363            Operand::Register {
364                operand: register, ..
365            } => register.unparse_tokens(tokens),
366            Operand::Immediate {
367                operand: immediate, ..
368            } => immediate.unparse_tokens(tokens),
369            Operand::Symbol { name: symbol, .. } => push_identifier(tokens, symbol),
370            Operand::VectorSymbolComponent {
371                symbol, component, ..
372            } => {
373                push_identifier(tokens, symbol);
374                push_directive(tokens, component);
375            }
376            Operand::SymbolOffset { symbol, offset, .. } => {
377                push_identifier(tokens, symbol);
378                tokens.push(PtxToken::Plus);
379                offset.unparse_tokens(tokens);
380            }
381        }
382    }
383}
384
385impl PtxUnparser for VectorOperand {
386    fn unparse_tokens(&self, tokens: &mut Vec<PtxToken>) {
387        tokens.push(PtxToken::LBrace);
388        match self {
389            VectorOperand::Vector1 { operand: item, .. } => item.unparse_tokens(tokens),
390            VectorOperand::Vector2 {
391                operands: items, ..
392            } => {
393                for (idx, item) in items.iter().enumerate() {
394                    if idx > 0 {
395                        tokens.push(PtxToken::Comma);
396                    }
397                    item.unparse_tokens(tokens);
398                }
399            }
400            VectorOperand::Vector3 {
401                operands: items, ..
402            } => {
403                for (idx, item) in items.iter().enumerate() {
404                    if idx > 0 {
405                        tokens.push(PtxToken::Comma);
406                    }
407                    item.unparse_tokens(tokens);
408                }
409            }
410            VectorOperand::Vector4 {
411                operands: items, ..
412            } => {
413                for (idx, item) in items.iter().enumerate() {
414                    if idx > 0 {
415                        tokens.push(PtxToken::Comma);
416                    }
417                    item.unparse_tokens(tokens);
418                }
419            }
420            VectorOperand::Vector8 {
421                operands: items, ..
422            } => {
423                for (idx, item) in items.iter().enumerate() {
424                    if idx > 0 {
425                        tokens.push(PtxToken::Comma);
426                    }
427                    item.unparse_tokens(tokens);
428                }
429            }
430        }
431        tokens.push(PtxToken::RBrace);
432    }
433}
434
435impl PtxUnparser for GeneralOperand {
436    fn unparse_tokens(&self, tokens: &mut Vec<PtxToken>) {
437        match self {
438            GeneralOperand::Vec {
439                operand: vector, ..
440            } => vector.unparse_tokens(tokens),
441            GeneralOperand::Single { operand, .. } => operand.unparse_tokens(tokens),
442        }
443    }
444}
445
446impl PtxUnparser for TexHandler2 {
447    fn unparse_tokens(&self, tokens: &mut Vec<PtxToken>) {
448        tokens.push(PtxToken::LBracket);
449        for (idx, item) in self.operands.iter().enumerate() {
450            if idx > 0 {
451                tokens.push(PtxToken::Comma);
452            }
453            item.unparse_tokens(tokens);
454        }
455        tokens.push(PtxToken::RBracket);
456    }
457}
458
459impl PtxUnparser for TexHandler3 {
460    fn unparse_tokens(&self, tokens: &mut Vec<PtxToken>) {
461        tokens.push(PtxToken::LBracket);
462        self.handle.unparse_tokens(tokens);
463        tokens.push(PtxToken::Comma);
464        self.sampler.unparse_tokens(tokens);
465        tokens.push(PtxToken::Comma);
466        self.coords.unparse_tokens(tokens);
467        tokens.push(PtxToken::RBracket);
468    }
469}
470
471impl PtxUnparser for TexHandler3Optional {
472    fn unparse_tokens(&self, tokens: &mut Vec<PtxToken>) {
473        tokens.push(PtxToken::LBracket);
474        self.handle.unparse_tokens(tokens);
475        tokens.push(PtxToken::Comma);
476        if let Some(sampler) = &self.sampler {
477            sampler.unparse_tokens(tokens);
478            tokens.push(PtxToken::Comma);
479        }
480        self.coords.unparse_tokens(tokens);
481        tokens.push(PtxToken::RBracket);
482    }
483}
484
485impl PtxUnparser for AddressBase {
486    fn unparse_tokens(&self, tokens: &mut Vec<PtxToken>) {
487        match self {
488            AddressBase::Register {
489                operand: register, ..
490            } => register.unparse_tokens(tokens),
491            AddressBase::Variable { symbol, .. } => symbol.unparse_tokens(tokens),
492        }
493    }
494}
495
496impl PtxUnparser for AddressOffset {
497    fn unparse_tokens(&self, tokens: &mut Vec<PtxToken>) {
498        match self {
499            AddressOffset::Register {
500                operand: register, ..
501            } => {
502                tokens.push(PtxToken::Plus);
503                register.unparse_tokens(tokens);
504            }
505            AddressOffset::Immediate {
506                sign,
507                value: immediate,
508                ..
509            } => {
510                // Design note: the offset separator is always `+`; a negative
511                // immediate adds its own `-`, producing ptxas syntax `+-N`.
512                tokens.push(PtxToken::Plus);
513                if matches!(sign, Sign::Negative { .. }) {
514                    tokens.push(PtxToken::Minus);
515                }
516                immediate.unparse_tokens(tokens);
517            }
518        }
519    }
520}
521
522impl PtxUnparser for AddressOperand {
523    fn unparse_tokens(&self, tokens: &mut Vec<PtxToken>) {
524        match self {
525            AddressOperand::Array { base, index, .. } => {
526                base.unparse_tokens(tokens);
527                tokens.push(PtxToken::LBracket);
528                index.unparse_tokens(tokens);
529                tokens.push(PtxToken::RBracket);
530            }
531            AddressOperand::ImmediateAddress { addr, .. } => {
532                tokens.push(PtxToken::LBracket);
533                addr.unparse_tokens(tokens);
534                tokens.push(PtxToken::RBracket);
535            }
536            AddressOperand::Offset { base, offset, .. } => {
537                tokens.push(PtxToken::LBracket);
538                base.unparse_tokens(tokens);
539                if let Some(offset) = offset {
540                    offset.unparse_tokens(tokens);
541                }
542                tokens.push(PtxToken::RBracket);
543            }
544        }
545    }
546}
547
548impl PtxUnparser for FunctionSymbol {
549    fn unparse_tokens(&self, tokens: &mut Vec<PtxToken>) {
550        push_identifier(tokens, &self.val);
551    }
552}
553
554impl PtxUnparser for VariableSymbol {
555    fn unparse_tokens(&self, tokens: &mut Vec<PtxToken>) {
556        push_identifier(tokens, &self.val);
557    }
558}
559
560impl PtxUnparser for crate::r#type::common::Instruction {
561    fn unparse_tokens(&self, tokens: &mut Vec<PtxToken>) {
562        self.unparse_tokens_mode(tokens, false);
563    }
564
565    fn unparse_tokens_mode(&self, tokens: &mut Vec<PtxToken>, spaced: bool) {
566        // Emit predicate if present
567        if let Some(predicate) = &self.predicate {
568            tokens.push(PtxToken::At);
569            if predicate.negated {
570                tokens.push(PtxToken::Exclaim);
571            }
572            predicate.operand.unparse_tokens_mode(tokens, spaced);
573            push_space(tokens, spaced);
574        }
575
576        // Emit the instruction
577        self.inst.unparse_tokens_mode(tokens, spaced);
578    }
579}
ptx_parser/unparser/common.rs

ptx_parser/unparser/
common.rs
