ptx_parser/type/
common.rs

1/* --------------------------------------------------- */
2/* -------------------- Basic Directives ------------- */
3/* --------------------------------------------------- */
4
5use crate::parser::Span;
6
7#[derive(Debug, Clone, PartialEq, Eq)]
8pub enum CodeLinkage {
9    /// `.visible`
10    Visible { span: Span },
11    /// `.extern`
12    Extern { span: Span },
13    /// `.weak`
14    Weak { span: Span },
15}
16
17impl CodeLinkage {
18    pub fn span(&self) -> Span {
19        match self {
20            CodeLinkage::Visible { span } => span.clone(),
21            CodeLinkage::Extern { span } => span.clone(),
22            CodeLinkage::Weak { span } => span.clone(),
23        }
24    }
25}
26
27#[derive(Debug, Clone, PartialEq, Eq)]
28pub enum DataLinkage {
29    /// `.visible`
30    Visible { span: Span },
31    /// `.extern`
32    Extern { span: Span },
33    /// `.weak`
34    Weak { span: Span },
35    /// `.common`
36    Common { span: Span },
37}
38
39impl DataLinkage {
40    pub fn span(&self) -> Span {
41        match self {
42            DataLinkage::Visible { span } => span.clone(),
43            DataLinkage::Extern { span } => span.clone(),
44            DataLinkage::Weak { span } => span.clone(),
45            DataLinkage::Common { span } => span.clone(),
46        }
47    }
48}
49
50#[derive(Debug, Clone, PartialEq, Eq)]
51pub enum CodeOrDataLinkage {
52    /// `.visible`
53    Visible { span: Span },
54    /// `.extern`
55    Extern { span: Span },
56    /// `.weak`
57    Weak { span: Span },
58    /// `.common`
59    Common { span: Span },
60}
61
62impl CodeOrDataLinkage {
63    pub fn span(&self) -> Span {
64        match self {
65            CodeOrDataLinkage::Visible { span } => span.clone(),
66            CodeOrDataLinkage::Extern { span } => span.clone(),
67            CodeOrDataLinkage::Weak { span } => span.clone(),
68            CodeOrDataLinkage::Common { span } => span.clone(),
69        }
70    }
71}
72
73#[derive(Debug, Clone, PartialEq, Eq)]
74pub enum TexType {
75    /// `.texref`
76    TexRef { span: Span },
77    /// `.samplerref`
78    SamplerRef { span: Span },
79    /// `.surfref`
80    SurfRef { span: Span },
81}
82
83impl TexType {
84    pub fn span(&self) -> Span {
85        match self {
86            TexType::TexRef { span } => span.clone(),
87            TexType::SamplerRef { span } => span.clone(),
88            TexType::SurfRef { span } => span.clone(),
89        }
90    }
91}
92
93/// Memory spaces
94#[derive(Debug, Clone, PartialEq, Eq)]
95pub enum AddressSpace {
96    Global { span: Span },
97    Const { span: Span },
98    Shared { span: Span },
99    Local { span: Span },
100    Param { span: Span },
101    Reg { span: Span },
102}
103
104impl AddressSpace {
105    pub fn span(&self) -> Span {
106        match self {
107            AddressSpace::Global { span } => span.clone(),
108            AddressSpace::Const { span } => span.clone(),
109            AddressSpace::Shared { span } => span.clone(),
110            AddressSpace::Local { span } => span.clone(),
111            AddressSpace::Param { span } => span.clone(),
112            AddressSpace::Reg { span } => span.clone(),
113        }
114    }
115}
116
117#[derive(Debug, Clone, PartialEq, Eq)]
118pub enum AttributeDirective {
119    /// `.unified(uuid1, uuid2)`
120    Unified { uuid1: u64, uuid2: u64, span: Span },
121    /// `.managed`
122    Managed { span: Span },
123}
124
125impl AttributeDirective {
126    pub fn span(&self) -> Span {
127        match self {
128            AttributeDirective::Unified { span, .. } => span.clone(),
129            AttributeDirective::Managed { span } => span.clone(),
130        }
131    }
132}
133
134#[derive(Debug, Clone, PartialEq, Eq)]
135pub enum DataType {
136    /// `.u8`
137    U8 { span: Span },
138    /// `.u16`
139    U16 { span: Span },
140    /// `.u32`
141    U32 { span: Span },
142    /// `.u64`
143    U64 { span: Span },
144    /// `.s8`
145    S8 { span: Span },
146    /// `.s16`
147    S16 { span: Span },
148    /// `.s32`
149    S32 { span: Span },
150    /// `.s64`
151    S64 { span: Span },
152    /// `.f16`
153    F16 { span: Span },
154    /// `.f16x2`
155    F16x2 { span: Span },
156    /// `.f32`
157    F32 { span: Span },
158    /// `.f64`
159    F64 { span: Span },
160    /// `.b8`
161    B8 { span: Span },
162    /// `.b16`
163    B16 { span: Span },
164    /// `.b32`
165    B32 { span: Span },
166    /// `.b64`
167    B64 { span: Span },
168    /// `.b128`
169    B128 { span: Span },
170    /// `.pred`
171    Pred { span: Span },
172}
173
174impl DataType {
175    pub fn span(&self) -> Span {
176        match self {
177            DataType::U8 { span } => span.clone(),
178            DataType::U16 { span } => span.clone(),
179            DataType::U32 { span } => span.clone(),
180            DataType::U64 { span } => span.clone(),
181            DataType::S8 { span } => span.clone(),
182            DataType::S16 { span } => span.clone(),
183            DataType::S32 { span } => span.clone(),
184            DataType::S64 { span } => span.clone(),
185            DataType::F16 { span } => span.clone(),
186            DataType::F16x2 { span } => span.clone(),
187            DataType::F32 { span } => span.clone(),
188            DataType::F64 { span } => span.clone(),
189            DataType::B8 { span } => span.clone(),
190            DataType::B16 { span } => span.clone(),
191            DataType::B32 { span } => span.clone(),
192            DataType::B64 { span } => span.clone(),
193            DataType::B128 { span } => span.clone(),
194            DataType::Pred { span } => span.clone(),
195        }
196    }
197}
198
199/* -------------------------------------------------- */
200/* -------------------- Math Basics ----------------- */
201/* -------------------------------------------------- */
202
203#[derive(Debug, Clone, PartialEq, Eq)]
204pub enum Sign {
205    Negative { span: Span },
206    Positive { span: Span },
207}
208
209impl Sign {
210    pub fn span(&self) -> Span {
211        match self {
212            Sign::Negative { span } => span.clone(),
213            Sign::Positive { span } => span.clone(),
214        }
215    }
216}
217
218/// Axis component for 3-component special registers (x/y/z)
219#[derive(Debug, Clone, PartialEq, Eq)]
220pub enum Axis {
221    /// No axis component present
222    None { span: Span },
223    X { span: Span },
224    Y { span: Span },
225    Z { span: Span },
226}
227
228impl Axis {
229    pub fn span(&self) -> Span {
230        match self {
231            Axis::None { span } => span.clone(),
232            Axis::X { span } => span.clone(),
233            Axis::Y { span } => span.clone(),
234            Axis::Z { span } => span.clone(),
235        }
236    }
237}
238
239/* --------------------------------------------------- */
240/* -------------------- Label ------------------------ */
241/* --------------------------------------------------- */
242
243/// Label name (e.g., `L__label_1`).
244#[derive(Debug, Clone, PartialEq, Eq)]
245pub struct Label {
246    pub name: String,
247    pub span: Span,
248}
249
250impl Label {
251    pub fn with_span(mut self, span: Span) -> Self {
252        self.span = span;
253        self
254    }
255}
256
257/* --------------------------------------------------- */
258/* -------------------- Special Registers ------------ */
259/* --------------------------------------------------- */
260
261#[derive(Debug, Clone, PartialEq, Eq)]
262pub enum SpecialRegister {
263    /// `%aggr_smem_size`
264    AggrSmemSize { span: Span },
265    /// `%dynamic_smem_size`
266    DynamicSmemSize { span: Span },
267    /// `%lanemask_gt`
268    LanemaskGt { span: Span },
269    /// `%reserved_smem_offset_begin`
270    ReservedSmemOffsetBegin { span: Span },
271    /// `%clock`
272    Clock { span: Span },
273    /// `%envreg0`, ..., `%envreg31`
274    Envreg { index: u8, span: Span },
275    /// `%lanemask_le`
276    LanemaskLe { span: Span },
277    /// `%reserved_smem_offset_cap`
278    ReservedSmemOffsetCap { span: Span },
279    /// `%clock64`
280    Clock64 { span: Span },
281    /// `%globaltimer`
282    Globaltimer { span: Span },
283    /// `%lanemask_lt`
284    LanemaskLt { span: Span },
285    /// `%reserved_smem_offset_end`
286    ReservedSmemOffsetEnd { span: Span },
287    /// `%cluster_ctaid` (optionally `.x/.y/.z`)
288    ClusterCtaid { axis: Axis, span: Span },
289    /// `%globaltimer_hi`
290    GlobaltimerHi { span: Span },
291    /// `%nclusterid`
292    Nclusterid { span: Span },
293    /// `%smid`
294    Smid { span: Span },
295    /// `%cluster_ctarank` (optionally `.x/.y/.z`)
296    ClusterCtarank { axis: Axis, span: Span },
297    /// `%globaltimer_lo`
298    GlobaltimerLo { span: Span },
299    /// `%nctaid` (optionally `.x/.y/.z`)
300    Nctaid { axis: Axis, span: Span },
301    /// `%tid` (optionally `.x/.y/.z`)
302    Tid { axis: Axis, span: Span },
303    /// `%cluster_nctaid` (optionally `.x/.y/.z`)
304    ClusterNctaid { axis: Axis, span: Span },
305    /// `%gridid`
306    Gridid { span: Span },
307    /// `%nsmid`
308    Nsmid { span: Span },
309    /// `%total_smem_size`
310    TotalSmemSize { span: Span },
311    /// `%cluster_nctarank` (optionally `.x/.y/.z`)
312    ClusterNctarank { axis: Axis, span: Span },
313    /// `%is_explicit_cluster`
314    IsExplicitCluster { span: Span },
315    /// `%ntid` (optionally `.x/.y/.z`)
316    Ntid { axis: Axis, span: Span },
317    /// `%warpid`
318    Warpid { span: Span },
319    /// `%clusterid`
320    Clusterid { span: Span },
321    /// `%laneid`
322    Laneid { span: Span },
323    /// `%nwarpid`
324    Nwarpid { span: Span },
325    /// `%WARPSZ`
326    WARPSZ { span: Span },
327    /// `%ctaid` (optionally `.x/.y/.z`)
328    Ctaid { axis: Axis, span: Span },
329    /// `%lanemask_eq`
330    LanemaskEq { span: Span },
331    /// `%pm0`, ..., `%pm7`
332    Pm { index: u8, span: Span },
333    /// `%pm0_64`, ..., `%pm7_64`
334    Pm64 { index: u8, span: Span },
335    /// `%current_graph_exec`
336    CurrentGraphExec { span: Span },
337    /// `%lanemask_ge`
338    LanemaskGe { span: Span },
339    /// `%reserved_smem_offset_0`, `%reserved_smem_offset_1`
340    ReservedSmemOffset { index: u8, span: Span },
341}
342
343impl SpecialRegister {
344    pub fn span(&self) -> Span {
345        match self {
346            SpecialRegister::AggrSmemSize { span } => span.clone(),
347            SpecialRegister::DynamicSmemSize { span } => span.clone(),
348            SpecialRegister::LanemaskGt { span } => span.clone(),
349            SpecialRegister::ReservedSmemOffsetBegin { span } => span.clone(),
350            SpecialRegister::Clock { span } => span.clone(),
351            SpecialRegister::Envreg { span, .. } => span.clone(),
352            SpecialRegister::LanemaskLe { span } => span.clone(),
353            SpecialRegister::ReservedSmemOffsetCap { span } => span.clone(),
354            SpecialRegister::Clock64 { span } => span.clone(),
355            SpecialRegister::Globaltimer { span } => span.clone(),
356            SpecialRegister::LanemaskLt { span } => span.clone(),
357            SpecialRegister::ReservedSmemOffsetEnd { span } => span.clone(),
358            SpecialRegister::ClusterCtaid { span, .. } => span.clone(),
359            SpecialRegister::GlobaltimerHi { span } => span.clone(),
360            SpecialRegister::Nclusterid { span } => span.clone(),
361            SpecialRegister::Smid { span } => span.clone(),
362            SpecialRegister::ClusterCtarank { span, .. } => span.clone(),
363            SpecialRegister::GlobaltimerLo { span } => span.clone(),
364            SpecialRegister::Nctaid { span, .. } => span.clone(),
365            SpecialRegister::Tid { span, .. } => span.clone(),
366            SpecialRegister::ClusterNctaid { span, .. } => span.clone(),
367            SpecialRegister::Gridid { span } => span.clone(),
368            SpecialRegister::Nsmid { span } => span.clone(),
369            SpecialRegister::TotalSmemSize { span } => span.clone(),
370            SpecialRegister::ClusterNctarank { span, .. } => span.clone(),
371            SpecialRegister::IsExplicitCluster { span } => span.clone(),
372            SpecialRegister::Ntid { span, .. } => span.clone(),
373            SpecialRegister::Warpid { span } => span.clone(),
374            SpecialRegister::Clusterid { span } => span.clone(),
375            SpecialRegister::Laneid { span } => span.clone(),
376            SpecialRegister::Nwarpid { span } => span.clone(),
377            SpecialRegister::WARPSZ { span } => span.clone(),
378            SpecialRegister::Ctaid { span, .. } => span.clone(),
379            SpecialRegister::LanemaskEq { span } => span.clone(),
380            SpecialRegister::Pm { span, .. } => span.clone(),
381            SpecialRegister::Pm64 { span, .. } => span.clone(),
382            SpecialRegister::CurrentGraphExec { span } => span.clone(),
383            SpecialRegister::LanemaskGe { span } => span.clone(),
384            SpecialRegister::ReservedSmemOffset { span, .. } => span.clone(),
385        }
386    }
387}
388
389/* --------------------------------------------------- */
390/* ------------------- Operands ---------------------- */
391/* --------------------------------------------------- */
392
393/// Texture handler with 2 operands, e.g. [%r1, %r2]
394#[derive(Debug, Clone, PartialEq, Eq)]
395pub struct TexHandler2 {
396    pub operands: [GeneralOperand; 2],
397    pub span: Span,
398}
399
400impl TexHandler2 {
401    pub fn with_span(mut self, span: Span) -> Self {
402        self.span = span;
403        self
404    }
405}
406
407/// Texture handler with optional sampler operand, e.g. `[tex, coords]` or `[tex, sampler, coords]`
408#[derive(Debug, Clone, PartialEq, Eq)]
409pub struct TexHandler3Optional {
410    pub handle: GeneralOperand,
411    pub sampler: Option<GeneralOperand>,
412    pub coords: GeneralOperand,
413    pub span: Span,
414}
415
416impl TexHandler3Optional {
417    pub fn with_span(mut self, span: Span) -> Self {
418        self.span = span;
419        self
420    }
421}
422
423/// Texture handler with optional sampler operand, e.g. `[tex, coords]` or `[tex, sampler, coords]`
424#[derive(Debug, Clone, PartialEq, Eq)]
425pub struct TexHandler3 {
426    pub handle: GeneralOperand,
427    pub sampler: GeneralOperand,
428    pub coords: GeneralOperand,
429    pub span: Span,
430}
431
432impl TexHandler3 {
433    pub fn with_span(mut self, span: Span) -> Self {
434        self.span = span;
435        self
436    }
437}
438
439#[derive(Debug, Clone, PartialEq, Eq)]
440pub enum GeneralOperand {
441    Vec { operand: VectorOperand, span: Span },
442    Single { operand: Operand, span: Span },
443}
444
445impl GeneralOperand {
446    pub fn span(&self) -> Span {
447        match self {
448            GeneralOperand::Vec { span, .. } => span.clone(),
449            GeneralOperand::Single { span, .. } => span.clone(),
450        }
451    }
452}
453
454#[derive(Debug, Clone, PartialEq, Eq)]
455pub enum VectorOperand {
456    /// {%r1}
457    Vector1 { operand: Operand, span: Span },
458    /// {%r1, %r2}
459    Vector2 { operands: [Operand; 2], span: Span },
460    /// {%r1, %r2, %r3}
461    Vector3 { operands: [Operand; 3], span: Span },
462    /// {%r1, %r2, %r3, %r4}
463    Vector4 { operands: [Operand; 4], span: Span },
464    /// {%r1, %r2, %r3, %r4, %r5, %r6, %r7, %r8}
465    Vector8 { operands: [Operand; 8], span: Span },
466}
467
468impl VectorOperand {
469    pub fn span(&self) -> Span {
470        match self {
471            VectorOperand::Vector1 { span, .. } => span.clone(),
472            VectorOperand::Vector2 { span, .. } => span.clone(),
473            VectorOperand::Vector3 { span, .. } => span.clone(),
474            VectorOperand::Vector4 { span, .. } => span.clone(),
475            VectorOperand::Vector8 { span, .. } => span.clone(),
476        }
477    }
478}
479
480#[derive(Debug, Clone, PartialEq, Eq)]
481pub enum Operand {
482    /// %r1
483    Register { operand: RegisterOperand, span: Span },
484    /// 0xffff
485    Immediate { operand: Immediate, span: Span },
486    /// foo
487    Symbol { name: String, span: Span },
488    /// foo + 4 (symbol + immediate offset)
489    SymbolOffset { symbol: String, offset: Immediate, span: Span },
490}
491
492impl Operand {
493    pub fn span(&self) -> Span {
494        match self {
495            Operand::Register { span, .. } => span.clone(),
496            Operand::Immediate { span, .. } => span.clone(),
497            Operand::Symbol { span, .. } => span.clone(),
498            Operand::SymbolOffset { span, .. } => span.clone(),
499        }
500    }
501}
502
503/// Register operand starting with % (e.g., `%r1`).
504#[derive(Debug, Clone, PartialEq, Eq)]
505pub struct RegisterOperand {
506    pub name: String,
507    pub span: Span,
508}
509
510impl RegisterOperand {
511    pub fn with_span(mut self, span: Span) -> Self {
512        self.span = span;
513        self
514    }
515}
516
517/// Predicate register names (e.g., `%p0`).
518#[derive(Debug, Clone, PartialEq, Eq)]
519pub struct PredicateRegister {
520    pub name: String,
521    pub span: Span,
522}
523
524impl PredicateRegister {
525    pub fn with_span(mut self, span: Span) -> Self {
526        self.span = span;
527        self
528    }
529}
530
531/// Representation of an address operand.
532#[derive(Debug, Clone, PartialEq, Eq)]
533pub enum AddressOperand {
534    /// base[immIndex]
535    Array { base: VariableSymbol, index: Immediate, span: Span },
536    /// Immediate address value, e.g., [0xffff], unsigned, 32-bit
537    ImmediateAddress { addr: Immediate, span: Span },
538    /// Offset address with optional displacement, e.g., [base + offset] and [base]
539    Offset { base: AddressBase, offset: Option<AddressOffset>, span: Span },
540}
541
542impl AddressOperand {
543    pub fn span(&self) -> Span {
544        match self {
545            AddressOperand::Array { span, .. } => span.clone(),
546            AddressOperand::ImmediateAddress { span, .. } => span.clone(),
547            AddressOperand::Offset { span, .. } => span.clone(),
548        }
549    }
550}
551
552/// Base location referenced by an address expression.
553#[derive(Debug, Clone, PartialEq, Eq)]
554pub enum AddressBase {
555    Register { operand: RegisterOperand, span: Span },
556    Variable { symbol: VariableSymbol, span: Span },
557}
558
559impl AddressBase {
560    pub fn span(&self) -> Span {
561        match self {
562            AddressBase::Register { span, .. } => span.clone(),
563            AddressBase::Variable { span, .. } => span.clone(),
564        }
565    }
566}
567
568/// Specific adjustment applied within a displacement term.
569#[derive(Debug, Clone, PartialEq, Eq)]
570pub enum AddressOffset {
571    Register { operand: RegisterOperand, span: Span },
572    Immediate { sign: Sign, value: Immediate, span: Span },
573}
574
575impl AddressOffset {
576    pub fn span(&self) -> Span {
577        match self {
578            AddressOffset::Register { span, .. } => span.clone(),
579            AddressOffset::Immediate { span, .. } => span.clone(),
580        }
581    }
582}
583
584/* --------------------------------------------------- */
585/* -------------------- Immediate -------------------- */
586/* --------------------------------------------------- */
587
588/// Immediate value represented as a string.
589/// TODO: Replace with appropriate numeric type later.
590#[derive(Debug, Clone, PartialEq, Eq)]
591pub struct Immediate {
592    pub value: String,
593    pub span: Span,
594}
595
596impl Immediate {
597    pub fn with_span(mut self, span: Span) -> Self {
598        self.span = span;
599        self
600    }
601}
602
603/* --------------------------------------------------- */
604/* -------------------- Symbols ---------------------- */
605/* --------------------------------------------------- */
606
607/// Function symbol
608#[derive(Debug, Clone, PartialEq, Eq)]
609pub struct FunctionSymbol {
610    pub name: String,
611    pub span: Span,
612}
613
614impl FunctionSymbol {
615    pub fn with_span(mut self, span: Span) -> Self {
616        self.span = span;
617        self
618    }
619}
620
621/// Variable symbol
622#[derive(Debug, Clone, PartialEq, Eq)]
623pub struct VariableSymbol {
624    pub name: String,
625    pub span: Span,
626}
627
628impl VariableSymbol {
629    pub fn with_span(mut self, span: Span) -> Self {
630        self.span = span;
631        self
632    }
633}
634
635/* --------------------------------------------------- */
636/* -------------------- Predicate -------------------- */
637/* --------------------------------------------------- */
638
639/// Predicate guard for conditional instruction execution
640#[derive(Debug, Clone, PartialEq)]
641pub struct Predicate {
642    pub negated: bool,
643    pub operand: Operand,
644    pub span: Span,
645}
646
647impl Predicate {
648    pub fn with_span(mut self, span: Span) -> Self {
649        self.span = span;
650        self
651    }
652}
653
654/* --------------------------------------------------- */
655/* -------------------- Instruction ------------------ */
656/* --------------------------------------------------- */
657
658/// Represents a complete instruction with optional label and predicate guard
659/// Format: [label:] [@{!}pred] instruction
660#[derive(Debug, Clone, PartialEq)]
661pub struct Instruction {
662    pub predicate: Option<Predicate>,
663    pub inst: crate::r#type::instruction::Inst,
664    pub span: Span,
665}
666
667impl Instruction {
668    pub fn with_span(mut self, span: Span) -> Self {
669        self.span = span;
670        self
671    }
672}
ptx_parser/type/common.rs

ptx_parser/type/
common.rs
