diff --git a/PERFORMANCE.md b/PERFORMANCE.md index 389e80a..54df883 100644 --- a/PERFORMANCE.md +++ b/PERFORMANCE.md @@ -8,9 +8,9 @@ Performance issues carry the GitHub label `performance` ([list](https://github.c | issue | what | experiment | status | | --- | --- | --- | --- | -| #81 (PR) | regime-A subtrees ran a whole KKT tree on one workgroup; flop limit `subtree_parallelism` | 0 | open PR | +| #81 (PR) | regime-A subtrees ran a whole KKT tree on one workgroup; flop limit `subtree_parallelism` | 0 | merged | | #82 | KKT refactorization and solve are level-bound after #81: 54–59 launches per refactorization, 62–85 per solve | 5, 6 | open | -| #75 | device LDLᵀ: serial pivot search, one workgroup per regime-B/C front, no vendor `sytrf` | 1, 2 | open, triaged | +| #75 | device LDLᵀ: serial pivot search, one workgroup per regime-B/C front, no vendor `sytrf` | 1, 2 | open, triaged; step 1 (cooperative pivot search) in a PR from `perf/exp1-pivot-search` | | #60 | regime-A follow-ups: CUDA timings (partly answered by experiment 0) and per-backend local-memory caps | 7 | open, triaged | | #25 | T25 performance pass (task) | 5, 6, 7 | open | @@ -109,6 +109,44 @@ A side effect: the default solve (`deterministic_mode = 0`) now uses its atomic Revised order: experiments 1–2 (LDLᵀ kernels, #75) and 5–6 (launches and solve) now have the largest measured payoff; experiment 3 (matching and scaling) remains the accuracy blocker on K2 (relres unchanged by this fix). +## Experiments 1–2 results (2026-10-05, RTX 4080) + +Issue #75 in three steps, one PR each, measured on the same machine with nothing else running. Refactorization: best of 5 warm runs of `bench/profile_phases.jl` (`*`: the `bench/compare.jl` median, for the matrices the profile skips); cuDSS from `bench/comparison/cudss.csv`; SDS Cholesky from the profile of `main`. The contract does not move: `piv` and `pivot_kind` equal `ref_ldlt!` on the test matrices on the CPU backend and CUDA, D and the panels within `panel_tol`. + +### Step 1: cooperative pivot search + +The pivot search of `front_ldlt_kernel!` and `subtree_ldlt_kernel!` ran on work item 1. Its column maxima are now workgroup reductions in up to three passes (λ and r with the column maximum; σ and the column-r maxima of Bunch–Kaufman; the fallback scan, one lane per candidate column), with the reference's tie-breaking (first maximum), so every decision is the serial one. Work item 1 takes the pivot from the values before the interchange, which keeps five barriers per step, and folds only the lanes that can hold data. The reference's fallback (`_best_1x1`) skips columns that cannot win and stops the threshold test at the first violation; the result is the same. The regime-A LDLᵀ kernel reserves 1024 B of local memory (`SUBTREE_LOCAL_RESERVE_LDLT`) for the reduction slots. + +LDLᵀ refactorization, `main` → step 1: + +| matrix | main ms | step 1 ms | step 1 / cuDSS | step 1 / SDS Cholesky | +| --- | --- | --- | --- | --- | +| lap2d_300 | 228 | 192 | 37.4× | 6.12× | +| lap3d_40 | 14,777* | 14,494* | 248× | 64.7× | +| HB/bcsstk17 | 206 | 167 | 49.5× | 6× | +| Boeing/bcsstk38 | 137 | 106 | 32.1× | 4.69× | +| GHS_psdef/apache2 | 92,576* | 97,296* | 199× | 94.5× | +| kkt_pglib_opf_case118_ieee_condensed_1 | 1.77 | 1.82 | 4.66× | 1.99× | +| kkt_pglib_opf_case118_ieee_condensed_10 | 1.72 | 1.79 | 5.13× | 2.04× | +| kkt_pglib_opf_case118_ieee_condensed_20 | 1.87 | 1.82 | 5.18× | 2.08× | +| kkt_pglib_opf_case118_ieee_k2_1 | 1.69 | 1.69 | 4.66× | | +| kkt_pglib_opf_case118_ieee_k2_10 | 1.62 | 1.68 | 4.74× | | +| kkt_pglib_opf_case118_ieee_k2_20 | 1.6 | 1.66 | 4.16× | | +| kkt_pglib_opf_case1354_pegase_condensed_1 | 10.9 | 7.2 | 10.5× | 2.46× | +| kkt_pglib_opf_case1354_pegase_condensed_10 | 10.9 | 7.3 | 11.4× | 2.5× | +| kkt_pglib_opf_case1354_pegase_condensed_20 | 11 | 7.32 | 10.4× | 2.51× | +| kkt_pglib_opf_case1354_pegase_k2_1 | 8.31 | 6.6 | 8.08× | | +| kkt_pglib_opf_case1354_pegase_k2_10 | 7.15 | 6.05 | 7.32× | | +| kkt_pglib_opf_case1354_pegase_k2_20 | 9.1 | 6.71 | 8.32× | | +| kkt_pglib_opf_case14_ieee_condensed_1 | 0.446 | 0.503 | 2.61× | 1.38× | +| kkt_pglib_opf_case14_ieee_condensed_10 | 0.443 | 0.52 | 2.26× | 1.53× | +| kkt_pglib_opf_case14_ieee_condensed_11 | 0.448 | 0.463 | 1.59× | 1.41× | +| kkt_pglib_opf_case14_ieee_k2_1 | 0.669 | 0.697 | 2.49× | | +| kkt_pglib_opf_case14_ieee_k2_10 | 0.722 | 0.708 | 2.98× | | +| kkt_pglib_opf_case14_ieee_k2_15 | 0.97 | 1.03 | 3.2× | | + +On the K2 dumps the inertia and `nperturbed` are those of `main` on all nine. On the T15 matrix `kkt_matrix(3000, 1000, 1e-8)` (fallback scans on most steps of its two large root fronts) CUDA goes from 33.9 s to 8.3 s; the KA CPU backend stays at 7.0 s against 2.6 s for `ref_ldlt!` (it runs the work items of a workgroup one after another, so a parallel search saves nothing there). The medium matrices gain 16–34%; the smallest KKT dumps (case14, w ≤ 16) lose up to 17% to the reductions' barriers; the two matrices with very large regime-C fronts (apache2, lap3d_40) are within ±5%, still one workgroup per front. + ## Recommendations: Prioritized Experiment Plan | # | Experiment | Payoff | Effort | Key measurement | Success criterion | diff --git a/src/numeric/ldlt.jl b/src/numeric/ldlt.jl index 5640e4e..7b96d5b 100644 --- a/src/numeric/ldlt.jl +++ b/src/numeric/ldlt.jl @@ -14,11 +14,14 @@ # front's packed contribution block on the update stack. # # Per pivot step (a loop over the `w` columns of the front with a uniform trip -# count; a 2×2 pivot leaves the last iterations idle): work item 1 chooses the -# pivot (serial search, the reference's `_choose_pivot`), the workgroup swaps -# rows/columns, work item 1 stores D (perturbing a tiny 1×1 pivot) and counts -# the statistics, the workgroup applies the rank-1/rank-2 update to the -# remaining fully-summed columns and then scales the pivot columns into L. +# count; a 2×2 pivot leaves the last iterations idle): the workgroup chooses the +# pivot (the reference's `_choose_pivot`, its column maxima as workgroup +# reductions with the reference's tie-breaking, in one to three passes; work +# item 1 decides and stores D, perturbing a tiny 1×1 pivot, and counts the +# statistics), the workgroup swaps rows/columns, applies the rank-1/rank-2 +# update to the remaining fully-summed columns, scales the pivot columns into +# L and reduces the next column (pass 1 of the next step). Five barriers per +# step when the first pass settles the pivot. # Only the lower triangle of the front is stored; the upper entries the # algorithm reads are `conj` (Hermitian) or plain (complex symmetric) mirrors. # Pivot decisions, control words and the per-front counts live in `@localmem`; @@ -41,7 +44,16 @@ const _LT_STEP = _ST_CTL + 3 # size of the pivot taken at _LT_K (1 or 2; 0 bef const _LT_C = _ST_CTL + 4 # chosen 1×1 pivot column (2×2: k) const _LT_R = _ST_CTL + 5 # 2×2 partner column (0: 1×1 pivot) const _LT_STAT = _ST_CTL + 5 # _LT_STAT + q: statistic q ∈ 1:5 (STAT_NPOS … STAT_N2X2) of the front -const _LT_CTL = _ST_CTL + 10 +const _LT_PHASE = _ST_CTL + 11 # pivot search state: 0 chosen, 2 Bunch–Kaufman pass 2 due, 3 fallback scan due +const _LT_RBK = _ST_CTL + 12 # Bunch–Kaufman candidate row r of pass 1 (0: none) +const _LT_CTL = _ST_CTL + 12 + +# `pv` slots: 1–4 the pivot (d, or the inverse 2×2 block e11, e12, e21, e22), 5–6 the new diagonal +# entries of the pivot block (d, or a and c), 7–8 λ and the column maximum of pass 1 +const _LT_NPV = 8 + +# lanes of the stage-1 combine of the cooperative pivot search (`NL` partial results → `_LT_S1` → 1) +const _LT_S1 = 16 # resolved pivoting parameters passed to the kernels (isbits); `H`: Hermitian (or real symmetric) vs complex # symmetric (a type parameter, so the kernels need no extra `Val` argument: launches with more than 32 @@ -110,21 +122,35 @@ end @inline _lt_front(fa, ctl, ::Val{false}) = @inbounds _PanelFront(fa, Int(ctl[_ST_LF]) - 1, Int(ctl[_ST_F])) # --------------------------------------------------------------------------- -# pivot choice (work item 1; the reference's `_colmax`, `_accept_1x1`, `_accept_2x2`, -# `_best_1x1`, `_choose_pivot` on the lower triangle) +# cooperative pivot search (the reference's `_choose_pivot` on the lower triangle, pivot for pivot) +# +# The column maxima of the search are reductions: every lane (work item `li ≤ NL`) reduces a strided +# share of the rows (or, in the fallback scan, of the candidate columns) into its slots of `red`/`redi`, +# `_lt_stage1!` folds them into `_LT_S1` slots and work item 1 into one (`_lt_final!`), then decides. +# `max` of absolute values is exact in any order (a NaN propagates as in the serial loop), and an +# arg-max keeps the first maximum (the smallest index on a tie, index 0 = none), so every decision +# equals the serial one. Slots: `red[1:NL]`, `red[NL .+ (1:NL)]`, `red[2NL .+ (1:NL)]` (values +# A, B, C), `redi[1:NL]`, `redi[NL .+ (1:NL)]` (indices of A and B). Per pivot step: +# +# * pass 1 (always; computed in the scale phase of the previous step, or before the first step): +# A/idx = λ and r (largest `|F[i, k]|` of the block rows `k+1:w`, first maximum), B = the column +# maximum of column k over rows `k+1:f` (the threshold test of a 1×1 pivot at k); +# * pass 2 (Bunch–Kaufman when `|a_kk| < αλ`): A = σ (column r over the block rows `k:w` but r), +# B = column r over rows `k:f` but r and k, C = column k over rows `k:f` but k and r; +# * pass 3 (the fallback when the choice fails the threshold): A/idx = the acceptable 1×1 pivot with +# the largest `|a_jj|` (`_best_1x1`, one lane per candidate column), B/idx = the largest `|a_jj|` +# (pivot type 'D'). -@inline function _lt_colmax(F, c, k, f, skip) - m = zero(real(eltype(F.a))) +# `F[i, c]` bound of the threshold test: `d ≥ u maxᵢ |F[i, c]|` over rows `k:f` but c, with an early +# exit (`u ≥ 0` is finite: `u · max = max(u · |F[i, c]|)`, and a NaN fails both forms) +@inline function _lt_threshold_ok(F, c, k, f, u, d) + m = zero(d) for i in k:f - (i == c || i == skip) && continue + i == c && continue m = max(m, _fabs(F, i, c)) + d >= u * m || return false end - return m -end - -@inline function _lt_accept_1x1(F, c, k, f, u, ε) - d = _fabs(F, c, c) - return d >= ε && d >= u * _lt_colmax(F, c, k, f, 0) + return d >= u * m end @inline _lt_det2(F, k, r, h::Val) = _fget(F, k, k) * _fget(F, r, r) - _fsym(F, k, r, h) * _fget(F, r, k) @@ -134,76 +160,340 @@ end return !iszero(det) && isfinite(det) && max(_fabs(F, k, k), _fabs(F, r, k), _fabs(F, r, r)) >= ε end -@inline function _lt_accept_2x2(F, k, r, f, u, ε, h::Val) +# the 2×2 pivot on (k, r) with the column maxima m1 (column k without rows k, r) and m2 (column r +# without rows r, k) is acceptable (the reference's `_accept_2x2`) +@inline function _lt_accept_2x2(F, k, r, m1, m2, u, ε, h::Val) _lt_nonsingular_2x2(F, k, r, ε, h) || return false det = _lt_det2(F, k, r, h) e11 = abs(_fget(F, r, r) / det) e12 = abs(_fsym(F, k, r, h) / det) e21 = abs(_fget(F, r, k) / det) e22 = abs(_fget(F, k, k) / det) - m1 = _lt_colmax(F, k, k, f, r) - m2 = _lt_colmax(F, r, k, f, k) return u * (e11 * m1 + e21 * m2) <= 1 && u * (e12 * m1 + e22 * m2) <= 1 end -@inline function _lt_best_1x1(F, k, w, f, u, ε) - best = 0 - bv = zero(real(eltype(F.a))) - for j in k:w - _lt_accept_1x1(F, j, k, f, u, ε) || continue - a = _fabs(F, j, j) - if best == 0 || a > bv - best = j - bv = a +# fold arg-max slot b into slot a (first maximum, index 0 = none) +@inline function _lt_argmax!(red, redi, a, b) + @inbounds begin + ib = redi[b] + if ib != 0 + ia = redi[a] + if ia == 0 || red[b] > red[a] || (red[b] == red[a] && ib < ia) + red[a] = red[b] + redi[a] = ib + end end end - return best + return nothing end -# `(c, 0)`: 1×1 pivot at column c; `(k, r)`: 2×2 pivot on columns k and r (as `_choose_pivot`) -@inline function _lt_choose(F, k, w, f, p::_LDLTDevice{R}, ε, h::Val) where {R} - p.ptype == _LT_PIVOT_NONE && return (k, 0) - u = p.u - if p.ptype == _LT_PIVOT_DIAGONAL - _lt_accept_1x1(F, k, k, f, u, ε) && return (k, 0) - best = _lt_best_1x1(F, k, w, f, u, ε) - best != 0 && return (best, 0) - big = k - for j in (k + 1):w - _fabs(F, j, j) > _fabs(F, big, big) && (big = j) +# fold lane b into lane a for the partial results of pass `P` +@inline function _lt_combine!(red, redi, a, b, ::Val{NL}, ::Val{P}) where {NL, P} + @inbounds if P == 1 + _lt_argmax!(red, redi, a, b) + red[NL + a] = max(red[NL + a], red[NL + b]) + elseif P == 2 + red[a] = max(red[a], red[b]) + red[NL + a] = max(red[NL + a], red[NL + b]) + red[2 * NL + a] = max(red[2 * NL + a], red[2 * NL + b]) + else + _lt_argmax!(red, redi, a, b) + _lt_argmax!(red, redi, NL + a, NL + b) + end + return nothing +end + +# lanes of pass `P` at column k that can hold a partial result (the others hold the neutral one): one per row +# below k (pass 1, with k the column the step advances to), per row from k (pass 2), per column from k (pass 3) +@inline _lt_nused(ctl, k, ::Val{P}) where {P} = + @inbounds P == 3 ? Int(ctl[_ST_W]) - k + 1 : Int(ctl[_ST_F]) - k + 1 + +# stage 1: lanes 1:_LT_S1 fold lanes _LT_S1+1:min(NL, nu) (nothing when NL ≤ _LT_S1) +@inline function _lt_stage1!(red, redi, ctl, li, nl::Val{NL}, pass::Val{P}) where {NL, P} + @inbounds if NL > _LT_S1 && li <= _LT_S1 + k = Int(ctl[_LT_K]) + (P == 1 ? Int(ctl[_LT_STEP]) : 0) + for b in (li + _LT_S1):_LT_S1:min(NL, _lt_nused(ctl, k, pass)) + _lt_combine!(red, redi, li, b, nl, pass) end - return (_fabs(F, big, big) >= ε ? big : k, 0) end - α = R(BUNCH_KAUFMAN_ALPHA) - λ = zero(R) - r = 0 - for i in (k + 1):w - a = _fabs(F, i, k) - if a > λ - λ = a - r = i + return nothing +end + +# work item 1: fold lanes 2:min(NL, _LT_S1, nu) into lane 1 (k: the column of the step) +@inline function _lt_final!(red, redi, ctl, k, nl::Val{NL}, pass::Val) where {NL} + for b in 2:min(NL, _LT_S1, _lt_nused(ctl, k, pass)) + _lt_combine!(red, redi, 1, b, nl, pass) + end + return nothing +end + +# pass 1 partials of the next column k = K + STEP (before the step that advances to it) +@inline function _lt_pass1!(fa, ctl, red, redi, li, ::Val{NL}, p::_LDLTDevice{R}, pk::Val) where {NL, R} + @inbounds if li <= NL + k = Int(ctl[_LT_K]) + Int(ctl[_LT_STEP]) + w = Int(ctl[_ST_W]) + λ = zero(R) + r = 0 + cm = zero(R) + if k <= w && p.ptype != _LT_PIVOT_NONE + F = _lt_front(fa, ctl, pk) + for i in (k + li):NL:Int(ctl[_ST_F]) + a = _fabs(F, i, k) + cm = max(cm, a) + if i <= w && a > λ + λ = a + r = i + end + end end + red[li] = λ + redi[li] = r % eltype(redi) + red[NL + li] = cm end - akk = _fabs(F, k, k) - c, r2 = k, 0 - if !(r == 0 || akk >= α * λ) - σ = zero(R) # largest off-diagonal of column r in the block - for i in k:w - i != r && (σ = max(σ, _fabs(F, i, r))) + return nothing +end + +# pass 2 partials (column r of the Bunch–Kaufman candidate, column k without row r) +@inline function _lt_pass2!(fa, ctl, red, redi, li, ::Val{NL}, ::_LDLTDevice{R}, pk::Val) where {NL, R} + @inbounds if li <= NL + k = Int(ctl[_LT_K]) + r = Int(ctl[_LT_RBK]) + w = Int(ctl[_ST_W]) + σ = zero(R) + mr = zero(R) + mk = zero(R) + F = _lt_front(fa, ctl, pk) + for i in (k + li - 1):NL:Int(ctl[_ST_F]) + if i != r + a = _fabs(F, i, r) + i <= w && (σ = max(σ, a)) + i != k && (mr = max(mr, a)) + i != k && (mk = max(mk, _fabs(F, i, k))) + end end + red[li] = σ + red[NL + li] = mr + red[2 * NL + li] = mk + end + return nothing +end + +# pass 3 partials: one lane per candidate column j of the block k:w +@inline function _lt_pass3!(fa, ctl, red, redi, li, ::Val{NL}, aux, p::_LDLTDevice{R}, pk::Val) where {NL, R} + @inbounds if li <= NL + k = Int(ctl[_LT_K]) + w = Int(ctl[_ST_W]) + f = Int(ctl[_ST_F]) + ε = _lt_eps(p, aux) + u = p.u + F = _lt_front(fa, ctl, pk) + bv = zero(R) + bi = 0 + gv = zero(R) + gi = 0 + for j in (k + li - 1):NL:w + a = _fabs(F, j, j) + if p.ptype == _LT_PIVOT_DIAGONAL && !isnan(a) && (gi == 0 || a > gv) + gv = a + gi = j + end + a >= ε || continue # not acceptable (or NaN) + bi != 0 && !(a > bv) && continue # cannot replace this lane's best + _lt_threshold_ok(F, j, k, f, u, a) || continue + bv = a + bi = j + end + IT = eltype(redi) + red[li] = bv + redi[li] = bi % IT + red[NL + li] = gv + redi[NL + li] = gi % IT + end + return nothing +end + +# work item 1: take the pivot (c, r) at column k (`r = 0`: 1×1 at column c; else 2×2 on k and r) from the +# values before the interchange: D (perturbing a tiny 1×1 pivot), the pivot kinds, the statistics, `pv` +# (the pivot, its inverse block and the new diagonal entries that `_lt_swap!` writes) and the step size +@inline function _lt_select!(fa, ctl, pv, d, pivot_kind, piv, psign, perm, aux, p::_LDLTDevice, c, r, pk::Val, + h::Val{H}) where {H} + @inbounds begin + n = length(perm) + IT = eltype(ctl) + T = eltype(fa) + k = Int(ctl[_LT_K]) + c0 = Int(ctl[_LT_C0]) + g = c0 + k - 1 + F = _lt_front(fa, ctl, pk) + ctl[_LT_C] = c % IT + ctl[_LT_R] = r % IT + ctl[_LT_PHASE] = zero(IT) + if r == 0 + ε = _lt_eps(p, aux) + x = _fget(F, c, c) + dk = H ? T(real(x)) : x + kind = PIVOT_KIND_1X1 + if !(abs(dk) >= ε) # tiny (or NaN): perturb + iszero(dk) && (ctl[_LT_STAT + STAT_NZERO] += one(IT)) + dk = _perturbation_sign(dk, Int(psign[perm[piv[c0 + c - 1]]]), H) * ε + kind = PIVOT_KIND_PERTURBED + ctl[_LT_STAT + STAT_NPERTURBED] += one(IT) + end + d[g] = dk + d[n + g] = zero(T) + pivot_kind[g] = kind + if H + if real(dk) > 0 + ctl[_LT_STAT + STAT_NPOS] += one(IT) + elseif real(dk) < 0 + ctl[_LT_STAT + STAT_NNEG] += one(IT) + end + end + pv[1] = dk + pv[5] = dk + ctl[_LT_STEP] = one(IT) + else + # after the interchange of k + 1 and r: F[k + 1, k] = F[r, k], F[k + 1, k + 1] = F[r, r] + a = _fget(F, k, k) + b = _fget(F, r, k) + c2 = _fget(F, r, r) + if H + a = T(real(a)) + c2 = T(real(c2)) + end + up = _cj(b, h) + det = a * c2 - up * b + pv[1] = c2 / det + pv[2] = -up / det + pv[3] = -b / det + pv[4] = a / det + pv[5] = a + pv[6] = c2 + d[g] = a + d[g + 1] = c2 + d[n + g] = b + d[n + g + 1] = zero(T) + pivot_kind[g] = PIVOT_KIND_2X2_FIRST + pivot_kind[g + 1] = PIVOT_KIND_2X2_SECOND + ctl[_LT_STAT + STAT_N2X2] += one(IT) + if H + if real(det) < 0 + ctl[_LT_STAT + STAT_NPOS] += one(IT) + ctl[_LT_STAT + STAT_NNEG] += one(IT) + elseif real(a) > 0 + ctl[_LT_STAT + STAT_NPOS] += IT(2) + else + ctl[_LT_STAT + STAT_NNEG] += IT(2) + end + end + ctl[_LT_STEP] = IT(2) + end + end + return nothing +end + +# work item 1 after pass 1: advance to column k, then take the pivot at k when the search ends here, +# else request pass 2 or the fallback scan +@inline function _lt_decide1!(fa, ctl, pv, red, redi, d, pivot_kind, piv, psign, perm, aux, + p::_LDLTDevice{R}, nl::Val{NL}, pk::Val, h::Val) where {R, NL} + @inbounds begin + IT = eltype(ctl) + k = Int(ctl[_LT_K]) + Int(ctl[_LT_STEP]) + ctl[_LT_K] = k % IT + ctl[_LT_STEP] = zero(IT) + ctl[_LT_PHASE] = zero(IT) + k <= Int(ctl[_ST_W]) || return nothing + if p.ptype == _LT_PIVOT_NONE + _lt_select!(fa, ctl, pv, d, pivot_kind, piv, psign, perm, aux, p, k, 0, pk, h) + return nothing + end + _lt_final!(red, redi, ctl, k, nl, Val(1)) + λ = red[1] + r = Int(redi[1]) + cm = red[NL + 1] + pv[7] = λ + pv[8] = cm + ctl[_LT_RBK] = r % IT + ctl[_LT_C] = k % IT # the Bunch–Kaufman choice so far: 1×1 at k + ctl[_LT_R] = zero(IT) + F = _lt_front(fa, ctl, pk) + ε = _lt_eps(p, aux) + akk = _fabs(F, k, k) + if p.ptype == _LT_PIVOT_BK && !(r == 0 || akk >= R(BUNCH_KAUFMAN_ALPHA) * λ) + ctl[_LT_PHASE] = IT(2) + elseif akk >= ε && akk >= p.u * cm + _lt_select!(fa, ctl, pv, d, pivot_kind, piv, psign, perm, aux, p, k, 0, pk, h) + else + ctl[_LT_PHASE] = IT(3) + end + end + return nothing +end + +# work item 1 after pass 2: the rest of the Bunch–Kaufman choice and its threshold test +@inline function _lt_decide2!(fa, ctl, pv, red, redi, d, pivot_kind, piv, psign, perm, aux, + p::_LDLTDevice{R}, nl::Val{NL}, pk::Val, h::Val) where {R, NL} + @inbounds begin + IT = eltype(ctl) + _lt_final!(red, redi, ctl, Int(ctl[_LT_K]), nl, Val(2)) + σ = red[1] + mr = red[NL + 1] + mk = red[2 * NL + 1] + λ = real(pv[7]) + cm = real(pv[8]) + k = Int(ctl[_LT_K]) + r = Int(ctl[_LT_RBK]) + F = _lt_front(fa, ctl, pk) + ε = _lt_eps(p, aux) + u = p.u + α = R(BUNCH_KAUFMAN_ALPHA) + akk = _fabs(F, k, k) + c, r2 = k, 0 if akk * σ >= α * λ^2 - # 1×1 at k - elseif _fabs(F, r, r) >= α * σ + ok = akk >= ε && akk >= u * cm + elseif (arr = _fabs(F, r, r); arr >= α * σ) c = r + ok = arr >= ε && arr >= u * max(mr, _fabs(F, k, r)) else r2 = r + ok = _lt_accept_2x2(F, k, r, mk, mr, u, ε, h) + end + if ok + _lt_select!(fa, ctl, pv, d, pivot_kind, piv, psign, perm, aux, p, c, r2, pk, h) + else + ctl[_LT_C] = c % IT + ctl[_LT_R] = r2 % IT + ctl[_LT_PHASE] = IT(3) end end - (r2 == 0 ? _lt_accept_1x1(F, c, k, f, u, ε) : _lt_accept_2x2(F, k, r2, f, u, ε, h)) && return (c, r2) - best = _lt_best_1x1(F, k, w, f, u, ε) - best != 0 && return (best, 0) - r2 != 0 && _lt_nonsingular_2x2(F, k, r2, ε, h) && return (c, r2) - return (c, 0) + return nothing +end + +# work item 1 after pass 3: the acceptable 1×1 pivot with the largest |a_jj|, else the Bunch–Kaufman +# choice when it is not tiny (or, for 'D', the largest diagonal), else column k perturbed +@inline function _lt_decide3!(fa, ctl, pv, red, redi, d, pivot_kind, piv, psign, perm, aux, + p::_LDLTDevice, nl::Val{NL}, pk::Val, h::Val) where {NL} + @inbounds begin + _lt_final!(red, redi, ctl, Int(ctl[_LT_K]), nl, Val(3)) + k = Int(ctl[_LT_K]) + best = Int(redi[1]) + c, r = best, 0 + if best == 0 + F = _lt_front(fa, ctl, pk) + ε = _lt_eps(p, aux) + if p.ptype == _LT_PIVOT_DIAGONAL + big = isnan(_fabs(F, k, k)) ? k : Int(redi[NL + 1]) + c = _fabs(F, big, big) >= ε ? big : k + else + c, r = Int(ctl[_LT_C]), Int(ctl[_LT_R]) + if r != 0 && !_lt_nonsingular_2x2(F, k, r, ε, h) + r = 0 + end + end + end + _lt_select!(fa, ctl, pv, d, pivot_kind, piv, psign, perm, aux, p, c, r, pk, h) + end + return nothing end # --------------------------------------------------------------------------- @@ -218,6 +508,8 @@ end ctl[_LT_STEP] = zero(IT) ctl[_LT_C] = zero(IT) ctl[_LT_R] = zero(IT) + ctl[_LT_PHASE] = zero(IT) + ctl[_LT_RBK] = zero(IT) for q in 1:5 ctl[_LT_STAT + q] = zero(IT) end @@ -236,35 +528,18 @@ end return nothing end -# work item 1: advance to the next column and choose its pivot -@inline function _lt_choose!(fa, ctl, aux, p::_LDLTDevice, pk::Val, h::Val) - @inbounds begin - IT = eltype(ctl) - k = Int(ctl[_LT_K]) + Int(ctl[_LT_STEP]) - ctl[_LT_K] = k % IT - ctl[_LT_STEP] = zero(IT) - w = Int(ctl[_ST_W]) - if k <= w - F = _lt_front(fa, ctl, pk) - c, r = _lt_choose(F, k, w, Int(ctl[_ST_F]), p, _lt_eps(p, aux), h) - ctl[_LT_C] = c % IT - ctl[_LT_R] = r % IT - end - end - return nothing -end - -# symmetric interchange of the columns/rows p < q of the front (lower triangle), as `_swap_front!` -@inline function _lt_swap!(fa, ctl, piv, li, ::Val{WG}, pk::Val, h::Val) where {WG} +# symmetric interchange of the columns/rows p < q of the front (lower triangle), as `_swap_front!`, then +# the new diagonal entries of the pivot block from `pv` (and the zero below a 2×2 block's diagonal) +@inline function _lt_swap!(fa, ctl, pv, piv, li, ::Val{WG}, pk::Val, h::Val) where {WG} @inbounds begin k = Int(ctl[_LT_K]) if k <= Int(ctl[_ST_W]) r = Int(ctl[_LT_R]) p = r == 0 ? k : k + 1 q = r == 0 ? Int(ctl[_LT_C]) : r - if p != q - F = _lt_front(fa, ctl, pk) - for i in li:WG:Int(ctl[_ST_F]) + F = _lt_front(fa, ctl, pk) + for i in li:WG:Int(ctl[_ST_F]) + if p != q if i < p x = _fget(F, p, i) _fset!(F, p, i, _fget(F, q, i)) @@ -285,89 +560,18 @@ end _fset!(F, i, q, x) end end - if li == 1 - c0 = Int(ctl[_LT_C0]) - x = piv[c0 + p - 1] - piv[c0 + p - 1] = piv[c0 + q - 1] - piv[c0 + q - 1] = x + if i == k + _fset!(F, k, k, pv[5]) + r == 0 || _fset!(F, k + 1, k, zero(eltype(fa))) # the 2×2 block lives in D, L's is I + elseif r != 0 && i == k + 1 + _fset!(F, k + 1, k + 1, pv[6]) end end - end - end - return nothing -end - -# work item 1: store the pivot block in D (perturbing a tiny 1×1 pivot), the pivot kinds and the -# statistics; `pv` receives d (1×1) or the entries (e11, e12, e21, e22) of the inverse 2×2 block -@inline function _lt_pivot!(fa, ctl, pv, d, pivot_kind, piv, psign, perm, aux, p::_LDLTDevice, pk::Val, - h::Val{H}) where {H} - @inbounds begin - n = length(perm) - IT = eltype(ctl) - T = eltype(fa) - k = Int(ctl[_LT_K]) - if k <= Int(ctl[_ST_W]) - F = _lt_front(fa, ctl, pk) - g = Int(ctl[_LT_C0]) + k - 1 - if ctl[_LT_R] == 0 - ε = _lt_eps(p, aux) - x = _fget(F, k, k) - dk = H ? T(real(x)) : x - kind = PIVOT_KIND_1X1 - if !(abs(dk) >= ε) # tiny (or NaN): perturb - iszero(dk) && (ctl[_LT_STAT + STAT_NZERO] += one(IT)) - dk = _perturbation_sign(dk, Int(psign[perm[piv[g]]]), H) * ε - kind = PIVOT_KIND_PERTURBED - ctl[_LT_STAT + STAT_NPERTURBED] += one(IT) - end - _fset!(F, k, k, dk) - d[g] = dk - d[n + g] = zero(T) - pivot_kind[g] = kind - if H - if real(dk) > 0 - ctl[_LT_STAT + STAT_NPOS] += one(IT) - elseif real(dk) < 0 - ctl[_LT_STAT + STAT_NNEG] += one(IT) - end - end - pv[1] = dk - ctl[_LT_STEP] = one(IT) - else - a = _fget(F, k, k) - b = _fget(F, k + 1, k) - c2 = _fget(F, k + 1, k + 1) - if H - a = T(real(a)) - c2 = T(real(c2)) - end - up = _cj(b, h) - _fset!(F, k, k, a) - _fset!(F, k + 1, k + 1, c2) - _fset!(F, k + 1, k, zero(T)) # the 2×2 block lives in D, L's block is the identity - det = a * c2 - up * b - pv[1] = c2 / det - pv[2] = -up / det - pv[3] = -b / det - pv[4] = a / det - d[g] = a - d[g + 1] = c2 - d[n + g] = b - d[n + g + 1] = zero(T) - pivot_kind[g] = PIVOT_KIND_2X2_FIRST - pivot_kind[g + 1] = PIVOT_KIND_2X2_SECOND - ctl[_LT_STAT + STAT_N2X2] += one(IT) - if H - if real(det) < 0 - ctl[_LT_STAT + STAT_NPOS] += one(IT) - ctl[_LT_STAT + STAT_NNEG] += one(IT) - elseif real(a) > 0 - ctl[_LT_STAT + STAT_NPOS] += IT(2) - else - ctl[_LT_STAT + STAT_NNEG] += IT(2) - end - end - ctl[_LT_STEP] = IT(2) + if li == 1 && p != q + c0 = Int(ctl[_LT_C0]) + x = piv[c0 + p - 1] + piv[c0 + p - 1] = piv[c0 + q - 1] + piv[c0 + q - 1] = x end end end @@ -536,7 +740,9 @@ statistics, updates its packed contribution block on the update stack li = @index(Local, Linear) G = @index(Group, Linear) ctl = @localmem IT (_LT_CTL,) - pv = @localmem TT (4,) + pv = @localmem TT (_LT_NPV,) + red = @localmem R (3 * WG,) + redi = @localmem IT (2 * WG,) if li == 1 @inbounds s = nodes[bm.first + _bm_node(bm, G) - 1] _lt_front_setup!(ctl, s, super_ptr, member_panels(front_ptr, _bm_gmember(bm, G), bm.nbatch), front_nrows, @@ -562,23 +768,57 @@ statistics, updates its packed contribution block on the update stack front_nrows, front_ncols, cb_ptr, child_ptr, child_list, relind_ptr, relind, Val(WG)) @synchronize end + _lt_pass1!(factor, ctl, red, redi, li, Val(WG), prm, Val(false)) + @synchronize for it in 1:ctl[_ST_W] - if li == 1 - _lt_choose!(factor, ctl, _mview(aux, _bm_gmember(bm, G), bm.nbatch), prm, Val(false), Val(HERM)) + if WG > _LT_S1 + _lt_stage1!(red, redi, ctl, li, Val(WG), Val(1)) + @synchronize end - @synchronize - _lt_swap!(factor, ctl, _mview(piv, _bm_gmember(bm, G), bm.nbatch), li, Val(WG), Val(false), Val(HERM)) - @synchronize if li == 1 k = _bm_gmember(bm, G) nb = bm.nbatch - _lt_pivot!(factor, ctl, pv, _mview(d, k, nb), _mview(pivot_kind, k, nb), _mview(piv, k, nb), psign, perm, - _mview(aux, k, nb), prm, Val(false), Val(HERM)) + _lt_decide1!(factor, ctl, pv, red, redi, _mview(d, k, nb), _mview(pivot_kind, k, nb), _mview(piv, k, nb), + psign, perm, _mview(aux, k, nb), prm, Val(WG), Val(false), Val(HERM)) end @synchronize + if ctl[_LT_PHASE] == 2 + _lt_pass2!(factor, ctl, red, redi, li, Val(WG), prm, Val(false)) + @synchronize + if WG > _LT_S1 + _lt_stage1!(red, redi, ctl, li, Val(WG), Val(2)) + @synchronize + end + if li == 1 + k = _bm_gmember(bm, G) + nb = bm.nbatch + _lt_decide2!(factor, ctl, pv, red, redi, _mview(d, k, nb), _mview(pivot_kind, k, nb), + _mview(piv, k, nb), psign, perm, _mview(aux, k, nb), prm, Val(WG), Val(false), Val(HERM)) + end + @synchronize + end + if ctl[_LT_PHASE] == 3 + _lt_pass3!(factor, ctl, red, redi, li, Val(WG), _mview(aux, _bm_gmember(bm, G), bm.nbatch), prm, + Val(false)) + @synchronize + if WG > _LT_S1 + _lt_stage1!(red, redi, ctl, li, Val(WG), Val(3)) + @synchronize + end + if li == 1 + k = _bm_gmember(bm, G) + nb = bm.nbatch + _lt_decide3!(factor, ctl, pv, red, redi, _mview(d, k, nb), _mview(pivot_kind, k, nb), + _mview(piv, k, nb), psign, perm, _mview(aux, k, nb), prm, Val(WG), Val(false), Val(HERM)) + end + @synchronize + end + _lt_swap!(factor, ctl, pv, _mview(piv, _bm_gmember(bm, G), bm.nbatch), li, Val(WG), Val(false), Val(HERM)) + @synchronize _lt_update!(factor, ctl, pv, li, Val(WG), Val(false), Val(HERM)) @synchronize _lt_scale!(factor, ctl, pv, li, Val(WG), Val(false)) + _lt_pass1!(factor, ctl, red, redi, li, Val(WG), prm, Val(false)) @synchronize end @inbounds s = nodes[bm.first + _bm_node(bm, G) - 1] @@ -652,7 +892,9 @@ stack. G = @index(Group, Linear) buf = @localmem TT (NE,) ctl = @localmem IT (_LT_CTL,) - pv = @localmem TT (4,) + pv = @localmem TT (_LT_NPV,) + red = @localmem R (3 * _LT_S1,) + redi = @localmem IT (2 * _LT_S1,) if li == 1 @inbounds t = trees[bm.first + _bm_node(bm, G) - 1] @inbounds ctl[_ST_FIRST] = subtree_ptr[t] @@ -676,23 +918,48 @@ stack. relind, li, Val(WG)) @synchronize end + _lt_pass1!(buf, ctl, red, redi, li, Val(_LT_S1), prm, Val(true)) + @synchronize for it in 1:ctl[_ST_W] - if li == 1 - _lt_choose!(buf, ctl, _mview(aux, _bm_gmember(bm, G), bm.nbatch), prm, Val(true), Val(HERM)) - end - @synchronize - _lt_swap!(buf, ctl, _mview(piv, _bm_gmember(bm, G), bm.nbatch), li, Val(WG), Val(true), Val(HERM)) - @synchronize if li == 1 mb = _bm_gmember(bm, G) nb = bm.nbatch - _lt_pivot!(buf, ctl, pv, _mview(d, mb, nb), _mview(pivot_kind, mb, nb), _mview(piv, mb, nb), psign, - perm, _mview(aux, mb, nb), prm, Val(true), Val(HERM)) + _lt_decide1!(buf, ctl, pv, red, redi, _mview(d, mb, nb), _mview(pivot_kind, mb, nb), + _mview(piv, mb, nb), psign, perm, _mview(aux, mb, nb), prm, Val(_LT_S1), Val(true), + Val(HERM)) + end + @synchronize + if ctl[_LT_PHASE] == 2 + _lt_pass2!(buf, ctl, red, redi, li, Val(_LT_S1), prm, Val(true)) + @synchronize + if li == 1 + mb = _bm_gmember(bm, G) + nb = bm.nbatch + _lt_decide2!(buf, ctl, pv, red, redi, _mview(d, mb, nb), _mview(pivot_kind, mb, nb), + _mview(piv, mb, nb), psign, perm, _mview(aux, mb, nb), prm, Val(_LT_S1), Val(true), + Val(HERM)) + end + @synchronize + end + if ctl[_LT_PHASE] == 3 + _lt_pass3!(buf, ctl, red, redi, li, Val(_LT_S1), _mview(aux, _bm_gmember(bm, G), bm.nbatch), prm, + Val(true)) + @synchronize + if li == 1 + mb = _bm_gmember(bm, G) + nb = bm.nbatch + _lt_decide3!(buf, ctl, pv, red, redi, _mview(d, mb, nb), _mview(pivot_kind, mb, nb), + _mview(piv, mb, nb), psign, perm, _mview(aux, mb, nb), prm, Val(_LT_S1), Val(true), + Val(HERM)) + end + @synchronize end + _lt_swap!(buf, ctl, pv, _mview(piv, _bm_gmember(bm, G), bm.nbatch), li, Val(WG), Val(true), Val(HERM)) @synchronize _lt_update!(buf, ctl, pv, li, Val(WG), Val(true), Val(HERM)) @synchronize _lt_scale!(buf, ctl, pv, li, Val(WG), Val(true)) + _lt_pass1!(buf, ctl, red, redi, li, Val(_LT_S1), prm, Val(true)) @synchronize end mb = _bm_gmember(bm, G) @@ -844,7 +1111,7 @@ function _launch_subtrees_ldlt!(N::Numeric{T}, S::Symbolic, nzval, first, count, kernel!(N.factor, N.stack, N.info, N.stats, N.d, N.piv, N.pivot_kind, N.psign, S.perm, N.aux, nzval, S.amap, S.amap_ptr, S.amap_src, S.group_nodes, bm, S.subtree_ptr, S.subtree_nodes, S.super_ptr, S.front_ptr, S.front_nrows, S.front_ncols, S.cb_ptr, S.local_front, S.local_cb, S.child_ptr, S.child_list, - S.relind_ptr, S.relind, prm, Val((LB - SUBTREE_LOCAL_RESERVE) ÷ sizeof(T)), Val(WG); + S.relind_ptr, S.relind, prm, Val((LB - SUBTREE_LOCAL_RESERVE_LDLT) ÷ sizeof(T)), Val(WG); ndrange = WG * count * bm.nact) return nothing end diff --git a/src/reference/ldlt.jl b/src/reference/ldlt.jl index c82b792..24cb859 100644 --- a/src/reference/ldlt.jl +++ b/src/reference/ldlt.jl @@ -243,10 +243,22 @@ function _colmax(F::AbstractMatrix{T}, c::Int, k::Int, skip::Int = 0) where {T} return m end +# `d ≥ u maxᵢ |F[i, c]|` over the remaining rows k:f but c, with an early exit: `u ≥ 0` is finite, so +# `u · max = max(u · |F[i, c]|)`, and a NaN fails both forms (the device kernels test the same way) +function _threshold_ok(F::AbstractMatrix, c::Int, k::Int, u, d) + m = zero(d) + for i in k:size(F, 1) + i == c && continue + m = max(m, abs(F[i, c])) + d >= u * m || return false + end + return d >= u * m +end + # acceptable 1×1 pivot at column c: not tiny and passing the threshold test function _accept_1x1(F, c, k, prm::_LDLTParams) d = abs(F[c, c]) - return d >= prm.eps && d >= prm.u * _colmax(F, c, k) + return d >= prm.eps && _threshold_ok(F, c, k, prm.u, d) end # the 2×2 block on columns (k, r) is nonsingular and not tiny @@ -264,12 +276,17 @@ function _accept_2x2(F, k, r, prm::_LDLTParams) return prm.u * (e11 * m1 + e21 * m2) <= 1 && prm.u * (e12 * m1 + e22 * m2) <= 1 end -# the acceptable 1×1 pivot of the block k:w with the largest |aⱼⱼ| (0 if none) +# the acceptable 1×1 pivot of the block k:w with the largest |aⱼⱼ| (0 if none; the first on a tie); +# a column that cannot beat the best so far is not tested function _best_1x1(F, k, w, prm::_LDLTParams) best = 0 + bv = zero(real(eltype(F))) for j in k:w - _accept_1x1(F, j, k, prm) || continue - (best == 0 || abs(F[j, j]) > abs(F[best, best])) && (best = j) + a = abs(F[j, j]) + a >= prm.eps || continue + best != 0 && !(a > bv) && continue + _threshold_ok(F, j, k, prm.u, a) || continue + best, bv = j, a end return best end diff --git a/src/solver.jl b/src/solver.jl index 291e34d..3635f79 100644 --- a/src/solver.jl +++ b/src/solver.jl @@ -376,7 +376,7 @@ function _symbolic!(solver::DirectSolver{T, INT}) where {T, INT} P = SymmetricPattern(solver.host_rowptr, solver.host_colval, A.nrows, solver.structure; view = _stored_view(solver), index = A.index) sp = supernode_partition(factor_pattern(P, ord), ord.perm, opts) - sc = build_schedule(sp, opts, T) + sc = build_schedule(sp, opts, T; reserve = subtree_local_reserve(solver.structure)) layout = build_layout(sp, sc) Sh = Symbolic(sp, sc, layout, solver.host_rowptr, solver.host_colval, A.nrows, solver.structure; view = _stored_view(solver), index = A.index) diff --git a/src/symbolic/maps.jl b/src/symbolic/maps.jl index 8c5cdf0..6a68282 100644 --- a/src/symbolic/maps.jl +++ b/src/symbolic/maps.jl @@ -245,7 +245,7 @@ function symbolic_analysis(A::CSR, structure, view = VIEW_FULL; opts::Options = pp = analysis_pairs(P, rowptr, colval, A.nzval, A.nrows, structure, opts; view, index = A.index) ord = compute_ordering(P, opts; T, pp.pairs, pp.candidates) sp = supernode_partition(factor_pattern(P, ord), ord.perm, opts) - sc = build_schedule(sp, opts, T) + sc = build_schedule(sp, opts, T; reserve = subtree_local_reserve(structure)) layout = build_layout(sp, sc) return Symbolic(sp, sc, layout, rowptr, colval, A.nrows, structure; view, index = A.index) end diff --git a/src/symbolic/schedule.jl b/src/symbolic/schedule.jl index 5f35560..7439635 100644 --- a/src/symbolic/schedule.jl +++ b/src/symbolic/schedule.jl @@ -39,6 +39,23 @@ const SUBTREE_LOCAL_SIZES = (8192, 16384, 32768, 49152) "Bytes of a regime-A kernel's local memory kept for its control words and pivot (not for fronts)." const SUBTREE_LOCAL_RESERVE = 256 +""" +Bytes of the regime-A LDLᵀ/LDLᴴ kernel's local memory kept for its control words, pivot block and the +reduction slots of the cooperative pivot search (structures `"S"`/`"H"`; at most 944 bytes for +`ComplexF64` with `Int64` maps). +""" +const SUBTREE_LOCAL_RESERVE_LDLT = 1024 + +""" + subtree_local_reserve(structure) -> Int + +[`SUBTREE_LOCAL_RESERVE_LDLT`](@ref) for the LDLᵀ/LDLᴴ structures `"S"`/`"H"`, else +[`SUBTREE_LOCAL_RESERVE`](@ref). +""" +subtree_local_reserve(structure) = + _structure(structure) in (STRUCTURE_SYMMETRIC, STRUCTURE_HERMITIAN) ? SUBTREE_LOCAL_RESERVE_LDLT : + SUBTREE_LOCAL_RESERVE + """ subtree_local_bytes(budget) -> Int @@ -50,14 +67,15 @@ instances serves every budget. subtree_local_bytes(budget::Integer) = foldl((acc, b) -> b <= budget ? b : acc, SUBTREE_LOCAL_SIZES; init = 0) """ - subtree_capacity(budget, elsize) -> Int + subtree_capacity(budget, elsize, reserve = SUBTREE_LOCAL_RESERVE) -> Int Entries of `elsize` bytes a regime-A subtree may keep in local memory under -`budget`: [`subtree_local_bytes`](@ref) minus `SUBTREE_LOCAL_RESERVE`, divided -by `elsize` (`0` when the budget is too small). +`budget`: [`subtree_local_bytes`](@ref) minus `reserve` (the kernel's other local +memory, [`subtree_local_reserve`](@ref)), divided by `elsize` (`0` when the +budget is too small). """ -subtree_capacity(budget::Integer, elsize::Integer) = - max(subtree_local_bytes(budget) - SUBTREE_LOCAL_RESERVE, 0) ÷ Int(elsize) +subtree_capacity(budget::Integer, elsize::Integer, reserve::Integer = SUBTREE_LOCAL_RESERVE) = + max(subtree_local_bytes(budget) - Int(reserve), 0) ÷ Int(elsize) """ packed_length(m) -> Int @@ -244,7 +262,8 @@ function front_flops(f::Integer, w::Integer) end """ - build_schedule(sp::SupernodePartition, opts::Options = Options(), ::Type{T} = Float64) -> Schedule + build_schedule(sp::SupernodePartition, opts::Options = Options(), ::Type{T} = Float64; + reserve = SUBTREE_LOCAL_RESERVE) -> Schedule Regime assignment, binning, levels and launch groups (PLAN §2.3 step 5) for the supernodes of `sp`, with byte budgets for the element type `T`: @@ -255,7 +274,8 @@ the supernodes of `sp`, with byte budgets for the element type `T`: and its subtree's serial stack peak (packed lower-triangular `f×f` fronts plus the packed `m×m` contribution blocks waiting for their parent, `m = f - w`) fits the [`subtree_capacity`](@ref) of the largest of - `opts.subtree_budgets`, and its subtree's flops ([`front_flops`](@ref)) + `opts.subtree_budgets` (local memory minus `reserve` bytes, + [`subtree_local_reserve`](@ref) of the structure), and its subtree's flops ([`front_flops`](@ref)) are at most `total ÷ opts.subtree_parallelism` (`0`: no flop limit), so that a large tree of small fronts is split into enough subtrees to fill the device instead of running on one workgroup; the maximal @@ -269,12 +289,13 @@ the supernodes of `sp`, with byte budgets for the element type `T`: (packed) contribution-block bytes stay `≤ opts.memory_budget` (a single front larger than the budget gets its own chunk; a negative budget means no chunking). """ -function build_schedule(sp::SupernodePartition, opts::Options = Options(), ::Type{T} = Float64) where {T} +function build_schedule(sp::SupernodePartition, opts::Options = Options(), ::Type{T} = Float64; + reserve::Integer = SUBTREE_LOCAL_RESERVE) where {T} ns = nsupernodes(sp) elsize = sizeof(T) cw, cr = opts.regime_c_width, opts.regime_c_rows budgets = sort(opts.subtree_budgets) - capacity = [subtree_capacity(b, elsize) for b in budgets] + capacity = [subtree_capacity(b, elsize, reserve) for b in budgets] maxcap = isempty(budgets) ? 0 : maximum(capacity) alg = opts.factorization_alg width = [snwidth(sp, s) for s in 1:ns] diff --git a/test/test_symbolic_schedule.jl b/test/test_symbolic_schedule.jl index 400d97d..c847db6 100644 --- a/test/test_symbolic_schedule.jl +++ b/test/test_symbolic_schedule.jl @@ -73,7 +73,7 @@ function check_schedule(S::SDS.Symbolic) nodes = sc.subtree_nodes[sc.subtree_ptr[t]:(sc.subtree_ptr[t + 1] - 1)] ok &= nodes[end] == sc.subtree_root[t] && all(v -> sc.subtree[v] == t, nodes) ok &= simulated_subtree_peak(sc, sp, nodes) == sc.subtree_peak[t] - cap(c) = SDS.subtree_capacity(sc.budgets[c], sc.elsize) * sc.elsize + cap(c) = SDS.subtree_capacity(sc.budgets[c], sc.elsize, SDS.subtree_local_reserve(S.structure)) * sc.elsize ok &= sc.subtree_peak[t] <= cap(sc.subtree_class[t]) <= sc.budgets[sc.subtree_class[t]] ok &= sc.subtree_class[t] == 1 || sc.subtree_peak[t] > cap(sc.subtree_class[t] - 1) # local layout: every front inside the peak, contribution blocks below their parent's front