Barrier Visualization

Barrier Visualization Report

When the user asks for a barrier visualization report, produce a structured analysis of barrier usage in the given IR (either from a file, an IR dump, or from running a compilation with MLIR_ENABLE_DUMP). The report has five sections. Use the IR directly as input -- read the file or dump and analyze it.

Report Format

Section 1: Partition Summary

Label each partition by its key ops -- the operations that differentiate it. Use short descriptive names. When multiple partitions contain similar ops, add qualifying detail.

Format as a table:

| Partition   | Role             | Key Ops                        | Warps |
|-------------|------------------|--------------------------------|-------|
| default     | Acc correction   | tmem_load, tmem_store          | 4     |
| partition0  | MMA              | tc_gen5_mma x2                 | 4     |
| partition1  | TMA loads (Q,K,V)| async_tma_copy_global_to_local | 1     |
| partition2  | Output store     | descriptor_store               | 1     |
| partition3  | Softmax (QK_1)   | tmem_load, exp2, reduce        | 2     |

How to identify key ops:

• MMA partition: contains tt.dot, warp_group_dot, tc_gen5_mma, or tc_gen5_mma_scaled
• TMA load partition: contains async_tma_copy_global_to_local or descriptor_load feeding local_alloc
• Store/epilogue partition: contains descriptor_store, tt.store, tmem_load at loop exit
• Softmax/reduction partition: contains tt.reduce, math.exp2, arith.maxf
• Accumulator correction: contains tmem_load + tmem_store (re-scaling accumulators)

When two partitions both do TMA loads, differentiate by what they load:

• "TMA load (Q, K)" vs "TMA load (V, scales)"
• Use loc metadata or tensor shapes to identify operand names when available

Section 2: Barrier Dependency Graph

Draw an ASCII diagram showing which partitions produce/consume through each barrier. Use arrows to show data flow direction.

Barrier Dependency Graph
========================

  Forward barriers:

  partition1 (TMA loads)
      |
      | barrier_expect + async_tma_copy (mbarrier, SMEM buffers A, B)
      v
  partition0 (MMA)
      |
      | tc_gen5_commit (mbarrier on TMEM result)
      v
  partition3/4 (Softmax)
      |
      | aref.put / aref.get  (SMEM buffer for P)
      v
  partition0 (MMA, 2nd use)
      |
      | tc_gen5_commit
      v
  partition2 (Output store)

  Backwards barriers (next-iteration dependencies):

  partition2 (Output store)
      |
      | TMEM token (backward): tmem_load token → next iter's tmem_store
      v
  partition0 (MMA, next iteration)

  partition0 (MMA)
      |
      | mbarrier phase (backward, implicit): phase tracking prevents
      |   TMA re-arrival until MMA has consumed the buffer
      v
  partition1 (TMA loads, next iteration)

For each arrow, annotate:

• The barrier mechanism type (see table below)
• What data flows across (buffer name or tensor shape)
• The direction: forward (producer → consumer) or backward (consumer → producer, signaling resource reuse)

Backwards-Direction Barriers

In persistent kernels (those with an outer tile loop), downstream partitions often need to signal upstream partitions that shared resources can be reused. These "backwards" barriers create cycles in the dependency graph.

Common backwards barriers:

• TMEM token chain: tmem_load (epilogue) produces a token consumed by tmem_store (MMA) in the next iteration — prevents zeroing the accumulator before the epilogue finishes reading it.
• consumer_release (legacy WS): Consumer releases the mbarrier slot, allowing the producer to re-acquire it for the next iteration.
• Phase-based mbarrier: Multi-buffered SMEM implicitly handles backwards sync — the producer can't re-arrive on a slot until the consumer has waited on it (phase flip).

Show backwards barriers as upward arrows or annotated return edges in the dependency graph. When a backwards token chain is expected but the SSA token is unused (not loop-carried), flag it as a potential issue.

Column-Packed TMEM Aliasing (full-overwrite producers)

This is a high-value, easy-to-miss class. The memory planner packs several small TMEM buffers into the spare columns of a larger allocation: they share one buffer.id but carry different buffer.offset values (e.g. a 128x128 QK accumulator at offset 0 with alpha/m_ij/l_i0 scalars packed at columns 64/65/66). Unlike merged barriers (one barrier protecting several buffers), each column-packed channel gets its own independent token. Every token is therefore individually arrive/wait-balanced, so the per-token checks in Sections 3-4 all pass even when the kernel races.

The hazard appears when the owner of the allocation (the channel whose alloc has NO buffer.offset) is produced by a full-overwrite producer — a tc_gen5_mma with useC=false / useAccumulator = false, which ZEROS the entire allocation (all columns) before writing. Such a producer clobbers every packed sibling's columns, so its producer-side acquire must wait on the consumer-release of every packed sibling, not just its own channel. If a packed sibling is consumed by a different partition (e.g. the default/correction partition reads alpha/m_ij/l_i0 after the inner loop) and there is no backward edge from that consumer to the owner's producer, the next-iteration MMA overwrites the scalars mid-read — a non-deterministic data race (the FA-fwd-persistent bug).

When auditing, for each buffer.id with column-packed members:

1. Identify the owner (alloc with no buffer.offset) and the packed siblings (buffer.offset > 0).
2. Check whether the owner's producer is a tc_gen5_mma with useC=false (4th operand %false, or useAccumulator traced to a constant false). If so, it overwrites ALL columns.
3. For each packed sibling whose consumer is in a different partition than the owner's producer, verify there is a backward producer_acquire / wait_barrier on that sibling's token in the owner-producer's partition, before the owner's overwrite. Mind the cadence: if the sibling is produced/consumed at the same loop level as the MMA, the wait sits right before the MMA; if the sibling is read at an outer level (e.g. a per-tile epilogue while the MMA runs in an inner KV loop), the wait must sit before the inner loop and use the sibling's outer-loop phase. A same-cadence wait on an outer-cadence barrier (or vice versa) deadlocks rather than racing.
4. Flag a race if any such back-edge is missing — the owner's barrier alone (gating only the owner channel's own consumer) is NOT sufficient. Siblings consumed within the owner-producer's own partition are safe (program order).

This check is invisible to arrive/wait-count balancing: the missing edge is an absent barrier across physically-aliased columns, not an imbalanced one. The compiler models the required edge via isFullOverwriteReuseOwner in CodePartitionUtility.cpp; the regression IR is test/Hopper/WarpSpecialization/ws_code_partition_tmem_packed_reuse_backward.mlir.

Emit a coverage table (enumerate absences, not just presences). The reason this class slips through is that reports describe the barriers that exist; force the analysis to enumerate the barriers that should exist. For each physical buffer.id whose owner has a full/partial-overwrite producer, emit one row per aliased buffer the write touches, and mark each ✓ ordered or ✗ MISSING:

Physical buffer.id = 8 (owner: QK accumulator, 128 cols)
  Writer: tc_gen5_mma useC=false  (task 1, inner loop)  write-extent: cols 0-127
  Aliased buffers overwritten:
    cols 0-63  QK result   consumer task 5 (gemm-internal)   ✓ ordered (QK backward)
    col  64    alpha       consumer task 0 (inner cadence)   ✓ own per-iter barrier
    col  65    m_ij        consumer task 0 (outer cadence)   ✗ MISSING backward edge
    col  66    l_i0        consumer task 0 (outer cadence)   ✗ MISSING backward edge

A ✗ is a race. Always print the table even when all cells are ✓ — the table is the artifact that makes an omission visible.

This pattern rule is a manual stand-in for a future executable coverage verifier (a triton-opt pass / doCodePartitionPost invariant that models physical layout, per-op write extent, and loop cadence). See third_party/nvidia/hopper/lib/Transforms/WarpSpecialization/docs/WSAliasingCoverage.proposal.md. When that verifier lands, this section becomes "run the verifier and interpret its output." To validate this rule today, run the skill against the pre-fix ws_code_partition_tmem_packed_reuse_backward.mlir (the back-edges removed) and confirm the coverage table reports ✗ for m_ij/l_i0.

Redundant cross-partition accumulator init (persistent-loop cross-tile hazard)

This is the opposite failure mode from a deadlock false-positive: an extra channel that should not exist at all, which the per-barrier checks happily report as "balanced and correct." A GEMM/attention accumulator (TMEM operand D) is normally zero-initialized implicitly by the MMA's useAccumulator=false on the first inner-loop iteration — no explicit store, no channel. If an explicit ttng.tmem_store <zero> into that same accumulator ALSO survives, and it lives in a different partition than the MMA (e.g. the epilogue/reduction partition zeroes while the gemm partition does the useAccumulator=false MMA), then the store becomes a redundant cross-partition channel: producer = the zeroing partition, consumer = the MMA. Each of its barriers is individually arrive/wait-balanced, so Sections 3–4 pass — but in a persistent kernel (outer scf.for or scf.while) that channel is carried across tiles and the zeroing partition races the MMA/read partition from one tile to the next → a non-deterministic hang (the static-persistent-while-GEMM bug).

The tell is structural divergence + redundancy, not an imbalance:

• The accumulator carries more handshake barriers than necessary — a correct accumulator needs only FULL (MMA-commit → reader-wait) and EMPTY (reader-arrive → MMA-reuse-wait). A redundant zero-store adds a second pair (reuse-wait before the store + ready-arrive after it), i.e. ~4 accumulator barriers / 2 tc_gen5_commits instead of ~2 / 1.
• An explicit tmem_store of a constant-zero tensor into an operand-D accumulator coexists with an MMA on the same accumulator whose useAccumulator 4th operand is %false (or a loop iter-arg whose init is false).

Detection. For each TMEM operand-D accumulator (buffer.id of the MMA's accumulator, typically a 4x…xf32 #tmem alloc):

1. Find its tc_gen5_mma writer and read whether useAccumulator is false on the first iteration (4th operand %false, or an inner-loop iter-arg with %false init → %true thereafter). If so, the MMA self-zeroes.
2. Look for an ttng.tmem_store of a constant-zero tensor (arith.constant dense<0.0…>) into that same accumulator.
3. If both exist, check whether the store and the MMA are in different partitions (compare async_task_id / the ttg.partition.types region).
4. Flag a cross-tile race/hang when the redundant zero-store exists in a different partition AND the enclosing structure is a persistent loop (scf.for or scf.while). The correct IR has the store removed entirely — the MMA's useAccumulator=false is the only initializer.

This is invisible to arrive/wait balancing because the defect is a channel that should not exist, with its own perfectly-balanced barriers — not a missing or imbalanced arrive. The compiler removes it in removeRedundantTmemZeroStores (WSCodePartition.cpp); the persistent while case additionally requires that pass to recognize scf::WhileOp as the outer loop (the zero-store sits directly in the while's after region) and to forward the store's dep token (getDep() → getToken()) when erasing — otherwise the store survives and the hang returns.

Emit a coverage row per operand-D accumulator (print it even when clean):

Operand-D accumulator buffer.id = 3 (TMEM 4x128x128xf32)
  MMA writer: tc_gen5_mma useAccumulator=false (first iter)  -> self-zeroes
  Explicit zero-store present? : YES  ttng.tmem_store <0>  (task 0, epilogue)
  MMA partition / store partition : task 1 (gemm) / task 0 (epilogue)  -> DIFFERENT
  Persistent outer loop : scf.while
  Verdict: ✗ REDUNDANT cross-partition zero-store -> cross-tile race/hang
           (expected: store removed; init via useAccumulator=false only)

A ✗ here means the redundant init channel must be removed. A clean accumulator prints Explicit zero-store present? : NO -> ✓ init via useAccumulator=false. To validate this rule, run the skill on the pre-fix static-persistent while-loop GEMM TTGIR (matmul_kernel_tma_static_persistent_ws_while): the broken version shows the explicit tmem_store <0> + 4 accumulator barriers / 2 commits and must report ✗; the fixed version has 0 stores / 2 barriers / 1 commit.

Barrier Mechanism Types