vyre 0.4.0

GPU compute intermediate representation with a standard operation library
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
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130

use crate::ir::{BufferDecl, DataType, Expr, Node, Program};
use crate::ops::AlgebraicLaw;
use crate::ops::OpSpec;

// WGSL lowering marker for `string.prefix_brace`.
//
// No special per-op lowering is needed. The normal IR lowerer handles the
// sequential prefix-scan composition.

pub const INPUTS: &[DataType] = &[DataType::Bytes];

pub const LAWS: &[AlgebraicLaw] = &[AlgebraicLaw::Bounded {
    lo: 0,
    hi: u32::MAX,
}];

pub const OUTPUTS: &[DataType] = &[DataType::U32, DataType::U32, DataType::U32];

/// Build the canonical prefix-scan program.
#[must_use]
pub fn prefix_brace_program() -> Program {
    Program::new(
        vec![
            BufferDecl::read("file_bytes", 0, DataType::Bytes),
            BufferDecl::output("brace_depths", 1, DataType::U32),
            BufferDecl::read_write("nested_depths", 2, DataType::U32),
            BufferDecl::read_write("newline_sums", 3, DataType::U32),
        ],
        [1, 1, 1],
        vec![
            Node::let_bind("file_size", Expr::buf_len("file_bytes")),
            // Always initialise the trailing-sentinel slot so empty inputs are
            // well-defined.
            Node::store("newline_sums", Expr::u32(0), Expr::u32(0)),
            Node::if_then(
                Expr::gt(Expr::var("file_size"), Expr::u32(0)),
                vec![
                    Node::let_bind("brace_depth", Expr::u32(0)),
                    Node::let_bind("nested_depth", Expr::u32(0)),
                    Node::let_bind("newline_count", Expr::u32(0)),
                    Node::loop_for(
                        "i",
                        Expr::u32(0),
                        Expr::var("file_size"),
                        vec![
                            Node::store("brace_depths", Expr::var("i"), Expr::var("brace_depth")),
                            Node::store("nested_depths", Expr::var("i"), Expr::var("nested_depth")),
                            Node::let_bind("byte", Expr::load("file_bytes", Expr::var("i"))),
                            // `{`
                            Node::if_then(
                                Expr::eq(Expr::var("byte"), Expr::u32(123)),
                                vec![
                                    Node::assign(
                                        "brace_depth",
                                        Expr::add(Expr::var("brace_depth"), Expr::u32(1)),
                                    ),
                                    Node::assign(
                                        "nested_depth",
                                        Expr::add(Expr::var("nested_depth"), Expr::u32(1)),
                                    ),
                                ],
                            ),
                            // `}`
                            Node::if_then(
                                Expr::eq(Expr::var("byte"), Expr::u32(125)),
                                vec![
                                    Node::if_then(
                                        Expr::gt(Expr::var("brace_depth"), Expr::u32(0)),
                                        vec![Node::assign(
                                            "brace_depth",
                                            Expr::sub(Expr::var("brace_depth"), Expr::u32(1)),
                                        )],
                                    ),
                                    Node::if_then(
                                        Expr::gt(Expr::var("nested_depth"), Expr::u32(0)),
                                        vec![Node::assign(
                                            "nested_depth",
                                            Expr::sub(Expr::var("nested_depth"), Expr::u32(1)),
                                        )],
                                    ),
                                ],
                            ),
                            // `(`
                            Node::if_then(
                                Expr::eq(Expr::var("byte"), Expr::u32(40)),
                                vec![Node::assign(
                                    "nested_depth",
                                    Expr::add(Expr::var("nested_depth"), Expr::u32(1)),
                                )],
                            ),
                            // `)`
                            Node::if_then(
                                Expr::eq(Expr::var("byte"), Expr::u32(41)),
                                vec![Node::if_then(
                                    Expr::gt(Expr::var("nested_depth"), Expr::u32(0)),
                                    vec![Node::assign(
                                        "nested_depth",
                                        Expr::sub(Expr::var("nested_depth"), Expr::u32(1)),
                                    )],
                                )],
                            ),
                            // `\n`
                            Node::if_then(
                                Expr::eq(Expr::var("byte"), Expr::u32(10)),
                                vec![Node::assign(
                                    "newline_count",
                                    Expr::add(Expr::var("newline_count"), Expr::u32(1)),
                                )],
                            ),
                            Node::store(
                                "newline_sums",
                                Expr::add(Expr::var("i"), Expr::u32(1)),
                                Expr::var("newline_count"),
                            ),
                        ],
                    ),
                ],
            ),
        ],
    )
}

/// Declarative operation specification for `string.prefix_brace`.
pub const SPEC: OpSpec = OpSpec::composition(
    "string.prefix_brace",
    INPUTS,
    OUTPUTS,
    LAWS,
    prefix_brace_program,
);