ptx_parser/type/
common.rs

1/* --------------------------------------------------- */
2/* -------------------- Basic Directives ------------- */
3/* --------------------------------------------------- */
4
5use crate::Spanned;
6use crate::parser::Span;
7
8#[derive(Debug, Clone, PartialEq, Eq, Spanned)]
9pub enum CodeLinkage {
10    /// `.visible`
11    Visible { span: Span },
12    /// `.extern`
13    Extern { span: Span },
14    /// `.weak`
15    Weak { span: Span },
16}
17
18#[derive(Debug, Clone, PartialEq, Eq, Spanned)]
19pub enum DataLinkage {
20    /// `.visible`
21    Visible { span: Span },
22    /// `.extern`
23    Extern { span: Span },
24    /// `.weak`
25    Weak { span: Span },
26    /// `.common`
27    Common { span: Span },
28}
29
30#[derive(Debug, Clone, PartialEq, Eq, Spanned)]
31pub enum AttributeDirective {
32    /// `.unified(uuid1, uuid2)`
33    Unified { uuid1: u64, uuid2: u64, span: Span },
34    /// `.managed`
35    Managed { span: Span },
36}
37
38#[derive(Debug, Clone, PartialEq, Eq, Spanned)]
39pub enum DataType {
40    /// `.u8`
41    U8 { span: Span },
42    /// `.u16`
43    U16 { span: Span },
44    /// `.u32`
45    U32 { span: Span },
46    /// `.u64`
47    U64 { span: Span },
48    /// `.s8`
49    S8 { span: Span },
50    /// `.s16`
51    S16 { span: Span },
52    /// `.s32`
53    S32 { span: Span },
54    /// `.s64`
55    S64 { span: Span },
56    /// `.f16`
57    F16 { span: Span },
58    /// `.f16x2`
59    F16x2 { span: Span },
60    /// `.f32`
61    F32 { span: Span },
62    /// `.f64`
63    F64 { span: Span },
64    /// `.b8`
65    B8 { span: Span },
66    /// `.b16`
67    B16 { span: Span },
68    /// `.b32`
69    B32 { span: Span },
70    /// `.b64`
71    B64 { span: Span },
72    /// `.b128`
73    B128 { span: Span },
74    /// `.pred`
75    Pred { span: Span },
76    /// `.texref`
77    TexRef { span: Span },
78    /// `.samplerref`
79    SamplerRef { span: Span },
80    /// `.surfref`
81    SurfRef { span: Span },
82}
83
84/* -------------------------------------------------- */
85/* -------------------- Math Basics ----------------- */
86/* -------------------------------------------------- */
87
88#[derive(Debug, Clone, PartialEq, Eq, Spanned)]
89pub enum Sign {
90    Negative { span: Span },
91    Positive { span: Span },
92}
93
94/// Axis component for 3-component special registers (x/y/z)
95#[derive(Debug, Clone, PartialEq, Eq, Spanned)]
96pub enum Axis {
97    /// No axis component present
98    None {
99        span: Span,
100    },
101    X {
102        span: Span,
103    },
104    Y {
105        span: Span,
106    },
107    Z {
108        span: Span,
109    },
110}
111
112/* --------------------------------------------------- */
113/* -------------------- Special Registers ------------ */
114/* --------------------------------------------------- */
115
116#[derive(Debug, Clone, PartialEq, Eq, Spanned)]
117pub enum SpecialRegister {
118    /// `%aggr_smem_size`
119    AggrSmemSize { span: Span },
120    /// `%dynamic_smem_size`
121    DynamicSmemSize { span: Span },
122    /// `%lanemask_gt`
123    LanemaskGt { span: Span },
124    /// `%reserved_smem_offset_begin`
125    ReservedSmemOffsetBegin { span: Span },
126    /// `%clock`
127    Clock { span: Span },
128    /// `%envreg0`, ..., `%envreg31`
129    Envreg { index: u8, span: Span },
130    /// `%lanemask_le`
131    LanemaskLe { span: Span },
132    /// `%reserved_smem_offset_cap`
133    ReservedSmemOffsetCap { span: Span },
134    /// `%clock64`
135    Clock64 { span: Span },
136    /// `%globaltimer`
137    Globaltimer { span: Span },
138    /// `%lanemask_lt`
139    LanemaskLt { span: Span },
140    /// `%reserved_smem_offset_end`
141    ReservedSmemOffsetEnd { span: Span },
142    /// `%cluster_ctaid` (optionally `.x/.y/.z`)
143    ClusterCtaid { axis: Axis, span: Span },
144    /// `%globaltimer_hi`
145    GlobaltimerHi { span: Span },
146    /// `%nclusterid`
147    Nclusterid { span: Span },
148    /// `%smid`
149    Smid { span: Span },
150    /// `%cluster_ctarank` (optionally `.x/.y/.z`)
151    ClusterCtarank { axis: Axis, span: Span },
152    /// `%globaltimer_lo`
153    GlobaltimerLo { span: Span },
154    /// `%nctaid` (optionally `.x/.y/.z`)
155    Nctaid { axis: Axis, span: Span },
156    /// `%tid` (optionally `.x/.y/.z`)
157    Tid { axis: Axis, span: Span },
158    /// `%cluster_nctaid` (optionally `.x/.y/.z`)
159    ClusterNctaid { axis: Axis, span: Span },
160    /// `%gridid`
161    Gridid { span: Span },
162    /// `%nsmid`
163    Nsmid { span: Span },
164    /// `%total_smem_size`
165    TotalSmemSize { span: Span },
166    /// `%cluster_nctarank` (optionally `.x/.y/.z`)
167    ClusterNctarank { axis: Axis, span: Span },
168    /// `%is_explicit_cluster`
169    IsExplicitCluster { span: Span },
170    /// `%ntid` (optionally `.x/.y/.z`)
171    Ntid { axis: Axis, span: Span },
172    /// `%warpid`
173    Warpid { span: Span },
174    /// `%clusterid`
175    Clusterid { span: Span },
176    /// `%laneid`
177    Laneid { span: Span },
178    /// `%nwarpid`
179    Nwarpid { span: Span },
180    /// `%WARPSZ`
181    WARPSZ { span: Span },
182    /// `%ctaid` (optionally `.x/.y/.z`)
183    Ctaid { axis: Axis, span: Span },
184    /// `%lanemask_eq`
185    LanemaskEq { span: Span },
186    /// `%pm0`, ..., `%pm7`
187    Pm { index: u8, span: Span },
188    /// `%pm0_64`, ..., `%pm7_64`
189    Pm64 { index: u8, span: Span },
190    /// `%current_graph_exec`
191    CurrentGraphExec { span: Span },
192    /// `%lanemask_ge`
193    LanemaskGe { span: Span },
194    /// `%reserved_smem_offset_0`, `%reserved_smem_offset_1`
195    ReservedSmemOffset { index: u8, span: Span },
196}
197
198/* --------------------------------------------------- */
199/* ------------------- Operands ---------------------- */
200/* --------------------------------------------------- */
201
202/// Texture handler with 2 operands, e.g. [%r1, %r2]
203#[derive(Debug, Clone, PartialEq, Eq, Spanned)]
204pub struct TexHandler2 {
205    pub operands: [GeneralOperand; 2],
206    pub span: Span,
207}
208
209/// Texture handler with optional sampler operand, e.g. `[tex, coords]` or `[tex, sampler, coords]`
210#[derive(Debug, Clone, PartialEq, Eq, Spanned)]
211pub struct TexHandler3Optional {
212    pub handle: GeneralOperand,
213    pub sampler: Option<GeneralOperand>,
214    pub coords: GeneralOperand,
215    pub span: Span,
216}
217
218/// Texture handler with optional sampler operand, e.g. `[tex, coords]` or `[tex, sampler, coords]`
219#[derive(Debug, Clone, PartialEq, Eq, Spanned)]
220pub struct TexHandler3 {
221    pub handle: GeneralOperand,
222    pub sampler: GeneralOperand,
223    pub coords: GeneralOperand,
224    pub span: Span,
225}
226
227#[derive(Debug, Clone, PartialEq, Eq, Spanned)]
228pub enum GeneralOperand {
229    Vec { operand: VectorOperand, span: Span },
230    Single { operand: Operand, span: Span },
231}
232
233#[derive(Debug, Clone, PartialEq, Eq, Spanned)]
234pub enum VectorOperand {
235    /// {%r1}
236    Vector1 { operand: Operand, span: Span },
237    /// {%r1, %r2}
238    Vector2 { operands: [Operand; 2], span: Span },
239    /// {%r1, %r2, %r3}
240    Vector3 { operands: [Operand; 3], span: Span },
241    /// {%r1, %r2, %r3, %r4}
242    Vector4 { operands: [Operand; 4], span: Span },
243    /// {%r1, %r2, %r3, %r4, %r5, %r6, %r7, %r8}
244    Vector8 { operands: [Operand; 8], span: Span },
245}
246
247#[derive(Debug, Clone, PartialEq, Eq, Spanned)]
248pub enum Operand {
249    /// %r1
250    Register {
251        operand: RegisterOperand,
252        span: Span,
253    },
254    /// 0xffff
255    Immediate { operand: Immediate, span: Span },
256    /// foo
257    Symbol { name: String, span: Span },
258    /// foo + 4 (symbol + immediate offset)
259    SymbolOffset {
260        symbol: String,
261        offset: Immediate,
262        span: Span,
263    },
264}
265
266/// Register operand starting with % (e.g., `%r1`).
267#[derive(Debug, Clone, PartialEq, Eq, Spanned)]
268pub struct RegisterOperand {
269    pub name: String,
270    /// Optional component suffix (e.g., `.x`, `.y`, `.z`, `.w`)
271    pub component: Option<String>,
272    pub span: Span,
273}
274
275/// Predicate register names (e.g., `%p0`).
276#[derive(Debug, Clone, PartialEq, Eq, Spanned)]
277pub struct PredicateRegister {
278    pub name: String,
279    pub span: Span,
280}
281
282/// Representation of an address operand.
283#[derive(Debug, Clone, PartialEq, Eq, Spanned)]
284pub enum AddressOperand {
285    /// base[immIndex]
286    Array {
287        base: VariableSymbol,
288        index: Immediate,
289        span: Span,
290    },
291    /// Immediate address value, e.g., [0xffff], unsigned, 32-bit
292    ImmediateAddress { addr: Immediate, span: Span },
293    /// Offset address with optional displacement, e.g., [base + offset] and [base]
294    Offset {
295        base: AddressBase,
296        offset: Option<AddressOffset>,
297        span: Span,
298    },
299}
300
301/// Base location referenced by an address expression.
302#[derive(Debug, Clone, PartialEq, Eq, Spanned)]
303pub enum AddressBase {
304    Register {
305        operand: RegisterOperand,
306        span: Span,
307    },
308    Variable {
309        symbol: VariableSymbol,
310        span: Span,
311    },
312}
313
314/// Specific adjustment applied within a displacement term.
315#[derive(Debug, Clone, PartialEq, Eq, Spanned)]
316pub enum AddressOffset {
317    Register {
318        operand: RegisterOperand,
319        span: Span,
320    },
321    Immediate {
322        sign: Sign,
323        value: Immediate,
324        span: Span,
325    },
326}
327
328/* --------------------------------------------------- */
329/* -------------------- Immediate -------------------- */
330/* --------------------------------------------------- */
331
332/// Immediate value representing a constant literal in PTX assembly.
333///
334/// Supports all PTX constant formats:
335///
336/// **Integer literals** (with optional `U` suffix for unsigned):
337/// - Hexadecimal: `0x1234`, `0x1234U`, `0XABCD`
338/// - Octal: `0777`, `0777U`, `0123`
339/// - Binary: `0b1010`, `0b1010U`, `0B0011`
340/// - Decimal: `42`, `42U`, `0`, `0U`, `123`
341///
342/// **Floating-point literals**:
343/// - Decimal float: `3.14`, `2.5`
344/// - Scientific notation: `1.5e10`, `3.2E-5`, `1e3`
345/// - Hex float (single-precision, 32-bit): `0f3f800000` (8 hex digits after `0f`)
346/// - Hex float (double-precision, 64-bit): `0d3ff0000000000000` (16 hex digits after `0d`)
347#[derive(Debug, Clone, PartialEq, Eq, Spanned)]
348pub struct Immediate {
349    /// The literal value as a string (includes prefix, digits, and optional 'U' suffix)
350    pub value: String,
351    pub span: Span,
352}
353
354/* --------------------------------------------------- */
355/* -------------------- Symbols ---------------------- */
356/* --------------------------------------------------- */
357
358/// Function symbol
359#[derive(Debug, Clone, PartialEq, Eq, Spanned)]
360pub struct FunctionSymbol {
361    pub val: String,
362    pub span: Span,
363}
364
365/// Variable symbol
366#[derive(Debug, Clone, PartialEq, Eq, Spanned)]
367pub struct VariableSymbol {
368    pub val: String,
369    pub span: Span,
370}
371
372/// Label name (e.g., `L__label_1`).
373#[derive(Debug, Clone, PartialEq, Eq, Spanned)]
374pub struct Label {
375    pub val: String,
376    pub span: Span,
377}
378
379/* --------------------------------------------------- */
380/* -------------------- Predicate -------------------- */
381/* --------------------------------------------------- */
382
383/// Predicate guard for conditional instruction execution
384#[derive(Debug, Clone, PartialEq, Spanned)]
385pub struct Predicate {
386    pub negated: bool,
387    pub operand: Operand,
388    pub span: Span,
389}
390
391/* --------------------------------------------------- */
392/* -------------------- Instruction ------------------ */
393/* --------------------------------------------------- */
394
395/// Represents a complete instruction with optional label and predicate guard
396/// Format: [label:] [@{!}pred] instruction
397#[derive(Debug, Clone, PartialEq, Spanned)]
398pub struct Instruction {
399    pub predicate: Option<Predicate>,
400    pub inst: crate::r#type::instruction::Inst,
401    pub span: Span,
402}