Restore upstream safety conditions on loop-parallel stores

[AntonOresten]

Problem

The token-ordering pass's loop-parallel store optimization (which lets a store in a for body consume the pre-loop token instead of a loop-carried one) dropped two of upstream Python's safety conditions:

1. is_iv_derived accepted any transitive IV dependence — including non-injective derivations like (i + 1) ÷ 2, i % k, min(i, c). Upstream's _filter_by_store_index requires the index to literally be the induction variable. Two iterations mapping to the same tile then store with no happens-before edge between them: a WAW race with a nondeterministic winner.
2. Upstream's may_alias_internally guard was dropped — an array whose layout maps two distinct in-bounds indices to the same memory (e.g. a zero stride, easily produced by adapting a broadcast-style wrapper) races even with an identity-IV index. The Julia port had no equivalent concept.

Fix

• is_iv_derived → injective-affine matcher: the index tuple must contain the IV composed only of addi/subi with a loop-invariant operand, muli by a nonzero constant, or exti. Julia's 1-based store always lowers the index through iv - 1, so upstream's bare-identity check would never fire here — the affine whitelist is what upstream's own TODO: allow more complex injective check anticipated. Anything else keeps the token carry (conservative = correct).
• Internal-overlap guard reinstated: ArraySpec gains a MayAliasInternally parameter. compute_array_spec derives it exactly from runtime sizes/strides (standard non-overlapping strided layout criterion — stricter than Python, which trusts a user-declared flag); hand-written specs default to false, matching upstream's signature default. get_parallel_stores resolves the store's alias root to its TileArray argument spec and bails unless the layout provably cannot self-overlap.
• sci.argtypes is threaded through the transform chain so the pass can consult argument specs.

Tests

New test/codegen/token_order.jl:

• identity-IV and 2i (affine) stores stay loop-parallel (store consumes the entry make_token, no iter_values)
• (i + 1) ÷ 2 keeps the loop-carried token + join_tokens
• may_alias_internally = true spec keeps the carry even with an identity IV
• two stores on one array per iteration keep carries
• unit tests for layout_may_alias_internally (zero/repeated/negative/padded strides) and the compute_array_spec derivation

Ran locally: codegen/* (all), types, analysis/dataflow, host/cache, device/core, examples/softmax — all pass on 2× RTX 6000 Ada.

🤖 Generated with Claude Code