Skip to content

Commit 9cc0573

Browse files
JuthoJutho
authored
make all caches behave the same (#244)
1 parent cabbf44 commit 9cc0573

9 files changed

Lines changed: 231 additions & 152 deletions

File tree

‎src/TensorKit.jl‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -88,6 +88,9 @@ export scalar, add!, contract!
8888
# truncation schemes
8989
export notrunc, truncerr, truncdim, truncspace, truncbelow
9090

91+
# cache management
92+
export empty_globalcaches!
93+
9194
# Imports
9295
#---------
9396
using TupleTools
@@ -134,6 +137,7 @@ using PackageExtensionCompat
134137
# Auxiliary files
135138
#-----------------
136139
include("auxiliary/auxiliary.jl")
140+
include("auxiliary/caches.jl")
137141
include("auxiliary/dicts.jl")
138142
include("auxiliary/iterators.jl")
139143
include("auxiliary/linalg.jl")

‎src/auxiliary/caches.jl‎

Lines changed: 155 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,155 @@
1+
const GLOBAL_CACHES = Pair{Symbol,Any}[]
2+
function empty_globalcaches!()
3+
foreach(empty! ∘ last, GLOBAL_CACHES)
4+
return nothing
5+
end
6+
7+
function global_cache_info(io::IO=stdout)
8+
for (name, cache) in GLOBAL_CACHES
9+
println(io, name, ":\t", LRUCache.cache_info(cache))
10+
end
11+
end
12+
13+
abstract type CacheStyle end
14+
struct NoCache <: CacheStyle end
15+
struct TaskLocalCache{D<:AbstractDict} <: CacheStyle end
16+
struct GlobalLRUCache <: CacheStyle end
17+
18+
const DEFAULT_GLOBALCACHE_SIZE = Ref(10^4)
19+
20+
function CacheStyle(args...)
21+
return GlobalLRUCache()
22+
end
23+
24+
macro cached(ex)
25+
Meta.isexpr(ex, :function) ||
26+
error("cached macro can only be used on function definitions")
27+
fcall = ex.args[1]
28+
if Meta.isexpr(fcall, :where)
29+
hasparams = true
30+
params = fcall.args[2:end]
31+
fcall = fcall.args[1]
32+
else
33+
hasparams = false
34+
end
35+
if Meta.isexpr(fcall, :(::))
36+
typed = true
37+
typeex = fcall.args[2]
38+
fcall = fcall.args[1]
39+
else
40+
typed = false
41+
end
42+
Meta.isexpr(fcall, :call) ||
43+
error("cached macro can only be used on function definitions")
44+
fname = fcall.args[1]
45+
fargs = fcall.args[2:end]
46+
fargnames = map(fargs) do arg
47+
if Meta.isexpr(arg, :(::))
48+
return arg.args[1]
49+
else
50+
return arg
51+
end
52+
end
53+
_fbody = ex.args[2]
54+
55+
# actual implenetation, with underscore name
56+
_fname = Symbol(:_, fname)
57+
_fcall = Expr(:call, _fname, fargs...)
58+
if hasparams
59+
_fcall = Expr(:where, _fcall, params...)
60+
end
61+
_fex = Expr(:function, _fcall, _fbody)
62+
63+
# implementation that chooses the cache style
64+
newfcall = fcall
65+
if hasparams
66+
newfcall = Expr(:where, newfcall, params...)
67+
end
68+
cachestylevar = gensym(:cachestyle)
69+
cachestyleex = Expr(:(=), cachestylevar,
70+
Expr(:call, :CacheStyle, fname, fargnames...))
71+
newfbody = Expr(:block,
72+
cachestyleex,
73+
Expr(:call, fname, fargnames..., cachestylevar))
74+
newfex = Expr(:function, newfcall, newfbody)
75+
76+
# nocache implementation
77+
fnocachecall = Expr(:call, fname, fargs..., :(::NoCache))
78+
if hasparams
79+
fnocachecall = Expr(:where, fnocachecall, params...)
80+
end
81+
fnocachebody = Expr(:call, _fname, fargnames...)
82+
if typed
83+
T = gensym(:T)
84+
fnocachebody = Expr(:block, Expr(:(=), T, typeex), Expr(:(::), fnocachebody, T))
85+
end
86+
fnocacheex = Expr(:function, fnocachecall, fnocachebody)
87+
88+
# tasklocal cache implementation
89+
Dvar = gensym(:D)
90+
flocalcachecall = Expr(:call, fname, fargs..., :(::TaskLocalCache{$Dvar}))
91+
if hasparams
92+
flocalcachecall = Expr(:where, flocalcachecall, params..., Dvar)
93+
else
94+
flocalcachecall = Expr(:where, flocalcachecall, Dvar)
95+
end
96+
localcachename = Symbol(:_tasklocal_, fname, :_cache)
97+
cachevar = gensym(:cache)
98+
getlocalcacheex = :($cachevar::$Dvar = get!(task_local_storage(), $localcachename) do
99+
return $Dvar()
100+
end)
101+
valvar = gensym(:val)
102+
if length(fargnames) == 1
103+
key = fargnames[1]
104+
else
105+
key = Expr(:tuple, fargnames...)
106+
end
107+
getvalex = :(get!($cachevar, $key) do
108+
return $_fname($(fargnames...))
109+
end)
110+
if typed
111+
T = gensym(:T)
112+
flocalcachebody = Expr(:block,
113+
getlocalcacheex,
114+
Expr(:(=), T, typeex),
115+
Expr(:(=), Expr(:(::), valvar, T), getvalex),
116+
Expr(:return, valvar))
117+
else
118+
flocalcachebody = Expr(:block,
119+
getlocalcacheex,
120+
Expr(:(=), valvar, getvalex),
121+
Expr(:return, valvar))
122+
end
123+
flocalcacheex = Expr(:function, flocalcachecall, flocalcachebody)
124+
125+
# # global cache implementation
126+
fglobalcachecall = Expr(:call, fname, fargs..., :(::GlobalLRUCache))
127+
if hasparams
128+
fglobalcachecall = Expr(:where, fglobalcachecall, params...)
129+
end
130+
globalcachename = Symbol(:GLOBAL_, uppercase(string(fname)), :_CACHE)
131+
getglobalcachex = Expr(:(=), cachevar, globalcachename)
132+
if typed
133+
T = gensym(:T)
134+
fglobalcachebody = Expr(:block,
135+
getglobalcachex,
136+
Expr(:(=), T, typeex),
137+
Expr(:(=), Expr(:(::), valvar, T), getvalex),
138+
Expr(:return, valvar))
139+
else
140+
fglobalcachebody = Expr(:block,
141+
getglobalcachex,
142+
Expr(:(=), valvar, getvalex),
143+
Expr(:return, valvar))
144+
end
145+
fglobalcacheex = Expr(:function, fglobalcachecall, fglobalcachebody)
146+
fglobalcachedef = Expr(:const,
147+
Expr(:(=), globalcachename,
148+
:(LRU{Any,Any}(; maxsize=DEFAULT_GLOBALCACHE_SIZE[]))))
149+
fglobalcacheregister = Expr(:call, :push!, :GLOBAL_CACHES,
150+
:($(QuoteNode(globalcachename)) => $globalcachename))
151+
152+
# # total expression
153+
return esc(Expr(:block, _fex, newfex, fnocacheex, flocalcacheex,
154+
fglobalcachedef, fglobalcacheregister, fglobalcacheex))
155+
end

‎src/fusiontrees/manipulations.jl‎

Lines changed: 37 additions & 60 deletions
Original file line numberDiff line numberDiff line change
@@ -526,10 +526,6 @@ function _recursive_repartition(f₁::FusionTree{I,N₁},
526526
end
527527
end
528528

529-
# transpose double fusion tree
530-
const transposecache = LRU{Any,Any}(; maxsize=10^5)
531-
const usetransposecache = Ref{Bool}(true)
532-
533529
"""
534530
transpose(f₁::FusionTree{I}, f₂::FusionTree{I},
535531
p1::NTuple{N₁, Int}, p2::NTuple{N₂, Int}) where {I, N₁, N₂}
@@ -548,28 +544,24 @@ function Base.transpose(f₁::FusionTree{I}, f₂::FusionTree{I},
548544
@assert length(f₁) + length(f₂) == N
549545
p = linearizepermutation(p1, p2, length(f₁), length(f₂))
550546
@assert iscyclicpermutation(p)
551-
if usetransposecache[]
552-
T = sectorscalartype(I)
553-
F₁ = fusiontreetype(I, N₁)
554-
F₂ = fusiontreetype(I, N₂)
555-
D = fusiontreedict(I){Tuple{F₁,F₂},T}
556-
return _get_transpose(D, (f₁, f₂, p1, p2))
557-
else
558-
return _transpose((f₁, f₂, p1, p2))
559-
end
547+
return fstranspose((f₁, f₂, p1, p2))
560548
end
561549

562-
@noinline function _get_transpose(::Type{D}, @nospecialize(key)) where {D}
563-
d::D = get!(transposecache, key) do
564-
return _transpose(key)
565-
end
566-
return d
567-
end
550+
const FSTransposeKey{I<:Sector,N₁,N₂} = Tuple{<:FusionTree{I},<:FusionTree{I},
551+
IndexTuple{N₁},IndexTuple{N₂}}
568552

569-
const TransposeKey{I<:Sector,N₁,N₂} = Tuple{<:FusionTree{I},<:FusionTree{I},
570-
IndexTuple{N₁},IndexTuple{N₂}}
553+
function _fsdicttype(I, N₁, N₂)
554+
F₁ = fusiontreetype(I, N₁)
555+
F₂ = fusiontreetype(I, N₂)
556+
T = sectorscalartype(I)
557+
return fusiontreedict(I){Tuple{F₁,F₂},T}
558+
end
571559

572-
function _transpose((f₁, f₂, p1, p2)::TransposeKey{I,N₁,N₂}) where {I<:Sector,N₁,N₂}
560+
@cached function fstranspose(key::FSTransposeKey{I,N₁,N₂})::_fsdicttype(I, N₁,
561+
N₂) where {I<:Sector,
562+
N₁,
563+
N₂}
564+
f₁, f₂, p1, p2 = key
573565
N = N₁ + N₂
574566
p = linearizepermutation(p1, p2, length(f₁), length(f₂))
575567
newtrees = repartition(f₁, f₂, N₁)
@@ -611,6 +603,14 @@ function _transpose((f₁, f₂, p1, p2)::TransposeKey{I,N₁,N₂}) where {I<:S
611603
return newtrees
612604
end
613605

606+
function CacheStyle(::typeof(fstranspose), k::FSTransposeKey{I}) where {I<:Sector}
607+
if FusionStyle(I) isa UniqueFusion
608+
return NoCache()
609+
else
610+
return GlobalLRUCache()
611+
end
612+
end
613+
614614
# COMPOSITE DUALITY MANIPULATIONS PART 2: Planar traces
615615
#-------------------------------------------------------------------
616616
# -> composite manipulations that depend on the duality (rigidity) and pivotal structure
@@ -1015,10 +1015,6 @@ function permute(f::FusionTree{I,N}, p::NTuple{N,Int}) where {I<:Sector,N}
10151015
end
10161016

10171017
# braid double fusion tree
1018-
const braidcache = LRU{Any,Any}(; maxsize=10^5)
1019-
const usebraidcache_abelian = Ref{Bool}(false)
1020-
const usebraidcache_nonabelian = Ref{Bool}(true)
1021-
10221018
"""
10231019
braid(f₁::FusionTree{I}, f₂::FusionTree{I},
10241020
levels1::IndexTuple, levels2::IndexTuple,
@@ -1043,42 +1039,15 @@ function braid(f₁::FusionTree{I}, f₂::FusionTree{I},
10431039
@assert length(f₁) + length(f₂) == N₁ + N₂
10441040
@assert length(f₁) == length(levels1) && length(f₂) == length(levels2)
10451041
@assert TupleTools.isperm((p1..., p2...))
1046-
if FusionStyle(f₁) isa UniqueFusion &&
1047-
BraidingStyle(f₁) isa SymmetricBraiding
1048-
if usebraidcache_abelian[]
1049-
T = Int # do we hardcode this ?
1050-
F₁ = fusiontreetype(I, N₁)
1051-
F₂ = fusiontreetype(I, N₂)
1052-
D = SingletonDict{Tuple{F₁,F₂},T}
1053-
return _get_braid(D, (f₁, f₂, levels1, levels2, p1, p2))
1054-
else
1055-
return _braid((f₁, f₂, levels1, levels2, p1, p2))
1056-
end
1057-
else
1058-
if usebraidcache_nonabelian[]
1059-
T = sectorscalartype(I)
1060-
F₁ = fusiontreetype(I, N₁)
1061-
F₂ = fusiontreetype(I, N₂)
1062-
D = FusionTreeDict{Tuple{F₁,F₂},T}
1063-
return _get_braid(D, (f₁, f₂, levels1, levels2, p1, p2))
1064-
else
1065-
return _braid((f₁, f₂, levels1, levels2, p1, p2))
1066-
end
1067-
end
1042+
return fsbraid((f₁, f₂, levels1, levels2, p1, p2))
10681043
end
1044+
const FSBraidKey{I<:Sector,N₁,N₂} = Tuple{<:FusionTree{I},<:FusionTree{I},
1045+
IndexTuple,IndexTuple,
1046+
IndexTuple{N₁},IndexTuple{N₂}}
10691047

1070-
@noinline function _get_braid(::Type{D}, @nospecialize(key)) where {D}
1071-
d::D = get!(braidcache, key) do
1072-
return _braid(key)
1073-
end
1074-
return d
1075-
end
1076-
1077-
const BraidKey{I<:Sector,N₁,N₂} = Tuple{<:FusionTree{I},<:FusionTree{I},
1078-
IndexTuple,IndexTuple,
1079-
IndexTuple{N₁},IndexTuple{N₂}}
1080-
1081-
function _braid((f₁, f₂, l1, l2, p1, p2)::BraidKey{I,N₁,N₂}) where {I<:Sector,N₁,N₂}
1048+
@cached function fsbraid(key::FSBraidKey{I,N₁,N₂})::_fsdicttype(I, N₁,
1049+
N₂) where {I<:Sector,N₁,N₂}
1050+
(f₁, f₂, l1, l2, p1, p2) = key
10821051
p = linearizepermutation(p1, p2, length(f₁), length(f₂))
10831052
levels = (l1..., reverse(l2)...)
10841053
local newtrees
@@ -1097,6 +1066,14 @@ function _braid((f₁, f₂, l1, l2, p1, p2)::BraidKey{I,N₁,N₂}) where {I<:S
10971066
return newtrees
10981067
end
10991068

1069+
function CacheStyle(::typeof(fsbraid), k::FSBraidKey{I}) where {I<:Sector}
1070+
if FusionStyle(I) isa UniqueFusion
1071+
return NoCache()
1072+
else
1073+
return GlobalLRUCache()
1074+
end
1075+
end
1076+
11001077
"""
11011078
permute(f₁::FusionTree{I}, f₂::FusionTree{I},
11021079
p1::NTuple{N₁, Int}, p2::NTuple{N₂, Int}) where {I, N₁, N₂}

‎src/spaces/homspace.jl‎

Lines changed: 11 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -233,18 +233,17 @@ struct FusionBlockStructure{I,N,F₁,F₂}
233233
fusiontreeindices::FusionTreeDict{Tuple{F₁,F₂},Int}
234234
end
235235

236-
abstract type CacheStyle end
237-
struct NoCache <: CacheStyle end
238-
struct TaskLocalCache{D<:AbstractDict} <: CacheStyle end
239-
struct GlobalLRUCache <: CacheStyle end
240-
241-
function CacheStyle(I::Type{<:Sector})
242-
return GlobalLRUCache()
236+
function fusionblockstructuretype(W::HomSpace)
237+
N₁ = length(codomain(W))
238+
N₂ = length(domain(W))
239+
N = N₁ + N₂
240+
I = sectortype(W)
241+
F₁ = fusiontreetype(I, N₁)
242+
F₂ = fusiontreetype(I, N₂)
243+
return FusionBlockStructure{I,N,F₁,F₂}
243244
end
244245

245-
fusionblockstructure(W::HomSpace) = fusionblockstructure(W, CacheStyle(sectortype(W)))
246-
247-
function fusionblockstructure(W::HomSpace, ::NoCache)
246+
@cached function fusionblockstructure(W::HomSpace)::fusionblockstructuretype(W)
248247
codom = codomain(W)
249248
dom = domain(W)
250249
N₁ = length(codom)
@@ -317,36 +316,8 @@ function _subblock_strides(subsz, sz, str)
317316
return Strided.StridedViews._computereshapestrides(subsz, sz_simplify...)
318317
end
319318

320-
function fusionblockstructure(W::HomSpace, ::TaskLocalCache{D}) where {D}
321-
cache::D = get!(task_local_storage(), :_local_tensorstructure_cache) do
322-
return D()
323-
end
324-
N₁ = length(codomain(W))
325-
N₂ = length(domain(W))
326-
N = N₁ + N₂
327-
I = sectortype(W)
328-
F₁ = fusiontreetype(I, N₁)
329-
F₂ = fusiontreetype(I, N₂)
330-
structure::FusionBlockStructure{I,N,F₁,F₂} = get!(cache, W) do
331-
return fusionblockstructure(W, NoCache())
332-
end
333-
return structure
334-
end
335-
336-
const GLOBAL_FUSIONBLOCKSTRUCTURE_CACHE = LRU{Any,Any}(; maxsize=10^4)
337-
# 10^4 different tensor spaces should be enough for most purposes
338-
function fusionblockstructure(W::HomSpace, ::GlobalLRUCache)
339-
cache = GLOBAL_FUSIONBLOCKSTRUCTURE_CACHE
340-
N₁ = length(codomain(W))
341-
N₂ = length(domain(W))
342-
N = N₁ + N₂
343-
I = sectortype(W)
344-
F₁ = fusiontreetype(I, N₁)
345-
F₂ = fusiontreetype(I, N₂)
346-
structure::FusionBlockStructure{I,N,F₁,F₂} = get!(cache, W) do
347-
return fusionblockstructure(W, NoCache())
348-
end
349-
return structure
319+
function CacheStyle(::typeof(fusionblockstructure), W::HomSpace)
320+
return GlobalLRUCache()
350321
end
351322

352323
# Diagonal ranges

0 commit comments

Comments
 (0)