From 954af07861606506989c2aba4d5506084c92a758 Mon Sep 17 00:00:00 2001 From: waynehacking8 Date: Fri, 19 Jun 2026 16:10:46 +0800 Subject: [PATCH] [CuTeDSL] Fix cute.is_major crash on List[int] mode is_major is typed `mode: Union[int, List[int]]`, but it built its lookup as `get(stride, mode=[mode])`, which double-lists a list argument. An int mode worked (`[0]`), but any List[int] became a nested list (`[[0, 1]]`) and raised: TypeError: '<=' not supported between instances of 'int' and 'list' so every list-mode call crashed. The two internal callers only ever pass an int (guarded by `mode == 0`), which is why it went unnoticed. Normalize an int mode to a list before the lookup, mirroring the existing pattern in E(). The int path is unchanged and List[int] modes now resolve. Adds test/python/CuTeDSL/test_is_major_list_mode.py. Co-authored-by: Claude Signed-off-by: waynehacking8 --- python/CuTeDSL/cutlass/cute/core.py | 4 +- .../python/CuTeDSL/test_is_major_list_mode.py | 44 +++++++++++++++++++ 2 files changed, 47 insertions(+), 1 deletion(-) create mode 100644 test/python/CuTeDSL/test_is_major_list_mode.py diff --git a/python/CuTeDSL/cutlass/cute/core.py b/python/CuTeDSL/cutlass/cute/core.py index c2648551cb..f2c8957580 100644 --- a/python/CuTeDSL/cutlass/cute/core.py +++ b/python/CuTeDSL/cutlass/cute/core.py @@ -2191,7 +2191,9 @@ def is_major( """ Check whether a mode in stride is the major mode. """ - first_stride = front(get(stride, mode=[mode], loc=loc, ip=ip), loc=loc, ip=ip) + if isinstance(mode, int): + mode = [mode] + first_stride = front(get(stride, mode=mode, loc=loc, ip=ip), loc=loc, ip=ip) if is_dynamic_expression(first_stride): return False return True if first_stride == 1 else False diff --git a/test/python/CuTeDSL/test_is_major_list_mode.py b/test/python/CuTeDSL/test_is_major_list_mode.py new file mode 100644 index 0000000000..471be82929 --- /dev/null +++ b/test/python/CuTeDSL/test_is_major_list_mode.py @@ -0,0 +1,44 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NvidiaProprietary +# +# Use of this software is governed by the terms and conditions of the +# NVIDIA End User License Agreement (EULA), available at: +# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html +# +# Any use, reproduction, disclosure, or distribution of this software +# and related documentation outside the scope permitted by the EULA +# is strictly prohibited. + +""" +Unit tests for cute.is_major mode argument handling. + +is_major is typed `mode: Union[int, List[int]]`, but it wrapped its argument +as `get(stride, mode=[mode])`, which double-listed a list input and raised +`TypeError: '<=' not supported between instances of 'int' and 'list'` for any +List[int] mode. A single int mode and its single-element list form must agree, +and a hierarchical list path must not crash. +""" + +import unittest + +from cutlass.cute import is_major + + +class TestIsMajorListMode(unittest.TestCase): + def test_single_mode_int_and_list_agree(self): + for stride in ((8, 1), (1, 8)): + for m in (0, 1): + self.assertEqual( + is_major([m], stride), + is_major(m, stride), + msg=f"mode {m} on stride {stride}", + ) + + def test_hierarchical_list_mode(self): + # A multi-index path into a nested stride must resolve, not crash. + self.assertIsInstance(is_major([0, 1], ((8, 1), 4)), bool) + self.assertIsInstance(is_major([0, 0], ((8, 1), 4)), bool) + + +if __name__ == "__main__": + unittest.main()