Add a GPU version of copytrito! (#504)

JuliaGPU · Dec 11, 2023 · effeef9 · effeef9
1 parent 6becb4f
commit effeef9
Show file tree

Hide file tree

Showing 2 changed files with 48 additions and 8 deletions.
diff --git a/src/host/linalg.jl b/src/host/linalg.jl
@@ -80,17 +80,16 @@ function Base.copyto!(A::Array{T,N}, B::Transpose{T, <:AbstractGPUArray{T,N}}) w
     copyto!(A, Transpose(Array(parent(B))))
 end
 
-
 ## copy upper triangle to lower and vice versa
 
-function LinearAlgebra.copytri!(A::AbstractGPUMatrix{T}, uplo::AbstractChar, conjugate::Bool=false) where T
+function LinearAlgebra.copytri!(A::AbstractGPUMatrix, uplo::AbstractChar, conjugate::Bool=false)
   n = LinearAlgebra.checksquare(A)
   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])
+          @inbounds _A[j,i] = conj(_A[i,j])
         end
         return
       end
@@ -99,7 +98,7 @@ function LinearAlgebra.copytri!(A::AbstractGPUMatrix{T}, uplo::AbstractChar, con
         I = @cartesianidx _A
         i, j = Tuple(I)
         if j > i
-          _A[j,i] = _A[i,j]
+          @inbounds _A[j,i] = _A[i,j]
         end
         return
       end
@@ -108,7 +107,7 @@ function LinearAlgebra.copytri!(A::AbstractGPUMatrix{T}, uplo::AbstractChar, con
         I = @cartesianidx _A
         i, j = Tuple(I)
         if j > i
-          _A[i,j] = conj(_A[j,i])
+          @inbounds _A[i,j] = conj(_A[j,i])
         end
         return
       end
@@ -117,7 +116,7 @@ function LinearAlgebra.copytri!(A::AbstractGPUMatrix{T}, uplo::AbstractChar, con
         I = @cartesianidx _A
         i, j = Tuple(I)
         if j > i
-          _A[i,j] = _A[j,i]
+          @inbounds _A[i,j] = _A[j,i]
         end
         return
       end
@@ -127,6 +126,36 @@ function LinearAlgebra.copytri!(A::AbstractGPUMatrix{T}, uplo::AbstractChar, con
   A
 end
 
+## copy a triangular part of a matrix to another matrix
+
+if isdefined(LinearAlgebra, :copytrito!)
+    function LinearAlgebra.copytrito!(B::AbstractGPUMatrix, A::AbstractGPUMatrix, uplo::AbstractChar)
+        LinearAlgebra.BLAS.chkuplo(uplo)
+        m,n = size(A)
+        m1,n1 = size(B)
+        (m1 < m || n1 < n) && throw(DimensionMismatch("B of size ($m1,$n1) should have at least the same number of rows and columns than A of size ($m,$n)"))
+        if uplo == 'U'
+            gpu_call(A, B) do ctx, _A, _B
+                I = @cartesianidx _A
+                i, j = Tuple(I)
+                if j >= i
+                    @inbounds _B[i,j] = _A[i,j]
+                end
+                return
+            end
+        else  # uplo == 'L'
+            gpu_call(A, B) do ctx, _A, _B
+                I = @cartesianidx _A
+                i, j = Tuple(I)
+                if j <= i
+                    @inbounds _B[i,j] = _A[i,j]
+                end
+                return
+            end
+        end
+        return B
+    end
+end
 
 ## triangular
 
@@ -146,7 +175,7 @@ function LinearAlgebra.tril!(A::AbstractGPUMatrix{T}, d::Integer = 0) where T
     I = @cartesianidx _A
     i, j = Tuple(I)
     if i < j - _d
-      _A[i, j] = 0
+      @inbounds _A[i, j] = zero(T)
     end
     return
   end
@@ -158,7 +187,7 @@ function LinearAlgebra.triu!(A::AbstractGPUMatrix{T}, d::Integer = 0) where T
     I = @cartesianidx _A
     i, j = Tuple(I)
     if j < i + _d
-      _A[i, j] = 0
+      @inbounds _A[i, j] = zero(T)
     end
     return
   end

diff --git a/test/testsuite/linalg.jl b/test/testsuite/linalg.jl
@@ -77,6 +77,17 @@
             end
         end
 
+        if isdefined(LinearAlgebra, :copytrito!)
+            @testset "copytrito!" begin
+                @testset for T in eltypes, uplo in ('L', 'U')
+                    n = 16
+                    A = rand(T,n,n)
+                    B = zeros(T,n,n)
+                    @test compare(copytrito!, AT, B, A, uplo)
+                end
+            end
+        end
+
         @testset "copyto! for triangular" begin
             for TR in (UpperTriangular, LowerTriangular)
                 @test compare(transpose!, AT, Array{Float32}(undef, 128, 32), rand(Float32, 32, 128))