chore: import upstream snapshot with attribution

tests/python/tirx/test_hint.py

@@ -0,0 +1,268 @@
+# 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.
+"""Tests for T.hint() — universal directive primitive for TIRx sketch language."""
+
+import tvm
+import tvm.script
+import tvm.testing
+from tvm.ir import assert_structural_equal
+from tvm.script import tirx as T
+from tvm.tirx import AttrStmt
+
+
+def from_source(code):
+    return tvm.script.from_source(code)
+
+
+def test_hint_statement():
+    """T.hint("msg") as a bare statement produces an AttrStmt with attr_key=tirx_hint."""
+
+    @T.prim_func
+    def func(A_ptr: T.handle) -> None:
+        _A = T.match_buffer(A_ptr, (64,), "float32", scope="global")
+        bx, by, bz = T.cta_id([1, 1, 1])
+        warp_id = T.warp_id([1])
+        lane_id = T.lane_id([32])
+        T.hint("persistent tile scheduler with L2 swizzle")
+        T.evaluate(0)
+
+    # Walk the IR to find the AttrStmt with tirx_hint
+    found = [False]
+
+    def visit(stmt):
+        if isinstance(stmt, AttrStmt) and stmt.attr_key == "tirx_hint":
+            # node is now a Map with "message" key
+            assert isinstance(stmt.node, tvm.ir.Map)
+            assert str(stmt.node["message"]) == "persistent tile scheduler with L2 swizzle"
+            found[0] = True
+
+    tvm.tirx.stmt_functor.post_order_visit(func.body, visit)
+    assert found[0], "Expected AttrStmt with attr_key='tirx_hint' not found"
+
+
+def test_hint_context_manager():
+    """with T.hint("msg"): scopes its body inside the AttrStmt."""
+
+    @T.prim_func
+    def func(A_ptr: T.handle) -> None:
+        _A = T.match_buffer(A_ptr, (64,), "float32", scope="global")
+        bx, by, bz = T.cta_id([1, 1, 1])
+        warp_id = T.warp_id([1])
+        lane_id = T.lane_id([32])
+        with T.hint("software pipeline, depth 4"):
+            T.evaluate(0)
+
+    found = [False]
+
+    def visit(stmt):
+        if isinstance(stmt, AttrStmt) and stmt.attr_key == "tirx_hint":
+            assert isinstance(stmt.node, tvm.ir.Map)
+            assert str(stmt.node["message"]) == "software pipeline, depth 4"
+            found[0] = True
+
+    tvm.tirx.stmt_functor.post_order_visit(func.body, visit)
+    assert found[0], "Expected AttrStmt with attr_key='tirx_hint' not found"
+
+
+def test_hint_with_attrs():
+    """T.hint("msg", key="value") passes structured attrs in Map node."""
+
+    @T.prim_func
+    def func(A_ptr: T.handle) -> None:
+        _A = T.match_buffer(A_ptr, (64,), "float32", scope="global")
+        bx, by, bz = T.cta_id([1, 1, 1])
+        warp_id = T.warp_id([1])
+        lane_id = T.lane_id([32])
+        T.hint("scheduler", mode="persistent", depth="4")
+        T.evaluate(0)
+
+    found = [False]
+
+    def visit(stmt):
+        if isinstance(stmt, AttrStmt) and stmt.attr_key == "tirx_hint":
+            assert isinstance(stmt.node, tvm.ir.Map)
+            assert str(stmt.node["message"]) == "scheduler"
+            assert str(stmt.node["mode"]) == "persistent"
+            assert str(stmt.node["depth"]) == "4"
+            found[0] = True
+
+    tvm.tirx.stmt_functor.post_order_visit(func.body, visit)
+    assert found[0], "Expected AttrStmt with attr_key='tirx_hint' not found"
+
+
+def test_hint_printer_roundtrip_statement():
+    """Verify T.hint("msg") prints as T.hint("msg") and roundtrips through script/parse."""
+
+    @T.prim_func
+    def func(A_ptr: T.handle) -> None:
+        _A = T.match_buffer(A_ptr, (64,), "float32", scope="global")
+        bx, by, bz = T.cta_id([1, 1, 1])
+        warp_id = T.warp_id([1])
+        lane_id = T.lane_id([32])
+        T.hint("persistent tile scheduler with L2 swizzle")
+        T.evaluate(0)
+
+    code = func.script()
+    assert 'hint("persistent tile scheduler with L2 swizzle")' in code
+    reparsed = from_source(code)
+    assert_structural_equal(func, reparsed)
+
+
+def test_hint_printer_roundtrip_context_manager():
+    """Verify with T.hint("msg"): prints correctly and roundtrips."""
+
+    @T.prim_func
+    def func(A_ptr: T.handle) -> None:
+        _A = T.match_buffer(A_ptr, (64,), "float32", scope="global")
+        bx, by, bz = T.cta_id([1, 1, 1])
+        warp_id = T.warp_id([1])
+        lane_id = T.lane_id([32])
+        with T.hint("software pipeline, depth 4"):
+            T.evaluate(0)
+
+    code = func.script()
+    assert 'hint("software pipeline, depth 4")' in code
+    reparsed = from_source(code)
+    assert_structural_equal(func, reparsed)
+
+
+def test_hint_printer_roundtrip_with_attrs():
+    """Verify T.hint("msg", key="val") prints with kwargs and roundtrips."""
+
+    @T.prim_func
+    def func(A_ptr: T.handle) -> None:
+        _A = T.match_buffer(A_ptr, (64,), "float32", scope="global")
+        bx, by, bz = T.cta_id([1, 1, 1])
+        warp_id = T.warp_id([1])
+        lane_id = T.lane_id([32])
+        T.hint("scheduler", mode="persistent")
+        T.evaluate(0)
+
+    code = func.script()
+    assert 'hint("scheduler"' in code
+    assert 'mode="persistent"' in code
+    reparsed = from_source(code)
+    assert_structural_equal(func, reparsed)
+
+
+def test_hint_keyword_arg_on_tx_op():
+    """Tx.op(..., hint="msg") stores hint in TilePrimitiveCall.config."""
+    from tvm.tirx.buffer import decl_buffer
+    from tvm.tirx.stmt import TilePrimitiveCall
+
+    A = decl_buffer((64, 64), "float32", scope="global")
+    A_sm = decl_buffer((64, 64), "float32", scope="shared")
+
+    op_call = TilePrimitiveCall(
+        A[0:64, 0:64],
+        A_sm[0:64, 0:64],
+        op=tvm.ir.Op.get("tirx.tile.copy"),
+        workspace={},
+        config={"hint": "3-input ptx"},
+    )
+    assert "hint" in op_call.config
+    assert str(op_call.config["hint"]) == "3-input ptx"
+
+
+def test_hint_keyword_arg_on_tx_op_roundtrip():
+    """Tx.op(..., hint="msg") roundtrips through printer/parser."""
+    from tvm.script.tirx import tile as Tx
+
+    @T.prim_func
+    def func(A_ptr: T.handle, B_ptr: T.handle):
+        A = T.match_buffer(A_ptr, [10], "float32", scope="global")
+        B = T.match_buffer(B_ptr, [10], "float32", scope="global")
+        Tx.add(B, A, T.float32(1), hint="use_fast_math")
+
+    code = func.script()
+    assert 'hint="use_fast_math"' in code
+    reparsed = from_source(code)
+    assert reparsed.script() == code
+    assert_structural_equal(func, reparsed)
+
+
+def test_hint_no_message():
+    """T.hint(access=...) with no message string."""
+
+    @T.prim_func
+    def func(A_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, (128,), "float32", scope="global")
+        bx, by, bz = T.cta_id([1, 1, 1])
+        warp_id = T.warp_id([1])
+        lane_id = T.lane_id([32])
+        T.hint(access=A[0:64])
+        T.evaluate(0)
+
+    found = [False]
+
+    def visit(stmt):
+        if isinstance(stmt, AttrStmt) and stmt.attr_key == "tirx_hint":
+            assert isinstance(stmt.node, tvm.ir.Map)
+            # Should have "access" key but no "message" key
+            assert "access" in stmt.node
+            assert "message" not in stmt.node
+            from tvm.tirx import BufferRegion
+
+            assert isinstance(stmt.node["access"], BufferRegion)
+            found[0] = True
+
+    tvm.tirx.stmt_functor.post_order_visit(func.body, visit)
+    assert found[0], "Expected AttrStmt with attr_key='tirx_hint' containing access not found"
+
+
+def test_hint_access_buffer_region():
+    """T.hint(access=A[region]) stores the BufferRegion structurally in the IR."""
+
+    @T.prim_func
+    def func(A_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, (128, 64), "float32", scope="global")
+        bx, by, bz = T.cta_id([2, 1, 1])
+        warp_id = T.warp_id([1])
+        lane_id = T.lane_id([32])
+        T.hint("partition", access=A[bx * 64 : (bx + 1) * 64, 0:64])
+        T.evaluate(0)
+
+    found = [False]
+
+    def visit(stmt):
+        if isinstance(stmt, AttrStmt) and stmt.attr_key == "tirx_hint":
+            assert isinstance(stmt.node, tvm.ir.Map)
+            assert str(stmt.node["message"]) == "partition"
+            assert "access" in stmt.node
+            from tvm.tirx import BufferRegion
+
+            assert isinstance(stmt.node["access"], BufferRegion)
+            br = stmt.node["access"]
+            assert br.buffer.name == "A"
+            assert len(br.region) == 2
+            found[0] = True
+
+    tvm.tirx.stmt_functor.post_order_visit(func.body, visit)
+    assert found[0], "Expected AttrStmt with structured BufferRegion access not found"
+
+
+if __name__ == "__main__":
+    test_hint_statement()
+    test_hint_context_manager()
+    test_hint_with_attrs()
+    test_hint_printer_roundtrip_statement()
+    test_hint_printer_roundtrip_context_manager()
+    test_hint_printer_roundtrip_with_attrs()
+    test_hint_keyword_arg_on_tx_op()
+    test_hint_keyword_arg_on_tx_op_roundtrip()
+    test_hint_no_message()
+    test_hint_access_buffer_region()