rust-gpu/crates/spirv-std/src/unique/builtin.rs at 00cacd0812b9ec229e1435bb3636e529536467b3 · Rust-GPU/rust-gpu · GitHub

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
use crate::unique::{Global, GlobalUniqueId, GlobalUniqueIndex, ScalarValue};
use core::arch::asm;
use glam::UVec3;

#[inline]
pub fn global_invocation_id() -> GlobalUniqueId {
    unsafe {
        let slot = UVec3::default();
        asm! {
            "%builtin = OpVariable typeof{result} Input",
            "OpDecorate %builtin BuiltIn GlobalInvocationId",
            "%result = OpLoad typeof*{result} %builtin",
            "OpStore {result} %result",
            result = in(reg) &slot,
        }
        GlobalUniqueId::new_unchecked(slot)
    }
}

// TODO: this built-in doesn't work and is deprecated
#[inline]
pub fn workgroup_size() -> ScalarValue<UVec3, Global> {
    unsafe {
        let slot = UVec3::default();
        asm! {
            "%builtin = OpVariable typeof{result} Input",
            "OpDecorate %builtin BuiltIn WorkgroupSize",
            "%result = OpLoad typeof*{result} %builtin",
            "OpStore {result} %result",
            result = in(reg) &slot,
        }
        ScalarValue::new(slot)
    }
}

#[inline]
pub fn num_workgroups() -> ScalarValue<UVec3, Global> {
    unsafe {
        let slot = UVec3::default();
        asm! {
            "%builtin = OpVariable typeof{result} Input",
            "OpDecorate %builtin BuiltIn NumWorkgroups",
            "%result = OpLoad typeof*{result} %builtin",
            "OpStore {result} %result",
            result = in(reg) &slot,
        }
        ScalarValue::new(slot)
    }
}

#[inline]
pub fn global_invocation_index() -> GlobalUniqueIndex {
    unsafe {
        let gid = global_invocation_id();
        let wg_cnt = workgroup_size();
        GlobalUniqueIndex::new_unchecked(gid.x + gid.y * wg_cnt.x + gid.z * wg_cnt.y * wg_cnt.x)
    }
}