vyre-libs 0.6.2

vyre Category A library ecosystem - pure-IR compositions over vyre-ops hardware primitives
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

use vyre::ir::{BufferAccess, BufferDecl, DataType, Expr, Node, Program};

use crate::builder::{
    build_indexed_map, strided_accumulate_child, INDEXED_MAP_OP_ID, STRIDED_ACCUMULATE_OP_ID,
};
use crate::harness::OpEntry;
use crate::region::wrap_anonymous;

fn u32s(words: &[u32]) -> Vec<u8> {
    vyre_primitives::wire::pack_u32_slice(words)
}

fn indexed_map_program() -> Program {
    build_indexed_map(
        INDEXED_MAP_OP_ID,
        vec![
            BufferDecl::storage("input", 0, BufferAccess::ReadOnly, DataType::U32).with_count(4),
            BufferDecl::output("out", 1, DataType::U32).with_count(4),
        ],
        "out",
        4,
        [4, 1, 1],
        |i| (i.clone(), Expr::add(Expr::load("input", i), Expr::u32(1))),
    )
}

fn strided_accumulate_program() -> Program {
    let tile = 4;
    let body = vec![
        Node::let_bind("local", Expr::LocalId { axis: 0 }),
        strided_accumulate_child(
            STRIDED_ACCUMULATE_OP_ID,
            tile,
            1,
            4,
            "acc",
            Expr::u32(0),
            "scratch",
            |idx, acc| Expr::add(acc, Expr::load("values", idx)),
        ),
        Node::barrier(),
        Node::if_then(
            Expr::eq(Expr::var("local"), Expr::u32(0)),
            vec![Node::store(
                "out",
                Expr::u32(0),
                Expr::load("scratch", Expr::u32(0)),
            )],
        ),
    ];
    Program::wrapped(
        vec![
            BufferDecl::storage("values", 0, BufferAccess::ReadOnly, DataType::U32).with_count(4),
            BufferDecl::workgroup("scratch", tile, DataType::U32),
            BufferDecl::output("out", 1, DataType::U32).with_count(1),
        ],
        [tile, 1, 1],
        vec![wrap_anonymous(STRIDED_ACCUMULATE_OP_ID, body)],
    )
}

inventory::submit! {
    OpEntry {
        id: INDEXED_MAP_OP_ID,
        build: indexed_map_program,
        test_inputs: Some(|| vec![vec![
            u32s(&[1, 2, 3, 4]),
        ]]),
        expected_output: Some(|| vec![vec![u32s(&[2, 3, 4, 5])]]),
        category: None,
    }
}

inventory::submit! {
    OpEntry {
        id: STRIDED_ACCUMULATE_OP_ID,
        build: strided_accumulate_program,
        test_inputs: Some(|| vec![vec![
            u32s(&[7, 11, 13, 17]),
        ]]),
        expected_output: Some(|| vec![vec![u32s(&[7])]]),
        category: None,
    }
}