Keep the whole dimension for a loop index that is not affine in the loop variable

dace/frontend/python/newast.py

@@ -320,6 +320,13 @@ def _disallow_stmt(visitor, node):
 ###############################################################
 
 
+def is_affine_in(expr, sym) -> bool:
+    """ Returns True if ``expr`` is at most linear in the symbol named like ``sym``. """
+    expr = expr.expr if isinstance(expr, symbolic.SymExpr) else sympy.sympify(expr)
+    matches = [f for f in expr.free_symbols if str(f) == str(sym)]
+    return not matches or (expr.is_polynomial(matches[0]) and sympy.degree(expr, matches[0]) <= 1)
+
+
 def _rescale_by_outer_steps(irng: subsets.Range, orng: subsets.Range):
     for n, ostep in enumerate(orng.strides()):
         if ostep == 1:
@@ -765,8 +772,9 @@ def _add_access(
                 for s, sr in self.symbols.items():
                     if s in symbolic.symlist(r).values():
                         ignore_indices.append(i)
-                        if any(t in self.sdfg.arrays or t in (str(sym) for sym in self.symbols)
-                               for t in sr.free_symbols):
+                        # Offsetting by f(min(sr)) needs f affine in s; otherwise keep the whole dimension.
+                        if (any(t in self.sdfg.arrays or t in (str(sym) for sym in self.symbols)
+                                for t in sr.free_symbols) or not all(is_affine_in(b, s) for b in r[:2])):
                             sym_rng.append(subsets.Range([(0, parent_array.shape[i] - 1, 1)]))
                             repl_dict = {}
                             break
@@ -3356,8 +3364,9 @@ def _add_access(
                 for s, sr in self.symbols.items():
                     if s in symbolic.symlist(r).values():
                         ignore_indices.append(i)
-                        if any(t in self.sdfg.arrays or t in (str(sym) for sym in self.symbols)
-                               for t in sr.free_symbols):
+                        # Offsetting by f(min(sr)) needs f affine in s; otherwise keep the whole dimension.
+                        if (any(t in self.sdfg.arrays or t in (str(sym) for sym in self.symbols)
+                                for t in sr.free_symbols) or not all(is_affine_in(b, s) for b in r[:2])):
                             sym_rng.append(subsets.Range([(0, parent_array.shape[i] - 1, 1)]))
                             repl_dict = {}
                             break

tests/python_frontend/nonaffine_loop_index_test.py

@@ -0,0 +1,42 @@
+# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved.
+""" Tests that a loop-dependent index that is not affine in the loop variable keeps the whole array dimension. """
+import numpy as np
+
+import dace
+
+N = dace.symbol('N')
+
+
+@dace.program
+def rotating_table_read(a: dace.float64[N, 3], tbl: dace.int32[3], out: dace.float64[N, 3], K: dace.int64):
+    for i in dace.map[0:N]:
+        for k in range(K):
+            for j in dace.map[0:3]:
+                for t in range(k + 1):
+                    out[i, j] = out[i, j] + a[i, j] * tbl[(j + t) % 3] + t
+
+
+def test_modulo_index_in_nested_loop_passes_the_whole_table():
+    sdfg = rotating_table_read.to_sdfg(simplify=False)
+    table_shapes = [
+        n.sdfg.arrays[e.dst_conn].shape for n, state in sdfg.all_nodes_recursive()
+        if isinstance(n, dace.nodes.NestedSDFG) for e in state.in_edges(n)
+        if n.sdfg.arrays[e.dst_conn].dtype == dace.int32
+    ]
+    assert table_shapes and all(shape == (3, ) for shape in table_shapes), table_shapes
+
+    rng = np.random.default_rng(0)
+    a = rng.random((5, 3))
+    tbl = np.array([2, 3, -1], dtype=np.int32)
+    out = np.zeros((5, 3))
+    sdfg(a=a, tbl=tbl, out=out, N=5, K=4)
+    ref = np.zeros((5, 3))
+    for k in range(4):
+        for j in range(3):
+            for t in range(k + 1):
+                ref[:, j] += a[:, j] * tbl[(j + t) % 3] + t
+    assert np.allclose(out, ref)
+
+
+if __name__ == '__main__':
+    test_modulo_index_in_nested_loop_passes_the_whole_table()