chore: import upstream snapshot with attribution

tests/python/tirx/test_verifier.py

@@ -0,0 +1,371 @@
+# 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.
+import pytest
+
+from tvm.script import tirx as T
+from tvm.script.tirx import tile as Tx
+from tvm.tirx.analysis import verify_tirx_well_formed as verify
+
+
+def test_root_scope():
+    # fmt: off
+    @T.prim_func(check_well_formed=False)
+    def test1() -> None:
+        T.device_entry()
+        pass
+
+    @T.prim_func(check_well_formed=False)
+    def test2() -> None:
+        pass
+
+    @T.prim_func(check_well_formed=False)
+    def test3() -> None:
+        pass
+
+    @T.prim_func(check_well_formed=False)
+    def test4() -> None:
+        T.device_entry()
+        pass
+
+        # fmt: on
+
+    verify(test1)
+    verify(test2)
+    verify(test3)
+    verify(test4)
+
+
+def test_nested_scope():
+    # fmt: off
+    @T.prim_func(check_well_formed=False)
+    def test1() -> None:
+        T.device_entry()
+        pass
+        pass
+
+    @T.prim_func(check_well_formed=False)
+    def test2() -> None:
+        T.device_entry()
+        pass
+
+    @T.prim_func(check_well_formed=False)
+    def test3() -> None:
+        T.device_entry()
+        pass
+    @T.prim_func(check_well_formed=False)
+    def test4() -> None:
+        T.device_entry()
+        pass
+        pass
+
+        # fmt: on
+
+    verify(test1)
+    verify(test2)
+    verify(test3)
+    verify(test4)
+
+
+def test_scope_id_consistency():
+    # fmt: off
+    @T.prim_func(check_well_formed=False)
+    def test1():
+        T.device_entry()
+        T.cta_id([32])
+        T.warp_id([4])
+        T.lane_id([32])
+        pass
+
+    @T.prim_func(check_well_formed=False)
+    def test2():
+        T.device_entry()
+        T.cta_id([32])
+        T.warp_id([4])
+        T.lane_id([32])
+        T.thread_id([128])
+        pass
+
+    @T.prim_func(check_well_formed=False)
+    def test3():
+        T.device_entry()
+        T.cta_id([32])
+        T.warp_id([2])
+        T.lane_id([32])
+        T.thread_id([128])
+        pass
+
+    @T.prim_func(check_well_formed=False)
+    def test4():
+        T.device_entry()
+        bx, by, bz = T.cta_id([8, 10, 12])
+        cbx, cby, cbz = T.cta_id_in_cluster([2, 2, 1])
+        clx, cly, clz = T.cluster_id([4, 5, 12])
+        T.evaluate(bx + by + bz)
+        T.evaluate(cbx + cby + cbz)
+        T.evaluate(clx + cly + clz)
+
+    @T.prim_func(check_well_formed=False)
+    def test5():
+        T.device_entry()
+        bx, by, bz = T.cta_id([8, 10, 12])
+        cbx, cby, cbz = T.cta_id_in_cluster([2, 2, 1])
+        clx, cly, clz = T.cluster_id([3, 5, 12])
+        T.evaluate(bx + by + bz)
+        T.evaluate(cbx + cby + cbz)
+        T.evaluate(clx + cly + clz)
+
+    @T.prim_func(check_well_formed=False)
+    def test6():
+        T.device_entry()
+        clx, cly, clz = T.cluster_id([4, 5, 12])
+        bx, by, bz = T.cta_id([8, 10, 12])
+        cbx, cby, cbz = T.cta_id_in_cluster([2, 2, 1])
+        T.evaluate(bx + by + bz)
+        T.evaluate(cbx + cby + cbz)
+        T.evaluate(clx + cly + clz)
+
+    @T.prim_func(check_well_formed=False)
+    def test7():
+        T.device_entry()
+        clx, cly, clz = T.cluster_id([3, 5, 12])
+        bx, by, bz = T.cta_id([8, 10, 12])
+        cbx, cby, cbz = T.cta_id_in_cluster([2, 2, 1])
+        T.evaluate(bx + by + bz)
+        T.evaluate(cbx + cby + cbz)
+        T.evaluate(clx + cly + clz)
+
+        # fmt: on
+
+    verify(test1)
+    verify(test2)
+    with pytest.raises(Exception, match="Inconsistent extents for scope"):
+        verify(test3)
+    verify(test4)
+    with pytest.raises(Exception, match="Inconsistent extents|non-divisible extents"):
+        verify(test5)
+    verify(test6)
+    with pytest.raises(Exception, match="Inconsistent extents|non-divisible extents"):
+        verify(test7)
+
+
+def test_layout():
+    ### TileLayout
+    # fmt: off
+    @T.prim_func(check_well_formed=False)
+    def test1():
+        T.device_entry()
+        T.cta_id([32])
+        T.warp_id([4])
+        T.lane_id([32])
+        A = T.alloc_buffer((2,), layout=T.TileLayout(T.S[2, 1]))
+
+        A[0] = 0
+        # fmt: on
+    verify(test1)
+
+    ### SwizzleLayout
+    # fmt: off
+    @T.prim_func(check_well_formed=False)
+    def test2():
+        T.device_entry()
+        T.cta_id([32])
+        T.warp_id([4])
+        T.lane_id([32])
+        A = T.alloc_buffer((512,), scope="shared", layout=T.SwizzleLayout(3, 3, 3))
+
+        A[0] = 0
+        # fmt: on
+    verify(test2)
+
+
+def test_host():
+    # fmt: off
+    @T.prim_func(check_well_formed=False)
+    def test1(A_ptr: T.handle):
+        A = T.match_buffer(A_ptr, (16, 16), dtype="float32", align=16)
+
+        A_map: T.let[T.handle("tensormap")] = T.tvm_stack_alloca("tensormap", 1)
+        T.call_packed("runtime.cuTensorMapEncodeTiled", A_map, "float32", 2, A.data, 16, 16, 64, 16, 16, 1, 1, 0, 0, 0, 0)  # noqa: E501
+
+        T.device_entry()
+        for blockIdx in T.thread_binding(1, thread="blockIdx.x"):
+            for threadIdx in T.thread_binding(128, thread="threadIdx.x"):
+                bar = T.alloc_buffer((1,), "uint64", scope="shared", align=8)
+                phase = T.alloc_buffer((1,), "int32", scope="local")
+                A_smem = T.alloc_buffer((16, 16), "float32", scope="shared", align=128)
+
+                phase[0] = 0
+                if threadIdx == 0:
+                    T.ptx.mbarrier.init(bar.data, 1)
+                    T.ptx.fence.proxy_async("shared::cta")
+                    T.ptx.cp_async.bulk.tensor.g2c(2, A_smem.data, bar.data, T.address_of(A_map), 0, 1, "", 0, 0)  # noqa: E501
+                    T.ptx.mbarrier.arrive.expect_tx(bar.data, 16*16*4)
+                T.ptx.mbarrier.try_wait(bar.data, phase[0])
+                phase[0] = phase[0] ^ 1
+                T.print_buffer(A_smem.data, "float32", False, False, 2, 16*16)
+        # fmt: on
+    verify(test1)
+
+
+def test_device_func():
+    # Per-call exec-scope migration: scope is now attached per op via the
+    # ``T.op[scope](...)`` subscription surface instead of a ``with T.cta():``
+    # region. ``test1`` exercises a per-call-scoped op; ``test2`` the plain
+    # (unscoped) op. The old multi-root-scope negative case asserted the removed
+    # "only one root scope" verifier rule and no longer has an equivalent, so it
+    # is dropped.
+    # fmt: off
+    @T.prim_func(check_well_formed=False)
+    def test1(A: T.Buffer((128,), "float32")):
+        T.device_entry()
+        T.cta_id([1])
+        T.thread_id([128])
+        Tx.cta.fill(A, 0.)
+
+    @T.prim_func(check_well_formed=False)
+    def test2(A: T.Buffer((128,), "float32")):
+        T.device_entry()
+        T.cta_id([128])
+        T.thread_id([128])
+        Tx.fill(A, 0.)
+    # fmt: on
+    verify(test1, device_func=True)
+    verify(test2, device_func=True)
+
+
+def test_preferred_cluster_validation():
+    # fmt: off
+    # Valid: cluster→cta with preferred_extents matching size
+    @T.prim_func(check_well_formed=False)
+    def test1() -> None:
+        T.device_entry()
+        cbx, cby = T.cta_id_in_cluster([2, 1], preferred=[2, 2])
+        tx = T.thread_id([128])
+        T.evaluate(cbx + cby + tx)
+
+        # Invalid: preferred size doesn't match extents size (caught at verify time)
+    @T.prim_func(check_well_formed=False)
+    def test2() -> None:
+        T.device_entry()
+        cbx, cby = T.cta_id_in_cluster([2, 1], preferred=[2])
+        tx = T.thread_id([128])
+        T.evaluate(cbx + cby + tx)
+        # fmt: on
+
+    verify(test1)
+    with pytest.raises(Exception, match="preferred_extents must have the same size"):
+        verify(test2)
+
+    # Invalid: preferred on a non-cluster→cta scope (caught at IR build time)
+    with pytest.raises(Exception):
+        # fmt: off
+        @T.prim_func(check_well_formed=False)
+        def test3() -> None:
+            T.device_entry()
+            bx = T.cta_id([128], preferred=[256])
+            tx = T.thread_id([128])
+            T.evaluate(bx + tx)
+            # fmt: on
+
+
+def test_scope_id_deferred_relaxed_at_construction():
+    """Deferred scope_id (no extents) must pass the well-formed check even when
+    no sibling provides enough info to resolve it -- strict resolution is
+    deferred to LowerTIRx."""
+
+    # fmt: off
+    @T.prim_func(check_well_formed=False)
+    def partial_only_cta():
+        T.device_entry()
+        bx = T.cta_id()           # deferred kernel→cta, no closure source
+        tx = T.thread_id([128])   # explicit
+        T.evaluate(bx + tx)
+
+    @T.prim_func(check_well_formed=False)
+    def all_deferred():
+        T.device_entry()
+        bx = T.cta_id()
+        wg = T.warpgroup_id()
+        warp = T.warp_id_in_wg()
+        lane = T.lane_id()
+        T.evaluate(bx + wg + warp + lane)
+
+    @T.prim_func(check_well_formed=False)
+    def mixed():
+        T.device_entry()
+                # kCtaWarp=4, kWarpThread=32 → kCtaThread=128 derivable.
+        T.warp_id([4])
+        T.lane_id([32])
+        T.thread_id()             # deferred kCtaThread, resolvable via closure
+        pass
+        # fmt: on
+
+        # All three accepted by well-formed: deferred extents are tolerated.
+    verify(partial_only_cta)
+    verify(all_deferred)
+    verify(mixed)
+
+
+def test_scope_id_deferred_consistency_still_enforced():
+    """Even with deferred defs, known-known consistency between sibling defs
+    must still be enforced by the closure check."""
+
+    # fmt: off
+    @T.prim_func(check_well_formed=False)
+    def inconsistent():
+        # 4 warps * 32 lanes = 128 threads, but explicit thread_id says 64 -> error.
+        T.device_entry()
+        T.cta_id([32])
+        T.warp_id([4])
+        T.lane_id([32])
+        T.thread_id()       # deferred (shouldn't shadow the conflict)
+        T.thread_id([64])   # conflicts with derived kCtaThread=128
+        pass
+        # fmt: on
+
+    with pytest.raises(Exception, match="Inconsistent extents for scope"):
+        verify(inconsistent)
+
+
+def test_scope_id_deferred_multi_var_rejected():
+    """Deferred form (no extents) requires exactly one Var. Multi-var defers
+    have no well-defined recovery from fused closure values."""
+
+    # The C++ ScopeIdDef ctor enforces this; constructing such a def from the
+    # parser path is not currently expressible (parser only emits single-Var
+    # deferred), but we exercise the FFI-level guard directly.
+    from tvm.tirx.exec_scope import ScopeIdDef
+    from tvm.tirx.expr import Var
+
+    # Single-Var deferred form is fine.
+    ScopeIdDef([Var("", "int32")], None, "kernel", "cta")
+
+    # Two-Var deferred should be rejected.
+    with pytest.raises(Exception, match="Deferred ScopeIdDef.*must define exactly one Var"):
+        ScopeIdDef([Var("", "int32"), Var("", "int32")], None, "kernel", "cta")
+
+
+if __name__ == "__main__":
+    test_root_scope()
+    test_nested_scope()
+    test_scope_id_consistency()
+    test_layout()
+    test_host()
+    test_device_func()
+    test_scope_id_deferred_relaxed_at_construction()
+    test_scope_id_deferred_consistency_still_enforced()
+    test_scope_id_deferred_multi_var_rejected()