Implement MutableArithmetics (#331)

* Implement MutableArithmetics

* Add usage example in README

Author: blegat

Project.toml

@@ -7,10 +7,12 @@ version = "2.0.9"
 [deps]
 Intervals = "d8418881-c3e1-53bb-8760-2df7ec849ed5"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+MutableArithmetics = "d8a4904e-b15c-11e9-3269-09a3773c0cb0"
 RecipesBase = "3cdcf5f2-1ef4-517c-9805-6587b60abb01"
 
 [compat]
 Intervals = "0.5, 1.0, 1.3"
+MutableArithmetics = "0.2.15"
 RecipesBase = "0.7, 0.8, 1"
 julia = "1"
 

README.md

@@ -106,6 +106,73 @@ julia> p + q
 ERROR: Polynomials must have same variable.
 ```
 
+#### Construction and Evaluation
+
+While polynomials of type `Polynomial` are mutable objects, operations such as
+`+`, `-`, `*`, always create new polynomials without modifying its arguments.
+The time needed for these allocations and copies of the polynomial coefficients
+may be noticeable in some use cases. This is amplified when the coefficients
+are for instance `BigInt` or `BigFloat` which are mutable themself.
+This can be avoided by modifying existing polynomials to contain the result
+of the operation using the [MutableArithmetics (MA) API](https://github.com/jump-dev/MutableArithmetics.jl).
+Consider for instance the following arrays of polynomials
+```julia
+using Polynomials
+d, m, n = 30, 20, 20
+p(d) = Polynomial(big.(1:d))
+A = [p(d) for i in 1:m, j in 1:n]
+b = [p(d) for i in 1:n]
+```
+In this case, the arrays are mutable objects for which the elements are mutable
+polynomials which have mutable coefficients (`BigInt`s).
+These three nested levels of mutable objects communicate with the MA
+API in order to reduce allocation.
+Calling `A * b` requires approximately 40 MiB due to 2 M allocations
+as it does not exploit any mutability. Using
+```julia
+using MutableArithmetics
+const MA = MutableArithmetics
+MA.operate(*, A, b)
+```
+exploits the mutability and hence only allocate approximately 70 KiB due to 4 k
+allocations. If the resulting vector is already allocated, e.g.,
+```julia
+z(d) = Polynomial([zero(BigInt) for i in 1:d])
+c = [z(2d - 1) for i in 1:m]
+```
+then we can exploit its mutability with
+```julia
+MA.mutable_operate!(MA.add_mul, c, A, b)
+```
+to reduce the allocation down to 48 bytes due to 3 allocations. These remaining
+allocations are due to the `BigInt` buffer used to store the result of
+intermediate multiplications. This buffer can be preallocated with
+```julia
+buffer = MA.buffer_for(MA.add_mul, typeof(c), typeof(A), typeof(b))
+MA.mutable_buffered_operate!(buffer, MA.add_mul, c, A, b)
+```
+then the second line is allocation-free.
+
+The `MA.@rewrite` macro rewrite an expression into an equivalent code that
+exploit the mutability of the intermediate results.
+For instance
+```julia
+MA.@rewrite(A1 * b1 + A2 * b2)
+```
+is rewritten into
+```julia
+c = MA.operate!(MA.add_mul, MA.Zero(), A1, b1)
+MA.operate!(MA.add_mul, c, A2, b2)
+```
+which is equivalent to
+```julia
+c = MA.operate(*, A1, b1)
+MA.mutable_operate!(MA.add_mul, c, A2, b2)
+```
+
+*Note that currently, only the `Polynomial` implements the API and it only
+implements part of it.*
+
 ### Integrals and Derivatives
 
 Integrate the polynomial `p` term by term, optionally adding a constant

src/Polynomials.jl

@@ -19,6 +19,7 @@ include("common.jl")
 
 # Polynomials
 include("polynomials/standard-basis.jl")
+include("polynomials/mutable-arithmetics.jl")
 include("polynomials/Polynomial.jl")
 include("polynomials/ImmutablePolynomial.jl")
 include("polynomials/SparsePolynomial.jl")

src/polynomials/Polynomial.jl

@@ -50,6 +50,7 @@ struct Polynomial{T <: Number, X} <: StandardBasisPolynomial{T, X}
 end
 
 @register Polynomial
+@register_mutable_arithmetic Polynomial
 
 
 

src/polynomials/mutable-arithmetics.jl

@@ -0,0 +1,67 @@
+using MutableArithmetics
+const MA = MutableArithmetics
+
+function _resize_zeros!(v::Vector, new_len)
+    old_len = length(v)
+    if old_len < new_len
+        resize!(v, new_len)
+        for i in (old_len + 1):new_len
+            v[i] = zero(eltype(v))
+        end
+    end
+end
+
+"""
+    add_conv(out::Vector{T}, E::Vector{T}, k::Vector{T})
+Returns the vector `out + fastconv(E, k)`. Note that only
+`MA.mutable_buffered_operate!` is implemented.
+"""
+function add_conv end
+
+# The buffer we need is the buffer needed by the `MA.add_mul` operation.
+# For instance, `BigInt`s need a `BigInt` buffer to store `E[x] * k[i]` before
+# adding it to `out[j]`.
+function MA.buffer_for(::typeof(add_conv), ::Type{Vector{T}}, ::Type{Vector{T}}, ::Type{Vector{T}}) where {T}
+    return MA.buffer_for(MA.add_mul, T, T, T)
+end
+function MA.mutable_buffered_operate!(buffer, ::typeof(add_conv), out::Vector{T}, E::Vector{T}, k::Vector{T}) where {T}
+    for x in eachindex(E)
+        for i in eachindex(k)
+            j = x + i - 1
+            out[j] = MA.buffered_operate!(buffer, MA.add_mul, out[j], E[x], k[i])
+        end
+    end
+    return out
+end
+
+"""
+    @register_mutable_arithmetic
+Register polynomial type (with vector based backend) to work with MutableArithmetics
+"""
+macro register_mutable_arithmetic(name)
+    poly = esc(name)
+    quote
+        MA.mutability(::Type{<:$poly}) = MA.IsMutable()
+
+        function MA.promote_operation(::Union{typeof(+), typeof(*)},
+                                      ::Type{$poly{S,X}}, ::Type{$poly{T,X}}) where {S,T,X}
+            R = promote_type(S,T)
+            return $poly{R,X}
+        end
+
+        function MA.buffer_for(::typeof(MA.add_mul),
+                               ::Type{<:$poly{T,X}},
+                               ::Type{<:$poly{T,X}}, ::Type{<:$poly{T,X}}) where {T,X}
+            V = Vector{T}
+            return MA.buffer_for(add_conv, V, V, V)
+        end
+
+        function MA.mutable_buffered_operate!(buffer, ::typeof(MA.add_mul),
+                                              p::$poly, q::$poly, r::$poly)
+            ps, qs, rs = coeffs(p), coeffs(q), coeffs(r)
+            _resize_zeros!(ps, length(qs) + length(rs) - 1)
+            MA.mutable_buffered_operate!(buffer, add_conv, ps, qs, rs)
+            return p
+        end
+    end
+end

test/Project.toml

@@ -1,5 +1,6 @@
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+MutableArithmetics = "d8a4904e-b15c-11e9-3269-09a3773c0cb0"
 OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
 RecipesBase = "3cdcf5f2-1ef4-517c-9805-6587b60abb01"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"

test/StandardBasis.jl

@@ -1209,3 +1209,27 @@ end
         end
     end
 end
+
+using MutableArithmetics
+const MA = MutableArithmetics
+
+function alloc_test(f, n)
+    f() # compile
+    @test n == @allocated f()
+end
+
+@testset "Mutable arithmetics" begin
+    d = m = n = 4
+    p(d) = Polynomial(big.(1:d))
+    z(d) = Polynomial([zero(BigInt) for i in 1:d])
+    A = [p(d) for i in 1:m, j in 1:n]
+    b = [p(d) for i in 1:n]
+    c = [z(2d - 1) for i in 1:m]
+    buffer = MA.buffer_for(MA.add_mul, typeof(c), typeof(A), typeof(b))
+    @test buffer isa BigInt
+    c = [z(2d - 1) for i in 1:m]
+    MA.mutable_buffered_operate!(buffer, MA.add_mul, c, A, b)
+    @test c == A * b
+    @test c == MA.operate(*, A, b)
+    @test 0 == @allocated MA.mutable_buffered_operate!(buffer, MA.add_mul, c, A, b)
+end