comparison src/LazyTensors/lazy_array.jl @ 1207:f1c2a4fa0ee1 performance/get_region_type_inference

Merge default
author Jonatan Werpers <jonatan@werpers.com>
date Fri, 03 Feb 2023 22:14:47 +0100
parents 6fd97b5ed3e5
children c5cf32fab163
comparison
equal deleted inserted replaced
919:b41180efb6c2 1207:f1c2a4fa0ee1
26 size::NTuple{D,Int} 26 size::NTuple{D,Int}
27 end 27 end
28 export LazyFunctionArray 28 export LazyFunctionArray
29 29
30 function LazyFunctionArray(f::F, size::NTuple{D,Int}) where {F<:Function,D} 30 function LazyFunctionArray(f::F, size::NTuple{D,Int}) where {F<:Function,D}
31 T = typeof(f(ones(D)...)) 31 T = typeof(f(ones(Int, D)...))
32 return LazyFunctionArray{F,T,D}(f,size) 32 return LazyFunctionArray{F,T,D}(f,size)
33 end 33 end
34 34
35 Base.size(lfa::LazyFunctionArray) = lfa.size 35 Base.size(lfa::LazyFunctionArray) = lfa.size
36 36
37 function Base.getindex(lfa::LazyFunctionArray{F,T,D}, I::Vararg{Int,D}) where {F,T,D} 37 function Base.getindex(lfa::LazyFunctionArray{F,T,D}, I::Vararg{Int,D}) where {F,T,D}
38 @boundscheck checkbounds(lfa, I...) 38 @boundscheck checkbounds(lfa, I...)
39 return lfa.f(I...) 39 return @inbounds lfa.f(I...)
40 end 40 end
41 41
42 42
43 """ 43 """
44 LazyElementwiseOperation{T,D,Op} <: LazyArray{T,D} 44 LazyElementwiseOperation{T,D,Op} <: LazyArray{T,D}
59 end 59 end
60 60
61 end 61 end
62 LazyElementwiseOperation{T,D,Op}(a::AbstractArray{T,D},b::T) where {T,D,Op} = LazyElementwiseOperation{T,D,Op}(a, LazyConstantArray(b, size(a))) 62 LazyElementwiseOperation{T,D,Op}(a::AbstractArray{T,D},b::T) where {T,D,Op} = LazyElementwiseOperation{T,D,Op}(a, LazyConstantArray(b, size(a)))
63 LazyElementwiseOperation{T,D,Op}(a::T,b::AbstractArray{T,D}) where {T,D,Op} = LazyElementwiseOperation{T,D,Op}(LazyConstantArray(a, size(b)), b) 63 LazyElementwiseOperation{T,D,Op}(a::T,b::AbstractArray{T,D}) where {T,D,Op} = LazyElementwiseOperation{T,D,Op}(LazyConstantArray(a, size(b)), b)
64 # TODO: Move Op to be the first parameter? Compare to Binary operations
65 64
66 Base.size(v::LazyElementwiseOperation) = size(v.a) 65 Base.size(v::LazyElementwiseOperation) = size(v.a)
67 66
68 evaluate(leo::LazyElementwiseOperation{T,D,:+}, I::Vararg{Int,D}) where {T,D} = leo.a[I...] + leo.b[I...] 67 Base.@propagate_inbounds evaluate(leo::LazyElementwiseOperation{T,D,:+}, I::Vararg{Int,D}) where {T,D} = leo.a[I...] + leo.b[I...]
69 evaluate(leo::LazyElementwiseOperation{T,D,:-}, I::Vararg{Int,D}) where {T,D} = leo.a[I...] - leo.b[I...] 68 Base.@propagate_inbounds evaluate(leo::LazyElementwiseOperation{T,D,:-}, I::Vararg{Int,D}) where {T,D} = leo.a[I...] - leo.b[I...]
70 evaluate(leo::LazyElementwiseOperation{T,D,:*}, I::Vararg{Int,D}) where {T,D} = leo.a[I...] * leo.b[I...] 69 Base.@propagate_inbounds evaluate(leo::LazyElementwiseOperation{T,D,:*}, I::Vararg{Int,D}) where {T,D} = leo.a[I...] * leo.b[I...]
71 evaluate(leo::LazyElementwiseOperation{T,D,:/}, I::Vararg{Int,D}) where {T,D} = leo.a[I...] / leo.b[I...] 70 Base.@propagate_inbounds evaluate(leo::LazyElementwiseOperation{T,D,:/}, I::Vararg{Int,D}) where {T,D} = leo.a[I...] / leo.b[I...]
72 71
73 # TODO: Make sure boundschecking is done properly and that the lenght of the vectors are equal 72 function Base.getindex(leo::LazyElementwiseOperation{T,D}, I::Vararg{Int,D}) where {T,D}
74 # NOTE: Boundschecking in getindex functions now assumes that the size of the 73 @boundscheck checkbounds(leo.a, I...)
75 # vectors in the LazyElementwiseOperation are the same size. If we remove the 74 return @inbounds evaluate(leo, I...)
76 # size assertion in the constructor we might have to handle
77 # boundschecking differently.
78 Base.@propagate_inbounds @inline function Base.getindex(leo::LazyElementwiseOperation{T,D}, I::Vararg{Int,D}) where {T,D}
79 @boundscheck if !checkbounds(Bool, leo.a, I...)
80 throw(BoundsError([leo], I...))
81 end
82 return evaluate(leo, I...)
83 end 75 end
84 76
85 # Define lazy operations for AbstractArrays. Operations constructs a LazyElementwiseOperation which 77 # Define lazy operations for AbstractArrays. Operations constructs a LazyElementwiseOperation which
86 # can later be indexed into. Lazy operations are denoted by the usual operator followed by a tilde 78 # can later be indexed into. Lazy operations are denoted by the usual operator followed by a tilde
87 Base.@propagate_inbounds +̃(a::AbstractArray{T,D}, b::AbstractArray{T,D}) where {T,D} = LazyElementwiseOperation{T,D,:+}(a,b) 79 Base.@propagate_inbounds +̃(a::AbstractArray{T,D}, b::AbstractArray{T,D}) where {T,D} = LazyElementwiseOperation{T,D,:+}(a,b)
99 Base.@propagate_inbounds *̃(a::T, b::AbstractArray{T,D}) where {T,D} = LazyElementwiseOperation{T,D,:*}(a,b) 91 Base.@propagate_inbounds *̃(a::T, b::AbstractArray{T,D}) where {T,D} = LazyElementwiseOperation{T,D,:*}(a,b)
100 Base.@propagate_inbounds /̃(a::T, b::AbstractArray{T,D}) where {T,D} = LazyElementwiseOperation{T,D,:/}(a,b) 92 Base.@propagate_inbounds /̃(a::T, b::AbstractArray{T,D}) where {T,D} = LazyElementwiseOperation{T,D,:/}(a,b)
101 93
102 94
103 95
104 # NOTE: Är det knas att vi har till exempel * istället för .* ?? 96 # Overload +,-,*,/ for LazyArrays
105 # Oklart om det ens går att lösa.. 97 # Element wise operation for `*` and `/` are not overloaded for due to conflicts with the behavior
98 # of regular `*` and `/` for AbstractArrays. Use tilde versions instead.
106 Base.@propagate_inbounds Base.:+(a::LazyArray{T,D}, b::LazyArray{T,D}) where {T,D} = a +̃ b 99 Base.@propagate_inbounds Base.:+(a::LazyArray{T,D}, b::LazyArray{T,D}) where {T,D} = a +̃ b
100 Base.@propagate_inbounds Base.:-(a::LazyArray{T,D}, b::LazyArray{T,D}) where {T,D} = a -̃ b
101
107 Base.@propagate_inbounds Base.:+(a::LazyArray{T,D}, b::AbstractArray{T,D}) where {T,D} = a +̃ b 102 Base.@propagate_inbounds Base.:+(a::LazyArray{T,D}, b::AbstractArray{T,D}) where {T,D} = a +̃ b
103 Base.@propagate_inbounds Base.:-(a::LazyArray{T,D}, b::AbstractArray{T,D}) where {T,D} = a -̃ b
104
108 Base.@propagate_inbounds Base.:+(a::AbstractArray{T,D}, b::LazyArray{T,D}) where {T,D} = a +̃ b 105 Base.@propagate_inbounds Base.:+(a::AbstractArray{T,D}, b::LazyArray{T,D}) where {T,D} = a +̃ b
109
110 Base.@propagate_inbounds Base.:-(a::LazyArray{T,D}, b::LazyArray{T,D}) where {T,D} = a -̃ b
111 Base.@propagate_inbounds Base.:-(a::LazyArray{T,D}, b::AbstractArray{T,D}) where {T,D} = a -̃ b
112 Base.@propagate_inbounds Base.:-(a::AbstractArray{T,D}, b::LazyArray{T,D}) where {T,D} = a -̃ b 106 Base.@propagate_inbounds Base.:-(a::AbstractArray{T,D}, b::LazyArray{T,D}) where {T,D} = a -̃ b
113 107
114 # Element wise operation for `*` and `\` are not overloaded due to conflicts with the behavior 108 Base.@propagate_inbounds Base.:+(a::LazyArray{T,D}, b::T) where {T,D} = a +̃ b
115 # of regular `*` and `/` for AbstractArrays. Use tilde versions instead. 109 Base.@propagate_inbounds Base.:-(a::LazyArray{T,D}, b::T) where {T,D} = a -̃ b
110
111 Base.@propagate_inbounds Base.:+(a::T, b::LazyArray{T,D}) where {T,D} = a +̃ b
112 Base.@propagate_inbounds Base.:-(a::T, b::LazyArray{T,D}) where {T,D} = a -̃ b
116 113
117 export +̃, -̃, *̃, /̃ 114 export +̃, -̃, *̃, /̃