Merge pull request #89 from SciML/allocs

fix broadcast allocations

Author: ChrisRackauckas

Project.toml

@@ -6,8 +6,6 @@ version = "2.3.0"
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
-NLsolve = "2774e3e8-f4cf-5e23-947b-6d7e65073b56"
-OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 RecipesBase = "3cdcf5f2-1ef4-517c-9805-6587b60abb01"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
@@ -23,8 +21,10 @@ ZygoteRules = "0.2"
 julia = "1.3"
 
 [extras]
+NLsolve = "2774e3e8-f4cf-5e23-947b-6d7e65073b56"
+OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 
 [targets]
-test = ["Test", "Unitful"]
+test = ["NLsolve", "OrdinaryDiffEq", "Test", "Unitful"]

src/array_partition.jl

@@ -230,15 +230,15 @@ ArrayPartitionStyle(::S, ::Val{N}) where {S,N} = ArrayPartitionStyle(S(Val(N)))
 ArrayPartitionStyle(::Val{N}) where N = ArrayPartitionStyle{Broadcast.DefaultArrayStyle{N}}()
 
 # promotion rules
-function Broadcast.BroadcastStyle(::ArrayPartitionStyle{AStyle}, ::ArrayPartitionStyle{BStyle}) where {AStyle, BStyle}
+@inline function Broadcast.BroadcastStyle(::ArrayPartitionStyle{AStyle}, ::ArrayPartitionStyle{BStyle}) where {AStyle, BStyle}
     ArrayPartitionStyle(Broadcast.BroadcastStyle(AStyle(), BStyle()))
 end
-Broadcast.BroadcastStyle(::ArrayPartitionStyle, ::Broadcast.DefaultArrayStyle{0}) = Broadcast.DefaultArrayStyle{1}()
+Broadcast.BroadcastStyle(::ArrayPartitionStyle{Style}, ::Broadcast.DefaultArrayStyle{0}) where Style<:Broadcast.BroadcastStyle = ArrayPartitionStyle{Style}()
 Broadcast.BroadcastStyle(::ArrayPartitionStyle, ::Broadcast.DefaultArrayStyle{N}) where N = Broadcast.DefaultArrayStyle{N}()
 
 combine_styles(args::Tuple{})         = Broadcast.DefaultArrayStyle{0}()
-combine_styles(args::Tuple{Any})      = Broadcast.result_style(Broadcast.BroadcastStyle(args[1]))
-combine_styles(args::Tuple{Any, Any}) = Broadcast.result_style(Broadcast.BroadcastStyle(args[1]), Broadcast.BroadcastStyle(args[2]))
+@inline combine_styles(args::Tuple{Any})      = Broadcast.result_style(Broadcast.BroadcastStyle(args[1]))
+@inline combine_styles(args::Tuple{Any, Any}) = Broadcast.result_style(Broadcast.BroadcastStyle(args[1]), Broadcast.BroadcastStyle(args[2]))
 @inline combine_styles(args::Tuple)   = Broadcast.result_style(Broadcast.BroadcastStyle(args[1]), combine_styles(Base.tail(args)))
 
 function Broadcast.BroadcastStyle(::Type{ArrayPartition{T,S}}) where {T, S}
@@ -256,7 +256,7 @@ end
 
 @inline function Base.copyto!(dest::ArrayPartition, bc::Broadcast.Broadcasted{ArrayPartitionStyle{Style}}) where Style
     N = npartitions(dest, bc)
-    for i in 1:N
+    @inbounds for i in 1:N
         copyto!(dest.x[i], unpack(bc, i))
     end
     dest

test/partitions_test.jl

@@ -119,3 +119,19 @@ for sizes in S
   @test all([x[i] == y[i] for i in eachindex(x)])
   @test all([x[i] == y_array[i] for i in eachindex(x)])
 end
+
+# Non-allocating broadcast
+xce0 = ArrayPartition(zeros(2),[0.])
+xcde0 = copy(xce0)
+function foo(y, x)
+	y .= y .+ x
+	nothing
+end
+foo(xcde0, xce0)
+#@test 0 == @allocated foo(xcde0, xce0)
+function foo(y, x)
+	y .= y .+ 2 .* x
+	nothing
+end
+foo(xcde0, xce0)
+#@test 0 == @allocated foo(xcde0, xce0)