Skip to content

Commit 122a77d

Browse files
luraessluraess
authored
Merge pull request #941
Refactor ldiv
2 parents 4ba6e0b + 876515b commit 122a77d

2 files changed

Lines changed: 67 additions & 29 deletions

File tree

‎src/solver/highlevel.jl‎

Lines changed: 35 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -104,7 +104,7 @@ for (fname, elty) in (
104104
@eval begin
105105
function ormqr!(
106106
side::Char, trans::Char, A::ROCMatrix{$elty},
107-
τ::ROCVector{$elty}, C::ROCVecOrMat{$elty},
107+
τ::ROCVector{$elty}, C::StridedROCVecOrMat{$elty},
108108
)
109109
$elty <: Complex && trans == 'T' && throw(ArgumentError(
110110
"rocSOLVER.ormqr! supports only 'N' or 'C' for Complex types, " *
@@ -557,51 +557,39 @@ end
557557
LinearAlgebra.qr!(A::ROCMatrix{T}) where T = QR(geqrf!(A)...)
558558

559559
LinearAlgebra.lmul!(
560-
A::QRPackedQ{T,<:ROCArray,<:ROCArray}, B::ROCVecOrMat{T},
560+
A::QRPackedQ{T,<:ROCArray,<:ROCArray}, B::StridedROCVecOrMat{T},
561561
) where T =
562562
ormqr!('L', 'N', A.factors, A.τ, B)
563563

564564
LinearAlgebra.lmul!(
565-
adjA::Adjoint{T,<:QRPackedQ{T,<:ROCArray,<:ROCArray}},
566-
B::ROCVecOrMat{T},
565+
adjA::LinearAlgebra.AdjointQ{<:Any,<:QRPackedQ{T,<:ROCArray,<:ROCArray}},
566+
B::StridedROCVecOrMat{T},
567567
) where T <: rocBLAS.ROCBLASReal =
568-
ormqr!('L', 'T', parent(adjA).factors, parent(adjA).τ, B)
568+
ormqr!('L', 'T', adjA.Q.factors, adjA.Q.τ, B)
569569

570570
LinearAlgebra.lmul!(
571-
adjA::Adjoint{T,<:QRPackedQ{T,<:ROCArray,<:ROCArray}},
572-
B::ROCVecOrMat{T},
571+
adjA::LinearAlgebra.AdjointQ{<:Any,<:QRPackedQ{T,<:ROCArray,<:ROCArray}},
572+
B::StridedROCVecOrMat{T},
573573
) where T <: rocBLAS.ROCBLASComplex =
574-
ormqr!('L', 'C', parent(adjA).factors, parent(adjA).τ, B)
575-
576-
LinearAlgebra.lmul!(
577-
trA::Transpose{T,<:QRPackedQ{T,<:ROCArray,<:ROCArray}},
578-
B::ROCVecOrMat{T},
579-
) where T <: rocBLAS.ROCBLASFloat =
580-
ormqr!('L', 'T', parent(trA).factors, parent(trA).τ, B)
574+
ormqr!('L', 'C', adjA.Q.factors, adjA.Q.τ, B)
581575

582576
LinearAlgebra.rmul!(
583-
A::ROCVecOrMat{T},
577+
A::StridedROCVecOrMat{T},
584578
B::QRPackedQ{T,<:ROCArray,<:ROCArray},
585579
) where T <: rocBLAS.ROCBLASFloat =
586580
ormqr!('R', 'N', B.factors, B.τ, A)
587581

588582
LinearAlgebra.rmul!(
589-
A::ROCVecOrMat{T},
590-
adjB::Adjoint{<:Any,<:QRPackedQ{T,<:ROCArray,<:ROCArray}},
583+
A::StridedROCVecOrMat{T},
584+
adjB::LinearAlgebra.AdjointQ{<:Any,<:QRPackedQ{T,<:ROCArray,<:ROCArray}},
591585
) where T <: rocBLAS.ROCBLASReal =
592-
ormqr!('R', 'T', parent(adjB).factors, parent(adjB).τ, A)
586+
ormqr!('R', 'T', adjB.Q.factors, adjB.Q.τ, A)
593587

594588
LinearAlgebra.rmul!(
595-
A::ROCVecOrMat{T},
596-
adjB::Adjoint{<:Any,<:QRPackedQ{T,<:ROCArray,<:ROCArray}},
589+
A::StridedROCVecOrMat{T},
590+
adjB::LinearAlgebra.AdjointQ{<:Any,<:QRPackedQ{T,<:ROCArray,<:ROCArray}},
597591
) where T <: rocBLAS.ROCBLASComplex =
598-
ormqr!('R', 'C', parent(adjB).factors, parent(adjB).τ, A)
599-
600-
LinearAlgebra.rmul!(
601-
A::ROCVecOrMat{T},
602-
trA::Transpose{<:Any,<:QRPackedQ{T,<:ROCArray,<:ROCArray}},
603-
) where T <: rocBLAS.ROCBLASFloat =
604-
ormqr!('R', 'T', parent(trA).factors, parent(adjB).τ, A)
592+
ormqr!('R', 'C', adjB.Q.factors, adjB.Q.τ, A)
605593

606594
function LinearAlgebra.ldiv!(_qr::QR, b::ROCVector)
607595
m, n = size(_qr)
@@ -626,6 +614,25 @@ function LinearAlgebra.ldiv!(x::ROCArray, _qr::QR, b::ROCArray)
626614
return x
627615
end
628616

617+
# Override \ for GPU QR to avoid stdlib's _zeros() allocating CPU arrays.
618+
function Base.:\(F::QR{T,<:ROCMatrix,<:ROCVector}, b::ROCVector{T}) where T
619+
ldiv!(F, copy(b))
620+
end
621+
622+
function Base.:\(F::QR{<:Any,<:ROCMatrix,<:ROCVector}, b::ROCVector)
623+
factors, τ, b = copy_rocblasfloat(F.factors, F.τ, b)
624+
ldiv!(QR(factors, τ), b)
625+
end
626+
627+
function Base.:\(F::QR{T,<:ROCMatrix,<:ROCVector}, B::ROCMatrix{T}) where T
628+
ldiv!(F, copy(B))
629+
end
630+
631+
function Base.:\(F::QR{<:Any,<:ROCMatrix,<:ROCVector}, B::ROCMatrix)
632+
factors, τ, B = copy_rocblasfloat(F.factors, F.τ, B)
633+
ldiv!(QR(factors, τ), B)
634+
end
635+
629636
# LAPACK
630637

631638
for elty in (:Float32, :Float64, :ComplexF32, :ComplexF64)
@@ -639,7 +646,7 @@ for elty in (:Float32, :Float64, :ComplexF32, :ComplexF64)
639646
LinearAlgebra.LAPACK.getrf!(A::ROCMatrix{$elty}) = rocSOLVER.getrf!(A)
640647
LinearAlgebra.LAPACK.getrf!(A::ROCMatrix{$elty}, ipiv::ROCVector{Cint}) = rocSOLVER.getrf!(A, ipiv)
641648
LinearAlgebra.LAPACK.getrs!(trans::Char, A::ROCMatrix{$elty}, ipiv::ROCVector{Cint}, B::ROCVecOrMat{$elty}) = rocSOLVER.getrs!(trans, A, ipiv, B)
642-
LinearAlgebra.LAPACK.ormqr!(side::Char, trans::Char, A::ROCMatrix{$elty}, tau::ROCVector{$elty}, C::ROCVecOrMat{$elty}) = rocSOLVER.ormqr!(side, trans, A, tau, C)
649+
LinearAlgebra.LAPACK.ormqr!(side::Char, trans::Char, A::ROCMatrix{$elty}, tau::ROCVector{$elty}, C::StridedROCVecOrMat{$elty}) = rocSOLVER.ormqr!(side, trans, A, tau, C)
643650
LinearAlgebra.LAPACK.orgqr!(A::ROCMatrix{$elty}, tau::ROCVector{$elty}) = rocSOLVER.orgqr!(A, tau)
644651
LinearAlgebra.LAPACK.gebrd!(A::ROCMatrix{$elty}) = rocSOLVER.gebrd!(A)
645652
LinearAlgebra.LAPACK.gesvd!(jobu::Char, jobvt::Char, A::ROCMatrix{$elty}) = rocSOLVER.gesvd!(jobu, jobvt, A)

‎test/hip_rocarray/solver.jl‎

Lines changed: 32 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,17 @@ p = 5
5050
@test collect(dQ * dR) ≈ A
5151
@test collect(dR * dQ') ≈ (R * Q')
5252

53+
# Explicit lmul!/rmul! coverage for AdjointQ paths (Julia 1.10+).
54+
dC = ROCMatrix{elty}(I, m, m)
55+
@test collect(lmul!(dF.Q, copy(dC))) ≈ lmul!(F.Q, Matrix{elty}(I, m, m))
56+
@test collect(lmul!(dF.Q', copy(dC))) ≈ lmul!(F.Q', Matrix{elty}(I, m, m))
57+
@test collect(rmul!(copy(dC), dF.Q)) ≈ rmul!(Matrix{elty}(I, m, m), F.Q)
58+
@test collect(rmul!(copy(dC), dF.Q')) ≈ rmul!(Matrix{elty}(I, m, m), F.Q')
59+
if elty <: Real
60+
@test collect(lmul!(transpose(dF.Q), copy(dC))) ≈ lmul!(transpose(F.Q), Matrix{elty}(I, m, m))
61+
@test collect(rmul!(copy(dC), transpose(dF.Q))) ≈ rmul!(Matrix{elty}(I, m, m), transpose(F.Q))
62+
end
63+
5364
A = rand(elty, n, m)
5465
dA = ROCArray(A)
5566
dF = qr(dA)
@@ -137,6 +148,26 @@ end
137148
db = ldiv!(dy, qr(dA), dx)
138149
@test Array(db) ≈ b
139150
end
151+
152+
@testset "qr(dA) \\ b, elty = $elty" for elty in [Float32, Float64, ComplexF32, ComplexF64]
153+
A = rand(elty, m, n)
154+
b = rand(elty, m)
155+
B = rand(elty, m, p)
156+
dA, db, dB = ROCArray(A), ROCArray(b), ROCArray(B)
157+
158+
@test Array(qr(dA) \ db) ≈ qr(A) \ b
159+
@test Array(qr(dA) \ dB) ≈ qr(A) \ B
160+
end
161+
162+
@testset "qr(dA) \\ b mixed eltypes" begin
163+
A = rand(Float32, m, n)
164+
b = rand(Float64, m)
165+
B = rand(Float64, m, p)
166+
dA, db, dB = ROCArray(A), ROCArray(b), ROCArray(B)
167+
168+
@test Array(qr(dA) \ db) ≈ qr(A) \ b rtol=sqrt(eps(Float32))
169+
@test Array(qr(dA) \ dB) ≈ qr(A) \ B rtol=sqrt(eps(Float32))
170+
end
140171
end
141172

142173
@testset "geqrf! -- omgqr!" begin
@@ -204,7 +235,7 @@ end
204235
@test B ≈ collect(d_B)
205236
end
206237

207-
238+
208239
end
209240

210241
@testset "sytrf!" begin

0 commit comments

Comments
 (0)