Mercurial > repos > public > sbplib_julia

comparison src/LazyTensors/lazy_array.jl @ 347:0a95a829176c performance/lazy_elementwise_operation

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
lazy_array.jl: Bring back the concrete types in the LazyElementwiseOperation struct. We went this to be able to avoid dynamic dispatch
author Jonatan Werpers <jonatan@werpers.com>
date Fri, 25 Sep 2020 21:51:14 +0200
parents 01b851161018
children f65809a26a17
comparison
equal deleted inserted replaced
344:f781d6da7d3d 347:0a95a829176c
21 Struct allowing for lazy evaluation of elementwise operations on AbstractArrays. 21 Struct allowing for lazy evaluation of elementwise operations on AbstractArrays.
22 22
23 A LazyElementwiseOperation contains two arrays together with an operation. 23 A LazyElementwiseOperation contains two arrays together with an operation.
24 The operations are carried out when the LazyElementwiseOperation is indexed. 24 The operations are carried out when the LazyElementwiseOperation is indexed.
25 """ 25 """
26 struct LazyElementwiseOperation{T,D,Op} <: LazyArray{T,D} 26 struct LazyElementwiseOperation{T,D,Op,T1<:AbstractArray{T,D},T2<:AbstractArray{T,D}} <: LazyArray{T,D}
27 a::AbstractArray{T,D} 27 a::T1
28 b::AbstractArray{T,D} 28 b::T2
29 29
30 function LazyElementwiseOperation{T,D,Op}(a::AbstractArray{T,D},b::AbstractArray{T,D}) where {T,D,Op} 30 function LazyElementwiseOperation{T,D,Op}(a::T1,b::T2) where {T,D,Op,T1<:AbstractArray{T,D},T2<:AbstractArray{T,D}}
31 @boundscheck if size(a) != size(b) 31 @boundscheck if size(a) != size(b)
32 throw(DimensionMismatch("dimensions must match")) 32 throw(DimensionMismatch("dimensions must match"))
33 end 33 end
34 return new{T,D,Op}(a,b) 34 return new{T,D,Op,T1,T2}(a,b)
35 end 35 end
36 36
37 LazyElementwiseOperation{T,D,Op}(a::AbstractArray{T,D},b::T) where {T,D,Op} = new{T,D,Op}(a, LazyConstantArray(b, size(a)))
38 LazyElementwiseOperation{T,D,Op}(a::T,b::AbstractArray{T,D}) where {T,D,Op} = new{T,D,Op}(LazyConstantArray(a, size(b)), b)
39 end 37 end
38 LazyElementwiseOperation{T,D,Op}(a::AbstractArray{T,D},b::T) where {T,D,Op} = LazyElementwiseOperation{T,D,Op}(a, LazyConstantArray(b, size(a)))
39 LazyElementwiseOperation{T,D,Op}(a::T,b::AbstractArray{T,D}) where {T,D,Op} = LazyElementwiseOperation{T,D,Op}(LazyConstantArray(a, size(b)), b)
40 # TODO: Move Op to be the first parameter? Compare to Binary operations 40 # TODO: Move Op to be the first parameter? Compare to Binary operations
41 41
42 Base.size(v::LazyElementwiseOperation) = size(v.a) 42 Base.size(v::LazyElementwiseOperation) = size(v.a)
43 43
44 evaluate(leo::LazyElementwiseOperation{T,D,:+}, I::Vararg{Int,D}) where {T,D} = leo.a[I...] + leo.b[I...] 44 evaluate(leo::LazyElementwiseOperation{T,D,:+}, I::Vararg{Int,D}) where {T,D} = leo.a[I...] + leo.b[I...]