add and benchmark typed_hvcat(SA, ::Val, ...)

to explore the benefits of JuliaLang/julia#36719

Author: simeonschaub

perf/hvcat_val.jl

@@ -0,0 +1,27 @@
+using StaticArray, BenchmarkTools
+
+let
+rows, cols = 4, 4
+_dims = Expr(:tuple, [cols for _ in 1:rows]...)
+
+for (f, wrap_val) in [(:f1, false), (:f2, true)]
+    dims = wrap_val ? :(Val{$_dims}()) : _dims
+    zeros_sa = :(Base.typed_hvcat(SA, $dims, $([0 for _ in 1:rows*cols]...)))
+    xs = [Symbol(:x, i) for i in 1:rows*cols]
+    is = [Symbol(:i, i) for i in 1:rows*cols]
+    is_sa = :(Base.typed_hvcat(SA, $dims, $(is...)))
+    @eval begin
+        function $f($(xs...))
+            r = $zeros_sa
+            for ($(is...),) in Iterators.product($(xs...))
+                r += $is_sa
+            end
+            r
+        end
+    end
+end
+
+xs = [:(1:2) for _ in 1:rows*cols]
+display(@eval @benchmark f1($(xs...)))
+display(@eval @benchmark f2($(xs...)))
+end

src/initializers.jl

@@ -59,3 +59,17 @@ end
 @inline Base.typed_hvcat(sa::Type{SA}, rows::Dims, xs::Number...)            = _SA_typed_hvcat(sa, rows, xs)
 @inline Base.typed_hvcat(sa::Type{SA{T}}, rows::Dims, xs::Number...) where T = _SA_typed_hvcat(sa, rows, xs)
 
+@generated function _SA_typed_hvcat(::Type{sa}, ::Val{rows}, xs) where {sa,rows}
+    M = rows[1]
+    if any(r->r != M, rows)
+        # @pure may not throw... probably. See
+        # https://discourse.julialang.org/t/can-pure-functions-throw-an-error/18459
+        return :(throw(ArgumentError("SA[...] matrix rows of length $_rows are inconsistent")))
+    end
+    msize = Size(M, length(rows))
+    # hvcat lowering is row major ordering, so we must transpose
+    :(Base.@_inline_meta; transpose($(similar_type(sa, msize))(xs)))
+end
+
+@inline Base.typed_hvcat(sa::Type{SA}, rows::Val{_rows}, xs::Number...) where {_rows} = _SA_typed_hvcat(sa, rows, xs)
+@inline Base.typed_hvcat(sa::Type{SA{T}}, rows::Val{_rows}, xs::Number...) where {T,_rows} = _SA_typed_hvcat(sa, rows, xs)