diff --git a/Project.toml b/Project.toml index 5aaaedb..0b2c0f5 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "DistributedArrays" uuid = "aaf54ef3-cdf8-58ed-94cc-d582ad619b94" -version = "0.6.8" +version = "0.6.9" [deps] Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b" diff --git a/src/broadcast.jl b/src/broadcast.jl index a689faa..e265f93 100644 --- a/src/broadcast.jl +++ b/src/broadcast.jl @@ -5,7 +5,7 @@ # We define a custom ArrayStyle here since we need to keep track of # the fact that it is Distributed and what kind of underlying broadcast behaviour # we will encounter. -struct DArrayStyle{Style <: BroadcastStyle} <: Broadcast.AbstractArrayStyle{Any} end +struct DArrayStyle{Style <: Union{Nothing,BroadcastStyle}} <: Broadcast.AbstractArrayStyle{Any} end DArrayStyle(::S) where {S} = DArrayStyle{S}() DArrayStyle(::S, ::Val{N}) where {S,N} = DArrayStyle(S(Val(N))) DArrayStyle(::Val{N}) where N = DArrayStyle{Broadcast.DefaultArrayStyle{N}}() @@ -50,6 +50,8 @@ function Base.similar(bc::Broadcasted{<:DArrayStyle{Style}}, ::Type{ElType}) whe end ## +# Ref https://docs.julialang.org/en/v1/manual/interfaces/#extending-in-place-broadcast-2 +# # We purposefully only specialise `copyto!`, # Broadcast implementation that defers to the underlying BroadcastStyle. We can't # assume that `getindex` is fast, furthermore we can't assume that the distribution of @@ -119,7 +121,7 @@ end # Distribute broadcast # TODO: How to decide on cuts -@inline bcdistribute(bc::Broadcasted{Style}) where Style = Broadcasted{DArrayStyle{Style}}(bc.f, bcdistribute_args(bc.args), bc.axes) +@inline bcdistribute(bc::Broadcasted{Style}) where Style<:Union{Nothing,BroadcastStyle} = Broadcasted{DArrayStyle{Style}}(bc.f, bcdistribute_args(bc.args), bc.axes) @inline bcdistribute(bc::Broadcasted{Style}) where Style<:DArrayStyle = Broadcasted{Style}(bc.f, bcdistribute_args(bc.args), bc.axes) # ask BroadcastStyle to decide if argument is in need of being distributed @@ -135,7 +137,7 @@ bcdistribute_args(args::Tuple{Any}) = (bcdistribute(args[1]),) bcdistribute_args(args::Tuple{}) = () # dropping axes here since recomputing is easier -@inline bclocal(bc::Broadcasted{DArrayStyle{Style}}, idxs) where Style = Broadcasted{Style}(bc.f, bclocal_args(_bcview(axes(bc), idxs), bc.args)) +@inline bclocal(bc::Broadcasted{DArrayStyle{Style}}, idxs) where Style<:Union{Nothing,BroadcastStyle} = Broadcasted{Style}(bc.f, bclocal_args(_bcview(axes(bc), idxs), bc.args)) # bclocal will do a view of the data and the copy it over # except when the data already is local diff --git a/test/darray.jl b/test/darray.jl index d97c7ca..a812b07 100644 --- a/test/darray.jl +++ b/test/darray.jl @@ -897,6 +897,13 @@ check_leaks() @test Array(g) == Array(a) .- Array(m) .* sin.(Array(c)) end + @testset "Broadcasting into DArray" begin + a .= ones(nrows, ncols) + @test all(isone, a) + a .= 3 .+ abs2.(@view(zeros(nrows, ncols + 5)[:, 6:end])) + @test all(x -> x == 3, a) + end + # @testset "lazy wrapped broadcast" begin # l = similar(a) # l[1:10, :] .= view(a, 1:10, : )