Skip to content

[DOCS][TIRX] Add in-kernel profiling (CudaProfiler) tutorial - #19895

Merged
spectrometerHBH merged 2 commits into
apache:mainfrom
spectrometerHBH:docs-tirx-cudaprofiler
Jun 27, 2026
Merged

[DOCS][TIRX] Add in-kernel profiling (CudaProfiler) tutorial#19895
spectrometerHBH merged 2 commits into
apache:mainfrom
spectrometerHBH:docs-tirx-cudaprofiler

Conversation

@spectrometerHBH

Copy link
Copy Markdown
Contributor

This adds an in-kernel profiling page to the TIRx native-basics CUDA section,
documenting the existing tvm.tirx.bench.CudaProfiler.

The page covers:

  • a minimal load / compute / store example using start / end / finalize
    markers and a user-supplied uint64 buffer;
  • decoding the record buffer on the host and exporting a Perfetto trace via
    export_to_perfetto_trace;
  • the record/tag encoding and the device code each call lowers to (a
    %globaltimer read, a leader-only global store, and a block fence);
  • usage notes: one leader per (block, group), buffer sizing, the 32-bit
    %globaltimer wrap, 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.rst toctree after "Compiling and inspecting". The
FlashAttention-4 timeline screenshot is served from tlc-pack/web-data
(images/tirx/tirx_cudaprofiler_fa4.png), matching the other tirx doc figures.

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.

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Import IntEnum instead of Enum to support the integer-based enum definition for Ev.

Suggested change
from enum import Enum
from enum import IntEnum

Comment on lines +57 to +60
class Ev(Enum):
Load = 0
Compute = 1
Store = 2

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

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.

Suggested change
class Ev(Enum):
Load = 0
Compute = 1
Store = 2
class Ev(IntEnum):
Load = 0
Compute = 1
Store = 2

Comment on lines +108 to +110
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)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

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.

Suggested change
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)

Comment on lines +126 to +127
elif event_type == 1:
spans.setdefault(block, []).append((EV_NAMES[event_idx], ts - opens[(block, event_idx)]))

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

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.

Suggested change
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.
@spectrometerHBH
spectrometerHBH merged commit 46a9a00 into apache:main Jun 27, 2026
11 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants