[BUG] Investigate potential compilation issues related to overloading `Base.cat`

[mtfishman]

@amilsted reported that as of this PR to NDTensors.SparseArraysBase: ITensor/ITensors.jl#1579 which adds Base.cat, loading NDTensors.jl/ITensors.jl along with DifferentialEquations.jl causes huge compile times. That might also be related to https://itensor.discourse.group/t/linearmaps-and-itensors-incompatibility/2216.

For NDTensors.jl/ITensors.jl, that issue should get fixed by ITensor/ITensors.jl#1601, since that PR removes SparseArraysBase.jl from NDTensors.jl in favor of the version in this repository. We need to check if this package has the same problem, and if so rewrite it to avoid that problem.

See also ITensor/BlockSparseArrays.jl#34.

[mtfishman]

I see locally that indeed loading SparseArraysBase.jl along with LinearMaps.jl takes an abnormally long time:

julia> @time using SparseArraysBase, LinearMaps
123.952115 seconds (263.01 M allocations: 10.517 GiB, 0.81% gc time)

(jl_AlNzTo) pkg> st
Status `/private/var/folders/qz/q22pzwm144z9fq57mpf1hfp40000gq/T/jl_AlNzTo/Project.toml`
  [7a12625a] LinearMaps v3.11.3
  [0d5efcca] SparseArraysBase v0.2.10

julia> versioninfo()
Julia Version 1.11.3
Commit d63adeda50d (2025-01-21 19:42 UTC)
Build Info:
  Official https://julialang.org/ release
Platform Info:
  OS: macOS (arm64-apple-darwin24.0.0)
  CPU: 10 × Apple M1 Max
  WORD_SIZE: 64
  LLVM: libLLVM-16.0.6 (ORCJIT, apple-m1)
Threads: 1 default, 0 interactive, 1 GC (on 8 virtual cores)

Not quite as long as what was seen in https://itensor.discourse.group/t/linearmaps-and-itensors-incompatibility/2216 when loading ITensors, but still it is way too long.

Loading SparseArraysBase.jl with DifferentialEquations.jl also takes a long time:

julia> @time using SparseArraysBase, DifferentialEquations
Precompiling DifferentialEquations...
  86 dependencies successfully precompiled in 108 seconds. 235 already precompiled.
116.417075 seconds (20.70 M allocations: 1.397 GiB, 0.54% gc time, 2 lock conflicts, 0.44% compilation time: 95% of which was recompilation)

(jl_G9MCxE) pkg> st
Status `/private/var/folders/qz/q22pzwm144z9fq57mpf1hfp40000gq/T/jl_G9MCxE/Project.toml`
  [0c46a032] DifferentialEquations v7.15.0
  [0d5efcca] SparseArraysBase v0.2.10

@amilsted identified the issue in NDTensors.SparseArraysBase as this definition of Base.cat for sparse arrays: https://github.com/ITensor/ITensors.jl/blob/v0.7.12/NDTensors/src/lib/SparseArraysBase/src/abstractsparsearray/cat.jl#L2, we have a similar overload of Base.cat in this repository derived from here: https://github.com/ITensor/DerivableInterfaces.jl/blob/main/src/abstractarrayinterface.jl#L328-L332. Probably it is causing a lot of invalidations, it is a vararg of big unions of abstract types. However, I don't know how else we can make the definition general, maybe we can do an @eval loop over a list of types? Ideally the Base version would be based on a trait like the design of Base.Broadcast or ArrayLayouts.jl so it could be more easily customized, but that isn't the case right now.

@lkdvos if you have time could you look into this?

[amilsted]

Seems relevant: JuliaLinearAlgebra/LinearMaps.jl#216

[amilsted]

If we want these kinds of cat methods, there might be a case for adding some corresponding methods to https://github.com/SciML/CommonWorldInvalidations.jl and then pulling in that package.

[lkdvos]

So I delved into this a bit deeper, and it seems like the main culprit in this case is not actually Base.cat, but rather a definition of Base.map! that is implicitly included in @derive ... AbstractArrayOps.
See here and here.

Before:

❯ julia --project -e '@time using SparseArraysBase, LinearMaps'
101.707268 seconds (263.14 M allocations: 10.525 GiB, 0.33% gc time, 0.02% compilation time)

After commenting out:

❯ julia --project -e '@time using SparseArraysBase, LinearMaps'
  0.328202 seconds (689.15 k allocations: 72.339 MiB, 15.83% gc time, 8.15% compilation time)

Note that the timings for DifferentialEquations you mention are not entirely fair, since you included the precompilation time of DifferentialEquations there. I don't actually see the long load times:

Before:

❯ julia --project -e '@time using SparseArraysBase, DifferentialEquations'
  5.183700 seconds (11.06 M allocations: 783.451 MiB, 8.18% gc time, 4.28% compilation time: 85% of which was recompilation)

After:

❯ julia --project -e '@time using SparseArraysBase, DifferentialEquations'
  4.937523 seconds (10.92 M allocations: 771.167 MiB, 8.47% gc time, 4.39% compilation time: 85% of which was recompilation)

Given this, it should be fairly straightforward to solve this issue, by rethinking the @derive ... AbstractArrayOps implementation slightly.

(Note that this already includes andyferris/Dictionaries.jl#158, which also helps out a bit)

[mtfishman]

@lkdvos thanks for investigating that, interesting that it is related to overloading Base.map!, hopefully that is easy to circumvent.

[lkdvos]

The map! invalidation turned out to be false alarm, and was a weird interplay of JuliaInterpreter (I'm assuming used in SnoopCompile) and not having a lot of packages loaded. Since that method gets re-specialized quite a lot anyways, it should not be a problem since the compiler gives up after 4 attempts.

The cat implementations did however show up if I first load SparseArraysBase and only then check what LinearMaps.jl invalidates, and removing this in favor of the _cat implementations here does seem to solve the issue.

After delving into that for a bit, it does seem to me like the Base approach to cat is generally not the greatest, since it is virtually impossible to hook into without fully replacing it. As @mtfishman mentioned, an approach which has elements of the broadcasted interface could provide more ways to customize concatenation behavior, and would also protect against accidentally bumping into issues like this. Something to keep in mind for the future, and for the implementations in LinearMaps.jl and OneHotArrays.jl...