cubecl_cpp/cuda/ptx/
asm.rs

1use core::{cell::Ref, fmt::Display};
2
3use cubecl_core::{
4    frontend::InputKind,
5    ir::{
6        AddressType, Scope,
7        dialect::{InlineAsmOp, InputSpecsAttr, MemoryClobbers, MemoryClobbersAttr},
8        interfaces::{MemoryEffect, MemoryEffects, TypedExt},
9        prelude::*,
10        types::VectorType,
11    },
12};
13use itertools::Itertools;
14use pliron::{
15    attribute::AttrObj,
16    builtin::attributes::{StringAttr, UnitAttr, VecAttr},
17    opts::dce::SideEffects,
18    printable::Printable,
19};
20
21use crate::{
22    cuda::cuda_op,
23    shared::{CppValue, lowering::LowerOp, scoped_block, ty::TypeExtCPP},
24    target::Cuda,
25};
26
27/// Inline PTX. Restricted to zero or one results because C++ semantics are too hard otherwise.
28/// Note that this does not *directly* map to PTX, because it actually destructures vectors to PTX
29/// vector expressions automatically. This means more than one register can be returned if it's part
30/// of a vector expression. To denote the difference, the syntax uses `$0`, `$1` etc for Pliron
31/// values, as opposed to the usual `%0`, `%1` etc for the PTX registers.
32#[pliron_op(
33    name = "cuda.inline_ptx",
34    format = "opt_attr($cuda_inline_ptx_volatile, $UnitAttr, label($volatile))
35    attr($cuda_inline_ptx_ptx, $StringAttr) ` : ` types(CharSpace(`,`)) ` : ` operands(CharSpace(`,`))
36    opt_attr($cuda_inline_ptx_clobbers, $VecAttr)
37    opt_attr($cuda_inline_ptx_in_spec, $InputSpecsAttr, label($in_spec))
38    opt_attr($cuda_inline_ptx_memory_clobbers, $MemoryClobbersAttr, label($memory_clobbers))",
39    attributes = (
40        cuda_inline_ptx_ptx: StringAttr,
41        cuda_inline_ptx_volatile: UnitAttr,
42        cuda_inline_ptx_clobbers: VecAttr,
43        cuda_inline_ptx_memory_clobbers: MemoryClobbersAttr,
44        cuda_inline_ptx_in_spec: InputSpecsAttr,
45    ),
46    verifier = "succ"
47)]
48pub struct InlinePtxOp;
49
50impl InlinePtxOp {
51    pub fn new(
52        ctx: &mut Context,
53        result_ty: Option<TypeHandle>,
54        ptx: impl Display,
55        inputs: Vec<Value>,
56    ) -> Self {
57        let op = Operation::new(
58            ctx,
59            Self::get_concrete_op_info(),
60            result_ty.into_iter().collect(),
61            inputs,
62            vec![],
63            0,
64        );
65        let op = Self { op };
66        op.set_attr_cuda_inline_ptx_ptx(ctx, ptx.to_string().into());
67        op.set_attr_cuda_inline_ptx_memory_clobbers(ctx, MemoryClobbers::Nomem.into());
68        op
69    }
70
71    pub fn new_volatile(
72        ctx: &mut Context,
73        result_ty: Option<TypeHandle>,
74        ptx: impl Display,
75        inputs: Vec<Value>,
76    ) -> Self {
77        let op = Self::new(ctx, result_ty, ptx, inputs);
78        op.set_attr_cuda_inline_ptx_volatile(ctx, UnitAttr::new());
79        op.set_attr_cuda_inline_ptx_memory_clobbers(ctx, MemoryClobbers::Nomem.into());
80        op
81    }
82
83    pub fn set_clobbers(&self, ctx: &Context, clobbers: Vec<String>) {
84        if !clobbers.is_empty() {
85            let clobbers = clobbers
86                .into_iter()
87                .map(StringAttr::new)
88                .map(|attr| -> AttrObj { Box::new(attr) });
89            self.set_attr_cuda_inline_ptx_clobbers(ctx, VecAttr(clobbers.collect()));
90        }
91    }
92
93    pub fn clobbers(&self, ctx: &Context) -> Vec<String> {
94        let clobbers = self
95            .get_attr_cuda_inline_ptx_clobbers(ctx)
96            .map(|it| it.0.clone())
97            .unwrap_or_default();
98        clobbers
99            .into_iter()
100            .map(|it| (*it.downcast::<StringAttr>().unwrap()).into())
101            .collect()
102    }
103
104    pub fn raw_ptx<'a>(&self, ctx: &'a Context) -> Ref<'a, str> {
105        Ref::map(self.get_attr_cuda_inline_ptx_ptx(ctx).unwrap(), |it| {
106            it.as_str()
107        })
108    }
109
110    pub fn is_volatile(&self, ctx: &Context) -> bool {
111        self.get_attr_cuda_inline_ptx_volatile(ctx).is_some()
112    }
113
114    pub fn inputs(&self, ctx: &Context) -> Vec<Value> {
115        self.get_operation().deref(ctx).operands().collect()
116    }
117
118    pub fn result(&self, ctx: &Context) -> Option<Value> {
119        self.get_operation().deref(ctx).results().next()
120    }
121}
122
123#[op_interface_impl]
124impl SideEffects for InlinePtxOp {
125    fn has_side_effects(&self, ctx: &Context) -> bool {
126        self.get_attr_cuda_inline_ptx_volatile(ctx).is_some()
127    }
128}
129
130#[op_interface_impl]
131impl MemoryEffects for InlinePtxOp {
132    fn memory_effects(&self, ctx: &Context) -> Vec<MemoryEffect> {
133        match &self
134            .get_attr_cuda_inline_ptx_memory_clobbers(ctx)
135            .unwrap()
136            .0
137        {
138            MemoryClobbers::Nomem => vec![],
139            MemoryClobbers::Readonly => vec![MemoryEffect::ReadAll],
140            MemoryClobbers::Explicit {
141                reads_spaces,
142                writes_spaces,
143            } => {
144                let mut out = vec![];
145                for space in reads_spaces.0.iter() {
146                    out.push(MemoryEffect::ReadAllInSpace(*space));
147                }
148                for space in writes_spaces.0.iter() {
149                    out.push(MemoryEffect::WriteAllInSpace(*space));
150                }
151                let specs = self.get_attr_cuda_inline_ptx_in_spec(ctx).unwrap();
152                for (value, spec) in self.inputs(ctx).into_iter().zip(specs.0.iter()) {
153                    match spec.kind {
154                        InputKind::MemIn => {
155                            out.push(MemoryEffect::Read(value));
156                        }
157                        InputKind::MemOut => {
158                            out.push(MemoryEffect::Write(value));
159                        }
160                        InputKind::MemInout => {
161                            out.push(MemoryEffect::Read(value));
162                            out.push(MemoryEffect::Write(value));
163                        }
164                        InputKind::In => {}
165                    }
166                }
167                out
168            }
169            MemoryClobbers::ReadWrite => vec![MemoryEffect::ReadAll, MemoryEffect::WriteAll],
170        }
171    }
172}
173
174#[macro_export]
175macro_rules! ptx_block {
176    ($($lines: expr)*) => {{
177        let mut out = String::from("{\n\t");
178        $(
179            out.push_str(&$lines);
180            out.push_str("\n\t");
181        )*
182        out.push_str("}");
183        out
184    }};
185}
186
187cuda_op!(InlinePtxOp, |op, ctx| {
188    let mut ptx = op.raw_ptx(ctx).to_owned();
189    let result = op.result(ctx);
190    let inputs = op.inputs(ctx);
191
192    let mut ptx_idx = 0;
193    let mut plir_idx = 0;
194
195    if let Some(result) = result {
196        ptx = insert_placeholders(ctx, &ptx, result.get_type(ctx), plir_idx, &mut ptx_idx);
197        plir_idx += 1;
198    }
199
200    for input in inputs.iter() {
201        ptx = insert_placeholders(ctx, &ptx, input.get_type(ctx), plir_idx, &mut ptx_idx);
202        plir_idx += 1;
203    }
204
205    let out_regs = result
206        .iter()
207        .flat_map(|val| flatten_result(ctx, val.get_type(ctx)))
208        .join(", ");
209    let input_regs = inputs
210        .iter()
211        .flat_map(|val| flatten_operand(ctx, *val))
212        .join(", ");
213
214    let volatile = if op.is_volatile(ctx) { "volatile" } else { "" };
215    let clobbers = if let Some(clobbers) = op.get_attr_cuda_inline_ptx_clobbers(ctx) {
216        let names = clobbers.0.iter();
217        let names = names
218            .map(|it| it.downcast_ref::<StringAttr>().unwrap().as_str())
219            .map(|name| format!(r#""{name}""#))
220            .join(", ");
221        format!(": {names}")
222    } else {
223        String::new()
224    };
225
226    let asm = format!("asm {volatile}({ptx:?} : {out_regs} : {input_regs} {clobbers});",);
227
228    if let Some(result) = result {
229        let block = scoped_block!(
230            format!("{} result;", result.get_type(ctx).to_cpp(ctx))
231            asm
232            "return result;"
233        );
234        format!("{} = {block};", result.fmt_left(ctx))
235    } else {
236        asm
237    }
238});
239
240fn flatten_result(ctx: &Context, ty: TypeHandle) -> Vec<String> {
241    if ty.is_vector(ctx) {
242        let vec = ty.vector_size(ctx);
243        let constraint = infer_constraint_letter(ctx, ty.scalar_ty(ctx));
244        (0..vec)
245            .map(|i| format!(r#""={constraint}"(result.i_{i})"#))
246            .collect()
247    } else {
248        let constraint = infer_constraint_letter(ctx, ty.get_type(ctx));
249        vec![format!(r#""={constraint}"(result)"#)]
250    }
251}
252
253fn flatten_operand(ctx: &Context, val: Value) -> Vec<String> {
254    if val.get_type(ctx).deref(ctx).is::<VectorType>() {
255        let vec = val.vector_size(ctx);
256        let constraint = infer_constraint_letter(ctx, val.scalar_ty(ctx));
257        (0..vec)
258            .map(|i| format!(r#""{constraint}"({}.i_{i})"#, val.name(ctx)))
259            .collect()
260    } else {
261        let constraint = infer_constraint_letter(ctx, val.get_type(ctx));
262        vec![format!(r#""{constraint}"({})"#, val.name(ctx))]
263    }
264}
265
266fn insert_placeholders(
267    ctx: &Context,
268    ptx: &str,
269    ty: TypeHandle,
270    plir_idx: usize,
271    ptx_idx: &mut usize,
272) -> String {
273    let pat = format!("${plir_idx}");
274    if !ptx.contains(&pat) {
275        panic!("Tried substituting argument {pat} in PTX string {ptx:?}, but it wasn't found.")
276    }
277    let substitute = if ty.deref(ctx).is::<VectorType>() {
278        let vec = ty.vector_size(ctx);
279        let mut placeholders = (0..vec).map(|i| format!("%{}", *ptx_idx + i));
280        let substitute = format!("{{{}}}", placeholders.join(", "));
281        *ptx_idx += vec;
282        substitute
283    } else {
284        let placeholder = format!("%{ptx_idx}");
285        *ptx_idx += 1;
286        placeholder
287    };
288    ptx.replace(&pat, &substitute)
289}
290
291fn infer_constraint_letter(ctx: &Context, ty: TypeHandle) -> char {
292    if ty.is_bool(ctx) {
293        'b'
294    } else if ty.is_int_of_width(ctx, 16) {
295        'h'
296    } else if ty.is_int_of_width(ctx, 32) {
297        'r'
298    } else if ty.is_int_of_width(ctx, 64) {
299        'l'
300    } else if ty.is_index(ctx) {
301        match ctx.address_type() {
302            AddressType::U32 => 'r',
303            AddressType::U64 => 'l',
304        }
305    } else if ty.is_float32(ctx) {
306        'f'
307    } else if ty.is_float64(ctx) {
308        'd'
309    } else if ty.is_ptr(ctx) {
310        'l'
311    } else {
312        panic!(
313         "The register type could not be deduced from Pliron type. The type {} is not supported. 
314Supported types are: bool, i16, i32, i64, f32, f64, pointers.
315Please use cube.reinterpret_cast if you have different type.
316See the constraints from here: https://docs.nvidia.com/cuda/inline-ptx-assembly/index.html#constraints",
317        ty.disp(ctx));
318    }
319}
320
321#[op_interface_impl]
322impl LowerOp<Cuda> for InlineAsmOp {
323    fn lower(&self, scope: &Scope) -> Vec<Value> {
324        let ctx = scope.ctx_mut();
325        let ptx = self.asm(ctx).as_str().to_owned();
326        let inputs = self.inputs(ctx);
327        let results = self
328            .get_operation()
329            .opt_result(ctx)
330            .map(|res| res.get_type(ctx));
331        let inline_ptx = InlinePtxOp::new(ctx, results, ptx, inputs);
332        let mem_clobbers = self.memory_clobbers(ctx).clone();
333        if !self.pure(ctx) {
334            inline_ptx.set_attr_cuda_inline_ptx_volatile(ctx, UnitAttr::new());
335        }
336        match &mem_clobbers {
337            MemoryClobbers::Nomem => {}
338            MemoryClobbers::Readonly
339            | MemoryClobbers::Explicit { .. }
340            | MemoryClobbers::ReadWrite => {
341                inline_ptx.set_attr_cuda_inline_ptx_clobbers(
342                    ctx,
343                    VecAttr(vec![Box::new(StringAttr::new("memory".into()))]),
344                );
345            }
346        }
347        inline_ptx.set_attr_cuda_inline_ptx_memory_clobbers(ctx, mem_clobbers.into());
348        inline_ptx.set_attr_cuda_inline_ptx_in_spec(ctx, self.in_specs(ctx).into());
349
350        inline_ptx
351            .get_operation()
352            .insert_before(ctx, self.get_operation());
353        inline_ptx.get_operation().results(ctx)
354    }
355}
cubecl_cpp/cuda/ptx/asm.rs

cubecl_cpp/cuda/ptx/
asm.rs
