chore: import upstream snapshot with attribution

tests/python/tirx-base/test_tir_intrin.py

@@ -0,0 +1,413 @@
+# 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: E712, F401
+import ctypes
+import math
+
+import numpy as np
+import pytest
+
+pytest.importorskip("scipy")
+
+import scipy
+
+import tvm
+import tvm.testing
+from tvm import te, tirx, topi
+from tvm.script import tirx as T
+from tvm.support import clang, utils
+
+
+def test_nearbyint():
+    m = te.var(
+        "m",
+    )
+    A = te.placeholder((m,), name="A")
+    A_rounded = te.compute((m,), lambda *i: tvm.tirx.nearbyint(A(*i)), name="A")
+
+    # Convert to TIR and create schedule
+    mod = te.create_prim_func([A, A_rounded])
+    sch = tvm.s_tir.Schedule(mod)
+
+    # Build from scheduled TIR
+    func = tvm.compile(sch.mod, target="llvm")
+
+    dev = tvm.cpu(0)
+    n = 10
+    a = tvm.runtime.tensor(np.random.uniform(high=100, size=n).astype(A.dtype.dtype), dev)
+    a_rounded = tvm.runtime.tensor(np.random.uniform(size=n).astype(A_rounded.dtype.dtype), dev)
+    func(a, a_rounded)
+    # Note that numpys rint rounds to nearest integer with
+    # ties to halfway is broken by rounding to even.
+    # So that 1.5 and 2.5 will round 2.
+    # This is the default rounding mode with libc as well.
+    # However one can set a different rounding mode and in that
+    # case numpy result might differ.
+    tvm.testing.assert_allclose(a_rounded.numpy(), np.rint(a.numpy()))
+
+
+def test_round_ties_to_even():
+    """Test that tir.round uses ties-to-even (banker's rounding) semantics."""
+    m = te.var("m")
+    A = te.placeholder((m,), name="A")
+    A_rounded = te.compute((m,), lambda *i: tvm.tirx.round(A(*i)), name="A")
+
+    mod = te.create_prim_func([A, A_rounded])
+    sch = tvm.s_tir.Schedule(mod)
+    func = tvm.compile(sch.mod, target="llvm")
+
+    dev = tvm.cpu(0)
+    # Midpoint values where ties-to-even and ties-away differ
+    test_values = np.array([0.5, 1.5, 2.5, 3.5, -0.5, -1.5, -2.5, -3.5], dtype="float32")
+    expected = np.array([0.0, 2.0, 2.0, 4.0, 0.0, -2.0, -2.0, -4.0], dtype="float32")
+
+    a = tvm.runtime.tensor(test_values, dev)
+    a_rounded = tvm.runtime.tensor(np.zeros(len(test_values), dtype="float32"), dev)
+    func(a, a_rounded)
+    tvm.testing.assert_allclose(a_rounded.numpy(), expected)
+
+
+def test_round_intrinsics_on_int():
+    i = tvm.tirx.Var("i", "int32")
+    for op in [tvm.tirx.round, tvm.tirx.trunc, tvm.tirx.ceil, tvm.tirx.floor, tvm.tirx.nearbyint]:
+        assert op(tvm.tirx.const(10, "int32")).value == 10
+        assert op(tvm.tirx.const(True, "bool")).value == True
+        assert op(i).same_as(i)
+
+    assert tvm.tirx.isnan(tvm.tirx.const(10, "int32")).value == False
+
+
+def test_unary_intrin():
+    test_funcs = [
+        (tvm.tirx.exp, lambda x: np.exp(x)),
+        (tvm.tirx.exp10, lambda x: np.power(10, x)),
+        (tvm.tirx.log2, lambda x: np.log2(x)),
+        (tvm.tirx.log10, lambda x: np.log10(x)),
+        (tvm.tirx.sinh, lambda x: np.sinh(x)),
+        (tvm.tirx.cosh, lambda x: np.cosh(x)),
+        (tvm.tirx.log1p, lambda x: np.log1p(x)),
+        (tvm.tirx.asin, lambda x: np.arcsin(x)),
+        (tvm.tirx.acos, lambda x: np.arccos(x)),
+        (tvm.tirx.atan, lambda x: np.arctan(x)),
+        (tvm.tirx.asinh, lambda x: np.arcsinh(x)),
+        (tvm.tirx.acosh, lambda x: np.arccosh(x)),
+        (tvm.tirx.atanh, lambda x: np.arctanh(x)),
+        (tvm.tirx.erf, lambda x: scipy.special.erf(x)),
+    ]
+
+    def run_test(tvm_intrin, np_func, atol=1e-5, rtol=1e-5):
+        m = te.var(
+            "m",
+        )
+        A = te.placeholder((m,), name="A")
+        B = te.compute((m,), lambda *i: tvm_intrin(A(*i)), name="B")
+
+        # Convert to TIR and create schedule
+        mod = te.create_prim_func([A, B])
+        sch = tvm.s_tir.Schedule(mod)
+
+        # Build from scheduled TIR
+        func = tvm.compile(sch.mod, target="llvm")
+
+        dev = tvm.cpu(0)
+        n = 10
+        a = tvm.runtime.tensor(np.random.uniform(0.1, 0.5, size=n).astype(A.dtype.dtype), dev)
+        b = tvm.runtime.tensor(np.random.uniform(size=n).astype(A.dtype.dtype), dev)
+        func(a, b)
+        tvm.testing.assert_allclose(b.numpy(), np_func(a.numpy()), atol=atol, rtol=rtol)
+
+        # Out-of-bounds test for asin/acos
+        name = tvm_intrin.__name__
+        if name in ("asin", "acos"):
+            # generate some values outside [-1, 1]
+            n = 8
+            out_np = np.concatenate(
+                [
+                    np.random.uniform(1.1, 2.0, size=n // 2),
+                    np.random.uniform(-2.0, -1.1, size=n // 2),
+                ]
+            ).astype(A.dtype.dtype)
+            a2 = tvm.runtime.tensor(out_np, dev)
+            b2 = tvm.runtime.tensor(np.empty_like(out_np), dev)
+            func(a2, b2)
+            # all outputs should be NaN
+            assert np.all(np.isnan(b2.numpy()))
+        if name == "exp":
+            n = 8
+            out_np = np.random.randint(-20, 20, size=n).astype(A.dtype.dtype)
+            a2 = tvm.runtime.tensor(out_np, dev)
+            b2 = tvm.runtime.tensor(np.empty_like(out_np), dev)
+            func(a2, b2)
+            assert b2.numpy().dtype == np.float32
+            # Verify correctness against NumPy exp
+            expected = np.exp(out_np.astype(np.float32))
+            tvm.testing.assert_allclose(b2.numpy(), expected, rtol=1e-5, atol=1e-5)
+
+    for func in test_funcs:
+        atol = rtol = 1e-3 if func[0].__name__ in ["asin", "acos", "atan"] else 1e-5
+        run_test(*func, atol, rtol)
+
+
+def test_asin_acos_boundary_values():
+    """Test asin and acos with boundary values and threshold switching."""
+    test_funcs = [
+        (tvm.tirx.asin, lambda x: np.arcsin(x)),
+        (tvm.tirx.acos, lambda x: np.arccos(x)),
+    ]
+
+    def run_test(tvm_intrin, np_func):
+        m = te.var("m")
+        A = te.placeholder((m,), name="A")
+        B = te.compute((m,), lambda *i: tvm_intrin(A(*i)), name="B")
+
+        mod = te.create_prim_func([A, B])
+        sch = tvm.s_tir.Schedule(mod)
+        func = tvm.compile(sch.mod, target="llvm")
+
+        dev = tvm.cpu(0)
+
+        # Test boundary values: ±1.0 (should use system library)
+        boundary_values = np.array([1.0, -1.0], dtype=np.float32)
+        a1 = tvm.runtime.tensor(boundary_values, dev)
+        b1 = tvm.runtime.tensor(np.empty_like(boundary_values), dev)
+        func(a1, b1)
+        tvm.testing.assert_allclose(b1.numpy(), np_func(boundary_values), atol=1e-5, rtol=1e-5)
+
+        # Test values at threshold: ±0.5 (should use system library)
+        threshold_values = np.array([0.5, -0.5], dtype=np.float32)
+        a2 = tvm.runtime.tensor(threshold_values, dev)
+        b2 = tvm.runtime.tensor(np.empty_like(threshold_values), dev)
+        func(a2, b2)
+        tvm.testing.assert_allclose(b2.numpy(), np_func(threshold_values), atol=1e-4, rtol=1e-4)
+
+        # Test values just below threshold: ±0.49 (should use Taylor series)
+        below_threshold_values = np.array([0.49, -0.49, 0.3, -0.3, 0.0], dtype=np.float32)
+        a3 = tvm.runtime.tensor(below_threshold_values, dev)
+        b3 = tvm.runtime.tensor(np.empty_like(below_threshold_values), dev)
+        func(a3, b3)
+        tvm.testing.assert_allclose(
+            b3.numpy(), np_func(below_threshold_values), atol=1e-3, rtol=1e-3
+        )
+
+        # Test out-of-domain values: should return NaN
+        out_of_domain = np.array([1.1, -1.1, 2.0, -2.0], dtype=np.float32)
+        a4 = tvm.runtime.tensor(out_of_domain, dev)
+        b4 = tvm.runtime.tensor(np.empty_like(out_of_domain), dev)
+        func(a4, b4)
+        assert np.all(np.isnan(b4.numpy())), "Out-of-domain inputs should return NaN"
+
+    for func in test_funcs:
+        run_test(*func)
+
+
+def test_binary_intrin():
+    test_funcs = [
+        (tvm.tirx.atan2, lambda x1, x2: np.arctan2(x1, x2)),
+        (tvm.tirx.nextafter, lambda x1, x2: np.nextafter(x1, x2)),
+        (tvm.tirx.copysign, lambda x1, x2: np.copysign(x1, x2)),
+        (tvm.tirx.hypot, lambda x1, x2: np.hypot(x1, x2)),
+    ]
+
+    def run_test(tvm_intrin, np_func):
+        m = te.var(
+            "m",
+        )
+        A = te.placeholder((m,), name="A")
+        B = te.placeholder((m,), name="B")
+        C = te.compute((m,), lambda *i: tvm_intrin(A(*i), B(*i)), name="C")
+
+        # Convert to TIR and create schedule
+        mod = te.create_prim_func([A, B, C])
+        sch = tvm.s_tir.Schedule(mod)
+
+        # Build from scheduled TIR
+        func = tvm.compile(sch.mod, target="llvm")
+
+        dev = tvm.cpu(0)
+        n = 10
+        a = tvm.runtime.tensor(np.random.uniform(0, 1, size=n).astype(A.dtype.dtype), dev)
+        b = tvm.runtime.tensor(np.random.uniform(0, 1, size=n).astype(B.dtype.dtype), dev)
+        c = tvm.runtime.tensor(np.random.uniform(size=n).astype(A.dtype.dtype), dev)
+        func(a, b, c)
+        tvm.testing.assert_allclose(c.numpy(), np_func(a.numpy(), b.numpy()), atol=1e-5, rtol=1e-5)
+
+    for func in test_funcs:
+        run_test(*func)
+
+
+def test_ldexp():
+    m = te.var(
+        "m",
+    )
+    A = te.placeholder((m,), name="A")
+    B = te.placeholder((m,), name="B", dtype="int32")
+    C = te.compute((m,), lambda *i: tvm.tirx.ldexp(A(*i), B(*i)), name="C")
+
+    # Convert to TIR and create schedule
+    mod = te.create_prim_func([A, B, C])
+    sch = tvm.s_tir.Schedule(mod)
+
+    # Build from scheduled TIR
+    func = tvm.compile(sch.mod, target="llvm")
+
+    dev = tvm.cpu(0)
+    n = 10
+    a = tvm.runtime.tensor(np.random.uniform(0, 1, size=n).astype(A.dtype.dtype), dev)
+    b = tvm.runtime.tensor(np.random.randint(0, 5, size=n).astype(B.dtype.dtype), dev)
+    c = tvm.runtime.tensor(np.random.uniform(size=n).astype(A.dtype.dtype), dev)
+    func(a, b, c)
+    tvm.testing.assert_allclose(c.numpy(), np.ldexp(a.numpy(), b.numpy()), atol=1e-5, rtol=1e-5)
+
+
+dtype = tvm.testing.parameter("int32", "int64")
+
+
+@pytest.mark.parametrize(
+    "target",
+    ["llvm", pytest.param({"kind": "vulkan", "from_device": 0}, marks=pytest.mark.gpu)],
+)
+def test_clz(target, dtype):
+    if not tvm.testing.device_enabled(target):
+        pytest.skip(f"{target} not enabled")
+    target = tvm.target.Target(target)
+    if (
+        target.kind.name == "vulkan"
+        and dtype == "int64"
+        and not target.attrs.get("supports_int64", False)
+    ):
+        pytest.xfail("Vulkan target does not support Int64 types")
+
+    def clz_np(x, dtype):
+        ceil_log2 = np.ceil(np.log2(x)).astype(dtype)
+        bits = int(dtype[-2:])
+        clz = bits - ceil_log2
+        clz[np.bitwise_and(x, x - 1) == 0] -= 1
+        return clz
+
+    m = te.var("m")
+    A = te.placeholder((m,), name="A", dtype=dtype)
+    B = te.compute((m,), lambda *i: tvm.tirx.clz(A(*i)), name="B")
+
+    # Convert to TIR and create schedule
+    mod = te.create_prim_func([A, B])
+    sch = tvm.s_tir.Schedule(mod)
+
+    # Apply scheduling primitives if target is Vulkan
+    if target.kind.name == "vulkan":
+        block = sch.get_sblock("B")
+        loop = sch.get_loops(block)[0]
+        bx, tx = sch.split(loop, factors=[None, 64])
+        sch.bind(bx, "blockIdx.x")
+        sch.bind(tx, "threadIdx.x")
+
+    # Build from scheduled TIR
+    func = tvm.compile(sch.mod, target=target)
+
+    def run_and_check():
+        dev = tvm.device(target.kind.name)
+        n = 10
+        highs = [10, 100, 1000, 10000, 100000, 1000000]
+
+        if dtype == "int64":
+            highs.append((1 << 63) - 1)
+
+        for high in highs:
+            a_np = np.random.randint(1, high=high, size=(n,), dtype=dtype)
+            a = tvm.runtime.tensor(a_np, dev)
+            b = tvm.runtime.tensor(np.zeros((n,)).astype("int32"), dev)
+            func(a, b)
+            ref = clz_np(a_np, dtype)
+            np.testing.assert_equal(b.numpy(), ref)
+
+    if target.kind.name == "llvm":
+        run_and_check()
+    else:
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@tvm.script.ir_module
+class Module:
+    @T.prim_func(s_tir=True)
+    def test_tir_fma(A: T.handle, B: T.handle, C: T.handle, d: T.handle) -> None:
+        # function attr dict
+        T.func_attr({"global_symbol": "test_fma", "tirx.noalias": True})
+        n = T.int32()
+        stride = T.int32()
+        stride_1 = T.int32()
+        stride_2 = T.int32()
+        stride_3 = T.int32()
+        A_1 = T.match_buffer(
+            A,
+            [n],
+            strides=[stride],
+            elem_offset=0,
+            align=64,
+            offset_factor=1,
+            buffer_type="auto",
+        )
+        B_1 = T.match_buffer(
+            B,
+            [n],
+            strides=[stride_1],
+            elem_offset=0,
+            align=64,
+            offset_factor=1,
+            buffer_type="auto",
+        )
+        C_1 = T.match_buffer(
+            C,
+            [n],
+            strides=[stride_2],
+            elem_offset=0,
+            align=64,
+            offset_factor=1,
+            buffer_type="auto",
+        )
+        d_1 = T.match_buffer(
+            d,
+            [n],
+            strides=[stride_3],
+            elem_offset=0,
+            align=64,
+            offset_factor=1,
+            buffer_type="auto",
+        )
+        # body
+        for i in T.serial(0, n):
+            d_1[(i * stride_3)] = (A_1[(i * stride)] * B_1[(i * stride_1)]) + C_1[(i * stride_2)]
+
+
+def test_fma():
+    opt = tvm.transform.Sequential(
+        [
+            tvm.tirx.transform.Apply(lambda f: f.with_attr("target", tvm.target.Target("llvm"))),
+            tvm.tirx.transform.LowerIntrin(),
+        ]
+    )
+    mod = opt(Module)
+    assert mod["test_tir_fma"].body.body.value.op.name == "tirx.call_llvm_pure_intrin"
+
+
+if __name__ == "__main__":
+    test_nearbyint()
+    test_unary_intrin()
+    test_round_intrinsics_on_int()
+    test_asin_acos_boundary_values()
+    test_binary_intrin()
+    test_ldexp()
+    test_clz()
+    test_fma()