diff --git a/python/tvm/backend/cuda/operator/tile_primitive/copy/_common.py b/python/tvm/backend/cuda/operator/tile_primitive/copy/_common.py index b1dd018f86ec..5ff8100c397d 100644 --- a/python/tvm/backend/cuda/operator/tile_primitive/copy/_common.py +++ b/python/tvm/backend/cuda/operator/tile_primitive/copy/_common.py @@ -320,7 +320,11 @@ def _extract_tile(layout, region): if isinstance(layout, ComposeLayout): return layout.tile_layout if isinstance(layout, SwizzleLayout): - extents = [int(end - start) for (start, end) in region] + # Region bounds may be constant-valued but remain as unfolded + # expressions after substitution. Simplify before converting to a + # Python integer; genuinely symbolic tile extents still raise. + analyzer = arith.Analyzer() + extents = [int(analyzer.simplify(end - start)) for (start, end) in region] return TileLayout(S[tuple(extents)]) return layout diff --git a/tests/python/tirx/test_layout.py b/tests/python/tirx/test_layout.py index 5099a5e32445..998eab97d0c1 100644 --- a/tests/python/tirx/test_layout.py +++ b/tests/python/tirx/test_layout.py @@ -1655,6 +1655,17 @@ def case_compose_slice_2d(): case_compose_slice_2d() +def test_cuda_copy_extract_swizzle_tile_simplifies_constant_region_extents(): + from tvm.tirx.cuda.operator.tile_primitive.copy._common import _extract_tile + + zero = tvm.tirx.Mul(tvm.tirx.IntImm("int32", 0), tvm.tirx.IntImm("int32", 64)) + end = tvm.tirx.Add(zero, tvm.tirx.IntImm("int32", 64)) + + tile = _extract_tile(SwizzleLayout(3, 3, 3), [(zero, end)]) + + assert [int(it.extent) for it in tile.shard] == [64] + + def test_apply_to_shape(): """``apply_to_shape`` should give per-shard coord, preferring per-dim split when the input shape aligns with the layout's grouping."""