https://github.com/JuliaLang/julia
Tip revision: 223e40f33c4fd16e100231dbef63d810497132fe authored by Yichao Yu on 22 November 2016, 14:29:50 UTC
More robust fenv_constants
More robust fenv_constants
Tip revision: 223e40f
broadcast.jl
# This file is a part of Julia. License is MIT: http://julialang.org/license
module Broadcast
using Base.Cartesian
using Base: promote_eltype_op, @get!, _msk_end, unsafe_bitgetindex, linearindices, tail, OneTo, to_shape
import Base: .+, .-, .*, ./, .\, .//, .==, .<, .!=, .<=, .÷, .%, .<<, .>>, .^
export broadcast, broadcast!, bitbroadcast, dotview
export broadcast_getindex, broadcast_setindex!
## Broadcasting utilities ##
# fallback routines for broadcasting with no arguments or with scalars
# to just produce a scalar result:
broadcast(f) = f()
broadcast(f, x::Number...) = f(x...)
# special cases for "X .= ..." (broadcast!) assignments
broadcast!(::typeof(identity), X::AbstractArray, x::Number) = fill!(X, x)
broadcast!(f, X::AbstractArray) = fill!(X, f())
broadcast!(f, X::AbstractArray, x::Number...) = fill!(X, f(x...))
function broadcast!{T,S,N}(::typeof(identity), x::AbstractArray{T,N}, y::AbstractArray{S,N})
check_broadcast_shape(size(x), size(y))
copy!(x, y)
end
## Calculate the broadcast shape of the arguments, or error if incompatible
# array inputs
broadcast_shape() = ()
broadcast_shape(A) = indices(A)
@inline broadcast_shape(A, B...) = broadcast_shape((), indices(A), map(indices, B)...)
# shape (i.e., tuple-of-indices) inputs
broadcast_shape(shape::Tuple) = shape
@inline broadcast_shape(shape::Tuple, shape1::Tuple, shapes::Tuple...) = broadcast_shape(_bcs((), shape, shape1), shapes...)
# _bcs consolidates two shapes into a single output shape
_bcs(out, ::Tuple{}, ::Tuple{}) = out
@inline _bcs(out, ::Tuple{}, newshape) = _bcs((out..., newshape[1]), (), tail(newshape))
@inline _bcs(out, shape, ::Tuple{}) = _bcs((out..., shape[1]), tail(shape), ())
@inline function _bcs(out, shape, newshape)
newout = _bcs1(shape[1], newshape[1])
_bcs((out..., newout), tail(shape), tail(newshape))
end
# _bcs1 handles the logic for a single dimension
_bcs1(a::Integer, b::Integer) = a == 1 ? b : (b == 1 ? a : (a == b ? a : throw(DimensionMismatch("arrays could not be broadcast to a common size"))))
_bcs1(a::Integer, b) = a == 1 ? b : (first(b) == 1 && last(b) == a ? b : throw(DimensionMismatch("arrays could not be broadcast to a common size")))
_bcs1(a, b::Integer) = _bcs1(b, a)
_bcs1(a, b) = _bcsm(b, a) ? b : (_bcsm(a, b) ? a : throw(DimensionMismatch("arrays could not be broadcast to a common size")))
# _bcsm tests whether the second index is consistent with the first
_bcsm(a, b) = a == b || length(b) == 1
_bcsm(a, b::Number) = b == 1
_bcsm(a::Number, b::Number) = a == b || b == 1
## Check that all arguments are broadcast compatible with shape
## Check that all arguments are broadcast compatible with shape
# comparing one input against a shape
check_broadcast_shape(::Tuple{}) = nothing
check_broadcast_shape(::Tuple{}, A::Union{AbstractArray,Number}) = check_broadcast_shape((), indices(A))
check_broadcast_shape(shp) = nothing
check_broadcast_shape(shp, A) = check_broadcast_shape(shp, indices(A))
check_broadcast_shape(::Tuple{}, ::Tuple{}) = nothing
check_broadcast_shape(shp, ::Tuple{}) = nothing
check_broadcast_shape(::Tuple{}, Ashp::Tuple) = throw(DimensionMismatch("cannot broadcast array to have fewer dimensions"))
function check_broadcast_shape(shp, Ashp::Tuple)
_bcsm(shp[1], Ashp[1]) || throw(DimensionMismatch("array could not be broadcast to match destination"))
check_broadcast_shape(tail(shp), tail(Ashp))
end
# comparing many inputs
@inline function check_broadcast_shape(shp, A, As...)
check_broadcast_shape(shp, A)
check_broadcast_shape(shp, As...)
end
## Indexing manipulations
# newindex(I, keep, Idefault) replaces a CartesianIndex `I` with something that
# is appropriate for a particular broadcast array/scalar. `keep` is a
# NTuple{N,Bool}, where keep[d] == true means that one should preserve
# I[d]; if false, replace it with Idefault[d].
@inline newindex(I::CartesianIndex, keep, Idefault) = CartesianIndex(_newindex(I.I, keep, Idefault))
@inline _newindex(I, keep, Idefault) =
(ifelse(keep[1], I[1], Idefault[1]), _newindex(tail(I), tail(keep), tail(Idefault))...)
@inline _newindex(I, keep::Tuple{}, Idefault) = () # truncate if keep is shorter than I
# newindexer(shape, A) generates `keep` and `Idefault` (for use by
# `newindex` above) for a particular array `A`, given the
# broadcast_shape `shape`
# `keep` is equivalent to map(==, indices(A), shape) (but see #17126)
newindexer(shape, x::Number) = (), ()
@inline newindexer(shape, A) = newindexer(shape, indices(A))
@inline newindexer(shape, indsA::Tuple{}) = (), ()
@inline function newindexer(shape, indsA::Tuple)
ind1 = indsA[1]
keep, Idefault = newindexer(tail(shape), tail(indsA))
(shape[1] == ind1, keep...), (first(ind1), Idefault...)
end
# Equivalent to map(x->newindexer(shape, x), As) (but see #17126)
map_newindexer(shape, ::Tuple{}) = (), ()
@inline function map_newindexer(shape, As)
A1 = As[1]
keeps, Idefaults = map_newindexer(shape, tail(As))
keep, Idefault = newindexer(shape, A1)
(keep, keeps...), (Idefault, Idefaults...)
end
# For output BitArrays
const bitcache_chunks = 64 # this can be changed
const bitcache_size = 64 * bitcache_chunks # do not change this
dumpbitcache(Bc::Vector{UInt64}, bind::Int, C::Vector{Bool}) =
Base.copy_to_bitarray_chunks!(Bc, ((bind - 1) << 6) + 1, C, 1, min(bitcache_size, (length(Bc)-bind+1) << 6))
## Broadcasting core
# nargs encodes the number of As arguments (which matches the number
# of keeps). The first two type parameters are to ensure specialization.
@generated function _broadcast!{K,ID,AT,nargs}(f, B::AbstractArray, keeps::K, Idefaults::ID, As::AT, ::Type{Val{nargs}})
quote
$(Expr(:meta, :noinline))
# destructure the keeps and As tuples
@nexprs $nargs i->(A_i = As[i])
@nexprs $nargs i->(keep_i = keeps[i])
@nexprs $nargs i->(Idefault_i = Idefaults[i])
@simd for I in CartesianRange(indices(B))
# reverse-broadcast the indices
@nexprs $nargs i->(I_i = newindex(I, keep_i, Idefault_i))
# extract array values
@nexprs $nargs i->(@inbounds val_i = A_i[I_i])
# call the function and store the result
@inbounds B[I] = @ncall $nargs f val
end
end
end
# For BitArray outputs, we cache the result in a "small" Vector{Bool},
# and then copy in chunks into the output
@generated function _broadcast!{K,ID,AT,nargs}(f, B::BitArray, keeps::K, Idefaults::ID, As::AT, ::Type{Val{nargs}})
quote
$(Expr(:meta, :noinline))
# destructure the keeps and As tuples
@nexprs $nargs i->(A_i = As[i])
@nexprs $nargs i->(keep_i = keeps[i])
@nexprs $nargs i->(Idefault_i = Idefaults[i])
C = Vector{Bool}(bitcache_size)
Bc = B.chunks
ind = 1
cind = 1
@simd for I in CartesianRange(indices(B))
# reverse-broadcast the indices
@nexprs $nargs i->(I_i = newindex(I, keep_i, Idefault_i))
# extract array values
@nexprs $nargs i->(@inbounds val_i = A_i[I_i])
# call the function and store the result
@inbounds C[ind] = @ncall $nargs f val
ind += 1
if ind > bitcache_size
dumpbitcache(Bc, cind, C)
cind += bitcache_chunks
ind = 1
end
end
if ind > 1
@inbounds C[ind:bitcache_size] = false
dumpbitcache(Bc, cind, C)
end
end
end
@inline function broadcast!{nargs}(f, B::AbstractArray, As::Vararg{Any,nargs})
shape = indices(B)
check_broadcast_shape(shape, As...)
keeps, Idefaults = map_newindexer(shape, As)
_broadcast!(f, B, keeps, Idefaults, As, Val{nargs})
B
end
# broadcast with computed element type
@generated function _broadcast!{K,ID,AT,nargs}(f, B::AbstractArray, keeps::K, Idefaults::ID, As::AT, ::Type{Val{nargs}}, iter, st, count)
quote
$(Expr(:meta, :noinline))
# destructure the keeps and As tuples
@nexprs $nargs i->(A_i = As[i])
@nexprs $nargs i->(keep_i = keeps[i])
@nexprs $nargs i->(Idefault_i = Idefaults[i])
while !done(iter, st)
I, st = next(iter, st)
# reverse-broadcast the indices
@nexprs $nargs i->(I_i = newindex(I, keep_i, Idefault_i))
# extract array values
@nexprs $nargs i->(@inbounds val_i = A_i[I_i])
# call the function
V = @ncall $nargs f val
S = typeof(V)
# store the result
if S <: eltype(B)
@inbounds B[I] = V
else
R = typejoin(eltype(B), S)
new = similar(B, R)
for II in take(iter, count)
new[II] = B[II]
end
new[I] = V
return _broadcast!(f, new, keeps, Idefaults, As, Val{nargs}, iter, st, count+1)
end
count += 1
end
return B
end
end
function broadcast_t(f, ::Type{Any}, As...)
shape = broadcast_shape(As...)
iter = CartesianRange(shape)
if isempty(iter)
return similar(Array{Any}, shape)
end
nargs = length(As)
keeps, Idefaults = map_newindexer(shape, As)
st = start(iter)
I, st = next(iter, st)
val = f([ As[i][newindex(I, keeps[i], Idefaults[i])] for i=1:nargs ]...)
B = similar(Array{typeof(val)}, shape)
B[I] = val
return _broadcast!(f, B, keeps, Idefaults, As, Val{nargs}, iter, st, 1)
end
@inline broadcast_t(f, T, As...) = broadcast!(f, similar(Array{T}, broadcast_shape(As...)), As...)
@inline broadcast(f, As...) = broadcast_t(f, promote_eltype_op(f, As...), As...)
# alternate, more compact implementation; unfortunately slower.
# also the `collect` machinery doesn't yet support arbitrary index bases.
#=
@generated function _broadcast{nargs}(f, keeps, As, ::Type{Val{nargs}}, iter)
quote
collect((@ncall $nargs f i->As[i][newindex(I, keeps[i])]) for I in iter)
end
end
function broadcast(f, As...)
shape = broadcast_shape(As...)
iter = CartesianRange(shape)
keeps, Idefaults = map_newindexer(shape, As)
naT = Val{nfields(As)}
_broadcast(f, keeps, Idefaults, As, naT, iter)
end
=#
@inline bitbroadcast(f, As...) = broadcast!(f, similar(BitArray, broadcast_shape(As...)), As...)
broadcast_getindex(src::AbstractArray, I::AbstractArray...) = broadcast_getindex!(similar(Array{eltype(src)}, broadcast_shape(I...)), src, I...)
@generated function broadcast_getindex!(dest::AbstractArray, src::AbstractArray, I::AbstractArray...)
N = length(I)
Isplat = Expr[:(I[$d]) for d = 1:N]
quote
@nexprs $N d->(I_d = I[d])
check_broadcast_shape(indices(dest), $(Isplat...)) # unnecessary if this function is never called directly
checkbounds(src, $(Isplat...))
@nexprs $N d->(@nexprs $N k->(Ibcast_d_k = indices(I_k, d) == OneTo(1)))
@nloops $N i dest d->(@nexprs $N k->(j_d_k = Ibcast_d_k ? 1 : i_d)) begin
@nexprs $N k->(@inbounds J_k = @nref $N I_k d->j_d_k)
@inbounds (@nref $N dest i) = (@nref $N src J)
end
dest
end
end
@generated function broadcast_setindex!(A::AbstractArray, x, I::AbstractArray...)
N = length(I)
Isplat = Expr[:(I[$d]) for d = 1:N]
quote
@nexprs $N d->(I_d = I[d])
checkbounds(A, $(Isplat...))
shape = broadcast_shape($(Isplat...))
@nextract $N shape d->(length(shape) < d ? OneTo(1) : shape[d])
@nexprs $N d->(@nexprs $N k->(Ibcast_d_k = indices(I_k, d) == 1:1))
if !isa(x, AbstractArray)
xA = convert(eltype(A), x)
@nloops $N i d->shape_d d->(@nexprs $N k->(j_d_k = Ibcast_d_k ? 1 : i_d)) begin
@nexprs $N k->(@inbounds J_k = @nref $N I_k d->j_d_k)
@inbounds (@nref $N A J) = xA
end
else
X = x
@nexprs $N d->(shapelen_d = length(shape_d))
@ncall $N Base.setindex_shape_check X shapelen
Xstate = start(X)
@inbounds @nloops $N i d->shape_d d->(@nexprs $N k->(j_d_k = Ibcast_d_k ? 1 : i_d)) begin
@nexprs $N k->(J_k = @nref $N I_k d->j_d_k)
x_el, Xstate = next(X, Xstate)
(@nref $N A J) = x_el
end
end
A
end
end
## elementwise operators ##
for op in (:÷, :%, :<<, :>>, :-, :/, :\, ://, :^)
@eval $(Symbol(:., op))(A::AbstractArray, B::AbstractArray) = broadcast($op, A, B)
end
.+(As::AbstractArray...) = broadcast(+, As...)
.*(As::AbstractArray...) = broadcast(*, As...)
# ## element-wise comparison operators returning BitArray ##
.==(A::AbstractArray, B::AbstractArray) = bitbroadcast(==, A, B)
.<(A::AbstractArray, B::AbstractArray) = bitbroadcast(<, A, B)
.!=(A::AbstractArray, B::AbstractArray) = bitbroadcast(!=, A, B)
.<=(A::AbstractArray, B::AbstractArray) = bitbroadcast(<=, A, B)
function broadcast_bitarrays(scalarf, bitf, A::AbstractArray{Bool}, B::AbstractArray{Bool})
local shape
try
shape = promote_shape(indices(A), indices(B))
catch
return bitbroadcast(scalarf, A, B)
end
F = BitArray(to_shape(shape))
Fc = F.chunks
Ac = BitArray(A).chunks
Bc = BitArray(B).chunks
if !isempty(Ac) && !isempty(Bc)
for i = 1:length(Fc) - 1
Fc[i] = (bitf)(Ac[i], Bc[i])
end
Fc[end] = (bitf)(Ac[end], Bc[end]) & _msk_end(F)
end
return F
end
biteq(a::UInt64, b::UInt64) = ~a $ b
bitlt(a::UInt64, b::UInt64) = ~a & b
bitneq(a::UInt64, b::UInt64) = a $ b
bitle(a::UInt64, b::UInt64) = ~a | b
.==(A::AbstractArray{Bool}, B::AbstractArray{Bool}) = broadcast_bitarrays(==, biteq, A, B)
.<(A::AbstractArray{Bool}, B::AbstractArray{Bool}) = broadcast_bitarrays(<, bitlt, A, B)
.!=(A::AbstractArray{Bool}, B::AbstractArray{Bool}) = broadcast_bitarrays(!=, bitneq, A, B)
.<=(A::AbstractArray{Bool}, B::AbstractArray{Bool}) = broadcast_bitarrays(<=, bitle, A, B)
function bitcache(op, A, B, refA, refB, l::Int, ind::Int, C::Vector{Bool})
left = l - ind + 1
@inbounds begin
for j = 1:min(bitcache_size, left)
C[j] = (op)(refA(A, ind), refB(B, ind))
ind += 1
end
C[left+1:bitcache_size] = false
end
return ind
end
# note: the following are not broadcasting, but need to be defined here to avoid
# ambiguity warnings
for (f, scalarf) in ((:.==, :(==)),
(:.< , :< ),
(:.!=, :!= ),
(:.<=, :<= ))
for (sigA, sigB, active, refA, refB) in ((:Any, :AbstractArray, :B,
:((A,ind)->A), :((B,ind)->B[ind])),
(:AbstractArray, :Any, :A,
:((A,ind)->A[ind]), :((B,ind)->B)))
shape = :(indices($active))
@eval begin
function ($f)(A::$sigA, B::$sigB)
P = similar(BitArray, $shape)
F = parent(P)
l = length(F)
l == 0 && return F
Fc = F.chunks
C = Array{Bool}(bitcache_size)
ind = first(linearindices($active))
cind = 1
for i = 1:div(l + bitcache_size - 1, bitcache_size)
ind = bitcache($scalarf, A, B, $refA, $refB, l, ind, C)
dumpbitcache(Fc, cind, C)
cind += bitcache_chunks
end
return P
end
end
end
end
## specialized element-wise operators for BitArray
(.^)(A::BitArray, B::AbstractArray{Bool}) = (B .<= A)
(.^)(A::AbstractArray{Bool}, B::AbstractArray{Bool}) = (B .<= A)
function bitcache_pow{T}(Ac::Vector{UInt64}, B::Array{T}, l::Int, ind::Int, C::Vector{Bool})
left = l - ind + 1
@inbounds begin
for j = 1:min(bitcache_size, left)
C[j] = unsafe_bitgetindex(Ac, ind) ^ B[ind]
ind += 1
end
C[left+1:bitcache_size] = false
end
return ind
end
function (.^){T<:Integer}(A::BitArray, B::Array{T})
local shape
try
shape = promote_shape(indices(A), indices(B))
catch
return bitbroadcast(^, A, B)
end
F = BitArray(to_shape(shape))
l = length(F)
l == 0 && return F
Ac = A.chunks
Fc = F.chunks
C = Array{Bool}(bitcache_size)
ind = 1
cind = 1
for i = 1:div(l + bitcache_size - 1, bitcache_size)
ind = bitcache_pow(Ac, B, l, ind, C)
dumpbitcache(Fc, cind, C)
cind += bitcache_chunks
end
return F
end
for (sigA, sigB) in ((BitArray, BitArray),
(AbstractArray{Bool}, BitArray),
(BitArray, AbstractArray{Bool}))
@eval function (.*)(A::$sigA, B::$sigB)
try
return BitArray(A) & BitArray(B)
catch
return bitbroadcast(&, A, B)
end
end
end
############################################################
# x[...] .= f.(y...) ---> broadcast!(f, dotview(x, ...), y...).
# The dotview function defaults to getindex, but we override it in
# a few cases to get the expected in-place behavior without affecting
# explicit calls to view. (All of this can go away if slices
# are changed to generate views by default.)
dotview(args...) = getindex(args...)
dotview(A::AbstractArray, args...) = view(A, args...)
dotview{T<:AbstractArray}(A::AbstractArray{T}, args...) = getindex(A, args...)
# avoid splatting penalty in common cases:
for nargs = 0:5
args = Symbol[Symbol("x",i) for i = 1:nargs]
eval(Expr(:(=), Expr(:call, :dotview, args...),
Expr(:call, :getindex, args...)))
eval(Expr(:(=), Expr(:call, :dotview, :(A::AbstractArray), args...),
Expr(:call, :view, :A, args...)))
end
end # module