tokitai-operator 0.1.0

Verified DL kernel compiler: formally-checked GEMM, p-adic, sheaf, contract-carrying ops. Paper-artifact grade.
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

//! `Representation` and `DeviceId`.
//!
//! `Representation` is the abstract data layout (dense CPU,
//! p-adic scalar, sparse, etc.). `DeviceId` is the per-device
//! identifier (the empty string for CPU; a HIP device id for
//! GPU). The `ObjectMeta::representation` field binds the
//! tensor to a specific device and layout.
//!
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct DeviceId(pub String);

impl DeviceId {
    pub fn cpu() -> Self {
        Self("cpu".to_string())
    }
}

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct RepresentationId(pub String);

impl RepresentationId {
    pub fn new(id: impl Into<String>) -> Self {
        Self(id.into())
    }
}

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum Layout {
    Scalar,
    Contiguous,
    Strided(Vec<isize>),
    Blocked { block_shape: Vec<usize> },
    CoverIndexed,
    PadicLimbArray,
    PadicBall,
    Symbolic,
}

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum Representation {
    DenseTensor { layout: Layout, device: DeviceId },
    SparseTensor { layout: Layout, device: DeviceId },
    LazyExpression,
    CoverIndexedSection { device: DeviceId },
    PadicScalar,
    Symbolic,
}

impl Representation {
    pub fn dense_cpu() -> Self {
        Self::DenseTensor {
            layout: Layout::Contiguous,
            device: DeviceId::cpu(),
        }
    }

    pub fn id(&self) -> RepresentationId {
        match self {
            Self::DenseTensor { layout, device } => {
                RepresentationId::new(format!("dense_tensor:{}:{}", layout.id(), device.0))
            }
            Self::SparseTensor { layout, device } => {
                RepresentationId::new(format!("sparse_tensor:{}:{}", layout.id(), device.0))
            }
            Self::LazyExpression => RepresentationId::new("lazy_expression"),
            Self::CoverIndexedSection { device } => {
                RepresentationId::new(format!("cover_indexed_section:{}", device.0))
            }
            Self::PadicScalar => RepresentationId::new("padic_scalar"),
            Self::Symbolic => RepresentationId::new("symbolic"),
        }
    }
}

impl Layout {
    pub fn id(&self) -> String {
        match self {
            Self::Scalar => "scalar".to_string(),
            Self::Contiguous => "contiguous".to_string(),
            Self::Strided(strides) => format!("strided:{strides:?}"),
            Self::Blocked { block_shape } => format!("blocked:{block_shape:?}"),
            Self::CoverIndexed => "cover_indexed".to_string(),
            Self::PadicLimbArray => "padic_limb_array".to_string(),
            Self::PadicBall => "padic_ball".to_string(),
            Self::Symbolic => "symbolic".to_string(),
        }
    }
}