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
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172

# type: ignore
import tvm
from tvm import relay
from calyx.py_ast import CompVar, Stdlib, CompInst, Cell, Invoke, CompPort
from calyx.utils import bits_needed
from typing import List
from dataclasses import dataclass

# Mapping from the tensor dimensions to the
# corresponding Calyx primitive.
NumDimsToCell = {
    0: Stdlib().register,
    1: Stdlib().seq_mem_d1,
    2: Stdlib().seq_mem_d2,
    3: Stdlib().seq_mem_d3,
    4: Stdlib().seq_mem_d4,
}


@dataclass
class DahliaFuncDef:
    """Necessary information to compute a Dahlia
    function definition."""

    function_id: str
    component_name: str
    dest: CompVar
    args: List[CompVar]
    attributes: tvm.ir.Attrs
    data_type: str
    component: CompInst


def get_dims(c: CompInst):
    """Mapping from memory to number of dimensions."""
    id = c.id
    id2dimensions = {
        "std_reg": 0,
        "seq_mem_d1": 1,
        "seq_mem_d2": 2,
        "seq_mem_d3": 3,
        "seq_mem_d4": 4,
    }
    assert id in id2dimensions, f"{id} not supported."
    return id2dimensions[id]


def get_dimension_sizes(c: CompInst) -> List[int]:
    """Given a cell `c`, returns the corresponding
    memory sizes.
    Example:
    comb_mem_d1(32, 8, 3) returns [8]."""
    dims = get_dims(c)
    return [c.args[i] for i in range(1, dims + 1)]


def get_addr_ports(c: CompInst):
    """Returns a list of (address, index size)
    for each address port in the component
    instance."""
    dims = get_dims(c)
    addresses = range(0, dims)
    indices = range(dims + 1, dims << 1 + 1)
    return [(f"addr{i}", c.args[n]) for (i, n) in zip(addresses, indices)]


def emit_invoke_control(
    decl: CompVar, dest: Cell, args: List[Cell], old_args=[], old_dest=None
) -> Invoke:
    """Returns the Invoke control."""
    ref_cells = []
    inputs = []

    def add_arg(cell):
        comp = cell.comp
        param = f"{cell.id.name}"
        arg = CompVar(cell.id.name)

        # If this is a constant or a register, connect the ports
        if any(p in comp.id for p in ["reg", "const"]):
            inputs.append((f"{param}", CompPort(arg, "out")))
        else:
            ref_cells.append((param, arg))

    # this function is similar to add_arg, but is for the case when we are
    # "reusing" a Dahlia Function (which will later be a Calyx component)
    # and therefore need to use the same parameter names as the previous invoke
    def add_arg2(arg_cell, param_cell):
        assert (
            arg_cell.comp == param_cell.comp
        ), "arg cell and param cell must be same component"
        comp = arg_cell.comp
        param = f"{param_cell.id.name}"
        arg = CompVar(arg_cell.id.name)

        # If this is a constant or a register, connect the ports
        if any(p in comp.id for p in ["reg", "const"]):
            inputs.append((f"{param}", CompPort(arg, "out")))
        else:
            ref_cells.append((param, arg))

    if len(old_args) == 0:
        for cell in args:
            add_arg(cell)
        add_arg(dest)
    else:
        # case for when we are "reusing" a Dahlia Function/Calyx component and
        # therefore need to make sure we're using the previous parameter names
        assert len(old_args) == len(
            args
        ), "we are reusing a dahlia function but the args are different lengths"
        assert old_dest is not None, "if using old_args must provide an old_dest too"
        for cell1, cell2 in zip(args, old_args):
            add_arg2(cell1, cell2)
        add_arg2(dest, old_dest)

    return Invoke(decl, inputs, [], ref_cells)


def get_dahlia_data_type(relay_type) -> str:
    """Gets the Dahlia data type from the given Relay type.
    It maps the types in the following manner:

    Relay   |        Dahlia
    --------|-------------------------------
     int    |   (`bit`, width)
     float  |   (`fix`, (width, width // 2))
    """
    width = get_bitwidth(relay_type)

    if "int" in relay_type.dtype:
        return f"bit<{width}>"
    if "float" in relay_type.dtype:
        return f"fix<{width}, {width // 2}>"
    assert 0, f"{relay_type} is not supported."


def get_bitwidth(relay_type) -> int:
    """Gets the bitwidth from a Relay type."""
    dtype = relay_type.dtype
    assert "int" in dtype or "float" in dtype, f"{relay_type} not supported."
    return int("".join(filter(str.isdigit, dtype)))


def get_memory(name: str, type: tvm.ir.Type) -> Cell:
    """Returns a Calyx memory for a given TVM type.
    For non-Tensor types, a register is returned.
    Otherwise, a memory with the corresponding dimension size
    is returned, if it exists in Calyx."""
    dims = type.concrete_shape
    # Bitwidth, along with sizes and index sizes (if it is a Tensor).
    args = [get_bitwidth(type)] + [d for d in dims] + [bits_needed(d) for d in dims]

    num_dims = len(dims)
    assert num_dims in NumDimsToCell, f"Memory of size {num_dims} not supported."

    return Cell(CompVar(name), NumDimsToCell[num_dims](*args), is_external=True)


def python2relay(func) -> str:
    """Used to lower Relay IR from the
    TVM Python library."""
    seq = tvm.transform.Sequential(
        [
            relay.transform.SimplifyExpr(),
            relay.transform.SimplifyInference(),
            relay.transform.ToANormalForm(),
        ]
    )
    mod_opt = tvm.IRModule.from_expr(func)
    mod_opt = seq(mod_opt)
    return mod_opt["main"]