calibrax.profiling.roofline

Roofline model analysis for JAX operations. Compares achieved arithmetic intensity against hardware peak throughput to identify whether a workload is compute-bound or memory-bound, and generates optimization recommendations.

Roofline analysis for JAX operations.

Identifies whether operations are compute-bound or memory-bound by comparing arithmetic intensity against hardware roofline limits, and generates optimization recommendations.

RooflineResult(*, arithmetic_intensity, critical_intensity, memory_bandwidth_utilization, flops_utilization, bottleneck, efficiency, execution_time_ms, recommendations=()) dataclass

Result of a roofline analysis on a JAX operation.

Attributes:

Name	Type	Description
arithmetic_intensity	float	Achieved FLOPs per byte of memory traffic.
critical_intensity	float	Hardware's ridge point (FLOPs/byte).
memory_bandwidth_utilization	float	Fraction of peak memory bandwidth used.
flops_utilization	float	Fraction of peak FLOPs achieved.
bottleneck	str	Either "memory_bandwidth" or "compute".
efficiency	float	Utilization of the binding resource.
execution_time_ms	float	Measured execution time in milliseconds.
recommendations	tuple[str, ...]	Optimization suggestions.

to_dict()

Serialize to a JSON-compatible dictionary.

from_dict(data) classmethod

Deserialize from a dictionary.

Parameters:

Name	Type	Description	Default
data	dict[str, Any]	Dictionary with roofline result fields.	required

Returns:

Type	Description
RooflineResult	Reconstructed RooflineResult instance.

RooflineAnalyzer(*, hardware_specs=detect_hardware_specs()) dataclass

Analyzes operation performance against hardware roofline limits.

Uses measured execution time and estimated FLOPs to determine whether an operation is compute-bound or memory-bound, and how efficiently it uses the available hardware resources.

Attributes:

Name	Type	Description
hardware_specs	dict[str, Any]	Hardware specification dictionary (auto-detected if not provided).

analyze_operation(func, inputs, *, flops_override=None)

Perform roofline analysis on a JAX operation.

Parameters:

Name	Type	Description	Default
func	Callable[..., Any]	JAX function to analyze.	required
inputs	list[Array]	Input arrays for the function.	required
flops_override	int | None	If provided, use this FLOP count instead of estimating. For accurate results, pass the output of FlopsCounter.count().	None

Returns:

Type	Description
RooflineResult	RooflineResult with bottleneck classification and recommendations.