Add extensible benchmark suite (TensorOperationsBenchmarks)

benchmark/Project.toml

@@ -0,0 +1,32 @@
+name = "TensorOperationsBenchmarks"
+uuid = "d983fc97-4e87-46ba-abd9-b4864e69d4dd"
+authors = ["Lukas Devos <lukas.devos@ugent.be>"]
+version = "0.1.0"
+
+[deps]
+ArgParse = "c7e460c6-2fb9-53a9-8c5b-16f535851c63"
+BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
+DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+PkgBenchmark = "32113eaa-f34f-5b0d-bd6c-c81e245fc73d"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+Strided = "5e0ebb24-38b0-5f93-81fe-25c709ecae67"
+TensorOperations = "6aa20fa7-93e2-5fca-9bc0-fbd0db3c71a2"
+
+[compat]
+ArgParse = "1"
+BenchmarkTools = "1"
+DataFrames = "1"
+LinearAlgebra = "1.10"
+PkgBenchmark = "0.2"
+Random = "1.10"
+Strided = "2.6"
+TensorOperations = "5.8"
+Test = "1"
+julia = "1.11"
+
+[extras]
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+
+[targets]
+test = ["Test"]

benchmark/README.md

@@ -0,0 +1,81 @@
+# TensorOperationsBenchmarks
+
+Extensible benchmark suite for TensorOperations.jl. Downstream packages (e.g. a
+symmetric/block-sparse tensor package) can plug in their own tensor type and run the same
+standardized shapes.
+
+## Running
+
+```julia
+julia --project=. -e '
+    using TensorOperationsBenchmarks
+    using TensorOperations: StridedNative, StridedBLAS
+    providers = [ArrayProvider{Float64}(; backend=StridedNative()),
+                 ArrayProvider{Float64}(; backend=StridedBLAS())]
+    suite = build_suite(providers)
+    results = run(suite)
+    rows = resultstable(results)
+'
+```
+
+Or for commit-to-commit regression comparison via PkgBenchmark. Both scripts are `@main` apps
+(Julia 1.11+), so `--help` works and `ARGS` are parsed the normal way:
+
+```
+julia --project=. scripts/run_benchmarks.jl --threads 1 4 --blas-threads 1 4
+julia --project=. scripts/show_benchmarks.jl results_t4_blas4_strided.json   # requires CairoMakie
+```
+
+## Categories
+
+- `:contract` -- generic pairwise contractions: a synthetic parametric shape family (tagged
+  `synthetic`) plus 24 real quantum-chemistry contractions (CCSD, CCSD(T), AO2MO, INTENSLI) from
+  the [TCCG benchmark](https://github.com/HPAC/tccg) (tagged `tccg`). Each synthetic shape also
+  comes in up to 4 label-order layouts (tagged `gemm_ready`/`a_permuted`/`b_permuted`/
+  `both_permuted`): `gemm_ready` is directly reshapeable to a BLAS call, the others interleave
+  open/contracted labels so no reshape or transpose flag suffices -- a real permutation is
+  required, which is what actually separates `StridedNative` from `StridedBLAS`.
+- `:permute` -- permutation-only (`tensorcopy!`) cost.
+- `:trace` -- partial and full traces.
+- `:mixed_precision` -- differing input/output element types (e.g. `Float32 x Float32 ->
+  Float64`, mixed real/complex).
+- `:network` -- multi-tensor-network motifs, tagged by `topic`: `mps` (MPS/MPO DMRG
+  effective-Hamiltonian, 1-site and 2-site "theta", swept over bond `D`), `ctmrg` (CTMRG
+  corner-growth step for 2D PEPS, swept over environment bond `chi`), `trg` (TRG plaquette
+  contraction, swept over bond `chi`).
+
+Not yet implemented, but addable without a redesign: MERA, contraction-order/path-finding timing.
+
+## Adding a category
+
+New file under `src/categories/`, define `mysizes -> Vector{BenchmarkCase}` building
+`ContractSpec`/`TraceSpec`/`AddSpec`/`NetworkSpec` values, `include` it, call
+`register_category!(:mycategory, mygenerator)`. Nothing else changes. Filtering uses
+`BenchmarkGroup`'s native tags (`suite[@tagged "..."]`, see `BenchmarkCase`'s docstring for how
+tags are derived) -- no bespoke filter API to learn.
+
+## Plugging in a downstream tensor type
+
+```julia
+struct MyProvider <: AbstractProvider end
+TensorOperations.scalartype(::MyProvider) = Float64
+TensorOperationsBenchmarks.randtensor(::MyProvider, labels, dims, T) = # build a random tensor
+# optional: backend(p), allocator(p), label(p), supports(p, category), rng(p)
+```
+
+Then `build_suite([MyProvider(), ArrayProvider{Float64}()])` compares directly against
+TensorOperations.jl's own backends. `ArrayProvider` shows the recommended `randtensor` pattern:
+allocate via `TensorOperations.tensoralloc` (goes through `allocator`) and fill from a stored,
+stateful `rng` field seeded once at construction, so runs are reproducible but tensors within a
+run still differ.
+
+## Threading and precision
+
+Thread config is not a `build_suite` axis -- setting it per-case would count the switch itself
+as part of the timing. Instead `set_threads!` is called once, at the process level, before a
+suite is built (`benchmarks.jl` reads `TOB_BLAS_THREADS`/`TOB_STRIDED_THREADS`, set by
+`run_benchmarks.jl --blas-threads/--strided-threads`). `with_threads(f, cfg)` is available for
+ad hoc comparisons, wrapping a whole `run(suite)` call and restoring afterwards.
+
+Mixed precision: set `TA`/`TB`/`TC` explicitly on a spec (see `mixed_precision.jl`); other
+categories leave them `nothing` (provider's default `scalartype`).

benchmark/benchmarks.jl

@@ -0,0 +1,17 @@
+# PkgBenchmark entrypoint: expects a top-level `const SUITE`. Thread counts come from env vars
+# (set by scripts/run_benchmarks.jl) and are applied once, before SUITE is built.
+using TensorOperationsBenchmarks
+using TensorOperations: StridedNative, StridedBLAS
+
+set_threads!(
+    ThreadConfig(;
+        blas = tryparse(Int, get(ENV, "TOB_BLAS_THREADS", "")),
+        strided = tryparse(Int, get(ENV, "TOB_STRIDED_THREADS", "")),
+    )
+)
+
+const ELTYPES = (Float64, ComplexF64)
+const BACKENDS = (StridedNative(), StridedBLAS())
+const PROVIDERS = [ArrayProvider{T}(; backend) for T in ELTYPES for backend in BACKENDS]
+
+const SUITE = build_suite(PROVIDERS)

benchmark/scripts/run_benchmarks.jl

@@ -0,0 +1,55 @@
+#!/usr/bin/env julia
+# CLI wrapper around PkgBenchmark.benchmarkpkg. `--threads` relaunches Julia per value (fixed
+# at startup); `--blas-threads`/`--strided-threads` set env vars benchmarks.jl reads.
+#   julia --project=. scripts/run_benchmarks.jl --threads 1 2 4 --blas-threads 1 4 --out results
+using Pkg
+Pkg.activate(joinpath(@__DIR__, ".."))
+
+using ArgParse
+using PkgBenchmark
+
+function parse_commandline(args)
+    s = ArgParseSettings(; description = "Run the TensorOperationsBenchmarks suite via PkgBenchmark.")
+    @add_arg_table! s begin
+        "--threads"
+        help = "outer Julia process thread count(s) to sweep (relaunches Julia per value)"
+        arg_type = Int
+        nargs = '*'
+        default = [Threads.nthreads()]
+        "--blas-threads"
+        help = "inner BLAS thread count(s) to sweep"
+        arg_type = Int
+        nargs = '*'
+        default = Int[]
+        "--strided-threads"
+        help = "inner Strided.jl thread count(s) to sweep"
+        arg_type = Int
+        nargs = '*'
+        default = Int[]
+        "--out"
+        help = "output file prefix (a suffix identifying the thread combo and `.json` are appended)"
+        default = "results"
+    end
+    return parse_args(args, s)
+end
+
+function (@main)(args)
+    opts = parse_commandline(args)
+    blascounts = isempty(opts["blas-threads"]) ? [nothing] : opts["blas-threads"]
+    stridedcounts = isempty(opts["strided-threads"]) ? [nothing] : opts["strided-threads"]
+
+    for nthreads in opts["threads"], blas in blascounts, strided in stridedcounts
+        @info "Running benchmarks" nthreads blas strided
+        withenv(
+            "TOB_BLAS_THREADS" => blas === nothing ? "" : string(blas),
+            "TOB_STRIDED_THREADS" => strided === nothing ? "" : string(strided),
+        ) do
+            cfg = BenchmarkConfig(; juliacmd = `julia -t $nthreads -O3`)
+            results = benchmarkpkg(dirname(@__DIR__), cfg)
+            outfile = "$(opts["out"])_t$(nthreads)_blas$(blas)_strided$(strided).json"
+            writeresults(outfile, results)
+            @info "Wrote $outfile"
+        end
+    end
+    return 0
+end

benchmark/scripts/show_benchmarks.jl

@@ -0,0 +1,55 @@
+#!/usr/bin/env julia
+# Plots time/GFLOPs-vs-size curves from a run_benchmarks.jl result JSON. Requires CairoMakie
+# (`Pkg.add("CairoMakie")` into this environment first -- not a package dependency, it's heavy).
+using Pkg
+Pkg.activate(joinpath(@__DIR__, ".."))
+
+using ArgParse
+using PkgBenchmark
+using CairoMakie
+using TensorOperationsBenchmarks
+
+function parse_commandline(args)
+    s = ArgParseSettings(; description = "Plot GFLOP/s-vs-size scaling curves from a PkgBenchmark result.")
+    @add_arg_table! s begin
+        "resultfile"
+        help = "path to a PkgBenchmark result JSON, as written by run_benchmarks.jl"
+        required = true
+        "--out"
+        help = "output image path (default: replace the input's extension with .png)"
+        default = nothing
+    end
+    return parse_args(args, s)
+end
+
+function (@main)(args)
+    opts = parse_commandline(args)
+    results = PkgBenchmark.readresults(opts["resultfile"])
+    group = PkgBenchmark.benchmarkgroup(results)
+
+    rows = resultstable(group)
+
+    fig = Figure(; size = (1000, 800))
+    categories = unique(r.category for r in eachrow(rows))
+    for (i, category) in enumerate(categories)
+        ax = Axis(
+            fig[fldmod1(i, 2)...]; xscale = log2, yscale = log10,
+            title = category, xlabel = "size", ylabel = "GFLOP/s"
+        )
+        catrows = filter(r -> r.category == category, rows)
+        for provider in unique(r.provider for r in eachrow(catrows))
+            provrows = filter(r -> r.provider == provider, catrows)
+            sort!(provrows; by = r -> get(r.params, :dim, get(r.params, :D, 0)))
+            xs = [get(r.params, :dim, get(r.params, :D, 0)) for r in eachrow(provrows)]
+            ys = [r.gflops for r in eachrow(provrows)]
+            lines!(ax, xs, ys; label = provider)
+            scatter!(ax, xs, ys)
+        end
+        axislegend(ax)
+    end
+
+    outfile = something(opts["out"], splitext(opts["resultfile"])[1] * ".png")
+    save(outfile, fig)
+    @info "Wrote $outfile"
+    return 0
+end

benchmark/src/TensorOperationsBenchmarks.jl

@@ -0,0 +1,49 @@
+module TensorOperationsBenchmarks
+
+using LinearAlgebra: BLAS
+using Strided: Strided
+using Random: Random, randn!
+using DataFrames: DataFrame
+using BenchmarkTools
+using TensorOperations
+using TensorOperations: DefaultBackend, DefaultAllocator, AbstractBackend
+
+# Specs (pure data) and the cost model computed from them.
+include("specs.jl")
+include("cost.jl")
+
+# The downstream extension points: what tensor type to run against, and how many threads.
+include("provider.jl")
+include("threading.jl")
+
+# The case registry, and turning a (spec, provider) pair into an executable benchmark.
+include("registry.jl")
+include("lowering.jl")
+
+# Suite assembly and reporting.
+include("suite.jl")
+include("report.jl")
+
+# Categories: tccg.jl/mps.jl/ctmrg.jl/trg.jl define generators merged by contract.jl/network.jl
+# (include order doesn't matter -- generators are only called after the whole module loads).
+include("categories/tccg.jl")
+include("categories/contract.jl")
+include("categories/permute.jl")
+include("categories/trace.jl")
+include("categories/mixed_precision.jl")
+include("categories/mps.jl")
+include("categories/ctmrg.jl")
+include("categories/trg.jl")
+include("categories/network.jl")
+
+export AbstractCaseSpec, AddSpec, TraceSpec, ContractSpec, NetworkSpec
+export flops, bytes
+export AbstractProvider, ArrayProvider, scalartype, randtensor, backend, allocator, label,
+    supports, rng
+export ThreadConfig, with_threads, set_threads!
+export BenchmarkCase, register_category!, REGISTRY, default_sizes, within_memory_budget
+export build_suite
+export resultstable
+export @tagged
+
+end # module

benchmark/src/categories/contract.jl

@@ -0,0 +1,69 @@
+# Pairwise contractions: a synthetic parametric shape family, plus the real TCCG equations
+# (tccg.jl) -- both just `ContractSpec`, merged into one category, tagged (`:synthetic`/`:tccg`)
+# for filtering via `@tagged` (see registry.jl). Sizes mix power-of-two with off-by-one values.
+#
+# Each shape gets up to 4 label-order "layouts": `(openA...,contract...)`/`(contract...,openB...)`
+# is directly reshapeable to GEMM (no data movement); interleaving open and contracted labels
+# (e.g. `[a1,c1,a2,c2]`) cannot be expressed as a single reshape or BLAS transpose flag, forcing
+# a real permutation -- exactly the case that separates StridedNative from StridedBLAS. Layouts
+# that coincide with the GEMM one (e.g. when a shape has ≤1 contracted index) are skipped.
+
+const CONTRACT_SHAPES = (
+    (1, 1, 1),   # matrix-vector-like
+    (2, 1, 2),   # single shared bond, several open legs each side
+    (2, 2, 2),   # GEMM-like, rank 4 total
+    (1, 3, 1),   # trace-heavy: many contracted, few open
+    (1, 0, 1),   # pure outer product, no contraction
+)
+
+function _interleave(a::Vector{Symbol}, b::Vector{Symbol})
+    n = min(length(a), length(b))
+    return vcat((Symbol[a[i], b[i]] for i in 1:n)..., a[(n + 1):end], b[(n + 1):end])
+end
+
+function _contract_layouts(openA, contract, openB)
+    gemmA, gemmB = vcat(openA, contract), vcat(contract, openB)
+    permA, permB = _interleave(openA, contract), _interleave(contract, openB)
+    candidates = (
+        (:gemm_ready, gemmA, gemmB),
+        (:a_permuted, permA, gemmB),
+        (:b_permuted, gemmA, permB),
+        (:both_permuted, permA, permB),
+    )
+    seen = Set{Tuple{Vector{Symbol}, Vector{Symbol}}}()
+    layouts = Tuple{Symbol, Vector{Symbol}, Vector{Symbol}}[]
+    for (layout, IA, IB) in candidates
+        (IA, IB) in seen && continue
+        push!(seen, (IA, IB))
+        push!(layouts, (layout, IA, IB))
+    end
+    return layouts
+end
+
+function _synthetic_contract_cases(sizes)
+    cases = BenchmarkCase[]
+    for dim in sizes
+        for (nopenA, ncontract, nopenB) in CONTRACT_SHAPES
+            openA = [Symbol("a", i) for i in 1:nopenA]
+            contract = [Symbol("c", i) for i in 1:ncontract]
+            openB = [Symbol("b", i) for i in 1:nopenB]
+            IC = vcat(openA, openB)
+            for (layout, IA, IB) in _contract_layouts(openA, contract, openB)
+                dims = Dict{Symbol, Int}(l => dim for l in vcat(IA, IB))
+                spec = ContractSpec(IA, IB, IC, dims)
+                id = "dim$(dim)_$(nopenA)_$(ncontract)_$(nopenB)_$(layout)"
+                within_memory_budget(spec, id) || continue
+                params = (; dim, nopenA, ncontract, nopenB, layout, source = :synthetic)
+                push!(cases, BenchmarkCase(:contract, id, params, spec))
+            end
+        end
+    end
+    return cases
+end
+
+_contract_cases(sizes) = vcat(_synthetic_contract_cases(sizes), _tccg_cases(sizes))
+
+register_category!(
+    :contract, _contract_cases;
+    sizes = (8, 12, 15, 16, 24, 32, 63, 96, 128, 200, 256)
+)

benchmark/src/categories/ctmrg.jl

@@ -0,0 +1,22 @@
+# CTMRG corner-growth step (boundary-MPS method for 2D PEPS): C-T-T-a, producing an unfused
+# rank-4 corner (chi,chi,D2,D2). Cost ~ O(chi^3*D2^3) (literature O(chi^3*D^6), D2=D^2).
+# `sizes` sweeps environment bond `chi`; PEPS bond `D` is fixed. Merged into `:network` (see
+# network.jl) tagged `params.topic = :ctmrg`.
+
+const CTMRG_PEPS_BOND = 3   # D
+
+function _ctmrg_case(chi)
+    D2 = CTMRG_PEPS_BOND^2
+    indexlists = [
+        [1, 2],            # C: (chi, chi)
+        [1, 3, -10],       # T_left: (chi, D2, chi)
+        [2, 4, -11],       # T_top: (chi, D2, chi)
+        [3, -12, 4, -13],  # a (double-layer PEPS tensor): (D2, D2, D2, D2)
+    ]
+    dims = Dict(1 => chi, 2 => chi, 3 => D2, 4 => D2, 10 => chi, 11 => chi, 12 => D2, 13 => D2)
+    spec = NetworkSpec(indexlists, dims; output = [-10, -11, -12, -13])
+    params = (; chi, D = CTMRG_PEPS_BOND, topic = :ctmrg)
+    return BenchmarkCase(:network, "ctmrg_corner_chi$(chi)", params, spec)
+end
+
+_ctmrg_cases(sizes) = [_ctmrg_case(chi) for chi in sizes]