cubecl_ir/dialect/
spirv.rs

1use cubecl_macros_internal::cube_op;
2
3use crate::{
4    CanMaterialize, Pure,
5    attributes::IndexAttr,
6    prelude::*,
7    types::spirv::{ClampMode, TensorLayoutType},
8};
9
10#[pliron_op(
11    name = "matrix_spirv.load_tensor",
12    operands = (buffer, layout: TensorLayoutType),
13    format,
14    verifier = "succ"
15)]
16#[op_interfaces(OneResultInterface)]
17#[op_traits(CanMaterialize)]
18pub struct LoadTensorOp;
19
20impl LoadTensorOp {
21    pub fn new(
22        ctx: &mut Context,
23        out_ty: TypeHandle,
24        buffer: Value,
25        layout: Value,
26        view: Option<Value>,
27    ) -> Self {
28        let mut operands = vec![buffer, layout];
29        operands.extend(view);
30        let op = Operation::new(
31            ctx,
32            Self::get_concrete_op_info(),
33            vec![out_ty],
34            operands,
35            vec![],
36            0,
37        );
38        Self { op }
39    }
40
41    pub fn buffer(&self, ctx: &Context) -> Value {
42        self.get_operation().deref(ctx).get_operand(0)
43    }
44
45    pub fn layout(&self, ctx: &Context) -> Value {
46        self.get_operation().deref(ctx).get_operand(1)
47    }
48
49    pub fn view(&self, ctx: &Context) -> Option<Value> {
50        let op = self.get_operation().deref(ctx);
51        if op.get_num_operands() > 2 {
52            Some(op.get_operand(2))
53        } else {
54            None
55        }
56    }
57}
58
59#[pliron_op(name = "matrix_spirv.store_tensor", format, verifier = "succ")]
60#[op_traits(CanMaterialize)]
61pub struct StoreTensorOp;
62
63impl StoreTensorOp {
64    pub fn new(
65        ctx: &mut Context,
66        buffer: Value,
67        matrix: Value,
68        layout: Value,
69        view: Option<Value>,
70    ) -> Self {
71        let mut operands = vec![buffer, matrix, layout];
72        operands.extend(view);
73        let op = Operation::new(
74            ctx,
75            Self::get_concrete_op_info(),
76            vec![],
77            operands,
78            vec![],
79            0,
80        );
81        Self { op }
82    }
83
84    pub fn buffer(&self, ctx: &Context) -> Value {
85        self.get_operation().deref(ctx).get_operand(0)
86    }
87
88    pub fn matrix(&self, ctx: &Context) -> Value {
89        self.get_operation().deref(ctx).get_operand(1)
90    }
91
92    pub fn layout(&self, ctx: &Context) -> Value {
93        self.get_operation().deref(ctx).get_operand(2)
94    }
95
96    pub fn view(&self, ctx: &Context) -> Option<Value> {
97        let op = self.get_operation().deref(ctx);
98        if op.get_num_operands() > 3 {
99            Some(op.get_operand(3))
100        } else {
101            None
102        }
103    }
104}
105
106#[pliron_op(name = "spirv.create_layout", format, attributes = (spirv_create_layout_rank: IndexAttr), verifier = "succ")]
107#[op_interfaces(NResultsInterface<1>, OneResultInterface)]
108#[op_traits(CanMaterialize, Pure)]
109pub struct CreateLayoutOp;
110
111impl CreateLayoutOp {
112    pub fn new(
113        ctx: &mut Context,
114        shape: Vec<Value>,
115        strides: Option<Vec<Value>>,
116        clamp_mode: ClampMode,
117    ) -> Self {
118        let rank = shape.len();
119        let out_ty = TensorLayoutType::get(ctx, rank, clamp_mode);
120        let mut operands = shape;
121        operands.extend(strides.into_iter().flatten());
122        let op = Self {
123            op: Operation::new(
124                ctx,
125                Self::get_concrete_op_info(),
126                vec![out_ty.into()],
127                operands,
128                vec![],
129                0,
130            ),
131        };
132        op.set_attr_spirv_create_layout_rank(ctx, rank.into());
133        op
134    }
135
136    pub fn shape(&self, ctx: &Context) -> Vec<Value> {
137        let rank = self.rank(ctx);
138        let op = self.get_operation().deref(ctx);
139        op.operands().take(rank).collect()
140    }
141
142    pub fn strides(&self, ctx: &Context) -> Option<Vec<Value>> {
143        let rank = self.rank(ctx);
144        let op = self.get_operation().deref(ctx);
145        if op.get_num_operands() > rank {
146            Some(op.operands().skip(rank).collect())
147        } else {
148            None
149        }
150    }
151
152    pub fn rank(&self, ctx: &Context) -> usize {
153        self.get_attr_spirv_create_layout_rank(ctx).unwrap().0
154    }
155}
156
157#[cube_op(name = "spirv.create_view")]
158#[result_ty(argument)]
159#[op_traits(CanMaterialize, Pure)]
160pub struct CreateViewOp {}
161
162#[pliron_op(
163    name = "spirv.slice_layout",
164    format,
165    attributes = (spirv_slice_layout_rank: IndexAttr),
166    verifier = "succ"
167)]
168#[op_interfaces(NResultsInterface<1>, OneResultInterface, OperandSegmentInterface)]
169#[op_traits(CanMaterialize, Pure)]
170pub struct SliceOp;
171
172impl SliceOp {
173    pub fn new(ctx: &mut Context, layout: Value, offsets: Vec<Value>, shape: Vec<Value>) -> Self {
174        let (operands, segment_sizes) =
175            Self::compute_segment_sizes(vec![vec![layout], offsets, shape]);
176        let out_ty = layout.get_type(ctx);
177        let op = Self {
178            op: Operation::new(
179                ctx,
180                Self::get_concrete_op_info(),
181                vec![out_ty],
182                operands,
183                vec![],
184                0,
185            ),
186        };
187        op.set_operand_segment_sizes(ctx, segment_sizes);
188        op
189    }
190
191    pub fn layout(&self, ctx: &Context) -> Value {
192        self.get_operation().deref(ctx).get_operand(0)
193    }
194
195    pub fn offsets(&self, ctx: &Context) -> Vec<Value> {
196        self.get_segment(ctx, 1)
197    }
198
199    pub fn shape(&self, ctx: &Context) -> Vec<Value> {
200        self.get_segment(ctx, 2)
201    }
202}
cubecl_ir/dialect/spirv.rs

cubecl_ir/dialect/
spirv.rs
