From af19f1593cf58cd000c3262dfd687d21e2d14986 Mon Sep 17 00:00:00 2001 From: Jian Weng Date: Wed, 12 Dec 2018 09:49:46 -0800 Subject: [PATCH 01/23] add a test case of value index --- tests/python/unittest/test_hybrid_script.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/python/unittest/test_hybrid_script.py b/tests/python/unittest/test_hybrid_script.py index 7efbbe43ee21..692220c0cc66 100644 --- a/tests/python/unittest/test_hybrid_script.py +++ b/tests/python/unittest/test_hybrid_script.py @@ -523,5 +523,3 @@ def kernel_b(b, a): test_value_index() # TODO: # test_inplace() - - From e37734fbb93f2e36a0e0cefb1e1b6829e913c72e Mon Sep 17 00:00:00 2001 From: Jian Weng Date: Wed, 12 Dec 2018 11:46:22 -0800 Subject: [PATCH 02/23] fix python2/3 incompat --- tests/python/unittest/test_hybrid_script.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/python/unittest/test_hybrid_script.py b/tests/python/unittest/test_hybrid_script.py index 692220c0cc66..4e97c253f058 100644 --- a/tests/python/unittest/test_hybrid_script.py +++ b/tests/python/unittest/test_hybrid_script.py @@ -496,6 +496,8 @@ def kernel_b(b, a): d = kernel_b(c, b) sch = tvm.create_schedule(d.op) module = tvm.build(sch, [a, d]) + # A bug to fix? + print(tvm.lower(sch, [a, b], simple_mode=True)) assert module np_a = numpy.arange(16).astype('int32') From 0b3b0abe0db8a58b0014fe98f6b67dd04fe9764d Mon Sep 17 00:00:00 2001 From: Jian Weng Date: Wed, 12 Dec 2018 13:13:54 -0800 Subject: [PATCH 03/23] fix floordiv --- tests/python/unittest/test_hybrid_script.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/python/unittest/test_hybrid_script.py b/tests/python/unittest/test_hybrid_script.py index 4e97c253f058..692220c0cc66 100644 --- a/tests/python/unittest/test_hybrid_script.py +++ b/tests/python/unittest/test_hybrid_script.py @@ -496,8 +496,6 @@ def kernel_b(b, a): d = kernel_b(c, b) sch = tvm.create_schedule(d.op) module = tvm.build(sch, [a, d]) - # A bug to fix? - print(tvm.lower(sch, [a, b], simple_mode=True)) assert module np_a = numpy.arange(16).astype('int32') From bf0590556f775a13a218d6fa39c40855d3aeb582 Mon Sep 17 00:00:00 2001 From: Jian Weng Date: Wed, 12 Dec 2018 13:13:54 -0800 Subject: [PATCH 04/23] fix floordiv --- python/tvm/hybrid/parser.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/tvm/hybrid/parser.py b/python/tvm/hybrid/parser.py index ba10dd8dde3c..a06db476c2b9 100644 --- a/python/tvm/hybrid/parser.py +++ b/python/tvm/hybrid/parser.py @@ -51,7 +51,7 @@ class HybridParser(ast.NodeVisitor): ast.Eq : operator.eq, ast.NotEq : operator.ne, ast.And : _all, - ast.Or : _any, + ast.Or : _any } From 88c2c2727d8e7f26fa9a3de32ca610883bc3802c Mon Sep 17 00:00:00 2001 From: Jian Weng Date: Wed, 12 Dec 2018 10:17:37 -0800 Subject: [PATCH 05/23] init codegen hybrid --- src/codegen/codegen_hybrid.cc | 0 src/codegen/codegen_hybrid.h | 0 2 files changed, 0 insertions(+), 0 deletions(-) create mode 100644 src/codegen/codegen_hybrid.cc create mode 100644 src/codegen/codegen_hybrid.h diff --git a/src/codegen/codegen_hybrid.cc b/src/codegen/codegen_hybrid.cc new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/src/codegen/codegen_hybrid.h b/src/codegen/codegen_hybrid.h new file mode 100644 index 000000000000..e69de29bb2d1 From 3f4cb1fea85b009e2f9d25b34693d2893ebded85 Mon Sep 17 00:00:00 2001 From: Jian Weng Date: Wed, 12 Dec 2018 11:36:36 -0800 Subject: [PATCH 06/23] refactor out func call --- python/tvm/hybrid/calls.py | 97 +++++++++++++++++++++++++++ python/tvm/hybrid/intrin.py | 15 +---- python/tvm/hybrid/parser.py | 128 +++++++++++++++++++----------------- 3 files changed, 164 insertions(+), 76 deletions(-) create mode 100644 python/tvm/hybrid/calls.py diff --git a/python/tvm/hybrid/calls.py b/python/tvm/hybrid/calls.py new file mode 100644 index 000000000000..9b94df7613d0 --- /dev/null +++ b/python/tvm/hybrid/calls.py @@ -0,0 +1,97 @@ +"""Intrinsics of TVM-Python Hybrid Script for Python compilation time""" + +import ast +from .. import api as _api +from .. import expr as _expr +from .. import make as _make +from ..ir_pass import Equal +from ..stmt import For +from .util import _internal_assert + +LOOP_INTRIN = { + 'range' : For.Serial, + 'unroll' : For.Unrolled, + 'parallel' : For.Parallel, + 'vectorize': For.Vectorized, +} + +def _range(visitor, func_id, args): + """Handling TVM loop types""" + n = len(args) + if n == 1: + low, ext = _api.const(0, dtype='int32'), visitor.visit(args[0]) + else: + _internal_assert(n == 2, "A loop intrinsic should only have 1 or 2 arguments!") + low, ext = visitor.visit(args[0]), visitor.visit(args[1]) + if not Equal(low, _api.const(0, dtype='int32')): + ext = ext - low + for_type = LOOP_INTRIN[func_id] + iter_var = None + return iter_var, low, ext, for_type + + +range = unroll = vectorize = parallel = _range #pylint: disable=invalid-name + + +def bind(visitor, func_id, args): + """Handling TVM thread binding""" + n = len(args) + _internal_assert(n == 2, "A loop bind should only have 2 arguments!") + _internal_assert(isinstance(args[0], ast.Str), \ + "A loop bind's first argument should be a string!") + _vn = args[0].s + iter_var = thread_axis(args[0].s) + low, ext = _api.const(0, dtype='int32'), visitor.visit(args[1]) + for_type = None + return iter_var, low, ext, for_type + + +def _math_intrin(visitor, func_id, args): + from .. import intrin + return getattr(intrin, func_id)(*[visitor.visit(arg) for arg in args]) + +sqrt = log = exp = tanh = sigmoid = power = popcount = _math_intrin + +def _min_max(self, func_id, args): + n = len(args) + _internal_assert(n == 2, "Max/Min function should have 2 elements") + a, b = self.visit(args[0]), self.visit(args[1]) + return getattr(_make, func_id.title())(a, b) + +min = max = _min_max + +def _allocate_tensor(visitor, func_id, args): + """Handling TVM tensor allocation""" + n = len(args) + _internal_assert(isinstance(args[0], ast.Tuple), \ + "allocate's first argument should be a tuple of shape!") + shape = tuple(visitor.visit(i) for i in args[0].elts) + if func_id == 'output_tensor': + _internal_assert(not visitor.loops_above, \ + "Are you sure to allocate a output buffer multiple times?") + for i in shape: + _internal_assert(isinstance(i, _expr.Expr), "The shape should be an expression") + if n > 1: + if isinstance(args[1], ast.Str): + dtype = args[1].s + else: + _internal_assert(isinstance(args[1], ast.Attribute), \ + "Unable to evaluate to get data type") + to_eval = args[1] + _internal_assert(isinstance(to_eval.value, ast.Name), \ + "Unable to evaluate the attribute to get data type") + _internal_assert(to_eval.attr == 'dtype', \ + "Only dtype attribute is supported so far") + dtype = visitor._get_buffer_from_id(to_eval.value.id).dtype + else: + dtype = 'float32' + if n > 2: + _internal_assert(isinstance(args[2], ast.Str), \ + "The data scope should be an string") + _internal_assert(func_id != 'output_tensor', "Output tensor cannot specify scope") + scope = args[2].s + else: + scope = 'global' if func_id != 'output_tensor' else 'output' + return (shape, dtype, scope) + +output_tensor = allocate = _allocate_tensor diff --git a/python/tvm/hybrid/intrin.py b/python/tvm/hybrid/intrin.py index 92e259585b7a..48e92a8bf5ac 100644 --- a/python/tvm/hybrid/intrin.py +++ b/python/tvm/hybrid/intrin.py @@ -1,7 +1,6 @@ -"""Intrinsics of TVM-Python Hybrid Script for Python runtime""" +"""Intrinsics of TVM-Python Hybrid Script for Python emulation runtime""" import numpy -from ..stmt import For class _range(object): """Base class of the loop ranges in hybrid script""" @@ -102,15 +101,3 @@ def sigmoid(x): 'sigmoid' : sigmoid, 'popcount' : popcount } - - -LOOP_INTRIN = { - 'range' : For.Serial, - 'unroll' : For.Unrolled, - 'parallel' : For.Parallel, - 'vectorize': For.Vectorized, - 'bind' : None -} - - -MATH_INTRIN = ['sqrt', 'log', 'exp', 'tanh', 'sigmoid', 'power', 'popcount'] diff --git a/python/tvm/hybrid/parser.py b/python/tvm/hybrid/parser.py index a06db476c2b9..883747f9bc87 100644 --- a/python/tvm/hybrid/parser.py +++ b/python/tvm/hybrid/parser.py @@ -5,7 +5,7 @@ import logging import sys from .util import make_nop, halide_imm_types, is_docstring, _internal_assert -from .intrin import LOOP_INTRIN, MATH_INTRIN +from . import calls from .var_decl import determine_variable_usage from ..api import thread_axis from ..api import all as _all @@ -121,7 +121,6 @@ def _get_buffer_from_id(self, s, for_provide=False): return self.alloc_buffers[s][0] - #pylint: disable=invalid-name, missing-docstring def visit_Module(self, node): _internal_assert(len(node.body) == 1, \ @@ -329,67 +328,72 @@ def visit_Call(self, node): # Yet, no function pointer supported _internal_assert(isinstance(node.func, ast.Name), \ "Only id-function function call is supported so far!") + func_id = node.func.id - n = len(node.args) - if func_id in LOOP_INTRIN.keys() and func_id != 'bind': - if n == 1: - low, ext = _api.const(0, dtype='int32'), self.visit(node.args[0]) - else: - _internal_assert(n == 2, "A loop intrinsic should only have 1 or 2 arguments!") - low, ext = self.visit(node.args[0]), self.visit(node.args[1]) - if not _ir_pass.Equal(low, _api.const(0, dtype='int32')): - ext = ext - low - for_type = LOOP_INTRIN[func_id] - iter_var = None - return iter_var, low, ext, for_type - elif func_id == 'bind': - _internal_assert(n == 2, "A loop bind should only have 2 arguments!") - _internal_assert(isinstance(node.args[0], ast.Str), \ - "A loop bind's first argument should be a string!") - _vn = node.args[0].s - iter_var = thread_axis(node.args[0].s) - low, ext = _api.const(0, dtype='int32'), self.visit(node.args[1]) - for_type = None - return iter_var, low, ext, for_type - elif func_id in MATH_INTRIN: - return getattr(intrin, func_id)(*[self.visit(arg) for arg in node.args]) - elif func_id in ['allocate', 'output_tensor']: - _internal_assert(isinstance(node.args[0], ast.Tuple), \ - "allocate's first argument should be a tuple of shape!") - shape = tuple(self.visit(i) for i in node.args[0].elts) - if func_id == 'output_tensor': - _internal_assert(not self.loops_above, \ - "Are you sure to allocate a output buffer multiple times?") - for i in shape: - _internal_assert(isinstance(i, _expr.Expr), "The shape should be an expression") - if n > 1: - if isinstance(node.args[1], ast.Str): - dtype = node.args[1].s - else: - _internal_assert(isinstance(node.args[1], ast.Attribute), \ - "Unable to evaluate to get data type") - to_eval = node.args[1] - _internal_assert(isinstance(to_eval.value, ast.Name), \ - "Unable to evaluate the attribute to get data type") - _internal_assert(to_eval.attr == 'dtype', \ - "Only dtype attribute is supported so far") - dtype = self._get_buffer_from_id(to_eval.value.id).dtype - else: - dtype = 'float32' - if n > 2: - _internal_assert(isinstance(node.args[2], ast.Str), \ - "The data scope should be an string") - _internal_assert(func_id != 'output_tensor', "Output tensor cannot specify scope") - scope = node.args[2].s - else: - scope = 'global' if func_id != 'output_tensor' else 'output' - return (shape, dtype, scope) - elif func_id == 'max' or func_id == 'min': - _internal_assert(n == 2, "Max/Min function should have 2 elements") - a, b = self.visit(node.args[0]), self.visit(node.args[1]) - return getattr(_make, func_id.title())(a, b) - else: - raise ValueError("Function call not supported yet!") + try: + return getattr(calls, func_id)(self, func_id, node.args) + except Exception as e: + raise e + #n = len(node.args) + #if func_id in LOOP_INTRIN.keys() and func_id != 'bind': + # if n == 1: + # low, ext = _api.const(0, dtype='int32'), self.visit(node.args[0]) + # else: + # _internal_assert(n == 2, "A loop intrinsic should only have 1 or 2 arguments!") + # low, ext = self.visit(node.args[0]), self.visit(node.args[1]) + # if not _ir_pass.Equal(low, _api.const(0, dtype='int32')): + # ext = ext - low + # for_type = LOOP_INTRIN[func_id] + # iter_var = None + # return iter_var, low, ext, for_type + #elif func_id == 'bind': + # _internal_assert(n == 2, "A loop bind should only have 2 arguments!") + # _internal_assert(isinstance(node.args[0], ast.Str), \ + # "A loop bind's first argument should be a string!") + # _vn = node.args[0].s + # iter_var = thread_axis(node.args[0].s) + # low, ext = _api.const(0, dtype='int32'), self.visit(node.args[1]) + # for_type = None + # return iter_var, low, ext, for_type + #elif func_id in MATH_INTRIN: + # return getattr(intrin, func_id)(*[self.visit(arg) for arg in node.args]) + #elif func_id in ['allocate', 'output_tensor']: + # _internal_assert(isinstance(node.args[0], ast.Tuple), \ + # "allocate's first argument should be a tuple of shape!") + # shape = tuple(self.visit(i) for i in node.args[0].elts) + # if func_id == 'output_tensor': + # _internal_assert(not self.loops_above, \ + # "Are you sure to allocate a output buffer multiple times?") + # for i in shape: + # _internal_assert(isinstance(i, _expr.Expr), "The shape should be an expression") + # if n > 1: + # if isinstance(node.args[1], ast.Str): + # dtype = node.args[1].s + # else: + # _internal_assert(isinstance(node.args[1], ast.Attribute), \ + # "Unable to evaluate to get data type") + # to_eval = node.args[1] + # _internal_assert(isinstance(to_eval.value, ast.Name), \ + # "Unable to evaluate the attribute to get data type") + # _internal_assert(to_eval.attr == 'dtype', \ + # "Only dtype attribute is supported so far") + # dtype = self._get_buffer_from_id(to_eval.value.id).dtype + # else: + # dtype = 'float32' + # if n > 2: + # _internal_assert(isinstance(node.args[2], ast.Str), \ + # "The data scope should be an string") + # _internal_assert(func_id != 'output_tensor', "Output tensor cannot specify scope") + # scope = node.args[2].s + # else: + # scope = 'global' if func_id != 'output_tensor' else 'output' + # return (shape, dtype, scope) + #elif func_id == 'max' or func_id == 'min': + # _internal_assert(n == 2, "Max/Min function should have 2 elements") + # a, b = self.visit(node.args[0]), self.visit(node.args[1]) + # return getattr(_make, func_id.title())(a, b) + #else: + # raise ValueError("Function call not supported yet!") def visit_For(self, node): From 9972c216777b894f1dca3a32a17bdd9e730f3f4b Mon Sep 17 00:00:00 2001 From: Jian Weng Date: Wed, 12 Dec 2018 14:48:34 -0800 Subject: [PATCH 07/23] remove comments --- python/tvm/hybrid/parser.py | 60 ------------------------------------- 1 file changed, 60 deletions(-) diff --git a/python/tvm/hybrid/parser.py b/python/tvm/hybrid/parser.py index 883747f9bc87..94608a1d71db 100644 --- a/python/tvm/hybrid/parser.py +++ b/python/tvm/hybrid/parser.py @@ -334,66 +334,6 @@ def visit_Call(self, node): return getattr(calls, func_id)(self, func_id, node.args) except Exception as e: raise e - #n = len(node.args) - #if func_id in LOOP_INTRIN.keys() and func_id != 'bind': - # if n == 1: - # low, ext = _api.const(0, dtype='int32'), self.visit(node.args[0]) - # else: - # _internal_assert(n == 2, "A loop intrinsic should only have 1 or 2 arguments!") - # low, ext = self.visit(node.args[0]), self.visit(node.args[1]) - # if not _ir_pass.Equal(low, _api.const(0, dtype='int32')): - # ext = ext - low - # for_type = LOOP_INTRIN[func_id] - # iter_var = None - # return iter_var, low, ext, for_type - #elif func_id == 'bind': - # _internal_assert(n == 2, "A loop bind should only have 2 arguments!") - # _internal_assert(isinstance(node.args[0], ast.Str), \ - # "A loop bind's first argument should be a string!") - # _vn = node.args[0].s - # iter_var = thread_axis(node.args[0].s) - # low, ext = _api.const(0, dtype='int32'), self.visit(node.args[1]) - # for_type = None - # return iter_var, low, ext, for_type - #elif func_id in MATH_INTRIN: - # return getattr(intrin, func_id)(*[self.visit(arg) for arg in node.args]) - #elif func_id in ['allocate', 'output_tensor']: - # _internal_assert(isinstance(node.args[0], ast.Tuple), \ - # "allocate's first argument should be a tuple of shape!") - # shape = tuple(self.visit(i) for i in node.args[0].elts) - # if func_id == 'output_tensor': - # _internal_assert(not self.loops_above, \ - # "Are you sure to allocate a output buffer multiple times?") - # for i in shape: - # _internal_assert(isinstance(i, _expr.Expr), "The shape should be an expression") - # if n > 1: - # if isinstance(node.args[1], ast.Str): - # dtype = node.args[1].s - # else: - # _internal_assert(isinstance(node.args[1], ast.Attribute), \ - # "Unable to evaluate to get data type") - # to_eval = node.args[1] - # _internal_assert(isinstance(to_eval.value, ast.Name), \ - # "Unable to evaluate the attribute to get data type") - # _internal_assert(to_eval.attr == 'dtype', \ - # "Only dtype attribute is supported so far") - # dtype = self._get_buffer_from_id(to_eval.value.id).dtype - # else: - # dtype = 'float32' - # if n > 2: - # _internal_assert(isinstance(node.args[2], ast.Str), \ - # "The data scope should be an string") - # _internal_assert(func_id != 'output_tensor', "Output tensor cannot specify scope") - # scope = node.args[2].s - # else: - # scope = 'global' if func_id != 'output_tensor' else 'output' - # return (shape, dtype, scope) - #elif func_id == 'max' or func_id == 'min': - # _internal_assert(n == 2, "Max/Min function should have 2 elements") - # a, b = self.visit(node.args[0]), self.visit(node.args[1]) - # return getattr(_make, func_id.title())(a, b) - #else: - # raise ValueError("Function call not supported yet!") def visit_For(self, node): From 67dcd49e329552a02340bd4563c615a76eb6eb37 Mon Sep 17 00:00:00 2001 From: Jian Weng Date: Wed, 12 Dec 2018 14:54:49 -0800 Subject: [PATCH 08/23] continue working... --- python/tvm/hybrid/parser.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/tvm/hybrid/parser.py b/python/tvm/hybrid/parser.py index 94608a1d71db..adae10b15cff 100644 --- a/python/tvm/hybrid/parser.py +++ b/python/tvm/hybrid/parser.py @@ -51,7 +51,7 @@ class HybridParser(ast.NodeVisitor): ast.Eq : operator.eq, ast.NotEq : operator.ne, ast.And : _all, - ast.Or : _any + ast.Or : _any, } From a88d37848b100161a8c5ae9d3bd29d2e83bbcd0e Mon Sep 17 00:00:00 2001 From: Jian Weng Date: Wed, 12 Dec 2018 15:03:27 -0800 Subject: [PATCH 09/23] cleanup stash --- python/tvm/hybrid/parser.py | 2 +- tests/python/unittest/test_hybrid_script.py | 19 +++++++++++++++++++ 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/python/tvm/hybrid/parser.py b/python/tvm/hybrid/parser.py index adae10b15cff..d02c24ff4a72 100644 --- a/python/tvm/hybrid/parser.py +++ b/python/tvm/hybrid/parser.py @@ -372,7 +372,7 @@ def visit_Return(self, node): _internal_assert(isinstance(i, ast.Name), "What do you return?") ids.append(i.id) _internal_assert(len(set(ids)) == len(ids), "Duplicated tensors in the return tuples") - if len(ids) != len(self.outputs): + if len(ids) < len(self.outputs): logging.log(logging.CRITICAL, '[Warning] Not all the output buffers returned!') self.outputs = [self._args[i] for i in ids] self.returned = True diff --git a/tests/python/unittest/test_hybrid_script.py b/tests/python/unittest/test_hybrid_script.py index 692220c0cc66..62006384baef 100644 --- a/tests/python/unittest/test_hybrid_script.py +++ b/tests/python/unittest/test_hybrid_script.py @@ -336,6 +336,7 @@ def blur(a): for dj in range(3): s += a[i-di, j-dj] b[i-2, j-2] = s / 9.0 + b[0, 0] += s return b a = tvm.placeholder((32, 32), 'float32', 'a') @@ -506,6 +507,24 @@ def kernel_b(b, a): module(tvm.ndarray.array(np_a), res) tvm.testing.assert_allclose(res.asnumpy(), ref) +def test_func_call(): + @tvm.hybrid.script + def kernel_a(a): + b = output_tensor((16, ), 'int32') + c = output_tensor((4, 4), 'int32') + for i in range(16): + b[i] = a[i] + 2 + c[i // 4, i % 4] = a[i] + 1 + return b, c + + @tvm.hybrid.script + def kernel_b(b, a): + c = output_tensor((4, 4), 'int32') + for i in range(4): + for j in range(4): + c[i, j] = a[i * 4 + j] * b[i, j] + return c + if __name__ == "__main__": From ff0bb819ad5ecb93ae651ad91b820eba4a7ff8b4 Mon Sep 17 00:00:00 2001 From: Jian Weng Date: Wed, 12 Dec 2018 16:01:19 -0800 Subject: [PATCH 10/23] more friendly user hints --- python/tvm/hybrid/var_decl.py | 2 ++ tests/python/unittest/test_hybrid_script.py | 23 ++++++++++----------- 2 files changed, 13 insertions(+), 12 deletions(-) diff --git a/python/tvm/hybrid/var_decl.py b/python/tvm/hybrid/var_decl.py index 27df87874377..ea15e7ad1f8f 100644 --- a/python/tvm/hybrid/var_decl.py +++ b/python/tvm/hybrid/var_decl.py @@ -75,6 +75,8 @@ def visit_Name(self, node): else: decl, loop, usage = self.status[node.id] usage.add(type(node.ctx)) + _internal_assert(loop in self.scope_level, + "%s is used out of the scope it is defined!" % node.id) self.status[node.id] = (decl, loop, usage) diff --git a/tests/python/unittest/test_hybrid_script.py b/tests/python/unittest/test_hybrid_script.py index 62006384baef..a64a674d8376 100644 --- a/tests/python/unittest/test_hybrid_script.py +++ b/tests/python/unittest/test_hybrid_script.py @@ -336,7 +336,6 @@ def blur(a): for dj in range(3): s += a[i-di, j-dj] b[i-2, j-2] = s / 9.0 - b[0, 0] += s return b a = tvm.placeholder((32, 32), 'float32', 'a') @@ -528,17 +527,17 @@ def kernel_b(b, a): if __name__ == "__main__": - test_outer_product() - test_fanout() - test_looptype() - test_if() - test_bind() - test_math_intrin() + #test_outer_product() + #test_fanout() + #test_looptype() + #test_if() + #test_bind() + #test_math_intrin() test_non_zero() - test_allocate() - test_upstream() - test_downstream() - test_const_param() - test_value_index() + #test_allocate() + #test_upstream() + #test_downstream() + #test_const_param() + #test_value_index() # TODO: # test_inplace() From d3e99d11c25dc7da338285f1ec867205d893fcbf Mon Sep 17 00:00:00 2001 From: Jian Weng Date: Wed, 12 Dec 2018 16:06:24 -0800 Subject: [PATCH 11/23] fix lint --- python/tvm/hybrid/calls.py | 12 ++++++++---- python/tvm/hybrid/parser.py | 2 -- 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/python/tvm/hybrid/calls.py b/python/tvm/hybrid/calls.py index 9b94df7613d0..c883cb008800 100644 --- a/python/tvm/hybrid/calls.py +++ b/python/tvm/hybrid/calls.py @@ -1,4 +1,5 @@ -"""Intrinsics of TVM-Python Hybrid Script for Python compilation time""" +"""Intrinsics of TVM-Python Hybrid Script for Python compilation time +semantic support.""" import ast from .. import api as _api @@ -8,6 +9,8 @@ from ..stmt import For from .util import _internal_assert +#pylint: disable=redefined-builtin + LOOP_INTRIN = { 'range' : For.Serial, 'unroll' : For.Unrolled, @@ -35,6 +38,7 @@ def _range(visitor, func_id, args): def bind(visitor, func_id, args): """Handling TVM thread binding""" + _internal_assert(func_id == "bind", "This function cannot be directly invoked!") n = len(args) _internal_assert(n == 2, "A loop bind should only have 2 arguments!") _internal_assert(isinstance(args[0], ast.Str), \ @@ -50,7 +54,7 @@ def _math_intrin(visitor, func_id, args): from .. import intrin return getattr(intrin, func_id)(*[visitor.visit(arg) for arg in args]) -sqrt = log = exp = tanh = sigmoid = power = popcount = _math_intrin +sqrt = log = exp = tanh = sigmoid = power = popcount = _math_intrin #pylint: disable=invalid-name def _min_max(self, func_id, args): n = len(args) @@ -58,7 +62,7 @@ def _min_max(self, func_id, args): a, b = self.visit(args[0]), self.visit(args[1]) return getattr(_make, func_id.title())(a, b) -min = max = _min_max +min = max = _min_max #pylint: disable=invalid-name def _allocate_tensor(visitor, func_id, args): """Handling TVM tensor allocation""" @@ -94,4 +98,4 @@ def _allocate_tensor(visitor, func_id, args): scope = 'global' if func_id != 'output_tensor' else 'output' return (shape, dtype, scope) -output_tensor = allocate = _allocate_tensor +output_tensor = allocate = _allocate_tensor #pylint: disable=invalid-name diff --git a/python/tvm/hybrid/parser.py b/python/tvm/hybrid/parser.py index d02c24ff4a72..9b503106a159 100644 --- a/python/tvm/hybrid/parser.py +++ b/python/tvm/hybrid/parser.py @@ -7,12 +7,10 @@ from .util import make_nop, halide_imm_types, is_docstring, _internal_assert from . import calls from .var_decl import determine_variable_usage -from ..api import thread_axis from ..api import all as _all from ..api import any as _any from .. import expr as _expr from .. import make as _make -from .. import intrin from .. import api as _api from .. import ir_pass as _ir_pass From 18a30a9ff72ccbbfedd87a0d8905c470862442a0 Mon Sep 17 00:00:00 2001 From: Jian Weng Date: Fri, 14 Dec 2018 00:15:55 -0800 Subject: [PATCH 12/23] add boolop test; inter-func call supported? --- tests/python/unittest/test_hybrid_script.py | 66 +++++++++++++-------- 1 file changed, 40 insertions(+), 26 deletions(-) diff --git a/tests/python/unittest/test_hybrid_script.py b/tests/python/unittest/test_hybrid_script.py index a64a674d8376..0d2d2d0e6dff 100644 --- a/tests/python/unittest/test_hybrid_script.py +++ b/tests/python/unittest/test_hybrid_script.py @@ -508,36 +508,50 @@ def kernel_b(b, a): def test_func_call(): @tvm.hybrid.script - def kernel_a(a): - b = output_tensor((16, ), 'int32') - c = output_tensor((4, 4), 'int32') - for i in range(16): - b[i] = a[i] + 2 - c[i // 4, i % 4] = a[i] + 1 - return b, c - - @tvm.hybrid.script - def kernel_b(b, a): - c = output_tensor((4, 4), 'int32') - for i in range(4): - for j in range(4): - c[i, j] = a[i * 4 + j] * b[i, j] - return c + def foo(a, b): + for i in range(10): + a[i] = i + 1.0 + for i in range(10): + b[i] = i + 1.0 + c = outer_product(10, 10, a, b) + d = output_tensor(c.shape, c.dtype) + for i in range(10): + for j in range(10): + d[i, j] = c[i, j] + i * j + return d + a = tvm.placeholder((10, ), name='a') + b = tvm.placeholder((10, ), name='b') + run_and_check(foo, [a, b]) +def test_bool(): + @tvm.hybrid.script + def foo(a): + b = output_tensor(a.shape, a.dtype) + b[0] = 1.2 + for i in range(1, a.shape[0] - 1): + if a[i] * a[i - 1] < a[i] or a[i] * a[i - 1] < a[i - 1]: + b[i] = a[i] + else: + b[i] = 0.0 + return b + a = tvm.placeholder((10, ), name='a') + run_and_check(foo, [a]) if __name__ == "__main__": - #test_outer_product() - #test_fanout() - #test_looptype() - #test_if() - #test_bind() - #test_math_intrin() + test_outer_product() + test_fanout() + test_looptype() + test_if() + test_bind() + test_math_intrin() test_non_zero() - #test_allocate() - #test_upstream() - #test_downstream() - #test_const_param() - #test_value_index() + test_allocate() + test_upstream() + test_downstream() + test_const_param() + test_value_index() + test_func_call() + test_bool() # TODO: # test_inplace() From b6036b4ecb9593ca9ac88e68831cc4c20096bddd Mon Sep 17 00:00:00 2001 From: Jian Weng Date: Fri, 14 Dec 2018 00:16:07 -0800 Subject: [PATCH 13/23] add boolop test; inter-func call supported? --- python/tvm/hybrid/api.py | 6 +- python/tvm/hybrid/calls.py | 61 +++++----- python/tvm/hybrid/parser.py | 212 +++++++++++++++++++++------------- python/tvm/hybrid/util.py | 17 +++ python/tvm/hybrid/var_decl.py | 13 ++- 5 files changed, 184 insertions(+), 125 deletions(-) diff --git a/python/tvm/hybrid/api.py b/python/tvm/hybrid/api.py index 5267731f4f52..0c1fedf4b0b8 100644 --- a/python/tvm/hybrid/api.py +++ b/python/tvm/hybrid/api.py @@ -24,17 +24,15 @@ def wrapped_func(func, *args, **kwargs): #pylint: disable=missing-docstring from .util import _enter_hybrid_runtime, _restore_runtime, _is_tvm_arg_types if _is_tvm_arg_types(args): src = _pruned_source(func) - parser = parse_python(src, args) + parser = parse_python(src, func.__globals__, args) input_tensors = [] for i in args: if isinstance(i, Tensor): input_tensors.append(i) - - op = _tvm_internal._HybridOp(parser.func_name, "HybridOp", None, input_tensors, + op = _tvm_internal._HybridOp(parser.func_name, "HybridOp", None, input_tensors, \ parser.outputs, parser.parsed_body) res = [op.output(i) for i in range(len(parser.outputs))] - return res[0] if len(res) == 1 else res intersect = _enter_hybrid_runtime(func) diff --git a/python/tvm/hybrid/calls.py b/python/tvm/hybrid/calls.py index c883cb008800..d094f957bfcb 100644 --- a/python/tvm/hybrid/calls.py +++ b/python/tvm/hybrid/calls.py @@ -5,6 +5,7 @@ from .. import api as _api from .. import expr as _expr from .. import make as _make +from ..container import Array from ..ir_pass import Equal from ..stmt import For from .util import _internal_assert @@ -18,14 +19,14 @@ 'vectorize': For.Vectorized, } -def _range(visitor, func_id, args): +def _range(func_id, args): """Handling TVM loop types""" n = len(args) if n == 1: - low, ext = _api.const(0, dtype='int32'), visitor.visit(args[0]) + low, ext = _api.const(0, dtype='int32'), args[0] else: _internal_assert(n == 2, "A loop intrinsic should only have 1 or 2 arguments!") - low, ext = visitor.visit(args[0]), visitor.visit(args[1]) + low, ext = args[0], args[1] if not Equal(low, _api.const(0, dtype='int32')): ext = ext - low for_type = LOOP_INTRIN[func_id] @@ -36,64 +37,54 @@ def _range(visitor, func_id, args): range = unroll = vectorize = parallel = _range #pylint: disable=invalid-name -def bind(visitor, func_id, args): +def bind(func_id, args): """Handling TVM thread binding""" _internal_assert(func_id == "bind", "This function cannot be directly invoked!") - n = len(args) - _internal_assert(n == 2, "A loop bind should only have 2 arguments!") + _internal_assert(len(args) == 2, "A loop bind should only have 2 arguments!") _internal_assert(isinstance(args[0], ast.Str), \ "A loop bind's first argument should be a string!") - _vn = args[0].s - iter_var = thread_axis(args[0].s) - low, ext = _api.const(0, dtype='int32'), visitor.visit(args[1]) + iter_var = thread_axis(args[0]) + low, ext = _api.const(0, dtype='int32'), args[1] for_type = None return iter_var, low, ext, for_type -def _math_intrin(visitor, func_id, args): +def _math_intrin(func_id, args): from .. import intrin - return getattr(intrin, func_id)(*[visitor.visit(arg) for arg in args]) + return getattr(intrin, func_id)(*args) sqrt = log = exp = tanh = sigmoid = power = popcount = _math_intrin #pylint: disable=invalid-name -def _min_max(self, func_id, args): - n = len(args) - _internal_assert(n == 2, "Max/Min function should have 2 elements") - a, b = self.visit(args[0]), self.visit(args[1]) - return getattr(_make, func_id.title())(a, b) + +def _min_max(func_id, args): + _internal_assert(len(args) == 2, "Max/Min function should have 2 elements") + return getattr(_make, func_id.title())(args[0], args[1]) + min = max = _min_max #pylint: disable=invalid-name -def _allocate_tensor(visitor, func_id, args): + +def _allocate_tensor(func_id, args): """Handling TVM tensor allocation""" n = len(args) - _internal_assert(isinstance(args[0], ast.Tuple), \ + _internal_assert(isinstance(_api.convert(args[0]), Array), \ "allocate's first argument should be a tuple of shape!") - shape = tuple(visitor.visit(i) for i in args[0].elts) - if func_id == 'output_tensor': - _internal_assert(not visitor.loops_above, \ - "Are you sure to allocate a output buffer multiple times?") + shape = args[0] for i in shape: _internal_assert(isinstance(i, _expr.Expr), "The shape should be an expression") if n > 1: - if isinstance(args[1], ast.Str): - dtype = args[1].s - else: - _internal_assert(isinstance(args[1], ast.Attribute), \ - "Unable to evaluate to get data type") - to_eval = args[1] - _internal_assert(isinstance(to_eval.value, ast.Name), \ - "Unable to evaluate the attribute to get data type") - _internal_assert(to_eval.attr == 'dtype', \ - "Only dtype attribute is supported so far") - dtype = visitor._get_buffer_from_id(to_eval.value.id).dtype + _internal_assert(isinstance(args[1], str), + "The data type should be an str") + _internal_assert(args[1].startswith('int') or args[1].startswith('float'), \ + "The data type should be either int or float!") + dtype = args[1] else: dtype = 'float32' if n > 2: - _internal_assert(isinstance(args[2], ast.Str), \ + _internal_assert(isinstance(args[2], str), \ "The data scope should be an string") _internal_assert(func_id != 'output_tensor', "Output tensor cannot specify scope") - scope = args[2].s + scope = args[2] else: scope = 'global' if func_id != 'output_tensor' else 'output' return (shape, dtype, scope) diff --git a/python/tvm/hybrid/parser.py b/python/tvm/hybrid/parser.py index 9b503106a159..121b1cf9f6de 100644 --- a/python/tvm/hybrid/parser.py +++ b/python/tvm/hybrid/parser.py @@ -4,11 +4,12 @@ import operator import logging import sys -from .util import make_nop, halide_imm_types, is_docstring, _internal_assert +from .util import make_nop, halide_imm_types, is_docstring, _internal_assert, replace_io from . import calls from .var_decl import determine_variable_usage from ..api import all as _all from ..api import any as _any +from ..tensor import Tensor, Operation from .. import expr as _expr from .. import make as _make from .. import api as _api @@ -60,7 +61,7 @@ class HybridParser(ast.NodeVisitor): } - def __init__(self, args, usage, func_name=None): + def __init__(self, args, usage, symbols, func_name=None): """ Parameters ---------- @@ -79,32 +80,49 @@ def __init__(self, args, usage, func_name=None): self.args = list(args) self.usage = usage.copy() self._args = {} # Dict maps arg name to actual arg instance (either a var or a buffer) - self.alloc_buffers = {} # Buffers formed by allocate instructions + self.alloc_buffers = {} # Buffers formed by explicit allocate instructions self.loops_above = {} # State variable that indicates loop levels above the current node - self.var_consts = {} # Variables that are determined as readonly in previous stage + self.variables = {} # The status of defined variables self.func_name = func_name # The name of the function to be lowered self.outputs = [] # Output tensors' name self.side_effect = set() # Tensors with side effects self.parsed_body = None # The parsed HalideIR body - self.returned = False + self.returned = False # If this function has a valid return + self.symbols = symbols # The global context def wrap_up_realize(self, node, body): """Wrap up all the variables which will no longer be used""" + pop_buf = [] + pop_var = [] for key, val in self.usage.items(): - if key in self.var_consts.keys(): - continue _, level, _ = val - if level == node: - if key in self._args.keys(): + if level != node: + continue + if key in self._args.keys(): + continue + if key in self.alloc_buffers.keys(): + _buf, _scope = self.alloc_buffers[key] + if _scope == 'output': continue - else: - _buf, _scope = self.alloc_buffers[key] - _domain = [_make.range_by_min_extent(0, i) for i in _buf.shape] - _dtype = _buf.dtype - _true = _api.convert(True) - body = _make.Realize(_buf.op, 0, _dtype, _domain, _true, body) - body = _make.AttrStmt(_buf.op, 'realize_scope', _api.convert(_scope), body) + pop_buf.append(key) + else: + _internal_assert(key in self.variables.keys(), + "Key should be either in one of args, buffers, and vars") + if not isinstance(self.variables[key], tuple): + continue + _buf, _scope = self.variables[key] + pop_var.append(key) + _domain = [_make.range_by_min_extent(0, i) for i in _buf.shape] + _dtype = _buf.dtype + _true = _api.convert(True) + body = _make.Realize(_buf.op, 0, _dtype, _domain, _true, body) + body = _make.AttrStmt(_buf.op, 'realize_scope', _api.convert(_scope), body) + + for elem in pop_buf: + self.alloc_buffers.pop(elem) + for elem in pop_var: + self.variables.pop(elem) return body @@ -130,13 +148,13 @@ def visit_FunctionDef(self, node): _internal_assert(len(node.args.args) == len(self.args), \ "The number of arguments passed to the \ function should be the same as it is defined!") + if self.func_name is None: + self.func_name = node.name for idx, arg in enumerate(node.args.args): _attr = 'id' if sys.version_info[0] < 3 else 'arg' # To make py2 and 3 compatible self._args[getattr(arg, _attr)] = self.args[idx] res = list_to_block(self.visit, node.body) res = self.wrap_up_realize(node, res) - if self.func_name is None: - self.func_name = node.name return res @@ -145,23 +163,22 @@ def visit_Expr(self, node): def visit_Name(self, node): - _id = node.id - if _id in self._args.keys() and isinstance(self._args[_id], (_expr.Var, _expr.ConstExpr)): - return self._args[_id] - elif _id in self.loops_above.keys(): - return self.loops_above[_id] - _internal_assert(_id not in self._args.keys(), \ - "This id %s should be handled in visit_Subscript!" % _id) - _internal_assert(_id in self.usage.keys(), \ - "This id %s is expected to be a defined variable!" % _id) - # Buffer - if _id in self.alloc_buffers.keys(): - _buf, _ = self.alloc_buffers[_id] - return _make.Call(_buf.dtype, _id, [_api.const(0)], _expr.Call.Halide, _buf.op, 0) - # Compilation time constant - _internal_assert(_id in self.var_consts.keys(), - "This id %s is expected to a compilation time constant!" % _id) - return self.var_consts[_id] + name = node.id + if name in self.loops_above.keys(): + return self.loops_above[name] + elif name in self.variables.keys(): + res = self.variables[name] + if isinstance(res, tuple): + buf = res[0] + if isinstance(node.ctx, ast.Load): + return _make.Call(buf.dtype, buf.name, [_api.const(0)], \ + _expr.Call.Halide, buf.op, buf.value_index) + return buf, [_api.const(0)] + if isinstance(node.ctx, ast.Load): + return res + return None + buf = self._get_buffer_from_id(name) + return buf def visit_Num(self, node): @@ -169,18 +186,36 @@ def visit_Num(self, node): def visit_AugAssign(self, node): - lhs = self.visit(node.target) + buf = self.visit(node.target) rhs = self.visit(node.value) - rhs = HybridParser._binop_maker[type(node.op)](lhs, rhs) - _internal_assert(isinstance(lhs, _expr.Call), \ - "The LHS of an AugAssign is supposed to be a call!") - return _make.Provide(lhs.func, 0, rhs, lhs.args) + if isinstance(buf, tuple): + _internal_assert(len(buf) == 2, "LHS is supposed to be (buf, args)!") + buf, args = buf + else: + args = [_api.const(0)] + _internal_assert(isinstance(buf, Tensor), "LHS is supposed to be Tensor!") + + read = _make.Call(buf.dtype, buf.name, args, _expr.Call.Halide, buf.op, buf.value_index) + value = HybridParser._binop_maker[type(node.op)](read, rhs) + + return _make.Provide(buf.op, 0, value, args) def visit_Assign(self, node): + rhs = self.visit(node.value) + if isinstance(rhs, Operation): + rmap = {} + _internal_assert(len(node.targets) == rhs.num_outputs, \ + "Unable to detuple the outs to targets") + for i in range(rhs.num_outputs): + _internal_assert(isinstance(node.targets[i], ast.Name), + "You should bind a pure name to the tensors") + self.alloc_buffers[node.targets[i].id] = (rhs.output(i), 'global') + rmap[rhs.outputs[i].op] = rhs.output(i) + return replace_io(rhs.body, rmap) + _internal_assert(len(node.targets) == 1, "So far only one-valued assignment is supported!") lhs = node.targets[0] - rhs = self.visit(node.value) if isinstance(rhs, _expr.Expr): rhs = _ir_pass.Simplify(rhs) if isinstance(lhs, ast.Name): @@ -191,65 +226,64 @@ def visit_Assign(self, node): "Loop variable cannot be overwritten!") decl, _, rw = self.usage[lhs] if decl == lhs_: - _internal_assert(lhs not in self.var_consts.keys(), \ - "A constant cannot be overwritten!") - _internal_assert(lhs not in self.alloc_buffers.keys(), \ + _internal_assert(lhs not in self.variables.keys() and + lhs not in self.alloc_buffers.keys(), \ "This value should not be defined before this point!") if isinstance(rhs, tuple): shape, dtype, scope = rhs ph = _api.placeholder(shape, dtype=dtype, name=lhs) - if scope != 'output': - self.alloc_buffers[lhs] = (ph, scope) - else: - self._args[lhs] = ph + self.alloc_buffers[lhs] = (ph, scope) + if scope == 'output': self.outputs.append(lhs) return make_nop() if isinstance(rhs, halide_imm_types) and ast.Store not in rw: - self.var_consts[lhs] = rhs + self.variables[lhs] = rhs else: ph = _api.placeholder((1, ), dtype=rhs.dtype, name=lhs) - self.alloc_buffers[lhs] = (ph, 'global') - if lhs in self.var_consts.keys(): - return make_nop() - _internal_assert(lhs in self.alloc_buffers.keys(), \ - "This variable should be defined before!") - tgt, _ = self.alloc_buffers[lhs] - return _make.Provide(tgt.op, 0, rhs, [_api.const(0, dtype=rhs.dtype)]) + self.variables[lhs] = (ph, 'global') + lhs = self.visit(lhs_) + if lhs is not None: + buf, args = lhs + return _make.Provide(buf.op, 0, rhs, args) + return make_nop() else: - lhs = self.visit(lhs) - _internal_assert(isinstance(lhs, _expr.Call), \ + lhs, args = self.visit(lhs) + _internal_assert(isinstance(lhs, Tensor), \ "An array access's LHS is expected to be a expr.Call!") - #TODO: support slice later - buf = self._get_buffer_from_id(lhs.name, for_provide=True) - return _make.Provide(buf.op, 0, rhs, lhs.args) + res = _make.Provide(lhs.op, lhs.value_index, rhs, args) + return res def visit_Index(self, node): if isinstance(node.value, ast.Tuple): - return [self.visit(i) for i in node.value.elts] + return self.visit(node.value) return [self.visit(node.value)] + def visit_Attribute(self, node): + _internal_assert(isinstance(node.value, ast.Name), \ + "For atrribute access, only both names are supported so far!") + buf = self._get_buffer_from_id(node.value.id) + return getattr(buf, node.attr) + + def visit_Subscript(self, node): args = self.visit(node.slice) if isinstance(node.value, ast.Name): array = node.value.id - _buf = self._get_buffer_from_id(array) - return _make.Call(_buf.dtype, array, args, _expr.Call.Halide, _buf.op, _buf.value_index) - - _internal_assert(isinstance(node.value, ast.Attribute), \ - "Only variable and attribute's subscript supported so far") - _internal_assert(isinstance(node.value.value, ast.Name), \ - "The root of array access is expect to be a id!") - _internal_assert(node.value.attr == "shape", \ - "Attribute access so far only 'shape' is supported!") + buf = self.visit(node.value) + if isinstance(node.ctx, ast.Load): + return _make.Call(buf.dtype, buf.name, args, _expr.Call.Halide, buf.op, buf.value_index) + else: + return buf, args + + shape = self.visit(node.value) _internal_assert(len(args) == 1, "For 'shape' access the argument should be only one!") args = args[0] #TODO: maybe support non-constant value later? _internal_assert(isinstance(args, (_expr.IntImm, _expr.UIntImm)), \ "So far only constant shape access supported!") - buf = self._get_buffer_from_id(node.value.value.id) - return buf.shape[args.value] + return shape[args.value] def visit_With(self, node): @@ -326,12 +360,17 @@ def visit_Call(self, node): # Yet, no function pointer supported _internal_assert(isinstance(node.func, ast.Name), \ "Only id-function function call is supported so far!") - + func_id = node.func.id + args = [self.visit(i) for i in node.args] try: - return getattr(calls, func_id)(self, func_id, node.args) - except Exception as e: - raise e + return getattr(calls, func_id)(func_id, args) + except AttributeError as e: + _internal_assert(func_id in self.symbols.keys(), \ + "The function called is not in the context either!") + outs = self.symbols[func_id](*args) + op = outs.op if isinstance(outs, Tensor) else outs[0].op + return op def visit_For(self, node): @@ -372,12 +411,20 @@ def visit_Return(self, node): _internal_assert(len(set(ids)) == len(ids), "Duplicated tensors in the return tuples") if len(ids) < len(self.outputs): logging.log(logging.CRITICAL, '[Warning] Not all the output buffers returned!') - self.outputs = [self._args[i] for i in ids] + self.outputs = [self.alloc_buffers[i][0] for i in ids] self.returned = True return make_nop() -def parse_python(src, args): + def visit_Tuple(self, node): + return tuple(self.visit(i) for i in node.elts) + + + def visit_Str(self, node): + return node.s + + +def parse_python(src, symbols, args): """The helper function of calling the AST visitor Parameters @@ -385,6 +432,9 @@ def parse_python(src, args): src : str The source code of the function to be parsed. + src : str + The symbol list of the global context of the function. + args : list of Tensors or Vars The argument lists to the function. It is NOT encouraged to write a function without arguments. @@ -396,8 +446,8 @@ def parse_python(src, args): The result Halide IR and the parser class instance. """ root = ast.parse(src) - var_usage = determine_variable_usage(root, args) - parser = HybridParser(args, var_usage) + var_usage = determine_variable_usage(root, args, symbols) + parser = HybridParser(args, var_usage, symbols) parser.parsed_body = parser.visit(root) _internal_assert(parser.returned, 'No valid return found in the function body!') return parser diff --git a/python/tvm/hybrid/util.py b/python/tvm/hybrid/util.py index 78106838f13e..a8eb25bb7111 100644 --- a/python/tvm/hybrid/util.py +++ b/python/tvm/hybrid/util.py @@ -10,6 +10,7 @@ from .. import api as _api from .. import make as _make from .. import expr as _expr +from .. import stmt as _stmt from ..tensor import Tensor @@ -86,3 +87,19 @@ def _restore_runtime(func, intersect): _globals.pop(elem) for k, v in intersect: _globals[k] = v + + +def replace_io(body, rmap): + from .. import ir_pass + + def replace(op): + if isinstance(op, _stmt.Provide) and op.func in rmap.keys(): + buf = rmap[op.func] + return _make.Provide(buf.op, op.value_index, op.value, op.args) + elif isinstance(op, _expr.Call) and op.func in rmap.keys(): + buf = rmap[op.func] + return _make.Call(buf.dtype, buf.name, op.args, _expr.Call.Halide, buf.op, buf.value_index) + return None + + return ir_pass.IRTransform(body, None, replace, ['Provide', 'Call']) + diff --git a/python/tvm/hybrid/var_decl.py b/python/tvm/hybrid/var_decl.py index ea15e7ad1f8f..eb893a7f22a1 100644 --- a/python/tvm/hybrid/var_decl.py +++ b/python/tvm/hybrid/var_decl.py @@ -10,12 +10,13 @@ class PyVariableUsage(ast.NodeVisitor): """The vistor class to determine the declaration, r/w status, and last use of each variable""" #pylint: disable=invalid-name #pylint: disable=missing-docstring - def __init__(self, args): + def __init__(self, args, symbols): self.status = {} self.scope_level = [] self._args = {} self.args = args self.aug_assign_ = False + self.symbols = symbols def visit_FunctionDef(self, node): @@ -43,8 +44,10 @@ def visit_Call(self, node): #No function pointer supported so far _internal_assert(isinstance(node.func, ast.Name), "Function call should be an id") func_id = node.func.id - _internal_assert(func_id in list(HYBRID_GLOBALS.keys()) + ['range', 'max', 'min'], \ - "Function call id not in intrinsics' list") + _internal_assert(func_id in list(HYBRID_GLOBALS.keys()) + \ + ['range', 'max', 'min'] + \ + list(self.symbols.keys()), \ + "Function call id not in intrinsics' list") for elem in node.args: self.visit(elem) @@ -80,8 +83,8 @@ def visit_Name(self, node): self.status[node.id] = (decl, loop, usage) -def determine_variable_usage(root, args): +def determine_variable_usage(root, args, symbols): """The helper function for calling the dedicated visitor.""" - visitor = PyVariableUsage(args) + visitor = PyVariableUsage(args, symbols) visitor.visit(root) return visitor.status From 208b4cfc4904b7a2ba1b487bb90c6a8259f616e7 Mon Sep 17 00:00:00 2001 From: Jian Weng Date: Fri, 14 Dec 2018 00:26:18 -0800 Subject: [PATCH 14/23] fix lint --- python/tvm/hybrid/parser.py | 11 +++++------ python/tvm/hybrid/util.py | 5 +++-- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/python/tvm/hybrid/parser.py b/python/tvm/hybrid/parser.py index 121b1cf9f6de..8ea18a7f350e 100644 --- a/python/tvm/hybrid/parser.py +++ b/python/tvm/hybrid/parser.py @@ -270,12 +270,11 @@ def visit_Attribute(self, node): def visit_Subscript(self, node): args = self.visit(node.slice) if isinstance(node.value, ast.Name): - array = node.value.id buf = self.visit(node.value) if isinstance(node.ctx, ast.Load): - return _make.Call(buf.dtype, buf.name, args, _expr.Call.Halide, buf.op, buf.value_index) - else: - return buf, args + return _make.Call(buf.dtype, buf.name, args, \ + _expr.Call.Halide, buf.op, buf.value_index) + return buf, args shape = self.visit(node.value) _internal_assert(len(args) == 1, "For 'shape' access the argument should be only one!") @@ -360,12 +359,12 @@ def visit_Call(self, node): # Yet, no function pointer supported _internal_assert(isinstance(node.func, ast.Name), \ "Only id-function function call is supported so far!") - + func_id = node.func.id args = [self.visit(i) for i in node.args] try: return getattr(calls, func_id)(func_id, args) - except AttributeError as e: + except AttributeError: _internal_assert(func_id in self.symbols.keys(), \ "The function called is not in the context either!") outs = self.symbols[func_id](*args) diff --git a/python/tvm/hybrid/util.py b/python/tvm/hybrid/util.py index a8eb25bb7111..aa86d55a6fcf 100644 --- a/python/tvm/hybrid/util.py +++ b/python/tvm/hybrid/util.py @@ -90,6 +90,7 @@ def _restore_runtime(func, intersect): def replace_io(body, rmap): + """Replacing tensors usage according to the dict given""" from .. import ir_pass def replace(op): @@ -98,8 +99,8 @@ def replace(op): return _make.Provide(buf.op, op.value_index, op.value, op.args) elif isinstance(op, _expr.Call) and op.func in rmap.keys(): buf = rmap[op.func] - return _make.Call(buf.dtype, buf.name, op.args, _expr.Call.Halide, buf.op, buf.value_index) + return _make.Call(buf.dtype, buf.name, op.args, \ + _expr.Call.Halide, buf.op, buf.value_index) return None return ir_pass.IRTransform(body, None, replace, ['Provide', 'Call']) - From 6cb3953f8b354782acb039af708abeb9f36526a7 Mon Sep 17 00:00:00 2001 From: Jian Weng Date: Fri, 14 Dec 2018 00:30:53 -0800 Subject: [PATCH 15/23] get rid of hybrid codegen --- src/codegen/codegen_hybrid.cc | 0 src/codegen/codegen_hybrid.h | 0 2 files changed, 0 insertions(+), 0 deletions(-) delete mode 100644 src/codegen/codegen_hybrid.cc delete mode 100644 src/codegen/codegen_hybrid.h diff --git a/src/codegen/codegen_hybrid.cc b/src/codegen/codegen_hybrid.cc deleted file mode 100644 index e69de29bb2d1..000000000000 diff --git a/src/codegen/codegen_hybrid.h b/src/codegen/codegen_hybrid.h deleted file mode 100644 index e69de29bb2d1..000000000000 From 9f1571a45d8bc823fdf5a2331451a89b4bf9749a Mon Sep 17 00:00:00 2001 From: Jian Weng Date: Fri, 14 Dec 2018 11:17:35 -0800 Subject: [PATCH 16/23] quick fix?! --- python/tvm/hybrid/parser.py | 11 ++++------- tests/python/unittest/test_hybrid_script.py | 2 +- 2 files changed, 5 insertions(+), 8 deletions(-) diff --git a/python/tvm/hybrid/parser.py b/python/tvm/hybrid/parser.py index 8ea18a7f350e..2c7d3e78a951 100644 --- a/python/tvm/hybrid/parser.py +++ b/python/tvm/hybrid/parser.py @@ -335,13 +335,10 @@ def visit_BoolOp(self, node): _internal_assert(isinstance(node.op, ast.Not), \ "Unary is supposed to be not!") return operator.not_(self.visit(node.values[0])) - elif n == 2: - _internal_assert(isinstance(node.op, (ast.And, ast.Or)), \ - "Binary is supposed to be and/or!") - values = [self.visit(i) for i in node.values] - return HybridParser._binop_maker[type(node.op)](*values) - else: - raise ValueError("This Bool Op is not supported yet!") + _internal_assert(isinstance(node.op, (ast.And, ast.Or)), \ + "Binary is supposed to be and/or!") + values = [self.visit(i) for i in node.values] + return HybridParser._binop_maker[type(node.op)](*values) def visit_UnaryOp(self, node): diff --git a/tests/python/unittest/test_hybrid_script.py b/tests/python/unittest/test_hybrid_script.py index 0d2d2d0e6dff..8c305d8d2ef9 100644 --- a/tests/python/unittest/test_hybrid_script.py +++ b/tests/python/unittest/test_hybrid_script.py @@ -530,7 +530,7 @@ def foo(a): b = output_tensor(a.shape, a.dtype) b[0] = 1.2 for i in range(1, a.shape[0] - 1): - if a[i] * a[i - 1] < a[i] or a[i] * a[i - 1] < a[i - 1]: + if a[i] * a[i - 1] < a[i] or a[i] * a[i - 1] < a[i - 1] or i * a[i] == a[i]: b[i] = a[i] else: b[i] = 0.0 From d182d2c7f33137dc5489011f17f1c4c4f7d9a1ce Mon Sep 17 00:00:00 2001 From: Jian Weng Date: Mon, 17 Dec 2018 10:44:54 -0800 Subject: [PATCH 17/23] fix review? --- python/tvm/hybrid/api.py | 2 +- python/tvm/hybrid/calls.py | 7 ++++--- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/python/tvm/hybrid/api.py b/python/tvm/hybrid/api.py index 0c1fedf4b0b8..d43217ca5dfc 100644 --- a/python/tvm/hybrid/api.py +++ b/python/tvm/hybrid/api.py @@ -30,7 +30,7 @@ def wrapped_func(func, *args, **kwargs): #pylint: disable=missing-docstring for i in args: if isinstance(i, Tensor): input_tensors.append(i) - op = _tvm_internal._HybridOp(parser.func_name, "HybridOp", None, input_tensors, \ + op = _tvm_internal._HybridOp(parser.func_name, "HybridOp", None, input_tensors, parser.outputs, parser.parsed_body) res = [op.output(i) for i in range(len(parser.outputs))] return res[0] if len(res) == 1 else res diff --git a/python/tvm/hybrid/calls.py b/python/tvm/hybrid/calls.py index d094f957bfcb..53473f06c28e 100644 --- a/python/tvm/hybrid/calls.py +++ b/python/tvm/hybrid/calls.py @@ -19,7 +19,7 @@ 'vectorize': For.Vectorized, } -def _range(func_id, args): +def _range(annotation, args): """Handling TVM loop types""" n = len(args) if n == 1: @@ -29,7 +29,7 @@ def _range(func_id, args): low, ext = args[0], args[1] if not Equal(low, _api.const(0, dtype='int32')): ext = ext - low - for_type = LOOP_INTRIN[func_id] + for_type = LOOP_INTRIN[annotation] iter_var = None return iter_var, low, ext, for_type @@ -65,7 +65,8 @@ def _min_max(func_id, args): def _allocate_tensor(func_id, args): - """Handling TVM tensor allocation""" + """Handling TVM tensor allocation. + You may refer hybrid.intrin.allocate for more details.""" n = len(args) _internal_assert(isinstance(_api.convert(args[0]), Array), \ "allocate's first argument should be a tuple of shape!") From ce28aa99421cbe72b944b5ae8495a598927777a6 Mon Sep 17 00:00:00 2001 From: Jian Weng Date: Mon, 17 Dec 2018 10:57:54 -0800 Subject: [PATCH 18/23] still do not know why load crashes --- python/tvm/hybrid/parser.py | 21 +++++++++++---------- 1 file changed, 11 insertions(+), 10 deletions(-) diff --git a/python/tvm/hybrid/parser.py b/python/tvm/hybrid/parser.py index 2c7d3e78a951..c3916abd3116 100644 --- a/python/tvm/hybrid/parser.py +++ b/python/tvm/hybrid/parser.py @@ -4,8 +4,9 @@ import operator import logging import sys -from .util import make_nop, halide_imm_types, is_docstring, _internal_assert, replace_io +from .util import _internal_assert from . import calls +from . import util from .var_decl import determine_variable_usage from ..api import all as _all from ..api import any as _any @@ -17,10 +18,10 @@ def list_to_block(visit, lst): """Convert a list of Python IR nodes to HalideIR Block""" - lst = [visit(stmt) for stmt in lst if not is_docstring(stmt)] - lst = [stmt for stmt in lst if not _ir_pass.Equal(stmt, make_nop())] + lst = [visit(stmt) for stmt in lst if not util.is_docstring(stmt)] + lst = [stmt for stmt in lst if not _ir_pass.Equal(stmt, util.make_nop())] if not lst: - return make_nop() + return util.make_nop() if len(lst) == 1: return lst[0] body = lst[0] @@ -212,7 +213,7 @@ def visit_Assign(self, node): "You should bind a pure name to the tensors") self.alloc_buffers[node.targets[i].id] = (rhs.output(i), 'global') rmap[rhs.outputs[i].op] = rhs.output(i) - return replace_io(rhs.body, rmap) + return util.replace_io(rhs.body, rmap) _internal_assert(len(node.targets) == 1, "So far only one-valued assignment is supported!") lhs = node.targets[0] @@ -235,8 +236,8 @@ def visit_Assign(self, node): self.alloc_buffers[lhs] = (ph, scope) if scope == 'output': self.outputs.append(lhs) - return make_nop() - if isinstance(rhs, halide_imm_types) and ast.Store not in rw: + return util.make_nop() + if isinstance(rhs, util.halide_imm_types) and ast.Store not in rw: self.variables[lhs] = rhs else: ph = _api.placeholder((1, ), dtype=rhs.dtype, name=lhs) @@ -245,7 +246,7 @@ def visit_Assign(self, node): if lhs is not None: buf, args = lhs return _make.Provide(buf.op, 0, rhs, args) - return make_nop() + return util.make_nop() else: lhs, args = self.visit(lhs) _internal_assert(isinstance(lhs, Tensor), \ @@ -305,7 +306,7 @@ def visit_If(self, node): if node.orelse: else_body = list_to_block(self.visit, node.orelse) else: - else_body = make_nop() + else_body = util.make_nop() return _make.IfThenElse(cond, if_body, else_body) @@ -409,7 +410,7 @@ def visit_Return(self, node): logging.log(logging.CRITICAL, '[Warning] Not all the output buffers returned!') self.outputs = [self.alloc_buffers[i][0] for i in ids] self.returned = True - return make_nop() + return util.make_nop() def visit_Tuple(self, node): From 1c54380c20a67cb151b8a29b72bbaaf4a15d3470 Mon Sep 17 00:00:00 2001 From: Jian Weng Date: Tue, 18 Dec 2018 09:51:57 -0800 Subject: [PATCH 19/23] fix runtime import error --- python/tvm/hybrid/calls.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/tvm/hybrid/calls.py b/python/tvm/hybrid/calls.py index 53473f06c28e..8fbf884f693c 100644 --- a/python/tvm/hybrid/calls.py +++ b/python/tvm/hybrid/calls.py @@ -6,7 +6,7 @@ from .. import expr as _expr from .. import make as _make from ..container import Array -from ..ir_pass import Equal +from .. import ir_pass from ..stmt import For from .util import _internal_assert @@ -27,7 +27,7 @@ def _range(annotation, args): else: _internal_assert(n == 2, "A loop intrinsic should only have 1 or 2 arguments!") low, ext = args[0], args[1] - if not Equal(low, _api.const(0, dtype='int32')): + if not ir_pass.Equal(low, _api.const(0, dtype='int32')): ext = ext - low for_type = LOOP_INTRIN[annotation] iter_var = None From 2460558aca1b9d2719c714a0ab0e2e8fb6f8d0eb Mon Sep 17 00:00:00 2001 From: Jian Weng Date: Tue, 18 Dec 2018 09:58:03 -0800 Subject: [PATCH 20/23] we do not need int32 --- python/tvm/hybrid/parser.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/tvm/hybrid/parser.py b/python/tvm/hybrid/parser.py index c3916abd3116..26b0e141d0db 100644 --- a/python/tvm/hybrid/parser.py +++ b/python/tvm/hybrid/parser.py @@ -378,7 +378,7 @@ def visit_For(self, node): if iter_var is None: _internal_assert(for_type is not None, "The loop bind function parse error!") offset = iter_var = _api.var(_name) - if not _ir_pass.Equal(low, _api.const(0, dtype='int32')): + if not _ir_pass.Equal(low, _api.const(0)): offset = iter_var + low self.loops_above[_name] = offset else: @@ -389,7 +389,7 @@ def visit_For(self, node): if for_type is None: res = _make.AttrStmt(iter_var, 'thread_extent', ext, _body) else: - res = _make.For(iter_var, _api.const(0, dtype='int32'), ext, for_type, 0, _body) + res = _make.For(iter_var, _api.const(0), ext, for_type, 0, _body) self.loops_above.pop(_name) return res From f0dad1e03e2987a2101f8a9d799af007930c5061 Mon Sep 17 00:00:00 2001 From: Jian Weng Date: Tue, 18 Dec 2018 10:51:33 -0800 Subject: [PATCH 21/23] fix the string thing... --- python/tvm/hybrid/calls.py | 2 +- tests/python/unittest/test_hybrid_script.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/python/tvm/hybrid/calls.py b/python/tvm/hybrid/calls.py index 8fbf884f693c..6118185d36e3 100644 --- a/python/tvm/hybrid/calls.py +++ b/python/tvm/hybrid/calls.py @@ -41,7 +41,7 @@ def bind(func_id, args): """Handling TVM thread binding""" _internal_assert(func_id == "bind", "This function cannot be directly invoked!") _internal_assert(len(args) == 2, "A loop bind should only have 2 arguments!") - _internal_assert(isinstance(args[0], ast.Str), \ + _internal_assert(isinstance(args[0], str), \ "A loop bind's first argument should be a string!") iter_var = thread_axis(args[0]) low, ext = _api.const(0, dtype='int32'), args[1] diff --git a/tests/python/unittest/test_hybrid_script.py b/tests/python/unittest/test_hybrid_script.py index 8c305d8d2ef9..f87c75f7929d 100644 --- a/tests/python/unittest/test_hybrid_script.py +++ b/tests/python/unittest/test_hybrid_script.py @@ -270,7 +270,7 @@ def test_bind(): return @script def vec_add(a, b): - c = output_tensor((1000, ), dtype='float32') + c = output_tensor((1000, ), 'float32') for tx in bind('threadIdx.x', 1000): c[tx] = a[tx] + b[tx] return c From 16965c778423886af8db72eeba29442d1bc7935e Mon Sep 17 00:00:00 2001 From: Jian Weng Date: Wed, 19 Dec 2018 12:22:20 -0800 Subject: [PATCH 22/23] fix gpu thing --- python/tvm/hybrid/calls.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/python/tvm/hybrid/calls.py b/python/tvm/hybrid/calls.py index 53473f06c28e..4ca3e972a5d0 100644 --- a/python/tvm/hybrid/calls.py +++ b/python/tvm/hybrid/calls.py @@ -41,10 +41,10 @@ def bind(func_id, args): """Handling TVM thread binding""" _internal_assert(func_id == "bind", "This function cannot be directly invoked!") _internal_assert(len(args) == 2, "A loop bind should only have 2 arguments!") - _internal_assert(isinstance(args[0], ast.Str), \ + _internal_assert(isinstance(args[0], str), \ "A loop bind's first argument should be a string!") - iter_var = thread_axis(args[0]) - low, ext = _api.const(0, dtype='int32'), args[1] + iter_var = _api.thread_axis(args[0]) + low, ext = _api.const(0), args[1] for_type = None return iter_var, low, ext, for_type From 2215d94bad0046b59ce8f09194b0a0a55fdafbc1 Mon Sep 17 00:00:00 2001 From: Jian Weng Date: Wed, 19 Dec 2018 12:59:31 -0800 Subject: [PATCH 23/23] fix lint --- python/tvm/hybrid/calls.py | 1 - 1 file changed, 1 deletion(-) diff --git a/python/tvm/hybrid/calls.py b/python/tvm/hybrid/calls.py index 1184d84aa932..730b56f58bd2 100644 --- a/python/tvm/hybrid/calls.py +++ b/python/tvm/hybrid/calls.py @@ -1,7 +1,6 @@ """Intrinsics of TVM-Python Hybrid Script for Python compilation time semantic support.""" -import ast from .. import api as _api from .. import expr as _expr from .. import make as _make