Add conjugate option to copytri! (#413)

Author: Red-Portal

src/host/linalg.jl

@@ -32,9 +32,18 @@ end
 
 ## copy upper triangle to lower and vice versa
 
-function LinearAlgebra.copytri!(A::AbstractGPUMatrix{T}, uplo::AbstractChar) where T
+function LinearAlgebra.copytri!(A::AbstractGPUMatrix{T}, uplo::AbstractChar, conjugate::Bool=false) where T
   n = LinearAlgebra.checksquare(A)
-  if uplo == 'U'
+  if uplo == 'U' && conjugate
+      gpu_call(A) do ctx, _A
+        I = @cartesianidx _A
+        i, j = Tuple(I)
+        if j > i
+          _A[j,i] = conj(_A[i,j])
+        end
+        return
+      end
+  elseif uplo == 'U' && !conjugate
       gpu_call(A) do ctx, _A
         I = @cartesianidx _A
         i, j = Tuple(I)
@@ -43,7 +52,16 @@ function LinearAlgebra.copytri!(A::AbstractGPUMatrix{T}, uplo::AbstractChar) whe
         end
         return
       end
-  elseif uplo == 'L'
+  elseif uplo == 'L' && conjugate
+      gpu_call(A) do ctx, _A
+        I = @cartesianidx _A
+        i, j = Tuple(I)
+        if j > i
+          _A[i,j] = conj(_A[j,i])
+        end
+        return
+      end
+  elseif uplo == 'L' && !conjugate
       gpu_call(A) do ctx, _A
         I = @cartesianidx _A
         i, j = Tuple(I)

test/testsuite/linalg.jl

@@ -63,8 +63,15 @@
 
     @testset "triangular" begin
         @testset "copytri!" begin
-            @testset for uplo in ('U', 'L')
-                @test compare(x -> LinearAlgebra.copytri!(x, uplo), AT, rand(Float32, 128, 128))
+            @testset for eltya in (Float32, Float64, ComplexF32, ComplexF64), uplo in ('U', 'L'), conjugate in (true, false)
+                n = 128
+                areal = randn(n,n)/2
+                aimg  = randn(n,n)/2
+                if !(eltya in eltypes)
+                    continue
+                end
+                a = convert(Matrix{eltya}, eltya <: Complex ? complex.(areal, aimg) : areal)
+                @test compare(x -> LinearAlgebra.copytri!(x, uplo, conjugate), AT, a)
             end
         end