run_benchmark.py

"""
Benchmark constant folding over a module of integer addition of constants using
xDSL's pattern rewriting machinery. The workload module is generated by the
``get_constant_folding_module()`` function, which is then transformed by the
``xdsl.transforms.test_constant_folding.TestConstantFoldingPass`` xDSL
transformation pass.

xDSL is a Python-native compiler framework built around SSA-based intermediate
representations. It re-implements many of MLIR's data structures and methods in
Python. This benchmark exercises the simple pattern rewriting transformation of
constant folding. This is a fair proxy for pattern rewriting transformations in
general, which are a major component of MLIR-like compilers in lowering passes.

Pattern rewriting in both xDSL and MLIR is a pointer-chasing, unstructured
workload, which makes it hard to optimise ahead-of-time. This diminishes the
traditional performance advantage of ahead-of-time compiled languages such as
C++ over dynamic languages such as Python -- making it an interesting benchmark.
More information about the design and impact of this benchmark can be found in
the Master's thesis ``Performance and Dynamism in User-extensible Compiler
Infrastructures'', which is `available on GitHub
<https://github.com/EdmundGoodman/masters-project>`_.
"""

import random

import pyperf

from xdsl.context import Context
from xdsl.dialects.arith import AddiOp, Arith, ConstantOp
from xdsl.dialects.builtin import (
    Builtin,
    IntegerAttr,
    ModuleOp,
    i32,
)
from xdsl.dialects.test import TestOp
from xdsl.ir import Operation
from xdsl.transforms.test_constant_folding import TestConstantFoldingPass


RANDOM_SEED = 0


def get_constant_folding_module(size: int = 100) -> ModuleOp:
    """Generate an integer addition constant folding workload of a given size.

    The output of running the command
    `print(WorkloadBuilder().constant_folding_module(size=5))` is shown
    below:

    ```mlir
    "builtin.module"() ({
        %0 = "arith.constant"() {"value" = 865 : i32} : () -> i32
        %1 = "arith.constant"() {"value" = 395 : i32} : () -> i32
        %2 = "arith.addi"(%1, %0) : (i32, i32) -> i32
        %3 = "arith.constant"() {"value" = 777 : i32} : () -> i32
        %4 = "arith.addi"(%3, %2) : (i32, i32) -> i32
        %5 = "arith.constant"() {"value" = 912 : i32} : () -> i32
        "test.op"(%4) : (i32) -> ()
    }) : () -> ()
    ```
    """
    assert size >= 0
    random.seed(RANDOM_SEED)
    ops: list[Operation] = []
    ops.append(ConstantOp(IntegerAttr(random.randint(1, 1000), i32)))
    for i in range(1, size + 1):
        if i % 2 == 0:
            ops.append(AddiOp(ops[i - 1], ops[i - 2]))
        else:
            ops.append(ConstantOp(IntegerAttr(random.randint(1, 1000), i32)))
    ops.append(TestOp([ops[(size // 2) * 2]]))
    return ModuleOp(ops)


def bench_xdsl_constant_folding(loops: int, size: int) -> float:
    """Benchmark constant folding integer additions with xDSL."""
    ctx = Context(allow_unregistered=True)
    ctx.load_dialect(Arith)
    ctx.load_dialect(Builtin)
    constant_folding_pass = TestConstantFoldingPass()
    constant_folding_workload = get_constant_folding_module(size)

    range_it = range(loops)
    t0 = pyperf.perf_counter()
    for _ in range_it:
        constant_folding_workload_mutable = constant_folding_workload.clone()
        constant_folding_pass.apply(ctx, constant_folding_workload_mutable)
    return pyperf.perf_counter() - t0


def add_cmdline_args(cmd, args):
    """Add command line arguments to a pyperf runner"""
    cmd.extend(("--size", str(args.size)))


if __name__ == "__main__":
    runner = pyperf.Runner(add_cmdline_args=add_cmdline_args)
    runner.metadata["description"] = "Benchmark xDSL constant folding integer addition"
    runner.argparser.add_argument(
        "--size",
        type=int,
        default=1000,
        help="Number of integer additions (default: 1000)",
    )
    args = runner.parse_args()
    runner.bench_time_func("xdsl_constant_fold", bench_xdsl_constant_folding, args.size)