chore: import upstream snapshot with attribution

tests/python/s_tir/base/test_sblock_dependence_info.py

@@ -0,0 +1,155 @@
+# 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=missing-function-docstring,missing-module-docstring
+# ruff: noqa: E501, F401
+import gc
+import sys
+
+import pytest
+
+import tvm
+import tvm.testing
+from tvm import tirx
+from tvm.ir import IRModule
+from tvm.s_tir import SBlockDependenceInfo
+from tvm.s_tir.sblock_scope import DepKind
+from tvm.script import tirx as T
+from tvm.tirx import PrimFunc
+from tvm.tirx.stmt_functor import post_order_visit
+
+# pylint: disable=no-member,invalid-name,unused-variable
+
+
+@T.prim_func(s_tir=True)
+def elementwise(a: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, (128, 128), "float32")
+    C = T.match_buffer(c, (128, 128), "float32")
+    B = T.sblock_alloc_buffer((128, 128), "float32")
+    for i, j in T.grid(128, 128):
+        with T.sblock("B"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            B[vi, vj] = A[vi, vj] * 2.0
+    for i, j in T.grid(128, 128):
+        with T.sblock("C"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            C[vi, vj] = B[vi, vj] + 1.0
+    for i, j in T.grid(128, 128):
+        with T.sblock("D"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            C[vi, vj] = B[vi, vj] + 1.0
+
+
+@T.prim_func(s_tir=True)
+def war_dependency(a: T.handle, b: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, (128, 128))
+    B = T.match_buffer(b, (128, 128))
+    C = T.match_buffer(c, (128, 128))
+
+    for i, j in T.grid(128, 128):
+        with T.sblock("C"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            C[vi, vj] = B[vi, vj] + 1.0
+        with T.sblock("B"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            B[vi, vj] = A[vi, vj] * 2.0
+
+
+@T.prim_func(s_tir=True)
+def matmul(a: T.handle, b: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, [128, 128])
+    B = T.match_buffer(b, [128, 128])
+    C = T.match_buffer(c, [128, 128])
+    for i, j in T.grid(128, 128):
+        with T.sblock("init"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            C[vi, vj] = T.float32(0)
+        for k in range(0, 128):
+            with T.sblock("update"):
+                vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+                C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk]
+
+
+# pylint: enable=no-member,invalid-name,unused-variable
+
+
+def get_sblocks(func: PrimFunc):
+    blocks = {}
+
+    def update_blocks(node):
+        if isinstance(node, tvm.tirx.SBlock):
+            blocks[node.name_hint] = node
+
+    # post_order_visit(func.body, lambda node: blocks[node.name_hint] = node if isinstance(node, tvm.tirx.SBlock) else None)
+    post_order_visit(func.body, update_blocks)
+    return blocks
+
+
+def _verify_dependence(dependence_info, src_block, dst_block, kind):
+    src_sref = dependence_info.get_sref(src_block)
+    dst_sref = dependence_info.get_sref(dst_block)
+    scope = dependence_info.get_sblock_scope(src_sref.parent)
+
+    def _find_dependence(deps):
+        for dep in deps:
+            if dep.src == src_sref and dep.dst == dst_sref and dep.kind == kind:
+                return dep
+        return None
+
+    def _get_dependency_kind_name(dep_kind):
+        if isinstance(dep_kind, int):
+            dep_kind = DepKind(dep_kind)
+        return dep_kind.name
+
+    # Check dependences by src
+    deps_by_src = scope.get_deps_by_src(src_sref)
+    dependence = _find_dependence(deps_by_src)
+    assert dependence, (
+        f"Expected a dependency with src block {src_block.name_hint} and dst block {dst_block.name_hint} of kind {kind.name}"
+    )
+
+    # Check dependences by dst
+    deps_by_dst = scope.get_deps_by_dst(dst_sref)
+    dependence = _find_dependence(deps_by_dst)
+    assert dependence, (
+        f"Expected a dependency with src block {src_block.name_hint} and dst block {dst_block.name_hint}"
+    )
+
+
+def test_RAW_dependences():
+    func = elementwise
+    dependence_info = SBlockDependenceInfo(func)
+    blocks = get_sblocks(func)
+    _verify_dependence(dependence_info, blocks["B"], blocks["C"], DepKind.RAW)
+
+
+def test_WAR_dependences():
+    func = war_dependency
+    dependence_info = SBlockDependenceInfo(func)
+    blocks = get_sblocks(func)
+    _verify_dependence(dependence_info, blocks["C"], blocks["B"], DepKind.WAR)
+
+
+def test_RAW_and_WAW_dependences():
+    func = matmul
+    dependence_info = SBlockDependenceInfo(func)
+    blocks = get_sblocks(func)
+    _verify_dependence(dependence_info, blocks["init"], blocks["update"], DepKind.RAW)
+    _verify_dependence(dependence_info, blocks["init"], blocks["update"], DepKind.WAW)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/s_tir/base/test_tir_data_layout.py

@@ -0,0 +1,176 @@
+# 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.
+"""Test layout and bijective-layout node"""
+
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.error import InternalError
+from tvm.topi.utils import get_const_tuple
+
+
+def test_layout():
+    layout = tvm.s_tir.slayout("NCHW16c")
+    assert layout is not None
+    assert isinstance(layout, tvm.s_tir.SLayout)
+
+    assert layout.factor_of("c") == 16
+    assert layout.factor_of("C") == 16
+    assert layout.factor_of("N") == -1
+
+    assert layout.index_of("N") == 0
+    assert layout.index_of("C") == 1
+    assert layout.index_of("H") == 2
+    assert layout.index_of("W") == 3
+    assert layout.index_of("16c") == 4
+    assert layout.index_of("O") == -1
+
+    assert "N" in layout
+    assert "C" in layout
+    assert "H" in layout
+    assert "W" in layout
+    assert "c" in layout
+    assert "O" not in layout
+
+    assert layout[0] == "N"
+    assert layout[1] == "C"
+    assert layout[2] == "H"
+    assert layout[3] == "W"
+    assert layout[4] == "16c"
+
+    layout = tvm.s_tir.slayout("OIHW[4o4i]")
+    assert layout is not None
+    assert isinstance(layout, tvm.s_tir.SLayout)
+
+    assert layout.factor_of("o") == 4
+    assert layout.factor_of("i") == 4
+    assert layout.factor_of("H") == -1
+    assert layout.factor_of("W") == -1
+    assert layout.factor_of("N") == -1
+
+    assert layout.index_of("O") == 0
+    assert layout.index_of("I") == 1
+    assert layout.index_of("H") == 2
+    assert layout.index_of("W") == 3
+    assert layout.index_of("4o4i") == 4
+    assert layout.index_of("i") == -1
+    assert layout.index_of("o") == -1
+
+    assert "O" in layout
+    assert "I" in layout
+    assert "H" in layout
+    assert "W" in layout
+    assert "4o4i" in layout
+    assert "i" in layout
+    assert "o" in layout
+
+    assert layout[0] == "O"
+    assert layout[1] == "I"
+    assert layout[2] == "H"
+    assert layout[3] == "W"
+    assert layout[4] == "4o4i"
+
+    with pytest.raises(InternalError):
+        layout = tvm.s_tir.slayout("[N4o]C")
+    with pytest.raises(InternalError):
+        layout = tvm.s_tir.slayout("[O4o]")
+    with pytest.raises(InternalError):
+        layout = tvm.s_tir.slayout("C4o")
+    with pytest.raises(InternalError):
+        layout = tvm.s_tir.slayout("OI[4o4i][]")
+    with pytest.raises(InternalError):
+        layout = tvm.s_tir.slayout("C4c[4c]")
+
+
+def test_layout_dtype():
+    layout_i32 = tvm.s_tir.slayout("NCHW")
+    assert layout_i32.axes[0].var.ty.dtype == "int32"
+    assert layout_i32.axes[0].dom.min.ty.dtype == "int32"
+    assert layout_i32.axes[0].dom.extent.ty.dtype == "int32"
+    assert layout_i32.axes[1].var.ty.dtype == "int32"
+    assert layout_i32.axes[1].dom.min.ty.dtype == "int32"
+    assert layout_i32.axes[1].dom.extent.ty.dtype == "int32"
+
+    layout_i64 = tvm.s_tir.slayout("NCHW", dtype="int64")
+    assert layout_i64.axes[2].var.ty.dtype == "int64"
+    assert layout_i64.axes[2].dom.min.ty.dtype == "int64"
+    assert layout_i64.axes[2].dom.extent.ty.dtype == "int64"
+    assert layout_i64.axes[3].var.ty.dtype == "int64"
+    assert layout_i64.axes[3].dom.min.ty.dtype == "int64"
+    assert layout_i64.axes[3].dom.extent.ty.dtype == "int64"
+
+    with pytest.raises(TypeError):
+        tvm.s_tir.slayout("NCHW", dtype="float32")
+    with pytest.raises(TypeError):
+        tvm.s_tir.slayout("NCHW", dtype=None)
+
+
+def test_bilayout_convertible():
+    # not convertible
+    assert tvm.s_tir.sbijective_layout("NCHW", "ABCD") is None
+    assert tvm.s_tir.sbijective_layout("__undef__", "NCHW") is None
+    assert tvm.s_tir.sbijective_layout("NCHW", "__undef__") is None
+    assert tvm.s_tir.sbijective_layout("__undef__", "__undef__") is None
+    assert tvm.s_tir.sbijective_layout("", "NCHW") is None
+    assert tvm.s_tir.sbijective_layout("NCHW", "") is None
+    assert tvm.s_tir.sbijective_layout("OIHW", "OIHW[4o4i]") is not None
+    assert tvm.s_tir.sbijective_layout("OIHW[2o4i]", "OIHW") is not None
+    assert tvm.s_tir.sbijective_layout("", "") is None
+    # convertible
+    assert tvm.s_tir.sbijective_layout("NCHW", "NCHW16c") is not None
+
+
+def test_bilayout_shape():
+    bilayout = tvm.s_tir.sbijective_layout("NCHW", "NCHW16c")
+    assert isinstance(bilayout, tvm.s_tir.SBijectiveLayout)
+
+    dst_shape = bilayout.forward_shape((1, 32, 7, 7))
+    assert get_const_tuple(dst_shape) == (1, 2, 7, 7, 16)
+
+    src_shape = bilayout.backward_shape(dst_shape)
+    assert get_const_tuple(src_shape) == (1, 32, 7, 7)
+
+    bilayout = tvm.s_tir.sbijective_layout("OIHW", "OIHW[4o4i]")
+
+    dst_shape = bilayout.forward_shape((64, 28, 7, 7))
+    assert get_const_tuple(dst_shape) == (16, 7, 7, 7, 16)
+
+    src_shape = bilayout.backward_shape((2, 11, 4, 4, 16))
+    assert get_const_tuple(src_shape) == (8, 44, 4, 4)
+
+
+def test_bilayout_index():
+    bilayout = tvm.s_tir.sbijective_layout("NCHW", "NCHW16c")
+
+    dst_index = bilayout.forward_index([0, 18, 6, 6])
+    assert get_const_tuple(dst_index) == (0, 1, 6, 6, 2)
+
+    src_index = bilayout.backward_index([0, 1, 6, 6, 2])
+    assert get_const_tuple(src_index) == (0, 18, 6, 6)
+
+    bilayout = tvm.s_tir.sbijective_layout("OIHW", "OIHW[4o4i]")
+
+    dst_index = bilayout.forward_index((63, 29, 7, 7))
+    assert get_const_tuple(dst_index) == (15, 7, 7, 7, 13)
+
+    src_index = bilayout.backward_index((4, 7, 4, 4, 13))
+    assert get_const_tuple(src_index) == (19, 29, 4, 4)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/s_tir/base/test_tir_te_extern_primfunc.py

@@ -0,0 +1,228 @@
+# 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: F401
+
+import sys
+
+import numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm import te
+from tvm.script import tirx as T
+
+# TODO(csullivan): Additional tests cases needed:
+# - PrimFunc with 1 arg, inplace update
+# - PrimFunc with buffer that uses custom storage_scope
+
+
+@T.prim_func(s_tir=True)
+def func_1(A: T.Buffer((16,), "float32"), C: T.Buffer((1,), "float32")):
+    for i in T.serial(
+        0,
+        16,
+    ):
+        with T.sblock():
+            B = T.sblock_alloc_buffer((1,), dtype="float32")
+            with T.sblock():
+                B[0] = A[i] * T.float32(2)
+            with T.sblock():
+                C[0] = C[0] + A[i] + B[0] + T.float32(1)
+                A[i] = B[0] + T.float32(1)
+
+
+def verify_func_1(module):
+    a_np = np.random.randint(low=-128, high=127, size=(16,)).astype(np.float32)
+    c_np = np.zeros((1,), dtype=np.float32)
+    a = tvm.runtime.tensor(a_np, device=tvm.cpu(0))
+    c = tvm.runtime.tensor(c_np, device=tvm.cpu(0))
+
+    module(a, c)
+    tvm.testing.assert_allclose(c_np + np.sum(3 * a_np + 1), c.numpy(), rtol=1e-4)
+    # also test in place update
+    tvm.testing.assert_allclose(a_np * 2 + 1, a.numpy(), rtol=1e-4)
+
+
+@T.prim_func(s_tir=True)
+def func_2(
+    C: T.Buffer((1,), "float32"), A: T.Buffer((16,), "float32"), D: T.Buffer((2,), "float32")
+):
+    for i in T.serial(
+        0,
+        16,
+    ):
+        with T.sblock():
+            B = T.sblock_alloc_buffer((1,), dtype="float32")
+            with T.sblock():
+                B[0] = A[i] * T.float32(2)
+            with T.sblock():
+                C[0] = C[0] + A[i] + B[0] + T.float32(1) + D[0]
+                A[i] = B[0] + T.float32(1) + D[1]
+
+
+def verify_func_2(module):
+    a_np = np.random.randint(low=-128, high=127, size=(16,)).astype(np.float32)
+    d_np = np.random.randint(low=-128, high=127, size=(2,)).astype(np.float32)
+    c_np = np.zeros((1,), dtype=np.float32)
+    a = tvm.runtime.tensor(a_np, device=tvm.cpu(0))
+    d = tvm.runtime.tensor(d_np, device=tvm.cpu(0))
+    c = tvm.runtime.tensor(c_np, device=tvm.cpu(0))
+
+    module(c, a, d)
+    tvm.testing.assert_allclose(c_np + np.sum(3 * a_np + 1 + d_np[0]), c.numpy(), rtol=1e-4)
+    tvm.testing.assert_allclose(a_np * 2 + 1 + d_np[1], a.numpy(), rtol=1e-4)
+
+
+@T.prim_func(s_tir=True)
+def func_3(
+    C: T.Buffer((1,), "float32"),
+    A: T.Buffer((16,), "float32"),
+    D: T.Buffer((2,), "float32"),
+    E: T.Buffer((16,), "float32"),
+    F: T.Buffer((16,), "float32"),
+):
+    for i in T.serial(
+        0,
+        16,
+    ):
+        with T.sblock():
+            B = T.sblock_alloc_buffer((1,), dtype="float32")
+            with T.sblock():
+                B[0] = A[i] * T.float32(2)
+            with T.sblock():
+                E[i] = A[i]
+                F[i] = E[i] + 1.0
+                C[0] = C[0] + A[i] + B[0] + T.float32(1) + D[0]
+                A[i] = B[0] + T.float32(1) + D[1]
+
+
+def verify_func_3(module):
+    a_np = np.random.randint(low=-128, high=127, size=(16,)).astype(np.float32)
+    d_np = np.random.randint(low=-128, high=127, size=(2,)).astype(np.float32)
+    c_np = np.zeros((1,), dtype=np.float32)
+    e_np = np.zeros((16,), dtype=np.float32)
+    f_np = np.zeros((16,), dtype=np.float32)
+    a = tvm.runtime.tensor(a_np, device=tvm.cpu(0))
+    d = tvm.runtime.tensor(d_np, device=tvm.cpu(0))
+    c = tvm.runtime.tensor(c_np, device=tvm.cpu(0))
+    e = tvm.runtime.tensor(e_np, device=tvm.cpu(0))
+    f = tvm.runtime.tensor(f_np, device=tvm.cpu(0))
+
+    module(c, a, d, e, f)
+    tvm.testing.assert_allclose(c_np + np.sum(3 * a_np + 1 + d_np[0]), c.numpy(), rtol=1e-4)
+    tvm.testing.assert_allclose(a_np * 2 + 1 + d_np[1], a.numpy(), rtol=1e-4)
+    tvm.testing.assert_allclose(a_np, e.numpy(), rtol=1e-4)
+    tvm.testing.assert_allclose(a_np + 1, f.numpy(), rtol=1e-4)
+
+
+@T.prim_func(s_tir=True)
+def func_4(
+    C: T.Buffer((1,), "float32"),
+    A: T.Buffer((16,), "float32"),
+    F: T.Buffer((16,), "float32"),
+    D: T.Buffer((2,), "float32"),
+    E: T.Buffer((16,), "float32"),
+):
+    for i in T.serial(
+        0,
+        16,
+    ):
+        with T.sblock():
+            B = T.sblock_alloc_buffer((1,), dtype="float32")
+            with T.sblock():
+                B[0] = A[i] * T.float32(2)
+            with T.sblock():
+                E[i] = A[i]
+                F[i] = E[i] + 1.0
+                C[0] = C[0] + A[i] + B[0] + T.float32(1) + D[0]
+                A[i] = B[0] + T.float32(1) + D[1]
+
+
+def verify_func_4(module):
+    a_np = np.random.randint(low=-128, high=127, size=(16,)).astype(np.float32)
+    d_np = np.random.randint(low=-128, high=127, size=(2,)).astype(np.float32)
+    c_np = np.zeros((1,), dtype=np.float32)
+    e_np = np.zeros((16,), dtype=np.float32)
+    f_np = np.zeros((16,), dtype=np.float32)
+    a = tvm.runtime.tensor(a_np, device=tvm.cpu(0))
+    d = tvm.runtime.tensor(d_np, device=tvm.cpu(0))
+    c = tvm.runtime.tensor(c_np, device=tvm.cpu(0))
+    e = tvm.runtime.tensor(e_np, device=tvm.cpu(0))
+    f = tvm.runtime.tensor(f_np, device=tvm.cpu(0))
+
+    module(c, a, f, d, e)
+    tvm.testing.assert_allclose(c_np + np.sum(3 * a_np + 1 + d_np[0]), c.numpy(), rtol=1e-4)
+    tvm.testing.assert_allclose(a_np * 2 + 1 + d_np[1], a.numpy(), rtol=1e-4)
+    tvm.testing.assert_allclose(a_np, e.numpy(), rtol=1e-4)
+    tvm.testing.assert_allclose(a_np + 1, f.numpy(), rtol=1e-4)
+
+
+_primfunc_cases = [
+    [func_1, ("A"), verify_func_1],
+    [func_2, ("C", "D"), verify_func_2],
+    [func_3, ("C", "A", "D", "E"), verify_func_3],
+    [func_4, ("C", "A", "D", "E"), verify_func_4],
+]
+
+
+class TestPrimFuncs:
+    @pytest.mark.parametrize("func,verify", [(case[0], case[2]) for case in _primfunc_cases])
+    def test_primfunc_call(self, func, verify):
+        target = tvm.target.Target("llvm")
+        func = tvm.compile(func, target=target)
+        verify(func)
+
+    @pytest.mark.parametrize("func,params,verify", _primfunc_cases)
+    def test_te_extern_call(self, func, params, verify):
+        ir_mod = tvm.IRModule.from_expr(func.with_attr("global_symbol", "main"))
+        prim_func = ir_mod["main"]
+
+        buf_name_map = {buf.name: buf for buf in func.buffer_map.values()}
+        input_tensors = [te.placeholder(buf_name_map[name].shape) for name in params]
+        output = te.extern_primfunc(input_tensors, prim_func)
+        rt_prim_func = te.create_prim_func(tensors_from_extern_op(output, prim_func))
+
+        target = tvm.target.Target("llvm")
+        func = tvm.compile(rt_prim_func, target=target)
+        verify(func)
+
+
+def tensors_from_extern_op(extern, func):
+    if isinstance(extern, list):
+        output_tensors = extern
+    else:
+        output_tensors = [extern]
+    output_buffers = []
+    input_buffers = []
+    input_tensors = []
+    for ext in output_tensors:
+        output_buffers.extend(ext.op.output_placeholders)
+        input_buffers.extend(ext.op.input_placeholders)
+        input_tensors.extend(ext.op.input_tensors)
+    input_binds = dict(zip(input_buffers, input_tensors))
+    output_binds = dict(zip(output_buffers, output_tensors))
+    buffer_to_tensor = {**input_binds, **output_binds}
+    ordered_tensors = []
+    for var in func.params:
+        buf = func.buffer_map[var]
+        ordered_tensors.append(buffer_to_tensor[buf])
+    return ordered_tensors
+
+
+if __name__ == "__main__":
+    tvm.testing.main()