[DSL][Runtime][Kernel] Add in-kernel wave tracing (fx.ktrace) for gfx942/gfx950 by Phil-amd · Pull Request #1162 · ROCm/FlyDSL

Phil-amd · 2026-09-18T06:35:40Z

Adds fx.ktrace: in-kernel wave tracing for gfx942/gfx950. Annotate phases in a
kernel, run it, and get a Perfetto timeline of the phases you named with the
nesting you wrote -- which is what ATT, working from instructions after the fact,
cannot give.

fx.ktrace.range_push("mainloop")
for k_tile in range(0, k_tiles, 1):
    fx.ktrace.range_push("k_tile", k_tile)
    ...
    fx.ktrace.range_pop()
fx.ktrace.range_pop()

FLYDSL_KTRACE_ENABLE=1 python3 your_script.py

Each annotation writes a timestamped 32-byte record into a device buffer, which
the host reads back and renders. Measured at ~40 ns per event on MI355X.

Why a dialect rather than emitting the records in the frontend

The frontend could build the record directly -- leader guard, slot claim, bounds
checks, stores -- and that is the smaller change. The problem is what any later
pass then sees: an atomic and six stores, with no op meaning "this is one trace
event". Keeping that meaning in the IR is what makes pass-level work possible at
all -- folding adjacent events, dropping a range that collapses, deciding in the
compiler rather than only at the Python _enabled() check.

The cost is one dialect (6 ops and a token type) plus one pass.

Key points

convert-fly-ktrace-to-rocdl runs just before convert-fly-to-rocdl,
forced from both sides: it emits fly.to_llvm_ptr, which only that pass
lowers, and scf.if, which must survive to convert-scf-to-cf.
One contended atomic per wave, confined to a single lane and hoisted to the
entry block; per event only an uncontended bump of that wave's own counter. An
unguarded claim runs on all 64 lanes and exhausts the buffer 64x early.
A token crossing a region boundary needs a dialect conversion, not an
in-place replacement: retyping a use does not retype an scf.for's result or
body argument, and the loop-carried range is the case the token form exists
for -- the common path, not a corner.
Configuration reaches the pass as module attributes, not re-read from the
environment, so one compilation cannot disagree with itself about the buffer it
was built against. Event ids are assigned by the pass and published as
fly_ktrace.event_names.
Disabled builds are byte-identical to unannotated ones -- verified, not
assumed: no fly_ktrace op, no atomic, no trace-buffer parameter.
The failure modes that matter here are silent ones. A dropped range_end,
a buffer smaller than the bound compiled against it, a wave past its 255-event
budget: each yields a plausible timeline rather than an error, so each is
rejected loudly instead.

Testing


Unit	59 tests, gfx950, no skips
FileCheck	7 tests; each confirmed to fail when its subject is perturbed
Annotated GEMM	`test_gemm_a16w16_gfx950.py -k acc_small_m` with tracing on: 20 passed, same as with it off
Device	a traced launch decodes with phase names intact and a loop-carried range pairing across every iteration

Untraced compiles were checked to produce a byte-identical pass pipeline.

Notes for review

gfx942/gfx950 only; a non-CDNA target is refused with a diagnostic naming ktrace.
Known limits, documented rather than papered over: one device per process, graph
replay is not traced correctly, and both ends of a range must sit under the same
guard. docs/ktrace_guide.md covers the API, the settings, the cost model
(including where the ~40 ns does not hold), and a "How it is implemented"
section pointing at the dialect and the pass.

coderfeli · 2026-09-20T09:37:57Z

 import flydsl.compiler as flyc
 import flydsl.expr as fx
-from flydsl.expr import const_expr, gpu, range_constexpr, rocdl
+from flydsl.expr import const_expr, gpu, ktrace, range_constexpr, rocdl


add an examples instead of chaning production kernels.

coderfeli · 2026-09-20T09:43:59Z

+    return max(int(env.ktrace.buffer_bytes) // _ktrace.RECORD_BYTES - 1, 1)
+
+
+_BUFFER = None


global buffer for all trace data?

coderfeli · 2026-09-20T09:53:32Z

+**Status: unstable.** The API, record format and output schema may change in any minor
+release. See [`api_stability.md`](api_stability.md).
+
+**Hardware: gfx942 (CDNA3) and gfx950 (CDNA4) only.** Other targets raise at trace time.


not only for cnda3, 4. There are 5, rdna...

This commit focuses on API and architecture design. Support for other architectures will be developed and tested separately, and submitted as standalone commits once validated.

coderfeli · 2026-09-20T09:55:26Z

+
+## Limitations
+
+- gfx942/gfx950 only.


This commit focuses on API and architecture design. Support for other architectures will be developed and tested separately, and submitted as standalone commits once validated.

coderfeli · 2026-09-20T09:58:54Z

    return False


+def _invalidate_ktrace_records() -> None:


move to ktrace.py?

Done — the three helpers that only touch ktrace's own state (invalidate_records, restore_names, require_matching_capacity) moved to ktrace.py; the rest stay in jit_function.py since they build MLIR attributes or pack the launch ABI. Also reverted the production GEMM in favour of examples/07-ktrace_wave_timeline.py, and moved the API to flydsl.expr.experimental.ktrace.

coderfeli · 2026-09-20T10:00:42Z

+
+
+@dsl_loc_tracing
+def s_memrealtime():


no need to expose so many apis in this file. and also set '_' for internal only

Done
Removed them — s_memrealtime, s_getreg_hw_id and s_getreg_xcc_id had no callers anywhere in the tree, since the conversion pass emits llvm.amdgcn.s.memrealtime / s.getreg itself after the dialect migration. Deleted along with their two private helpers and the test file that only covered them, which also drops three names from fx.rocdl.*.

coderfeli · 2026-09-20T10:01:35Z

+_BUFFER = None
+
+
+def ensure_buffer():


not safe to in this way.

Done
The consequence reproduces: four threads reaching a first traced launch together each call hipMalloc, so three buffers leak for the life of the process and every record written into them is never read — trace data lost silently. Fixed with a double-checked lock following the existing pattern in compiler/jit_executor.py (the common already-allocated path still takes no lock), plus a regression test confirmed to fail with assert 4 == 1 when the lock is removed rather than passing vacuously.

coderfeli · 2026-09-20T10:02:58Z

@@ -0,0 +1,443 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2025 FlyDSL Project Contributors


move to exp ?

Done — it's now flydsl.expr.experimental.ktrace, so the import path itself marks the feature as unstable rather than relying on the absence of all. Graduating it later becomes a visible move out of experimental/ instead of a one-line change.

coderfeli · 2026-09-20T10:04:42Z

+_hip = None
+
+
+def _hip_lib():


why need this?

Done — Delete

Peter9606 · 2026-09-23T00:29:44Z

+kernel entry, where each wave claims its whole slot range at once; per event the
+cost is a timestamp, an uncontended bump of that wave's own counter, and the
+record store.
+"""


I like keeping a high-level name here. range_push / mark / collect
is the right author-facing surface — the same annotations should work on
another GPU later without rewriting kernels. Iluvatar can do an in-kernel
phase timeline too, and so can NV; the cheap way to get there is a
portable API plus a per-backend lowering, which is already the reason this
is a dialect instead of inlined stores.

The part that would make that follow-up hard is that this module currently
ships the CDNA record as if it were the API: xcc_id, wave slots,
s_memrealtime, HIP ktrace_buffer, and convert-fly-ktrace-to-rocdl
are all visible from the public entry points. A second backend then has
to either fake XCD topology or break the 32-byte layout.

Would you consider splitting it, even if gfx942/gfx950 stays the only
implemented target in this PR?

Keep fx.experimental.ktrace (and fly_ktrace.*) as the annotation
contract: names, push/pop vs start/end, payload, host collect/summary.

Treat clock source, workgroup identity, buffer allocation, and the 32-byte
encoding as backend details behind FlyKtraceToROCDL / the HIP runtime.

Then a later FlyKtraceToIXDL (or NV) only has to implement those details.
Kernel authors keep the same high-level names, which is the point of not
calling this fx.rocdl.ktrace.

This is still an experimental API, so the focus for now is getting the frontend surface and its stability right. Worth noting that most of it is already backend-agnostic — 6 of the 8 constants (the five KIND_* and UNPAIRED_SLOT) are consumed by the portable decoder in utils/ktrace_trace.py; only RECORD_BYTES and EVENTS_PER_WAVE are ROCm-shaped.

Peter9606 · 2026-09-23T00:30:24Z

+imported by the package ``__init__`` -- doing so would pull ``flydsl.expr`` into
+every ``import flydsl.runtime``, which today needs none of it.
+"""
+


Related to the API split: could the extra kernel argument and the
ktrace compile hint be a backend hook rather than always-on in the
shared JIT path?

ROCm can keep today's HIP buffer + convert-fly-ktrace-to-rocdl. Other
targets would register their own buffer + conversion, or fail closed with
"ktrace is not implemented here" instead of compiling the annotations
away. That way FLYDSL_KTRACE_ENABLE=1 stays a high-level switch, and
enabling it on a backend that has no lowering is loud rather than an
empty Perfetto file.

This is already fail-closed, and loudly — the check is in the conversion pass, so an unsupported arch fails at compile time rather than producing an empty trace. On a gfx1100 module:

$ fly-opt --convert-fly-ktrace-to-rocdl unsupported_arch.mlir
error: ktrace: in-kernel timestamps are not supported on target arch 'gfx1100';
supported: gfx942 (CDNA3) and gfx950 (CDNA4)
There's no path where tracing silently no-ops.

Kernel authors annotate phases with fx.experimental.ktrace -- mark, range_push/range_pop and range_start/range_end; each annotation records a timestamped 32-byte event that a post-processing step renders as a Perfetto timeline. The annotations are MLIR ops, not inlined code. The frontend emits one fly_ktrace op per call and convert-fly-ktrace-to-rocdl expands it into the record-writing form. That is the point of the change: keeping the annotation in the IR lets a pass reason about a trace event instead of an atomic and six stores. The pass runs just before convert-fly-to-rocdl, forced from both sides -- it emits fly.to_llvm_ptr, which only that pass lowers, and scf.if, which must survive until convert-scf-to-cf. Disabled by default: without FLYDSL_KTRACE_ENABLE=1, or the ktrace compile hint, an annotated kernel compiles byte-identically to an unannotated one. The expansion's invariants -- one contended atomic per wave, the record written by the first active lane, unsigned bounds compares -- are documented where they are emitted, in lib/Conversion/FlyKtraceToROCDL.

Phil-amd force-pushed the feature/ktrace-dialect-squashed branch 2 times, most recently from c140368 to 447758e Compare September 18, 2026 06:48

Phil-amd self-assigned this Sep 18, 2026

Phil-amd requested review from coderfeli and sjfeng1999 September 18, 2026 06:52

Phil-amd closed this Sep 18, 2026

Phil-amd reopened this Sep 18, 2026

Phil-amd force-pushed the feature/ktrace-dialect-squashed branch 6 times, most recently from 56c4143 to 96f01ec Compare September 20, 2026 05:59

coderfeli reviewed Sep 20, 2026

View reviewed changes

Comment thread docs/api_stability.md Outdated

coderfeli reviewed Sep 20, 2026

View reviewed changes

Phil-amd force-pushed the feature/ktrace-dialect-squashed branch 6 times, most recently from ae834af to 4ecc5b1 Compare September 21, 2026 11:19

Peter9606 reviewed Sep 23, 2026

View reviewed changes

Phil-amd force-pushed the feature/ktrace-dialect-squashed branch from 4ecc5b1 to 9e64d4f Compare September 23, 2026 23:28

		return max(int(env.ktrace.buffer_bytes) // _ktrace.RECORD_BYTES - 1, 1)


		_BUFFER = None

		@@ -0,0 +1,443 @@
		# SPDX-License-Identifier: Apache-2.0
		# Copyright (c) 2025 FlyDSL Project Contributors



		@dsl_loc_tracing
		def s_memrealtime():

		_hip = None


		def _hip_lib():

Conversation

Phil-amd commented Sep 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Why a dialect rather than emitting the records in the frontend

Key points

Testing

Notes for review

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Phil-amd Sep 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Phil-amd commented Sep 18, 2026 •

edited

Loading

Phil-amd Sep 21, 2026 •

edited

Loading