Improve source-location granularity for ROCprof ATT traces

[jhinpan]

Shout out to FlyDSL team's effort first.

We are looking at ROCprof ATT source mapping for a FlashAttention kernel. Debug
info is present in the capture, but the line attribution is still too coarse to
navigate the ping-pong / double-buffer schedule from ROCprof Compute Viewer.

References

• Trace bundle: https://github.com/jhinpan/flydsl-kernel-profiling/tree/21f15497bb4a00c81fcb4283e0442977d98a4f71/examples/flash_attn_func
• Analysis report: https://github.com/jhinpan/flydsl-kernel-profiling/blob/21f15497bb4a00c81fcb4283e0442977d98a4f71/examples/flash_attn_func/REPORT.md
• Big ATT code.json: https://github.com/jhinpan/flydsl-kernel-profiling/blob/21f15497bb4a00c81fcb4283e0442977d98a4f71/examples/flash_attn_func/att_viewer/big/ui_output_agent_38430_dispatch_44/code.json
• Captured source file: https://github.com/jhinpan/flydsl-kernel-profiling/blob/21f15497bb4a00c81fcb4283e0442977d98a4f71/examples/flash_attn_func/att_viewer/big/ui_output_agent_38430_dispatch_44/source_0_flash_attn_func.py
• Related source-map change we noticed: 9f29c0d

What we confirmed

• The trace was captured from a fresh FLYDSL_RUNTIME_CACHE_DIR.
• FLYDSL_DEBUG_ENABLE_DEBUG_INFO=1 was set before discovery and ATT capture.
• The captured code object has non-empty DWARF line tables.
• code.json maps 2069 / 2070 ISA rows to Python source.

So this does not look like the earlier "missing debug info because a no-debug
HSACO was cached first" failure mode.

Observed behavior

The Python source clearly expresses the schedule we want to inspect:

• K alternates between current/next LDS buffers:
  https://github.com/jhinpan/flydsl-kernel-profiling/blob/21f15497bb4a00c81fcb4283e0442977d98a4f71/examples/flash_attn_func/att_viewer/big/ui_output_agent_38430_dispatch_44/source_0_flash_attn_func.py#L643-L698
• V is waited on / made visible before GEMM2:
  https://github.com/jhinpan/flydsl-kernel-profiling/blob/21f15497bb4a00c81fcb4283e0442977d98a4f71/examples/flash_attn_func/att_viewer/big/ui_output_agent_38430_dispatch_44/source_0_flash_attn_func.py#L970-L980
• GEMM2 pre-reads V packs before issuing MFMA:
  https://github.com/jhinpan/flydsl-kernel-profiling/blob/21f15497bb4a00c81fcb4283e0442977d98a4f71/examples/flash_attn_func/att_viewer/big/ui_output_agent_38430_dispatch_44/source_0_flash_attn_func.py#L1067-L1112

However, many hot instructions in code.json are attributed only to broad
wrapper or kernel-body locations, for example:

• flash_attn_func.py:257, the kernel body/decorator area:
  https://github.com/jhinpan/flydsl-kernel-profiling/blob/21f15497bb4a00c81fcb4283e0442977d98a4f71/examples/flash_attn_func/att_viewer/big/ui_output_agent_38430_dispatch_44/source_0_flash_attn_func.py#L257
• flash_attn_func.py:283, the _mfma helper:
  https://github.com/jhinpan/flydsl-kernel-profiling/blob/21f15497bb4a00c81fcb4283e0442977d98a4f71/examples/flash_attn_func/att_viewer/big/ui_output_agent_38430_dispatch_44/source_0_flash_attn_func.py#L282-L283

That makes it difficult to jump from ROCprof back to the actual scheduling
sites around coop_dma_k, _next_k_buf_id, coop_dma_v, wait/barrier
placement, LDS reads, and GEMM2 V pre-read.

Questions

1. For helpers that eventually emit DMA-to-LDS loads, s_waitcnt, barriers,
   LDS reads, or MFMA calls, should the public wrapper capture the user call
   site once and pass loc down to the leaf ODS / inline-asm emitters?
2. If a loc is already provided, should nested helper wrappers avoid calling
   _caller_location() again, so they do not remap the op to helper-internal
   lines?
3. For multi-op helpers, is it better to attribute all emitted ops to the outer
   user call site, or would FlyDSL prefer a more structured way to annotate
   sub-ops within the helper?
4. Are there specific wrappers in the DMA / LDS / MFMA path where maintainers
   would welcome targeted source-location propagation patches?

We are not asking to add @traced_op everywhere blindly. We would like to
improve ROCprof source navigation for double-buffer scheduling while avoiding
misleading mappings or behavior changes from eager argument unwrapping. Much thanks!!

[jhinpan]

We dug into this and put up a fix: #593.

Root cause. The schedule emits its DMA / MFMA / wait / barrier ops from nested helper closures (coop_dma_k, mfma_acc→_mfma, _waitcnt_vm_n, _read_v_pack), so traced_op's _caller_location(depth=1) resolves to the helper body — that's why every matmul lands on :283 and the DMA/waits fall through to :257. Quantified on the committed big/dispatch_44: 91.9 % of kernel ISA on :257, 6.3 % on :283, only 7 distinct source lines.

Fix (#593). A source_loc scope that pins the user call-site location once at the outer helper; traced_op honors the pin (loc or _pinned_loc() or _caller_location(1)) so untraced ODS builders inherit it via MLIR's ambient location and traced leaves via the pin. Plus loc-aware gpu.barrier / s_waitcnt / sched_* for the inline sync calls. Five one-line @source_loc_scope decorators on the kernel; no @traced_op on the multi-arg helpers (avoids the eager-unwrap you flagged).

Answering directly:

1. Yes — capture the call site once at the wrapper and let everything inherit it; no per-leaf loc=.
2. Yes — traced_op skips _caller_location when a scope/loc is already in effect, so nested helpers stop remapping to their own lines.
3. Attribute all of a helper's ops to the outer user call site, via a with-scope / decorator. That's what's actionable for the schedule.
4. The DMA/LDS/MFMA leaves are covered by the scope + the loc-aware sync wrappers.

Result (recaptured, after bundle, code.json): distinct mapped lines 7 → 20, MFMA un-collapsed from :283 to GEMM1/GEMM2 call sites (740/741, 1111/1112), and the K DMA double-buffer / V-visibility waits / GEMM2 V pre-read now resolve to the lines that express them. Pure metadata: debug-off ISA is byte-identical and perf is unchanged (372 TFLOPS).

The residual on the function line is backend-inserted s_nop/s_waitcnt (no source line) plus the inline softmax/index arithmetic — the latter is what 9f29c0de handles at the arith layer, complementary to this. We can rebase #593 onto that branch or fold the arith-loc work in — your call.

[jhinpan]

Update: extended the fix to the softmax compute as well (PR #593, 36e9f0f6) — @source_loc_scope on the _fadd/_fsub/_fmul/_fmax helpers, so the online-softmax running-max / exp-sum and the o_acc rescale now map to lines 922–961. Distinct flash_attn_func.py lines 6 → 29. Refreshed before/after bundle + numbers: REPORT.md.

Straight about what's left: ~54 % of sampled weight still sits on the function line — ~35 % backend s_nop/s_waitcnt (no source line), ~7 % exp2 + select (ArithValue framework methods), ~12 % kernel ops not yet scoped (the epilogue O global_store is the biggest). So the schedule and the softmax compute are navigable; full per-instruction coverage would need a bit more scoping plus the 9f29c0de arith layer.

[jhinpan]

Solved by #586