Profiling

Call-path profiling, derived throughput, and unbiased GPU benchmarking for Enigma kernels.

The Enigma profiler attaches every measurement to the chain of active user scopes (a call-path tree), so the same kernel launched from two places stays distinguishable. Scopes carry semantic work counts the hardware cannot know, and the profiler derives throughput (GFLOP/s, GB/s) from measured GPU time.

Profiling never changes generated MSL, and every API is a no-op when no profiler is active — the disabled path allocates nothing.

Quick start

import enigma
import numpy as np

@enigma.kernel
def vector_add(A: enigma.f32, B: enigma.f32, C: enigma.f32):
    tid = enigma.thread_position_in_grid
    C[tid] = A[tid] + B[tid]

compiled = enigma.compile(vector_add)
rt = enigma.MetalRuntime()
n = 1 << 20
a = np.random.randn(n).astype(np.float32)
b = np.random.randn(n).astype(np.float32)

with enigma.profile() as prof:
    with enigma.scope("step", metrics={"bytes": 12.0 * n}):
        rt.execute(compiled, [a, b], n * 4, grid=(n, 1, 1), threads=(256, 1, 1))

print(prof.key_averages().table(sort_by="gpu_total_us"))
print(prof.tree())
prof.export_hatchet("profile.json")     # call tree (Hatchet literal format)
prof.export_chrome_trace("trace.json")  # chrome://tracing / Perfetto

How it works

1. enigma.profile() installs a Profiler in a contextvar. The runtime checks it once per call; when empty, dispatch takes the untimed fast path.
2. enigma.scope(name, metrics=...) pushes name onto a contextvar scope stack and records one event on exit. Scopes nest; the stack is the call path.
3. When a profiler is active, dispatch goes through the Swift runtime’s timed path, which reads Metal command-buffer GPU timestamps — so gpu_time_us is GPU time, not CPU wall clock.
4. Profiler.add_event stamps the current scope stack onto each event (call_path) and runs any registered kernel hook to attach metrics.
5. Aggregation (key_averages), the call tree (tree), and the exporters all read the same event list.

Dispatch metadata

Profiled gpu_dispatch events carry extra metadata (visible in events() and Chrome traces):

Field	Meaning
scheduling_us	CPU-side command scheduling time
queue_wait_us	Commit-to-GPU-start latency
encoder_gpu_us	Encoder GPU time from stage-boundary counter sampling (only when runtime.supports_stage_counters() is true)
max_threads_per_threadgroup, thread_execution_width, static_threadgroup_memory_bytes	Pipeline-state limits
threadgroup_occupancy	Your threadgroup size ÷ the pipeline’s maxTotalThreadsPerThreadgroup; values well below 1.0 suggest the launch underfills the pipeline

One-shot vs prepared dispatch

execute() is one-shot: its table includes library load, pipeline creation, buffer creation, readback, and release. PreparedKernel.dispatch() reuses all of that.

Never compare an execute_total row against a prepared_dispatch row as if they measure the same thing.

Scopes, metrics, derived throughput

Attach semantic work counts to a scope and the profiler derives throughput from measured GPU time:

with enigma.profile() as prof:
    with enigma.scope("gemm", metrics={"flops": 2 * M * N * K,
                                       "bytes": 4 * (M * K + K * N + M * N)}):
        prepared.dispatch(grid, threads)

print(prof.key_averages().table())   # gains GFLOP/s and GB/s columns

The metric names flops and bytes drive the derivation; any other keys pass through to exports untouched.

Kernel hooks

Register the work formula once; every profiled dispatch of that kernel gets metrics automatically:

enigma.register_kernel_hook(
    "gemm_kernel",
    lambda *, kernel_name, grid, threads: {
        "flops": 2.0 * M * N * K,
        "bytes": 4.0 * (M * K + K * N + M * N),
    },
)
...
enigma.unregister_kernel_hook("gemm_kernel")

Call-path analysis

The same kernel called from different contexts stays distinguishable:

with enigma.profile() as prof:
    with enigma.scope("attention"):
        gemm_prepared.dispatch(grid, threads)
    with enigma.scope("mlp"):
        gemm_prepared.dispatch(grid, threads)

print(prof.key_averages(group_by="call_path").table())
# attention/gpu_dispatch:gemm_kernel ...
# mlp/gpu_dispatch:gemm_kernel ...
print(prof.tree())

export_hatchet("profile.json") writes the tree as Hatchet literal JSON — load it with hatchet.GraphFrame.from_literal(json.load(f)) to query hotspots or diff two profiles.

Unbiased benchmarking on unified memory

Apple Silicon shares one physical memory between CPU and GPU. That removes explicit transfers but creates two biases:

• Cold-start bias — the first dispatches pay pipeline creation, driver work, and page-residency costs. Timing them overstates kernel cost.
• Warm-cache bias — after a few iterations the working set is resident in the shared cache, so tight repeat loops on small buffers report bandwidth no cold-data workload will see. If you are measuring DRAM bandwidth, size the working set well beyond the chip’s last-level cache.

enigma.benchmark_kernel is built around these constraints:

prepared = rt.prepare(compiled, [a, b], n * 4)
bench = enigma.benchmark_kernel(
    prepared,
    grid=(n, 1, 1), threads=(256, 1, 1),
    repeat=100, warmup=10,
    flops=None, bytes_moved=12.0 * n,
)
print(bench.summary())
# vector_add: 100 calls  min 6.91us  p50 7.12us  mean 7.40us  p90 8.05us  max 12.3us  176.40 GB/s (p50)
prepared.release()

• Times come from Metal GPU timestamps only — no profiler events, no Python timing inside the measured region.
• Warmup runs on the untimed fast path.
• The full distribution is reported. Compare kernels by median (robust to thermal/scheduler outliers); min is best-case; p90 exposes tail variance.

Use benchmark_kernel to compare kernels; use profile_kernel / enigma.profile() to understand where time goes.

Xcode GPU capture (.gputrace)

In-shader clock intrinsics do not exist on Apple Silicon, so per-instruction timing is not available from Python. For intra-kernel analysis (instruction mix, memory traffic, occupancy timelines), hand the profiled region to Xcode’s GPU debugger instead:

rt = enigma.MetalRuntime()                      # registers the capture backend
with enigma.profile(capture="run.gputrace") as prof:
    prepared.dispatch(grid, threads)
# open run.gputrace in Xcode

Run the script with capture enabled:

MTL_CAPTURE_ENABLED=1 python my_script.py

• The path must end in .gputrace (ValueError otherwise).
• Capture starts when the profile() context enters and stops when it exits; everything dispatched inside is in the trace.
• Without MTL_CAPTURE_ENABLED=1 (or before any MetalRuntime exists) entering the context raises RuntimeError — capture never fails silently.

See the Profiler API reference for full signatures.