Merge pull request #372 from JuliaParallel/jps/at-spawn-broadcast

at-spawn: Add support for broadcasting

Author: jpsamaroo

src/thunk.jl

@@ -367,9 +367,20 @@ macro spawn(exs...)
     _par(ex; lazy=false, opts=opts)
 end
 
+struct ExpandedBroadcast{F} end
+(eb::ExpandedBroadcast{F})(args...) where F =
+    Base.materialize(Base.broadcasted(F, args...))
+replace_broadcast(ex) = ex
+function replace_broadcast(fn::Symbol)
+    if startswith(string(fn), '.')
+        return :($ExpandedBroadcast{$(Symbol(string(fn)[2:end]))}())
+    end
+    return fn
+end
+
 function _par(ex::Expr; lazy=true, recur=true, opts=())
     if ex.head == :call && recur
-        f = ex.args[1]
+        f = replace_broadcast(ex.args[1])
         args = ex.args[2:end]
         opts = esc.(opts)
         if lazy

test/thunk.jl

@@ -61,6 +61,11 @@ end
         @test fetch(c) == 20
     end
     @test Dagger.Sch.EAGER_CONTEXT[] isa Context
+    @testset "broadcast" begin
+        A, B = rand(4), rand(4)
+        @test fetch(@spawn A .+ B) ≈ A .+ B
+        @test fetch(@spawn A .* B) ≈ A .* B
+    end
     @testset "waiting" begin
         a = @spawn sleep(1)
         @test !isready(a)