Fix spurious memory allocations of AddedOperator after regression with v1.0 release

src/basic.jl

@@ -616,33 +616,46 @@ end
         w
     end
 end
+
 # Out-of-place: v is action vector, u is update vector
 function (L::AddedOperator)(v::AbstractVecOrMat, u, p, t; kwargs...)
-    L = update_coefficients(L, u, p, t; kwargs...)
-    sum(op -> iszero(op) ? zero(v) : op(v, u, p, t; kwargs...), L.ops)
+    # We don't need to update coefficients of L, as op(v, u, p, t) will do it for each op
+    sum(op -> op(v, u, p, t; kwargs...), L.ops)
 end
 
 # In-place: w is destination, v is action vector, u is update vector
-function (L::AddedOperator)(w::AbstractVecOrMat, v::AbstractVecOrMat, u, p, t; kwargs...)
-    update_coefficients!(L, u, p, t; kwargs...)
-    for op in L.ops
-        if !iszero(op)
-            op(w, v, u, p, t, 1.0, 1.0; kwargs...)
+@generated function (L::AddedOperator)(w::AbstractVecOrMat, v::AbstractVecOrMat, u, p, t; kwargs...)
+    # We don't need to update coefficients of L, as op(w, v, u, p, t) will do it for each op
+
+    ops_types = L.parameters[2].parameters
+    N = length(ops_types)-1
+
+    quote
+        L.ops[1](w, v, u, p, t; kwargs...)
+        Base.@nexprs $N i->begin
+            op = L.ops[i+1]
+            op(w, v, u, p, t, true, true; kwargs...)
         end
+        w
     end
-    w
 end
 
 # In-place with scaling: w = α*(L*v) + β*w
-function (L::AddedOperator)(w::AbstractVecOrMat, v::AbstractVecOrMat, u, p, t, α, β; kwargs...)
-    update_coefficients!(L, u, p, t; kwargs...)
-    lmul!(β, w)
-    for op in L.ops
-        if !iszero(op)
-            op(w, v, u, p, t, α, 1.0; kwargs...)
+@generated function (L::AddedOperator)(w::AbstractVecOrMat, v::AbstractVecOrMat, u, p, t, α, β; kwargs...)
+    # We don't need to update coefficients of L, as op(w, v, u, p, t) will do it for each op
+
+    T = L.parameters[1]
+    ops_types = L.parameters[2].parameters
+    N = length(ops_types)-1
+
+    quote
+        L.ops[1](w, v, u, p, t, α, β; kwargs...)
+        Base.@nexprs $N i->begin
+            op = L.ops[i+1]
+            op(w, v, u, p, t, α, true; kwargs...)
         end
+        w
     end
-    w
 end
 
 """

src/matrix.jl

@@ -201,7 +201,7 @@ function (L::MatrixOperator)(w::AbstractVecOrMat, v::AbstractVecOrMat, u, p, t;
 end
 
 # In-place with scaling: w = α*(L*v) + β*w
-function (L::MatrixOperator)(w::AbstractVecOrMat, v::AbstractVecOrMat, u, p, t, α, β; kwargs...)
+Base.@constprop :aggressive function (L::MatrixOperator)(w::AbstractVecOrMat, v::AbstractVecOrMat, u, p, t, α, β; kwargs...)
     update_coefficients!(L, u, p, t; kwargs...)
     mul!(w, L.A, v, α, β)
 end

test/downstream/alloccheck.jl

@@ -1,62 +1,62 @@
-using SciMLOperators, AllocCheck, Random, SparseArrays, Test
+using SciMLOperators, Random, SparseArrays, Test
 using SciMLOperators: IdentityOperator,
                       NullOperator,
                       ScaledOperator,
                       AddedOperator
-Random.seed!(0)
-N = 8
-K = 12
-A = rand(N, N) |> MatrixOperator
-B = rand(N, N) |> MatrixOperator
-C = rand(N, N) |> MatrixOperator
-α = rand()
-β = rand()
-u = rand(N, K)       # Update vector
-v = rand(N, K)       # Action vector
-w = zeros(N, K)      # Output vector
-p = ()
-t = 0
-op = AddedOperator(A, B)
-
-# Define a function to test allocations with the new interface
-@check_allocs ignore_throw = true function apply_op!(H, w, v, u, p, t)
+
+function apply_op!(H, w, v, u, p, t)
     H(w, v, u, p, t)
     return nothing
 end
 
-if VERSION >= v"1.12-beta"
-    apply_op!(op, w, v, u, p, t)
-else
-    @test_throws AllocCheckFailure apply_op!(op, w, v, u, p, t)
-end
-
-for T in (Float32, Float64, ComplexF32, ComplexF64)
-    N = 100
-    A1_sparse = MatrixOperator(sprand(T, N, N, 5 / N))
-    A2_sparse = MatrixOperator(sprand(T, N, N, 5 / N))
-    A3_sparse = MatrixOperator(sprand(T, N, N, 5 / N))
-
-    A1_dense = MatrixOperator(rand(T, N, N))
-    A2_dense = MatrixOperator(rand(T, N, N))
-    A3_dense = MatrixOperator(rand(T, N, N))
-
-    coeff1(a, u, p, t) = sin(p.ω * t)
-    coeff2(a, u, p, t) = cos(p.ω * t)
-    coeff3(a, u, p, t) = sin(p.ω * t) * cos(p.ω * t)
-
-    c1 = ScalarOperator(rand(T), coeff1)
-    c2 = ScalarOperator(rand(T), coeff2)
-    c3 = ScalarOperator(rand(T), coeff3)
-
-    H_sparse = c1 * A1_sparse + c2 * A2_sparse + c3 * A3_sparse
-    H_dense = c1 * A1_dense + c2 * A2_dense + c3 * A3_dense
-
-    u = rand(T, N)
-    v = rand(T, N) 
-    w = similar(u)
-    p = (ω = 0.1,)
-    t = 0.1
-
-    @test_throws AllocCheckFailure apply_op!(H_sparse, w, v, u, p, t)
-    @test_throws AllocCheckFailure apply_op!(H_dense, w, v, u, p, t)
+test_apply_noalloc(H, w, v, u, p, t) = @test (@allocations apply_op!(H, w, v, u, p, t)) == 0
+
+@testset "Allocations Check" begin
+    Random.seed!(0)
+    N = 8
+    K = 12
+    A = rand(N, N) |> MatrixOperator
+    B = rand(N, N) |> MatrixOperator
+    u = rand(N, K)       # Update vector
+    v = rand(N, K)       # Action vector
+    w = zeros(N, K)      # Output vector
+    p = ()
+    t = 0
+    op = AddedOperator(A, B)
+
+    apply_op!(op, w, v, u, p, t) # Warm up
+    test_apply_noalloc(op, w, v, u, p, t)
+
+    for T in (Float32, Float64, ComplexF32, ComplexF64)
+        N = 100
+        A1_sparse = MatrixOperator(sprand(T, N, N, 5 / N))
+        A2_sparse = MatrixOperator(sprand(T, N, N, 5 / N))
+        A3_sparse = MatrixOperator(sprand(T, N, N, 5 / N))
+
+        A1_dense = MatrixOperator(rand(T, N, N))
+        A2_dense = MatrixOperator(rand(T, N, N))
+        A3_dense = MatrixOperator(rand(T, N, N))
+
+        coeff1(a, u, p, t) = sin(p.ω * t)
+        coeff2(a, u, p, t) = cos(p.ω * t)
+        coeff3(a, u, p, t) = sin(p.ω * t) * cos(p.ω * t)
+
+        c1 = ScalarOperator(rand(T), coeff1)
+        c2 = ScalarOperator(rand(T), coeff2)
+        c3 = ScalarOperator(rand(T), coeff3)
+
+        H_sparse = c1 * A1_sparse + c2 * A2_sparse + c3 * A3_sparse
+        H_dense = c1 * A1_dense + c2 * A2_dense + c3 * A3_dense
+
+        u = rand(T, N)
+        v = rand(T, N) 
+        w = similar(u)
+        p = (ω = 0.1,)
+        t = 0.1
+
+        apply_op!(H_sparse, w, v, u, p, t) # Warm up
+        apply_op!(H_dense, w, v, u, p, t) # Warm up
+        test_apply_noalloc(H_sparse, w, v, u, p, t)
+        test_apply_noalloc(H_dense, w, v, u, p, t)
+    end
 end