Permalink
Browse files

Add return-type annotation to ngenerate/Cartesian method cache

  • Loading branch information...
1 parent 27e2628 commit 867a86158f171c572d97d738b6e0ae4dac24e706 @timholy timholy committed Feb 3, 2014
Showing with 38 additions and 43 deletions.
  1. +2 −2 base/broadcast.jl
  2. +7 −7 base/cartesian.jl
  3. +19 −20 base/multidimensional.jl
  4. +10 −14 base/reducedim.jl
View
@@ -109,7 +109,7 @@ broadcast!_function(f::Function) = (B, As...) -> broadcast!(f, B, As...)
broadcast_function(f::Function) = (As...) -> broadcast(f, As...)
broadcast_getindex(src::AbstractArray, I::AbstractArray...) = broadcast_getindex!(Array(eltype(src), broadcast_shape(I...)), src, I...)
-@ngenerate N function broadcast_getindex!(dest::AbstractArray, src::AbstractArray, I::NTuple{N, AbstractArray}...)
+@ngenerate N typeof(dest) function broadcast_getindex!(dest::AbstractArray, src::AbstractArray, I::NTuple{N, AbstractArray}...)
check_broadcast_shape(size(dest), I...) # unnecessary if this function is never called directly
checkbounds(src, I...)
@nloops N i dest d->(@nexprs N k->(j_d_k = size(I_k, d) == 1 ? 1 : i_d)) begin
@@ -119,7 +119,7 @@ broadcast_getindex(src::AbstractArray, I::AbstractArray...) = broadcast_getindex
dest
end
-@ngenerate N function broadcast_setindex!(A::AbstractArray, x, I::NTuple{N, AbstractArray}...)
+@ngenerate N typeof(A) function broadcast_setindex!(A::AbstractArray, x, I::NTuple{N, AbstractArray}...)
checkbounds(A, I...)
shape = broadcast_shape(I...)
@nextract N shape d->(length(shape) < d ? 1 : shape[d])
View
@@ -4,19 +4,19 @@ export @ngenerate, @nloops, @nref, @ncall, @nexprs, @nextract, @nall, @ntuple, n
### @ngenerate, for auto-generation of separate versions of functions for different dimensionalities
# Examples (deliberately trivial):
-# @ngenerate N myndims{T,N}(A::Array{T,N}) = N
+# @ngenerate N returntype myndims{T,N}(A::Array{T,N}) = N
# or alternatively
# function gen_body(N::Int)
# quote
# return $N
# end
# end
-# eval(ngenerate(:N, :(myndims{T,N}(A::Array{T,N})), gen_body))
+# eval(ngenerate(:N, returntypeexpr, :(myndims{T,N}(A::Array{T,N})), gen_body))
# The latter allows you to use a single gen_body function for both ngenerate and
# when your function maintains its own method cache (e.g., reduction or broadcasting).
#
# Special syntax for function prototypes:
-# @ngenerate N function myfunction(A::AbstractArray, I::NTuple{N, Int}...)
+# @ngenerate N returntype function myfunction(A::AbstractArray, I::NTuple{N, Int}...)
# for N = 3 translates to
# function myfunction(A::AbstractArray, I_1::Int, I_2::Int, I_3::Int)
# and for the generic (cached) case as
@@ -30,16 +30,16 @@ export @ngenerate, @nloops, @nref, @ncall, @nexprs, @nextract, @nall, @ntuple, n
const CARTESIAN_DIMS = 2 # FIXME: increase after testing is complete
-macro ngenerate(itersym, funcexpr)
+macro ngenerate(itersym, returntypeexpr, funcexpr)
isfuncexpr(funcexpr) || error("Requires a function expression")
- esc(ngenerate(itersym, funcexpr.args[1], N->sreplace!(copy(funcexpr.args[2]), itersym, N)))
+ esc(ngenerate(itersym, returntypeexpr, funcexpr.args[1], N->sreplace!(copy(funcexpr.args[2]), itersym, N)))
end
generate1(itersym, prototype, bodyfunc, N::Int, varname, T) =
Expr(:function, spliceint!(sreplace!(resolvesplat!(copy(prototype), varname, T, N), itersym, N)),
resolvesplats!(bodyfunc(N), varname, N))
-function ngenerate(itersym, prototype, bodyfunc, dims=1:CARTESIAN_DIMS, makecached::Bool = true)
+function ngenerate(itersym, returntypeexpr, prototype, bodyfunc, dims=1:CARTESIAN_DIMS, makecached::Bool = true)
varname, T = get_splatinfo(prototype, itersym)
# Generate versions for specific dimensions
fdim = [generate1(itersym, prototype, bodyfunc, N, varname, T) for N in dims]
@@ -71,7 +71,7 @@ function ngenerate(itersym, prototype, bodyfunc, dims=1:CARTESIAN_DIMS, makecach
_F_
end)
end
- $(dictname)[$itersym]($(fargs...))
+ ($(dictname)[$itersym]($(fargs...)))::$returntypeexpr
end)
Expr(:block, fdim..., quote
let $dictname = Dict{Int,Function}()
View
@@ -1,29 +1,27 @@
### From array.jl
-@ngenerate N function _checksize(A::AbstractArray, I::NTuple{N, Any}...)
+@ngenerate N Nothing function checksize(A::AbstractArray, I::NTuple{N, Any}...)
@nexprs N d->(size(A, d) == length(I_d) || throw(DimensionMismatch("Index $d has length $(length(I_d)), but size(A, $d) = $(size(A,d))")))
nothing
end
-checksize(A, I) = (_checksize(A, I); return nothing)
-checksize(A, I, J) = (_checksize(A, I, J); return nothing)
-checksize(A, I...) = (_checksize(A, I...); return nothing)
unsafe_getindex(v::Real, ind::Integer) = v
unsafe_getindex(v::Ranges, ind::Integer) = first(v) + (ind-1)*step(v)
unsafe_getindex(v::AbstractArray, ind::Integer) = v[ind]
# Version that uses cartesian indexing for src
-@ngenerate N function _getindex!(dest::Array, src::AbstractArray, I::NTuple{N,Union(Int,AbstractVector)}...)
+@ngenerate N typeof(dest) function _getindex!(dest::Array, src::AbstractArray, I::NTuple{N,Union(Int,AbstractVector)}...)
checksize(dest, I...)
k = 1
@nloops N i dest d->(@inbounds j_d = unsafe_getindex(I_d, i_d)) begin
@inbounds dest[k] = (@nref N src j)
k += 1
end
+ dest
end
# Version that uses linear indexing for src
-@ngenerate N function _getindex!(dest::Array, src::Array, I::NTuple{N,Union(Int,AbstractVector)}...)
+@ngenerate N typeof(dest) function _getindex!(dest::Array, src::Array, I::NTuple{N,Union(Int,AbstractVector)}...)
checksize(dest, I...)
stride_1 = 1
@nexprs N d->(stride_{d+1} = stride_d*size(src,d))
@@ -33,6 +31,7 @@ end
@inbounds dest[k] = src[offset_0]
k += 1
end
+ dest
end
getindex!(dest, src, I) = (checkbounds(src, I); _getindex!(dest, src, to_index(I)); return dest)
@@ -57,7 +56,7 @@ function getindex(A::Array, I::Union(Real,AbstractVector), J::Union(Real,Abstrac
end
-@ngenerate N function _setindex!(A::Array, x, I::NTuple{N,Union(Int,AbstractArray)}...)
+@ngenerate N typeof(A) function _setindex!(A::Array, x, I::NTuple{N,Union(Int,AbstractArray)}...)
stride_1 = 1
@nexprs N d->(stride_{d+1} = stride_d*size(A,d))
@nexprs N d->(offset_d = 1) # really only need offset_$N = 1
@@ -75,6 +74,7 @@ end
k += 1
end
end
+ A
end
function setindex!(A::Array, x, I::Union(Real,AbstractArray), J::Union(Real,AbstractArray))
@@ -94,7 +94,7 @@ function setindex!(A::Array, x, I::Union(Real,AbstractArray)...)
end
-@ngenerate N function findn{T,N}(A::AbstractArray{T,N})
+@ngenerate N NTuple{N,Vector{Int}} function findn{T,N}(A::AbstractArray{T,N})
nnzA = countnz(A)
@nexprs N d->(I_d = Array(Int, nnzA))
k = 1
@@ -127,7 +127,7 @@ function gen_getindex_body(N::Int)
end
end
-eval(ngenerate(:N, :(getindex{T}(s::SubArray{T,N}, ind::Integer)), gen_getindex_body, 2:5, false))
+eval(ngenerate(:N, nothing, :(getindex{T}(s::SubArray{T,N}, ind::Integer)), gen_getindex_body, 2:5, false))
function gen_setindex!_body(N::Int)
@@ -145,19 +145,18 @@ function gen_setindex!_body(N::Int)
end
end
-eval(ngenerate(:N, :(setindex!{T}(s::SubArray{T,N}, v, ind::Integer)), gen_setindex!_body, 2:5, false))
+eval(ngenerate(:N, nothing, :(setindex!{T}(s::SubArray{T,N}, v, ind::Integer)), gen_setindex!_body, 2:5, false))
### from abstractarray.jl
-@ngenerate N function _fill!{T,N}(A::AbstractArray{T,N}, x)
+@ngenerate N typeof(A) function fill!{T,N}(A::AbstractArray{T,N}, x)
@nloops N i A begin
@inbounds (@nref N A i) = x
end
+ A
end
-fill!(A::AbstractArray, x) = (_fill!(A, x); return A)
-
### from bitarray.jl
@@ -173,7 +172,7 @@ function getindex(B::BitArray, I0::Range1)
end
# TODO: extend to I:Union(Real,AbstractArray)... (i.e. not necessarily contiguous)
-@ngenerate N function getindex(B::BitArray, I0::Range1, I::NTuple{N,Union(Real,Range1)}...)
+@ngenerate N BitArray{length(I)+1} function getindex(B::BitArray, I0::Range1, I::NTuple{N,Union(Real,Range1)}...)
ndims(B) < N+1 && error("wrong number of dimensions")
checkbounds(B, I0, I...)
X = BitArray(index_shape(I0, I...))
@@ -207,7 +206,7 @@ end
return X
end
-@ngenerate N function getindex(B::BitArray, I::NTuple{N,Union(Real,AbstractVector)}...)
+@ngenerate N BitArray{length(I)} function getindex(B::BitArray, I::NTuple{N,Union(Real,AbstractVector)}...)
checkbounds(B, I...)
@nexprs N d->(I_d = to_index(I_d))
X = BitArray(index_shape(I...))
@@ -237,7 +236,7 @@ function setindex!(B::BitArray, X::BitArray, I0::Range1)
end
# TODO: extend to I:Union(Real,AbstractArray)... (i.e. not necessarily contiguous)
-@ngenerate N function setindex!(B::BitArray, X::BitArray, I0::Range1, I::NTuple{N,Union(Real,Range1)}...)
+@ngenerate N typeof(B) function setindex!(B::BitArray, X::BitArray, I0::Range1, I::NTuple{N,Union(Real,Range1)}...)
ndims(B) != N+1 && error("wrong number of dimensions in assigment")
I0 = to_index(I0)
lI = length(I0)
@@ -275,7 +274,7 @@ end
return B
end
-@ngenerate N function setindex!(B::BitArray, X::AbstractArray, I::NTuple{N,Union(Real,AbstractArray)}...)
+@ngenerate N typeof(B) function setindex!(B::BitArray, X::AbstractArray, I::NTuple{N,Union(Real,AbstractArray)}...)
checkbounds(B, I...)
@nexprs N d->(I_d = to_index(I_d))
nel = 1
@@ -294,7 +293,7 @@ end
return B
end
-@ngenerate N function setindex!(B::BitArray, x, I::NTuple{N,Union(Real,AbstractArray)}...)
+@ngenerate N typeof(B) function setindex!(B::BitArray, x, I::NTuple{N,Union(Real,AbstractArray)}...)
x = convert(Bool, x)
checkbounds(B, I...)
@nexprs N d->(I_d = to_index(I_d))
@@ -305,7 +304,7 @@ end
return B
end
-@ngenerate N function findn{N}(B::BitArray{N})
+@ngenerate N NTuple{N,Vector{Int}} function findn{N}(B::BitArray{N})
nnzB = countnz(B)
I = ntuple(N, x->Array(Int, nnzB))
if nnzB > 0
@@ -322,7 +321,7 @@ end
for (V, PT, BT) in [((:N,), BitArray, BitArray), ((:T,:N), Array, StridedArray)]
@eval begin
- @ngenerate N function permutedims!{$(V...)}(P::$PT{$(V...)}, B::$BT{$(V...)}, perm)
+ @ngenerate N typeof(P) function permutedims!{$(V...)}(P::$PT{$(V...)}, B::$BT{$(V...)}, perm)
dimsB = size(B)
(length(perm) == N && isperm(perm)) || error("no valid permutation of dimensions")
dimsP = size(P)
View
@@ -76,8 +76,7 @@ function reducedim!(f::Function, R, A)
else
func = reducedim_cache[key]
end
- func(R, A)
- R
+ func(R, A)::typeof(R)
end
end # let reducedim_cache
@@ -92,6 +91,8 @@ function gen_reduction_body(N, f::Function)
(size(R, i) == size(A, i) || size(R, i) == 1) || throw(DimensionMismatch("Reduction on array of size $(size(A)) with output of size $(size(R))"))
end
@nextract $N sizeR d->size(R,d)
+ # If we're reducing along dimension 1, for efficiency we can make use of a temporary.
+ # Otherwise, keep the result in R so that we traverse A in storage order.
if size(R, 1) < size(A, 1)
@nloops $N i d->(d>1? (1:size(A,d)) : (1:1)) d->(j_d = sizeR_d==1 ? 1 : i_d) begin
@inbounds tmp = (@nref $N R j)
@@ -105,6 +106,7 @@ function gen_reduction_body(N, f::Function)
@inbounds (@nref $N R j) = ($F)((@nref $N R j), (@nref $N A i))
end
end
+ R
end
end
@@ -115,25 +117,19 @@ reduction_init{T}(A::AbstractArray, region, initial::T) = fill!(similar(A,T,redu
# For performance, these bypass reducedim_cache
all(A::AbstractArray{Bool}, region) = all!(reduction_init(A,region,true), A)
-all!(R, A) = (_all!(R,A); return R)
-eval(ngenerate(:N, :(_all!{N}(R::AbstractArray, A::AbstractArray{Bool,N})), N->gen_reduction_body(N, &)))
+eval(ngenerate(:N, :(typeof(R)), :(all!{N}(R::AbstractArray, A::AbstractArray{Bool,N})), N->gen_reduction_body(N, &)))
any(A::AbstractArray{Bool}, region) = any!(reduction_init(A,region,false), A)
-any!(R, A) = (_any!(R,A); return R)
-eval(ngenerate(:N, :(_any!{N}(R::AbstractArray, A::AbstractArray{Bool,N})), N->gen_reduction_body(N, |)))
+eval(ngenerate(:N, :(typeof(R)), :(any!{N}(R::AbstractArray, A::AbstractArray{Bool,N})), N->gen_reduction_body(N, |)))
maximum{T}(A::AbstractArray{T}, region) =
isempty(A) ? similar(A,reduced_dims0(A,region)) : maximum!(reduction_init(A,region,typemin(T)), A)
-maximum!(R, A) = (_maximum!(R,A); return R)
-eval(ngenerate(:N, :(_maximum!{T,N}(R::AbstractArray, A::AbstractArray{T,N})), N->gen_reduction_body(N, scalarmax)))
+eval(ngenerate(:N, :(typeof(R)), :(maximum!{T,N}(R::AbstractArray, A::AbstractArray{T,N})), N->gen_reduction_body(N, scalarmax)))
minimum{T}(A::AbstractArray{T}, region) =
isempty(A) ? similar(A,reduced_dims0(A,region)) : minimum!(reduction_init(A,region,typemax(T)), A)
-minimum!(R, A) = (_minimum!(R,A); return R)
-eval(ngenerate(:N, :(_minimum!{T,N}(R::AbstractArray, A::AbstractArray{T,N})), N->gen_reduction_body(N, scalarmin)))
+eval(ngenerate(:N, :(typeof(R)), :(minimum!{T,N}(R::AbstractArray, A::AbstractArray{T,N})), N->gen_reduction_body(N, scalarmin)))
sum{T}(A::AbstractArray{T}, region) = sum!(reduction_init(A,region,zero(T)), A)
sum(A::AbstractArray{Bool}, region) = sum!(reduction_init(A,region,0), A)
-sum!(R, A) = (_sum!(R,A); return R)
-eval(ngenerate(:N, :(_sum!{T,N}(R::AbstractArray, A::AbstractArray{T,N})), N->gen_reduction_body(N, +)))
+eval(ngenerate(:N, :(typeof(R)), :(sum!{T,N}(R::AbstractArray, A::AbstractArray{T,N})), N->gen_reduction_body(N, +)))
prod{T}(A::AbstractArray{T}, region) = prod!(reduction_init(A,region,one(T)), A)
-prod!(R, A) = (_prod!(R,A); return R)
-eval(ngenerate(:N, :(_prod!{T,N}(R::AbstractArray, A::AbstractArray{T,N})), N->gen_reduction_body(N, *)))
+eval(ngenerate(:N, :(typeof(R)), :(prod!{T,N}(R::AbstractArray, A::AbstractArray{T,N})), N->gen_reduction_body(N, *)))
prod(A::AbstractArray{Bool}, region) = error("use all() instead of prod() for boolean arrays")

0 comments on commit 867a861

Please sign in to comment.