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    /// foo + 4 (symbol + immediate offset)
260    SymbolOffset {
261        symbol: String,
262        offset: Immediate,
263        span: Span,
264    },
265}
266
267/// Register operand starting with % (e.g., `%r1`).
268#[derive(Debug, Clone, PartialEq, Eq, Spanned, Serialize)]
269pub struct RegisterOperand {
270    pub name: String,
271    /// Optional component suffix (e.g., `.x`, `.y`, `.z`, `.w`)
272    pub component: Option<String>,
273    pub span: Span,
274}
275
276/// Predicate register names (e.g., `%p0`).
277#[derive(Debug, Clone, PartialEq, Eq, Spanned, Serialize)]
278pub struct PredicateRegister {
279    pub name: String,
280    pub span: Span,
281}
282
283/// Representation of an address operand.
284#[derive(Debug, Clone, PartialEq, Eq, Spanned, Serialize)]
285pub enum AddressOperand {
286    /// base[immIndex]
287    Array {
288        base: VariableSymbol,
289        index: Immediate,
290        span: Span,
291    },
292    /// Immediate address value, e.g., [0xffff], unsigned, 32-bit
293    ImmediateAddress { addr: Immediate, span: Span },
294    /// Offset address with optional displacement, e.g., [base + offset] and [base]
295    Offset {
296        base: AddressBase,
297        offset: Option<AddressOffset>,
298        span: Span,
299    },
300}
301
302/// Base location referenced by an address expression.
303#[derive(Debug, Clone, PartialEq, Eq, Spanned, Serialize)]
304pub enum AddressBase {
305    Register {
306        operand: RegisterOperand,
307        span: Span,
308    },
309    Variable {
310        symbol: VariableSymbol,
311        span: Span,
312    },
313}
314
315/// Specific adjustment applied within a displacement term.
316#[derive(Debug, Clone, PartialEq, Eq, Spanned, Serialize)]
317pub enum AddressOffset {
318    Register {
319        operand: RegisterOperand,
320        span: Span,
321    },
322    Immediate {
323        sign: Sign,
324        value: Immediate,
325        span: Span,
326    },
327}
328
329/* --------------------------------------------------- */
330/* -------------------- Immediate -------------------- */
331/* --------------------------------------------------- */
332
333/// Immediate value representing a constant literal in PTX assembly.
334///
335/// Supports all PTX constant formats:
336///
337/// **Integer literals** (with optional `U` suffix for unsigned):
338/// - Hexadecimal: `0x1234`, `0x1234U`, `0XABCD`
339/// - Octal: `0777`, `0777U`, `0123`
340/// - Binary: `0b1010`, `0b1010U`, `0B0011`
341/// - Decimal: `42`, `42U`, `0`, `0U`, `123`
342///
343/// **Floating-point literals**:
344/// - Decimal float: `3.14`, `2.5`
345/// - Scientific notation: `1.5e10`, `3.2E-5`, `1e3`
346/// - Hex float (single-precision, 32-bit): `0f3f800000` (8 hex digits after `0f`)
347/// - Hex float (double-precision, 64-bit): `0d3ff0000000000000` (16 hex digits after `0d`)
348#[derive(Debug, Clone, PartialEq, Eq, Spanned, Serialize)]
349pub struct Immediate {
350    /// The literal value as a string (includes prefix, digits, and optional 'U' suffix)
351    pub value: String,
352    pub span: Span,
353}
354
355/* --------------------------------------------------- */
356/* -------------------- Symbols ---------------------- */
357/* --------------------------------------------------- */
358
359/// Function symbol
360#[derive(Debug, Clone, PartialEq, Eq, Spanned, Serialize)]
361pub struct FunctionSymbol {
362    pub val: String,
363    pub span: Span,
364}
365
366/// Variable symbol
367#[derive(Debug, Clone, PartialEq, Eq, Spanned, Serialize)]
368pub struct VariableSymbol {
369    pub val: String,
370    pub span: Span,
371}
372
373/// Label name (e.g., `L__label_1`).
374#[derive(Debug, Clone, PartialEq, Eq, Spanned, Serialize)]
375pub struct Label {
376    pub val: String,
377    pub span: Span,
378}
379
380/* --------------------------------------------------- */
381/* -------------------- Predicate -------------------- */
382/* --------------------------------------------------- */
383
384/// Predicate guard for conditional instruction execution
385#[derive(Debug, Clone, PartialEq, Spanned, Serialize)]
386pub struct Predicate {
387    pub negated: bool,
388    pub operand: Operand,
389    pub span: Span,
390}
391
392/* --------------------------------------------------- */
393/* -------------------- Instruction ------------------ */
394/* --------------------------------------------------- */
395
396/// Represents a complete instruction with optional label and predicate guard
397/// Format: [label:] [@{!}pred] instruction
398#[derive(Debug, Clone, PartialEq, Spanned, Serialize)]
399pub struct Instruction {
400    pub predicate: Option<Predicate>,
401    pub inst: crate::r#type::instruction::Inst,
402    pub span: Span,
403}