centralize QR/LQ gaugefixing (#203)

Author: lkdvos

ext/MatrixAlgebraKitGenericLinearAlgebraExt.jl

@@ -3,7 +3,7 @@ module MatrixAlgebraKitGenericLinearAlgebraExt
 using MatrixAlgebraKit
 using MatrixAlgebraKit: sign_safe, check_input, diagview, gaugefix!, one!, zero!, default_fixgauge
 using MatrixAlgebraKit: GLA, Driver
-import MatrixAlgebraKit: gesvd!, heev!
+import MatrixAlgebraKit: qr_householder!, qr_null_householder!, gesvd!, heev!
 using GenericLinearAlgebra: svd!, svdvals!, eigen!, eigvals!, Hermitian, qr!
 using LinearAlgebra: I, Diagonal, lmul!
 
@@ -53,7 +53,7 @@ function heev!(::GLA, A::AbstractMatrix, Dd::AbstractVector, V::AbstractMatrix;
     return Dd, V
 end
 
-function MatrixAlgebraKit.qr_householder!(
+function qr_householder!(
         driver::MatrixAlgebraKit.GLA, A::AbstractMatrix, Q::AbstractMatrix, R::AbstractMatrix;
         positive::Bool = true, pivoted::Bool = false, blocksize::Int = 0
     )
@@ -68,36 +68,20 @@ function MatrixAlgebraKit.qr_householder!(
 
     # compute QR
     Q̃, R̃ = qr!(A)
-    lmul!(Q̃, MatrixAlgebraKit.one!(Q))
-
-    if positive
-        @inbounds for j in 1:minmn
-            s = sign_safe(R̃[j, j])
-            @simd for i in 1:m
-                Q[i, j] *= s
-            end
-        end
-    end
+    lmul!(Q̃, one!(Q))
 
     if computeR
-        if positive
-            @inbounds for j in n:-1:1
-                @simd for i in 1:min(minmn, j)
-                    R[i, j] = R̃[i, j] * conj(sign_safe(R̃[i, i]))
-                end
-                @simd for i in (min(minmn, j) + 1):size(R, 1)
-                    R[i, j] = zero(eltype(R))
-                end
-            end
-        else
-            R[1:minmn, :] .= R̃
-            MatrixAlgebraKit.zero!(@view(R[(minmn + 1):end, :]))
-        end
+        copyto!(view(R, 1:minmn, :), R̃)
+        zero!(view(R, (minmn + 1):size(R, 1), :))
+        positive && gaugefix!(qr_householder!, Q, R, diagview(R̃))
+    elseif positive
+        gaugefix!(qr_householder!, Q, nothing, diagview(R̃))
     end
+
     return Q, R
 end
 
-function MatrixAlgebraKit.qr_null_householder!(
+function qr_null_householder!(
         driver::MatrixAlgebraKit.GLA, A::AbstractMatrix, N::AbstractMatrix;
         positive::Bool = true, pivoted::Bool = false, blocksize::Int = 0
     )

src/MatrixAlgebraKit.jl

@@ -99,8 +99,6 @@ include("interface/schur.jl")
 include("interface/polar.jl")
 include("interface/orthnull.jl")
 
-include("common/gauge.jl") # needs to be defined after the functions are
-
 include("implementations/projections.jl")
 include("implementations/truncation.jl")
 include("implementations/qr.jl")
@@ -113,6 +111,8 @@ include("implementations/schur.jl")
 include("implementations/polar.jl")
 include("implementations/orthnull.jl")
 
+include("common/gauge.jl") # needs to be defined after the functions are
+
 include("pullbacks/qr.jl")
 include("pullbacks/lq.jl")
 include("pullbacks/eig.jl")

src/common/gauge.jl

@@ -12,6 +12,28 @@ is real and positive.
 _argmaxabs(x) = reduce(_largest, x; init = zero(eltype(x)))
 _largest(x, y) = abs(x) < abs(y) ? y : x
 
+function gaugefix!(::typeof(qr_householder!), Q, R, Rd)
+    ax = Base.OneTo(length(Rd))
+    Qf = view(Q, axes(Q, 1), ax)
+    Qf .*= sign_safe.(transpose(Rd))
+    if !isnothing(R)
+        Rf = view(R, ax, axes(R, 2))
+        Rf .*= conj.(sign_safe.(Rd))
+    end
+    return Q, R
+end
+
+function gaugefix!(::typeof(lq_householder!), L, Q, Ld)
+    ax = Base.OneTo(length(Ld))
+    Qf = view(Q, ax, axes(Q, 2))
+    Qf .*= sign_safe.(Ld)
+    if !isnothing(L)
+        Lf = view(L, axes(L, 1), ax)
+        Lf .*= conj.(sign_safe.(transpose(Ld)))
+    end
+    return L, Q
+end
+
 function gaugefix!(::Union{typeof(eig_full!), typeof(eigh_full!), typeof(gen_eig_full!)}, V::AbstractMatrix)
     for j in axes(V, 2)
         v = view(V, :, j)

src/implementations/lq.jl

@@ -163,27 +163,14 @@ function lq_householder!(
         end
     end
 
-    if positive # already fix Q even if we do not need L
-        @inbounds for j in 1:n
-            @simd for i in 1:minmn
-                s = sign_safe(A[i, i])
-                Q[i, j] *= s
-            end
-        end
-    end
-
     if computeL
         L̃ = lowertriangular!(view(A, axes(L)...))
-        if positive
-            @inbounds for j in 1:minmn
-                s = conj(sign_safe(L̃[j, j]))
-                @simd for i in j:m
-                    L̃[i, j] = L̃[i, j] * s
-                end
-            end
-        end
+        positive && gaugefix!(lq_householder!, L̃, Q, diagview(A))
         copyto!(L, L̃)
+    else
+        gaugefix!(lq_householder!, nothing, Q, diagview(A))
     end
+
     return L, Q
 end
 function lq_householder!(

src/implementations/qr.jl

@@ -174,43 +174,16 @@ function qr_householder!(
         end
     end
 
-    if positive # already fix Q even if we do not need R
-        if driver === LAPACK()
-            @inbounds for j in 1:minmn
-                s = sign_safe(A[j, j])
-                @simd for i in 1:m
-                    Q[i, j] *= s
-                end
-            end
-        else
-            # guaranteed τ exists and no longer needed
-            τ .= sign_safe.(diagview(A))
-            Qf = view(Q, 1:m, 1:minmn) # first minmn columns of Q
-            Qf .= Qf .* transpose(τ)
-        end
-    end
-
     if computeR
-        R̃ = uppertriangular!(view(A, axes(R)...))
-        if positive
-            if driver === LAPACK()
-                @inbounds for j in n:-1:1
-                    @simd for i in 1:min(minmn, j)
-                        R̃[i, j] = R̃[i, j] * conj(sign_safe(R̃[i, i]))
-                    end
-                end
-            else
-                R̃f = view(R̃, 1:minmn, 1:n) # first minmn rows of R
-                R̃f .= conj.(τ) .* R̃f
-            end
-        end
-        if !pivoted
-            copyto!(R, R̃)
-        else
-            # probably very inefficient in terms of memory access
-            copyto!(view(R, :, jpvt), R̃)
-        end
+        # we need to first copy then gaugefix - avoiding aliasing between R and Rd for broadcast
+        Rd = diagview(A)
+        Rf = pivoted ? view(R, :, jpvt) : R
+        copyto!(Rf, uppertriangular!(view(A, axes(R)...)))
+        positive && gaugefix!(qr_householder!, Q, Rf, Rd)
+    elseif positive
+        gaugefix!(qr_householder!, Q, nothing, diagview(A))
     end
+
     return Q, R
 end
 function qr_householder!(