chore: import upstream snapshot with attribution

tests/python/tirx-base/test_tir_base.py

@@ -0,0 +1,197 @@
+# 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.
+# ruff: noqa: E711, F841
+import itertools
+
+import numpy as np
+import pytest
+
+import tvm
+from tvm import tirx
+from tvm.ir.transform import PassContext
+from tvm.script import tirx as T
+
+
+def build_tir_func(func):
+    func = func.with_attr("global_symbol", "main")
+    pass_ctx = PassContext.current()
+    if pass_ctx.config.get("tirx.noalias", True):
+        func = func.with_attr("tirx.noalias", True)
+    mod = tvm.IRModule({"main": func})
+    func = tvm.compile(mod)
+    return func
+
+
+def test_scalar_add():
+    # All these types should be interchangeable with each other
+    # E.g. float16 + float32 upconverts the float16 --> float32
+    # Meanwhile if an int or float or together the int will be
+    # cast to the float type.
+    lhs_types = ["float32", "float16", "int32", "int64"]
+    rhs_types = ["float32", "float16"]
+    for lhs_type, rhs_type in itertools.product(lhs_types, rhs_types):
+        # Input vars should be float32, we will cast to test for upcasting between them
+        lhs_input = tirx.Var("lhs", "float32")
+        rhs_input = tirx.Var("rhs", "float32")
+        lhs = tirx.Cast(lhs_type, lhs_input)
+        rhs = tirx.Cast(rhs_type, rhs_input)
+        output = lhs + rhs
+        output = tirx.ret(output)
+        output = tirx.Evaluate(output)
+        func = tirx.PrimFunc([lhs_input, rhs_input], output)
+        func = build_tir_func(func)
+        out = func(1.0, 2.0)
+        assert out == 3.0
+
+
+def assignment_helper(store_dtype, value_dtype):
+    store = tirx.Var("store", dtype=store_dtype)
+    value = tirx.Var("value", dtype=value_dtype)
+    tirx.Let(store, value, body=store)
+
+
+def test_fail_implicit_downcasts_same_type():
+    # These lists should be sorted
+    bits = [8, 16, 32, 64]
+    for type in ["float", "int", "uint"]:
+        for i in range(len(bits) - 1):
+            with pytest.raises(RuntimeError):
+                assignment_helper(
+                    store_dtype=f"{type}{bits[i]}", value_dtype=f"{type}{bits[i + 1]}"
+                )
+
+
+def test_cast_between_types():
+    # We should only be able to assign values with the same types
+    bits = [16, 32]
+    types = ["float", "int", "uint"]
+    for store_type, store_bits, value_type, value_bits in itertools.product(
+        types, bits, types, bits
+    ):
+        store_dtype = f"{store_type}{store_bits}"
+        value_dtype = f"{value_type}{value_bits}"
+        if store_dtype == value_dtype:
+            assignment_helper(store_dtype, value_dtype)
+        else:
+            # TODO: we might want to allow casts between uint and int types
+            with pytest.raises(RuntimeError):
+                assignment_helper(store_dtype, value_dtype)
+
+
+def test_ret_const():
+    a = tirx.const(0)
+    b = tirx.ret(a)
+    b = tirx.Evaluate(b)
+    func = tirx.PrimFunc([], b)
+    func = build_tir_func(func)
+    out = func()
+    assert out == 0
+
+
+def test_control_flow_jump():
+    @T.prim_func(s_tir=True)
+    def func(a: T.float32, b: T.float32):
+        if True:
+            T.evaluate(T.ret(a))
+        T.evaluate(T.ret(b))
+
+    func = build_tir_func(func)
+    out = func(1.0, 2.0)
+    assert out == 1.0
+
+
+def test_break_loop():
+    @T.prim_func(s_tir=True)
+    def func(In: T.Buffer((2,), "int32"), Out: T.Buffer((2,), "int32")):
+        Out[0] = 0
+        Out[1] = 1
+        for i in range(10):
+            for j in range(10):
+                if i * 10 + j == In[0]:
+                    Out[0] = i + j
+                    break
+            if Out[0] > 0:
+                break
+        while Out[1] > 0:
+            Out[1] = Out[1] + 1
+            if Out[1] > In[1]:
+                break
+
+    func = build_tir_func(func)
+    a = np.asarray([49, 8], "int32")
+    b = np.zeros([2], "int32")
+    if not hasattr(b, "__dlpack__"):
+        return
+    func(a, b)
+    assert b[0] == 13
+    assert b[1] == 9
+
+
+def test_continue_loop():
+    @T.prim_func(s_tir=True)
+    def func(Out: T.Buffer((2,), "int32")):
+        T.func_attr({"global_symbol": "main"})
+        Out[0] = 0
+        Out[1] = 0
+        for i in range(10):
+            for j in range(10):
+                if (i * 10 + j) % 3 != 0:
+                    continue
+                Out[0] = Out[0] + 1
+        k = T.decl_buffer([], "int32")
+        k[()] = 0
+        while k[()] < Out[0]:
+            k[()] = k[()] + 1
+            if k[()] % 6 == 0:
+                Out[1] = Out[1] + 1
+                continue
+
+    func = build_tir_func(func)
+    b = np.zeros([2], "int32")
+    if not hasattr(b, "__dlpack__"):
+        return
+    func(b)
+    assert b[0] == 34
+    assert b[1] == 5
+
+
+def test_exception():
+    with pytest.raises(TypeError):
+        x = tirx.Var(name=1, dtype="int")
+
+
+def test_eq_ops():
+    # NOTE: the `== None` / `!= None` below are intentional and must NOT be
+    # rewritten as `is None` / `is not None`. This test exercises the overloaded
+    # `__eq__` / `__ne__` operators on `IntImm` / `StringImm`; the `is` operators
+    # bypass those overloads and would defeat the test.
+    a = tirx.IntImm("int8", 1)
+    with pytest.raises(ValueError):
+        assert a != None
+    with pytest.raises(ValueError):
+        assert not a == None
+    b = tirx.StringImm("abc")
+    assert b != None
+    assert not b == None
+
+
+if __name__ == "__main__":
+    test_scalar_add()
+    test_ret_const()
+    test_control_flow_jump()
+    test_exception()
+    test_eq_ops()