Add iris.bench benchmarking framework

[mawad-amd]

Summary

• Add iris.bench, a declarative GPU benchmarking framework inspired by NVBench and Google Benchmark
• Write ~25 lines instead of ~350 to benchmark a kernel — framework handles distributed init, parameter sweeps, timing via do_bench(), and output formatting
• Supports console table, JSON, and CSV output formats
• Axis sweeps with cartesian product, CLI overrides (--axis_M=pow2:8:13), and skip filters (--skip_dtype=bf16)
• num_ranks is a special axis that controls process spawning via elastic_launch
• Includes sample all-gather benchmark at benchmark/examples/bench_all_gather.py

Usage

@bench.register
@bench.axis("M", bench.power_of_two(8, 13))
@bench.axis("dtype", [torch.float16, torch.float32])
def all_gather(state, ctx):
    inp = ctx.zeros((state["M"], 1024), dtype=state["dtype"])
    out = ctx.zeros((ctx.get_num_ranks() * state["M"], 1024), dtype=state["dtype"])
    state.set_bytes((ctx.get_num_ranks() - 1) * state["M"] * 1024 * inp.element_size())
    state.exec(lambda: ctx.ccl.all_gather(out, inp, config=config))

if __name__ == "__main__":
    bench.main()

python bench_all_gather.py --axis_num_ranks=2 --axis_M=1024,4096

Test plan

• Tested on MI300X (alola cluster, 2 ranks)
• Console table output with bandwidth
• JSON output format
• CSV output format
• Axis overrides (explicit lists, pow2 ranges)
• Skip filters
• File output (--benchmark_out)
• CI green

🤖 Generated with Claude Code

[Copilot]

Pull request overview

Adds iris.bench, a declarative distributed GPU benchmarking framework with axis sweeps, CLI overrides, and multi-format output.

Changes:

• Introduces core benchmark registry/types (State, decorators, range helpers) and a distributed runner using elastic_launch.
• Adds console/JSON/CSV result formatting plus CLI parsing for axis overrides and skip filters.
• Includes a sample all-gather benchmark and exposes iris.bench from the top-level package.

Reviewed changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 8 comments.

Show a summary per file

File	Description
iris/bench/_runner.py	Implements distributed execution, CLI parsing, and output formatting for benchmarks.
iris/bench/_core.py	Defines benchmark/state abstractions, registry, decorators, and axis-range helpers.
iris/bench/init.py	Public API exports and extensive usage documentation for the new framework.
iris/init.py	Exposes bench as a top-level iris submodule.
benchmark/examples/bench_all_gather.py	Demonstrates an all-gather benchmark written with iris.bench.