LinearAlgebra: use band index in structured broadcast #54075
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This adds a new
BandIndex
type internal toLinearAlgebra
that parallels aCartesianIndex
, and stores a band index and a linear index of an element along that band. If the index of the band is a compile-time constant (which is often the case withDiagonal
/Bidiagonal
/Tridiagonal
matrices), constant-propagation may eliminate branches in indexing into the matrix, and directly forward the indexing to the corresponding diagonal. This is particularly important in broadcasting for these matrices, which acts band-wise.An example of an improvement in performance with this PR:
This makes the broadcast operation as fast as the sum, where the latter adds the diagonals directly.
I'm not 100% certain why this works as well as it does, as the constant band index may get lost in the
newindex
computation. I suspect branch prediction somehow works around this and preserves the constant.