cubecl-ir 0.11.0-pre.1

Intermediate representation for CubeCL
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99

use crate::{ClampMode, TypeHash};
use crate::{OperationArgs, OperationCode};
use alloc::string::ToString;
use alloc::vec::Vec;
use core::fmt::Display;

use crate::OperationReflect;

use super::Value;

#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
#[derive(Debug, Clone, TypeHash, PartialEq, Eq, Hash, OperationCode)]
#[operation(opcode_name = TensorIndexingOpCode)]
/// Operations available on a barrier
pub enum TensorIndexingOps {
    CreateLayout {
        shape: Vec<Value>,
        strides: Option<Vec<Value>>,
        clamp_mode: ClampMode,
    },
    CreateView,
    Slice {
        layout: Value,
        offsets: Vec<Value>,
        shape: Vec<Value>,
    },
}

impl OperationReflect for TensorIndexingOps {
    type OpCode = TensorIndexingOpCode;

    fn op_code(&self) -> Self::OpCode {
        self.__match_opcode()
    }

    fn args(&self) -> Option<Vec<Value>> {
        match self {
            TensorIndexingOps::CreateLayout { .. }
            | TensorIndexingOps::CreateView
            | TensorIndexingOps::Slice { .. } => None,
        }
    }

    fn sanitize_args(&mut self, scope: &crate::Scope) {
        match self {
            TensorIndexingOps::CreateLayout { shape, strides, .. } => {
                shape.sanitize_args_ptr(scope);
                strides.sanitize_args_ptr(scope);
            }
            TensorIndexingOps::CreateView => {}
            TensorIndexingOps::Slice { offsets, shape, .. } => {
                offsets.sanitize_args_ptr(scope);
                shape.sanitize_args_ptr(scope);
            }
        }
    }
}

impl Display for TensorIndexingOps {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        match self {
            TensorIndexingOps::CreateLayout {
                shape,
                strides,
                clamp_mode,
            } => {
                let shape: Vec<_> = shape.iter().map(|it| it.to_string()).collect();
                let strides: Option<Vec<_>> = strides
                    .as_ref()
                    .map(|strides| strides.iter().map(|it| it.to_string()).collect());
                let strides = strides
                    .map(|strides| alloc::format!("[{}]", strides.join(", ")))
                    .unwrap_or("None".into());
                write!(
                    f,
                    "create_layout([{}], strides: {strides}, clamp_mode: {clamp_mode:?}",
                    shape.join(", ")
                )
            }
            TensorIndexingOps::CreateView => {
                write!(f, "create_view()")
            }
            TensorIndexingOps::Slice {
                layout,
                offsets,
                shape,
            } => {
                let offsets = offsets.iter().map(|it| it.to_string()).collect::<Vec<_>>();
                let shape = shape.iter().map(|it| it.to_string()).collect::<Vec<_>>();
                write!(
                    f,
                    "slice({layout}, offsets: [{}], shape: [{}])",
                    offsets.join(", "),
                    shape.join(", ")
                )
            }
        }
    }
}