From fb89c4d47d9e88733df1d95a7bcbee57a68d5675 Mon Sep 17 00:00:00 2001 From: adityas-amd Date: Sat, 18 Jul 2026 00:05:29 -0400 Subject: [PATCH 1/3] Fix Vector list()/iteration hang (#793) Vector defined __getitem__ but neither __iter__ nor __len__, so list(Vec(buffer_load(vec_width>1))) fell back to CPython's legacy sequence protocol (v[0], v[1], ...). That protocol only stops on IndexError, and the integer branch of __getitem__ never raised one, so tracing spun forever emitting vector.extract ops and never finalized any IR. - Add __len__ (returns numel) and __iter__ (yields the numel lanes). - Bounds-check the integer branch of __getitem__: normalize negative indices and raise IndexError when out of range. - Add tests/unit/test_vector_iteration.py regression coverage. --- python/flydsl/expr/typing.py | 14 ++++++ tests/unit/test_vector_iteration.py | 75 +++++++++++++++++++++++++++++ 2 files changed, 89 insertions(+) create mode 100644 tests/unit/test_vector_iteration.py diff --git a/python/flydsl/expr/typing.py b/python/flydsl/expr/typing.py index eaaa3384a..da27c9559 100644 --- a/python/flydsl/expr/typing.py +++ b/python/flydsl/expr/typing.py @@ -1843,6 +1843,10 @@ def __getitem__(self, idx): if idx is None: return self if isinstance(idx, int): + if idx < 0: + idx += self.numel + if not 0 <= idx < self.numel: + raise IndexError(f"vector index {idx} out of range for numel {self.numel}") res = _vector.ExtractOp(self, static_position=[idx], dynamic_position=[]).result return self._dtype(res) if isinstance(idx, (Numeric, ArithValue, ir.Value)): @@ -1880,6 +1884,16 @@ def __getitem__(self, idx): return self._build_result(res, res_shape, row_major=True) raise TypeError(f"unsupported index type: {type(idx)}") + def __len__(self) -> int: + return self.numel + + def __iter__(self): + # Without this, list()/unpacking falls back to the legacy sequence + # protocol (self[0], self[1], ...), which spins forever because the + # int branch of __getitem__ used to never raise IndexError (see #793). + for i in range(self.numel): + yield self[i] + def _build_result(self, value, shape, *, row_major=False) -> "Vector": shape = self._canonical_shape(shape) flat_ty = self.make_type(shape, self._dtype) diff --git a/tests/unit/test_vector_iteration.py b/tests/unit/test_vector_iteration.py new file mode 100644 index 000000000..7c76434f3 --- /dev/null +++ b/tests/unit/test_vector_iteration.py @@ -0,0 +1,75 @@ +#!/usr/bin/env python3 + +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 FlyDSL Project Contributors + +"""Regression tests for Vector iteration / list() conversion (ROCm/FlyDSL#793). + +Previously ``list(Vec(buffer_load(vec_width>1)))`` hung during tracing: Vector +defined ``__getitem__`` but neither ``__iter__`` nor ``__len__``, so CPython fell +back to the legacy sequence protocol (``v[0], v[1], ...``) which only stops on +``IndexError`` -- and the integer branch never raised one, emitting an unbounded +stream of ``vector.extract`` ops. +""" + +import pytest + +from flydsl._mlir import ir +from flydsl._mlir.dialects import func +from flydsl.expr.typing import Vector + + +def _in_vector_func(numel, body): + """Run ``body(v)`` on a Vector over a ``vector`` func argument.""" + with ir.Context() as ctx: + ctx.allow_unregistered_dialects = True + with ir.Location.unknown(ctx): + module = ir.Module.create() + with ir.InsertionPoint(module.body): + vec_ty = ir.VectorType.get([numel], ir.F32Type.get()) + f = func.FuncOp("t", ir.FunctionType.get([vec_ty], [])) + with ir.InsertionPoint(f.add_entry_block()): + (arg,) = list(f.entry_block.arguments) + result = body(Vector(arg)) + func.ReturnOp([]) + return result + + +@pytest.mark.l0_backend_agnostic +def test_list_of_vector_terminates(): + """list(Vec) yields exactly numel scalar values instead of hanging (#793).""" + vals = _in_vector_func(4, list) + assert len(vals) == 4 + + +@pytest.mark.l0_backend_agnostic +def test_len_of_vector(): + assert _in_vector_func(4, len) == 4 + + +@pytest.mark.l0_backend_agnostic +def test_vector_unpacking(): + n = _in_vector_func(3, lambda v: len([x for x in v])) + assert n == 3 + + +@pytest.mark.l0_backend_agnostic +def test_int_index_out_of_range_raises(): + def body(v): + with pytest.raises(IndexError): + v[v.numel] + return True + + assert _in_vector_func(4, body) + + +@pytest.mark.l0_backend_agnostic +def test_negative_index(): + """v[-1] resolves to the last lane; too-negative raises IndexError.""" + def body(v): + assert v[-1] is not None + with pytest.raises(IndexError): + v[-(v.numel + 1)] + return True + + assert _in_vector_func(4, body) From 6b82202fbede91590bdc11d683c03c009bd397f3 Mon Sep 17 00:00:00 2001 From: adityas-amd Date: Sat, 18 Jul 2026 00:40:46 -0400 Subject: [PATCH 2/3] Apply black formatting to test_vector_iteration.py Add blank line after the docstring in test_negative_index to satisfy the repo's black style check (pre-checks CI). --- tests/unit/test_vector_iteration.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/unit/test_vector_iteration.py b/tests/unit/test_vector_iteration.py index 7c76434f3..2502fe452 100644 --- a/tests/unit/test_vector_iteration.py +++ b/tests/unit/test_vector_iteration.py @@ -66,6 +66,7 @@ def body(v): @pytest.mark.l0_backend_agnostic def test_negative_index(): """v[-1] resolves to the last lane; too-negative raises IndexError.""" + def body(v): assert v[-1] is not None with pytest.raises(IndexError): From b3bb5fd508c13a312e0d777f08ba0efbe14ad802 Mon Sep 17 00:00:00 2001 From: adityas-amd Date: Sat, 18 Jul 2026 00:43:19 -0400 Subject: [PATCH 3/3] Address review: real tuple-unpacking test + original index in error - Split the iteration test: keep list-comprehension coverage as test_vector_iteration and add test_vector_unpacking exercising the tuple-unpacking path (a, b, c = v). - Report the caller-supplied index in the out-of-range IndexError instead of the negative-normalized value (e.g. v[-5] now says -5, not -1). --- python/flydsl/expr/typing.py | 3 ++- tests/unit/test_vector_iteration.py | 11 ++++++++++- 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/python/flydsl/expr/typing.py b/python/flydsl/expr/typing.py index da27c9559..4b6ebcbd9 100644 --- a/python/flydsl/expr/typing.py +++ b/python/flydsl/expr/typing.py @@ -1843,10 +1843,11 @@ def __getitem__(self, idx): if idx is None: return self if isinstance(idx, int): + orig_idx = idx if idx < 0: idx += self.numel if not 0 <= idx < self.numel: - raise IndexError(f"vector index {idx} out of range for numel {self.numel}") + raise IndexError(f"vector index {orig_idx} out of range for numel {self.numel}") res = _vector.ExtractOp(self, static_position=[idx], dynamic_position=[]).result return self._dtype(res) if isinstance(idx, (Numeric, ArithValue, ir.Value)): diff --git a/tests/unit/test_vector_iteration.py b/tests/unit/test_vector_iteration.py index 2502fe452..218413e0f 100644 --- a/tests/unit/test_vector_iteration.py +++ b/tests/unit/test_vector_iteration.py @@ -48,11 +48,20 @@ def test_len_of_vector(): @pytest.mark.l0_backend_agnostic -def test_vector_unpacking(): +def test_vector_iteration(): n = _in_vector_func(3, lambda v: len([x for x in v])) assert n == 3 +@pytest.mark.l0_backend_agnostic +def test_vector_unpacking(): + def body(v): + a, b, c = v + return len((a, b, c)) + + assert _in_vector_func(3, body) == 3 + + @pytest.mark.l0_backend_agnostic def test_int_index_out_of_range_raises(): def body(v):