chore: import upstream snapshot with attribution

tests/python/tirx/test_alloc_pool.py

@@ -0,0 +1,117 @@
+# 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 CUDA allocation pool validation."""
+
+import pytest
+
+from tvm.tirx.cuda.lang.alloc_pool import _validate_mma_alloc_shape
+from tvm.tirx.cuda.operator.tile_primitive.tma_utils import SwizzleMode
+
+# ---------------------------------------------------------------------------
+# alloc_mma shape validation: bad inputs raise actionable ValueError instead of
+# the opaque "Divide by zero" diagnostic that ``Layout.tile_to`` would emit.
+# ---------------------------------------------------------------------------
+
+
+class TestAllocMmaValidationRowBytes:
+    """row width (cols * itemsize) must be a positive multiple of swizzle atom bytes."""
+
+    def test_bf16_32cols_128b_swizzle_too_narrow(self):
+        # The exact case that bit gdn-prefill v1_0 / v1_2 (eval R10).
+        # Row = 32 * 2B = 64B < 128B atom.
+        with pytest.raises(ValueError, match=r"64B rows.*128B swizzle atom"):
+            _validate_mma_alloc_shape((128, 32), "bfloat16", SwizzleMode.SWIZZLE_128B_ATOM)
+
+    def test_error_suggests_smaller_swizzle(self):
+        try:
+            _validate_mma_alloc_shape((128, 32), "bfloat16", SwizzleMode.SWIZZLE_128B_ATOM)
+        except ValueError as e:
+            assert "SWIZZLE_64B_ATOM" in str(e), f"missing fix-it hint: {e}"
+        else:
+            pytest.fail("should have raised")
+
+    def test_error_suggests_widening_cols(self):
+        try:
+            _validate_mma_alloc_shape((128, 32), "bfloat16", SwizzleMode.SWIZZLE_128B_ATOM)
+        except ValueError as e:
+            assert "multiple of 64 elements" in str(e), f"missing widen hint: {e}"
+        else:
+            pytest.fail("should have raised")
+
+    def test_fp32_16cols_128b_swizzle_too_narrow(self):
+        # Row = 16 * 4B = 64B < 128B atom.
+        with pytest.raises(ValueError, match=r"64B rows.*128B swizzle atom"):
+            _validate_mma_alloc_shape((128, 16), "float32", SwizzleMode.SWIZZLE_128B_ATOM)
+
+    def test_3d_shape_validates_last_dim(self):
+        # Validation must consider shape[-1], not shape[0].
+        with pytest.raises(ValueError, match=r"64B rows"):
+            _validate_mma_alloc_shape((2, 128, 32), "bfloat16", SwizzleMode.SWIZZLE_128B_ATOM)
+
+
+class TestAllocMmaValidationRowCount:
+    """rows (shape[-2]) must be a positive multiple of the 8-row atom."""
+
+    def test_rows_below_atom_rejected(self):
+        with pytest.raises(ValueError, match=r"shape\[-2\]=4.*multiple of 8"):
+            _validate_mma_alloc_shape((4, 64), "bfloat16", SwizzleMode.SWIZZLE_128B_ATOM)
+
+    def test_rows_not_multiple_of_8_rejected(self):
+        with pytest.raises(ValueError, match=r"shape\[-2\]=12.*multiple of 8"):
+            _validate_mma_alloc_shape((12, 64), "bfloat16", SwizzleMode.SWIZZLE_128B_ATOM)
+
+
+class TestAllocMmaValidationRank:
+    """rank-1 shapes cannot be tiled with a 2-D swizzle atom."""
+
+    def test_rank_one_rejected(self):
+        with pytest.raises(ValueError, match=r"fewer than 2 dimensions"):
+            _validate_mma_alloc_shape((128,), "bfloat16", SwizzleMode.SWIZZLE_128B_ATOM)
+
+
+class TestAllocMmaValidationValid:
+    """combinations that should succeed must not be rejected."""
+
+    @pytest.mark.parametrize(
+        "shape,dtype,mode",
+        [
+            # The fix path the agent should pick when row_bytes >= 128.
+            ((128, 64), "bfloat16", SwizzleMode.SWIZZLE_128B_ATOM),
+            ((128, 128), "bfloat16", SwizzleMode.SWIZZLE_128B_ATOM),
+            # Or downgrade to a swizzle whose atom matches the row.
+            ((128, 32), "bfloat16", SwizzleMode.SWIZZLE_64B_ATOM),
+            ((128, 16), "bfloat16", SwizzleMode.SWIZZLE_32B_ATOM),
+            # 3-D request validates the last two dims only.
+            ((2, 128, 64), "bfloat16", SwizzleMode.SWIZZLE_128B_ATOM),
+            # fp32 with row width >= atom.
+            ((128, 32), "float32", SwizzleMode.SWIZZLE_128B_ATOM),
+            # fp8 (1B) with row width >= atom.
+            ((128, 128), "float8_e4m3", SwizzleMode.SWIZZLE_128B_ATOM),
+        ],
+    )
+    def test_valid_combinations_accepted(self, shape, dtype, mode):
+        _validate_mma_alloc_shape(shape, dtype, mode)
+
+    def test_swizzle_none_skips_validation(self):
+        # SWIZZLE_NONE has no atom — even otherwise-bad shapes are allowed.
+        _validate_mma_alloc_shape((128, 32), "bfloat16", SwizzleMode.SWIZZLE_NONE)
+        _validate_mma_alloc_shape((3, 5), "bfloat16", SwizzleMode.SWIZZLE_NONE)
+        _validate_mma_alloc_shape((128,), "bfloat16", SwizzleMode.SWIZZLE_NONE)
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])