Broadcasting matrix-vector multiplication over specific axis for n-dimensional arrays

General Usage Performance
1

I have a matrix-vector multiplication to fill axis=2 of my 3D array a (where basis is type Matrix):

@inbounds for n ∈ 1:size(a, 1), s ∈ 1:size(a, 3)
    a[n, :, s] = @views basis * b[n, :, s]
end

The question is how can I broadcast this over axis=3, so that the loop only goes through n ∈ 1:size(a, 1).
Is there a performance benefit on doing so?

Thanks!

2

There’s nothing wrong with writing a loop. Done properly, it is usually at least as fast as broadcasting. However, you’ll get better performance from in-place multiplication mul! here (since you already have the output space allocated). Also, as you were perhaps alluding to, there’s no need to loop over s as that part can be handled by using matrix-matrix multiplication instead of matrix-vector.

using LinearAlgebra
for n in axes(a,1)
    @views mul!(a[n,:,:], basis, b[n,:,:])
end

I removed @inbounds because there isn’t a check here to ensure that axes(a,1) == axes(b,1) and also because it is unlikely to significantly change performance.

You can write the above via a somewhat arcane broadcast:

using LinearAlgebra
mul!.(eachslice(a;dims=1), Ref(basis), eachslice(b;dims=1))

Although I wouldn’t say this is easier to read that the loop and I don’t expect it to perform noticeably differently. Note that I have broadcasted mul! by using mul!., sliced the arrays by iterating the first dimension, and used Ref(basis) to not broadcast over that.

2 Likes
3

Yeap, of course the matrix-matrix multiplication would do this already… my bad.
And thanks for the insightful answer and the performance tips using mul! :slight_smile:

4

Tangentially related, but this function may be useful to you. It is for selecting 1D slices in an ND array along a dimension.

1 Like