Generic mapping over abstractarrays

Problem: write f(a::AbstractArray) = ... such that

  • the output is an abstractarray of the same type as a
  • out[i, j, ...] = a[i, j, ...] + i*j*k*...
  • f is non-allocating with StaticArrays. But you’re not allowed to import StaticArrays.

Is there a solution?

you are not allowed to allocate out?

With a normal Array, of course there will be an allocation. But with SMatrix, there shouldn’t be.

The point is that with just regular arrays, this is easy:

f(a::AbstractMatrix) = [a[i, j]+i*j for i in 1:size(a, 1), for j in 1:size(a, 2)]

but this will return a Matrix when passed an SMatrix, which is not what I want. And, of course, it doesn’t support arbitrarily-many dimensions.

There’s map(x->x+2, arr), which works nicely with both Array and static arrays. Is there a map_indices(ind->f(ind), arr), which constructs a similar array to arr?

What about map!(f, similar(a), arr)?

That turns a stack-allocated SVector input into a heap-allocated MVector output.

You’re asking to create a fixed-size array similar to a given array (but distinct from the given array) at runtime, if I understand correctly? Without preallocating or having already created a SArray and passing that into the function, I don’t think it’s possible.

But that’s fine. I want f([1 2; 3 4]) == [2 4; 5 8] and f(@SMatrix([1 2; 3 4])) === @SMatrix [2 4; 5 8]. Without having a dependency on StaticArrays

I just want a function that performs generic mapping over arrays (which map does), but also gives me the cartesian index of each element.

There is not a solution. The solution is to treat static arrays and scalars similarly (the immutables), and other AbstractArrays similarly (the mutables). If you then have one generic code for immutable and one for mutables then you can cover everything I’ve ran into so far. This is how DiffEq supports both BTW: it doubles each internal time step algorithm for these two forms.

I don’t clearly understand, if scalars are immutables, yet I can reassign it with another value. If an (fixed-sized) array is immutable, why i cannot reassign its elements with another values? Is there some confusion between immutability in terms of fized size, and in terms of array elements? And is there some link between heap /stack and mutable/immutable?

You can reassign it but not mutate it. With scalars, you cannot write a function change!(a) and expect it to change the value of a since it passes the value. Instead, you’d need to do a = change(a) for immutables. While for mutables in many cases change!(a) is more performant, so if you want to optimize both cases you need both forms. You could do things like ensure your mutating functions always return the value, so then the calling code could do a = change!(a), but then change! itself still needs at least two code paths to handle the very different behaviors that’s going on between the mutable and immutable arrays.

Thanks! So, is there any problems with the absence of mutable, fixed-sized, stack-allocated array?
Also, is it correct that mutable = heap-allocated?

julia> using BenchmarkTools

julia> f(X) = map((x,I) -> x + prod(Tuple(I)), X, CartesianIndices(X))
f (generic function with 1 method)

julia> @btime f($([1. 2; 3 4]))
  81.636 ns (3 allocations: 160 bytes)
2×2 Array{Float64,2}:
 2.0  4.0
 5.0  8.0

julia> using StaticArrays

julia> @btime f($(@SArray [1. 2; 3 4]))
  4.287 ns (0 allocations: 0 bytes)
2×2 SArray{Tuple{2,2},Float64,2,4}:
 2.0  4.0
 5.0  8.0
6 Likes

This is mindblowing to me:

julia> @code_llvm f(@SArray [1. 2; 3 4])

define void @julia_f_35288({ [4 x double] }* noalias nocapture sret, { [4 x double] } addrspace(11)* nocapture nonnull readonly dereferenceable(32)) {
top:
 %2 = bitcast { [4 x double] } addrspace(11)* %1 to <4 x double> addrspace(11)*
 %3 = load <4 x double>, <4 x double> addrspace(11)* %2, align 8
 %4 = fadd <4 x double> %3, <double 1.000000e+00, double 2.000000e+00, double 2.000000e+00, double 4.000000e+00>
 %5 = bitcast { [4 x double] }* %0 to <4 x double>*
 store <4 x double> %4, <4 x double>* %5, align 8
 ret void
}
1 Like

can you explain a bit what does this mean? StaticArray is smart/hacky?

No, this exists. It’s MArray in StaticArrays.jl.

In Julia, yes, at least for now. That’s not always true in general. C and Fortran allow for stack-allocated mutable buffers for example. It’s just not available in Julia quite yet. And it’s not always true in Julia too, since some escape analysis can optimize some cases IIRC. And as the compiler continues to get smarter, this connection between mutability and heap allocations will only lesson.

1 Like

No, it’s that the <4 x double> operations are SIMD operations computing 4 floating points at a time. If you wanted to hand write the assembly and knew to do SIMD, you’d write pretty much that, yet Julia does it automatically. That doesn’t mean that it’s the most optimized LLVM IR that’s possible, but it does mean it generates better LLVM IR than I could write myself, which makes me bow down to the compiler gods.

2 Likes

Fantastic! Thank you.