cubecl_ir/
tensor_indexing.rs

1use crate::{ClampMode, TypeHash};
2use crate::{OperationArgs, OperationCode};
3use alloc::string::ToString;
4use alloc::vec::Vec;
5use core::fmt::Display;
6
7use crate::OperationReflect;
8
9use super::Value;
10
11#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
12#[derive(Debug, Clone, TypeHash, PartialEq, Eq, Hash, OperationCode)]
13#[operation(opcode_name = TensorIndexingOpCode)]
14/// Operations available on a barrier
15pub enum TensorIndexingOps {
16    CreateLayout {
17        shape: Vec<Value>,
18        strides: Option<Vec<Value>>,
19        clamp_mode: ClampMode,
20    },
21    CreateView,
22    Slice {
23        layout: Value,
24        offsets: Vec<Value>,
25        shape: Vec<Value>,
26    },
27}
28
29impl OperationReflect for TensorIndexingOps {
30    type OpCode = TensorIndexingOpCode;
31
32    fn op_code(&self) -> Self::OpCode {
33        self.__match_opcode()
34    }
35
36    fn args(&self) -> Option<Vec<Value>> {
37        match self {
38            TensorIndexingOps::CreateLayout { .. }
39            | TensorIndexingOps::CreateView
40            | TensorIndexingOps::Slice { .. } => None,
41        }
42    }
43
44    fn sanitize_args(&mut self, scope: &crate::Scope) {
45        match self {
46            TensorIndexingOps::CreateLayout { shape, strides, .. } => {
47                shape.sanitize_args_ptr(scope);
48                strides.sanitize_args_ptr(scope);
49            }
50            TensorIndexingOps::CreateView => {}
51            TensorIndexingOps::Slice { offsets, shape, .. } => {
52                offsets.sanitize_args_ptr(scope);
53                shape.sanitize_args_ptr(scope);
54            }
55        }
56    }
57}
58
59impl Display for TensorIndexingOps {
60    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
61        match self {
62            TensorIndexingOps::CreateLayout {
63                shape,
64                strides,
65                clamp_mode,
66            } => {
67                let shape: Vec<_> = shape.iter().map(|it| it.to_string()).collect();
68                let strides: Option<Vec<_>> = strides
69                    .as_ref()
70                    .map(|strides| strides.iter().map(|it| it.to_string()).collect());
71                let strides = strides
72                    .map(|strides| alloc::format!("[{}]", strides.join(", ")))
73                    .unwrap_or("None".into());
74                write!(
75                    f,
76                    "create_layout([{}], strides: {strides}, clamp_mode: {clamp_mode:?}",
77                    shape.join(", ")
78                )
79            }
80            TensorIndexingOps::CreateView => {
81                write!(f, "create_view()")
82            }
83            TensorIndexingOps::Slice {
84                layout,
85                offsets,
86                shape,
87            } => {
88                let offsets = offsets.iter().map(|it| it.to_string()).collect::<Vec<_>>();
89                let shape = shape.iter().map(|it| it.to_string()).collect::<Vec<_>>();
90                write!(
91                    f,
92                    "slice({layout}, offsets: [{}], shape: [{}])",
93                    offsets.join(", "),
94                    shape.join(", ")
95                )
96            }
97        }
98    }
99}
cubecl_ir/tensor_indexing.rs

cubecl_ir/
tensor_indexing.rs
