fix test

Author: Rachmanino

testing/python/transform/test_tilelang_transform_lower_hopper_intrin.py

@@ -28,18 +28,19 @@ def test_lower_hopper_intrin_barrier():
     def before():
         with T.Kernel(8):
             _ = T.launch_thread("threadIdx.x", 128)
-            T.create_list_of_mbarrier(128, 128, 128, 128)
+            T.call_intrin("handle", tir.op.Op.get("tl.create_list_of_mbarrier"), 128, 128, 128, 128)
 
     @T.prim_func
     def after():
         with T.Kernel(8):
             v_1 = T.launch_thread("threadIdx.x", 128)
-            T.evaluate(tir.Call("handle", "tir.create_barriers", [4]))
-            with T.If(v_1 == 0), T.Then():
-                T.evaluate(tir.Call("handle", "tir.ptx_init_barrier_thread_count", [T.get_mbarrier(0), 128]))
-                T.evaluate(tir.Call("handle", "tir.ptx_init_barrier_thread_count", [T.get_mbarrier(1), 128]))
-                T.evaluate(tir.Call("handle", "tir.ptx_init_barrier_thread_count", [T.get_mbarrier(2), 128]))
-                T.evaluate(tir.Call("handle", "tir.ptx_init_barrier_thread_count", [T.get_mbarrier(3), 128]))
+            mbarrier = T.alloc_barrier([128, 128, 128, 128])
+            with T.If(tir.Call("bool", tir.op.Op.get("tl.tl_shuffle_elect"), [0])), T.Then():
+                T.evaluate(tir.Call("handle", "tir.ptx_init_barrier_thread_count", [T.call_intrin("handle", tir.op.Op.get("tl.get_mbarrier"), 0), 128]))
+                T.evaluate(tir.Call("handle", "tir.ptx_init_barrier_thread_count", [T.call_intrin("handle", tir.op.Op.get("tl.get_mbarrier"), 1), 128]))
+                T.evaluate(tir.Call("handle", "tir.ptx_init_barrier_thread_count", [T.call_intrin("handle", tir.op.Op.get("tl.get_mbarrier"), 2), 128]))
+                T.evaluate(tir.Call("handle", "tir.ptx_init_barrier_thread_count", [T.call_intrin("handle", tir.op.Op.get("tl.get_mbarrier"), 3), 128]))
+            T.evaluate(tir.Call("handle", tir.op.Op.get("tl.ptx_fence_barrier_init"), []))
             T.evaluate(tir.Call("handle", "tir.tvm_storage_sync", ["shared"]))
 
     _check(before, after)

testing/python/transform/test_tilelang_transform_warp_specialized.py

@@ -76,45 +76,45 @@ def after(A: T.Tensor((M, K), dtype), B: T.Tensor((K, N), dtype)):
         A_shared = T.decl_buffer((3, 1, 8, 256), T.float16, scope="shared.dyn")
         B_shared = T.decl_buffer((3, 1, 4, 512), T.float16, scope="shared.dyn")
         C_local = T.decl_buffer((32,), scope="local")
-        T.create_list_of_mbarrier(128, 128, 128, 128, 128, 128)
+        T.call_intrin("handle", tir.op.Op.get("tl.create_list_of_mbarrier"), 128, 128, 128, 128, 128, 128)
         T.attr([128, 128], "kWarpSpecializationScope", 0)
         if v >= 128:
             T.set_max_nreg(24, 0)
             for k in range(16):
-                T.mbarrier_wait_parity(T.get_mbarrier(k % 3 + 3), T.bitwise_xor(k // 3 % 2, 1))
+                T.call_intrin("handle", tir.op.Op.get("tl.mbarrier_wait_parity"), T.call_intrin("handle", tir.op.Op.get("tl.get_mbarrier"), k % 3 + 3), T.bitwise_xor(k // 3 % 2, 1))
                 if v - 128 == 0:
-                    T.mbarrier_expect_tx(T.get_mbarrier(k % 3), 4096)
+                    T.call_intrin("handle", tir.op.Op.get("tl.mbarrier_expect_tx"), T.call_intrin("handle", tir.op.Op.get("tl.get_mbarrier"), k % 3), 4096)
                 if v - 128 == 0:
                     T.tma_load(
                         T.create_tma_descriptor(6, 2, A.data, 512, 512, 2, 1024, 32, 64, 1, 1, 0, 2, 2, 0),
-                        T.get_mbarrier(k % 3),
+                        T.call_intrin("handle", tir.op.Op.get("tl.get_mbarrier"), k % 3),
                         T.tvm_access_ptr(T.type_annotation(T.float16), A_shared.data, k % 3 * 2048, 2048, 2),
                         k * 32,
                         by * 64,
                     )
                 if v - 128 == 0:
-                    T.mbarrier_expect_tx(T.get_mbarrier(k % 3), 4096)
+                    T.call_intrin("handle", tir.op.Op.get("tl.mbarrier_expect_tx"), T.call_intrin("handle", tir.op.Op.get("tl.get_mbarrier"), k % 3), 4096)
                 if v - 128 == 0:
                     T.tma_load(
                         T.create_tma_descriptor(6, 2, B.data, 512, 512, 2, 1024, 64, 32, 1, 1, 0, 3, 2, 0),
-                        T.get_mbarrier(k % 3),
+                        T.call_intrin("handle", tir.op.Op.get("tl.get_mbarrier"), k % 3),
                         T.tvm_access_ptr(T.type_annotation(T.float16), B_shared.data, k % 3 * 2048, 2048, 2),
                         bx * 64,
                         k * 32,
                     )
-                T.evaluate(tir.Call("handle", "tir.ptx_arrive_barrier", [T.get_mbarrier(k % 3)]))
+                T.evaluate(tir.Call("handle", "tir.ptx_arrive_barrier", [T.call_intrin("handle", tir.op.Op.get("tl.get_mbarrier"), k % 3)]))
         else:
             T.set_max_nreg(240, 1)
             for k in range(16):
-                T.mbarrier_wait_parity(T.get_mbarrier(k % 3), k // 3 % 2)
+                T.call_intrin("handle", tir.op.Op.get("tl.mbarrier_wait_parity"), T.call_intrin("handle", tir.op.Op.get("tl.get_mbarrier"), k % 3), k // 3 % 2)
                 T.call_extern(
                     "handle",
                     "tl::gemm_ss<64, 64, 32, 4, 1, 0, 0>",
                     T.tvm_access_ptr(T.type_annotation(T.float16), A_shared.data, k % 3 * 2048, 2048, 1),
                     T.tvm_access_ptr(T.type_annotation(T.float16), B_shared.data, k % 3 * 2048, 2048, 1),
                     T.tvm_access_ptr(T.type_annotation(T.float32), C_local.data, 0, 32, 3),
                 )
-                T.evaluate(tir.Call("handle", "tir.ptx_arrive_barrier", [T.get_mbarrier(k % 3 + 3)]))
+                T.evaluate(tir.Call("handle", "tir.ptx_arrive_barrier", [T.call_intrin("handle", tir.op.Op.get("tl.get_mbarrier"), k % 3 + 3)]))
 
     _check(before, after)