nxpu_ir/
global.rs

1//! Global variables, address spaces, and resource bindings.
2
3use crate::arena::Handle;
4use crate::types::Type;
5
6/// Bitflags for storage buffer access modes.
7#[derive(Clone, Copy, Debug, Hash, Eq, PartialEq)]
8pub struct StorageAccess(u32);
9
10impl StorageAccess {
11    /// No access.
12    pub const EMPTY: Self = Self(0);
13    /// Read access.
14    pub const LOAD: Self = Self(1);
15    /// Write access.
16    pub const STORE: Self = Self(2);
17
18    /// Returns `true` if `self` contains all flags in `other`.
19    pub fn contains(self, other: Self) -> bool {
20        self.0 & other.0 == other.0
21    }
22
23    /// Returns `true` if no flags are set.
24    pub fn is_empty(self) -> bool {
25        self.0 == 0
26    }
27}
28
29impl std::ops::BitOr for StorageAccess {
30    type Output = Self;
31    fn bitor(self, rhs: Self) -> Self {
32        Self(self.0 | rhs.0)
33    }
34}
35
36impl std::ops::BitOrAssign for StorageAccess {
37    fn bitor_assign(&mut self, rhs: Self) {
38        self.0 |= rhs.0;
39    }
40}
41
42/// Memory address space for variables.
43#[derive(Clone, Copy, Debug, Hash, Eq, PartialEq)]
44pub enum AddressSpace {
45    /// Function-local storage.
46    Function,
47    /// Module-scope private storage.
48    Private,
49    /// Workgroup shared storage.
50    Workgroup,
51    /// Uniform buffer (read-only).
52    Uniform,
53    /// Storage buffer with specified access.
54    Storage { access: StorageAccess },
55}
56
57/// `@group(N) @binding(N)` resource binding.
58#[derive(Clone, Copy, Debug, Hash, Eq, PartialEq)]
59pub struct ResourceBinding {
60    /// Bind group index.
61    pub group: u32,
62    /// Binding index within the group.
63    pub binding: u32,
64}
65
66/// Built-in shader inputs/outputs.
67#[derive(Clone, Copy, Debug, Hash, Eq, PartialEq)]
68pub enum BuiltIn {
69    /// `@builtin(global_invocation_id)` — `vec3<u32>`
70    GlobalInvocationId,
71    /// `@builtin(local_invocation_id)` — `vec3<u32>`
72    LocalInvocationId,
73    /// `@builtin(local_invocation_index)` — `u32`
74    LocalInvocationIndex,
75    /// `@builtin(workgroup_id)` — `vec3<u32>`
76    WorkgroupId,
77    /// `@builtin(num_workgroups)` — `vec3<u32>`
78    NumWorkgroups,
79}
80
81/// A binding for a function argument or result.
82#[derive(Clone, Copy, Debug, Hash, Eq, PartialEq)]
83pub enum Binding {
84    /// A built-in shader variable.
85    BuiltIn(BuiltIn),
86    /// A user-defined location.
87    Location { location: u32 },
88}
89
90/// A module-scope variable.
91#[derive(Clone, Debug)]
92pub struct GlobalVariable {
93    /// Optional variable name.
94    pub name: Option<String>,
95    /// Address space (uniform, storage, etc.).
96    pub space: AddressSpace,
97    /// Optional resource binding (`@group(N) @binding(N)`).
98    pub binding: Option<ResourceBinding>,
99    /// The type of this variable.
100    pub ty: Handle<Type>,
101    /// Optional initializer expression.
102    pub init: Option<Handle<crate::Expression>>,
103    /// Optional memory layout annotation for tensor data.
104    pub layout: Option<crate::types::MemoryLayout>,
105}
106
107#[cfg(test)]
108mod tests {
109    use super::*;
110
111    #[test]
112    fn storage_access_flags() {
113        let read = StorageAccess::LOAD;
114        let write = StorageAccess::STORE;
115        let rw = read | write;
116        assert!(rw.contains(read));
117        assert!(rw.contains(write));
118        assert!(!read.contains(write));
119        assert!(!StorageAccess::EMPTY.contains(read));
120        assert!(StorageAccess::EMPTY.is_empty());
121    }
122
123    #[test]
124    fn storage_access_bitor_assign() {
125        let mut access = StorageAccess::LOAD;
126        access |= StorageAccess::STORE;
127        assert!(access.contains(StorageAccess::LOAD));
128        assert!(access.contains(StorageAccess::STORE));
129    }
130
131    #[test]
132    fn address_space_storage() {
133        let space = AddressSpace::Storage {
134            access: StorageAccess::LOAD | StorageAccess::STORE,
135        };
136        if let AddressSpace::Storage { access } = space {
137            assert!(access.contains(StorageAccess::LOAD));
138            assert!(access.contains(StorageAccess::STORE));
139        } else {
140            panic!("expected Storage");
141        }
142    }
143
144    #[test]
145    fn resource_binding() {
146        let binding = ResourceBinding {
147            group: 0,
148            binding: 3,
149        };
150        assert_eq!(binding.group, 0);
151        assert_eq!(binding.binding, 3);
152    }
153}
nxpu_ir/global.rs

nxpu_ir/
global.rs
