polyglot_optima / server /tools /hardware_profiler.py
Swastikr's picture
Upload folder using huggingface_hub
2bf863e verified
"""Tool 1/9: get_hardware_profile.
Returns the hardware profile for the current episode along with the precomputed
Roofline bound. The profile is sampled at reset() time and frozen for the episode;
this tool just exposes it to the agent.
Roofline math (per plan §10):
simd_w = {"SSE4.2": 4, "AVX2": 8, "AVX-512": 16, "NEON": 4, "none": 1}
peak_flops = cores × freq_ghz × simd_w × 2 (FMA = 2 ops/cycle)
peak_bandwidth_flops = bandwidth_gbs × 0.5 (rough flop-per-byte ceiling)
roofline_bound = min(peak_flops, peak_bandwidth_flops)
"""
from __future__ import annotations
from typing import Any
SIMD_WIDTH = {
"SSE4.2": 4,
"AVX2": 8,
"AVX-512": 16,
"NEON": 4,
"none": 1,
}
def roofline_bound(hw: dict[str, Any]) -> float:
"""Compute the Roofline-model peak GFLOPS for a hardware profile."""
simd_w = SIMD_WIDTH.get(hw["simd"], 1)
peak_flops = hw["cores"] * hw["freq_ghz"] * simd_w * 2
peak_bw = hw["bw_gbs"] * 0.5
return float(min(peak_flops, peak_bw))
def get_hardware_profile_tool(tool_args: dict[str, Any], state) -> dict[str, Any]:
"""Return the episode's hardware profile + Roofline bound.
No arguments — the profile is fixed at episode start.
"""
hw = state.hardware_profile
return {
"id": hw.get("id", "unknown"),
"cores": hw["cores"],
"freq_ghz": hw["freq_ghz"],
"l1_kb": hw["l1_kb"],
"simd": hw["simd"],
"bandwidth_gbs": hw["bw_gbs"],
"roofline_bound_gflops": roofline_bound(hw),
# Extra context the agent may use
"simd_width_floats": SIMD_WIDTH.get(hw["simd"], 1),
"bytes_per_flop_threshold": 1.0 / max(roofline_bound(hw), 0.001),
}
__all__ = ["get_hardware_profile_tool", "roofline_bound", "SIMD_WIDTH"]