cubecl_cpp/cuda/
dialect.rs

1use cubecl_core::ir::{
2    dialect::synchronization::{SyncAsyncProxyOp, SyncOp, SyncScope},
3    prelude::*,
4};
5
6use crate::{shared::signature::op_includes, target::Cuda};
7
8macro_rules! cuda_op {
9    ($ty: ty, $impl: expr) => {
10        #[pliron::derive::op_interface_impl]
11        impl $crate::shared::operation::OpToCPP<$crate::target::Cuda> for $ty {
12            fn to_cpp(&self, ctx: &pliron::context::Context) -> String {
13                $crate::shared::closure_inference_hack::<$ty, String>(self, ctx, $impl)
14            }
15        }
16    };
17}
18pub(super) use cuda_op;
19
20macro_rules! cuda_op_with_out {
21    ($ty: ty, $impl: expr) => {
22        #[pliron::derive::op_interface_impl]
23        impl $crate::shared::operation::OpToCPP<$crate::target::Cuda> for $ty {
24            fn to_cpp(&self, ctx: &pliron::context::Context) -> String {
25                use cubecl_core::ir::prelude::*;
26                use $crate::shared::CppValue;
27                let op = $crate::shared::closure_inference_hack::<$ty, String>(self, ctx, $impl);
28                let out = self.get_result(ctx).fmt_left(ctx);
29                format!("{out} = {op};\n")
30            }
31        }
32    };
33}
34pub(super) use cuda_op_with_out;
35
36macro_rules! ptx_with_out {
37    ($ty: ty, $ptx: expr, $pred: expr) => {
38        #[op_interface_impl]
39        impl $crate::shared::lowering::LowerOp<$crate::target::Cuda> for $ty {
40            fn should_lower(&self, ctx: &pliron::context::Context) -> bool {
41                $crate::shared::closure_inference_hack::<$ty, bool>(self, ctx, $pred)
42            }
43            fn lower(&self, scope: &cubecl_core::ir::Scope) -> Vec<pliron::value::Value> {
44                use cubecl_core::ir::dialect::base::OperationPtrExt;
45                use pliron::{op::Op, r#type::Typed};
46                let ctx = scope.ctx_mut();
47                let ptx = $crate::shared::closure_inference_hack::<$ty, String>(self, ctx, $ptx);
48                let op = $crate::cuda::ptx::InlinePtxOp::new(
49                    ctx,
50                    Some(self.get_result(ctx).get_type(ctx)),
51                    ptx,
52                    self.get_operation().operands(ctx),
53                );
54                scope.register(&op);
55                vec![op.result(ctx).unwrap()]
56            }
57        }
58    };
59    ($ty: ty, $ptx: expr) => {
60        ptx_with_out!($ty, $ptx, |_, _| true);
61    };
62}
63pub(super) use ptx_with_out;
64
65op_includes!(Cuda, [SyncAsyncProxyOp] => "cuda/barrier");
66
67cuda_op!(SyncOp, |op, ctx| {
68    match op.scope(ctx).0 {
69        SyncScope::Plane => "__syncwarp();\n",
70        SyncScope::Cube => "__syncthreads();\n",
71        // `__syncthreads` orders this block's memory for this block. The fence is what carries
72        // the writes out to device scope, so a cube that synchronizes later reads them.
73        SyncScope::Device => "__threadfence();\n__syncthreads();\n",
74        SyncScope::Unit => "",
75    }
76    .into()
77});
78
79cuda_op!(SyncAsyncProxyOp, |_, _| {
80    "cuda::device::experimental::fence_proxy_async_shared_cta();".into()
81});
82
83pub(crate) const COMPLEX_HELPERS: &str = r#"
84__device__ __host__ inline cuFloatComplex operator+(cuFloatComplex a, cuFloatComplex b) { return cuCaddf(a, b); }
85__device__ __host__ inline cuFloatComplex operator-(cuFloatComplex a, cuFloatComplex b) { return cuCsubf(a, b); }
86__device__ __host__ inline cuFloatComplex operator*(cuFloatComplex a, cuFloatComplex b) { return cuCmulf(a, b); }
87__device__ __host__ inline cuFloatComplex operator/(cuFloatComplex a, cuFloatComplex b) { return cuCdivf(a, b); }
88__device__ __host__ inline cuFloatComplex operator-(cuFloatComplex a) { return make_cuFloatComplex(-cuCrealf(a), -cuCimagf(a)); }
89__device__ __host__ inline bool operator==(cuFloatComplex a, cuFloatComplex b) { return cuCrealf(a)==cuCrealf(b) && cuCimagf(a)==cuCimagf(b); }
90__device__ __host__ inline bool operator!=(cuFloatComplex a, cuFloatComplex b) { return !(a==b); }
91__device__ __host__ inline cuDoubleComplex operator+(cuDoubleComplex a, cuDoubleComplex b) { return cuCadd(a, b); }
92__device__ __host__ inline cuDoubleComplex operator-(cuDoubleComplex a, cuDoubleComplex b) { return cuCsub(a, b); }
93__device__ __host__ inline cuDoubleComplex operator*(cuDoubleComplex a, cuDoubleComplex b) { return cuCmul(a, b); }
94__device__ __host__ inline cuDoubleComplex operator/(cuDoubleComplex a, cuDoubleComplex b) { return cuCdiv(a, b); }
95__device__ __host__ inline cuDoubleComplex operator-(cuDoubleComplex a) { return make_cuDoubleComplex(-cuCreal(a), -cuCimag(a)); }
96__device__ __host__ inline bool operator==(cuDoubleComplex a, cuDoubleComplex b) { return cuCreal(a)==cuCreal(b) && cuCimag(a)==cuCimag(b); }
97__device__ __host__ inline bool operator!=(cuDoubleComplex a, cuDoubleComplex b) { return !(a==b); }
98__device__ __host__ inline float abs(cuFloatComplex a) { return hypotf(cuCrealf(a), cuCimagf(a)); }
99__device__ __host__ inline double abs(cuDoubleComplex a) { return hypot(cuCreal(a), cuCimag(a)); }
100__device__ __host__ inline cuFloatComplex exp(cuFloatComplex a) { const float x=cuCrealf(a), y=cuCimagf(a), ex=expf(x); return make_cuFloatComplex(ex*cosf(y), ex*sinf(y)); }
101__device__ __host__ inline cuDoubleComplex exp(cuDoubleComplex a) { const double x=cuCreal(a), y=cuCimag(a), ex=exp(x); return make_cuDoubleComplex(ex*cos(y), ex*sin(y)); }
102__device__ __host__ inline cuFloatComplex log(cuFloatComplex a) { const float x=cuCrealf(a), y=cuCimagf(a); return make_cuFloatComplex(logf(hypotf(x,y)), atan2f(y,x)); }
103__device__ __host__ inline cuDoubleComplex log(cuDoubleComplex a) { const double x=cuCreal(a), y=cuCimag(a); return make_cuDoubleComplex(log(hypot(x,y)), atan2(y,x)); }
104__device__ __host__ inline cuFloatComplex sin(cuFloatComplex a) { const float x=cuCrealf(a), y=cuCimagf(a); return make_cuFloatComplex(sinf(x)*coshf(y), cosf(x)*sinhf(y)); }
105__device__ __host__ inline cuDoubleComplex sin(cuDoubleComplex a) { const double x=cuCreal(a), y=cuCimag(a); return make_cuDoubleComplex(sin(x)*cosh(y), cos(x)*sinh(y)); }
106__device__ __host__ inline cuFloatComplex cos(cuFloatComplex a) { const float x=cuCrealf(a), y=cuCimagf(a); return make_cuFloatComplex(cosf(x)*coshf(y), -sinf(x)*sinhf(y)); }
107__device__ __host__ inline cuDoubleComplex cos(cuDoubleComplex a) { const double x=cuCreal(a), y=cuCimag(a); return make_cuDoubleComplex(cos(x)*cosh(y), -sin(x)*sinh(y)); }
108__device__ __host__ inline cuFloatComplex sqrt(cuFloatComplex a) { const float x=cuCrealf(a), y=cuCimagf(a), r=hypotf(x,y); if(x>=0.0f){ const float re=sqrtf(0.5f*(r+x)); return make_cuFloatComplex(re,re==0.0f?0.0f:y/(2.0f*re)); } const float im=copysignf(sqrtf(0.5f*(r-x)),y); return make_cuFloatComplex(im==0.0f?0.0f:y/(2.0f*im),im); }
109__device__ __host__ inline cuDoubleComplex sqrt(cuDoubleComplex a) { const double x=cuCreal(a), y=cuCimag(a), r=hypot(x,y); if(x>=0.0){ const double re=sqrt(0.5*(r+x)); return make_cuDoubleComplex(re,re==0.0?0.0:y/(2.0*re)); } const double im=copysign(sqrt(0.5*(r-x)),y); return make_cuDoubleComplex(im==0.0?0.0:y/(2.0*im),im); }
110__device__ __host__ inline cuFloatComplex tanh(cuFloatComplex a) { const float x2=2.0f*cuCrealf(a), y2=2.0f*cuCimagf(a), d=coshf(x2)+cosf(y2); return make_cuFloatComplex(sinhf(x2)/d,sinf(y2)/d); }
111__device__ __host__ inline cuDoubleComplex tanh(cuDoubleComplex a) { const double x2=2.0*cuCreal(a), y2=2.0*cuCimag(a), d=cosh(x2)+cos(y2); return make_cuDoubleComplex(sinh(x2)/d,sin(y2)/d); }
112__device__ __host__ inline cuFloatComplex pow(cuFloatComplex a, cuFloatComplex b) { return exp(b*log(a)); }
113__device__ __host__ inline cuDoubleComplex pow(cuDoubleComplex a, cuDoubleComplex b) { return exp(b*log(a)); }
114"#;
cubecl_cpp/cuda/dialect.rs

cubecl_cpp/cuda/
dialect.rs
