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
100
101
102
103

use super::{Elem, Item};
use serde::{Deserialize, Serialize};

#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq)]
#[allow(missing_docs)]
pub enum Variable {
    GlobalInputArray(u16, Item),
    GlobalScalar(u16, Elem),
    GlobalOutputArray(u16, Item),
    Local(u16, Item, u8),
    LocalScalar(u16, Elem, u8),
    ConstantScalar(f64, Elem),
    SharedMemory(u16, Item, u32),
    LocalArray(u16, Item, u8, u32),
    Id,
    LocalInvocationIndex,
    LocalInvocationIdX,
    LocalInvocationIdY,
    LocalInvocationIdZ,
    WorkgroupIdX,
    WorkgroupIdY,
    WorkgroupIdZ,
    GlobalInvocationIdX,
    GlobalInvocationIdY,
    GlobalInvocationIdZ,
    Rank,
    WorkgroupSizeX,
    WorkgroupSizeY,
    WorkgroupSizeZ,
    NumWorkgroupsX,
    NumWorkgroupsY,
    NumWorkgroupsZ,
}

impl Variable {
    pub(crate) fn index(&self) -> Option<u16> {
        match self {
            Variable::GlobalInputArray(idx, _) => Some(*idx),
            Variable::GlobalScalar(idx, _) => Some(*idx),
            Variable::Local(idx, _, _) => Some(*idx),
            Variable::LocalScalar(idx, _, _) => Some(*idx),
            Variable::GlobalOutputArray(idx, _) => Some(*idx),
            Variable::ConstantScalar(_, _) => None,
            Variable::SharedMemory(idx, _, _) => Some(*idx),
            Variable::LocalArray(idx, _, _, _) => Some(*idx),
            Variable::Id => None,
            Variable::LocalInvocationIndex => None,
            Variable::LocalInvocationIdX => None,
            Variable::LocalInvocationIdY => None,
            Variable::LocalInvocationIdZ => None,
            Variable::Rank => None,
            Variable::WorkgroupIdX => None,
            Variable::WorkgroupIdY => None,
            Variable::WorkgroupIdZ => None,
            Variable::GlobalInvocationIdX => None,
            Variable::GlobalInvocationIdY => None,
            Variable::GlobalInvocationIdZ => None,
            Variable::WorkgroupSizeX => None,
            Variable::WorkgroupSizeY => None,
            Variable::WorkgroupSizeZ => None,
            Variable::NumWorkgroupsX => None,
            Variable::NumWorkgroupsY => None,
            Variable::NumWorkgroupsZ => None,
        }
    }
    pub(crate) fn item(&self) -> Item {
        match self {
            Variable::GlobalInputArray(_, item) => *item,
            Variable::GlobalOutputArray(_, item) => *item,
            Variable::GlobalScalar(_, elem) => Item::Scalar(*elem),
            Variable::Local(_, item, _) => *item,
            Variable::LocalScalar(_, elem, _) => Item::Scalar(*elem),
            Variable::ConstantScalar(_, elem) => Item::Scalar(*elem),
            Variable::SharedMemory(_, item, _) => *item,
            Variable::LocalArray(_, item, _, _) => *item,
            Variable::Id => Item::Scalar(Elem::UInt),
            Variable::Rank => Item::Scalar(Elem::UInt),
            Variable::LocalInvocationIndex => Item::Scalar(Elem::UInt),
            Variable::LocalInvocationIdX => Item::Scalar(Elem::UInt),
            Variable::LocalInvocationIdY => Item::Scalar(Elem::UInt),
            Variable::LocalInvocationIdZ => Item::Scalar(Elem::UInt),
            Variable::WorkgroupIdX => Item::Scalar(Elem::UInt),
            Variable::WorkgroupIdY => Item::Scalar(Elem::UInt),
            Variable::WorkgroupIdZ => Item::Scalar(Elem::UInt),
            Variable::GlobalInvocationIdX => Item::Scalar(Elem::UInt),
            Variable::GlobalInvocationIdY => Item::Scalar(Elem::UInt),
            Variable::GlobalInvocationIdZ => Item::Scalar(Elem::UInt),
            Variable::WorkgroupSizeX => Item::Scalar(Elem::UInt),
            Variable::WorkgroupSizeY => Item::Scalar(Elem::UInt),
            Variable::WorkgroupSizeZ => Item::Scalar(Elem::UInt),
            Variable::NumWorkgroupsX => Item::Scalar(Elem::UInt),
            Variable::NumWorkgroupsY => Item::Scalar(Elem::UInt),
            Variable::NumWorkgroupsZ => Item::Scalar(Elem::UInt),
        }
    }
}

// Useful with the gpu! macro.
impl From<&Variable> for Variable {
    fn from(value: &Variable) -> Self {
        *value
    }
}