Lazy broadcasting macro

README.md

@@ -81,6 +81,18 @@ julia> B = BroadcastArray(+, A, 2);
 julia> B == A .+ 2
 true
 ```
+Such arrays can also be created using the macro `@lazy` which acts on ordinary 
+broadcasting expressions, or the macro `@lazydot` which applies `@.` to add dots first: 
+```julia
+julia> C = rand(1000)';
+
+julia> D = @lazy exp.(C)
+
+julia> E = @lazydot 2 + log(C)
+
+julia> @btime sum(@lazy C .* C'; dims=1) #  1.438 ms (5 allocations: 7.64 MiB) without @lazy
+  74.425 μs (7 allocations: 8.08 KiB)
+```
 
 ## Multiplication
 

REQUIRE

@@ -1,3 +1,4 @@
 julia 0.7
 FillArrays 0.3
 StaticArrays 0.8.3
+MacroTools

src/lazybroadcasting.jl

@@ -37,9 +37,91 @@ getindex(B::BroadcastArray{<:Any,1}, kr::AbstractVector{<:Integer}) =
 
 copy(bc::Broadcasted{<:LazyArrayStyle}) = BroadcastArray(bc)
 
-# issue 16: sum(b, dims=(1,2,3)) faster than sum(b)
-Base._sum(b::BroadcastArray{T,N}, ::Colon) where {T,N} = first(Base._sum(b, ntuple(identity, N)))
-Base._prod(b::BroadcastArray{T,N}, ::Colon) where {T,N} = first(Base._prod(b, ntuple(identity, N)))
+# Replacement for #18.
+# Could extend this to other similar reductions in Base... or apply at lower level? 
+# for (fname, op) in [(:sum, :add_sum), (:prod, :mul_prod),
+#                     (:maximum, :max), (:minimum, :min),
+#                     (:all, :&),       (:any, :|)]
+function Base._sum(f, A::BroadcastArray, ::Colon)
+    bc = A.broadcasted
+    T = Broadcast.combine_eltypes(f ∘ bc.f, bc.args) 
+    out = zero(T)
+    @simd for I in eachindex(bc)
+        @inbounds out += f(bc[I])
+    end
+    out
+end
+function Base._prod(f, A::BroadcastArray, ::Colon)
+    bc = A.broadcasted
+    T = Broadcast.combine_eltypes(f ∘ bc.f, bc.args) 
+    out = one(T)
+    @simd for I in eachindex(bc)
+        @inbounds out *= f(bc[I])
+    end
+    out
+end
+
+# Macros for lazy broadcasting, #21 WIP
+# based on @dawbarton  https://discourse.julialang.org/t/19641/20
+# and @tkf            https://github.com/JuliaLang/julia/issues/19198#issuecomment-457967851
+# and @chethega      https://github.com/JuliaLang/julia/pull/30939
+
+export @lazy, @lazydot, @□, @⊡
+
+lazy(::Any) = error("function `lazy` must be called with a dot")
+struct LazyCast{T}
+    value::T
+end
+Broadcast.broadcasted(::typeof(lazy), x) = LazyCast(x)
+Broadcast.materialize(x::LazyCast) = BroadcastArray(x.value)
+
+lazyhelp = """
+    @lazy A .+ B     == @□ A .+ B
+    @lazydot A + B   == @⊡ A + B
+
+Macros for creating lazy `BroadcastArray`s: `@lazy` expects a broadcasting expression, 
+while `@lazydot` applies `@.` first. Short forms are typed `@\\square` and `@\\boxdot` 
+(or perhaps `@z` & `@ż` except that `ż` seems hard to enter at the REPL).
+"""
+
+@doc lazyhelp
+macro lazy(ex)
+    checkex(ex)
+    :( lazy.($(esc(ex))) )
+end
+@doc lazyhelp
+macro □(ex)
+    checkex(ex)
+    :( lazy.($(esc(ex))) )
+end
+
+@doc lazyhelp
+macro lazydot(ex)
+    checkex(ex, "@lazydot")
+    :( @. lazy($(esc(ex))) )
+end
+@doc lazyhelp
+macro ⊡(ex)
+    checkex(ex, "@lazydot")
+    :( @. lazy($(esc(ex))) )
+end
+
+using MacroTools 
+
+function checkex(ex, name="@lazy")
+    if @capture(ex, (arg__,) = val_ ) 
+        if arg[2]==:dims
+            throw(ArgumentError("$name is capturing keyword arguments, try with `; dims = $val` instead of a comma"))
+        else
+            throw(ArgumentError("$name is probably capturing capturing keyword arguments, needs a single expression"))
+        end
+    end
+    if @capture(ex, (arg_,rest__) ) 
+        throw(ArgumentError("$name is capturing more than one expression, try $name($arg) with brackets"))
+    end
+    ex
+end
+
 
 BroadcastStyle(::Type{<:BroadcastArray{<:Any,N}}) where N = LazyArrayStyle{N}()
 BroadcastStyle(L::LazyArrayStyle{N}, ::StaticArrayStyle{N}) where N = L

test/runtests.jl

@@ -164,14 +164,28 @@ end
 
 @testset "BroadcastArray" begin
     A = randn(6,6)
+
     B = BroadcastArray(exp, A)
+    B′ = @lazy exp.(A)
+    B′′ = @lazydot exp(A)
     @test Matrix(B) == exp.(A)
+    @test Matrix(B′) == exp.(A)
+    @test Matrix(B′′) == exp.(A)
 
     C = BroadcastArray(+, A, 2)
+    C′ = @lazy A .+ 2
+    C′′ = @lazydot A + 2
     @test C == A .+ 2
+    @test C′ == A .+ 2
+    @test C′′ == A .+ 2
+
     D = BroadcastArray(+, A, C)
+    D′ = @lazy A + C
+    D′′ = @lazydot A + C
     @test D == A + C
-
+    @test D′ == A + C
+    @test D′′ == A + C
+
     @test sum(B) ≈ sum(exp, A)
     @test sum(C) ≈ sum(A .+ 2)