From c692e3e81fd6c319066638967035313ce129f472 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Thu, 29 Oct 2026 09:55:16 -0400 Subject: [PATCH] =?UTF-8?q?Add=20`permutedims!(C,=20A,=20p,=20=CE=B1,=20?= =?UTF-8?q?=CE=B2)`=20with=20a=20direct=20path=20for=20small=20arrays?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Computes `C = β * C + α * permutedims(A, p)`, as used by TensorOperations. Arrays that fit in a single block skip the layout planning, blocking and threading machinery: the dimensions are ordered by the strides of the destination, fused, and passed to the kernel directly. A nonzero `β` is handled in the same pass instead of a separate reduction kernel. Co-Authored-By: Claude Opus 5.5 --- src/Strided.jl | 1 + src/permute.jl | 102 +++++++++++++++++++++++++++++++++++++++++++++ test/gpu.jl | 4 ++ test/othertests.jl | 36 ++++++++++++++++ 4 files changed, 143 insertions(+) create mode 100644 src/permute.jl diff --git a/src/Strided.jl b/src/Strided.jl index 5b5e319..ac43231 100644 --- a/src/Strided.jl +++ b/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") diff --git a/src/permute.jl b/src/permute.jl new file mode 100644 index 0000000..4487454 --- /dev/null +++ b/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 diff --git a/test/gpu.jl b/test/gpu.jl index 9a19d3d..11a92dd 100644 --- a/test/gpu.jl +++ b/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 diff --git a/test/othertests.jl b/test/othertests.jl index 6240dea..f2034f2 100644 --- a/test/othertests.jl +++ b/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))