minor fix

Author: kurisu6912

testing/python/transform/test_tilelang_transform_layout_inference.py

@@ -16,6 +16,7 @@ def test_loop_tail_split(block_M, block_N, block_K, threads, vec_load_b, dtype):
     K = tvm.te.var("k")
 
     def before():
+
         @T.prim_func
         def main(B: T.Tensor((K, N), dtype),):
             with T.Kernel(T.ceildiv(N, block_N), threads=threads) as (bx):
@@ -35,9 +36,11 @@ def main(B: T.Tensor((K, N), dtype),):
                                            t // (block_N // vec_load_b), bx * block_N + t %
                                            (block_N // vec_load_b) * (block_N // vec_load_b) + vec],
                                          T.float16(0))
+
         return tvm.IRModule({'main': main})
 
     def after():
+
         @T.prim_func
         def main(B: T.Tensor((K, N), dtype),):
             with T.Kernel(T.ceildiv(N, block_N), threads=threads) as (bx):
@@ -73,6 +76,7 @@ def main(B: T.Tensor((K, N), dtype),):
                                                t // (block_N // vec_load_b),
                                                bx * block_N + t % (block_N // vec_load_b) *
                                                (block_N // vec_load_b) + vec], T.float16(0))
+
         return tvm.IRModule({'main': main})
 
     with tvm.target.Target(auto_target):

testing/python/transform/test_tilelang_transform_lower_tile_op.py

@@ -16,15 +16,18 @@ def test_loop_tail_split(block_M, block_N, block_K, threads, vec_load_b, dtype):
     K = tvm.te.var("k")
 
     def before():
+
         @T.prim_func
         def main(B: T.Tensor((K, N), dtype),):
             with T.Kernel(T.ceildiv(N, block_N), threads=threads) as (bx):
                 B_shared = T.alloc_shared((block_K, block_N), dtype)
                 for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=3):
                     T.copy(B[k * block_K, bx * block_N], B_shared)
+
         return tvm.IRModule({'main': main})
 
     def after():
+
         @T.prim_func
         def main(B: T.Tensor((K, N), dtype),):
             with T.Kernel(T.ceildiv(N, block_N), threads=threads) as (bx):
@@ -60,6 +63,7 @@ def main(B: T.Tensor((K, N), dtype),):
                                                t // (block_N // vec_load_b),
                                                bx * block_N + t % (block_N // vec_load_b) *
                                                (block_N // vec_load_b) + vec], T.float16(0))
+
         return tvm.IRModule({'main': main})
 
     with tvm.transform.PassContext():