heterocl.tvm.tag module¶
Tag class for TVM operators.
-
class
TagScope
(tag)[source]¶ 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)[source]¶ 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
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)))