cubecl_cpp/shared/
kernel.rs

1use crate::shared::ty::TypeExtCPP;
2
3use cubecl_core::ir::metadata::Info;
4use cubecl_runtime::kernel::{BufferIOAttr, Visibility};
5use pliron::context::Context;
6
7use core::fmt::{Display, Write};
8
9pub struct ComputeKernel {
10    pub shared_memory_size: usize,
11    pub buffers: Vec<Visibility>,
12    /// What the kernel does with each buffer binding, by buffer position —
13    /// the four-state answer the launch path's taint bookkeeping consumes.
14    /// Unlike [`buffers`](Self::buffers) this is never widened or collapsed.
15    pub io: Vec<BufferIOAttr>,
16    /// The emitted source, rendered once during `compile_ir` where emission errors can still
17    /// fail the compilation.
18    pub source: String,
19}
20
21impl Display for ComputeKernel {
22    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
23        f.write_str(&self.source)
24    }
25}
26
27pub fn type_definitions(f: &mut dyn Write, long: &str) -> std::fmt::Result {
28    writeln!(f, "typedef unsigned int uint32_t;")?;
29    writeln!(f, "typedef unsigned char uint8_t;")?;
30    writeln!(f, "typedef unsigned short uint16_t;")?;
31    writeln!(f, "typedef unsigned {long} int uint64_t;")?;
32
33    writeln!(f, "typedef signed char int8_t;")?;
34    writeln!(f, "typedef signed short int16_t;")?;
35    writeln!(f, "typedef signed int int32_t;")?;
36    writeln!(f, "typedef signed {long} int int64_t;")?;
37
38    Ok(())
39}
40
41/// Define a minimal version of C++'s `std::array` so we can match Rust semantics on arrays.
42pub fn define_array_polyfill(f: &mut dyn Write) -> core::fmt::Result {
43    writeln!(
44        f,
45        "
46template <typename T, size_t N>
47struct array {{
48    T data[N];
49    __device__ T& operator[](size_t i) {{ return data[i]; }}
50    __device__ const T& operator[](size_t i) const {{ return data[i]; }}
51}};\n"
52    )
53}
54
55pub fn define_tensormap_opaque(f: &mut dyn Write) -> core::fmt::Result {
56    f.write_str(
57        "
58typedef struct CUtensorMap_st {
59alignas(128) unsigned long long int opaque[16];
60} CUtensorMap;\n",
61    )
62}
63
64pub fn type_info_definition_sized(
65    f: &mut dyn Write,
66    ctx: &Context,
67    info: &Info,
68) -> std::fmt::Result {
69    let scalars = info
70        .scalars
71        .iter()
72        .map(|field| {
73            let ty = field.ty.to_type(ctx).to_cpp(ctx);
74            format!("{ty} scalars_{}[{}];", field.ty, field.padded_size(ctx))
75        })
76        .collect::<Vec<_>>()
77        .join("\n");
78    let static_meta = info
79        .sized_meta
80        .as_ref()
81        .map(|field| {
82            format!(
83                "{} static_meta[{}];",
84                field.ty.to_type(ctx).to_cpp(ctx),
85                field.padded_size(ctx)
86            )
87        })
88        .unwrap_or_default();
89    write!(
90        f,
91        "
92struct info_st {{
93    {scalars}{static_meta}
94}};
95"
96    )
97}
cubecl_cpp/shared/kernel.rs

cubecl_cpp/shared/
kernel.rs
