Set up the benchmark suite

Project.toml

@@ -1,6 +1,6 @@
 name = "TensorAlgebra"
 uuid = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a"
-version = "0.13.1"
+version = "0.13.2"
 authors = ["ITensor developers <support@itensor.org> and contributors"]
 
 [workspace]

benchmark/Project.toml

@@ -1,3 +1,9 @@
 [deps]
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 TensorAlgebra = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a"
+
+[sources.TensorAlgebra]
+path = ".."
+
+[compat]
+BenchmarkTools = "1"

benchmark/README.md

@@ -0,0 +1,34 @@
+# TensorAlgebra benchmarks
+
+A [`BenchmarkTools.jl`](https://github.com/JuliaCI/BenchmarkTools.jl) suite for
+TensorAlgebra. The benchmarks live in the `SUITE` global defined in
+`benchmarks.jl`, covering the core contraction entry points (`contract`,
+`contract!`) across a range of dimensions. Small dimensions track the fixed
+per-call label bookkeeping, large ones the BLAS call, so a regression in either
+regime shows up.
+
+The suite is also run once as a smoke test in the package tests (`test_benchmarks.jl`),
+so it cannot silently fall out of sync with the API.
+
+## Running the suite
+
+```bash
+julia --project=benchmark -e '
+    include("benchmark/benchmarks.jl")
+    using BenchmarkTools
+    display(run(SUITE; verbose = true))'
+```
+
+## Comparing two revisions
+
+With [`AirspeedVelocity.jl`](https://github.com/MilesCranmer/AirspeedVelocity.jl)
+installed in a shared environment, compare the working tree against `main`:
+
+```bash
+benchpkg TensorAlgebra --rev=dirty,main -o benchmark/results/
+benchpkgtable TensorAlgebra --rev=dirty,main -i benchmark/results/
+```
+
+Timings on shared or CI machines are noisy. Compare allocation counts and relative
+numbers rather than absolute times, and run on a quiet machine when absolute
+timings matter.

benchmark/benchmarks.jl

@@ -1,26 +1,50 @@
-using BenchmarkTools
-using TensorAlgebra
+using BenchmarkTools: @benchmarkable, BenchmarkGroup
+using TensorAlgebra: contract, contract!
 
-SUITE = BenchmarkGroup()
+# Benchmarks of the core contraction entry points across a range of dimensions.
+# Small dimensions are dominated by the fixed per-call label bookkeeping, large
+# ones by the BLAS call, so a regression in either regime is visible. The `SUITE`
+# global is the entry point `AirspeedVelocity.jl`'s `benchpkg` expects.
 
-const CONTRACTIONS_PATH = joinpath(@__DIR__, "benchmark_specs", "randomTCs.dat")
+const SUITE = BenchmarkGroup()
 
-include("contractions.jl")
+const DIMS = (4, 16, 64)
 
-# Contraction benchmarks
-# ----------------------
-contraction_suite = SUITE["contractions"] = BenchmarkGroup()
+contract_suite = SUITE["contract"] = BenchmarkGroup()
 
-Ts = (Float64, ComplexF64)
-algs = (TensorAlgebra.Matricize(),)
+# Matrix multiply, one shared index: C[i,k] = A[i,j] B[j,k].
+matmul = contract_suite["matmul"] = BenchmarkGroup()
+for d in DIMS
+    matmul[d] = @benchmarkable(
+        contract(A, (:i, :j), B, (:j, :k)),
+        setup = (A = randn($d, $d); B = randn($d, $d)),
+    )
+end
+
+# Rank-3 contraction over two shared indices: C[i,l] = A[i,j,k] B[k,j,l].
+rank3 = contract_suite["rank3"] = BenchmarkGroup()
+for d in (4, 16)
+    rank3[d] = @benchmarkable(
+        contract(A, (:i, :j, :k), B, (:k, :j, :l)),
+        setup = (A = randn($d, $d, $d); B = randn($d, $d, $d)),
+    )
+end
 
-for alg in algs
-    alg_suite = contraction_suite[alg] = BenchmarkGroup()
-    for T in Ts
-        alg_suite[T] = BenchmarkGroup()
+# Full contraction to a scalar: c = A[i,j] B[i,j].
+scalar = contract_suite["scalar"] = BenchmarkGroup()
+for d in DIMS
+    scalar[d] = @benchmarkable(
+        contract(A, (:i, :j), B, (:i, :j)),
+        setup = (A = randn($d, $d); B = randn($d, $d)),
+    )
+end
 
-        for (i, line) in enumerate(eachline(CONTRACTIONS_PATH))
-            alg_suite[T][i] = generate_contract_benchmark(line; T, alg)
-        end
-    end
+# In-place matrix multiply into a preallocated destination.
+contract!_suite = SUITE["contract!"] = BenchmarkGroup()
+matmul! = contract!_suite["matmul"] = BenchmarkGroup()
+for d in DIMS
+    matmul![d] = @benchmarkable(
+        contract!(C, (:i, :k), A, (:i, :j), B, (:j, :k)),
+        setup = (A = randn($d, $d); B = randn($d, $d); C = zeros($d, $d)),
+    )
 end

test/Project.toml

@@ -1,6 +1,7 @@
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
+BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
 EllipsisNotation = "da5c29d0-fa7d-589e-88eb-ea29b0a81949"
 ITensorPkgSkeleton = "3d388ab1-018a-49f4-ae50-18094d5f71ea"
@@ -23,6 +24,7 @@ path = ".."
 [compat]
 Adapt = "4"
 Aqua = "0.8.9"
+BenchmarkTools = "1"
 BlockArrays = "1.6.1"
 EllipsisNotation = "1.8"
 ITensorPkgSkeleton = "0.3.42"

test/test_benchmarks.jl

@@ -0,0 +1,11 @@
+using BenchmarkTools: warmup
+using Test: @test, @testset
+
+# Run the benchmark suite once, so it cannot fall out of sync with the API without
+# CI noticing. This is a smoke test, not a regression check (for that, compare
+# revisions with `benchpkg`, see `benchmark/README.md`).
+include(joinpath(@__DIR__, "..", "benchmark", "benchmarks.jl"))
+
+@testset "benchmarks (smoke run)" begin
+    @test (warmup(SUITE); true)
+end