Updated a few additional transform_layout usages from main

apache · May 23, 2022 · a7251f2 · a7251f2
1 parent 90e0798
commit a7251f2
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/tests/python/unittest/test_tir_schedule_tensorize_ldmatrix_mma.py b/tests/python/unittest/test_tir_schedule_tensorize_ldmatrix_mma.py
@@ -177,9 +177,9 @@ def tile_wmma_fragment(block_read, height, width):
     else:
         loop_b = tile_wmma_fragment(B_warp, k_inner, 16)
 
-    sch.transform_layout(A_warp, 0, "write", index_map_A)
-    sch.transform_layout(B_warp, 0, "write", index_map_B)
-    sch.transform_layout(C_warp, 0, "read", index_map_C)
+    sch.transform_layout(A_warp, ("write",0) index_map_A)
+    sch.transform_layout(B_warp, ("write",0) index_map_B)
+    sch.transform_layout(C_warp, ("read",0) index_map_C)
 
     sch.tensorize(loop_a, ldmatrix_a_intrin)
     sch.tensorize(loop_b, ldmatrix_b_intrin)