Here’s a crazy proposal for working with mutables and immutables more generically. Sorry it ended up a bit long.
One great feature of Julia is the ability to perform generic programming: write just one function method that works for a large variety of inputs, including some that maybe you weren’t considering when you first wrote it. So long as the input types conform to the operators and function you call, all will work well. For instance, a Taylor series evaluation function for numbers would probably work perfectly well for square matrices.
So let’s say we had a Taylor expansion evaluation (note, this might not be optimal):
function taylor_exp(x, coeffs)
tmp = one(x)
out = zero(x)
for c in coeffs
out = out + c*tmp
tmp = tmp*x
end
return out
end
Here, coeffs
can be any iterable collection, and the method works fine for x
being either a number or a square matrix (or any other type supporting +
, *
, zero
and one
), however while for numbers this is close to optimal there is a big problem if x
is of a mutable (or non-isbits immutable) type.
The problem is that every call to out = out + c*tmp
would be allocating, making my code (possibly quite signficantly) slower than it needs to be. We can fix this in certain cases, e.g. using fused mutating broadcast, so that out .= out .+ c*tmp
but then this won’t work with immutable numbers (or the immutable arrays from StaticArrays.jl). I admit, there is also the problem of c*tmp
being allocating and tmp = tmp*x
being allocating (I have a better example below) - but this highlights a bigger, structural problem, that in certain cases I can’t write generic and performant code for both immutable and mutable types. Adding to the pain is that it is not easy to write just two methods - one for heap-allocated, mutable objects and another for stack-allocated, mutable objects - because it is difficult to catch the distinction at dispatch (I would have to write a bunch of methods for types that I know are immutable, and another bunch of methods for types that I know are mutable, or some such madness).
So how to fix this? Let’s take a better example than the above, for which (at work) I have already written optimal code for both cases. Say we have a collection of 3-vectors representing points in space, and we want to find the 3x3 “covariance” matrix for how they are scattered in space (say, because its Eigenvalues tell us something interesting about the distribution of points - whether they are linear, planar or spread out in 3D). If we didn’t have to worry about allocation, then I would write this code (for simplicity I also assume the mean of the points
is at the origin and ignore normalization):
function covariance_immutable(points)
cov = zero(first(points)*first(points)')
for p in points
cov = cov + p*p'
end
return cov
end
In fact, this code works perfectly well if points
is a collection of static vectors from StaticArrays.jl, such as a Vector{SVector{3,Float64}}
, and this function is much, much slower using Base.Array
because it reallocates cov
many, many times in the loop (and it is just a small array, so allocation and garbage collection costs dominate the actual computations).
To remove the allocations we can use:
function covariance_mutable(points)
cov = zero(first(points)*first(points)')
for p in points
cov .= cov .+ p.*p'
end
return cov
end
Let’s assume for a moment that we have inline non-isbits immutables so that p'
is a non-allocating view of p
. Then broadcast loop fusion would eliminate all the allocation within the loop, and only the output matrix cov
would be heap-allocated in the first line. Then this function becomes competitive to the above when p
is a Base.Vector
or StaticArrays.MVector
(or more importantly, when cov
is a Base.Matrix
or StaticArrays.MMatrix
), but it no longer works for SVector
.
While I could implement the first method as default and implement the second wherever I know it works, it would be preferable if I could just write one method. Here I propose one idiom for assigments which are overwriting for mutables and work for immutables, as an alternative to x = y
or x .= y
(or the new x <- y
):
- Introduce
:=
such thatx := y
is lowered tox = Base.set(x,y)
, where types can implement setting as they need to (via some function likeBase.set!
). The idea ofx := y
is that the memory allocated tox
is overwritten in place (whether that is the stack or the heap) - if that is not possible then an error would be thrown (e.g. arrays of incompatible size). - For immutable
x
we would simply haveset(x,y) = y
. - For mutable
x
we would do the equivalent ofset(x,y) = (set!(x, y); return x)
(but it would be up to the types to implement this correctly, e.g. composite types/structs could do this too, not just arrays, but the default behavior ofset!
could bebroadcast!
or maybesetindex!
). - We could go futher and implement
broadcast_set
and do dot-call fusion with:=
, e.g.x := y .+ z
.
Then we could end up with:
function covariance_generic(points)juliua
cov = zero(first(points)*first(points)')
for p in points
cov := cov .+ p.*p'
end
return cov
end
which works well in both cases!
Here’s a possible implementation of set
:
@generated function Base.set(x, y)
if x.mutable
return :(Base.set!(x,y); return x)
else
return :y
end
end
where it is the resposibility of type owners to define a specialized method for Base.set!
if their type is mutable (maybe with a fallback of broadcast!
or setindex!
(or vice-versa), or some generated function for structs).
(Please note: I don’t think the lack of genericism with mutable/immutable objects is limited to just geometry and static arrays as I’ve focussed on in the examples - I’m bringing this up because I feel that this could be useful more broadly in the language. It is also possible that better compiler technology could re-use allocated memory in cases like shown in the example, making the first “functional” approach work well also for mutable structs and arrays, but that sounds tricky when e.g. there are no static guarantees as to the size of the array being compatible).