[CI] Fix CI Script and Broken Tests

Co-authored-by: Shengjie Liu <Shengjie.Liu@armchina.com>

python/tvm/script/parser/tir/parser.py

@@ -21,6 +21,7 @@
 from typing import Any
 
 import tvm
+from tvm import relax, relay
 from tvm.ir import GlobalVar, PrimType
 from tvm.tir import Buffer, IterVar, PrimExpr, Var
 
@@ -454,7 +455,7 @@ def visit_expr_stmt(self: Parser, node: doc.Expr) -> None:
         T.evaluate(res)
     elif isinstance(res, (int, bool)):
         T.evaluate(tvm.tir.const(res))
-    elif isinstance(res, (tvm.relay.Call, tvm.relax.Call)) and not res.args:
+    elif isinstance(res, (relay.Call, relax.Call)) and not res.args:
         # Using GlobalVar.__call__ with no arguments is ambiguous, as
         # each IR has a different function Call representation.  If
         # this occurs, convert to the TIR representation.

src/tir/ir/tir_visitor_with_path.cc

@@ -318,7 +318,17 @@ void TIRVisitorWithPath::VisitStmt_(const BlockNode* op, ObjectPath path) {
     for (size_t i = 0; i < op->match_buffers.size(); i++) {
       auto buf = op->match_buffers[i]->buffer;
       auto buffer_path = match_path->ArrayIndex(i)->Attr("buffer");
+      auto buffer_strides_path = buffer_path->Attr("strides");
       context.push_back(WithDef(buf->data, buffer_path->Attr("data")));
+      // Define buffer strides and elem_offset if they are vars
+      if (const auto* v = buf->elem_offset.as<VarNode>()) {
+        context.push_back(WithDef(GetRef<Var>(v), buffer_path->Attr("elem_offset")));
+      }
+      for (size_t i = 0; i < buf->strides.size(); ++i) {
+        if (const auto* v = buf->strides.as<VarNode>()) {
+          context.push_back(WithDef(GetRef<Var>(v), buffer_strides_path->ArrayIndex(i)));
+        }
+      }
       context.push_back(WithDef(buf, buffer_path));
     }
   }

tests/python/codegen/test_target_codegen_cuda.py

@@ -1069,6 +1069,8 @@ def check_cuda(n, lanes):
     check_cuda(64, 2)
 
 
+@tvm.testing.requires_gpu
+@tvm.testing.requires_cuda
 def test_cuda_thread_sync_inside_condition():
     @T.prim_func
     def func1(A: T.Buffer((4, 4), "float32")) -> None:

tests/python/meta_schedule/test_meta_schedule_postproc_rewrite_tensorize.py

@@ -296,7 +296,7 @@ def main(
                                     W_shared[j, k_o * 4 : k_o * 4 + 4],
                                 )
                                 T.writes(compute_local[i, j])
-                                T.block_attr({"meta_schedule.auto_tensorize": "dp4a"})
+                                T.block_attr({"meta_schedule.auto_tensorize": "dp4a_s8s8s32"})
                                 with T.init():
                                     with T.block("compute_init"):
                                         T.reads()

tests/python/runtime/test_runtime_trace.py

@@ -45,6 +45,7 @@ def check_assign(dtype):
             x.shape, lambda i, j, k: tvm.tir.trace([y[i][j][k]], "tvm.tir.trace_callback2")
         )
         s = te.create_schedule(z.op)
+        te.create_prim_func([x, y, z]).show()
         f = tvm.build(s, [x, y, z], "llvm")
 
         xnd = tvm.nd.array(np.ones((n, n, n), dtype=x.dtype))
@@ -212,10 +213,4 @@ def check_assign(dtype):
 
 
 if __name__ == "__main__":
-    test_trace_expr_assign()
-    test_trace_expr_sum_generated()
-    test_trace_expr_sum_custom()
-    test_trace_expr_sum_args()
-    test_trace_default_action()
-    test_trace_can_change_traced_value_int()
-    test_trace_can_change_traced_value_float()
+    tvm.testing.main()

tests/python/te/test_te_create_primfunc.py

@@ -244,8 +244,8 @@ def tir_extern(a: T.handle, b: T.handle, c: T.handle) -> None:
     C = T.match_buffer(c, (128, 128), elem_offset=off3)
     # body
     with T.block("C"):
-        T.reads([A[0:128, 0:128], B[0:128, 0:128]])
-        T.writes([C[0:128, 0:128]])
+        T.reads()
+        T.writes()
         T.evaluate(
             T.tvm_call_packed(
                 "tvm.contrib.cblas.matmul",
@@ -778,8 +778,8 @@ def tir_extern(var_A: T.handle, var_B: T.handle, var_P: T.handle, var_C: T.handl
         P = T.match_buffer(var_P, [1], dtype="float32", offset_factor=1)
         C = T.match_buffer(var_C, [128, 128], dtype="float32", offset_factor=1)
         with T.block("C"):
-            T.reads(A[0:128, 0:128], B[0:128, 0:128], P[0])
-            T.writes(C[0:128, 0:128])
+            T.reads()
+            T.writes()
             T.call_extern("myfunc", A.data, B.data, C.data, P[0], dtype="")
 
     _check_workload(te_extern, tir_extern)

tests/python/tir-schedule/test_tir_schedule_rfactor.py

@@ -1582,6 +1582,7 @@ def test_reduction_rfactor_argmax_body_bufferstore_value_not_var():
         s.rfactor(ki, 1)
 
 
+@pytest.mark.xfail(reason="The input IR is not well-formed")
 def test_reduction_rfactor_argmax_body_bufferstore_value_unbound_var():
     s = tir.Schedule(argmax_split_body_bufferstore_value_unbound_var, debug_mask="all")
     argmax = s.get_block("argmax")

tests/python/tir-schedule/test_tir_schedule_transform_layout.py

@@ -739,7 +739,10 @@ def expected(A: T.Buffer((14, 32), "int32")):
                 with T.init():
                     B[vi, vj] = T.if_then_else(vi == 3 and 2 <= vj, 0, 0, dtype="int32")
                 B[vi, vj] = T.if_then_else(
-                    vi == 3 and 2 <= vj, 0, B[vi, vj] + A[vi * 4 + vj, vk], dtype="int32"
+                    vi == 3 and 2 <= vj,
+                    0,
+                    B[vj // 4 + vi, vj % 4] + A[vi * 4 + vj, vk],
+                    dtype="int32",
                 )
 
 
@@ -764,7 +767,7 @@ def expected(A: T.Buffer((14, 32), "int32")):
             for k in T.serial(32):
                 with T.block("block"):
                     B[i, j] = T.if_then_else(
-                        i == 3 and 2 <= j, 0, B[i, j] + A[i * 4 + j, k], dtype="int32"
+                        i == 3 and 2 <= j, 0, B[j // 4 + i, j % 4] + A[i * 4 + j, k], dtype="int32"
                     )
 
 

tests/scripts/task_python_unittest.sh

@@ -54,7 +54,7 @@ TEST_FILES=(
   "usmp"
 )
 
-for TEST_FILE in ${TEST_FILES}; do
+for TEST_FILE in ${TEST_FILES[@]}; do
     run_pytest ctypes ${TEST_FILE}-0, tests/python/${TEST_FILE}
     run_pytest cython ${TEST_FILE}-1, tests/python/${TEST_FILE}
 done