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()