Skip to content

fix(analysis): make access relations state what their Ops read - #210

Merged
zhen8838 merged 25 commits into
tile-ai:mainfrom
zhen8838:fix/access-relation-contracts-207
Oct 5, 2026
Merged

zhen8838 merged 25 commits into
tile-ai:mainfrom
zhen8838:fix/access-relation-contracts-207

Conversation

@zhen8838

@zhen8838 zhen8838 commented Oct 5, 2026 •

Copy link
Copy Markdown
Collaborator

Why

  • Fixes fix(ir): make five access relations state what their Ops read #207
  • Five registered access relations stated reads and writes their Ops do not perform: full_like read its template, ShapeOf/Rank/Local had wrong boundary ranks, RoPE accepted q/k of different ranks until isl failed, the Reduce workspace was addressed by source coordinates, and Copy rejected shard shapes its verifier accepts.
  • Placing them correctly needed one IR→isl conversion API and a real layout placement (positions_of dropped merged axes and folded out-of-range coordinates onto legal positions).
  • Correct numbers then exposed two existing defects: MatMul gave a replicated result its lhs layout domain, and the first-solution allocator seed overshot the live-byte bound (including a false smem capacity error in the schedule tutorial).

What

  • One carrier: AccessRelation(relation, values). A handler returns a flat tuple, one per argument then one per result field; AffineAccess/BoundaryRelation/AccessRelations and the relation_of/restricted_access/boundary_maps unpacking helpers are removed.
  • ir/isl_interop.py: one IslParamValues dictionary; dim_to_isl_pw_aff, shape_to_isl_set, isl_to_dim, and layout_to_isl_map (flattened layout positions, nested shards per layer, mesh coordinates as parameters). Parameters are matched by value identity and named p<number>, never off a value's name; iterating renames each Op's parameters by identity into new relations before combining them.
  • projected places shard boundaries with layout_to_isl_map, clips logical coordinates first, and fixes a requested unit (device, default 0).
  • readnone_relations/unread_access for metadata Ops; full_like reads nothing and writes everything (cost evaluator synced).
  • RoPE rejects q/k rank mismatch at typeinfer; Reduce workspace reaches one slot per warp and verify checks its size; Copy relates same-buffer shard shapes by a row-major reshape.
  • MatMul type inference builds an unsharded result layout over its own shape.
  • Allocation seed: first-fit in fixed orders, each checked against every address constraint; the most compact valid seed bounds the search, and an invalid one never does.
  • Extended existing tests with hand-computed guards (tile bytes, 2·M·N·K flops, live-byte peaks, independent WGMMA max(tensor, smem) time); regenerated affected goldens and tutorials.

Contract

  • docs/spec/types.md §11 (new ISL interoperability section), visitor-registry.md, semantic-analysis.md: AccessRelation tuple order, fixed p<number> names, identity within one conversion context, layout placement, readnone relations.
  • docs/spec/hir.md, tir.md: RoPE same-rank q/k; Reduce workspace size; Copy between shard shapes; FullLike reads only its template's type.
  • docs/spec/analysis.md: a seed is an internal candidate placement (addresses and peak); one that bounds the allocator must meet every address constraint.
  • Memory reports change where old numbers were wrong: whole smem tiles, result-sized MatMul buffers, lower placement peaks; KNOWN_OVER_BOUND is now empty.

Risk

  • FullLike no longer aliases its template in allocation (its old proof relied on the removed read); peaks still meet the live bound.
  • The allocator seed improves fixed cases, not every graph; the solver still stops at its first feasible solution.
  • test_wgmma_performance...[n=16] agrees with the model at ceil(10.93 ns) vs ceil(10.09 ns) because of different clock references.

Comment thread docs/spec/analysis.md
Comment thread docs/spec/types.md Outdated
Comment thread src/tilefoundry/ir/hir/nn/matmul.py Outdated
Comment thread src/tilefoundry/ir/hir/schedule.py Outdated
Comment thread src/tilefoundry/ir/isl_interop.py Outdated
Comment thread src/tilefoundry/ir/tir/memory/copy.py
Comment thread tests/analysis/test_analyze_at_a_size.py Outdated
Comment thread tests/analysis/test_analyze_by_hand.py Outdated
Comment thread tests/ops/ir/test_local.py Outdated
…sRelation tuple

One AccessRelation (relation + values) replaces AffineAccess, BoundaryRelation and AccessRelations. A handler returns one relation per argument, then one per result field; readers split by len(call.args). relation_of, restricted_access, boundary_maps and _affine_boundaries are gone, and callers read .relation or pick a transformed map's bindings where they compose it.
_fresh_name no longer reads a DimVar, Var or mesh axis name or cleans it with a regex; _by_identity gives a value that cannot keep its name the first free p<number> instead of <name>_<suffix>. The isl utility test now covers unspellable names, taken p names, identity across accesses, alias merging, swaps and reuse of one relation by two Ops.
Delete _unsharded_layout and choose the result layout where MatMul infers it. Move the isl_interop contract out of the dim section into its own types section and point its readers there, define the allocation seed, drop the persistent_gemm_tiled module alias, and state the 64 KiB smem expectation as the sum of its three bf16 tiles.
@zhen8838
zhen8838 merged commit f71ed0e into tile-ai:main Oct 5, 2026
2 checks passed
@zhen8838
zhen8838 deleted the fix/access-relation-contracts-207 branch October 5, 2026 10:20
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

fix(ir): make five access relations state what their Ops read

1 participant