Skip to content

perf: invert index maps when slicing and assigning block tensors - #72

Merged
lkdvos merged 3 commits into
mainfrom
indexing
Aug 26, 2026
Merged

perf: invert index maps when slicing and assigning block tensors#72
lkdvos merged 3 commits into
mainfrom
indexing

Conversation

@lkdvos

@lkdvos lkdvos commented Aug 25, 2026

Copy link
Copy Markdown
Member

Closes #67.

1. Slicing getindex is no longer quadratic

getindex scanned the whole source index grid once per stored block, making it
O(nnz × prod(length.(indices))) and materializing the slice region as an
Array{CartesianIndex}. The index maps are now inverted per dimension, so each stored block
finds its destinations in O(1).

Banded sparse block tensor (nnz = n, one block per row), sliced to every other channel —
t[1:2:n, 1:1, 1:1, 1:2:n], minimum of 20 runs:

n nnz before after speedup alloc before alloc after
16 16 0.0045 ms 0.0015 ms 3 kB 2 kB
32 32 0.0211 ms 0.0020 ms 11× 9 kB 2 kB
64 64 0.1452 ms 0.0036 ms 40× 33 kB 4 kB
128 128 1.1344 ms 0.0052 ms 218× 127 kB 8 kB
256 256 8.9455 ms 0.0107 ms 836× 504 kB 15 kB
512 512 71.1836 ms 0.0188 ms 3786× 2.01 MB 28 kB

Dense BlockTensorMap, same slice:

n blocks before after speedup
16 256 0.0435 ms 0.0131 ms
32 1024 0.5348 ms 0.0272 ms 20×
64 4096 7.9382 ms 0.0874 ms 91×

The "before" column reproduces the measurements in #67. The old column grows ~8× per doubling
of n, the new one linearly.

SparseTensorArray had a second copy of the same logic for slicing the parent array; both call
sites now share _copyslice!.

Bugs fixed along the way

Normalizing indices with Base.to_indices before inverting them cures four separate problems
in the old index handling. All four were reproduced on main before the change.

Repeated indices dropped blocks — as uninitialized memory for a dense destination.
findfirst returns only the first match, so the second copy was never written, and a dense
destination comes from BlockTensorMap{TT}(undef, P):

julia> b = t[[1, 1, 3], :, :];
julia> norm.((b[1, 1, 1], b[2, 1, 1], b[3, 1, 1]))
(1.9422665505663932, 2.7242541828772e-310, 1.4410118868645043)   # before
(1.9997422044279123, 1.9997422044279123, 1.7876528228084503)     # after

This is the one intended behaviour change: repeated indices now duplicate the selected tensors,
as they do for AbstractArray. Note ndims(t[:, 1:2, [1, 1]]) == 3 already appears in
docs/src/blocktensors.md, so the documented example was returning garbage.

Single-index slicing of one-dimensional block tensors returned uninitialized blocks.
pairs(A::AbstractVector) uses linear indices, so nonzero_pairs yielded Int keys for a 1-d
dense parent while the scan compared against CartesianIndex{1} — nothing ever matched. Fixed
at the root in nonzero_pairs (first commit), which also makes it consistent with
nonzero_keys.

julia> t1 = rand(V  one(V)); norm(t1), norm(t1[:]), norm(t1[1:2])
(1.428, 2.4847601550714e-310, 2.4848926097841e-310)   # before
(2.132, 2.1318787907021446, 1.791801404015716)        # after

Integer indices other than Int threw. ind isa Int missed the other BitIntegers:

julia> t[UInt(1), 1:2, 1]              # before: BoundsError
julia> t[Int8(1), Int8(1), Int8(1)]    # before: MethodError: no method matching space(::TensorMapSpace{ComplexSpace, 3, 0})

Logical masks threw on the parent-array path. Base.LogicalIndex <: AbstractVector{Int}, so
it reached a findfirst that needs getindex, which LogicalIndex does not define:

julia> parent(st)[[true, false, true], :, :]   # before: CanonicalIndexError

2. Slice assignment no longer scales with the region

The mirror-image problem, on the setindex! side. All three SparseTensorArray copyto!
methods walked the assigned region rather than the stored blocks, so t[inds...] = v and cat
cost O(prod(length.(inds))) no matter how little was stored. Same inverted index maps, reused.

Banded sparse tensors, t[inds...] = src over the interior:

shape region nnz before after speedup
2-d, n = 64 3 844 64 0.150 ms 0.015 ms 10×
2-d, n = 256 64 516 256 2.267 ms 0.036 ms 63×
2-d, n = 512 260 100 512 13.752 ms 0.051 ms 269×
4-d, n = 24 234 256 24 9.459 ms 0.009 ms 1040×
4-d, n = 48 4 477 456 48 140.997 ms 0.013 ms 10 600×
narrow (1:2, :, :, 1:2 into n = 512) 4 512 0.0045 ms 0.0044 ms

copyto!(dst, view(parent(t), inds...)) likewise: 14.582 ms → 0.041 ms at a 260 100-block
region (357×).

The narrow row is the point of the design. A small assignment into a large sparse array is
already cheap, and must not become proportional to the destination's stored blocks. So
_deletemissing! sweeps whichever of the assigned region and the stored entries is smaller,
making the whole thing O(nnz(src) + min(region, nnz(dst))) — never worse than before, in any
shape. copyto!(t::SparseTensorArray, ::SubArray) needs no sweep at all: the destination spans
exactly the viewed region, so everything not copied is dropped.

A third bug: the nonzero_* accessors on SparseTensorArray

nonzero_keys and nonzero_length fell through to the AbstractArray fallbacks
(eachindex(IndexCartesian(), A) and length(A)) and so reported every index rather than the
stored ones. nonzero_pairs was accidentally correct via Base.pairs(::SparseTensorArray) until
the first commit here routed it through pairs(IndexCartesian(), A), which bypasses that
specialization.

Two visible consequences, both now covered by tests:

  • slicing the parent array densified its result — parent(st)[1:3, :, :] came back with every
    block stored;
  • region-to-region copyto! stored explicitly-zero blocks, because
    Rsrc[I] in nonzero_keys(src) was always true.

Fixed with the four obvious specializations, mirroring SparseBlockTensorMap's.

Notes for review

  • The Vararg{Strided.SliceIndex} twins are kept deliberately. They are not redundant
    specializations: they disambiguate against TensorKit's own
    getindex(::AbstractTensorMap, ::Vararg{SliceIndex}) / setindex! (abstracttensor.jl:540
    and :552), which slice the StridedView of the trivial fusion tree. t[:, :, :] resolves
    to the twin, so deleting either one turns it into a MethodError. Only the bodies were
    collapsed; the setindex! dedup is pure code motion.
  • Error behaviour is otherwise unchanged: t[:] and t[[1]] still throw ArgumentError,
    t[:, :] and out-of-bounds indices still throw BoundsError (including under @inbounds,
    since SumSpace indexing bounds-checks unconditionally), and all existing @inferred tests
    still pass.
  • _invert_index deliberately carries no @inbounds/@propagate_inbounds, so an out-of-range
    index can never become an unchecked write.
  • Stepped and reversed ranges fall through to the generic CSR inversion rather than getting a
    divrem specialization — it measured no faster and is the most error-prone arithmetic in the
    design.
  • The copyto! rewrite is verified semantics-preserving. The stored-key sets after slice
    assignment, copyto! from a view, region-to-region copyto!, and cat are byte-identical to
    the previous implementation once the accessors are fixed (compared over a seeded batch of
    randomized cases, dumping sorted key tuples and diffing).
  • One deliberate difference: with a repeated index in an assignment target
    (t[[1, 1], :] = v), the old code's result depended on iteration order, since it could write
    and then delete the same parent entry. The new code deletes only entries the source has no
    copy of, then writes — deterministic.
  • Region-to-region copyto! still does not clear destination entries where the source is
    structurally zero. That is pre-existing behaviour and cat depends on it, so it is left
    alone.

test/abstracttensor/indexing.jl grew from 62 to ~240 lines. The dense and sparse copies are now
one loop, and it covers repeated indices, logical masks, BitVector, reversed and stepped ranges,
empty slices, non-Int integers, 1-d block tensors, the parent-array path (including that it does
not densify), the nonzero_* accessors, randomized slice-assignment against an explicit
key-set/identity reference, copyto! from a view, region-to-region copyto! with a stepped
region, and cat.

Two complexity guards, both checked to fail against the old code:

  • slicing: allocation-based (@allocated < 2 MB on a 260 100-block region; the old Rsrc alone
    is ~8 MB), so it cannot be timing-flaky;
  • assignment: allocation gives no signal there (the old cost is Dict traffic, not allocation), so
    it is a wall-clock bound with a deliberately huge margin — a 4.5 M-block region where the old
    code takes 141 ms and the new one 0.013 ms, asserted under 20 ms.

🤖 Generated with Claude Code

@lkdvos lkdvos changed the title perf: invert index maps when slicing block tensors perf: invert index maps when slicing and assigning block tensors Aug 25, 2026
@codecov

codecov Bot commented Aug 25, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 95.12195% with 4 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
src/tensors/abstractblocktensor/abstractarray.jl 83.33% 3 Missing ⚠️
src/auxiliary/sliceindices.jl 96.87% 1 Missing ⚠️
Files with missing lines Coverage Δ
src/BlockTensorKit.jl 100.00% <ø> (ø)
src/auxiliary/sparsetensorarray.jl 84.00% <100.00%> (+55.02%) ⬆️
src/tensors/abstractblocktensor/sparsity.jl 50.00% <100.00%> (ø)
src/auxiliary/sliceindices.jl 96.87% <96.87%> (ø)
src/tensors/abstractblocktensor/abstractarray.jl 69.79% <83.33%> (+26.19%) ⬆️

... and 2 files with indirect coverage changes

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

lkdvos and others added 3 commits August 25, 2026 21:04
`pairs(A::AbstractVector)` uses linear indices, so `nonzero_pairs` yielded `Int`
keys for the parent of a one-dimensional dense `BlockTensorMap`, inconsistent
with `nonzero_keys`. Consumers comparing against `CartesianIndex` then silently
matched nothing, which made `t[:]` and `t[1:2]` return uninitialized blocks.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Slicing `getindex` scanned the whole source index grid once per stored block,
making it `O(nnz * prod(length.(indices)))` and materializing the slice region
as an `Array{CartesianIndex}`. Invert the index maps per dimension instead, so
each stored block finds its destinations in `O(1)`:

    n    nnz   before      after     speedup
    16    16   0.0045 ms   0.0015 ms      3x
    64    64   0.1452 ms   0.0036 ms     40x
    256  256   8.9455 ms   0.0107 ms    836x
    512  512  71.1836 ms   0.0188 ms   3786x

(banded sparse block tensor, sliced to every other channel; the dense
`BlockTensorMap` variant goes from 7.94 ms to 0.087 ms at n = 64.)

`SparseTensorArray` had a second copy of the same logic for slicing the parent
array; both now share `_copyslice!`.

Along the way this fixes four bugs in the old index handling:

- Repeated indices dropped blocks instead of duplicating them, and for a dense
  destination the dropped slot was returned as uninitialized memory. They now
  duplicate, as they do for `AbstractArray`.
- Single-index slicing of one-dimensional block tensors (`t[:]`, `t[1:2]`)
  returned uninitialized blocks.
- `Integer` indices other than `Int` threw: `t[UInt(1), 1:2, 1]` a
  `BoundsError`, `t[Int8(1), Int8(1), Int8(1)]` a `MethodError`.
- Logical masks threw `CanonicalIndexError` on the parent array path,
  `parent(t)[[true, false, true], :, :]`.

The `Vararg{Strided.SliceIndex}` methods are kept: they disambiguate against
TensorKit's own `getindex`/`setindex!` for `AbstractTensorMap`, so removing
them turns `t[:, :, :]` into a `MethodError`.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The three `SparseTensorArray` `copyto!` methods all walked the assigned region
rather than the stored blocks, so `t[inds...] = v` and `cat` cost
`O(prod(length.(inds)))` regardless of how little was stored — the mirror image
of the `getindex` problem in #67. Reuse the inverted index maps instead:

    shape        region      nnz    before      after     speedup
    2d n=64        3_844       64    0.150 ms   0.015 ms      10x
    2d n=256      64_516      256    2.267 ms   0.036 ms      63x
    2d n=512     260_100      512   13.752 ms   0.051 ms     269x
    4d n=24      234_256       24    9.459 ms   0.009 ms    1040x
    4d n=48    4_477_456       48  140.997 ms   0.013 ms   10600x
    narrow             4      512    0.005 ms   0.004 ms       1x

`copyto!(dst, view(parent(t), inds...))` improves likewise, 14.6 ms to 0.041 ms
for a 260_100-block region.

The narrow row matters: a small assignment into a large sparse array must not
become proportional to the destination's stored blocks. `_deletemissing!`
therefore sweeps whichever of the two is smaller, making it
`O(nnz(src) + min(region, nnz(dst)))` — never worse than before.

`copyto!(t::SparseTensorArray, ::SubArray)` needs no sweep at all: the
destination spans exactly the viewed region, so everything not copied is
dropped.

Also fix the `nonzero_*` accessors for `SparseTensorArray`, which fell through
to the `AbstractArray` fallbacks and reported every index rather than the stored
ones. `nonzero_keys` and `nonzero_length` were always wrong this way;
`nonzero_pairs` became wrong when it started going through
`pairs(IndexCartesian(), A)`, which bypasses `Base.pairs(::SparseTensorArray)`.
The visible effects were that slicing the parent array densified its result and
that `copyto!` over regions stored explicitly-zero blocks.

Verified semantics-preserving: the stored-key sets after slice assignment,
`copyto!` from a view, region-to-region `copyto!` and `cat` are identical to the
previous implementation once the accessors are fixed.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
@lkdvos
lkdvos merged commit 8aa4d68 into main Aug 26, 2026
28 checks passed
@lkdvos
lkdvos deleted the indexing branch August 26, 2026 08:18
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Slicing getindex is quadratic: replace the per-block findfirst scan with inverse index maps

1 participant