Skip to main content

ptx_parser/type/
common.rs

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