diff --git a/docs/src/darray.md b/docs/src/darray.md index 12c86ec80..6409744ef 100644 --- a/docs/src/darray.md +++ b/docs/src/darray.md @@ -705,7 +705,7 @@ From `LinearAlgebra`: - `mul!` (In-place Matrix-Matrix and Matrix-Vector multiply) - `cholesky`/`cholesky!` (In-place/Out-of-place Cholesky factorization) - `lu`/`lu!` (In-place/Out-of-place LU factorization (`NoPivot` and `RowMaximum`)) -- `qr`/`qr!` (In-place/Out-of-place QR factorization) +- `qr`/`qr!` (In-place/Out-of-place tiled Compact-WY QR; keywords `ib` and `p` for inner block size and CAQR domains) - `\`/`ldiv!` (In-place/Out-of-place Linear solving with LU, Cholesky, QR, and SVD factorizations) - `inv` (Out-of-place matrix inversion, including via SVD) - `svd`/`svd!`/`svdvals!` (In-place/Out-of-place Singular Value Decomposition) diff --git a/src/array/qr.jl b/src/array/qr.jl index fca2f4eaa..4409cc1c5 100644 --- a/src/array/qr.jl +++ b/src/array/qr.jl @@ -774,14 +774,38 @@ function qr_measure_workspace(A::DMatrix{T}, ib::Int) where {T<: Number} lm, ln end -function _qr_impl!(A::DMatrix{T}; ib::Int=1, p::Int=1) where {T<:Number} +""" + _default_qr_ib(A; p=1) -> Int + +Choose a Compact-WY inner block size for tiled QR. + +`ib == 1` is numerically fine but runs NextLA's unblocked kernels and scales +poorly. A modest inner block (capped at the tile width) restores BLAS-3 +intensity inside `geqrt!`/`tsmqr!` without lengthening the panel critical path. + +CAQR (`p > 1`) currently only produces correct results for `ib == 1`, so that +path keeps the unblocked default until the tree kernels are fixed for `ib > 1`. +""" +function _default_qr_ib(A::DMatrix; p::Int=1) + p > 1 && return 1 + nb = Int(A.partitioning.blocksize[2]) + nb <= 0 && return 1 + return min(32, nb) +end + +function _qr_impl!(A::DMatrix{T}; ib::Union{Int,Nothing}=nothing, p::Int=1) where {T<:Number} p >= 1 || throw(ArgumentError("p must be >= 1, got $p")) - lm, ln = qr_measure_workspace(A, ib) + ib_eff = something(ib, _default_qr_ib(A; p=p)) + ib_eff >= 1 || throw(ArgumentError("ib must be >= 1, got $ib_eff")) + if p > 1 && ib_eff > 1 + throw(ArgumentError("CAQR (p > 1) currently requires ib=1, got ib=$ib_eff")) + end + lm, ln = qr_measure_workspace(A, ib_eff) nb = A.partitioning.blocksize[2] irregular = _use_irregular_qr_tiling(A) Tm_cols = p > 1 ? 2 * ln : ln - Tm = DArray{T}(undef, Blocks(ib, nb), (lm, Tm_cols)) + Tm = DArray{T}(undef, Blocks(ib_eff, nb), (lm, Tm_cols)) if irregular # Hierarchical irregular path (no CAQR tree metadata). @@ -795,8 +819,22 @@ function _qr_impl!(A::DMatrix{T}; ib::Int=1, p::Int=1) where {T<:Number} return QRCompactWY(A, Tm) end -function LinearAlgebra.qr!(A::DMatrix{T}; ib::Int=1, p::Int=1) where {T<:Number} +""" + qr!(A::DMatrix; ib=nothing, p=1) -> QRCompactWY + +In-place tiled Compact-WY QR factorization of a distributed matrix. + +# Keywords +- `ib`: Compact-WY inner block size. For flat tile QR (`p=1`), defaults to + `min(32, tilewidth)` — much faster than `ib=1` for multi-threaded runs. + CAQR (`p>1`) currently requires `ib=1`. +- `p`: CAQR domain count (`p=1` is flat tile QR). `p>1` needs uniform square + tiling with `mt % p == 0`; on shared-memory multi-threading it is rarely faster + than `p=1`. +""" +function LinearAlgebra.qr!(A::DMatrix{T}; ib::Union{Int,Nothing}=nothing, p::Int=1) where {T<:Number} p >= 1 || throw(ArgumentError("p must be >= 1, got $p")) + ib === nothing || ib >= 1 || throw(ArgumentError("ib must be >= 1, got $ib")) if !_use_irregular_qr_tiling(A) return _qr_impl!(A; ib=ib, p=p) diff --git a/test/array/linalg/qr.jl b/test/array/linalg/qr.jl index 736ac1ae6..df4226461 100644 --- a/test/array/linalg/qr.jl +++ b/test/array/linalg/qr.jl @@ -20,7 +20,7 @@ function check_qr_residuals(A_col, DQ, DR; tol=2.0) end # ====================================================================== -# 1. Basic qr(DA) — varying shapes and block sizes (p=1, ib=1) +# 1. Basic qr(DA) — varying shapes and block sizes (p=1, default ib) # ====================================================================== @testset "Tile QR: $T" for T in (Float64, ComplexF64) @testset "Square matrices" begin @@ -510,6 +510,15 @@ end @test_throws ArgumentError cageqrf!(DA_copy, Tm; p=2) end +# ====================================================================== +# 18b. CAQR currently requires ib=1 +# ====================================================================== +@testset "CAQR rejects ib>1" begin + A = rand(Float64, 128, 64) + DA = distribute(A, Blocks(32, 32)) + @test_throws ArgumentError qr!(copy(DA); ib=32, p=2) +end + # ====================================================================== # 19. Edge case — single tile # ======================================================================