cubecl_spirv/
tensor_indexing.rs

1use cubecl_core::ir::{self as core, ClampMode, TensorIndexingOps};
2use rspirv::spirv::Capability;
3
4use crate::{SpirvCompiler, SpirvTarget};
5
6impl<T: SpirvTarget> SpirvCompiler<T> {
7    pub fn compile_tensor_indexing(&mut self, op: TensorIndexingOps, out: Option<core::Value>) {
8        self.capabilities.insert(Capability::TensorAddressingNV);
9        let out = self.compile_value(out.unwrap());
10        match op {
11            TensorIndexingOps::CreateLayout {
12                shape,
13                strides,
14                clamp_mode,
15            } => {
16                let out_id = self.write_id(&out);
17                let ty = out.item().id(self);
18
19                let shape = shape
20                    .into_iter()
21                    .map(|it| self.compile_value(it))
22                    .collect::<Vec<_>>();
23                let shape = shape.iter().map(|it| self.read(it)).collect::<Vec<_>>();
24
25                let strides = strides.map(|s| {
26                    s.into_iter()
27                        .map(|it| self.compile_value(it))
28                        .collect::<Vec<_>>()
29                });
30                let strides = strides
31                    .as_ref()
32                    .map(|s| s.iter().map(|it| self.read(it)).collect::<Vec<_>>());
33
34                let mut layout = self.create_tensor_layout_nv(ty, None).unwrap();
35                // Write straight to out if strides and clamp is default
36                let result_id = match (&strides, clamp_mode) {
37                    (None, clamp_mode) if !matches!(clamp_mode, ClampMode::Constant(_)) => {
38                        Some(out_id)
39                    }
40                    (None, ClampMode::Constant(0)) => Some(out_id),
41                    _ => None,
42                };
43                layout = self
44                    .tensor_layout_set_dimension_nv(ty, result_id, layout, shape)
45                    .unwrap();
46                let result_id = match clamp_mode {
47                    ClampMode::Constant(0) => Some(out_id),
48                    ClampMode::Constant(_) => None,
49                    _ => Some(out_id),
50                };
51                if let Some(strides) = strides {
52                    layout = self
53                        .tensor_layout_set_stride_nv(ty, result_id, layout, strides)
54                        .unwrap();
55                }
56                match clamp_mode {
57                    ClampMode::Constant(0) => {}
58                    ClampMode::Constant(val) => {
59                        let val = self.const_u32(val);
60                        layout = self
61                            .tensor_layout_set_clamp_value_nv(ty, Some(out_id), layout, val)
62                            .unwrap();
63                    }
64                    _ => {}
65                }
66                self.write(&out, layout);
67            }
68            TensorIndexingOps::CreateView => {
69                let out_id = self.write_id(&out);
70                let ty = out.item().id(self);
71
72                self.create_tensor_view_nv(ty, Some(out_id)).unwrap();
73                self.write(&out, out_id);
74            }
75            TensorIndexingOps::Slice {
76                layout,
77                offsets,
78                shape,
79            } => {
80                let out_id = self.write_id(&out);
81                let ty = out.item().id(self);
82
83                let layout = self.compile_value(layout);
84                let layout = self.read(&layout);
85                let args = offsets
86                    .into_iter()
87                    .zip(shape)
88                    .flat_map(|(offset, shape)| {
89                        let offset = self.compile_value(offset);
90                        let shape = self.compile_value(shape);
91                        [offset, shape]
92                    })
93                    .collect::<Vec<_>>();
94                let args = args.iter().map(|it| self.read(it)).collect::<Vec<_>>();
95                self.tensor_layout_slice_nv(ty, Some(out_id), layout, args)
96                    .unwrap();
97                self.write(&out, out_id);
98            }
99        }
100    }
101}
cubecl_spirv/tensor_indexing.rs

cubecl_spirv/
tensor_indexing.rs
