Add permutedims!(C, A, p, α, β) with a direct path for small arrays

src/Strided.jl

@@ -53,6 +53,7 @@ end
 
 include("linalg.jl")
 include("mapreduce.jl")
+include("permute.jl")
 include("broadcast.jl")
 include("macros.jl")
 include("convert.jl")

src/permute.jl

@@ -0,0 +1,102 @@
+"""
+    permutedims!(C::StridedView, A::StridedView, p, α::Number, β::Number)
+
+Compute `C = β * C + α * permutedims(A, p)` in a single pass, without reading `C` if `β`
+is zero.
+"""
+function Base.permutedims!(
+        C::StridedView{<:Any, N}, A::StridedView{<:Any, N}, p, α::Number, β::Number
+    ) where {N}
+    perm = _permtuple(p, Val(N))
+    dims = size(C)
+    TupleTools.getindices(size(A), perm) == dims ||
+        throw(DimensionMismatch("destination size $dims does not match the permuted source"))
+    any(iszero, dims) && return C
+    sA = TupleTools.getindices(strides(A), perm)
+    if !(_isdense_cpu(C) && _isdense_cpu(A)) ||
+            prod(dims) * (sizeof(eltype(C)) + sizeof(eltype(A))) > BLOCKMEMORYSIZE
+        arrays = (C, StridedView(parent(A), dims, sA, offset(A), A.op))
+        if iszero(β)
+            _mapreduce_order!(_Scale(α), nothing, nothing, dims, arrays)
+        else
+            _mapreduce_order!(_Scale(α), +, _Scale(β), dims, arrays)
+        end
+    elseif iszero(β)
+        _permute_small!(_Scale(α), dims, (C, A), (strides(C), sA))
+    else
+        _permute_small!(_ScaleAdd(α, β), dims, (C, C, A), (strides(C), strides(C), sA))
+    end
+    return C
+end
+
+_permtuple(p::NTuple{N, Int}, ::Val{N}) where {N} =
+    isperm(p) ? p : throw(ArgumentError("$p is not a valid permutation"))
+_permtuple(p, ::Val{N}) where {N} = _permtuple(ntuple(i -> Int(p[i]), Val(N)), Val(N))
+
+# the direct kernel call is only valid for arrays in CPU memory
+@static if isdefined(Core, :Memory)
+    _isdense_cpu(a::StridedView) = parent(a) isa Union{Array, Memory}
+else
+    _isdense_cpu(a::StridedView) = parent(a) isa Array
+end
+
+struct _Scale{T}
+    α::T
+end
+(s::_Scale)(a) = s.α * a
+struct _ScaleAdd{Tα, Tβ}
+    α::Tα
+    β::Tβ
+end
+(s::_ScaleAdd)(c, a) = s.β * c + s.α * a
+
+# Arrays that fit in a block need no blocking or threading, and their loop order hardly
+# matters: loop in the memory order of the destination and call the kernel directly.
+function _permute_small!(f, dims::Dims{N}, arrays, allstrides) where {N}
+    order = _sortperm_small(ntuple(i -> dims[i] == 1 ? typemax(Int) : abs(allstrides[1][i]), Val(N)))
+    odims = TupleTools.getindices(dims, order)
+    ostrides = map(s -> TupleTools.getindices(s, order), allstrides)
+    kdims, kstrides = _fuse_compact(odims, ostrides)
+    _mapreduce_kernel!(f, nothing, nothing, kdims, kdims, arrays, kstrides, map(offset, arrays))
+    return nothing
+end
+
+# stable sort permutation for a short tuple; returns the identity without work if sorted
+function _sortperm_small(keys::NTuple{N, Int}) where {N}
+    issorted(keys) && return ntuple(identity, Val(N))
+    return invperm(ntuple(i -> _stablerank(keys, i), Val(N)))
+end
+function _stablerank(keys::NTuple{N, Int}, i::Int) where {N}
+    r = 1
+    for j in 1:N
+        r += (keys[j] < keys[i]) | ((keys[j] == keys[i]) & (j < i))
+    end
+    return r
+end
+
+# Fuse dimension `i` into `i - 1` when it is contiguous with it in all arrays, as
+# `_fusedims`, but drop the fused dimensions instead of leaving singletons in their place.
+# Whether a dimension fuses only depends on the original sizes and strides.
+function _fuse_compact(dims::NTuple{N, Int}, allstrides) where {N}
+    fuse = ntuple(i -> i > 1 && all(s -> s[i] == dims[i - 1] * s[i - 1], allstrides), Val(N))
+    heads = ntuple(k -> _kthhead(fuse, k), Val(N))
+    newdims = ntuple(k -> _headlength(dims, fuse, heads[k]), Val(N))
+    newstrides = map(s -> ntuple(k -> heads[k] == 0 ? 0 : s[heads[k]], Val(N)), allstrides)
+    return newdims, newstrides
+end
+function _kthhead(fuse::NTuple{N, Bool}, k::Int) where {N}
+    for i in 1:N
+        k -= !fuse[i]
+        k == 0 && return i
+    end
+    return 0
+end
+function _headlength(dims::NTuple{N, Int}, fuse::NTuple{N, Bool}, h::Int) where {N}
+    h == 0 && return 1
+    d = dims[h]
+    for j in (h + 1):N
+        fuse[j] || break
+        d *= dims[j]
+    end
+    return d
+end

test/gpu.jl

@@ -38,6 +38,8 @@ end
         @test compare(adjoint!, AT, A1, A2)
         @test compare(transpose!, AT, A1, A2)
         @test compare((x, y) -> permutedims!(x, y, (2, 1)), AT, A1, A2)
+        @test compare((x, y) -> permutedims!(x, y, (2, 1), 0.5, 0), AT, A1, A2)
+        @test compare((x, y) -> permutedims!(x, y, (2, 1), 0.5, 0.3), AT, A1, A2)
 
         B1 = A1[4:4:end, 1:4:end]
         B2 = A2[4:4:end, 1:4:end]
@@ -46,6 +48,8 @@ end
         @test compare(adjoint!, AT, B1, B2)
         @test compare(transpose!, AT, B1, B2)
         @test compare((x, y) -> permutedims!(x, y, (2, 1)), AT, B1, B2)
+        @test compare((x, y) -> permutedims!(x, y, (2, 1), 0.5, 0), AT, B1, B2)
+        @test compare((x, y) -> permutedims!(x, y, (2, 1), 0.5, 0.3), AT, B1, B2)
     end
 end
 

test/othertests.jl

@@ -257,6 +257,42 @@ end
     end
 end
 
+@testset "permutedims! with α and β" begin
+    for T in (Float64, ComplexF64), N in 1:6, trial in 1:10
+        sz = ntuple(_ -> rand(1:(N > 3 ? 4 : 40)), N)
+        p = Tuple(randperm(N))
+        A = randn(T, map(i -> sz[i], invperm(p)))
+        C0 = randn(T, sz)
+        for (α, β) in ((1, 0), (0.5, 0), (2.0, 1), (0.3 + 0.2im, 0.7))
+            T <: Real && α isa Complex && continue
+            C = copy(C0)
+            permutedims!(StridedView(C), StridedView(A), p, α, β)
+            @test C ≈ β * C0 + α * permutedims(A, p)
+        end
+        # strided and conjugated views
+        Ab = randn(T, map(i -> 2 * sz[i], invperm(p)))
+        Cb = randn(T, map(x -> 2x, sz))
+        As = conj(StridedView(Ab)[map(n -> 1:2:(2n), map(i -> sz[i], invperm(p)))...])
+        Cs = StridedView(Cb)[map(n -> 2:2:(2n), sz)...]
+        Cref = 0.4 * Array(Cs) + 1.5 * permutedims(Array(As), p)
+        permutedims!(Cs, As, p, 1.5, 0.4)
+        @test Array(Cs) ≈ Cref
+    end
+    # C is not read if β is zero
+    C = fill(NaN, 3, 4)
+    A = rand(4, 3)
+    permutedims!(StridedView(C), StridedView(A), (2, 1), 1, 0)
+    @test C == permutedims(A, (2, 1))
+    # arrays that do not fit in a block
+    A = randn(40, 50, 30)
+    C = randn(30, 40, 50)
+    C0 = copy(C)
+    permutedims!(StridedView(C), StridedView(A), (3, 1, 2), 2.0, 0.5)
+    @test C ≈ 0.5 * C0 + 2 * permutedims(A, (3, 1, 2))
+    @test_throws DimensionMismatch permutedims!(StridedView(zeros(2, 3)), StridedView(zeros(2, 3)), (2, 1), 1, 0)
+    @test_throws ArgumentError permutedims!(StridedView(zeros(2, 2)), StridedView(zeros(2, 2)), (1, 1), 1, 0)
+end
+
 @testset "@strided macro" begin
     @testset for T in (Float32, Float64, ComplexF32, ComplexF64)
         A1, A2, A3 = rand(T, (10,)), rand(T, (10, 10)), rand(T, (10, 10, 10))