# Lazy vcat of matrices just before multiplication

**URL:** <https://discourse.julialang.org/t/lazy-vcat-of-matrices-just-before-multiplication/112387>\
**Category:** Performance\
**Created:** [April 1, 2024, 8:03pm UTC](https://discourse.julialang.org/t/lazy-vcat-of-matrices-just-before-multiplication/112387 "2024-04-01T20:03:01Z")\
**Posts on this page:** 4\
**Page:** 1

<div class="post-metadata">

**Author:** ![Tomas\_Pevny](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/tomas_pevny/32/25466_2.png) [@Tomas\_Pevny](https://discourse.julialang.org/u/Tomas_Pevny)\
**Post date:** [April 1, 2024, 8:03pm UTC](https://discourse.julialang.org/t/lazy-vcat-of-matrices-just-before-multiplication/112387/1 "2024-04-01T20:03:01Z")

</div>

Dear All,

I am implementing neural network for hyper-multigraph (this is not a joke, the graph has multiple types of edges of arities higher than 3). This is “fairly” straightforward extension of the Mill.jl framework. Now the graph I process has about 10^5 vertices and because I have multiple types of edges, I concatenate output of message pass over each type of edge and then feed it to the dense layer. Looking at the profiler, I have noticed that vertical concatenation of matrices takes some time, therefore I have thought that I would do it lazily.

So to be concrete, I want to compute something like

```julia
x = randn(16,105625)
y = randn(16,105625)
z = randn(16,105625)
w = randn(16,3*16)

w * vcat(x,y,z)

```

without explicitly constructing the output of `vcat`.

I have found `LazyArrays`, which do have this functionality, but propagate the lazy computation further and do not lead to speedup, when I use them as `Array(w * LazyArrays.Vcat(x,y,z))` (concrete measurements are below).

Finally, I have written something myself as

```julia
struct LazyVCat{N,T<:AbstractMatrix}
	xs::NTuple{N,T}
	function LazyVCat(xs::T) where {T}
		n = size(first(xs),2)
		all(n == size(x,2) for x in xs) || error("all matrices should be of the same time")
	end
end

import Base.*
using LinearAlgebra
function *(A::Matrix, B::LazyVCat)
	T = eltype(A)
	o = similar(A, size(A,1), size(first(B.xs),2))
	offset = 0
	for x in B.xs
		LinearAlgebra.mul!(o, view(w, :,offset+1:offset+size(x,1)), x,one(T),one(T))
		offset += size(x,1)
	end
	o
end

```

which gives me the speedup I was looking for.

The output of benchmarktools is below

```julia
using LinearAlgebra, BenchmarkTools, LazyArrays
julia> @benchmark Array(w * LazyArrays.Vcat(x,y,z))
BenchmarkTools.Trial: 195 samples with 1 evaluation.
 Range (min … max): 18.112 ms … 109.953 ms ┊ GC (min … max): 0.00% … 82.59%
 Time (median): 25.102 ms ┊ GC (median): 18.01%
 Time (mean ± σ): 25.652 ms ± 8.802 ms ┊ GC (mean ± σ): 21.05% ± 12.97%

                █ ▄ ▄ ▁
  ▇█▅▅▅▆▃▇▄▆▄▁▄▇█▆█▃▆▅▆██▇▇█▆▆█▆█▄▃▄▄▄▆▆▄▄▅▃▃▁▁▃▃▁▁▁▁▁▁▁▁▁▁▁▁▃ ▃
  18.1 ms Histogram: frequency by time 37.7 ms <

 Memory estimate: 51.58 MiB, allocs estimate: 161.

julia> @benchmark w * vcat(x,y,z)
BenchmarkTools.Trial: 198 samples with 1 evaluation.
 Range (min … max): 17.451 ms … 131.668 ms ┊ GC (min … max): 0.00% … 83.29%
 Time (median): 23.815 ms ┊ GC (median): 20.01%
 Time (mean ± σ): 25.258 ms ± 11.110 ms ┊ GC (mean ± σ): 22.70% ± 15.27%

  ▄▇▂ █▂ ▆▃▂
  ███▄▄▅▄▆███▅▅▅███▅▅▅▃▃▃▁▁▁▁▁▁▁▃▃▁▁▁▁▁▃▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▃ ▃
  17.5 ms Histogram: frequency by time 57 ms <

 Memory estimate: 51.57 MiB, allocs estimate: 4.

julia> @benchmark w * L
BenchmarkTools.Trial: 495 samples with 1 evaluation.
 Range (min … max): 8.240 ms … 20.951 ms ┊ GC (min … max): 0.00% … 38.86%
 Time (median): 8.735 ms ┊ GC (median): 0.00%
 Time (mean ± σ): 10.085 ms ± 2.322 ms ┊ GC (mean ± σ): 11.81% ± 15.18%

   ▆█▇▄▁ ▁ ▃▄▃▂▁
  █████████▇█▇▄▆▅▄▄▁▁▄▄▁▅▄▁▁▁▁▄▁▁▅▁▁▅▅█████▅▇▄▅▅▄▄▁▄▁▁▁▁▄▄▁▄▄ ▇
  8.24 ms Histogram: log(frequency) by time 16.2 ms <

 Memory estimate: 12.89 MiB, allocs estimate: 23.

```

Now of course I can finish the implementation by writing a `rrule` for ChainRules and do all the stuff around. But I do not want to repeat someone else’s work, therefore I want to ask, if someone know about a package, which has implemented this, is fast, and works with AD.

Thanks for answers in advance.

---

<div class="post-metadata">

**Author:** ![Dan](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/dan/32/42581_2.png) [@Dan](https://discourse.julialang.org/u/Dan)\
**Post date:** [April 1, 2024, 10:45pm UTC](https://discourse.julialang.org/t/lazy-vcat-of-matrices-just-before-multiplication/112387/2 "2024-04-01T22:45:51Z")

</div>

You might be looking for:

> **[GitHub - mcabbott/LazyStack.jl: 🥞](https://github.com/mcabbott/LazyStack.jl)**
>
> 🥞. Contribute to mcabbott/LazyStack.jl development by creating an account on GitHub.

Perhaps mcabbott will confirm, but it does mention ChainRules compatibility.

---

<div class="post-metadata">

**Author:** ![mcabbott](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mcabbott/32/6603_2.png) [@mcabbott](https://discourse.julialang.org/u/mcabbott)\
**Post date:** [April 2, 2024, 3:54am UTC](https://discourse.julialang.org/t/lazy-vcat-of-matrices-just-before-multiplication/112387/3 "2024-04-02T03:54:36Z")

</div>

This `LazyVCat` looks sensible to me, modulo typos, I don’t know of a packaged version. Although if you can supply a pre-allocated space for `vcat(x,y,z)`, one call to `*` may be quicker?

```julia
julia> @btime $w * vcat($x,$y,$z);
  min 14.890 ms, mean 16.542 ms (6 allocations, 51.57 MiB)

julia> @btime $w * $L;
  min 4.006 ms, mean 4.769 ms (21 allocations, 12.89 MiB)

julia> @btime $w * $(vcat(x,y,z));
  min 1.509 ms, mean 2.917 ms (3 allocations, 12.89 MiB)

```

---

<div class="post-metadata">

**Author:** ![Tomas\_Pevny](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/tomas_pevny/32/25466_2.png) [@Tomas\_Pevny](https://discourse.julialang.org/u/Tomas_Pevny)\
**Post date:** [April 2, 2024, 5:16am UTC](https://discourse.julialang.org/t/lazy-vcat-of-matrices-just-before-multiplication/112387/4 "2024-04-02T05:16:11Z")

</div>

Thanks a lot for the answer. I cannot preallocate, as the number of vertices differ between states.
