heterocl.tvm.tag module

Tag class for TVM operators.

class TagScope(tag)

Bases: object

Tag scope object to set tag for operators, working as context manager and decorator both. See also tag_scope.

current = None

tag_scope(tag)

The operator tag scope.

Parameters

tag (str) – The tag name.

Returns

tag_scope – The tag scope object, which can be used as decorator or context manger.

Return type

TagScope

Example

n = tvm.var('n')
m = tvm.var('m')
l = tvm.var('l')
A = tvm.placeholder((n, l), name='A')
B = tvm.placeholder((m, l), name='B')
k = tvm.reduce_axis((0, l), name='k')

with tvm.tag_scope(tag='matmul'):
    C = tvm.compute((n, m), lambda i, j: tvm.sum(A[i, k] * B[j, k], axis=k))

# or use tag_scope as decorator
@tvm.tag_scope(tag="conv")
def compute_relu(data):
    return tvm.compute(data.shape, lambda *i: tvm.select(data(*i) < 0, 0.0, data(*i)))