diff --git a/tests/unit/test_symbolic_client.py b/tests/unit/test_symbolic_client.py index ffd99612..37bda65c 100644 --- a/tests/unit/test_symbolic_client.py +++ b/tests/unit/test_symbolic_client.py @@ -142,6 +142,20 @@ def test_loop_sites_in_different_files_get_distinct_iterator_state(): assert f"loop_i_{site_b}" in str(ctx_b.iterator_constraint) +def test_loop_index_int_is_the_current_iteration(): + client = _LoopSiteSymbolicClient() + iterable = RangeWrapper(range(2), length=2, start=0, stop=2, step=1) + site = LoopSite(7, loop_file_token("/src/kernel.py")) + + client._loop_hook_before(site, iterable) + idx = SymbolicExpr.from_value(client._loop_hook_iter_overrider(site, 1)) + # Triton 3.6 indexes a tuple through int(tensor.handle.data); the loop + # index's Z3 form is a free variable, so int() must use the iteration. + assert isinstance(idx, SymbolicExpr) + assert int(idx.data) == 1 + client._loop_hook_after(site) + + # ======== Reduce Operations Tests ========= diff --git a/tilelens/clients/symbolic_engine.py b/tilelens/clients/symbolic_engine.py index af435c88..0207d99d 100644 --- a/tilelens/clients/symbolic_engine.py +++ b/tilelens/clients/symbolic_engine.py @@ -306,10 +306,10 @@ def coerce_int(val: Any) -> int: ) def __int__(self) -> int: - int_val, _ = self.symbolic_expr.eval() - if isinstance(int_val, list): - int_val = int_val[0] - return self.coerce_int(int_val) + # Triton 3.6's tensor.__index__ calls int(handle.data) where newer + # releases call handle.data.squeeze(). Concretize on both paths so a + # loop index or program id yields its current value, not a free Z3 var. + return self.coerce_int(self._scalar_data()) def __bool__(self) -> bool: return bool(self._scalar_data().item())