Skip to content

Commit 3392c90

Browse files
lkdvoslkdvos
authored
Performance improvements (#43)
* do not make sparse arrays dense * properly annotate inbounds * bump version * avoid dualizing spaces to compute `size`
1 parent 7d3fd12 commit 3392c90

5 files changed

Lines changed: 30 additions & 54 deletions

File tree

‎Project.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
name = "BlockTensorKit"
22
uuid = "5f87ffc2-9cf1-4a46-8172-465d160bd8cd"
3-
version = "0.3.6"
3+
version = "0.3.7"
44
authors = ["Lukas Devos <ldevos98@gmail.com> and contributors"]
55

66
[deps]

‎src/auxiliary/sparsetensorarray.jl‎

Lines changed: 16 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -35,33 +35,44 @@ Base.keys(A::SparseTensorArray) = keys(A.data)
3535
Base.values(A::SparseTensorArray) = values(A.data)
3636

3737
TensorKit.space(A::SparseTensorArray) = A.space
38+
TensorKit.codomain(A::SparseTensorArray) = codomain(space(A))
39+
TensorKit.domain(A::SparseTensorArray) = domain(space(A))
40+
41+
TensorKit.numout(A::SparseTensorArray) = numout(eltype(A))
42+
TensorKit.numout(::Type{T}) where {T <: SparseTensorArray} = numout(eltype(A))
43+
TensorKit.numin(A::SparseTensorArray) = numin(eltype(A))
44+
TensorKit.numin(::Type{T}) where {T <: SparseTensorArray} = numin(eltype(A))
3845

3946
# AbstractArray interface
4047
# -----------------------
41-
Base.size(A::SparseTensorArray) = ntuple(i -> length(space(A)[i]), ndims(A))
48+
Base.size(A::SparseTensorArray) = ntuple(Base.Fix1(size, A), ndims(A))
49+
function Base.size(A::SparseTensorArray, i::Int)
50+
1 ≤ i ≤ ndims(A) || throw(ArgumentError("Invalid number of dimensions"))
51+
return i <= numout(A) ? length(codomain(A)[i]) : length(domain(A)[i - numout(A)])
52+
end
4253

43-
@inline function Base.getindex(
54+
@propagate_inbounds function Base.getindex(
4455
A::SparseTensorArray{S, N₁, N₂, T, N}, I::Vararg{Int, N}
4556
) where {S, N₁, N₂, T, N}
4657
@boundscheck checkbounds(A, I...)
4758
return @inbounds get(A.data, CartesianIndex(I)) do
4859
return fill!(similar(T, eachspace(A)[I...]), zero(scalartype(T)))
4960
end
5061
end
51-
@inline function getindex!(
62+
@propagate_inbounds function getindex!(
5263
A::SparseTensorArray{S, N₁, N₂, T, N}, I::CartesianIndex{N}
5364
) where {S, N₁, N₂, T, N}
5465
@boundscheck checkbounds(A, I)
5566
return @inbounds get!(A.data, I) do
5667
return fill!(similar(T, eachspace(A)[I]), zero(scalartype(T)))
5768
end
5869
end
59-
@inline function getindex!(
70+
@propagate_inbounds function getindex!(
6071
A::SparseTensorArray{S, N₁, N₂, T, N}, I::Vararg{Int, N}
6172
) where {S, N₁, N₂, T, N}
6273
return getindex!(A, CartesianIndex(I))
6374
end
64-
@inline function Base.setindex!(
75+
@propagate_inbounds function Base.setindex!(
6576
A::SparseTensorArray{S, N₁, N₂, T, N}, v, I::Vararg{Int, N}
6677
) where {S, N₁, N₂, T, N}
6778
@boundscheck begin

‎src/tensors/abstractblocktensor/abstractarray.jl‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -80,13 +80,13 @@ function checkspaces(t::AbstractBlockTensorMap)
8080
end
8181

8282
# scalar indexing is dispatched through:
83-
@inline Base.getindex(t::AbstractBlockTensorMap, I::Vararg{Int, N}) where {N} =
83+
@propagate_inbounds Base.getindex(t::AbstractBlockTensorMap, I::Vararg{Int, N}) where {N} =
8484
getindex(parent(t), I...)
85-
@inline Base.getindex(t::AbstractBlockTensorMap, I::CartesianIndex{N}) where {N} =
85+
@propagate_inbounds Base.getindex(t::AbstractBlockTensorMap, I::CartesianIndex{N}) where {N} =
8686
getindex(parent(t), I)
87-
@inline getindex!(t::AbstractBlockTensorMap, I::Vararg{Int, N}) where {N} =
87+
@propagate_inbounds getindex!(t::AbstractBlockTensorMap, I::Vararg{Int, N}) where {N} =
8888
getindex!(parent(t), I...)
89-
@inline getindex!(t::AbstractBlockTensorMap, I::CartesianIndex{N}) where {N} =
89+
@propagate_inbounds getindex!(t::AbstractBlockTensorMap, I::CartesianIndex{N}) where {N} =
9090
getindex!(parent(t), I)
9191

9292
# slicing getindex needs to correctly allocate output blocktensor:
@@ -115,7 +115,7 @@ Base.@propagate_inbounds function Base.getindex(
115115
end
116116

117117
# disambiguate:
118-
Base.@propagate_inbounds function Base.getindex(
118+
@propagate_inbounds function Base.getindex(
119119
t::AbstractBlockTensorMap, indices::Vararg{Strided.SliceIndex}
120120
)
121121
V = space(eachspace(t)[indices...])
@@ -138,7 +138,7 @@ Base.@propagate_inbounds function Base.getindex(
138138
end
139139

140140
# TODO: check if this fallback is fair
141-
@inline Base.setindex!(t::AbstractBlockTensorMap, v::AbstractTensorMap, args...) = (
141+
@propagate_inbounds Base.setindex!(t::AbstractBlockTensorMap, v::AbstractTensorMap, args...) = (
142142
setindex!(parent(t), v, args...); t
143143
)
144144

‎src/tensors/indexmanipulations.jl‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@ function TK.add_transform!(
3333
end
3434
scale!(tdst, β)
3535
p = (p₁..., p₂...)
36-
for (I, v) in nonzero_pairs(tsrc)
36+
@inbounds for (I, v) in nonzero_pairs(tsrc)
3737
I′ = CartesianIndex(TT.getindices(I.I, p))
3838
tdst[I′] = TK.add_transform!(
3939
tdst[I′], v, (p₁, p₂), fusiontreetransform, α, One(), backend...
@@ -55,7 +55,7 @@ function TK.add_transform!(
5555
end
5656
scale!(tdst, β)
5757
p = (p₁..., p₂...)
58-
for (I, v) in nonzero_pairs(tsrc)
58+
@inbounds for (I, v) in nonzero_pairs(tsrc)
5959
I′ = CartesianIndex(TT.getindices(I.I, p))
6060
tdst[I′] = TK.add_transform!(
6161
tdst[I′], v, (p₁, p₂), fusiontreetransform, α, One(), backend...
@@ -121,7 +121,7 @@ for f! in (:add_permute!, :add_transpose!)
121121
end
122122
scale!(tdst, β)
123123
p = (p₁..., p₂...)
124-
for (I, v) in nonzero_pairs(tsrc)
124+
@inbounds for (I, v) in nonzero_pairs(tsrc)
125125
I′ = CartesianIndex(TT.getindices(I.I, p))
126126
tdst[I′] = TK.$f!(tdst[I′], v, (p₁, p₂), α, One(), backend...)
127127
end
@@ -140,7 +140,7 @@ for f! in (:add_permute!, :add_transpose!)
140140
end
141141
scale!(tdst, β)
142142
p = (p₁..., p₂...)
143-
for (I, v) in nonzero_pairs(tsrc)
143+
@inbounds for (I, v) in nonzero_pairs(tsrc)
144144
I′ = CartesianIndex(TT.getindices(I.I, p))
145145
tdst[I′] = TK.$f!(tdst[I′], v, (p₁, p₂), α, One(), backend...)
146146
end
@@ -193,7 +193,7 @@ function TK.add_braid!(
193193
end
194194
scale!(tdst, β)
195195
p = (p₁..., p₂...)
196-
for (I, v) in nonzero_pairs(tsrc)
196+
@inbounds for (I, v) in nonzero_pairs(tsrc)
197197
I′ = CartesianIndex(TT.getindices(I.I, p))
198198
tdst[I′] = TK.add_braid!(tdst[I′], v, (p₁, p₂), levels, α, One(), backend...)
199199
end
@@ -213,7 +213,7 @@ function TK.add_braid!(
213213
end
214214
scale!(tdst, β)
215215
p = (p₁..., p₂...)
216-
for (I, v) in nonzero_pairs(tsrc)
216+
@inbounds for (I, v) in nonzero_pairs(tsrc)
217217
I′ = CartesianIndex(TT.getindices(I.I, p))
218218
tdst[I′] = TK.add_braid!(tdst[I′], v, (p₁, p₂), levels, α, One(), backend...)
219219
end

‎src/tensors/tensoroperations.jl‎

Lines changed: 1 addition & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -36,50 +36,15 @@ for TTA in (:AbstractTensorMap, :AbstractBlockTensorMap), TTB in (:AbstractTenso
3636
end
3737
end
3838

39-
function TO.tensoralloc_contract(
40-
TC,
41-
A::AbstractBlockTensorMap, pA::Index2Tuple, conjA::Bool,
42-
B::AbstractBlockTensorMap, pB::Index2Tuple, conjB::Bool,
43-
pAB::Index2Tuple,
44-
istemp::Val = Val(false),
45-
allocator = TO.DefaultAllocator(),
46-
)
47-
ttype = TO.tensorcontract_type(TC, A, pA, conjA, B, pB, conjB, pAB)
48-
structure = TO.tensorcontract_structure(A, pA, conjA, B, pB, conjB, pAB)
49-
TT = eltype(ttype)
50-
51-
if isabstracttype(TT)
52-
# do not allocate, use undef allocator
53-
E, S, N1, N2 = scalartype(TT), spacetype(TT), numout(structure), numin(structure)
54-
if issparse(A) && issparse(B)
55-
return SparseBlockTensorMap{AbstractTensorMap{E, S, N1, N2}}(
56-
undef, codomain(structure), domain(structure)
57-
)
58-
else
59-
return BlockTensorMap{AbstractTensorMap{E, S, N1, N2}}(
60-
undef, codomain(structure), domain(structure)
61-
)
62-
end
63-
else
64-
return tensoralloc(ttype, structure, istemp, allocator)
65-
end
66-
end
67-
68-
function promote_blocktype(::Type{TT}, ::Type{A₁}, ::Type{A₂}) where {TT, A₁, A₂}
69-
N = similarblocktype(A₁, TT)
70-
@assert N === similarblocktype(A₂, TT) "incompatible block types"
71-
return N
72-
end
73-
7439
function similarblocktype(::Type{A}, ::Type{TT}) where {A, TT}
7540
return Core.Compiler.return_type(similar, Tuple{A, Type{TT}, NTuple{numind(TT), Int}})
7641
end
7742

78-
# By default, make "dense" allocations
7943
function TO.tensoralloc(
8044
::Type{BT}, structure::TensorMapSumSpace, istemp::Val, allocator = TO.DefaultAllocator()
8145
) where {BT <: AbstractBlockTensorMap}
8246
C = BT(undef_blocks, structure)
47+
issparse(C) && return C # don't fill up sparse blocks
8348
blockallocator(V) = TO.tensoralloc(eltype(C), V, istemp, allocator)
8449
map!(blockallocator, parent(C), eachspace(C))
8550
return C

0 commit comments

Comments
 (0)