xdsl: benchmark convert_type with vs without @cache (context: xdslproject/xdsl#5909)

bench_convert_type.py

#!/usr/bin/env python3
# /// script
# requires-python = ">=3.12"
# dependencies = ["xdsl @ git+https://github.com/xdslproject/xdsl@d099652fb82c1abb47ed1880e644e44a620d970d", "llvmlite>=0.47"]
# ///
# Benchmark: does @cache on convert_type help?
# Run: uv run bench_convert_type.py
# Context: https://github.com/xdslproject/xdsl/pull/5909

import time
from functools import cache
from statistics import mean, stdev

import llvmlite.ir as ir

from xdsl.backend.llvm.convert_type import _TYPE_CONVERTERS
from xdsl.dialects.builtin import (
    ArrayAttr,
    ComplexType,
    Float16Type,
    Float32Type,
    Float64Type,
    FunctionType,
    IntAttr,
    IntegerType,
    TupleType,
    VectorType,
    f32,
    f64,
    i32,
    i64,
)
from xdsl.dialects.llvm import (
    LLVMArrayType,
    LLVMFunctionType,
    LLVMPointerType,
    LLVMStructType,
    LLVMVoidType,
)
from xdsl.ir import Attribute
from xdsl.utils.exceptions import LLVMTranslationException

DISTINCT_TYPES: list[Attribute] = [
    IntegerType(1),
    IntegerType(8),
    IntegerType(16),
    IntegerType(32),
    IntegerType(64),
    Float16Type(),
    Float32Type(),
    Float64Type(),
    LLVMVoidType(),
    LLVMPointerType(),
    LLVMPointerType(IntAttr(1)),
    LLVMPointerType(IntAttr(2)),
    LLVMArrayType(16, i32),
    LLVMArrayType(64, f32),
    LLVMStructType.from_type_list([i32, f32]),
    LLVMStructType.from_type_list([i64, f64, i32]),
    VectorType(i32, [4]),
    VectorType(f32, [8]),
    ComplexType(f32),
    ComplexType(f64),
    TupleType(ArrayAttr([i32, f32])),
    LLVMFunctionType([i32, f32], i64),
    FunctionType.from_lists([i32, f32], [i64]),
]

# 200 functions x 21 types = 4200 calls, simulates convert_module
N = len(DISTINCT_TYPES)
WORKLOAD = [DISTINCT_TYPES[(fn + arg) % N] for fn in range(200) for arg in range(21)]


def make_convert_type(*, use_cache: bool):
    converters = dict(_TYPE_CONVERTERS)

    def impl(type_attr: Attribute) -> ir.Type:
        try:
            return converters[type(type_attr)](type_attr)
        except KeyError:
            raise LLVMTranslationException(f"Type not supported: {type_attr}")

    return cache(impl) if use_cache else impl


def bench(label: str, use_cache: bool, repeats: int = 500, warmups: int = 5) -> float:
    for _ in range(warmups):
        fn = make_convert_type(use_cache=use_cache)
        for t in WORKLOAD:
            fn(t)
    times: list[float] = []
    for _ in range(repeats):
        fn = make_convert_type(use_cache=use_cache)
        start = time.perf_counter()
        for t in WORKLOAD:
            fn(t)
        times.append(time.perf_counter() - start)
    avg = mean(times)
    sd = stdev(times)
    print(f"  {label}: {avg*1e6:8.1f} us  (+/- {sd*1e6:.1f} us)  [{repeats} repeats, {len(WORKLOAD)} calls each]")
    return avg


if __name__ == "__main__":
    print(f"Workload: {len(WORKLOAD)} convert_type calls per iteration\n")
    avg_cached = bench("cached  ", use_cache=True)
    avg_uncached = bench("uncached", use_cache=False)
    ratio = avg_cached / avg_uncached
    print(f"\n  ratio (cached / uncached): {ratio:.2f}x")
    if ratio > 0.9:
        print("  -> @cache provides no meaningful speedup")

Output (Apple M3 Pro, Python 3.12):

Workload: 4200 convert_type calls per iteration

  cached  :   1802.8 us  (+/- 327.3 us)  [500 repeats, 4200 calls each]
  uncached:   1868.5 us  (+/- 70.2 us)  [500 repeats, 4200 calls each]

  ratio (cached / uncached): 0.96x
  -> @cache provides no meaningful speedup