mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-11-01 08:23:24 +08:00
36 lines
868 B
Python
36 lines
868 B
Python
import tilelang
|
|
from tilelang import tvm as tvm
|
|
from tvm.ir import IRModule
|
|
import tilelang.testing
|
|
import tilelang.language as T
|
|
|
|
|
|
def merge_if_test():
|
|
@T.prim_func
|
|
def main():
|
|
A = T.alloc_fragment((1,), "float16")
|
|
B = T.alloc_fragment((1,), "float16")
|
|
C = T.alloc_fragment((1,), "float16")
|
|
D = T.alloc_fragment((1,), "float16")
|
|
if A[0] == 0:
|
|
A[0] = 0
|
|
if B[0] == 0:
|
|
B[0] = 0
|
|
if C[0] == 0:
|
|
C[0] = 0
|
|
if D[0] == 0:
|
|
D[0] = 0
|
|
|
|
return main
|
|
|
|
|
|
def test_merge_if():
|
|
func = merge_if_test()
|
|
original_module = IRModule.from_expr(func)
|
|
transformed = tilelang.transform.MergeIfStmt()(original_module)
|
|
tvm.ir.assert_structural_equal(original_module["main"], transformed["main"], True)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|