taplib can't express dynamic shapes — represent access patterns in the compiler

[hunhoffe]

BLUF

Today all DMA tiling lives in the Python frontend (taplib) as compile-time integer math over concrete shapes. The dynamic-runtime-sequences work (#3222) lets one xclbin serve many problem sizes by carrying sizes/strides/offsets as runtime SSA values. taplib cannot produce those — its whole job is to enumerate the access pattern numerically before the program runs. So as #3222 lands, the high-level IRON data-movement API can't drive the dynamic designs #3222 exists to enable. (taplib stays perfectly good for static shapes — this is strictly about the dynamic ones.)

I don't want to fork the programming examples into static vs. dynamic versions unless I absolutely have to. One IRON source should compile to a static schedule when shapes are constants and a runtime schedule when they aren't — not two parallel copies of every example to maintain.

Scope: Implementation can start once #3222 Phase 2 lands (it's the lowering target); posting now for design input. This is an RFC, not a finished plan.

The one key insight

The runtime-valued unit is the loop, not the descriptor. taplib does two jobs: (a) compute one descriptor's sizes/strides/offset, and (b) iterate a sequence of them via a host-side Python for loop. The dynamic version of (b) is just an scf.for with a runtime trip count whose body emits a descriptor with an induction-derived offset.

That is exactly what the affine dialect models: affine's symbol = #3222's runtime scalar; an affine map with constant symbols folds to today's static pattern (near-free static≡dynamic equivalence for the affine-expressible subset — data-dependent indices like runtime modulus still need work, see open Q3), and with runtime symbols stays parametric and lowers to a real loop. So the proposal is to represent access patterns as affine maps + scf loops in the compiler, lowered to BDs — not to keep enumerating them in Python.

Proposal sketch

• A thin IRON-facing op (aiex.tile_access, strawman) carries one access pattern; its dims may be constants or runtime i32 SSA values.
• A lowering pass expands it to affine maps + scf loops → Dynamic runtime sequences via standalone TXN encoding + EmitC C++ generation #3222's dynamic BD encoder (constant dims fold to the existing static descriptors).
• TensorTiler2D becomes a thin emitter of this, so the same IRON code serves static and dynamic shapes.
• Static output stays byte-identical to today (guarded by Dynamic runtime sequences via standalone TXN encoding + EmitC C++ generation #3222's equivalence oracle) — existing examples don't change behavior.
• Reuse mlir-air's affine/scf representation; do not reuse its lowering (it unrolls everything to constants — the opposite of what we need).

First steps (incremental, mirroring #3222's ethos)

• Step 0 — dynamic passthrough (runtime length). One tensor, one runtime dim, no tiling. Exercises the runtime-BD spine (op → affine → dynamic BD) end-to-end on hardware. Deliberately trivial: it proves the plumbing, not the architecture.
• Step 1 — dynamic-length tiled copy. One runtime scf.for whose body emits a descriptor with an induction-derived offset. This is the real proof: it's the smallest design that actually exercises loop-not-descriptor and a runtime offset map — the crux of this whole proposal. Static shapes must stay byte-identical to today's taplib output here.

GEMM-class tiling (multiple interleaved tilers, roll/unroll mix) comes only after Step 1 works.

Where I'd like input

• Is mlir-air's affine tiling representation reusable as I describe (representation yes, unroll-to-constant lowering no), or is it more coupled to unrolling than I think? (See mlir-air notes below.)
• Does a thin aiex.tile_access + lowering pass fit cleanly alongside the existing AIEX DMA lowering (AIEDmaToNpu.cpp / AIEDMATasksToNPU.cpp)?

Open questions (need group input)

1. Roll vs. unroll classification. Some loops must unroll (e.g. a loop over columns indexing a distinct per-column FIFO) and some should roll into a runtime scf.for (problem-size loops). Inferring from int-vs-Value bound type is insufficient (see the whole_array case below) — this needs a real answer.
2. How much of taplib's generality must survive dynamically? Likely a small subset (GEMM-style M/K/N) needs runtime support first; most patterns can stay static-only. What's the minimal dynamic surface?
3. Runtime modulus isn't affine. Patterns using % len(tiles) with a runtime operand degrade from affine to arith.remui, losing fusion/analysis. How much do we care, and which patterns hit it?

---

Evidence: why taplib can't follow #3222 (grounded in code)

TensorAccessPattern (python/helpers/taplib/tap.py) stores sizes/strides/offset as concrete np.int32 sequences. TensorTiler2D (python/helpers/taplib/tensortiler2d.py) generates them with intrinsically compile-time operations:

• divisibility/guards: tensor_dim % tile_dim != 0, tile_step > tensor_dim // tile_dim
• step enumeration: num_steps = int(np.prod(steps_per_dim))
• per-step math: step_num % steps_per_row, col_idx // tile_step_height

A runtime M/N is an SSA i32, not a Python int — it can't flow through %, np.prod, int(...), or range(). The pattern is materialized eagerly at trace time (iron/runtime/data.py:53), then handed to DMATask (iron/runtime/dmatask.py) frozen. No path exists for a runtime scalar to reach the tiler.

Evidence: what #3222 establishes that this builds on

• SSA-operand npu ops ([dyn-seq P1.2] Carry npu scalar op fields as SSA operands #3225): sizes/strides/offsets already flow as SSA, not attributes.
• Dynamic BD-word encoder (Phase 2): computes d0/d1/d2/iter size+stride, repeat_count, masks/shifts from runtime values — byte-identical to static for constants. This is the backend this work lowers onto.
• static≡dynamic equivalence oracle: a constant-fed dynamic lowering must be byte-identical to the static one. A tiling op inherits this gate for free.
• Legalization boundary: Dynamic runtime sequences via standalone TXN encoding + EmitC C++ generation #3222 scopes out runtime legalization (illegal-wrap splitting, >4-D unroll). This work respects the same line — derive legal patterns within hardware bounds, reject/guard otherwise.

Evidence: mlir-air proves the representation works — but its lowering is the wrong half

mlir-air's airrt-to-npu (mlir/lib/Conversion/AIRRtToNpuPass.cpp) already represents tiling as affine/scf, then resolves it all to constants:

• unrolls everything: unrollAffineFors/unrollSCFFors → loopUnrollFull (lines 2106/2147/1375), scf::forallToForLoop first (1443).
• DmaToNpuPattern (~360-435) is structurally constant-dependent: calls getConstantIntValue() on every offset/length/stride, defaults to 0 when non-constant. This is the design, not an edge.
• silent-failure trap (1456-1465): after unrolling, affine.apply(constant) doesn't auto-fold, so offsets silently default to 0 — air patches it with an explicit AffineApplyOp canonicalization.

Reuse: the affine/scf representation, forallToForLoop normalization, post-unroll metadata renumbering. Don't reuse: the unroll-to-constant lowering (erwei noted on #3222 that airrt-to-npu's goal is to satisfy a fully-static runtime sequence). Named risk: the affine.apply(constant) fold trap is a real silent-zero mode — any affine lowering here must carry the fold step and test a non-zero induction-derived offset.

Worked hard case: whole-array matmul (why roll-vs-unroll is first-class)

programming_examples/basic/matrix_multiplication/whole_array/whole_array.py:305-353:

for tb in range(ceildiv(M//m//n_aie_rows, tb_max_n_rows)):   # problem-size → roll
    for pingpong in [0, 1]:                                   # scheduling
        for col in range(n_aie_cols):                         # hardware (per-col FIFO) → unroll
            rt.drain(C_l2l3_fifos[col].cons(), C, tap=C_tiles[c_index]); c_index += 1
            for tile_row in range(current_tb_n_rows):         # problem-size → roll
                tile_offset = ((row_base + tile_row)*n_shim_mem_A + col) % len(A_tiles)
                rt.fill(A_l3l2_fifos[col].prod(), A, tap=A_tiles[tile_offset])
                rt.fill(B_l3l2_fifos[col].prod(), B, tap=B_tiles[col])

• A "tiler yields tiles in order" API cannot express this — A/B/C are consumed interleaved and in permuted order (running counter; multi-IV index with% len). An in-order single-tiler surface only covers simple cases (passthrough, elementwise).
• An "explicit loop + tile_at(expr)" API does — the loop nest is the schedule and stays explicit; only indexing changes.
• col must unroll (distinct SSA FIFO per iteration — same constraint as --dynamic-ObjFifos Investigate if --dynamic-ObjFifos would be a reasonable default #2441); tb/tile_row should roll when M is dynamic. Note col is an int yet must not roll — which is why int-vs-Value inference is insufficient and roll/unroll needs a real design (open question 1).
• % len(A_tiles) is affine only while len is constant; dynamic M makes it arith.remui (open question 3).

Keeping the Python declarative (the surface options)

The MLIR substrate must not leak into user code — and IRON already proves the pattern: range_ (python/helpers/dialects/scf.py, _for) is a Python generator that emits scf.for and already accepts a runtime Value bound. The access-pattern API should extend that same value-type polymorphism to tiling.

Surfaces, increasing power / decreasing cleanliness:

• A — tiler yields handles (cleanest): for a_tile in A.tiles(tile=(m,k), dims=(M,N)). Static dims enumerate; runtime dims emit an scf.for. Good for in-order cases only.
• B — explicit loop, tap inside (general): with rt.tiled_range(M, step=m) as i: rt.fill(..., tap=A.tile_at(i)).
• C — affine-map-as-lambda (escape hatch): A.tile_at(lambda i,j: i*stride + j). Don't lead with it.

Likely answer: A as convenience for simple kernels, B as the general mechanism for hand-scheduled ones, C as an escape hatch.

Relationship to existing work

• Depends on Dynamic runtime sequences via standalone TXN encoding + EmitC C++ generation #3222 (the compiler-side enabler).
• Validated by [WIP] Share BD size/stride verification between AIE and AIEX dialects (#2566) #3109 (shared BD size/stride/iteration verification).
• Answers the access-pattern question from [experiment] Support unplaced TileOps #2265 (closed/superseded for tile
  placement; this is its live descendant).
• Enables Support for fast dynamic dispatch #2236 (fast dynamic dispatch) and Investigate if --dynamic-ObjFifos would be a reasonable default #2441 (--dynamic-ObjFifos):
  runtime-loop data movement keeps loops rolled instead of materialized, and lets
  transfers be retiled/fused without recompiling — one xclbin, many schedules.