Revert "[FRONTEND] [HYBRID] Augmented assign operator supported! (#1459)" (#1466)

This reverts commit 33245b8c09009fa2c080e1f42e786ede745de963.

python/tvm/hybrid/parser.py

@@ -15,7 +15,7 @@ from .. import ir_pass as _ir_pass
 
 def list_to_block(visit, lst):
     """Convert a list of Python IR nodes to HalideIR Block"""
-    lst = [visit(i) for i in lst]
+    lst = list(map(visit, lst))
     lst = [stmt for stmt in lst if not _ir_pass.Equal(stmt, make_nop())]
     if not lst:
         return make_nop()
@@ -162,13 +162,6 @@ class HybridParser(ast.NodeVisitor):
     def visit_Num(self, node):
         return _api.const(node.n)
 
-    def visit_AugAssign(self, node):
-        lhs = self.visit(node.target)
-        rhs = self.visit(node.value)
-        rhs = HybridParser._binop_maker[type(node.op)](lhs, rhs)
-        if not isinstance(lhs, _expr.Call):
-            raise ValueError("The LHS of an AugAssign is supposed to be a call!")
-        return _make.Provide(lhs.func, 0, rhs, lhs.args)
 
     def visit_Assign(self, node):
         if len(node.targets) != 1:

python/tvm/hybrid/var_decl.py

@@ -14,7 +14,6 @@ class PyVariableUsage(ast.NodeVisitor):
         self.scope_level = []
         self._args = {}
         self.args = args
-        self.aug_assign_ = False
 
 
     def visit_FunctionDef(self, node):
@@ -49,12 +48,6 @@ class PyVariableUsage(ast.NodeVisitor):
             self.visit(elem)
 
 
-    def visit_AugAssign(self, node):
-        self.aug_assign_ = True
-        self.generic_visit(node)
-        self.aug_assign_ = False
-
-
     def visit_Name(self, node):
         # If it is from the argument list or loop variable, we do not worry about it!
         if node.id in self._args.keys():
@@ -69,8 +62,6 @@ class PyVariableUsage(ast.NodeVisitor):
         if node.id not in self.status.keys():
             if not isinstance(node.ctx, ast.Store):
                 raise ValueError('In Python, "first store" indicates "declaration"')
-            if self.aug_assign_:
-                raise ValueError('"First store" cannot be an AugAssign')
             self.status[node.id] = (node, self.scope_level[-1], set())
         else:
             decl, loop, usage = self.status[node.id]

tests/python/unittest/test_hybrid_script.py

@@ -38,9 +38,7 @@ def run_and_check(func, args, outs, var_dict={}, target='llvm'):
     module(*nd_args)
 
     for nd, np in to_check:
-        numpy.testing.assert_allclose(nd.asnumpy(), np, rtol=1e-3, atol=1e-3)
-
-    return module
+        numpy.testing.assert_allclose(nd.asnumpy(), np, rtol=1e-5, atol=1e-5)
 
 
 @script
@@ -85,7 +83,7 @@ def test_outer_product():
     func = tvm.lower(ir, [n, m, a, b, c])
     func = tvm.build(func)
 
-    run_and_check(outer_product, [n, m, a, b, c], [c], {n: 99, m: 101})
+    run_and_check(outer_product, [n, m, a, b, c], [c], {n: 999, m: 1001})
 
     for key, _ in HYBRID_GLOBALS.items():
         assert key not in globals().keys()
@@ -167,31 +165,19 @@ def test_fanout():
     run_and_check(fanout, [n, a, b], [b], {n: 10})
 
 
-def test_failure():
-    try:
-        @script
-        def failure():
+@script
+def failure():
     for i in range(1, 100):
         i = 0
-        tvm.hybrid.parse(failure, [])
-    except IOError as err:
-        assert sys.version_info[0] == 2
-        print('[Warning] Case test_failure.0 is skipped by Python2 because "%s"' % str(err))
-    except ValueError as err:
-        assert str(err) == 'You CAN NEVER overwrite a loop variable!'
 
+def test_failure():
     try:
-        @tvm.hybrid.script
-        def augdefine():
-            for i in range(10):
-                es += 0
-        tvm.hybrid.parse(augdefine, [])
+        tvm.hybrid.parse(failure, [])
     except IOError as err:
         assert sys.version_info[0] == 2
-        print('[Warning] Case test_failure.1 is skipped by Python2 because "%s"' % str(err))
-    except ValueError as err:
-        assert str(err) == '"First store" cannot be an AugAssign'
-
+        print('[Warning] Case test_failure is skipped by Python2 because "%s"' % str(err))
+    except Exception as err:
+        assert str(err) == 'You CAN NEVER overwrite a loop variable!'
 
 
 def test_looptype():
@@ -294,7 +280,7 @@ def test_non_zero():
                 s = 0.0
                 for di in range(3):
                     for dj in range(3):
-                        s += a[i-di, j-dj]
+                        s = s + a[i-di, j-dj]
                 b[i-2, j-2] = s / 9.0
     try:
         a = tvm.placeholder((32, 32), 'float32', 'a')
@@ -329,39 +315,29 @@ def test_allocate():
 
     a = tvm.placeholder((32, 32), 'float32', 'a')
     b = tvm.placeholder((30, 30), 'float32', 'b')
+
     run_and_check(blur2d, [a, b], [b])
 
     if tvm.gpu().exist:
         @tvm.hybrid.script
-        def shared_gemm(a, b, c):
-            for io in bind('blockIdx.x', 8):
-                for ii in bind('blockIdx.y', 8):
-                    shared_b = allocate((64, 64), 'float32', 'shared')
-                    for k in range(64):
-                        shared_b[io * 8 + ii, k] = b[io * 8 + ii, k]
-                    for jo in bind('threadIdx.y', 8):
-                        for ji in bind('threadIdx.x', 8):
-                            for k in range(64):
-                                c[io*8+ii, jo*8+ji] += a[io*8+ii, k] * shared_b[k, jo*8+ji]
-
-        a = tvm.placeholder((64, 64), dtype='float32', name='a')
-        b = tvm.placeholder((64, 64), dtype='float32', name='b')
-        c = tvm.placeholder((64, 64), dtype='float32', name='c')
-        module = run_and_check(shared_gemm, [a, b, c], [c], target='cuda')
-        assert "__syncthreads()" in module.imported_modules[0].get_source()
+        def share_vec_add(a, b, c):
+            shared = allocate((256, ), 'float32', 'shared')
+            for i in bind("threadIdx.x", 256):
+                shared[i] = a[i]
+            local = allocate((256, ), 'float32', 'local')
+            for i in bind("threadIdx.x", 256):
+                local[i] = b[i]
+            for i in bind("threadIdx.x", 256):
+                c[i] = shared[i] + local[i]
+
+        a = tvm.placeholder((256, ), dtype='float32', name='a')
+        b = tvm.placeholder((256, ), dtype='float32', name='b')
+        c = tvm.placeholder((256, ), dtype='float32', name='c')
+        run_and_check(share_vec_add, [a, b, c], [c], target='cuda')
     else:
         print('[Warning] No GPU found! Skip shared mem test!')
 
 
-def test_augassign():
-    @tvm.hybrid.script
-    def augassign(a):
-        for i in range(a.shape[0]):
-            a[i] += 1.0
-    a = tvm.placeholder((16, ), dtype='float32', name='a')
-    run_and_check(augassign, [a], [a])
-
-
 if __name__ == "__main__":
     test_outer_product()
     test_fanout()
@@ -372,5 +348,4 @@ if __name__ == "__main__":
     test_math_intrin()
     test_non_zero()
     test_allocate()
-    test_augassign()