AI-bench: Unified AI Benchmarking Suite

A benchmarking framework for evaluating AI kernel implementations across multiple backends (PyTorch, Triton, Helion) and devices (CPU, XPU).

Installation

The project is using uv package manager.

uv can be installed locally using:

pip install uv

The project can be installed with appropriate device-specific extensions using:

# CPU only
uv sync --extra cpu

# CPU + XPU
uv sync --extra xpu

Usage

Command Line Interface

After installation, the ai-bench command is available:

# Show help
ai-bench --help

# PyTorch on CPU (default)
ai-bench

# PyTorch on XPU
ai-bench --xpu

# PyTorch compile on XPU
ai-bench --xpu --torch-compile

# Triton on XPU
ai-bench --xpu --triton

# Helion on XPU
ai-bench --xpu --helion

# Benchmark mode (with timing)
ai-bench --xpu --bench

# Log results to CSV
ai-bench --xpu --bench --csv results.csv --note "baseline run"

Using Scripts Directly

Run KernelBench problems with different backends and devices:

# PyTorch on CPU (default)
python infra/scripts/run_kernel_bench.py

# PyTorch on XPU
python infra/scripts/run_kernel_bench.py --xpu

# PyTorch compile on XPU
python infra/scripts/run_kernel_bench.py --xpu --torch-compile

# Triton on XPU
python infra/scripts/run_kernel_bench.py --xpu --triton

# Helion on XPU
python infra/scripts/run_kernel_bench.py --xpu --helion

# Benchmark mode (with timing)
python infra/scripts/run_kernel_bench.py --xpu --bench

# Triton benchmark on XPU
python infra/scripts/run_kernel_bench.py --xpu --triton --bench

As a Library

import ai_bench
import torch

# Configure paths if running outside project root
ai_bench.configure(
    specs_dir="/path/to/specs",
    kernels_dir="/path/to/kernels",
)

# Create benchmark runner
runner = ai_bench.KernelBenchRunner(
    spec_type=ai_bench.SpecKey.V_BENCH_GPU,
    device=torch.device("xpu"),
    backend=ai_bench.Backend.PYTORCH,
    flops_unit=ai_bench.FlopsUnit.TFLOPS,
    mem_bw_unit=ai_bench.MemBwUnit.GBS,
    csv_path="results.csv",
)
runner.run_kernels()

CSV Logging

Benchmark results can be logged to a CSV file using the --csv option:

# Log results to CSV
python infra/scripts/run_kernel_bench.py --xpu --triton --bench --csv results.csv

# Add a note to identify the run
python infra/scripts/run_kernel_bench.py --xpu --triton --bench --csv results.csv --note "BMG card test"

The CSV file includes the following columns:

• kernel_name: Name of the kernel
• kernel_type: Backend used (pytorch/triton)
• problem_level: KernelBench problem level
• flops: Number of floating-point operations
• flops_val: Computed FLOPS value
• flops_unit: FLOPS unit (GFLOPS/TFLOPS)
• flops_note: FLOPS measurement annotation (see 'Notes legend')
• mem_bytes: Number of memory bytes transferred - input reads + output writes
• mem_bw_val: Computed memory bandwidth value
• mem_bw_unit: Memory bandwidth unit (MB/s or GB/s)
• mem_note: Memory measurement annotation (see 'Notes legend')
• time_us: Execution time in microseconds
• input_values: Input dimensions as JSON
• note: User-provided note

Additionally, any environment variables prefixed with AIBENCH_ are automatically captured and included in the CSV output. This is useful for recording system configuration:

# Set environment variables for tracking
export AIBENCH_CARD="BMG"
export AIBENCH_SYSTEM="TestRig1"
python infra/scripts/run_kernel_bench.py --xpu --triton --bench --csv results.csv

Notes legend:

• ⚠️: estimated value, use with caution

Command Line Options

Testing

Run tests with pytest:

pytest -v

Linting

The project uses pre-commit to run various checks automatically.

All checks can be run using:

pre-commit run -a

Config variables

Environment variables used for project configuration:

License

MIT License - see LICENSE for details.