[ANN] InplacePermutations.jl

Hello,

This started as a fun exploration, I don’t have a specific usecase for it but maybe somebody else will.

This is a very small package that introduces a permutedims!(array, perm) variant that permutes an array reusing the same underlaying memory buffer without needing auxiliary buffers/allocations.
It needs to known array sizes and permutation tuple at compile time for best performance.

using InplacePermutations

A = rand(8, 6, 2, 2)
Aorig = copy(A)                                  # keep a pristine reference for comparison

B = InplacePermutations.permutedims!(A, (2, 4, 1, 3))   # mutates A in place

B == permutedims(Aorig, (2, 4, 1, 3))        # true
size(B)                                      # (6, 2, 8, 2)

Baseline Base.permutedims! with a (36,24) array

BenchmarkTools.Trial: 10000 samples with 259 evaluations per sample.
 Range (min … max):  297.297 ns … 692.085 ns  β”Š GC (min … max): 0.00% … 0.00%
 Time  (median):     297.938 ns               β”Š GC (median):    0.00%
 Time  (mean Β± Οƒ):   302.942 ns Β±  15.778 ns  β”Š GC (mean Β± Οƒ):  0.00% Β± 0.00%

  β–ˆβ–ƒβ– ▁ ▁▂▃▁                                                    ▁
  β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‡β–‡β–ˆβ–ˆβ–‡β–†β–‡β–‡β–†β–…β–†β–†β–…β–…β–…β–…β–…β–…β–†β–…β–…β–…β–…β–…β–…β–…β–„β–†β–…β–…β–†β–†β–„β–†β–†β–†β–…β–…β–†β–…β–„β–…β–…β–…β–…β–…β–„β–…β–„β–ƒ β–ˆ
  297 ns        Histogram: log(frequency) by time        377 ns <

 Memory estimate: 0 bytes, allocs estimate: 0.

Same array with array size and permutation known at compile time

BenchmarkTools.Trial: 10000 samples with 671 evaluations per sample.
 Range (min … max):  187.841 ns … 325.633 ns  β”Š GC (min … max): 0.00% … 0.00%
 Time  (median):     188.462 ns               β”Š GC (median):    0.00%
 Time  (mean Β± Οƒ):   191.987 ns Β±   8.242 ns  β”Š GC (mean Β± Οƒ):  0.00% Β± 0.00%

  β–ˆβ–†β–‚β–β–β–β–‚β–„β–‚β–β– ▁   ▁                                             ▁
  β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‡β–‡β–ˆβ–ˆβ–‡β–‡β–‡β–‡β–†β–†β–†β–†β–‡β–†β–‡β–†β–‡β–‡β–ˆβ–‡β–‡β–‡β–‡β–‡β–‡β–‡β–‡β–‡β–‡β–†β–…β–…β–†β–…β–…β–†β–…β–†β–…β–…β–…β–…β–…β–† β–ˆ
  188 ns        Histogram: log(frequency) by time        226 ns <

 Memory estimate: 0 bytes, allocs estimate: 0.

There are some caveats:

  • Works best on small ish arrays
  • Can be compilation heavy with big arrays (tunable trade-off between compilation time and speed)
  • New sizes/perm combinations will generate a new function everytime.

More details in the repo.

PS.: currently being registered. Not married to the name, if you have opinions let me know!

Would InplacePermuteDims be a better name? Seems narrower and less suggestive of permutations in general?

For Intel/AMD processors with AVX-512 and small arrays (up to 2048 bit), one could use the fast vector indexing from SmallCollections.jl, at least if the array size is known at compile time. As an example, here is an in-place transpose:

julia> using InplacePermutations, StaticArrays, SmallCollections, Chairmarks

julia> a = rand(MMatrix{16,16,Int8});

julia> @b smalltranspose!($a), InplacePermutations.permutedims!($a, Val((2, 1)))
(11.503 ns, 66.851 ns)

julia> a = rand(MMatrix{17,15,Int8});

julia> @b smalltranspose!($a), InplacePermutations.permutedims!($a, Val((2, 1)))
(15.091 ns, 127.801 ns)

The 2048 bit limit is to avoid excessive code generation. I didn’t try to improve smalltranspose! for variable matrix sizes.

Code
function smalltranspose!(a::AbstractMatrix{T}) where T
    N = length(a)
    @assert N <= 256
    p = Ptr{FixedVector{N,T}}(pointer(a))
    v = unsafe_load(p)
    u = FixedVector{N}(transpose(LinearIndices(a))) .% UInt8
    w = @inbounds v[u]
    unsafe_store!(p, w)
    reshape(a, axes(a, 2), axes(a, 1))
end

Yeah you are probably right. Good call.

This is cool! Although this is still allocating an intermediate array w correct? In this experiment I wanted to see if I could get rid of any intermediate buffer.

The purely in-place is definitely not the best way of doing this speed wise. As you just showed leveraging SIMD is much better and by construction my approach is definitely not very amenable to that.

I’m not sure if there’s ever a situation where you can’t spare the extra memory to allocate a destination array and this package could become useful. But just in case!

No, w is an allocation-free FixedVector, analogous to SVector. (The @b macro would otherwise report an allocation.) For variable size one should probably use an allocation-free SmallVector (with a maximal length).

I see, that’s super interesting, Thanks! I’ll definitely play a little bit with that library and see if we can merge the two approaches