diff --git a/src/host/linalg.jl b/src/host/linalg.jl index 36bb8582..a2a99019 100644 --- 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 index 24d6b8a6..da858608 100644 --- 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))