vyre-driver-wgpu 0.7.1

wgpu backend for vyre IR - implements VyreBackend, owns GPU runtime, buffer pool, pipeline cache
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
100
101
102
103
104
105
106
107
108
109

//! Failure-oriented tests for subgroup support detection in lowering.
//!
//! Pins that subgroup builtins are emitted only when the program
//! actually contains subgroup ops, and that capability reporting
//! matches the emitted shader.

use vyre::ir::{BinOp, BufferDecl, DataType, Expr, Node, Program};
use vyre::VyreBackend;
use vyre_driver_wgpu::emit;
use vyre_driver_wgpu::WgpuBackend;

fn lower_wgsl(program: &Program) -> String {
    emit::lower(program).expect("Fix: test program must lower to WGSL")
}

#[test]
fn subgroup_builtins_absent_when_program_has_no_subgroup_ops() {
    let program = Program::wrapped(
        vec![BufferDecl::output("out", 0, DataType::U32)],
        [1, 1, 1],
        vec![Node::store("out", Expr::u32(0), Expr::u32(1))],
    );
    let wgsl = lower_wgsl(&program);
    assert!(
        !wgsl.contains("subgroup"),
        "Fix: program without subgroup ops must not emit subgroup builtins or intrinsics. WGSL:\n{wgsl}"
    );
    assert!(
        !wgsl.contains("SubgroupInvocationId"),
        "Fix: subgroup builtin must not appear when unused. WGSL:\n{wgsl}"
    );
    assert!(
        !wgsl.contains("SubgroupSize"),
        "Fix: subgroup builtin must not appear when unused. WGSL:\n{wgsl}"
    );
}

#[test]
fn subgroup_builtins_present_when_program_uses_subgroup_ops() {
    let program = Program::wrapped(
        vec![BufferDecl::output("out", 0, DataType::U32)],
        [1, 1, 1],
        vec![Node::store(
            "out",
            Expr::u32(0),
            Expr::BinOp {
                op: BinOp::WaveReduce,
                left: Box::new(Expr::u32(7)),
                right: Box::new(Expr::u32(0)),
            },
        )],
    );
    let wgsl = lower_wgsl(&program);
    assert!(
        wgsl.contains("subgroup"),
        "Fix: program with WaveReduce must emit subgroup intrinsics. WGSL:\n{wgsl}"
    );
}

#[test]
fn subgroup_shuffle_lowers_to_subgroup_builtin() {
    let program = Program::wrapped(
        vec![BufferDecl::output("out", 0, DataType::U32)],
        [1, 1, 1],
        vec![Node::store(
            "out",
            Expr::u32(0),
            Expr::subgroup_shuffle(Expr::u32(1), Expr::u32(2)),
        )],
    );
    let wgsl = lower_wgsl(&program);
    assert!(
        wgsl.contains("subgroupShuffle"),
        "Fix: subgroup_shuffle must lower to subgroupShuffle intrinsic. WGSL:\n{wgsl}"
    );
}

#[test]
fn subgroup_local_id_builtin_present_when_used() {
    let program = Program::wrapped(
        vec![BufferDecl::output("out", 0, DataType::U32)],
        [1, 1, 1],
        vec![Node::store("out", Expr::u32(0), Expr::subgroup_local_id())],
    );
    let wgsl = lower_wgsl(&program);
    assert!(
        wgsl.contains("@builtin(subgroup_invocation_id)"),
        "Fix: subgroup_local_id() must cause subgroup_invocation_id builtin to be declared. WGSL:\n{wgsl}"
    );
}

#[test]
fn subgroup_size_matches_backend_report() {
    let backend = WgpuBackend::acquire().expect(
        "WgpuBackend::acquire failed on a machine that must have a GPU. \
         Fix: repair adapter/driver probing instead of skipping subgroup capability checks.",
    );

    let reported = backend.subgroup_size();
    let info = backend.adapter_info();

    if let Some(size) = reported {
        assert!(
            size > 0,
            "Fix: reported subgroup size must be positive when Some. Adapter: {}",
            info.name
        );
    }
}