.. Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements. See the NOTICE file distributed with this work for additional information regarding copyright ownership. The ASF licenses this file to you under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at .. http://www.apache.org/licenses/LICENSE-2.0 .. Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. In-kernel profiling with CudaProfiler ===================================== Once a kernel is correct and you have seen how it compiles (see :doc:`compiling`), the next question is usually *where the cycles go*. Host-side timers and ``nsys`` tell you how long a launch took, but not how that time splits across the regions *inside* one kernel — the TMA loads, the mainloop MMAs, the softmax, the epilogue. ``tvm.tirx.bench.CudaProfiler`` is a lightweight, in-kernel event tracer for exactly this. You bracket regions of device code with ``start`` / ``end`` markers; at runtime one leader thread per block stamps the GPU global timer into a buffer you pass in as an ordinary kernel argument. After the launch you read the buffer back and decode it into per-region durations or a Perfetto timeline. It is *not* zero cost — every event is a ``%globaltimer`` read plus a global store, and every thread in the region pays a block fence — so it is a profiling/debugging tool, not something you leave on in production. The kernel ---------- The kernel below brackets a ``load`` / ``compute`` / ``store`` sequence. The ``compute`` region runs a 4000-iteration FMA loop so it clearly dominates. Events are a plain ``enum.Enum`` whose integer values start at 0 and index a names list. .. code-block:: python from enum import Enum import numpy as np import tvm from tvm.script import tirx as T from tvm.tirx.bench import CudaProfiler, export_to_perfetto_trace NUM_BLOCKS, BLOCK, NUM_GROUPS = 4, 128, 1 WRITE_STRIDE = NUM_BLOCKS * NUM_GROUPS # >= number of (block, group) lanes PROF_SIZE = 4096 # uint64 slots in the profiler buffer N = NUM_BLOCKS * BLOCK class Ev(Enum): Load = 0 Compute = 1 Store = 2 EV_NAMES = ["load", "compute", "store"] @T.prim_func def profiled_kernel(out_ptr: T.handle, inp_ptr: T.handle, prof_ptr: T.handle): out = T.match_buffer(out_ptr, (N,), "float32") inp = T.match_buffer(inp_ptr, (N,), "float32") prof = T.match_buffer(prof_ptr, (PROF_SIZE,), "uint64") T.device_entry() bid = T.cta_id([NUM_BLOCKS]) tid = T.thread_id([BLOCK]) idx = bid * BLOCK + tid # Construct the profiler inside the kernel; only the leader thread writes. p = CudaProfiler(prof, write_stride=WRITE_STRIDE, num_groups=NUM_GROUPS, default_leader=(tid == 0)) p.init(0) # group_id = 0; also stamps the buffer header at slot 0 p.start(Ev.Load) x: T.f32 = inp[idx] p.end(Ev.Load) p.start(Ev.Compute) acc: T.f32 = T.float32(0) for _ in range(4000): acc = acc * T.float32(1.0001) + x p.end(Ev.Compute) p.start(Ev.Store) out[idx] = acc p.end(Ev.Store) p.finalize() # mark this (block, group) lane done Run it and read the trace ------------------------- Allocate a zeroed ``uint64`` buffer, pass it as the last argument, then read it back. Each record is one ``uint64``: the high 32 bits are the timestamp, the low 32 bits a packed tag, so decoding is plain bit-twiddling on the host. .. code-block:: python dev = tvm.cuda(0) exe = tvm.compile(tvm.IRModule({"main": profiled_kernel}), target=tvm.target.Target("cuda"), tir_pipeline="tirx") 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) exe(out, inp, prof) dev.sync() prof_np = prof.numpy() opens, spans = {}, {} for i in range(1, len(prof_np)): word = int(prof_np[i]) if word == 0: continue ts, tag = word >> 32, word & 0xFFFFFFFF block = (tag >> 12) // NUM_GROUPS event_idx, event_type = (tag >> 2) & 0x3FF, tag & 0x3 # 0=start 1=end 2=instant 3=finalize if event_type == 0: opens[(block, event_idx)] = ts elif event_type == 1: spans.setdefault(block, []).append((EV_NAMES[event_idx], ts - opens[(block, event_idx)])) for block in sorted(spans): print(f"block {block}:", ", ".join(f"{n}={d}ns" for n, d in spans[block])) export_to_perfetto_trace(prof_np, "cudaprofiler.perfetto-trace", EV_NAMES) Durations are stable to within a few percent (they shift with GPU clocks):: block 0: load=32ns, compute=8704ns, store=64ns block 1: load=96ns, compute=8704ns, store=64ns block 2: load=96ns, compute=8704ns, store=64ns block 3: load=96ns, compute=8704ns, store=64ns ``export_to_perfetto_trace`` writes ``cudaprofiler.perfetto-trace`` from the same records; drop it onto https://ui.perfetto.dev for an interactive timeline. Because the timestamps come from the global ``%globaltimer`` (not a per-SM cycle counter), events from different blocks share one time axis and are directly comparable. On a real kernel ---------------- The same markers, sprinkled through a warp-specialized FlashAttention-4 kernel (one ``group`` per warp-group via ``num_groups``), produce a per-warp-group timeline of the whole pipeline: .. figure:: https://raw.githubusercontent.com/tlc-pack/web-data/main/images/tirx/tirx_cudaprofiler_fa4.png :align: center :alt: FlashAttention-4 in-kernel timeline in Perfetto One CTA of an FA4 forward kernel. ``group_0`` issues the TMA loads (``issue-tma-*``), ``group_3`` / ``group_4`` run the softmax pipeline (``softmax-max`` / ``-exp2`` / ``-sum``), and ``group_5`` runs the ``correction`` — the overlap between the producer and consumer warp-groups is exactly what intra-kernel profiling is for. The API ------- Construct the profiler **inside** the kernel body and call four methods: * ``init(group_id)`` — once per thread; ``group_id`` selects the sub-track and stamps the buffer header at slot 0. * ``start(event_type, leader=None)`` / ``end(event_type, leader=None)`` — open and close a region. Every thread executes them, but only the leader stores a record. * ``finalize(leader=None)`` — write a terminal record for this lane. Constructor arguments: * ``profiler_buffer`` — the ``uint64`` buffer you pass into the kernel. * ``write_stride`` — how far each leader advances between writes. Must be ``>=`` the number of ``(block, group)`` lanes so per-lane streams never collide; ``NUM_BLOCKS * NUM_GROUPS`` is the tight value, a persistent-grid kernel uses ``num_sms * num_groups``. * ``num_groups`` — independent sub-tracks per block. Use ``1`` for a plain kernel; in a warp-specialized kernel give each warp-group its own ``group_id`` and leader so their timelines don't mix. * ``default_leader`` — the predicate for the one writing thread (override per call with ``leader=``). * ``profiler_enabled`` — pass ``False`` (or a false-y ``PrimExpr``) to turn every method into a no-op, so you can leave the markers in and compile them out. ``CudaProfiler`` emits ``start`` / ``end`` / ``finalize``; ``instant`` (event type 2) is reserved in the wire format and understood by the decoder, but there is no method that produces one. Groups and granularity ---------------------- A block's threads are partitioned into ``num_groups`` logical *groups*, and the trace's unit is one ``(block, group)`` lane — each becomes its own track. The partition is yours: a group can be a warp-group, a single warp, or any set of threads, and it does **not** have to align to a warp (the recording path has no warp-collective op — just a predicated per-thread store and a block fence). Two rules: * a thread joins a group by calling ``init(group_id)``, which points *its* write cursor at that group's lane; * exactly one thread per group is the leader and actually writes — pick it with a predicate that is true for one thread in the group, and it must be a thread that called ``init`` for that group. Because each leader has its own cursor, one ``start`` / ``end`` statement records into *every* group at once: each leader stamps its own lane. **Groups as warp-groups.** A 256-thread block is two warp-groups; give each its own ``group_id`` and make its first thread the leader. Here the two warp-groups do different amounts of compute, so their tracks have different durations: .. code-block:: python NUM_GROUPS = 2 p = CudaProfiler(prof, write_stride=NUM_BLOCKS * NUM_GROUPS, num_groups=NUM_GROUPS, default_leader=(tid % 128 == 0)) # first thread of each warp-group if tid < 128: p.init(0) else: p.init(1) # ... load ... p.start(Ev.Compute) if tid < 128: for _ in range(1000): # warp-group 0: light acc = acc * T.float32(1.0001) + x else: for _ in range(5000): # warp-group 1: heavy acc = acc * T.float32(1.0001) + x p.end(Ev.Compute) :: block 0 group 0: load=96ns, compute=3040ns, store=64ns block 0 group 1: load=96ns, compute=10816ns, store=64ns block 1 group 0: load=96ns, compute=3072ns, store=64ns block 1 group 1: load=128ns, compute=10784ns, store=64ns **Groups that are not warp multiples.** A 128-thread block split 48 / 48 / 32 works the same way — the leaders are the base thread of each group, and the 48-thread groups (1.5 warps, crossing warp boundaries) each record a correct track: .. code-block:: python NUM_GROUPS = 3 # groups [0, 48) [48, 96) [96, 128) p = CudaProfiler(prof, write_stride=NUM_BLOCKS * NUM_GROUPS, num_groups=NUM_GROUPS, default_leader=((tid == 0) | (tid == 48) | (tid == 96))) if tid < 48: p.init(0) elif tid < 96: p.init(1) else: p.init(2) :: block 0 group 0: load=96ns, compute=4544ns, store=64ns # 48 threads (1.5 warps) block 0 group 1: load=64ns, compute=4512ns, store=96ns # 48 threads, crosses warp lines block 0 group 2: load=64ns, compute=4576ns, store=64ns # 32 threads What each call wraps -------------------- The methods are thin wrappers around the ``T.cuda.timer_*`` intrinsics, which lower to small ``__device__`` helpers emitted into the generated CUDA. The profiler keeps two per-thread ``"local"`` scratch slots — the running tag and write cursor — and every record is written by: .. code-block:: c++ // tvm_builtin_get_timestamp() == asm("mov.u32 %0, %globaltimer_lo;") profiler_buffer[profiler_write_offset[0]] = ((uint64_t)tvm_builtin_get_timestamp() << 32) | (profiler_tag[0] | event_bits); profiler_write_offset[0] += profiler_write_stride; // global store; only the leader runs this ``init`` computes ``BLOCK_GROUP_IDX = block_idx * num_groups + group_id``, writes the header ``profiler_buffer[0] = ((uint64_t)num_groups << 32) | num_blocks`` from block 0 / ``threadIdx.x == 0``, and seeds this lane's cursor to ``1 + BLOCK_GROUP_IDX`` and tag to ``BLOCK_GROUP_IDX << 12``. ``start`` writes the record (``event_bits = (event << 2) | 0``) then ``__threadfence_block()``; ``end`` fences then writes (``| 1``); ``finalize`` fences then writes ``0x3``. The fence runs on *every* thread in the region, only the store is leader-only — that fence is what brackets the region's memory traffic, and why the markers perturb the kernel. Usage notes and caveats ----------------------- * **Zero the buffer before the launch.** The decoder treats ``0`` as "empty" and reads the grid shape from slot 0, which only block 0 / thread 0 writes. * **Exactly one leader per (block, group).** Each thread keeps its own cursor, initialized to ``1 + block_group``; two leaders in the same lane write the same offsets and clobber each other. Use ``tid == 0`` or lane 0 of the group's leader warp. * **Call ``init`` once, before any ``start``.** It seeds each thread's tag and cursor; without it both are garbage. * **Size ``write_stride`` and the buffer together.** The largest slot a lane touches is ``1 + block_group + (records_per_lane - 1) * write_stride``; over-allocate, unused slots stay ``0`` and are skipped. * **``%globaltimer_lo`` is only the low 32 bits of the nanosecond timer.** It wraps about every 4.29 s (``2**32`` ns), so a region straddling a wrap decodes to a bogus duration. Resolution is coarse (tens of ns), so very short regions read 0 or a single tick. * **No payload.** ``start`` / ``end`` record only a timestamp and the event id; encode anything extra in the event id (a distinct ``Ev`` member) or in ``num_groups``. * **It is not free.** Two stores plus two block fences per region. Profile, read the numbers, then build with ``profiler_enabled=False``.