cubecl_cpp/cuda/ptx/
mod.rs

1use cubecl_core::{
2    self as cubecl,
3    frontend::barrier::Barrier,
4    intrinsic,
5    ir::{
6        AddressSpace,
7        interfaces::{TypeExt, aliasing::AliasingOp},
8        prelude::*,
9        types::{
10            PointerType,
11            barrier::{BarrierLevel, BarrierType},
12            cuda::TensorMapType,
13        },
14    },
15    prelude::*,
16};
17
18mod asm;
19mod mma;
20
21pub use asm::*;
22pub use mma::*;
23use pliron::{
24    builtin::types::{IntegerType, Signedness},
25    printable::Printable,
26    verify_err,
27};
28
29pub mod copy_async;
30pub mod tma_load_im2col;
31
32use crate::{cuda::cuda_op_with_out, shared::ty::TypeExtCPP};
33
34/// Equivalent of `__cvta_generic_to_shared`, required when a PTX instruction uses a specific
35/// `.shared` modifier. It should only be used to cast a pointer for use in that specific context,
36/// and using it without adding the `.shared` modifier will break. Using shared addresses in generic
37/// instructions will also break, which is why this isn't automatically applied in `InlinePtxOp`.
38#[cube_op(name = "cuda.generic_to_shared", format = "$0 ` : ` type($0)")]
39#[result_ty(fixed = IntegerType::get(ctx, 32, Signedness::Unsigned).to_handle())]
40pub struct GenericToSharedOp {
41    ptr: Value,
42}
43
44#[op_interface_impl]
45impl AliasingOp for GenericToSharedOp {
46    fn source_ptr(&self, ctx: &Context) -> Option<Value> {
47        Some(self.ptr(ctx))
48    }
49}
50
51cuda_op_with_out!(GenericToSharedOp, |op, ctx| {
52    let ptr = op.ptr(ctx).name(ctx);
53    format!("__cvta_generic_to_shared({ptr})")
54});
55
56#[cube]
57pub fn generic_to_shared<T: CubePrimitive>(ptr: *const T) -> u32 {
58    intrinsic!(|scope| {
59        let cvt = GenericToSharedOp::new(scope.ctx_mut(), unsafe { *ptr }.value(scope));
60        scope.register_with_result(&cvt).into()
61    })
62}
63
64/// Returns the shared address directly because it's always combined with `generic_to_shared` in
65/// practice, and it makes the types easier to deal with.
66#[cube_op(
67    name = "cuda.barrier_native_handle",
68    format = "$0 ` : ` type($0)",
69    verifier = "custom"
70)]
71#[result_ty(fixed = IntegerType::get(ctx, 32, Signedness::Unsigned).to_handle())]
72#[op_interfaces(OperandNOfType<0, PointerType>)]
73pub struct BarrierNativeHandleOp {
74    bar_ptr: Value,
75}
76
77#[op_interface_impl]
78impl AliasingOp for BarrierNativeHandleOp {
79    fn source_ptr(&self, ctx: &Context) -> Option<Value> {
80        Some(self.bar_ptr(ctx))
81    }
82}
83
84impl Verify for BarrierNativeHandleOp {
85    fn verify(&self, ctx: &Context) -> Result<()> {
86        let loc = self.loc(ctx);
87        let barrier_ty = self.bar_ptr(ctx).get_type(ctx).as_ptr(ctx);
88
89        if !barrier_ty.inner.deref(ctx).is::<BarrierType>() {
90            let expected = PointerType::get(
91                ctx,
92                BarrierType::get(ctx, BarrierLevel::Cube).into(),
93                AddressSpace::Shared,
94            )
95            .to_handle();
96            verify_err!(
97                loc,
98                OperandNOfTypeError::AllOperandsOfTypeVerifyErr(
99                    expected.disp(ctx).to_string(),
100                    self.bar_ptr(ctx).get_type(ctx).disp(ctx).to_string()
101                )
102            )?;
103        }
104        Ok(())
105    }
106}
107
108cuda_op_with_out!(BarrierNativeHandleOp, |op, ctx| {
109    let ptr = op.bar_ptr(ctx).name(ctx);
110    format!("__cvta_generic_to_shared(cuda::device::barrier_native_handle(*{ptr}))")
111});
112
113#[cube]
114pub fn barrier_native_handle(bar: &Barrier) -> u32 {
115    intrinsic!(|scope| {
116        let handle = BarrierNativeHandleOp::new(scope.ctx_mut(), bar.value(scope));
117        scope.register_with_result(&handle).into()
118    })
119}
120
121/// Tensor maps are weird because they're passed as a reference but consumed by pointer. Don't want
122/// to add a generic `ReferenceOp`, so for now just do a tensor-map specific reference. That way
123/// it doesn't need to deal with Metal address space nonsense. Returns the raw address as `u64`
124/// since it's always used for PTX anyways, and that coerces pointers to u64.
125#[cube_op(name = "cuda.tensor_map_addr", format = "$0 ` : ` type($0)")]
126#[result_ty(fixed = IntegerType::get(ctx, 64, Signedness::Unsigned).to_handle())]
127#[op_interfaces(OperandNOfType<0, TensorMapType>)]
128pub struct TensorMapAddrOp {
129    tensor_map: Value,
130}
131
132cuda_op_with_out!(TensorMapAddrOp, |op, ctx| {
133    let tensor_map = op.tensor_map(ctx).name(ctx);
134    let out_ty = op.get_result(ctx).get_type(ctx).to_cpp(ctx);
135    format!("reinterpret_cast<{out_ty}>(&{tensor_map})")
136});
137
138#[cube]
139pub fn tensor_map_address<T: CubePrimitive, K: TensorMapKind>(tensor_map: &TensorMap<T, K>) -> u64 {
140    intrinsic!(|scope| {
141        let addr = TensorMapAddrOp::new(scope.ctx_mut(), tensor_map.value(scope));
142        scope.register_with_result(&addr).into()
143    })
144}
cubecl_cpp/cuda/ptx/mod.rs

cubecl_cpp/cuda/ptx/
mod.rs
