Refresh the benchmark suite and implement low hanging performance fixes (AI driven PR)

.gitignore

@@ -4,7 +4,8 @@
 docs/build/
 docs/site/
 *.jld
-benchmark/*.md
+benchmark/*.json
+benchmark/results_*.md
 src/.DS_Store
 Manifest.toml
 Manifest-v*.*.toml

benchmark/Project.toml

@@ -0,0 +1,9 @@
+[deps]
+BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
+BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
+
+[compat]
+BenchmarkTools = "1"
+julia = "1.10"

benchmark/README.md

@@ -0,0 +1,23 @@
+# Benchmarks
+
+Set up the benchmark environment from the repository root:
+
+```julia
+julia --project=benchmark -e 'using Pkg; Pkg.develop(path="."); Pkg.instantiate()'
+```
+
+Run and save the suite under a descriptive name:
+
+```julia
+julia --project=benchmark -L benchmark/runbenchmarks.jl -e 'run_benchmarks("main")'
+```
+
+Generate a standalone report or compare two saved runs:
+
+```julia
+julia --project=benchmark -L benchmark/runbenchmarks.jl -e 'generate_report("main")'
+julia --project=benchmark -L benchmark/runbenchmarks.jl -e 'generate_report("main", "branch")'
+```
+
+Benchmark parameters and results are stored as ignored JSON files under
+`benchmark/`.

benchmark/generate_report.jl

@@ -40,7 +40,11 @@ function printreport(io::IO, results; iscomparisonjob::Bool = false)
     return nothing
 end
 
-idrepr(id) = (str = repr(id); str[searchindex(str, '['):end])
+function idrepr(id)
+    str = repr(id)
+    firstbracket = findfirst(==('['), str)
+    return isnothing(firstbracket) ? str : str[firstbracket:end]
+end
 intpercent(p) = string(ceil(Int, p * 100), "%")
 resultrow(ids, t::BenchmarkTools.Trial) = resultrow(ids, minimum(t))
 

benchmark/runbenchmarks.jl

@@ -1,7 +1,6 @@
 using BlockArrays
 using BenchmarkTools
-using FileIO
-using JLD
+using LinearAlgebra
 
 include("generate_report.jl")
 
@@ -10,6 +9,10 @@ const SUITE = BenchmarkGroup()
 g = addgroup!(SUITE, "indexing")
 # g_block = addgroup!(SUITE, "blockindexing")
 g_size = addgroup!(SUITE, "size")
+g_metadata = addgroup!(SUITE, "metadata")
+g_product = addgroup!(SUITE, "product")
+g_broadcast = addgroup!(SUITE, "broadcast")
+g_reduction = addgroup!(SUITE, "reduction")
 
 for n = (5,)
     for BT in (BlockArray, BlockedArray)
@@ -30,31 +33,68 @@ for n = (5,)
     end
 end
 
+blockkron_vector = Ref(BlockKron(1:10_000, 1:4, 1:3))
+blockkron_matrix = Ref(BlockKron(
+    reshape(1:800_000, 1_000, 800), reshape(1:20, 4, 5), reshape(1:6, 3, 2),
+))
+g_metadata["BlockKron", "axes", "vector"] = @benchmarkable axes($blockkron_vector[])
+g_metadata["BlockKron", "axes", "matrix"] = @benchmarkable axes($blockkron_matrix[])
+
+blocklasts_a = collect(2:2:20_000)
+blocklasts_b = collect(3:3:30_000)
+g_metadata["sortedunion", "vectors"] =
+    @benchmarkable BlockArrays.sortedunion($blocklasts_a, $blocklasts_b)
+
+khatri_block_sizes = fill(2, 10)
+khatri_a = BlockArray(randn(20, 20), khatri_block_sizes, khatri_block_sizes)
+khatri_b = BlockArray(randn(20, 20), khatri_block_sizes, khatri_block_sizes)
+g_product["khatri_rao", "10x10", "2x2 blocks"] =
+    @benchmarkable khatri_rao($khatri_a, $khatri_b)
+
+broadcast_block_sizes = fill(4, 32)
+broadcast_blocked_a = BlockedArray(randn(128, 128), broadcast_block_sizes, broadcast_block_sizes)
+broadcast_blocked_b = BlockedArray(randn(128, 128), broadcast_block_sizes, broadcast_block_sizes)
+broadcast_blocked_dest = similar(broadcast_blocked_a)
+g_broadcast["in-place", "BlockedArray", "matching"] =
+    @benchmarkable $broadcast_blocked_dest .= $broadcast_blocked_a .+ $broadcast_blocked_b
+
+broadcast_block_a = BlockArray(randn(128, 128), broadcast_block_sizes, broadcast_block_sizes)
+broadcast_block_b = BlockArray(randn(128, 128), broadcast_block_sizes, broadcast_block_sizes)
+broadcast_block_dest = similar(broadcast_block_a)
+g_broadcast["in-place", "BlockArray", "matching"] =
+    @benchmarkable $broadcast_block_dest .= $broadcast_block_a .+ $broadcast_block_b
+
+reduction_block_sizes = fill(4, 32)
+reduction_block_array = BlockArray(randn(128, 128), reduction_block_sizes, reduction_block_sizes)
+g_reduction["sum", "BlockArray", "matrix"] = @benchmarkable sum($reduction_block_array)
+g_reduction["sum(abs2)", "BlockArray", "matrix"] =
+    @benchmarkable sum(abs2, $reduction_block_array)
+g_reduction["norm", "BlockArray", "matrix"] = @benchmarkable norm($reduction_block_array)
+
 
 function run_benchmarks(name, tagfilter = @tagged ALL)
-    paramspath = joinpath(dirname(@__FILE__), "params.jld")
+    paramspath = joinpath(@__DIR__, "params.json")
     if !isfile(paramspath)
         println("Tuning benchmarks...")
         tune!(SUITE, verbose=true)
-        JLD.save(paramspath, "SUITE", params(SUITE))
+        BenchmarkTools.save(paramspath, params(SUITE))
     end
-    loadparams!(SUITE, JLD.load(paramspath, "SUITE"), :evals, :samples)
+    loadparams!(SUITE, only(BenchmarkTools.load(paramspath)), :evals, :samples)
     results = run(SUITE[tagfilter], verbose = true, seconds = 2)
-    JLD.save(joinpath(dirname(@__FILE__), name * ".jld"), "results", results)
+    BenchmarkTools.save(joinpath(@__DIR__, name * ".json"), results)
 end
 
 function generate_report(v1, v2)
-    v1_res = load(joinpath(dirname(@__FILE__), v1 * ".jld"), "results")
-    v2_res = load(joinpath(dirname(@__FILE__), v2 * ".jld"), "results")
-    open(joinpath(dirname(@__FILE__), "results_$(v1)_$(v2).md"), "w") do f
+    v1_res = only(BenchmarkTools.load(joinpath(@__DIR__, v1 * ".json")))
+    v2_res = only(BenchmarkTools.load(joinpath(@__DIR__, v2 * ".json")))
+    open(joinpath(@__DIR__, "results_$(v1)_$(v2).md"), "w") do f
         printreport(f, judge(minimum(v1_res), minimum(v2_res)); iscomparisonjob = true)
     end
 end
 
 function generate_report(v1)
-    v1_res = load(joinpath(dirname(@__FILE__), v1 * ".jld"), "results")
-    open(joinpath(dirname(@__FILE__), "results_$(v1).md"), "w") do f
+    v1_res = only(BenchmarkTools.load(joinpath(@__DIR__, v1 * ".json")))
+    open(joinpath(@__DIR__, "results_$(v1).md"), "w") do f
         printreport(f, minimum(v1_res); iscomparisonjob = false)
     end
 end
-

src/blockaxis.jl

@@ -735,6 +735,34 @@ Base.BroadcastStyle(::Type{<:AbstractBlockedUnitRange{<:Any,R}}) where R = _broa
 ###
 
 const OneToCumsum{T<:Integer} = RangeCumsum{T,Base.OneTo{T}}
+
+# Base has no public vector merge that exploits sorted inputs, so merge block
+# boundaries directly instead of hashing and then sorting them with `union`.
+function sortedunion(a::StridedVector{<:Integer}, b::StridedVector{<:Integer})
+    T = promote_type(eltype(a), eltype(b))
+    result = Vector{T}(undef, length(a) + length(b))
+    ia, ib = firstindex(a), firstindex(b)
+    lasta, lastb = lastindex(a), lastindex(b)
+    nresult = 0
+
+    @inbounds while ia <= lasta || ib <= lastb
+        value = if ib > lastb || (ia <= lasta && !isless(b[ib], a[ia]))
+            value = a[ia]
+            ia += 1
+            value
+        else
+            value = b[ib]
+            ib += 1
+            value
+        end
+        if iszero(nresult) || !isequal(result[nresult], value)
+            nresult += 1
+            result[nresult] = value
+        end
+    end
+    return resize!(result, nresult)
+end
+
 sortedunion(a::OneToCumsum, ::OneToCumsum) = a
 function sortedunion(a::RangeCumsum{<:Any,<:AbstractRange}, b::RangeCumsum{<:Any,<:AbstractRange})
     @assert a == b

src/blockbroadcast.jl

@@ -33,7 +33,6 @@ BroadcastStyle(::BlockStyle{M}, ::BlockedStyle{N}) where {M,N} = BlockStyle(Val(
 BroadcastStyle(::BlockedStyle{M}, ::BlockStyle{N}) where {M,N} = BlockStyle(Val(max(M,N)))
 
 
-# sortedunion can assume inputs are already sorted so this could be improved
 maybeinplacesort!(v::StridedVector) = sort!(v)
 maybeinplacesort!(v) = sort(v)
 sortedunion(a,b) = maybeinplacesort!(union(a,b))
@@ -42,6 +41,16 @@ sortedunion(a::AbstractUnitRange, b::AbstractUnitRange) = min(first(a),first(b))
 combine_blockaxes(a, b) = _BlockedUnitRange(sortedunion(blocklasts(a), blocklasts(b)))
 combine_blockaxes(a::BlockedOneTo, b::BlockedOneTo) = BlockedOneTo(sortedunion(blocklasts(a), blocklasts(b)))
 
+function combine_blockaxes(a::BlockedUnitRange{T,Vector{T}},
+                           b::BlockedUnitRange{T,Vector{T}}) where T
+    return blockisequal(a, b) ? a : _BlockedUnitRange(sortedunion(blocklasts(a), blocklasts(b)))
+end
+
+function combine_blockaxes(a::BlockedOneTo{T,CS}, b::BlockedOneTo{T,CS}) where
+                           {T<:Integer,CS<:Union{Vector{T},Base.OneTo{T}}}
+    return blockisequal(a, b) ? a : BlockedOneTo(sortedunion(blocklasts(a), blocklasts(b)))
+end
+
 Base.Broadcast.axistype(a::AbstractBlockedUnitRange, b::AbstractBlockedUnitRange) = length(b) == 1 ? a : combine_blockaxes(a, b)
 Base.Broadcast.axistype(a::AbstractBlockedUnitRange, b) = length(b) == 1 ? a : combine_blockaxes(a, b)
 Base.Broadcast.axistype(a, b::AbstractBlockedUnitRange) = length(b) == 1 ? a : combine_blockaxes(a, b)
@@ -176,31 +185,26 @@ copyto!(dest::AbstractArray,
     _generic_blockbroadcast_copyto!(dest, bc)
 
 # type-stable version of _bview.(args, K)
-__bview(args::Tuple{}, K) = ()
-__bview(args::Tuple, K) = tuple(_bview(args[1],K), __bview(tail(args), K)...)
+__bview(args::Tuple{}, K...) = ()
+__bview(args::Tuple, K...) = tuple(_bview(args[1], K...), __bview(tail(args), K...)...)
 
-function _fast_blockbradcast_copyto!(dest, bc)
-    @inbounds for K in blockaxes(bc)[1]
-        broadcast!(bc.f, view(dest,K), __bview(bc.args, K)...)
+function _fast_blockbroadcast_copyto!(dest, bc)
+    @inbounds for K in Iterators.product(blockaxes(bc)...)
+        broadcast!(bc.f, view(dest, K...), __bview(bc.args, K...)...)
     end
     dest
 end
 
-_hasscalarlikevec() = false
-_hasscalarlikevec(a, b...) = _hasscalarlikevec(b...)
-_hasscalarlikevec(a::AbstractVector, b...) = size(a,1) == 1 || _hasscalarlikevec(b...)
-
-blockisequalorscalar(ax, ::Number) = true
-blockisequalorscalar(ax, a) = blockisequal(ax, Base.axes1(a))
+blockisequalorscalar(ax::Tuple, ::Number) = true
+blockisequalorscalar(ax::Tuple, a) = blockisequal(ax, axes(a))
 
-function copyto!(dest::AbstractVector,
-        bc::Broadcasted{<:AbstractBlockStyle{1}, <:Any, <:Any, Args}) where {Args <: Tuple}
-    _hasscalarlikevec(bc.args...) && return _generic_blockbroadcast_copyto!(dest, bc)
-    ax = axes(dest,1)
+function copyto!(dest::AbstractArray{<:Any,N},
+        bc::Broadcasted{<:AbstractBlockStyle{N},<:Any,<:Any,Args}) where {N,Args<:Tuple}
+    ax = axes(dest)
     for a in bc.args
         blockisequalorscalar(ax, a) || return _generic_blockbroadcast_copyto!(dest, bc)
     end
-    return _fast_blockbradcast_copyto!(dest, bc)
+    return _fast_blockbroadcast_copyto!(dest, bc)
 end
 @inline function Broadcast.instantiate(bc::Broadcasted{Style}) where {Style <:BlockStyle}
     bcf = Broadcast.instantiate(Broadcast.flatten(Broadcasted{Nothing}(bc.f, bc.args, bc.axes)))
@@ -214,6 +218,13 @@ _removeblocks(a::Adjoint) = _removeblocks(parent(a))'
 _removeblocks(a::Transpose) = transpose(_removeblocks(parent(a)))
 _removeblocks(a::SubArray{<:Any,N,<:BlockedArray}) where N = view(_removeblocks(parent(a)), map(_removeblocks, parentindices(a))...)
 _removeblocks(a) = a
+
+function copyto!(dest::BlockedArray{<:Any,N},
+        bc::Broadcasted{BlockedStyle{N},<:Any,<:Any,Args}) where {N,Args<:Tuple}
+    copyto!(parent(dest), _removeblocks(bc))
+    return dest
+end
+
 copy(bc::Broadcasted{BlockedStyle{N}}) where N = BlockedArray(Broadcast.materialize(_removeblocks(bc)), axes(bc))
 
 for op in (:+, :-, :*)

src/blockproduct.jl

@@ -6,23 +6,14 @@ References
 * Khatri, C. G., and Rao, C. Radhakrishna (1968) Solutions to Some Functional Equations and Their Applications to Characterization of Probability Distributions. Sankhya: Indian J. Statistics, Series A 30, 167–180.
 """
 function khatri_rao(A::AbstractBlockMatrix, B::AbstractBlockMatrix)
-    #
-    Ablksize = blocksize(A)
-    Bblksize = blocksize(B)
-
-    @assert Ablksize == Bblksize "A and B must have the same blocksize"
-
-    kblk = []
-    for iblk in blockaxes(A,1)
-        kblk_j = []
-        for _jblk in blockaxes(A,2)
-            Ablk = A[iblk, _jblk]
-            Bblk = B[iblk, _jblk]
-            push!(kblk_j, kron(Ablk, Bblk))
-        end
-        push!(kblk, tuple(kblk_j...))
+    @assert blocksize(A) == blocksize(B) "A and B must have the same blocksize"
+
+    product = Iterators.product(blockaxes(A)...)
+    result_blocks = map(product) do block_index
+        K, J = block_index
+        kron(view(A, K, J), view(B, K, J))
     end
-    mortar(kblk...)
+    return mortar(result_blocks)
 end
 
 function khatri_rao(A::AbstractMatrix, B::AbstractMatrix)
@@ -51,15 +42,11 @@ size(K::BlockKron, j::Int) = prod(size.(K.args, j))
 size(a::BlockKron{<:Any,1}) = (size(a,1),)
 size(a::BlockKron{<:Any,2}) = (size(a,1), size(a,2))
 
-function axes(K::BlockKron{<:Any,1})
-    A,B = K.args
-    (blockedrange(fill(prod(size.(tail(K.args),1)), size(K.args[1],1))),)
-end
-
-function axes(K::BlockKron{<:Any,2})
-    A,B = K.args
-    blockedrange.((fill(prod(size.(tail(K.args),1)), size(K.args[1],1)),
-                   fill(prod(size.(tail(K.args),2)), size(K.args[1],2))))
+function axes(K::BlockKron{<:Any,N}) where N
+    ntuple(Val(N)) do dim
+        blocklength = prod(size.(tail(K.args), dim))
+        blockedrange(Fill(blocklength, size(K.args[1], dim)))
+    end
 end
 
 kron_getindex((A,)::Tuple{AbstractVector}, k::Integer) = A[k]

src/blockreduce.jl

@@ -17,5 +17,25 @@ Base.mapfoldl(f::F, op::OP, B::BlockedArray; kw...) where {F, OP} =
 Base.mapreduce(f::F, op::OP, B::BlockedArray; kw...) where {F, OP} =
     mapreduce(f, op, B.blocks; kw...)
 
+function Base.mapreduce(f::F, op::typeof(Base.add_sum), B::BlockArray;
+                        dims=:, kw...) where F
+    if dims isa Colon && !isempty(B)
+        nonemptyblocks = Iterators.filter(!isempty, B.blocks)
+        return mapreduce(block -> mapreduce(f, op, block), op, nonemptyblocks; kw...)
+    end
+    return invoke(mapreduce, Tuple{Any,Any,AbstractArray}, f, op, B; dims, kw...)
+end
+Base.mapreduce(f::F, op::typeof(Base.add_sum), B::BlockVector;
+               dims=:, kw...) where F =
+    invoke(mapreduce, Tuple{F,typeof(op),BlockArray}, f, op, B; dims, kw...)
+
+_norm_hypot(x, y) = (ismissing(x) || isnan(x)) ? x :
+                    (ismissing(y) || isnan(y)) ? y : hypot(x, y)
+
+function LinearAlgebra.norm2(B::BlockArray)
+    isempty(B.blocks) && return float(norm(zero(eltype(B))))
+    return mapreduce(norm, _norm_hypot, B.blocks)
+end
+
 # support sum, need to return something analogous to Base.OneTo(1) but same type
 Base.reduced_index(::BR) where BR<:AbstractBlockedUnitRange = convert(BR, Base.OneTo(1))