calyx 0.7.1

Compiler Infrastructure for Hardware Accelerator Generation
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
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153

import tvm
from tvm import relay
import relay_visitor
import relay_utils


def tensor_add():
    x = relay.var("x", relay.TensorType((2, 4), "int32"))
    y = relay.var("y", relay.TensorType((2, 4), "int32"))
    return relay.Function([x, y], relay.subtract(x, y))


def expand_dims():
    x = relay.var("x", shape=[512], dtype="int32")
    return relay.Function([x], relay.expand_dims(x, axis=1, num_newaxis=2))


def batch_flatten():
    x = relay.var("x", relay.TensorType((2, 5, 5), "int32"))
    return relay.Function([x], relay.nn.batch_flatten(x))


def batch_matmul():
    x = relay.var("x", shape=[1, 3, 3], dtype="float32")
    y = relay.var("y", shape=[1, 3, 3], dtype="float32")
    return relay.Function([x, y], relay.nn.batch_matmul(x, y))


def bias_add():
    x = relay.var("x", shape=[2, 4], dtype="float32")
    bias = relay.var("bias", shape=[4], dtype="float32")
    return relay.Function([x, bias], relay.nn.bias_add(data=x, bias=bias))


def relu():
    x = relay.var("x", shape=[2, 4], dtype="int32")
    return relay.Function([x], relay.nn.relu(x))


def dense():
    x = relay.var("x", shape=[1, 4096], dtype="int32")
    y = relay.var("y", shape=[10, 4096], dtype="int32")
    return relay.Function([x, y], relay.nn.dense(x, y, units=10))


def softmax():
    x = relay.var("x", shape=[1, 10], dtype="float32")
    return relay.Function([x], relay.nn.softmax(x))


def max_pool2d():
    data = relay.var("data", shape=[2, 2, 4, 4], dtype="int32")
    return relay.Function(
        [data],
        relay.nn.max_pool2d(
            data, padding=[0, 0, 0, 0], strides=[2, 2], pool_size=[2, 2]
        ),
    )


def conv2d():
    d = relay.var("data", shape=[5, 512, 14, 14], dtype="int32")
    w = relay.var("weight", shape=[512, 512, 3, 3], dtype="int32")
    return relay.Function(
        [d, w],
        relay.nn.conv2d(d, w, padding=[1, 1, 1, 1], channels=512, kernel_size=[3, 3]),
    )


def mlp_net():
    """The MLP test from Relay."""
    from tvm.relay.testing import mlp

    return mlp.get_net(1)


def vgg_net():
    """The VGG test from Relay."""
    from tvm.relay.testing import vgg

    return vgg.get_net(
        batch_size=5,
        image_shape=(3, 224, 224),
        num_classes=10,
        dtype="float32",
        num_layers=13,
        batch_norm=True,
    )


FUNCTIONS = {
    "tensor_add": tensor_add,
    "expand_dims": expand_dims,
    "batch_flatten": batch_flatten,
    "batch_matmul": batch_matmul,
    "bias_add": bias_add,
    "relu": relu,
    "dense": dense,
    "softmax": softmax,
    "conv2d": conv2d,
    "max_pool2d": max_pool2d,
    "mlp_net": mlp_net,
    "vgg_net": vgg_net,
}


def pretty_print_functions():
    """Pretty prints the available functions."""
    half = len(FUNCTIONS) // 2
    keys = list(FUNCTIONS.keys())
    for (f1, f2) in zip(keys[:half], keys[half:]):
        whitespace = (16 - len(f1)) * " "
        print(f"- {f1}{whitespace} - {f2}")


def run_example():
    import sys

    """Runs the example.
    Displays Relay IR if `-r` is found.
    Displays Calyx otherwise."""
    input = sys.argv[1:]
    if "-h" in input or not input:
        print(
            f"""
help  -h    Displays available functions to play with.
relay -r    Displays the Relay IR. Displays Calyx otherwise.

Available functions:"""
        )
        pretty_print_functions()
        return

    # See if the command line contains a correct function name.
    func_name = input[0]
    func = FUNCTIONS[func_name]() if func_name in FUNCTIONS.keys() else None
    if func is None:
        print(f"Function `{func_name}` is not a supported.")
        pretty_print_functions()
        return

    if "-r" in input:
        # Dump the Relay IR.
        print(relay_utils.python2relay(func))
    else:
        mod = tvm.IRModule.from_expr(func)
        relay_ir = mod["main"]
        # Compile and dump the Calyx.
        print(relay_visitor.emit_calyx(relay_ir))


if __name__ == "__main__":
    run_example()