AI-bench: Unified AI Benchmarking Suite

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

	PyTorch	Triton	Helion	MLIR	Gluon	SYCL
CPU	✅	✅1	❌	✅2	❌	❌
XPU	✅	✅	✅	✅2*	✅	✅
CUDA	✅	✅	✅	❌	❌	❌

✅ - Supported ⚠️ - Partially implemented ❌ - Unsupported

¹ via triton-lang/triton-cpu
² via llvm/lighthouse

*NOTE: MLIR XPU requires custom LLVM build with Intel GPU support enabled.
Override PYTHONPATH to point to mlir_core packages and ensure LLVM libs can be found.
See Lighthouse XeGPU example for more details.

Project Docs

• Problem Specification Files

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 and backend extensions using:

# CPU only
uv sync --extra cpu

# CPU + XPU
uv sync --extra xpu

# CPU + CUDA
uv sync --extra cuda

# CPU + Triton-CPU backend (builds Triton from source)
uv sync --extra cpu --extra triton-cpu

# CPU + MLIR backend
uv sync --extra cpu --extra mlir

Usage

Command Line Interface

After installation, the ai-bench command is available:

# Show help
ai-bench --help

# PyTorch on CPU (default)
ai-bench

# Triton on CPU
ai-bench --triton

# MLIR on CPU
ai-bench --mlir

# PyTorch on XPU
ai-bench --xpu

# PyTorch on CUDA GPU
ai-bench --cuda

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

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

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

# Gluon on XPU
ai-bench --xpu --gluon

# SYCL on XPU (requires env setup, see below)
ai-bench --xpu --sycl

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

# Benchmark a custom spec variant
ai-bench --variant custom-variant

# Validate a custom spec variant without benchmarking
ai-bench --variant custom-variant --ci

# Run only float32 variants
ai-bench --bench --dtype float32

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

# Run a single kernel with a problem specification
ai-bench --kernel /path/to/kernel.py /path/to/spec.yaml

# Run a single kernel with a problem specification on XPU
ai-bench --kernel /path/to/kernel.py /path/to/spec.yaml --xpu

Using ai-bench-compare command, KernelBench performance can be compared across multiple backends:

# Show help
ai-bench-compare --help

# Comparison for the given problem on CPU (default)
ai-bench-compare --problem level1/1_Square_matrix_multiplication_

# Compare PyTorch and Triton backends on XPU
ai-bench-compare --problem level2/99_Matmul_GELU_Softmax --backend pytorch triton --xpu

# Compare only bfloat16 variants
ai-bench-compare --problem level1/1_Square_matrix_multiplication_ --dtype bfloat16

Optionally, custom CLI autocompletion is available for certain scripts. It can be activated using:

activate-global-python-argcomplete --user

As a Library

import ai_bench
import torch

# Create a single kernel benchmark
kernel_runner = ai_bench.KernelRunner(
    spec_type=ai_bench.SpecKey.V_BENCH_CPU,
    device=torch.device("cpu"),
    backend=ai_bench.Backend.PYTORCH,
    flops_unit=ai_bench.FlopsUnit.TFLOPS,
    mem_bw_unit=ai_bench.MemBwUnit.GBS,
)
kernel_runner.run_kernel_spec("path/to/kernel.py", "path/to/spec.yaml")

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

# Create KernelBench benchmark runner
kb_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",
)
kb_runner.run_kernels()

CSV Logging

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

# Log results to CSV
ai-bench --xpu --triton --bench --csv results.csv

# Add a note to identify the run
ai-bench --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
• dtype: Variant's data type if specified
• 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"
ai-bench --xpu --triton --bench --csv results.csv

Notes legend:

• ⚠️: estimated value, use with caution

Command Line Options

Option	Description
--kernel KERNEL_PATH SPEC_PATH	Run a kernel with a spec (default: KernelBench)
--xpu	Run on Intel XPU (default: CPU)
--cuda	Run on Nvidia GPU (default: CPU)
--triton	Use Triton backend (default: PyTorch eager)
--torch-compile	Use PyTorch compile mode (default: PyTorch eager)
--helion	Use Helion backend (default: PyTorch eager)
--mlir	Use MLIR backend (default: PyTorch eager)
--gluon	Use Gluon backend (default: PyTorch eager)
--sycl	Use SYCL backend (default: PyTorch eager)
--bench	Run benchmarks with timing (default: CI validation)
--ci	Force a single validation run only for any variant
--variant VARIANT	Run only specified spec VARIANT
--dtype DTYPE	Run only variants with specified DTYPE
--gflops	Report GFLOPS (default: TFLOPS)
--mbs	Report MB/s (default: GB/s)
--csv PATH	Log results to specified CSV file
--note TEXT	Add a note to CSV output for identifying runs
--specs-dir PATH	Path to specs directory (CLI only)
--kernels-dir PATH	Path to kernels directory (CLI only)
--triton-kernels-dir PATH	Path to Triton kernels directory (CLI only)
--helion-kernels-dir PATH	Path to Helion kernels directory (CLI only)
--mlir-kernels-dir PATH	Path to MLIR kernels directory (CLI only)
--gluon-kernels-dir PATH	Path to Gluon kernels directory (CLI only)
--sycl-kernels-dir PATH	Path to SYCL kernels directory (CLI only)
--env-file PATH	Path to .env file (default: auto-detect)
--no-env	Disable loading .env config

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

SYCL Backend Setup

The SYCL backend compiles and runs C++ CUTLASS kernels via icpx. It requires:

1. Intel oneAPI DPC++ compiler (icpx) — install via Intel oneAPI Base Toolkit
2. Intel GPU drivers (Level Zero runtime)
3. sycl-tla headers — git clone https://github.com/intel/sycl-tla

Set the following environment variables before running:

export AIBENCH_SYCL_COMPILER=icpx
export AIBENCH_SYCL_INCLUDE=/path/to/sycl-tla/include:/path/to/sycl-tla/tools/util/include:/path/to/sycl-tla/examples/common
ai-bench --xpu --sycl --bench

A helper script env/sycl.sh is provided for convenience:

export SYCL_TLA_DIR=/path/to/sycl-tla
source env/sycl.sh
ai-bench --xpu --sycl --bench

Config variables

Environment variables used for project configuration:

Variable	Description
AIBENCH_LOG=INFO|DEBUG|...	Globally overrides logging level
AIBENCH_SPECS_DIR	Path to specs directory
AIBENCH_KERNELS_DIR	Path to PyTorch kernels directory
AIBENCH_TRITON_KERNELS_DIR	Path to Triton kernels directory
AIBENCH_HELION_KERNELS_DIR	Path to Helion kernels directory
AIBENCH_MLIR_KERNELS_DIR	Path to MLIR kernels directory
AIBENCH_MLIR_SCHEDULES_DIR	Path to MLIR CPU pipeline schedules (YAML descriptors) directory
AIBENCH_MLIR_LIB_PATH	Paths to MLIR shared libraries (colon separated)
AIBENCH_MLIR_DUMP	Dump imported MLIR IR
AIBENCH_MLIR_DUMP_OBJ	Dump jitted MLIR to an object file
AIBENCH_GLUON_KERNELS_DIR	Path to Gluon kernels directory
AIBENCH_SYCL_KERNELS_DIR	Path to SYCL kernels directory
AIBENCH_SYCL_COMPILER	Path to SYCL compiler (default: icpx)
AIBENCH_SYCL_INCLUDE	Colon-separated include paths for CUTLASS/SYCL headers
AIBENCH_SYCL_FLAGS	Extra compiler flags (space-separated)
AIBENCH_WARMUP	Override number of warmup iterations
AIBENCH_REP	Override number of timed iterations
AIBENCH_CPU_MIN_CACHE_NUKE_MIB	Minimum memory size (in MiB) for a cache-nuking GEMM between timed iterations on CPU

License

MIT License - see LICENSE for details.