Add GPU benchmarking and profiling appendix

README.md

@@ -25,7 +25,9 @@ vehicle is **TIRx** (Tensor IR next), a Python DSL for writing GPU kernels at th
   persistent scheduling, warp specialization, and 2-CTA clusters.
 - **Part IV — Flash Attention 4.** A complete attention kernel built from the Part III techniques:
   two MMAs with softmax between them, online-softmax rescaling, causal masking, and GQA.
-- **Reference.** TIRx language reference and compiler internals.
+- **Appendices.** TIRx language reference, [reproducible GPU benchmarking and
+  profiling](appendix/benchmarking_gpu_kernels.md), compiler internals, and asynchronous-kernel
+  debugging.
 
 ## Build the book locally
 

appendix/iket_example.py

@@ -0,0 +1,73 @@
+"""Minimal TIRx workload with IKET ranges for two warp roles."""
+
+from pathlib import Path
+
+import numpy as np
+
+import tvm
+from tvm.script import tirx as T
+from tvm.tirx.cuda import iket
+
+
+N = 256
+ELEMS_PER_LANE = 8
+
+
+@T.prim_func
+def warp_role_example(inp: T.Buffer((N,), "float32"), out: T.Buffer((N,), "float32")):
+    T.device_entry()
+    profiler = iket.IketProfiler()
+    warp_id = T.warp_id([2])
+    lane = T.lane_id([32])
+    shared = T.alloc_buffer((N,), "float32", scope="shared")
+
+    profiler.mark("kernel_start", warp_id)
+    if warp_id == 0:
+        profiler.range_push("producer_load")
+        for i in T.serial(ELEMS_PER_LANE, unroll=False):
+            index = i * 32 + lane
+            shared[index] = inp[index]
+        profiler.range_pop()
+
+    profiler.range_push("wait_for_data")
+    T.cuda.cta_sync()
+    profiler.range_pop()
+
+    if warp_id == 1:
+        profiler.range_push("consumer_compute")
+        for i in T.serial(ELEMS_PER_LANE, unroll=False):
+            index = i * 32 + lane
+            out[index] = shared[index] * T.float32(2) + T.float32(1)
+        profiler.range_pop()
+
+
+def profile_workload():
+    target = tvm.target.Target({"kind": "cuda", "arch": "sm_100a"})
+    executable = tvm.compile(warp_role_example, target=target, tir_pipeline="tirx")
+    module = executable.jit()
+
+    input_numpy = np.arange(N, dtype=np.float32)
+    inp = tvm.runtime.tensor(input_numpy, device=tvm.cuda())
+    out = tvm.runtime.empty((N,), "float32", device=tvm.cuda())
+    module.main(inp, out)
+    tvm.cuda().sync()
+
+    expected = input_numpy * 2 + 1
+    np.testing.assert_array_equal(out.numpy(), expected)
+
+
+def main():
+    result = iket.run(
+        profile_workload,
+        output_dir=Path("reports/iket-warp-roles"),
+        postprocess="all",
+        clobber=True,
+        timeout=600.0,
+    )
+    print(f"IKET output directory: {result.output_dir}")
+    for path in (*result.json_traces, *result.perfetto_traces, *result.html_reports):
+        print(f"IKET artifact: {path}")
+
+
+if __name__ == "__main__":
+    main()

appendix/index.md

@@ -1,11 +1,12 @@
 (chap_appendix)=
 # Overview
 
-The main text runs through Parts I–IV. The Reference section collects material you may want to consult while reading:
+The main text runs through Parts I–IV. The appendices collect material you may want to consult while reading:
 
 | Need | Where |
 |------|-----|
 | Look up a TIRx language feature | **{ref}`chap_language_reference`** |
+| Measure, compare, and profile GPU kernels reproducibly | **{ref}`chap_benchmarking`** |
 | Compiler internals (the lowering pipeline) | **{ref}`chap_arch`** |
 | Debug asynchronous GEMM/FA hangs, crashes, wrong results, or slowdowns | **{ref}`chap_warp_spec_debug`** |
 

appendix/nsys_example.py

@@ -0,0 +1,119 @@
+"""Reusable CUDA workload for the benchmarking and profiling appendix."""
+
+import argparse
+from statistics import median
+
+import torch
+
+
+def make_workload(size: int):
+    a = torch.randn((size, size), dtype=torch.bfloat16, device="cuda")
+    b = torch.randn((size, size), dtype=torch.bfloat16, device="cuda")
+    c = torch.empty((size, size), dtype=torch.bfloat16, device="cuda")
+    output = torch.empty_like(c)
+
+    def run():
+        with torch.cuda.nvtx.range("BF16 GEMM"):
+            torch.mm(a, b, out=c)
+        with torch.cuda.nvtx.range("ReLU"):
+            torch.clamp_min(c, 0, out=output)
+
+    def validate():
+        torch.set_float32_matmul_precision("highest")
+        expected = torch.clamp_min(torch.mm(a.float(), b.float()), 0).to(output.dtype)
+        torch.testing.assert_close(output, expected, rtol=2e-2, atol=1e-2)
+
+    return run, validate
+
+
+def run_once_for_profiler(run, *, warmup_calls: int):
+    for _ in range(warmup_calls):
+        run()
+    torch.cuda.synchronize()
+
+    cudart = torch.cuda.cudart()
+    cudart.cudaProfilerStart()
+    with torch.cuda.nvtx.range("target operation"):
+        run()
+        torch.cuda.synchronize()
+    cudart.cudaProfilerStop()
+
+
+def measure_event_us(run, *, warmup_calls: int, samples: int):
+    for _ in range(warmup_calls):
+        run()
+    torch.cuda.synchronize()
+
+    values = []
+    for _ in range(samples):
+        start = torch.cuda.Event(enable_timing=True)
+        end = torch.cuda.Event(enable_timing=True)
+        start.record()
+        run()
+        end.record()
+        end.synchronize()
+        values.append(start.elapsed_time(end) * 1e3)
+    return values
+
+
+def collect_proton(run, *, warmup_calls: int, profile_calls: int, output: str):
+    import triton.profiler as proton
+
+    for _ in range(warmup_calls):
+        run()
+    torch.cuda.synchronize()
+
+    session = proton.start(output, context="shadow", data="tree")
+    if session is None:
+        raise RuntimeError("Proton session could not be created")
+    try:
+        with proton.scope("target_operation"):
+            for _ in range(profile_calls):
+                run()
+        torch.cuda.synchronize()
+    finally:
+        proton.finalize(session)
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    mode = parser.add_mutually_exclusive_group()
+    mode.add_argument("--profile-once", action="store_true")
+    mode.add_argument("--event-samples", type=int, default=0)
+    mode.add_argument("--proton-calls", type=int, default=0)
+    parser.add_argument("--size", type=int, default=4096)
+    parser.add_argument("--warmup-calls", type=int, default=500)
+    parser.add_argument("--proton-output", default="operator")
+    args = parser.parse_args()
+
+    run, validate = make_workload(args.size)
+    run()
+    torch.cuda.synchronize()
+    validate()
+
+    if args.profile_once:
+        run_once_for_profiler(run, warmup_calls=args.warmup_calls)
+    elif args.event_samples:
+        values = measure_event_us(
+            run,
+            warmup_calls=args.warmup_calls,
+            samples=args.event_samples,
+        )
+        print(
+            f"median={median(values):.3f} us, "
+            f"min={min(values):.3f} us, max={max(values):.3f} us"
+        )
+    elif args.proton_calls:
+        collect_proton(
+            run,
+            warmup_calls=args.warmup_calls,
+            profile_calls=args.proton_calls,
+            output=args.proton_output,
+        )
+    else:
+        run()
+        torch.cuda.synchronize()
+
+
+if __name__ == "__main__":
+    main()

chapter_gemm_advanced/index.md

@@ -861,7 +861,7 @@ def hgemm_v9(M, N, K):
 
 ## End-to-End Results
 
-The table below follows the progression from the naive baseline to the warp-specialized cluster kernel, with cuBLAS included as a reference. The measurements use an NVIDIA B200 with `M=N=K=4096`, fp16 inputs, locked clocks, and 1,000 timed iterations per measured version:
+The table below follows the progression from the naive baseline to the warp-specialized cluster kernel, with cuBLAS included as a reference. The measurements use an NVIDIA B200 with `M=N=K=4096`, fp16 inputs, locked clocks, and 1,000 timed iterations per measured version. New measurements and reproduction attempts should follow the full protocol in {ref}`chap_benchmarking`.
 
 | Step | Technique | Time | Speedup |
 |------|-----------|------|---------|

chapter_gemm_basics/index.md

@@ -320,6 +320,11 @@ tflops = 2 * M * N * K / ms / 1e9
 print(f"Performance: {ms:.3f} ms, {tflops:.1f} TFLOPS")
 ```
 
+This short timing loop is sufficient for a smoke measurement, but it is not a complete experimental
+protocol. For reported results, follow {ref}`chap_benchmarking`: define the timing boundary, collect
+multiple samples, state the cache and clock policy, and separate unprofiled latency measurements from
+Proton or Nsight Compute analysis.
+
 ### Limits of the Single-Tile Kernel
 
 The kernel is correct, but it still has a narrow operating range:

chapter_performance/index.md

@@ -338,3 +338,7 @@ For a memory-bound kernel, focus on reducing data movement and making transfers
 bandwidth ceiling. For a compute-bound kernel, focus on reducing idle time in the compute units. The roofline model does not
 produce the final implementation, but it prevents effort from being spent on resources that are not
 the bottleneck.
+
+Roofline interpretation starts from a trustworthy measurement. The practical workflow for timing a
+kernel, locating its expensive launches with Proton, and testing a hardware hypothesis with Nsight
+Compute is collected in {ref}`chap_benchmarking`.