Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Implement lazy addition #9

Merged
merged 3 commits into from
Nov 26, 2018
Merged

Implement lazy addition #9

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
3e3ed39
Implement lazy addition
tkf Nov 24, 2018
d0c6234
Support B .= Mul(Add(A1, A2), B)
tkf Nov 25, 2018
75836a8
Unexport Add
tkf Nov 26, 2018
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 25 additions & 0 deletions src/lazybroadcasting.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,20 @@ BroadcastStyle(::Type{<:BroadcastArray{<:Any,N}}) where N = LazyArrayStyle{N}()
BroadcastStyle(L::LazyArrayStyle{N}, ::StaticArrayStyle{N}) where N = L
BroadcastStyle(::StaticArrayStyle{N}, L::LazyArrayStyle{N}) where N = L

"""
BroadcastLayout(f, layouts)

is returned by `MemoryLayout(A)` if a matrix `A` is a `BroadcastArray`.
`f` is a function that broadcast operation is applied and `layouts` is
a tuple of `MemoryLayout` of the broadcasted arguments.
"""
struct BroadcastLayout{F, LAY} <: MemoryLayout
f::F
layouts::LAY
end

MemoryLayout(A::BroadcastArray) = BroadcastLayout(A.broadcasted.f, MemoryLayout.(A.broadcasted.args))

## scalar-range broadcast operations ##
# Ranges already support smart broadcasting
for op in (+, -, big)
Expand Down Expand Up @@ -76,3 +90,14 @@ broadcasted(::LazyArrayStyle{N}, ::typeof(*), a::AbstractArray{T,N}, b::Zeros{V,
broadcast(DefaultArrayStyle{N}(), *, a, b)
broadcasted(::LazyArrayStyle{N}, ::typeof(*), a::Zeros{T,N}, b::AbstractArray{V,N}) where {T,V,N} =
broadcast(DefaultArrayStyle{N}(), *, a, b)

const Add = BroadcastArray{<:Any, <:Any, <:Broadcasted{<:Any, <:Any, typeof(+)}}
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Add is an AbstractMarix so maybe AddArray is better as it would be consistent with MulArray? (Though I prefer seeing Add in my code than AddArray or AddMatrix.)

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hmm, this is a tough one. What we could do is the following:

struct AddStyle <: BroadcastStyle end
const Add = Broadcasted{<:AddStyle, <:Any, typeof(+)}
Add(A...) = Broadcasted{AddStyle}(+, A)

where the AddStyle means it's not type piracy.

But I'm not convinced Add should exist at all: what about Minus? Is Add(A,B) that much better than, broadcasted(+, A, B), which also forces it to not materialize?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So maybe just use BroadcastArray like in this PR at the moment?

ref: #8 (comment)

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sure, it's fine for now.

const AddVector = Add{<:Any, 1}
const AddMatrix = Add{<:Any, 2}

"""
Add(A1, A2, …, AN)

A lazy representation of `A1 .+ A2 .+ … .+ AN`; i.e., a shorthand for `BroadcastArray(+, A1, A2, …, AN)`.
"""
Add(As...) = BroadcastArray(+, As...)
16 changes: 16 additions & 0 deletions src/linalg/blasmul.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -197,6 +197,22 @@ function materialize!(M::MatMulVecAdd)
default_blasmul!(α, A, B, iszero(β) ? false : β, C)
end

for MulAdd_ in [MatMulMatAdd, MatMulVecAdd]
# `MulAdd{<:BroadcastLayout{typeof(+)}}` cannot "win" against
# `MatMulMatAdd` and `MatMulVecAdd` hence `@eval`:
@eval function materialize!(M::$MulAdd_{<:BroadcastLayout{typeof(+)}})
α, A, B, β, C = M.α, M.A, M.B, M.β, M.C
if C ≡ B
B = copy(B)
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I noticed that you are testing B .= Mul(A,B) so I added this branch but is it a good idea? First, aliasing can happen through view etc. so using Base.mightalias would be safer? Also, wouldn't it be better to just error out if we can't guarantee it is allocation-free?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

is it a good idea?

I'm not sure. I think I added that line just for completeness, I don't think we necessary need it.

so using Base.mightalias would be safer?

Yes, but when I wrote that line I didn't know about mightalias.

Also, wouldn't it be better to just error out if we can't guarantee it is allocation-free?

Possibly. I believe broadcasting has an unaliascopy for this situation, but this does allocate... I'm happy with either behaviour so maybe its safest to throw an error until someone complains about it?

Copy link
Member Author

@tkf tkf Nov 26, 2018 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

broadcasting has an unaliascopy for this situation, but this does allocate

Ah, yes. So throwing error is incompatible with standard broadcasting behavior... I missed that point. I wonder what was the reasoning for unaliascopy fallback. Maybe it was something like "if you benchmark it you find the GC".

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We can change this here and elsewhere to use unaliascopy at a later date.

end
lmul!(β, C)
for A in A.broadcasted.args
C .= α .* Mul(A, B) .+ C
end
C
end
end

# make copy to make sure always works
@inline function _gemv!(tA, α, A, x, β, y)
if x ≡ y
Expand Down
3 changes: 3 additions & 0 deletions src/linalg/lazymul.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -64,3 +64,6 @@ macro lazylmul(Typ)
LinearAlgebra.lmul!(A::$Typ, x::StridedMatrix) = copyto!(x, LazyArrays.Mul(A,x))
end)
end


@lazymul AddMatrix
8 changes: 7 additions & 1 deletion test/memorylayouttests.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@ using LazyArrays, LinearAlgebra, FillArrays, Test
UnitUpperTriangularLayout, LowerTriangularLayout,
UnitLowerTriangularLayout, ScalarLayout,
hermitiandata, symmetricdata, FillLayout, ZerosLayout,
VcatLayout
VcatLayout, BroadcastLayout, Add

struct FooBar end
struct FooNumber <: Number end
Expand Down Expand Up @@ -142,4 +142,10 @@ struct FooNumber <: Number end
@test MemoryLayout(Vcat(Ones(10),Zeros(10))) == VcatLayout((FillLayout(), ZerosLayout()))
@test MemoryLayout(Vcat([1.],Zeros(10))) == VcatLayout((DenseColumnMajor(), ZerosLayout()))
end

@testset "BroadcastArray" begin
A = [1.0 2; 3 4]
@test MemoryLayout(Add(A, Fill(0, (2, 2)), Zeros(2, 2))) ==
BroadcastLayout(+, (DenseColumnMajor(), FillLayout(), ZerosLayout()))
end
end
95 changes: 94 additions & 1 deletion test/multests.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
using Test, LinearAlgebra, LazyArrays, StaticArrays, FillArrays
import LazyArrays: MulAdd, MemoryLayout, DenseColumnMajor, DiagonalLayout, SymTridiagonalLayout
import LazyArrays: MulAdd, MemoryLayout, DenseColumnMajor, DiagonalLayout, SymTridiagonalLayout, Add
import Base.Broadcast: materialize, materialize!

@testset "Mul" begin
Expand Down Expand Up @@ -640,6 +640,9 @@ import Base.Broadcast: materialize, materialize!
Ac = A'
blasnoalloc(c, 2.0, Ac, x, 3.0, y)
@test @allocated(blasnoalloc(c, 2.0, Ac, x, 3.0, y)) == 0
Aa = Add(A, Ac)
blasnoalloc(c, 2.0, Aa, x, 3.0, y)
@test_broken @allocated(blasnoalloc(c, 2.0, Aa, x, 3.0, y)) == 0
end

@testset "multi-argument mul" begin
Expand Down Expand Up @@ -670,3 +673,93 @@ import Base.Broadcast: materialize, materialize!
@test Matrix(M) ≈ A^2
end
end


@testset "Add" begin
@testset "gemv Float64" begin
for A in (Add(randn(5,5), randn(5,5)),
Add(randn(5,5), view(randn(9, 5), 1:2:9, :))),
b in (randn(5), view(randn(5),:), view(randn(5),1:5), view(randn(9),1:2:9))

à = copy(A)
c = similar(b)

c .= Mul(A,b)
@test c ≈ Ã*b ≈ BLAS.gemv!('N', 1.0, Ã, b, 0.0, similar(c))

copyto!(c, Mul(A,b))
@test c ≈ Ã*b ≈ BLAS.gemv!('N', 1.0, Ã, b, 0.0, similar(c))

b̃ = copy(b)
copyto!(b̃, Mul(A,b̃))
@test c ≈ b̃

c .= 2.0 .* Mul(A,b)
@test c ≈ BLAS.gemv!('N', 2.0, Ã, b, 0.0, similar(c))

c = copy(b)
c .= Mul(A,b) .+ c
@test c ≈ BLAS.gemv!('N', 1.0, Ã, b, 1.0, copy(b))

c = copy(b)
c .= Mul(A,b) .+ 2.0 .* c
@test c ≈ BLAS.gemv!('N', 1.0, Ã, b, 2.0, copy(b))

c = copy(b)
c .= 2.0 .* Mul(A,b) .+ c
@test c ≈ BLAS.gemv!('N', 2.0, Ã, b, 1.0, copy(b))

c = copy(b)
c .= 3.0 .* Mul(A,b) .+ 2.0 .* c
@test c ≈ BLAS.gemv!('N', 3.0, Ã, b, 2.0, copy(b))

d = similar(c)
c = copy(b)
d .= 3.0 .* Mul(A,b) .+ 2.0 .* c
@test d ≈ BLAS.gemv!('N', 3.0, Ã, b, 2.0, copy(b))
end
end

@testset "gemm" begin
for A in (Add(randn(5,5), randn(5,5)),
Add(randn(5,5), view(randn(9, 5), 1:2:9, :))),
B in (randn(5,5), view(randn(5,5),:,:), view(randn(5,5),1:5,:),
view(randn(5,5),1:5,1:5), view(randn(5,5),:,1:5))

à = copy(A)
C = similar(B)

C .= Mul(A,B)
@test C ≈ BLAS.gemm!('N', 'N', 1.0, Ã, B, 0.0, similar(C))

B .= Mul(A,B)
@test C ≈ B

C .= 2.0 .* Mul(A,B)
@test C ≈ BLAS.gemm!('N', 'N', 2.0, Ã, B, 0.0, similar(C))

C = copy(B)
C .= Mul(A,B) .+ C
@test C ≈ BLAS.gemm!('N', 'N', 1.0, Ã, B, 1.0, copy(B))


C = copy(B)
C .= Mul(A,B) .+ 2.0 .* C
@test C ≈ BLAS.gemm!('N', 'N', 1.0, Ã, B, 2.0, copy(B))

C = copy(B)
C .= 2.0 .* Mul(A,B) .+ C
@test C ≈ BLAS.gemm!('N', 'N', 2.0, Ã, B, 1.0, copy(B))


C = copy(B)
C .= 3.0 .* Mul(A,B) .+ 2.0 .* C
@test C ≈ BLAS.gemm!('N', 'N', 3.0, Ã, B, 2.0, copy(B))

d = similar(C)
C = copy(B)
d .= 3.0 .* Mul(A,B) .+ 2.0 .* C
@test d ≈ BLAS.gemm!('N', 'N', 3.0, Ã, B, 2.0, copy(B))
end
end
end