wehub-resource-sync
311 lines
13 KiB
ReStructuredText
311 lines
13 KiB
ReStructuredText
.. 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``.
|