Profiler

Call-path profiling, kernel hooks, derived throughput, and unbiased GPU benchmarking.

Every profiling API is a no-op when no profiler is active, and profiling never changes generated MSL. See the Profiling guide for concepts and worked examples.

---

enigma.profile()

Create a Profiler context manager. Installs the profiler in a contextvar so runtime dispatch routes through the timed GPU path; outside the context, dispatch takes the untimed fast path.

prof: Profiler = enigma.profile(
    record_shapes: bool = False,
    profile_memory: bool = False,
    capture: str | Path | None = None,
)

Parameters

Parameter	Type	Default	Description
record_shapes	bool	False	Record tensor shapes on events
profile_memory	bool	False	Record input/output byte counts
capture	str | Path | None	None	Write an Xcode .gputrace for the region; path must end in .gputrace

Returns

A Profiler context manager. With capture set, entering the context starts GPU capture and requires MTL_CAPTURE_ENABLED=1 plus an existing MetalRuntime — it raises RuntimeError otherwise (never fails silently), and ValueError if the path does not end in .gputrace.

---

enigma.scope()

Context manager that profiles a named region and stamps the call path onto every event recorded inside it. Scopes nest; the stack of active scopes is the call path. No-op when no profiler is active.

with enigma.scope(
    name: str,
    *,
    metrics: dict[str, float] | None = None,
    category: str = "scope",
    **metadata,
):
    ...

Parameters

Parameter	Type	Default	Description
name	str	required	Scope label, becomes a call-path node
metrics	dict | None	None	Semantic work counts, e.g. {"flops": ..., "bytes": ...}
category	str	"scope"	Event category
**metadata	any	—	Extra fields passed through to exports

The flops and bytes metric keys drive derived GFLOP/s and GB/s; any other keys pass through to exports untouched. The scope stack is restored even if the body raises.

---

enigma.record_function()

Alias of scope with category="python", for marking Python-side regions.

with enigma.record_function(name: str, *, category: str = "python", **metadata):
    ...

---

enigma.get_active_profiler()

Return the Profiler active in the current context, or None.

prof: Profiler | None = enigma.get_active_profiler()

---

enigma.register_kernel_hook()

Attach a metrics hook to every profiled dispatch of kernel_name. The hook fires only when the event has no metrics already, and is called as hook(kernel_name=..., grid=..., threads=...); it should return a metrics dict such as {"flops": ..., "bytes": ...}.

enigma.register_kernel_hook(
    kernel_name: str,
    hook: Callable[..., dict[str, float] | None],
)

Parameters

Parameter	Type	Description
kernel_name	str	Kernel to attach metrics to
hook	callable	Returns a metrics dict (or None)

---

enigma.unregister_kernel_hook()

Remove a previously registered hook. Safe to call when no hook is registered.

enigma.unregister_kernel_hook(kernel_name: str)

---

enigma.benchmark_kernel()

Benchmark a PreparedKernel using Metal GPU timestamps only — no profiler events and no Python timing inside the measured region. Warmup dispatches run on the untimed fast path to absorb pipeline creation, driver work, and unified-memory page residency.

bench: KernelBenchmark = enigma.benchmark_kernel(
    prepared: PreparedKernel,
    *,
    grid: tuple[int, int, int],
    threads: tuple[int, int, int],
    repeat: int = 100,
    warmup: int = 10,
    flops: float | None = None,
    bytes_moved: float | None = None,
)

Parameters

Parameter	Type	Default	Description
prepared	PreparedKernel	required	Pre-allocated kernel
grid	tuple	required	(gx, gy, gz) grid dimensions
threads	tuple	required	(tx, ty, tz) threads per threadgroup
repeat	int	100	Timed runs
warmup	int	10	Discarded warmup runs
flops	float | None	None	FLOPs per dispatch, for derived GFLOP/s
bytes_moved	float | None	None	Bytes per dispatch, for derived GB/s

Returns

A KernelBenchmark. Raises ValueError if repeat < 1 or warmup < 0.

---

enigma.profile_kernel()

Profile repeated dispatches of a prepared kernel and return aggregated rows. Use this to understand where time goes; use benchmark_kernel to compare kernels.

result: ProfilerResult = enigma.profile_kernel(
    prepared: PreparedKernel,
    *,
    grid: tuple[int, int, int],
    threads: tuple[int, int, int],
    repeat: int = 50,
    warmup: int = 5,
    profile_memory: bool = True,
)

Parameters

Parameter	Type	Default	Description
prepared	PreparedKernel	required	Pre-allocated kernel
grid	tuple	required	(gx, gy, gz) grid dimensions
threads	tuple	required	(tx, ty, tz) threads per threadgroup
repeat	int	50	Profiled runs
warmup	int	5	Warmup runs (untimed)
profile_memory	bool	True	Record byte counts

Returns

A ProfilerResult. Raises ValueError if repeat < 1 or warmup < 0.

---

Profiler

The object returned by enigma.profile().

Method	Description
events()	List of recorded ProfilerEvents
key_averages(*, group_by="name")	Aggregate into a ProfilerResult; group_by is "name" or "call_path"
tree()	Render the call-path tree with inclusive CPU/GPU times
export_hatchet(path)	Write the call tree as Hatchet literal JSON
export_chrome_trace(path)	Write a chrome://tracing / Perfetto trace

Load a Hatchet export with hatchet.GraphFrame.from_literal(json.load(f)) to query hotspots or diff two profiles.

---

ProfilerResult

Aggregated view over events, returned by key_averages().

Method	Description
rows(*, sort_by="cpu_total_us")	Sorted list of ProfilerRows
by_name(name)	The row (or merged rows) for name; raises KeyError if absent
table(*, sort_by="cpu_total_us")	Formatted text table; gains GFLOP/s and GB/s columns when metrics exist

---

ProfilerRow

Property	Type	Description
name	str	Event name or call path
kernel_name	str | None	Kernel, when applicable
calls	int	Number of aggregated events
cpu_total_us	float	Total CPU (wall) time
gpu_total_us	float	Total GPU time
cpu_avg_us	float	Mean CPU time per call
gpu_avg_us	float	Mean GPU time per call
input_bytes	int	Summed input bytes
output_bytes	int	Summed output bytes
metrics	dict	Summed semantic metrics
gflops_per_s	float	From flops over GPU time (CPU fallback)
gbytes_per_s	float	From bytes over GPU time (CPU fallback)

---

ProfilerEvent

A single timed event. Frozen dataclass.

Field / Property	Type	Description
name	str	Event name
category	str	"scope", "python", "metal", …
start_ns, end_ns	int	CPU timestamps
kernel_name	str | None	Dispatched kernel
grid, threads	tuple | None	Launch dimensions
input_bytes, output_bytes	int	Byte counts
buffer_count	int	Number of GPU buffers involved
gpu_time_us	float	GPU time from command-buffer timestamps
call_path	tuple[str, ...]	Active scope stack at record time
metrics	dict	Semantic work counts
metadata	dict	Dispatch metadata (scheduling, occupancy, …)
duration_us	float	CPU duration
label	str	name or name:kernel_name
node_path	tuple[str, ...]	Full call-tree path including this node

---

KernelBenchmark

Returned by benchmark_kernel. Frozen dataclass; times come from GPU timestamps only.

Property / Method	Type	Description
kernel_name	str	Benchmarked kernel
calls	int	Number of timed runs
gpu_times_us	tuple[float, ...]	Per-call GPU times
min_us	float	Best-case time
median_us	float	Median (p50) — use for comparisons
mean_us	float	Mean time
p90_us	float	90th-percentile (tail variance)
max_us	float	Worst-case time
gflops_per_s	float	From flops over median time
gbytes_per_s	float	From bytes_moved over median time
summary()	str	One-line distribution + throughput