cubecl_cpp/shared/
body.rs

1use super::{Dialect, Instruction, Variable, barrier::BarrierOps, pipeline::PipelineOps};
2use std::fmt::Display;
3
4/// A body is composed of a list of [instructions](Instruction).
5#[derive(Debug, Clone)]
6pub struct Body<D: Dialect> {
7    pub instructions: Vec<Instruction<D>>,
8    pub shared_memories: Vec<super::SharedMemory<D>>,
9    pub pipelines: Vec<PipelineOps<D>>,
10    pub barriers: Vec<BarrierOps<D>>,
11    pub const_arrays: Vec<super::ConstArray<D>>,
12    pub local_arrays: Vec<super::LocalArray<D>>,
13}
14
15impl<D: Dialect> Display for Body<D> {
16    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
17        D::compile_bindings_body(f, self)?;
18
19        for shared in &self.shared_memories {
20            D::compile_shared_memory_declaration(f, shared)?;
21        }
22
23        for pipeline in self.pipelines.iter() {
24            writeln!(f, "{pipeline}")?;
25        }
26        for barrier in self.barriers.iter() {
27            writeln!(f, "{barrier}")?;
28        }
29
30        for const_array in self.const_arrays.iter() {
31            f.write_fmt(format_args!(
32                "const {} arrays_{}[{}] = {{",
33                const_array.item, const_array.index, const_array.size
34            ))?;
35            let elem = const_array.item.elem;
36            for value in const_array.values.iter().copied() {
37                let value = match value {
38                    Variable::ConstantScalar(value, _) => Variable::ConstantScalar(value, elem),
39                    _ => unreachable!("Value is always constant"),
40                };
41                f.write_fmt(format_args!("{value},"))?;
42            }
43            f.write_str("};\n")?;
44        }
45
46        // Local arrays
47        for array in self.local_arrays.iter() {
48            write!(
49                f,
50                "{} l_arr_{}[{}];\n\n",
51                array.item, array.index, array.size
52            )?;
53        }
54
55        D::compile_wmma_local_variables(f)?;
56
57        for ops in self.instructions.iter() {
58            write!(f, "{ops}")?;
59        }
60
61        Ok(())
62    }
63}
cubecl_cpp/shared/body.rs

cubecl_cpp/shared/
body.rs
