[DOCS][TIRX] Add in-kernel profiling (CudaProfiler) tutorial - #19895
Conversation
Document tvm.tirx.bench.CudaProfiler in the TIRx native-basics CUDA section: a minimal start/end/finalize example, record-buffer decoding and Perfetto export, the tag encoding and generated device code, and usage caveats.
There was a problem hiding this comment.
Code Review
This pull request introduces documentation for in-kernel profiling using CudaProfiler in TVM's TIRX. The review feedback highlights several critical issues in the provided Python code examples, including the use of standard Enum instead of IntEnum (which causes TVM FFI compatibility issues), the use of the non-existent tvm.runtime.tensor API instead of tvm.nd.array, and a potential KeyError during trace decoding when handling unmatched end events.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
|
|
||
| .. code-block:: python | ||
|
|
||
| from enum import Enum |
| class Ev(Enum): | ||
| Load = 0 | ||
| Compute = 1 | ||
| Store = 2 |
There was a problem hiding this comment.
Using enum.Enum for Ev might cause TVM compilation or FFI errors because standard Enum members are not instances of int and cannot be automatically converted by TVM's FFI. Changing Ev to inherit from enum.IntEnum ensures that the enum members behave as integers and are compatible with TVM Script and FFI.
| class Ev(Enum): | |
| Load = 0 | |
| Compute = 1 | |
| Store = 2 | |
| class Ev(IntEnum): | |
| Load = 0 | |
| Compute = 1 | |
| Store = 2 |
| inp = tvm.runtime.tensor(np.ones(N, "float32"), device=dev) | ||
| out = tvm.runtime.tensor(np.zeros(N, "float32"), device=dev) | ||
| prof = tvm.runtime.tensor(np.zeros(PROF_SIZE, "uint64"), device=dev) |
There was a problem hiding this comment.
In TVM, tvm.runtime.tensor is not a standard API and will raise an AttributeError. Use tvm.nd.array instead to allocate and initialize the NDArrays.
| inp = tvm.runtime.tensor(np.ones(N, "float32"), device=dev) | |
| out = tvm.runtime.tensor(np.zeros(N, "float32"), device=dev) | |
| prof = tvm.runtime.tensor(np.zeros(PROF_SIZE, "uint64"), device=dev) | |
| inp = tvm.nd.array(np.ones(N, "float32"), device=dev) | |
| out = tvm.nd.array(np.zeros(N, "float32"), device=dev) | |
| prof = tvm.nd.array(np.zeros(PROF_SIZE, "uint64"), device=dev) |
| elif event_type == 1: | ||
| spans.setdefault(block, []).append((EV_NAMES[event_idx], ts - opens[(block, event_idx)])) |
There was a problem hiding this comment.
If the profiler buffer wraps or contains incomplete trace data, an end event might be processed without a matching start event in opens. Using opens.pop((block, event_idx), None) avoids a potential KeyError and safely ignores unmatched end events.
| elif event_type == 1: | |
| spans.setdefault(block, []).append((EV_NAMES[event_idx], ts - opens[(block, event_idx)])) | |
| elif event_type == 1: | |
| start_ts = opens.pop((block, event_idx), None) | |
| if start_ts is not None: | |
| spans.setdefault(block, []).append((EV_NAMES[event_idx], ts - start_ts)) |
Explain the block -> groups -> one-leader-per-group model and the per-(block, group) track unit, with two runnable configs: groups as warp-groups, and non-warp-multiple (48/48/32) groups.
This adds an in-kernel profiling page to the TIRx native-basics CUDA section,
documenting the existing
tvm.tirx.bench.CudaProfiler.The page covers:
start/end/finalizemarkers and a user-supplied
uint64buffer;export_to_perfetto_trace;%globaltimerread, a leader-only global store, and a block fence);(block, group), buffer sizing, the 32-bit%globaltimerwrap, and the per-region cost.The example is tested end-to-end on a CUDA GPU (B200, sm_100). It is wired into
the
native_basics.rsttoctree after "Compiling and inspecting". TheFlashAttention-4 timeline screenshot is served from
tlc-pack/web-data(
images/tirx/tirx_cudaprofiler_fa4.png), matching the other tirx doc figures.