Implement permutedims without specializing on the permutation (#383)

Co-authored-by: Tim Besard <tim.besard@gmail.com>

Author: GiggleLiu

.buildkite/pipeline.yml

@@ -2,7 +2,7 @@ steps:
   - label: "CUDA.jl"
     plugins:
       - JuliaCI/julia#v1:
-          version: 1.6
+          version: 1.7
       - JuliaCI/julia-coverage#v1:
           codecov: true
     command: |
@@ -24,7 +24,7 @@ steps:
   - label: "oneAPI.jl"
     plugins:
       - JuliaCI/julia#v1:
-          version: 1.6
+          version: 1.7
       - JuliaCI/julia-coverage#v1:
           codecov: true
     command: |

src/host/linalg.jl

@@ -180,26 +180,35 @@ LinearAlgebra.lmul!(a::Number, B::AbstractGPUArray) = generic_lmul!(a, B)
 
 
 ## permutedims
+LinearAlgebra.permutedims!(dest::AbstractGPUArray, src::AbstractGPUArray, perm) =
+    permutedims!(dest, src, Tuple(perm))
 
 function LinearAlgebra.permutedims!(dest::AbstractGPUArray, src::AbstractGPUArray,
-                                    perm::NTuple)
+                                    perm::NTuple{N}) where N
     Base.checkdims_perm(dest, src, perm)
-    function permutedims_kernel(ctx, dest, src, ::Val{perm}) where {perm}
-        I = @cartesianidx src
-        @inbounds begin
-            J = CartesianIndex(map(i->I[i], perm))
-            dest[J] = src[I]
-        end
+
+    # get the new strides of destination tensor
+    dest_strides = ntuple(k->k==1 ? 1 : prod(i->size(dest, i), 1:k-1), N)
+    dest_strides_perm = ntuple(i->dest_strides[findfirst(==(i), perm)], N)
+
+    function permutedims_kernel(ctx, dest, src, dest_strides_perm)
+        # find the cartesian index in source tensor
+        LI = @linearidx src
+        I = @inbounds CartesianIndices(src)[LI]
+
+        # the corresponding linear index in the destination tensor
+        dest_index = map_index(I.I, dest_strides_perm)
+        @inbounds dest[dest_index] = src[LI]
         return
     end
-    gpu_call(permutedims_kernel, dest, src, Val(perm))
+    gpu_call(permutedims_kernel, dest, src, dest_strides_perm)
     return dest
 end
 
-# TODO: implementation without the memory copy
-LinearAlgebra.permutedims!(dest::AbstractGPUArray, src::AbstractGPUArray, perm) =
-    permutedims!(dest, src, Tuple(perm))
-
+# get linear index from cartesian indices and strides.
+@inline @generated function map_index(I::NTuple{N}, dest_strides::NTuple{N,T}) where {N,T}
+    Expr(:call, :+, one(T), [:(@inbounds (I[$i]-1) * dest_strides[$i]) for i in 1:N]...)
+end
 
 ## norm
 

test/testsuite/linalg.jl

@@ -15,6 +15,10 @@
         @test compare(x -> permutedims(x, (2, 1, 3)), AT, rand(Float32, 4, 5, 6))
         @test compare(x -> permutedims(x, (3, 1, 2)), AT, rand(Float32, 4, 5, 6))
         @test compare(x -> permutedims(x, [2,1,4,3]), AT, randn(ComplexF32,3,4,5,1))
+        # high dimensional tensor
+        @static if VERSION >= v"1.7"
+            @test compare(x -> permutedims(x, 18:-1:1), AT, rand(Float32, 4, [2 for _ = 2:18]...))
+        end
     end
 
     @testset "symmetric" begin