fix vm test

apache · Apr 2, 2020 · bac5c24 · bac5c24
1 parent 84674f2
commit bac5c24
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 4 deletions.
diff --git a/tests/python/unittest/test_runtime_extension.py b/tests/python/unittest/test_runtime_extension.py
@@ -39,8 +39,8 @@ def test_dltensor_compatible():
         A[i + 1] = A[i] + 1
     stmt = ib.get()
     fapi = tvm.tir.ir_pass.MakeAPI(stmt, "arange", [Ab], 0, True)
-    fapi = tvm.tir.ir_pass.LowerTVMBuiltin(fapi)
     mod = tvm.testing.LoweredFuncsToIRModule([fapi])
+    mod = tvm.tir.transform.LowerTVMBuiltin()(mod)
     f = tvm.target.codegen.build_module(mod, "stackvm")
     a = tvm.nd.array(np.zeros(10, dtype=dtype))
     aview = MyTensorView(a)

diff --git a/tests/python/unittest/test_target_codegen_vm_basic.py b/tests/python/unittest/test_target_codegen_vm_basic.py
@@ -58,7 +58,6 @@ def test_stack_vm_loop():
 
     stmt = ib.get()
     fapi = tvm.tir.ir_pass.MakeAPI(stmt, "ramp", [Ab], 0, True)
-    fapi = tvm.tir.ir_pass.LowerTVMBuiltin(fapi)
     a = tvm.nd.array(np.zeros(10, dtype=dtype))
     def check(f):
         f(a)
@@ -81,7 +80,6 @@ def test_stack_vm_cond():
 
     stmt = ib.get()
     fapi = tvm.tir.ir_pass.MakeAPI(stmt, "test", [Ab], 0, True)
-    fapi = tvm.tir.ir_pass.LowerTVMBuiltin(fapi)
     def check(f):
         a = tvm.nd.array(np.zeros(10, dtype=dtype))
         f(a)
@@ -101,7 +99,6 @@ def test_vm_parallel():
         A[i] = A[i] + 1
     stmt = ib.get()
     fapi = tvm.tir.ir_pass.MakeAPI(stmt, "ramp", [Ab], 0, True)
-    fapi = tvm.tir.ir_pass.LowerTVMBuiltin(fapi)
     def check(f):
         a = tvm.nd.array(np.zeros(10, dtype=dtype))
         f(a)