Hespas is a distributed ML performance estimation tool built on StableHLO. It takes a StableHLO workload representation of an ML model, splits it into compute and communication regions, estimates compute performance using analytical, simulation, or profiling-based backends, and outputs Chakra execution traces for network simulation with ASTRA-sim.
One workload representation. Multiple simulation backends. Multiple hardware targets.
git clone https://github.com/imec-int/hespas.git
cd hespas
pip install .Or with uv:
uv syncGenerate a Chakra trace using the roofline estimator:
hespas_chakra_gen tests/fixtures/configs/config_roofline_a100.jsonUse a hardware preset with a custom workload:
hespas_chakra_gen configs/nodes/A100_PCIe_40GB/roofline/config.json \
--mlir_file tests/fixtures/mlir/jax_example.mlir \
--output output/jax_example_a100| Backend | Type | Hardware Required |
|---|---|---|
| Roofline | Analytical | None |
| XLA | Profiling | NVIDIA GPU + CUDA |
| IREE | Compilation + profiling | CPU or GPU |
| Command | Description |
|---|---|
hespas_chakra_gen |
Main tool — split, estimate, and generate Chakra traces |
hespas_mlir_splitter |
Split MLIR into compute/communication modules |
hespas_mlir_analyzer |
Analyze StableHLO operator statistics |
Evaluating Cross-Architecture Performance Modeling of Distributed ML Workloads Using StableHLO
Jonas Svedas, Nathan Laubeuf, Ryan Harvey, Arjun Singh, Changhai Man, Abubakr Nada, Tushar Krishna, James Myers, Debjyoti Bhattacharjee
IEEE International Symposium on Performance Analysis of Systems and Software (ISPASS), 2026
@inproceedings{svedas2026stablehlo,
title = {Evaluating Cross-Architecture Performance Modeling of Distributed ML Workloads Using StableHLO},
author = {Svedas, Jonas and Laubeuf, Nathan and Harvey, Ryan and Singh, Arjun and Man, Changhai and Nada, Abubakr and Krishna, Tushar and Myers, James and Bhattacharjee, Debjyoti},
year = {2026},
booktitle = {Proceedings of ISPASS 2026},
organization = {IEEE}
}This work is funded by the Advanced Research + Invention Agency (ARIA).
MIT — © imec
