Unnecessary Set allocation in sectorequal/issetequal for Trivial sectors

[VinceNeede]

sectorequal/issetequal for Trivial-sector spaces materializes a full hash Set to compare "0 or 1 elements"

Summary

TensorKit.sectorequal, for Trivial-sector spaces, falls through to
Base.issetequal's fully generic fallback. That fallback materializes a
full hash Set (a Dict plus 3 backing Memory arrays) on each side,
just to compare two TensorKit.OneOrNoneIterator{Trivial} objects — each
of which can hold at most one element. This shows up just from
constructing ordinary TensorMaps (no contraction, no TensorOperations
even required — see reproducer below), and in our own downstream package
accounted for roughly 35-40% of total allocations in two structurally
different algorithms built on TensorKit.

Minimal reproducer

Bisected down to the simplest trigger found: no contraction is
needed at all. Repeatedly constructing two structurally-distinct
(transposed-shape) TensorMaps — nothing done with them beyond that —
is enough, with fixed (not even varying) dimensions:

using TensorKit

function build_pair()
    A = randn(ComplexF64, ℂ^3 ← ℂ^4)
    B = randn(ComplexF64, ℂ^4 ← ℂ^3)
    return A, B
end

for _ in 1:100
    build_pair()
end

Allocation-by-type breakdown over these 100 iterations (full script,
including the Profile.Allocs setup, in the collapsible section below):

800 allocs  (62.5 KiB)  <-  Dict{Trivial, Nothing}
800 allocs  (25.0 KiB)  <-  Memory{UInt8}
800 allocs  (25.0 KiB)  <-  Vector{Tuple{Any, Any}}
800 allocs  (12.5 KiB)  <-  Memory{Trivial}
800 allocs  (12.5 KiB)  <-  Memory{Nothing}
600 allocs  (14.1 KiB)  <-  TensorKit.DegeneracyStructure{2}
200 allocs  (9.4 KiB)   <-  TensorKit.Hashed{TensorMapSpace{ComplexSpace, 1, 1}, typeof(TensorKit.sectorhash), typeof(TensorKit.sectorequal)}

(That last line is worth noting on its own — sectorequal appears
directly in the type name of another frequently-allocated object here,
independent confirmation alongside the stack trace below.)

Full script (includes the one-vs-two-structures comparison referenced below)

using TensorKit
using Profile

"Build ONE structure repeatedly -- if caching explains the pattern, this should be cheap (same HomSpace every time)."
function build_one_repeatedly(n::Int)
    for _ in 1:n
        randn(ComplexF64, ℂ^3 ← ℂ^4)
    end
end

"Build TWO distinct structures repeatedly (the confirmed-reproducing case) -- for direct comparison against build_one_repeatedly."
function build_pair()
    A = randn(ComplexF64, ℂ^3 ← ℂ^4)
    B = randn(ComplexF64, ℂ^4 ← ℂ^3)
    return A, B
end

"Construct `n` such pairs in a row -- nothing done with them beyond that."
function build_many_pairs(n::Int)
    for _ in 1:n
        build_pair()
    end
end

"Profile allocations made while running `f`, returning the raw Profile.Allocs results."
function profile_allocations(f)
    Profile.Allocs.clear()
    Profile.Allocs.@profile sample_rate = 1.0 f()
    return Profile.Allocs.fetch()
end

"Summarize allocation results as (count, bytes) per allocated type."
function count_allocations_by_type(results)
    counts = Dict{String,Int}()
    bytes = Dict{String,Int}()
    for alloc in results.allocs
        tname = string(alloc.type)
        counts[tname] = get(counts, tname, 0) + 1
        bytes[tname] = get(bytes, tname, 0) + alloc.size
    end
    return counts, bytes
end

"Print the `top` most frequently allocated types, by count."
function print_top_allocations(counts, bytes; top=10)
    for (tname, cnt) in sort(collect(counts); by=last, rev=true)[1:min(top, end)]
        kib = round(bytes[tname] / 1024, digits=1)
        println(cnt, " allocs  (", kib, " KiB)  <-  ", tname)
    end
end

"Print the deepest stack frames for one `Dict{Trivial,...}` allocation, if any occurred."
function print_dict_trivial_stacktrace(results)
    dict_allocs = filter(a -> occursin("Dict{Trivial", string(a.type)), results.allocs)
    isempty(dict_allocs) && return println("No Dict{Trivial,...} allocations found.")
    println("Stack trace for a Dict{Trivial,...} allocation:")
    for frame in first(dict_allocs).stacktrace[2:min(9, end)]
        println("    ", frame)
    end
end

# --- Run it: compare ONE distinct structure vs TWO, repeated many times ---

build_one_repeatedly(5)  # warm up
build_many_pairs(5)

println("=== ONE structure, repeated 100x ===")
results_one = profile_allocations(() -> build_one_repeatedly(100))
counts_one, bytes_one = count_allocations_by_type(results_one)
print_top_allocations(counts_one, bytes_one)
println()
print_dict_trivial_stacktrace(results_one)

println()
println("=== TWO distinct structures, each repeated 100x ===")
results_two = profile_allocations(() -> build_many_pairs(100))
counts_two, bytes_two = count_allocations_by_type(results_two)
print_top_allocations(counts_two, bytes_two)
println()
print_dict_trivial_stacktrace(results_two)

Where it comes from

Stack trace for one of the Dict{Trivial,Nothing} allocations (captured
via Profile.Allocs on the MWE above):

ijl_gc_small_alloc at gc-stock.c:777
Dict at dict.jl:80 [inlined]
Set at set.jl:45 [inlined]
Set at set.jl:47 [inlined]
_Set at set.jl:59 [inlined]
Set at set.jl:58 [inlined]
issetequal(a::TensorKit.OneOrNoneIterator{Trivial}, b::TensorKit.OneOrNoneIterator{Trivial}) at abstractset.jl:527
sectorequal at vectorspaces.jl:368 [inlined]

Confirmed to also reproduce (with essentially the same allocation
pattern) via a plain @tensor contraction of the two operands, and via
a larger 5-operand @tensoropt contraction with drifting dimensions
(mimicking a real MPS/MPO environment-tensor construction) — the
above is simply the simplest case found, not the only one.

Likely mechanism: competing entries in TensorKit's own structural cache

We don't have full visibility into TensorKit's internals, but a direct
comparison points at a specific trigger. TensorKit's own changelog
(v0.13) describes structural information for a TensorMap/HomSpace
being "cached in a global (or task local) dictionary" so it doesn't need
recomputing on every new TensorMap. Testing this directly:

• Repeatedly constructing one unchanging HomSpace (ℂ^3←ℂ^4, same
  shape every call, 100 calls in a row) — no Dict{Trivial,...}
  allocations at all.
• Repeatedly constructing two distinct HomSpaces (ℂ^3←ℂ^4 and
  ℂ^4←ℂ^3, each individually unchanging, 100 calls in a row) — the
  full Dict{Trivial,...} pattern shown above.

This suggests the expensive path is specifically hit when two or more
distinct structures compete for entries in that cache (e.g. to
disambiguate a hash collision, or to confirm a lookup key is/isn't
already present) — a single, repeatedly-reused structure apparently
never needs this comparison at all. We haven't confirmed this against
the actual cache implementation, so please treat it as a lead rather
than a diagnosis — you'll know in a few minutes of looking at the source
whether this is right.

Suggested fix direction

OneOrNoneIterator{G} can only ever hold 0 or 1 elements, so equality is
expressible without any allocation:

function Base.issetequal(a::TensorKit.OneOrNoneIterator{G}, b::TensorKit.OneOrNoneIterator{G}) where {G}
    ea = isempty(a) ? nothing : only(a)
    eb = isempty(b) ? nothing : only(b)
    return ea == eb
end

We prototyped this (as a local, deliberately-temporary type-piracy patch
in our own package, since it's not something we can fix from downstream)
and confirmed it eliminates the Dict{Trivial,Nothing} allocation
pattern almost entirely (~480x reduction in our own benchmark, from
~50,000 to ~100 allocations for a representative workload). We're not
proposing this exact snippet as the final fix -- you'll know better than
we do where it belongs / whether OneOrNoneIterator itself should
implement this more directly rather than specializing issetequal -- but
wanted to share a working confirmation that it's fixable cheaply.

Impact observed

• Real-world effect on wall-clock time was modest in our testing (a few
  percent) -- the dominant costs in our workload were FLOP-bound
  (BLAS/LAPACK), so this overhead is a smaller fraction of total time at
  larger problem sizes. We'd expect the relative impact to be larger for
  workloads with many small operations, or in memory-constrained/
  multi-threaded settings where GC pressure matters more than raw pause
  time.
• This is triggered by core space-comparison logic (sectorequal,
  used internally whenever Trivial-sector HomSpaces need to be
  compared/cached), not anything specific to our own package or usage
  pattern -- so we'd expect it to be structural to any Trivial-sector
  TensorKit workload with more than one distinct tensor shape in play,
  not particular to us.

Environment

• TensorKit.jl version: v0.17.0
• Julia version: 1.12.6
• OS: Linux (x86_64-linux-gnu)

Happy to help however is useful -- turn this into a PR, test against a
different fix approach, etc.

[Jutho]

Thanks. I'll look into this. There isn't much point into caching the homspace structures for trivial spaces at all.

[Jutho]

What is weird is that TensorKit.Hashed{TensorMapSpace{ComplexSpace, 1, 1}, typeof(TensorKit.sectorhash), typeof(TensorKit.sectorequal)} is a bits type, so it shouldn't cause any allocations at all.

[lkdvos]

I haven't investigated this in detail, but one other possible cause is that the global dict is typed Any, so there needs to be some boxing of variables before we restore type stability further up the call stack which may cause allocations. In that case we might have to just go back to specializing for trivial tensors, which is unfortunately something we've been trying to move away from

[lkdvos]

Will close this for now, I think the main allocations have been dealt with, and if necessary we can consider more Trivial specializations in the future.