Merge pull request #402 from JuliaParallel/jps/diffeq-support

DArray: Improve DiffEq support

Author: jpsamaroo

Diff for: src/array/alloc.jl

@@ -68,6 +68,14 @@ Base.zeros(p::Blocks, t::Type, dims::Integer...) = zeros(p, t, dims)
 Base.zeros(p::Blocks, dims::Integer...) = zeros(p, Float64, dims)
 Base.zeros(p::Blocks, dims::Tuple) = zeros(p, Float64, dims)
 
+function Base.zero(x::DArray{T,N}) where {T,N}
+    dims = ntuple(i->x.domain.indexes[i].stop, N)
+    sd = first(x.subdomains)
+    part_size = ntuple(i->sd.indexes[i].stop, N)
+    a = zeros(Blocks(part_size...), T, dims)
+    return cached_stage(Context(global_context()), a)
+end
+
 function sprand(p::Blocks, m::Integer, n::Integer, sparsity::Real)
     s = rand(UInt)
     f = function (idx, t,sz)

Diff for: src/array/darray.jl

@@ -163,6 +163,19 @@ function Base.isequal(x::ArrayOp, y::ArrayOp)
     x === y
 end
 
+function Base.similar(x::DArray{T,N,F}) where {T,N,F}
+    alloc(idx, sz) = Array{T,N}(undef, sz)
+    thunks = [Dagger.@spawn alloc(i, size(x)) for (i, x) in enumerate(x.subdomains)]
+    return DArray{T,N,F}(x.domain, x.subdomains, thunks, x.concat)
+end
+
+Base.copy(x::DArray{T,N,F}) where {T,N,F} =
+    cached_stage(Context(global_context()), map(identity, x))::DArray{T,N,F}
+
+# Because OrdinaryDiffEq uses `Base.promote_op(/, ::DArray, ::Real)`
+Base.:(/)(x::DArray{T,N,F}, y::U) where {T<:Real,U<:Real,N,F} =
+    (x ./ y)::DArray{Base.promote_op(/, T, U),N,F}
+
 """
     view(c::DArray, d)
 

Diff for: src/array/operators.jl

@@ -1,4 +1,3 @@
-
 import Base: exp, expm1, log, log10, log1p, sqrt, cbrt, exponent,
              significand, sin, sinpi, cos, cospi, tan, sec, cot, csc,
              sinh, cosh, tanh, coth, sech, csch,
@@ -51,7 +50,7 @@ BroadcastStyle(::DaggerBroadcastStyle, ::BroadcastStyle) = DaggerBroadcastStyle(
 BroadcastStyle(::BroadcastStyle, ::DaggerBroadcastStyle) = DaggerBroadcastStyle()
 
 function Base.copy(b::Broadcast.Broadcasted{<:DaggerBroadcastStyle})
-    BCast(b)
+    cached_stage(Context(global_context()), BCast(b))::DArray
 end
 
 function stage(ctx::Context, node::BCast)

Diff for: test/array.jl

@@ -47,6 +47,30 @@ end
     @test collect(X1) .+ 1 == collect(X2)
 end
 
+@testset "copy/similar" begin
+    X1 = fetch(ones(Blocks(10, 10), 100, 100))
+    X2 = copy(X1)
+    X3 = similar(X1)
+    @test typeof(X1) === typeof(X2) === typeof(X3)
+    @test collect(X1) == collect(X2)
+    @test collect(X1) != collect(X3)
+end
+
+@testset "DiffEq support" begin
+    X = fetch(ones(Blocks(10), 100))
+    X0 = zero(X)
+    @test typeof(X) === typeof(X0)
+    @test all(collect(X0) .== 0)
+    @testset for T in (Int8, Int, Float32, Float64)
+        DT = DArray{Base.promote_op(/, Float64, T), 1, typeof(cat)}
+        @test Base.promote_op(/, typeof(X), T) === DT
+        y = T(2)
+        Xd = X / y
+        @test typeof(Xd) === DT
+        @test collect(Xd) == collect(X) ./ y
+    end
+end
+
 @testset "sum" begin
     X = ones(Blocks(10, 10), 100, 100)
     @test sum(X) == 10000
@@ -70,6 +94,16 @@ end
     @test mean(Y) == 0
 end
 
+@testset "broadcast" begin
+    X1 = fetch(rand(Blocks(10), 100))
+    X2 = X1 .* 3.4
+    @test typeof(X1) === typeof(X2)
+    @test collect(X1) .* 3.4 == collect(X2)
+    X3 = X1 .+ X1
+    @test typeof(X1) === typeof(X3)
+    @test collect(X1) .* 2 == collect(X3)
+end
+
 @testset "distributing an array" begin
     function test_dist(X)
         X1 = Distribute(Blocks(10, 20), X)