fix(analysis): terminate when an insert_slice offset is a product of a loop induction variable

[zhen8838]

Why

• analyze does not finish on a program whose insert_slice destination offset is a product of a loop induction variable. The same program with that one offset made literal finishes in 7 s.
• It is not a big program: one mesh, one row loop, one strip-group loop of 3 trips, 48 matmuls.

R0_literal              7.4 s   ideal-ns=139407
R1_weight_only         19.0 s   ideal-ns=418219      # sb only in the ConstTensor window
R2_store_only        >300.0 s   (timeout)            # sb only in the insert_slice offset
R3_both              >300.0 s   (timeout)

What

• sb = group * NS with group a range induction variable.
• Used as a window base on a ConstTensor (w[sb:sb + NS, ...]) it costs 2.6x and completes.
• Used as an insert_slice destination offset (insert_slice(acc, o, (sb, base, 0))) it does not complete.
• base = t + 0 in the same offset tuple is fine, and sb = group + 0 in the same position is fine, so it is the product and the destination together, not symbolic offsets as such.

Contract

• No behavioural contract: the analysis is correct when it terminates. What changes is whether a legal program can be analysed at all.
• Consequence for authors: a wide projection cannot be written as one stage looping over its strip groups; it has to be split into one stage per group so the offset stays literal. That is a real structural constraint imposed by an analysis cost.

Risk

• Likely the polyhedral footprint of a write whose base is a non-affine-looking product; if so, the fix is to recognise induction * constant as affine rather than to widen the solver budget.
• --memory is required for --roofline/--performance, so there is no flag combination that avoids it.

Reproducer

tools/issues/analyze_symbolic_offset.py, inlined — four cases from one template, differing only in __WIN__ / __STORE__:

from __future__ import annotations
from tilefoundry import func, module
from tilefoundry.dsl import ConstTensor, DimVar, Mesh, Tensor, tf
from tilefoundry.dsl.tf import *  # noqa: F401, F403
from tilefoundry.ir.types.shard import Topology
from tilefoundry.target import CudaTarget

H, GRID, HP, NHS = 2048, 128, 128, 16
NS, NM, RTG, NGROUP = 32, 4, 256, 3
NOUT = NGROUP * NS
S = DimVar("seq_len", 0, 8193)
_H200, _CTA = CudaTarget("nvidia.h200_sxm"), Topology("cta", GRID)


@module(entry="proj", target=_H200, topologies=(_CTA,))
class Case:
    @func
    def proj(x: Tensor[(NHS, S, HP), "bf16"],
             w: ConstTensor[(NOUT, H, HP), "bf16"],
             out: Tensor[(NOUT, S, HP), "bf16"]) -> Tensor[(NOUT, S, HP), "bf16"]:
        with Mesh(("cta",), layout=(NS, NM), names=("strip", "tile")) as m:
            acc = out
            for t in tile(S, RTG):
                base = t + 0
                for gp in range(NGROUP):
                    gi = gp + 0
                    sb = gi * NS
                    xs = tf.reshard(x[0:1, base:base + RTG, :],
                                    (1, RTG @ m.tile, HP), "smem")
                    ws = tf.reshard(w[__WIN__:__WIN__ + NS, 0:HP, :],
                                    (NS @ m.strip, HP, HP), "smem")
                    o = tf.matmul(xs, ws)
                    for s1 in range(1, NHS, 1):
                        si = s1 + 0
                        kb = si * HP
                        xk = tf.reshard(x[si:si + 1, base:base + RTG, :],
                                        (1, RTG @ m.tile, HP), "smem")
                        wk = tf.reshard(w[__WIN__:__WIN__ + NS, kb:kb + HP, :],
                                        (NS @ m.strip, HP, HP), "smem")
                        o = o + tf.matmul(xk, wk)
                    acc = tf.insert_slice(acc, o, (__STORE__, base, 0))
            return acc

case	__WIN__	__STORE__
R0_literal	0	0
R1_weight_only	sb	0
R2_store_only	0	sb
R3_both	sb	sb

tilefoundry analyze case.py:Case /tmp/r.txt \
  --compute-cost --memory --roofline --performance --dim seq_len=8192

tilefoundry 0.0.2.dev43+g45032b967 (45032b9), target nvidia.h200_sxm.

[zhen8838]

Re-ran the reproducer with a 900 s budget instead of 300 s; both symbolic-store cases are still unfinished at 900 s, so it is not a slow-but-finite cost:

R0_literal                7.4 s  ideal-ns=139407
R1_weight_only           19.0 s  ideal-ns=418219
R2_store_only         >900.0 s  (timeout)
R3_both               >900.0 s  (timeout)

[zhen8838]

Not fixed in #123 — it wants a maintainer's call, and the PR's Risk section says why. Referenced there so it does not get lost.

[zhen8838]

Verified fixed on main at 5c763de (0.0.2.dev37). All four cases from the
issue's template, --compute-cost --memory --roofline --performance --dim seq_len=8192, 400 s budget:

case	__WIN__	__STORE__	before	now	ideal-ns now
R0_literal	0	0	7.4 s	3.81 s	139407
R1_weight_only	sb	0	19.0 s	3.85 s	418219
R2_store_only	0	sb	>900 s	3.80 s	139407
R3_both	sb	sb	>900 s	3.90 s	418219

Two things worth recording beyond "it terminates":

• The ideal-ns column pairs as R0=R2 and R1=R3, so the insert_slice
  destination offset being symbolic no longer changes the answer at all — only
  the window base does, which is the side that legitimately reads more weight.
• The fix took the shape this issue's Risk section predicted. analysis/affine.py
_loop_term now recognises induction * constant as affine directly
  (the BinaryKind.MUL branch, affine.py:94-108: it folds the constant into
  stride * factor), rather than widening a solver budget. loop_affine_term
  reaches the footprint through analysis/scope.py:269.

So the structural constraint this issue described — a wide projection having to
be split into one stage per group to keep the offset literal — is lifted.

Closing as fixed.