Skip to content

Commit

Permalink
fix cat invalidations
Browse files Browse the repository at this point in the history
This patch removes the invalidation on cat and thus reduces the OneHotArrays loading time from 4.5s to 0.5s (the normal status)
  • Loading branch information
johnnychen94 authored Sep 26, 2023
1 parent 897948b commit 0956320
Showing 1 changed file with 22 additions and 7 deletions.
29 changes: 22 additions & 7 deletions src/blockmap.jl
Original file line number Diff line number Diff line change
Expand Up @@ -505,13 +505,28 @@ for k in 1:8 # is 8 sufficient?
mapargs = ntuple(n ->:($(Symbol(:A, n))), Val(k-1))
# yields (:LinearMap(A1), :LinearMap(A2), ..., :LinearMap(A(k-1)))

@eval function Base.cat($(Is...), $L, As::MapOrVecOrMat...; dims::Dims{2})
if dims == (1,2)
return BlockDiagonalMap(convert_to_lmaps($(mapargs...))...,
$(Symbol(:A, k)),
convert_to_lmaps(As...)...)
else
throw(ArgumentError("dims keyword in cat of LinearMaps must be (1,2)"))
@static if VERSION >= v"1.8"
# Dispatching on `cat` makes compiler hard to infer types and causes invalidations
# after https://github.com/JuliaLang/julia/pull/45028
# Here we instead dispatch on _cat
@eval function Base._cat(dims, $(Is...), $L, As...)
if dims == (1,2)
return BlockDiagonalMap(convert_to_lmaps($(mapargs...))...,
$(Symbol(:A, k)),
convert_to_lmaps(As...)...)
else
throw(ArgumentError("dims keyword in cat of LinearMaps must be (1,2)"))
end
end
else
@eval function Base.cat($(Is...), $L, As...; dims::Dims{2})
if dims == (1,2)
return BlockDiagonalMap(convert_to_lmaps($(mapargs...))...,

Check warning on line 524 in src/blockmap.jl

View check run for this annotation

Codecov / codecov/patch

src/blockmap.jl#L522-L524

Added lines #L522 - L524 were not covered by tests
$(Symbol(:A, k)),
convert_to_lmaps(As...)...)
else
throw(ArgumentError("dims keyword in cat of LinearMaps must be (1,2)"))

Check warning on line 528 in src/blockmap.jl

View check run for this annotation

Codecov / codecov/patch

src/blockmap.jl#L528

Added line #L528 was not covered by tests
end
end
end
end
Expand Down

0 comments on commit 0956320

Please sign in to comment.