Add LLIR source mapping from LLVM debug metadata by wychi · Pull Request #434 · meta-pytorch/tritonparse

wychi · 2026-09-15T06:55:24Z

Fixes #61.

Summary

The LLIR stage is registered with supports_source_mapping=True and parser_id="generic_loc" on both the NVIDIA and AMD adapters, but _parse_generic_loc only understands MLIR #loc directives. LLVM IR records locations as !dbg !N references into !DILocation metadata, so the parser silently matches nothing and source_mappings ships without an llir key — the stage is advertised as mapped and never is, with no warning.

This change adds extract_llvm_dbg_mappings() / _parse_llvm_dbg, registers the parser as llvm_dbg on the base adapter (it is LLVM, not vendor specific), and points both LLIR stage descriptors at it.

The `inlinedAt` part is the bit worth reviewing

!DILocation may carry inlinedAt, meaning the line belongs to an inlined callee in another file. tl.sum nests two frames deep:

standard.py:273  inlined at  standard.py:313  inlined at  your_kernel.py:210

Taking the innermost line naively attributes library code to the user's file. Following the existing callsite convention in extract_loc_definitions, file/line describe the callee and the call site is recorded separately; inlined_at_file / inlined_at_line give the root of the chain, which is the line in the user's kernel.

Scope

#61 lists the mappings that exist today — Python↔TTIR, Python↔TTGIR, Python↔PTX, TTIR↔TTGIR, TTIR↔PTX, TTGIR↔PTX — with LLIR absent from every pair, which is exactly what the broken generic_loc registration produces. It links tritonparse/extract_source_mappings.py; that module is now a CLI wrapper and the parsers have since moved to parse/ir_parser.py, where this change lands.

The issue asks only for LLIR line mapping. Because the cross-stage join is already generic, wiring up the parser also yields LLIR↔TTIR, LLIR↔TTGIR and LLIR↔PTX/SASS (NVIDIA) or LLIR↔AMDGCN (ROCm), in both directions, at no extra cost. Both are measured below.

Test plan

Runs on CPU from a clean checkout; the example trace inlines file_content for every stage.

python -m pytest tests/cpu -q
TRITONPARSE_FB_MODE=0 python -m tritonparse.cli parse \
    tests/example_output/logs --out /tmp/out --overwrite

Note. The checked-in tests/example_output/parsed_output_* files are from an older tritonparse and are not a valid baseline — re-parse the raw log on each branch.

1. Tests

720 passed, 1 skipped, 44 subtests (712 on main; +8).

TestLLVMDebugMappings covers a plain !DILocation, a two-level inlinedAt chain, !dbg pointing at a DISubprogram rather than a location, and non-LLVM input.

TestStageParserContract adds two guards against the class of bug this PR fixes: a mappable stage must name a parser that can actually read its syntax, and the allowed syntax_id → parser_id table is exhaustive, so a stage whose syntax is absent fails rather than passes by default. Replayed onto main, both fail and name the original bug.

2. A/B on the example trace

main has no llir key, and no entry in any stage carries llir_lines. This branch:

kernel	llir mapped	callsites
`matmul_kernel` w1	149/293	6
`matmul_kernel` w1	195/355	6
`fused_op_kernel` w4	27/74	0
`fused_op_kernel` w4	29/78	0

The existing cross-stage join then picks LLIR up in both directions with no further change. For matmul_kernel w1, entries gaining llir_lines:

stage	python	ttir	ttgir	ptx	sass
gained `llir_lines`	24/29	155/181	160/186	174/216	372/379

The callsite split matches the kernels: matmul_kernel calls tl.cdiv, which inlines; fused_op_kernel calls nothing that does.

3. Blast radius

2616 non-llir entries byte-identical to main. The only changes elsewhere are additive — the llir_lines key, and one new python key (303, tl.zeros, which emits no PTX or SASS instruction but has explicit zero-init in LLVM IR). An audit of both branches for structural invariants and per-line TTIR construct agreement reports identical counts throughout: this adds a stage rather than altering the existing ones.

4. AMD

On an 8× gfx950 host (torch 2.9.1+rocm6.4, pytorch-triton-rocm 3.4.0), real launches and a real trace:

source_mappings = [amdgcn, llir, python, ttgir, ttir] — 63 llir entries, 44 callsites, 137 amdgcn entries gaining llir_lines.

Other arches are checkable without the corresponding hardware, since Triton's AMD backend compiles ahead of time for an explicit GPUTarget:

arch		mapped	callsites
gfx90a	CDNA2	72	42
gfx942	CDNA3	72	42
gfx950	CDNA4	72	42

Inlined frames resolve as on NVIDIA. This is the expected result: the metadata comes from the shared pass lib/Target/LLVMIR/LLVMDIScope.cpp, which both backends invoke, and the parser never inspects the target triple.

Not covered: gfx1100 (RDNA3), where the AMD backend segfaults during compilation regardless of this change.

5. UI

Code Comparison, a pane set to LLIR, clicking python 287 pid = tl.program_id(axis=0):

	screenshot	LLIR pane
before	https://pxl.cl/dfrv4	nothing highlighted — no llir mapping exists
after	https://pxl.cl/dfrvc	L9 `%12 = tail call i32 @llvm.nvvm.read.ptx.sreg.ctaid.x()`

No frontend change needed: CodeComparisonView already reads llir_lines and getViewableStages() already admits the stage.

6. Lint

ruff check and ruff format --check clean.

meta-codesync · 2026-09-15T17:30:21Z

@wychi has imported this pull request. If you are a Meta employee, you can view this in D120178767.

Summary: The LLIR stage is registered with `supports_source_mapping=True` and `parser_id="generic_loc"` on both the NVIDIA and AMD adapters, but `_parse_generic_loc` only understands MLIR `#loc` directives. LLVM IR records locations as `!dbg !N` references into `!DILocation` metadata, so the parser silently matches nothing and `source_mappings` ships without an `llir` key -- the stage is advertised as mapped and never is, with no warning. Add `extract_llvm_dbg_mappings()` / `_parse_llvm_dbg`, register the parser as `llvm_dbg` on the base adapter (it is LLVM, not vendor specific), and point both LLIR stage descriptors at it. `!DILocation` may carry `inlinedAt`, meaning the line belongs to an inlined callee in another file. `tl.sum` nests two frames deep (standard.py -> standard.py -> the user's kernel), so taking the innermost line naively attributes library code to the user's file. Following the existing callsite convention in `extract_loc_definitions`, `file`/`line` describe the callee and the call site is recorded separately; `inlined_at_file`/`inlined_at_line` give the root of the chain, which is the line in the user's kernel. Fixes #61. That issue lists the mappings that exist today -- Python<->TTIR, Python<->TTGIR, Python<->PTX, TTIR<->TTGIR, TTIR<->PTX, TTGIR<->PTX -- with LLIR absent from every pair, which is exactly what the broken `generic_loc` registration produces. It links `tritonparse/extract_source_mappings.py`; that module is now a CLI wrapper and the parsers have since moved to `parse/ir_parser.py`, where this change lands. The issue asks only for LLIR line mapping. Because the cross-stage join is already generic, wiring up the parser also yields LLIR<->TTIR, LLIR<->TTGIR and LLIR<->PTX/SASS (NVIDIA) or LLIR<->AMDGCN (ROCm), in both directions, at no extra cost. Both are measured below. Test Plan: Runs on CPU from a clean checkout; the example trace inlines `file_content` for every stage. (The checked-in `parsed_output_*` files are from an older tritonparse and are not a valid baseline -- re-parse the raw log on each branch.) python -m pytest tests/cpu -q TRITONPARSE_FB_MODE=0 python -m tritonparse.cli parse \ tests/example_output/logs --out /tmp/out --overwrite Tests: 720 passed, 1 skipped, 44 subtests (712 on main; +8). TestLLVMDebugMappings covers a plain !DILocation, a two-level inlinedAt chain, `!dbg` on a DISubprogram, and non-LLVM input. TestStageParserContract is the pair of guards above; replayed onto main both fail, naming the original bug. A/B on the example trace. main has no `llir` key and no entry in any stage carries `llir_lines`. This branch: kernel llir mapped callsites matmul_kernel/w1 149/293 6 matmul_kernel/w1 195/355 6 fused_op_kernel/w4 27/74 0 fused_op_kernel/w4 29/78 0 The existing cross-stage join then picks LLIR up in both directions with no further change -- for matmul_kernel/w1, entries gaining `llir_lines`: python 24/29, ttir 155/181, ttgir 160/186, ptx 174/216, sass 372/379. The callsite split matches the kernels: matmul_kernel calls `tl.cdiv`, which inlines; fused_op_kernel calls nothing that does. Blast radius: 2616 non-llir entries byte-identical to main. The only changes elsewhere are additive -- the `llir_lines` key, and one new python key (303, `tl.zeros`, which emits no PTX or SASS instruction but has explicit zero-init in LLVM IR). An audit of both branches for structural invariants and per-line TTIR construct agreement reports identical counts throughout: this adds a stage rather than altering the existing ones. AMD, on an 8x gfx950 host (torch 2.9.1+rocm6.4, pytorch-triton-rocm 3.4.0): `source_mappings` = [amdgcn, llir, python, ttgir, ttir], 63 llir entries, 44 callsites, 137 amdgcn entries gaining `llir_lines`. Other arches are checkable without hardware, since Triton's AMD backend compiles ahead of time for an explicit GPUTarget -- gfx90a, gfx942 and gfx950 each give 72 mapped lines and 42 callsites, with inlined frames resolving as on NVIDIA. Expected: the metadata comes from the shared pass lib/Target/LLVMIR/LLVMDIScope.cpp, which both backends invoke, and the parser never inspects the target triple. Not covered: gfx1100 (RDNA3), where the AMD backend segfaults during compilation regardless of this change. UI (Code Comparison, a pane set to LLIR, clicking python 287 `pid = tl.program_id(axis=0)`): | | screenshot | LLIR pane | | ------ | -------------------- | --------- | | before | https://pxl.cl/dfrv4 | nothing highlighted -- no llir mapping exists | | after | https://pxl.cl/dfrvc | L9 `%12 = tail call i32 @llvm.nvvm.read.ptx.sreg.ctaid.x()` | No frontend change needed: `CodeComparisonView` already reads `llir_lines` and `getViewableStages()` already admits the stage. Lint: `ruff check` and `ruff format --check` clean.

meta-codesync · 2026-09-16T22:12:50Z

@wychi merged this pull request in 512addb.

Summary: Inlined Triton helpers — `tl.cdiv`, `tl.sum`, `tl.dot` — are mapped to the wrong source line. A mapping entry's `file`/`line` describe the **callee**, a line in a Triton library file, and `create_python_mapping` keys purely on `info["line"]`. So the arithmetic `tl.cdiv` generates is filed under `standard.py:43` and is unreachable from the kernel's source pane: ``` py288: num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) filed under python 288: L106: #loc5 = loc("generate_examples.py":288:21) <- declarations only L148: #loc56 = loc("num_pid_m"(#loc5)) filed under standard.py: L21: %num_pid_m = arith.addi %M, %c15_i32 <- the actual code L22: %num_pid_m_2 = arith.divsi %num_pid_m, %c16_i32 ``` Based on `main` and **independent of #434 — every path here is pre-existing and broken today. #434 added `inlined_at_file` / `inlined_at_line` for LLIR but left them unread; this makes them mean something for every stage. The two can land in either order. ## Three commits **1 — `Map inlined code back to the line the user wrote`** - **MLIR** `extract_loc_definitions`: a second pass walks each callsite's `callsite_caller` chain to the outermost frame and records `inlined_at_*`. Two passes are needed because during the first, a callsite's caller may not be in `locations` yet. The walk stops on `nxt is None` rather than a falsy check — `""` is the key of the bare `#loc`, so `loc(callsite(#locN at #loc))` is a real caller. - **SASS** `_iter_sass_instructions`: nvdisasm emits one `//## File` comment per inline frame, innermost first, and the parser kept only the innermost — discarding the user's call site. It now tracks the last frame too, plus `is_callsite`, which SASS entries never carried. - `create_python_mapping` prefers `inlined_at_line`, tested with `is not None` rather than `or`: a line of 0 is falsy and would silently re-file the entry under the callee. - `inlined_at_*` is recorded only when the walk reaches a frame that is itself not a callsite. Breaking out on a missing, cyclic or unresolvable caller leaves the cursor on a callee — a library line — and writing that as the call site would be the very mis-attribution this pass removes. A broken chain leaves the entry plain, matching the PTX path when a `.file` index will not resolve. **2 — `Fix three .loc defects on the inlined-code path`** *Inlined `.loc` was never matched.* `PTX_LOC_PATTERN` required `//` right after the column, but inlined directives carry extra fields first: ``` .loc 2 43 13, function_name $L__info_string0, inlined_at 1 288 21 // standard.py:43:13 ``` Instructions under them inherited whatever `.loc` came before, making PTX under-join on inlined code. On the example trace all 38 affected entries landed on the right python line *anyway*, because the preceding `.loc` happened to be the call site — luck, not correctness. *`inlined_at` names the immediate caller, not the root.* Every other parser resolves to the outermost frame; taking the field verbatim does not. ptxas output for a two-level `tl.sum`: ``` .loc 2 273 12, ... inlined_at 2 313 12 standard.py -> standard.py .loc 2 313 12, ... inlined_at 1 8 9 standard.py -> the kernel ``` The inner frame gets filed under `standard.py:313` — a library line — while MLIR, SASS and LLIR file the same instruction under the kernel line, and the stages silently disagree. The `.loc`'s own file index is now captured so a frame → caller map can be built, and the chain is walked to the root with a cycle guard. **Found in review by slp_generic**; the 17-kernel corpus audit missed it because `tl.cdiv` inlines one level deep, where the immediate caller and the root are the same line. *`.loc <file> 0 <col>`* is DWARF for "no line information"; read literally it invented a mapping to python line 0, so the entry looked mapped and resolved to nothing. Now skipped, and it clears the current mapping so the next instruction doesn't inherit a stale one. This check sits after the `is_ptx`/`is_amdgcn` branch, so it **applies to AMDGCN too** — previously true but neither advertised nor tested. `inlined_at_file` and `inlined_at_line` are written together or not at all. Setting the line while the `.file` index fails to resolve would let `create_python_mapping`, which keys on the line alone, file the entry under a caller line with nothing having checked the file it came from. **3 — `Highlight the user's line when clicking inlined IR`** (`website/src`) The reverse direction. `calculatePythonLines` reads the entry's own `file`/`line`, so `mapping.file.includes(pythonInfo.file_path)` is false for inlined code — the file is `language/standard.py`, not the kernel on screen — and the function returns `[]`. It now prefers `inlined_at_*`, requiring **both** fields. Plain entries carry neither and are untouched. The Monaco rewrite moved this function from `CodeComparisonView.tsx` to `monaco/comparisonMapping.ts` and ported the defect with it, documenting the file match as "legacy direction preserved". It dropped the range check that used to reject the callee line a second time, so only one guard needs changing now. Pull Request resolved: #435 Test Plan: CPU only; the example trace inlines `file_content` for every stage. ``` python -m pytest tests/cpu -q TRITONPARSE_FB_MODE=0 python -m tritonparse.cli parse tests/example_output/logs --out /tmp/out --overwrite ``` > The checked-in `tests/example_output/parsed_output_*` files are from an older tritonparse and are **not** a valid baseline — re-parse the raw log on each branch. **Tests** — `main` 712 → +commit 1 `717` → +commit 2 `725`. Commit 3 is frontend-only: `npm run test:unit` 48/48, and reverting its fix fails exactly one case (47/48). `ruff check`/`format` clean per commit, no tabs in fixtures; `tsc`, `eslint`, `vite build` clean at the tip. Commit 1 adds a broken-chain case; commit 2 adds cases covering the two-level chain resolving to the outermost frame, a cyclic chain terminating, an unresolvable `.file` index yielding no call site rather than a half-attributed one, the tab-separated form ptxas actually emits, and an AMDGCN line-0 directive. **Commit 1**, `matmul_kernel`, python 288: | | before | after | |---|---|---| | entries | ttir 2, ttgir 2, ptx 9, sass 4 | ttir 4, ttgir 4, ptx 9, **sass 8** | | `is_callsite` | ttir 6, ttgir 6, sass 0 | ttir 6, ttgir 6, **sass 10** | No `llir` column: that needs #434 present. What stays at `standard.py:43` is metadata, correctly — `//## File` comments and `#loc` definitions *declare* that location, and a declaration has no call site. **Commit 2** — ptx entries with `line == 0` 18 → 0; ptx `is_callsite` 0 → 38; entries with a caller line but no caller file 0; python map containing key 0 in 2 compilations → 0. | cross-stage join | before | after | |---|---|---| | sass → ptx | 931/967 | 967/967 | | ttgir → ptx | 408/489 | 436/489 | | ttir → ptx | 398/479 | 426/479 | The ptx denominator drops by the 18 phantom line-0 entries. The two-level chain fix is verified separately, by compiling a `tl.sum` kernel ahead of time for sm_90: every inlined PTX entry resolves to the single `tl.sum` line in the kernel rather than to `standard.py:313`. LLIR's `inlinedAt` chain on the same kernel confirms the expected root, `273 → 313 → kernel`. **Corpus** — 17 kernels: the two from `tests/example_output/logs` plus 13 traced on an H100 covering elementwise, branchy, reduction, matmul and TMA shapes. Phantom line-0 entries: **143 on main, 0 here.** **Commit 3** — replayed both resolver versions over every IR entry: | | llir | ptx | sass | ttgir | ttir | |---|---|---|---|---|---| | entries | 400 | 588 | 967 | 489 | 479 | | resolve before | 388 | 550 | 931 | 461 | 451 | | resolve after | 400 | 588 | 951 | 473 | 463 | | newly fixed | 12 | 38 | 20 | 12 | 12 | | regressed / retargeted | 0 | 0 | 0 | 0 | 0 | 94 newly resolve, nothing regresses, and no entry that already resolved now points elsewhere — purely additive. (The `llir` column needs #434; the other four do not.) ## Screenshots Clicking python 288 `tl.cdiv(M, BLOCK_SIZE_M)` — python → IR: | | screenshot | | |---|---|---| | before | https://pxl.cl/dfqM4 | 5 highlights, every one a `#loc` declaration at the bottom of the IR; no instruction highlighted | | after | https://pxl.cl/dfqMj | 9 highlights, adding the `arith.addi` / `arith.divsi` that `tl.cdiv` actually produced | Clicking TTIR L21 `%num_pid_m = arith.addi %M, %c15_i32` — IR → python, the mirror case: | | screenshot | | |---|---|---| | before | https://pxl.cl/dgnCL | ttir 21, ttgir 25/27 highlighted; python pane nothing, and it does not scroll | | after | https://pxl.cl/dgnD8 | the same IR highlights, plus python 288 `num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)`, scrolled into view | The `Loaded from:` banner identifies which build is shown. The second pair is on the Monaco UI at the commits in this stack; the first predates the rewrite, which does not change python → IR behaviour. Reviewed By: FindHao Differential Revision: D120381731 Pulled By: wychi fbshipit-source-id: 43ceb380c4fbdae642f4906f0edc319bd219a246

meta-cla Bot added the CLA Signed This label is managed by the Meta Open Source bot. label Sep 15, 2026

wychi force-pushed the llir-source-mapping branch from b8098c6 to 90bc327 Compare September 15, 2026 17:27

wychi requested a review from FindHao September 15, 2026 17:29

wychi force-pushed the llir-source-mapping branch from 90bc327 to fec2f37 Compare September 16, 2026 08:09

wychi mentioned this pull request Sep 16, 2026

Fix source mapping for inlined Triton helpers #435

Closed

FindHao approved these changes Sep 16, 2026

View reviewed changes

meta-codesync Bot closed this in 512addb Sep 16, 2026

meta-codesync Bot added the Merged label Sep 16, 2026

wychi deleted the llir-source-mapping branch September 16, 2026 22:12

FindHao mentioned this pull request Sep 17, 2026

Support more intermediate IRs #220

Open

Navigation Menu

Repository navigation

Add LLIR source mapping from LLVM debug metadata - #434

Add LLIR source mapping from LLVM debug metadata#434
wychi wants to merge 1 commit into
mainfrom
llir-source-mapping

wychi commented Sep 15, 2026 •

edited

Loading

Uh oh!

meta-codesync Bot commented Sep 15, 2026

Uh oh!

meta-codesync Bot commented Sep 16, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

wychi commented Sep 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

The inlinedAt part is the bit worth reviewing

Scope

Test plan

1. Tests

2. A/B on the example trace

3. Blast radius

4. AMD

5. UI

6. Lint

Uh oh!

meta-codesync Bot commented Sep 15, 2026

Uh oh!

meta-codesync Bot commented Sep 16, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

wychi commented Sep 15, 2026 •

edited

Loading

The `inlinedAt` part is the bit worth reviewing