chore: import upstream snapshot with attribution

tests/python/tvmscript/test_tvmscript_ir_builder_tir.py

@@ -0,0 +1,510 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=invalid-name, missing-docstring
+# ruff: noqa: F401, F841
+"""Unittests for tvm.script.ir_builder.tirx"""
+
+import numpy as np
+import pytest
+
+import tvm
+import tvm.runtime
+import tvm.testing
+from tvm import tirx
+from tvm.ir.base import assert_structural_equal
+from tvm.script.ir_builder import IRBuilder
+from tvm.script.ir_builder import tirx as T
+
+
+def test_ir_builder_tir_primfunc_base():
+    with IRBuilder() as ib:
+        with T.prim_func(s_tir=True):
+            T.evaluate(0)
+
+    # the prim_func generated by IRBuilder
+    prim_func_actual = ib.get()
+
+    # the expected prim_func
+    prim_func_expected = tirx.PrimFunc(
+        params=[],
+        body=tirx.Evaluate(0),
+        ret_type=None,
+        buffer_map=None,
+        attrs=tvm.ir.make_node("ir.DictAttrs", s_tir=True),
+    )
+
+    # Check if the generated ir is expected
+    assert_structural_equal(prim_func_actual, prim_func_expected, map_free_vars=True)
+
+
+def test_ir_builder_tir_primfunc_complete():
+    with IRBuilder() as ib:
+        with T.prim_func(s_tir=True):
+            T.arg("a", T.handle())
+            T.arg("b", T.int64())
+            T.arg("c", T.Buffer((128, 128), "float32"))
+            d = T.arg("d", T.handle())
+            e = T.arg("e", T.Buffer((1024,), "int8"))
+            T.func_attr({"key": "value"})
+            T.func_ret(tvm.ir.PrimType("int64"))
+            buffer_d = T.match_buffer(d, (64, 64), "int64")
+            T.evaluate(0)
+
+    # the prim_func generated by IRBuilder
+    prim_func_actual = ib.get()
+
+    # the expected prim_func
+    c_handle, c_buffer = (
+        tirx.Var("c_handle", tvm.ir.PointerType(tvm.ir.PrimType("void"))),
+        tirx.decl_buffer((128, 128), "float32", name="c", layout=None),
+    )
+    d_handle, d_buffer = (
+        tirx.Var("d", tvm.ir.PointerType(tvm.ir.PrimType("void"))),
+        tirx.decl_buffer((64, 64), "int64", name="d", layout=None),
+    )
+    e_handle, e_buffer = (
+        tirx.Var("e_handle", tvm.ir.PointerType(tvm.ir.PrimType("void"))),
+        tirx.decl_buffer((1024,), "int8", name="e", layout=None),
+    )
+    prim_func_expected = tirx.PrimFunc(
+        params=[
+            tirx.Var("a", tvm.ir.PointerType(tvm.ir.PrimType("void"))),
+            tirx.Var("b", "int64"),
+            c_handle,
+            d_handle,
+            e_handle,
+        ],
+        body=tirx.Evaluate(0),
+        ret_type=tvm.ir.PrimType("int64"),
+        buffer_map={c_handle: c_buffer, d_handle: d_buffer, e_handle: e_buffer},
+        attrs=tvm.ir.make_node("ir.DictAttrs", key="value", s_tir=True),
+    )
+
+    # Check if the generated ir is expected
+    assert_structural_equal(prim_func_actual, prim_func_expected, map_free_vars=True)
+
+
+def test_ir_builder_tir_block_base():
+    with IRBuilder() as ib:
+        with T.sblock("block"):
+            T.evaluate(0)
+
+    # the block generated by IRBuilder
+    block_realize_actual = ib.get()
+
+    # the expected block
+    block_expected = tirx.SBlock(
+        iter_vars=[],
+        reads=[],
+        writes=[],
+        name_hint="block",
+        body=tirx.Evaluate(0),
+        alloc_buffers=None,
+        match_buffers=None,
+        annotations={"tirx.script_parsing_detect_access": tirx.IntImm("int64", 3)},
+    )
+    block_realize_expected = tirx.SBlockRealize(
+        iter_values=[],
+        predicate=True,
+        block=block_expected,
+    )
+
+    # Check if the generated ir is expected
+    assert_structural_equal(block_realize_actual, block_realize_expected, map_free_vars=True)
+
+
+def test_ir_builder_tir_block_complete():
+    with IRBuilder() as ib:
+        a = T.int64()
+        b = T.Buffer((128, 128), "float32")
+        c = T.Buffer((128, 128), "float32")
+        d = T.int32()
+        e = T.Buffer((128, 128), "float32")
+        f = T.int32()
+        with T.sblock("block"):
+            T.where(a > 1)
+            T.reads(b[0:16, 0:16])
+            T.writes(c[d:128, d:128])
+            T.sblock_attr({"key": "value"})
+            T.sblock_alloc_buffer((128, 128), "float32")
+            T.match_buffer(e[0:32, 0:32], (32, 32), "float32")
+            T.axis.spatial(128, f)
+            T.evaluate(0)
+
+    # the block generated by IRBuilder
+    block_realize_actual = ib.get()
+
+    # the expected block
+    var_a = tirx.Var("a", "int64")
+    buffer_b = tirx.decl_buffer((128, 128), "float32", name="b")
+    buffer_c = tirx.decl_buffer((128, 128), "float32", name="c")
+    var_d = tirx.Var("d", "int32")
+    buffer_e = tirx.decl_buffer((128, 128), "float32", name="c")
+    var_f = tirx.Var("f", "int32")
+    block_expected = tirx.SBlock(
+        iter_vars=[tirx.IterVar((0, 128), tirx.Var("", "int32"), iter_type=tirx.IterVar.DataPar)],
+        reads=[buffer_b[0:16, 0:16]],
+        writes=[buffer_c[var_d:128, var_d:128]],
+        name_hint="block",
+        body=tirx.Evaluate(0),
+        alloc_buffers=[tirx.decl_buffer((128, 128), "float32")],
+        match_buffers=[
+            tirx.MatchBufferRegion(tirx.decl_buffer((32, 32), "float32"), buffer_e[0:32, 0:32])
+        ],
+        annotations={"key": "value"},
+    )
+    block_realize_expected = tirx.SBlockRealize(
+        iter_values=[var_f],
+        predicate=var_a > 1,
+        block=block_expected,
+    )
+
+    # Check if the generated ir is expected
+    assert_structural_equal(block_realize_actual, block_realize_expected, map_free_vars=True)
+
+
+def test_ir_builder_tir_axis():
+    with IRBuilder() as ib:
+        a = T.int32()
+        b = T.int32()
+        c = T.int32()
+        d = T.int32()
+        with T.sblock("block"):
+            T.axis.spatial(8, a)
+            T.axis.reduce(16, b)
+            T.axis.scan(32, c)
+            T.axis.opaque(64, d)
+            T.evaluate(0)
+
+    # the block generated by IRBuilder
+    block_realize_actual = ib.get()
+
+    # the expected block
+    var_a = tirx.Var("a", "int32")
+    var_b = tirx.Var("b", "int32")
+    var_c = tirx.Var("c", "int32")
+    var_d = tirx.Var("d", "int32")
+    block_expected = tirx.SBlock(
+        iter_vars=[
+            tirx.IterVar((0, 8), tirx.Var("", "int32"), iter_type=tirx.IterVar.DataPar),
+            tirx.IterVar((0, 16), tirx.Var("", "int32"), iter_type=tirx.IterVar.CommReduce),
+            tirx.IterVar((0, 32), tirx.Var("", "int32"), iter_type=tirx.IterVar.Ordered),
+            tirx.IterVar((0, 64), tirx.Var("", "int32"), iter_type=tirx.IterVar.Opaque),
+        ],
+        reads=[],
+        writes=[],
+        name_hint="block",
+        body=tirx.Evaluate(0),
+        annotations={"tirx.script_parsing_detect_access": tirx.IntImm("int64", 3)},
+    )
+    block_realize_expected = tirx.SBlockRealize(
+        iter_values=[var_a, var_b, var_c, var_d],
+        predicate=True,
+        block=block_expected,
+    )
+
+    # Check if the generated ir is expected
+    assert_structural_equal(block_realize_actual, block_realize_expected, map_free_vars=True)
+
+
+def test_ir_builder_tir_for():
+    with IRBuilder() as ib:
+        with T.serial(128) as a:
+            with T.parallel(64) as b:
+                with T.vectorized(32) as c:
+                    with T.unroll(16) as d:
+                        with T.thread_binding(8, thread="threadIdx.x") as e:
+                            T.evaluate(0)
+
+    # the for generated by IRBuilder
+    for_actual = ib.get()
+
+    # the expected for
+    thread_binding_expected = tirx.For(
+        loop_var=tirx.Var("", "int32"),
+        min=0,
+        extent=8,
+        kind=tirx.ForKind.THREAD_BINDING,
+        body=tirx.Evaluate(0),
+        thread_binding=tirx.IterVar(
+            None, tirx.Var("", "int32"), tirx.IterVar.ThreadIndex, "threadIdx.x"
+        ),
+    )
+    unroll_expected = tirx.For(
+        loop_var=tirx.Var("", "int32"),
+        min=0,
+        extent=16,
+        kind=tirx.ForKind.UNROLLED,
+        body=thread_binding_expected,
+    )
+    vectorized_expected = tirx.For(
+        loop_var=tirx.Var("", "int32"),
+        min=0,
+        extent=32,
+        kind=tirx.ForKind.VECTORIZED,
+        body=unroll_expected,
+    )
+    parallel_expected = tirx.For(
+        loop_var=tirx.Var("", "int32"),
+        min=0,
+        extent=64,
+        kind=tirx.ForKind.PARALLEL,
+        body=vectorized_expected,
+    )
+    for_expected = tirx.For(
+        loop_var=tirx.Var("", "int32"),
+        min=0,
+        extent=128,
+        kind=tirx.ForKind.SERIAL,
+        body=parallel_expected,
+    )
+
+    # Check if the generated ir is expected
+    assert_structural_equal(for_actual, for_expected, map_free_vars=True)
+
+
+def test_ir_builder_tir_for_uint():
+    with IRBuilder() as ib:
+        with T.serial(tirx.const(128, "uint32")) as a:
+            T.evaluate(0)
+
+    # the for generated by IRBuilder
+    for_actual = ib.get()
+
+    for_expected = tirx.For(
+        loop_var=tirx.Var("", "uint32"),
+        min=tirx.const(0, "uint32"),
+        extent=tirx.const(128, "uint32"),
+        kind=tirx.ForKind.SERIAL,
+        body=tirx.Evaluate(0),
+    )
+
+    # Check if the generated ir is expected
+    assert_structural_equal(for_actual, for_expected, map_free_vars=True)
+
+
+def test_ir_builder_tir_assert():
+    with IRBuilder() as ib:
+        with T.Assert(T.int32() == 0, message="a is 0"):
+            T.evaluate(0)
+    # the assert generated by IRBuilder
+    assert_actual = ib.get()
+
+    # AssertStmt is a leaf. The frame emits the assert and then the body stmts as siblings.
+    assert_expected = tirx.SeqStmt(
+        [
+            tirx.AssertStmt(
+                T.int32() == 0,
+                tirx.StringImm("RuntimeError"),
+                [tirx.StringImm("a is 0")],
+            ),
+            tirx.Evaluate(0),
+        ]
+    )
+
+    # Check if the generated ir is expected
+    assert_structural_equal(assert_actual, assert_expected, map_free_vars=True)
+
+
+def test_ir_builder_tir_bind():
+    # Test that T.bind emits a flat Bind statement and returns the Var.
+    with IRBuilder() as ib:
+        v = T.bind(tirx.IntImm("int32", 2))
+    # the let binding generated by IRBuilder
+    let_actual = ib.get()
+
+    # Bind is now flat (no body), so a single Bind stmt is emitted.
+    let_expected = tirx.Bind(T.int32(), tirx.IntImm("int32", 2))
+
+    # Check if the generated ir is expected
+    assert_structural_equal(let_actual, let_expected, map_free_vars=True)
+    # Check that the returned value is a Var
+    assert isinstance(v, tirx.Var)
+
+
+def test_ir_builder_tir_thread():
+    with IRBuilder() as ib:
+        with T.prim_func(s_tir=True):
+            brow = T.env_thread("blockIdx.y")
+            with T.launch_thread(brow, 1):
+                T.evaluate(0)
+
+    # the prim_func generated by IRBuilder
+    ir_actual = ib.get()
+
+    # the expected prim_func
+    iter_var = tirx.IterVar((0, 1), "v", iter_type=1, thread_tag="blockIdx.y")
+    attr_stmt = tirx.AttrStmt(iter_var, "thread_extent", 1, tirx.Evaluate(0))
+    func = tirx.PrimFunc([], attr_stmt).with_attr("s_tir", True)
+
+    # Check if the generated ir is expected
+    assert_structural_equal(ir_actual, func, map_free_vars=True)
+
+
+def test_ir_builder_tir_allocate():
+    with IRBuilder() as ib:
+        with T.prim_func(s_tir=True):
+            T.func_name("test")
+            buf = T.alloc_buffer([10], "float32", scope="local")
+            T.evaluate(1)
+
+    # the allocate generated by IRBuilder
+    ir_actual = ib.get()
+    body = ir_actual.body
+
+    # AllocBuffer is flat: body should be a SeqStmt with [AllocBuffer, Evaluate(1)]
+    assert isinstance(body, tirx.SeqStmt), f"Expected SeqStmt but got {type(body)}"
+    assert len(body) == 2
+    assert isinstance(body[0], tirx.AllocBuffer)
+    assert isinstance(body[1], tirx.Evaluate)
+    assert body[1].value.value == 1
+
+
+def test_ir_builder_tir_while():
+    with IRBuilder() as ib:
+        with T.While(T.int32() > 0):
+            T.evaluate(0)
+
+    # the while generated by IRBuilder
+    ir_actual = ib.get()
+
+    # the expected while
+    ir_expected = tirx.While(tirx.Var("x", "int32") > 0, tirx.Evaluate(0))
+
+    # Check if the generated ir is expected
+    assert_structural_equal(ir_actual, ir_expected, map_free_vars=True)
+
+
+def test_ir_builder_tir_if_then_else():
+    with IRBuilder() as ib:
+        with T.If(T.int32() < 12):
+            with T.Then():
+                T.evaluate(T.int32(0))
+            with T.Else():
+                T.evaluate(T.int32(1))
+
+    # the if_then_else generated by IRBuilder
+    ir_actual = ib.get()
+
+    # the expected if_then_else
+    ir_expected = tirx.IfThenElse(
+        tirx.Var("c", "int32") < 12,
+        tirx.Evaluate(tirx.IntImm("int32", 0)),
+        tirx.Evaluate(tirx.IntImm("int32", 1)),
+    )
+
+    # Check if the generated ir is expected
+    assert_structural_equal(ir_actual, ir_expected, map_free_vars=True)
+
+
+def test_ir_builder_tir_buffer_store():
+    buffer_a = T.Buffer((10, 10), "float32")
+    i = T.int32()
+    with IRBuilder() as ib:
+        T.buffer_store(buffer_a, 0.1, [0, i])
+
+    # the buffer store generated by IRBuilder
+    ir_actual = ib.get()
+
+    # the expected buffer store
+    ir_expected = tirx.BufferStore(buffer_a, 0.1, [0, i])
+
+    # Check if the generated ir is expected
+    assert_structural_equal(ir_actual, ir_expected, map_free_vars=True)
+
+
+def test_ir_builder_tir_buffer_store_scalable_vec():
+    buffer_a = T.Buffer((30,), "float32")
+    value = T.broadcast(0.11, 4 * tvm.tirx.vscale())
+    index = T.ramp(0, 1, 4 * tvm.tirx.vscale())
+
+    with IRBuilder() as ib:
+        T.buffer_store(buffer_a, value, [index])
+
+    # the buffer store generated by IRBuilder
+    ir_actual = ib.get()
+
+    # the expected buffer store
+    ir_expected = tirx.BufferStore(buffer_a, value, [index])
+
+    # Check if the generated ir is expected
+    assert_structural_equal(ir_actual, ir_expected, map_free_vars=True)
+
+
+def test_ir_builder_tir_buffer_store_predicate():
+    buffer_a = T.Buffer((30,), "float32")
+    value = T.broadcast(0.11, T.vscale() * 4)
+    index = T.ramp(0, 1, T.vscale() * 4)
+    predicate = T.broadcast(T.bool(True), T.vscale() * 4)
+
+    with IRBuilder() as ib:
+        T.buffer_store(buffer_a, value, [index], predicate)
+
+    ir_actual = ib.get()
+    ir_expected = tirx.BufferStore(buffer_a, value, [index], predicate)
+    assert_structural_equal(ir_actual, ir_expected, map_free_vars=True)
+
+
+def test_ir_builder_tir_evaluate():
+    with IRBuilder() as ib:
+        T.evaluate(0)
+    # the evaluate generated by IRBuilder
+    eval_actual = ib.get()
+
+    # the expected evaluate
+    eval_expected = tirx.Evaluate(0)
+
+    # Check if the generated ir is expected
+    assert_structural_equal(eval_actual, eval_expected, map_free_vars=True)
+
+
+def test_ir_builder_tir_decl_buffer():
+    with IRBuilder() as ib:
+        with T.prim_func(s_tir=True):
+            T.func_name("test")
+            buf = T.decl_buffer([128, 128], "float32")
+            T.evaluate(1)
+
+    # the decl_buffer generated by IRBuilder
+    ir_actual = ib.get()
+    body = ir_actual.body
+
+    # decl_buffer without data emits AllocBuffer (flat): body should be SeqStmt
+    assert isinstance(body, tirx.SeqStmt), f"Expected SeqStmt but got {type(body)}"
+    assert len(body) == 2
+    assert isinstance(body[0], tirx.AllocBuffer)
+    assert isinstance(body[1], tirx.Evaluate)
+    assert body[1].value.value == 1
+
+
+def test_ir_builder_tir_inline():
+    with IRBuilder() as ib:
+        m, n = T.meta_var(1), T.meta_var(2)
+        a, b = T.meta_var([3, 4])
+        T.evaluate(m.value + n.value + a.value + b.value)
+    # the evaluate generated by IRBuilder
+    eval_actual = ib.get()
+
+    # the expected evaluate
+    eval_expected = tirx.Evaluate(10)
+
+    # Check if the generated ir is expected
+    assert_structural_equal(eval_actual, eval_expected, map_free_vars=True)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()