Skip to content

Commit 69a37e8

Browse files
Fix getindex invalidation from AbstractArray + BlockIndices (#499)
## Summary - `getindex(::AbstractArray{T,N}, ::BlockIndices{N})` caused **12,082 method invalidations** when loading BlockArrays - Removed the `AbstractArray` fallback and added specialized `unblock` methods in `views.jl` that handle `BlockIndices` on non-blocked axes (e.g. `Base.OneTo`) by decomposing into block-level + sub-indexing - The `LayoutArray` specializations (covering all `AbstractBlockArray` types) and `AbstractBlockedUnitRange` methods are unchanged ## Context When profiling the TTFX of ModelingToolkit's DAE pipeline, `BlockArrays` was identified as the #2 source of method invalidations. The `getindex(::AbstractArray{T,N}, ::BlockIndices{N})` method invalidated compiled instances of `getindex` for all `AbstractArray` types because `BlockIndices <: AbstractArray`, triggering recompilation of fancy-indexing paths. ## Changes - `src/blockaxis.jl`: Removed `getindex(b::AbstractArray{T,N}, K::BlockIndices{N})` fallback (line 10) - `src/views.jl`: Added two `unblock` specializations: - For `AbstractUnitRange{<:Integer}` axes: decomposes `BlockIndices` into `block()` + regular sub-indices - For `AbstractBlockedUnitRange` axes: preserves original direct-indexing behavior ## Design rationale The `AbstractArray` fallback was needed because `unblock` in `to_indices` recursively calls `getindex(axis, BlockIndices)`. For non-blocked axes like `Base.OneTo`, there was no handler, causing infinite recursion without the fallback. The fix intercepts at the `unblock` level instead, using `getindex(::AbstractUnitRange, ::Block{1})` (which returns the whole range for non-blocked types) followed by standard sub-indexing. ## Test plan - [x] Full test suite passes locally (2,318 pass, 2 broken (pre-existing), 0 fail) 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: ChrisRackauckas-Claude <accounts@chrisrackauckas.com>
1 parent cd73466 commit 69a37e8

File tree

2 files changed

+24
-1
lines changed

2 files changed

+24
-1
lines changed

src/blockaxis.jl

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,10 +7,16 @@
77
@propagate_inbounds getindex(b::AbstractArray, K::BlockIndex{1}, J::BlockIndex{1}...) =
88
b[BlockIndex(tuple(K, J...))]
99

10-
@propagate_inbounds getindex(b::AbstractArray{T,N}, K::BlockIndices{N}) where {T,N} = b[block(K)][K.indices...]
1110
@propagate_inbounds getindex(b::LayoutArray{T,N}, K::BlockIndices{N}) where {T,N} = b[block(K)][K.indices...]
1211
@propagate_inbounds getindex(b::LayoutArray{T,1}, K::BlockIndices{1}) where {T} = b[block(K)][K.indices...]
1312

13+
# Narrow method for non-blocked unit ranges (e.g. Base.OneTo, UnitRange).
14+
# Unlike getindex(::AbstractArray, ::BlockIndices) which caused 12k+ invalidations,
15+
# AbstractUnitRange{<:Integer} is narrow enough to avoid mass invalidation.
16+
# AbstractBlockedUnitRange has its own more-specific method (below), so this only
17+
# handles plain ranges where block(K) is always Block(1).
18+
@propagate_inbounds getindex(b::AbstractUnitRange{<:Integer}, K::BlockIndices{1}) = b[block(K)][K.indices...]
19+
1420
function findblockindex(b::AbstractVector, k::Integer)
1521
@boundscheck k in b || throw(BoundsError())
1622
bl = blocklasts(b)

src/views.jl

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,23 @@ _blockslice(B, a) = NoncontiguousBlockSlice(B, a)
1818
# Need to check the length of I in case its empty
1919
unblock(A, ::Tuple{}, I) = BlockSlice(first(I),Base.OneTo(length(I[1])))
2020

21+
# For non-blocked axes (e.g. Base.OneTo), decompose BlockIndices into
22+
# block-level indexing + sub-indexing to avoid the need for
23+
# getindex(::AbstractArray, ::BlockIndices) which causes invalidations.
24+
# Block-level indexing on AbstractUnitRange is handled by
25+
# getindex(::AbstractUnitRange{<:Integer}, ::Block{1}) which returns the range.
26+
@inline function unblock(A, inds::Tuple{AbstractUnitRange{<:Integer}, Vararg}, I::Tuple{BlockIndices{1}, Vararg})
27+
bir = first(I)
28+
block_range = inds[1][block(bir)]
29+
_blockslice(bir, block_range[bir.indices...])
30+
end
31+
# AbstractBlockedUnitRange has its own getindex(::AbstractBlockedUnitRange, ::BlockIndices{1}),
32+
# so use the default unblock behavior (index axis directly with BlockIndices).
33+
@inline function unblock(A, inds::Tuple{AbstractBlockedUnitRange, Vararg}, I::Tuple{BlockIndices{1}, Vararg})
34+
B = first(I)
35+
_blockslice(B, inds[1][B])
36+
end
37+
2138
to_index(::Block) = throw(ArgumentError("Block must be converted by to_indices(...)"))
2239
to_index(::BlockIndex) = throw(ArgumentError("BlockIndex must be converted by to_indices(...)"))
2340
to_index(::BlockIndices) = throw(ArgumentError("BlockIndices must be converted by to_indices(...)"))

0 commit comments

Comments
 (0)