Enhance tuple unpacking

Signed-off-by: Greg Bonik <gbonik@nvidia.com>

Author: gbonik

changelog.d/tuple-unpacking.md

@@ -0,0 +1,9 @@
+<!--- SPDX-FileCopyrightText: Copyright (c) <2025> NVIDIA CORPORATION & AFFILIATES. All rights reserved. -->
+<!--- SPDX-License-Identifier: Apache-2.0 -->
+
+- Added a missing check for unpacking a tuple with too many values. For example, ``a, b = 1, 2, 3``
+  now raises an error, instead of silently discarding the extra value.
+- Added support for unpacking nested tuples (e.g, ``a, (b, c) = t``) , as well as using square
+  brackets for unpacking (e.g., ``[a, b] = 1, 2``).
+- Fixed the missing column indicator in error messages when the underlined text is only one
+  character wide.

src/cuda/tile/_exception.py

@@ -99,7 +99,7 @@ def _format_location_frame(loc: Loc) -> str:
 
     visual_col = _wcwidth(line_bytes[:loc.col].decode())
     if end_col == loc.col + 1:
-        end_visual_col = visual_col
+        end_visual_col = visual_col + 1
         cols_str = f"col {visual_col + 1}"
     else:
         end_visual_col = _wcwidth(line_bytes[:end_col].decode())

src/cuda/tile/_ir/hir.py

@@ -221,6 +221,7 @@ def __call__(self, x: Operand) -> str:
 def if_else(cond, then_block, else_block, /): ...
 def loop(body, iterable, /): ...  # infinite if `iterable` is None
 def build_tuple(*items): ...  # Makes a tuple (i.e. returns `items`)
+def unpack(iterable, expected_len, /): ...
 def identity(x): ...   # Identity function (i.e. returns `x`)
 def store_var(name, value, /): ...  # Store into a named variable
 def load_var(name, /): ...  # Load from a named variable

src/cuda/tile/_ir/ops.py

@@ -1341,6 +1341,22 @@ def build_tuple(items: tuple[Var, ...]) -> Var:
     return res
 
 
+@impl(hir.unpack)
+def unpack_impl(iterable: Var, expected_len: Var) -> Var:
+    ty = iterable.get_type()
+    # Don't use the require_tuple_type() helper because we'd like to customize the error message
+    if not isinstance(ty, TupleTy):
+        raise TileTypeError("Expected a tuple", iterable.loc)
+    expected_len = require_constant_int(expected_len)
+    if len(ty.value_types) != expected_len:
+        few_many = "few" if len(ty.value_types) < expected_len else "many"
+        raise TileValueError(f"Too {few_many} values to unpack"
+                             f" (expected {expected_len}, got {len(ty.value_types)})")
+    # Return the input tuple. If we add support for additional iterables,
+    # the idea is to cast them to a tuple here.
+    return iterable
+
+
 class Unary(TypedOperation):
     def __init__(self, fn: str, operand: Var,
                  rounding_mode: Optional[RoundingMode], flush_to_zero: bool,

src/cuda/tile/_passes/ast2hir.py

@@ -409,13 +409,11 @@ def _do_assign(value: hir.Operand, target, ctx: _Context):
     with ctx.change_loc(target):
         if isinstance(target, ast.Name):
             ctx.store(target.id, value)
-        elif isinstance(target, ast.Tuple):
+        elif isinstance(target, ast.Tuple | ast.List):
+            value = ctx.call(hir.unpack, (value, len(target.elts)))
             for i, el in enumerate(target.elts):
-                with ctx.change_loc(el):
-                    if not isinstance(el, ast.Name):
-                        raise ctx.unsupported_syntax()
-                    item_var = ctx.call(operator.getitem, (value, i), )
-                    ctx.store(el.id, item_var)
+                item_var = ctx.call(operator.getitem, (value, i), )
+                _do_assign(item_var, el, ctx)
         else:
             raise ctx.unsupported_syntax()
 

test/test_unpacking.py

@@ -0,0 +1,42 @@
+# SPDX-FileCopyrightText: Copyright (c) <2026> NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
+import pytest
+import torch.cuda
+
+import cuda.tile as ct
+from cuda.tile import TileValueError
+
+
+def test_too_few_values_to_unpack():
+    @ct.kernel
+    def kernel():
+        t = 1, 2
+        a, b, c = t
+    with pytest.raises(TileValueError, match="Too few values to unpack"):
+        ct.launch(torch.cuda.current_stream(), (1,), kernel, ())
+
+
+def test_too_many_values_to_unpack():
+    @ct.kernel
+    def kernel():
+        t = 1, 2, 3
+        a, b = t
+    with pytest.raises(TileValueError, match="Too many values to unpack"):
+        ct.launch(torch.cuda.current_stream(), (1,), kernel, ())
+
+
+def test_unpack_nested_tuple():
+    @ct.kernel
+    def kernel(x):
+        t = (1, (2, 3, 4)), 5
+        [(a, [b, c, d]), e] = t
+        ct.scatter(x, 0, a)
+        ct.scatter(x, 1, b)
+        ct.scatter(x, 2, c)
+        ct.scatter(x, 3, d)
+        ct.scatter(x, 4, e)
+
+    x = torch.zeros((5,), dtype=torch.int32, device="cuda")
+    ct.launch(torch.cuda.current_stream(), (1,), kernel, (x,))
+    assert x.tolist() == [1, 2, 3, 4, 5]