Function create_test_statements

tests/python/tirx/transform/test_stmt_functor.py:621–695

Create test statements for various TIR constructs.

()

def create_test_statements():
    """Create test statements for various TIR constructs."""
    x = tir.Var("x", "int32")
    tir.Var("y", "int32")

    # IntImm
    int_imm = tir.IntImm("int32", 10)

    # Simple expression
    add_expr = tir.Add(x, int_imm)

    # Evaluate
    evaluate_stmt = tir.Evaluate(add_expr)

    # Bind + SeqStmt (was LetStmt)
    let_stmt = tir.SeqStmt([tir.Bind(x, int_imm), evaluate_stmt])

    # For loop
    for_loop = tir.For(x, 0, 10, tir.ForKind.SERIAL, evaluate_stmt)

    # While loop
    while_loop = tir.While(tir.LT(x, int_imm), evaluate_stmt)

    # Buffer operations
    buffer_var = tir.Var("buf", "handle")
    buffer = tir.decl_buffer((10,), "int32", buffer_var.name)
    buffer_store = tir.BufferStore(buffer, add_expr, [int_imm])

    # Sequence of statements
    seq_stmt = tir.SeqStmt([evaluate_stmt, for_loop])

    # Block with iteration variables
    iter_var = tir.IterVar(Range(0, 10), x, 0)
    block = tir.SBlock([iter_var], [], [], "block", evaluate_stmt)
    block_realize = tir.SBlockRealize([int_imm], tir.IntImm("bool", 1), block)

    # IfThenElse statement
    if_then_else = tir.IfThenElse(tir.LT(x, int_imm), evaluate_stmt, evaluate_stmt)

    # Break and continue statements inside a for loop
    @T.prim_func
    def func(A: T.Buffer((10,), "int32")):
        for x in range(10):
            A[x] = x + 1
            if x == 5:
                break
            continue

    # DeclBuffer
    buffer_decl = tir.DeclBuffer(T.buffer((10,), "int32"), evaluate_stmt)

    # TilePrimitiveCall — extract the TilePrimitiveCall from the kernel body, then wrap in an SBlock
    @T.prim_func
    def op_call(A: T.Buffer((10,), "int32"), B: T.Buffer((10,), "int32")):
        T.device_entry()
        Tx.add(A, B, 1.0)

        # op_call.body is the tirx.device_entry AttrStmt, op_call.body.body is TilePrimitiveCall