Cholesky honors pivot_epsilon: static perturbation of near-zero pivots

bench/fact_fused.jl

@@ -411,6 +411,7 @@ function _sorted_subtree_launch(N::SDS.Numeric{T}, S, nzval, list, ::Val{LB}, ::
     k!(N.factor, N.stack, N.info, nzval, S.amap, S.amap_ptr, S.amap_src, list, bm, S.subtree_ptr,
        S.subtree_nodes, 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,
+       real(T)(SDS.resolved_pivot_epsilon(SDS.Options(), real(T))),   # production default (#127)
        Val((LB - SDS.SUBTREE_LOCAL_RESERVE) ÷ sizeof(T)), Val(WG); ndrange = WG * length(list) * bm.nact)
     return nothing
 end

bench/fact_split.jl

@@ -248,6 +248,7 @@ function _sorted_subtree_launch2(N::SDS.Numeric{T}, S, nzval, list, bm, k!, ::Va
     k!(N.factor, N.stack, N.info, nzval, S.amap, S.amap_ptr, S.amap_src, list, bm, S.subtree_ptr,
        S.subtree_nodes, 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,
+       real(T)(SDS.resolved_pivot_epsilon(SDS.Options(), real(T))),   # production default (#127)
        Val((LB - SDS.SUBTREE_LOCAL_RESERVE) ÷ sizeof(T)), Val(WG); ndrange = WG * length(list) * bm.nact)
     return nothing
 end

src/dense/fallback/potrf.jl

@@ -113,15 +113,19 @@ function ka_chol_diag_info!(info::AbstractVector{Int32}, L::AbstractArray)
     return info
 end
 
-@kernel function _ka_chol_check_kernel!(info, A, n, idx)
+@kernel function _ka_chol_check_kernel!(info, A, n, idx, dmin)
     li = @index(Local, Linear)
     b = @index(Group, Linear)
     if li == 1
         @inbounds if info[idx + b - 1] == 0
             st = Int32(0)
             for j in 1:n
                 d = _get(A, j, j, b)
-                if !(isfinite(real(d)) && real(d) > 0 && iszero(imag(d)))
+                # dmin > 0: the Cholesky pivot-epsilon floor (issue #127) — potrf accepts
+                # any positive pivot, so a pivot that became tiny DURING elimination is
+                # only visible here, as a factor diagonal at or below sqrt(epsilon)
+                ok = dmin > zero(dmin) ? real(d) >= dmin : real(d) > zero(real(d))
+                if !(isfinite(real(d)) && ok && iszero(imag(d)))
                     st = Int32(j)
                     break
                 end
@@ -140,10 +144,10 @@ the first `j` whose diagonal entry `L[j, j]` is not a finite positive real
 (still 0 if there is none); a nonzero `info[idx]` is kept. Asynchronous (see
 [`potrf_info!`](@ref)).
 """
-function ka_chol_check_info!(info::AbstractVector{Int32}, idx::Integer, L::AbstractMatrix)
+function ka_chol_check_info!(info::AbstractVector{Int32}, idx::Integer, L::AbstractMatrix, dmin::Real = 0.0)
     1 <= idx <= length(info) || throw(DimensionMismatch("info index $idx outside 1:$(length(info))"))
     kernel! = _ka_chol_check_kernel!(KernelAbstractions.get_backend(L), 1)
-    kernel!(info, L, size(L, 1), Int(idx); ndrange = 1)
+    kernel!(info, L, size(L, 1), Int(idx), real(eltype(L))(dmin); ndrange = 1)
     return info
 end
 
@@ -153,12 +157,13 @@ end
 Batched form: the check of `ka_chol_check_info!(info, offset + b, L[:, :, b])`
 for every member `b` of the 3-D strided batch `L`, one launch.
 """
-function ka_chol_check_info!(info::AbstractVector{Int32}, offset::Integer, L::AbstractArray{<:Any, 3})
+function ka_chol_check_info!(info::AbstractVector{Int32}, offset::Integer, L::AbstractArray{<:Any, 3},
+                             dmin::Real = 0.0)
     nb = size(L, 3)
     0 <= offset && offset + nb <= length(info) ||
         throw(DimensionMismatch("info has length $(length(info)) < offset $offset + batch count $nb"))
     nb == 0 && return info
     kernel! = _ka_chol_check_kernel!(KernelAbstractions.get_backend(L), 1)
-    kernel!(info, L, size(L, 1), Int(offset) + 1; ndrange = nb)
+    kernel!(info, L, size(L, 1), Int(offset) + 1, real(eltype(L))(dmin); ndrange = nb)
     return info
 end

src/dense/interface.jl

@@ -287,16 +287,19 @@ function potrf_info!(uplo, A::AbstractMatrix, info::AbstractVector{Int32}, idx::
 end
 
 # `potrf_info!` with an already resolved implementation `p` (no capability lookup)
-function _potrf_info_impl!(p::Symbol, ul::Char, A::AbstractMatrix, info::AbstractVector{Int32}, idx::Integer)
+function _potrf_info_impl!(p::Symbol, ul::Char, A::AbstractMatrix, info::AbstractVector{Int32}, idx::Integer,
+                           dmin::Real = 0.0)
     if p === :vendor
         vendor_potrf_info!(ul, A, info, idx)
-        ka_chol_check_info!(info, idx, A)
+        ka_chol_check_info!(info, idx, A, dmin)
     elseif p === :generic
         cholesky!(Hermitian(A, ul == 'L' ? :L : :U), NoPivot(); check = false)
         fill!(view(info, idx:idx), Int32(0))
-        ka_chol_check_info!(info, idx, A)
+        ka_chol_check_info!(info, idx, A, dmin)
     else
         ka_potrf!(ul, A, info; offset = idx - 1)
+        # ka_potrf validates its own pivots at > 0; the epsilon floor needs the pass too
+        dmin > 0 && ka_chol_check_info!(info, idx, A, dmin)
     end
     return info
 end

src/fused.jl

@@ -992,7 +992,7 @@ end
                                     cb_done, @Const(snparent), @Const(front_ptr), @Const(front_nrows),
                                     @Const(front_ncols), @Const(cb_ptr), @Const(child_ptr),
                                     @Const(child_list), @Const(relind_ptr), @Const(relind), base,
-                                    @Const(wbase), wupd, wpdone2, wtrsml, wcbl, scap, eac, d, @Const(super_ptr),
+                                    @Const(wbase), wupd, wpdone2, wtrsml, wcbl, scap, eac, epsc, d, @Const(super_ptr),
                                     ::Val{LDL}, ::Val{WG}) where {LDL, WG}
     @uniform TT = eltype(factor)
     @uniform RT = real(eltype(factor))
@@ -1103,13 +1103,13 @@ end
                 if LDL
                     _front_ldl_update!(sh, st, piv, j, v, li, front_ncols, d, g0, Val(WG))
                 else
-                    SDS._front_chol_update!(sh, st, piv, j, v, li, front_ncols, Val(64), Val(WG))
+                    SDS._front_chol_update!(sh, st, piv, j, v, li, front_ncols, epsc, Val(64), Val(WG))
                 end
                 @synchronize
                 if LDL
                     _front_ldl_scale!(sh, st, piv, j, v, li, front_ncols, Val(WG))
                 else
-                    SDS._front_chol_scale!(sh, st, piv, j, v, li, front_ncols, Val(64), Val(WG))
+                    SDS._front_chol_scale!(sh, st, piv, j, v, li, front_ncols, epsc, Val(64), Val(WG))
                 end
                 @synchronize
             end
@@ -1365,8 +1365,10 @@ end
                             piv[1] = one(RT)
                         end
                     else
-                        if st[1] == Int32(0) && dj > zero(RT)
-                            piv[1] = sqrt(dj)
+                        # same pivot-epsilon semantics as the narrow kernels (#127):
+                        # [0, epsc] perturbs to epsc, negative/NaN fails the front
+                        if st[1] == Int32(0) && (dj > epsc || (epsc > zero(RT) && dj >= zero(RT)))
+                            piv[1] = sqrt(max(dj, epsc))
                             sh[SDS._packed(j, j, 64)] = piv[1]
                         elseif st[1] == Int32(0)
                             st[1] = Int32(c0 + j)
@@ -1765,13 +1767,13 @@ end
 end
 
 # subtree launcher with a custom (sorted) list, local memory sized by eltype
-function launch_subtrees_sorted!(N, S, nzval, list, lb, backend, ldl::Bool = false)
+function launch_subtrees_sorted!(N, S, nzval, list, lb, backend, ldl::Bool = false, epsc::Real = 0.0)
     SDS._with_local_bytes(Int(lb)) do LBv
-        _sorted_subtree_launch(N, S, nzval, list, Val(ldl), LBv, Val(SDS.SUBTREE_WORKGROUP), backend)
+        _sorted_subtree_launch(N, S, nzval, list, Val(ldl), LBv, Val(SDS.SUBTREE_WORKGROUP), backend, epsc)
     end
 end
 function _sorted_subtree_launch(N::SDS.Numeric{T}, S, nzval, list, ::Val{LDL}, ::Val{LB}, ::Val{WG},
-                                backend) where {T, LDL, LB, WG}
+                                backend, epsc::Real = 0.0) where {T, LDL, LB, WG}
     bm = SDS.batch_map(N; first = 1)
     if LDL
         # the schedule sized the "S" budget classes with the LDLT reserve: match it
@@ -1785,7 +1787,7 @@ function _sorted_subtree_launch(N::SDS.Numeric{T}, S, nzval, list, ::Val{LDL}, :
         k! = SDS.subtree_cholesky_kernel!(backend, WG)
         k!(N.factor, N.stack, N.info, nzval, S.amap, S.amap_ptr, S.amap_src, list, bm, S.subtree_ptr,
            S.subtree_nodes, 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,
+           S.child_ptr, S.child_list, S.relind_ptr, S.relind, real(T)(epsc),
            Val((LB - SDS.SUBTREE_LOCAL_RESERVE) ÷ sizeof(T)), Val(WG); ndrange = WG * length(list) * bm.nact)
     end
     return nothing
@@ -2021,6 +2023,9 @@ function build_fused_plan(s; ldl::Bool = false, inkernel_wides::Union{Nothing, B
             wcb_b = _dvec(backend, wcb_b), wupd = KA.zeros(backend, Int32, max(nwslots, 1)),
             wtl = KA.zeros(backend, Int32, max(nwslots, 1)), wpd = KA.zeros(backend, Int32, max(nwslots, 1)),
             wcb = KA.zeros(backend, Int32, ns), wgf, scap = Int32(scap), eac = Int32(ea_chunk), ldl,
+            # captured at plan build: a later setparam!("pivot_epsilon") changes the stock
+            # path but not an already-built fused plan (rebuild the plan to change it)
+            chol_eps = ldl ? 0.0 : Float64(SDS.resolved_pivot_epsilon(s.options, real(T))),
             ldl_potrf = ldl && wides_potrf,
             sptr_h = Array(S.super_ptr),
             ldlw = KA.zeros(backend, T,
@@ -2050,7 +2055,7 @@ function refact_fused!(s, nzval, fp)
     fill!(fp.wpd, Int32(0))
     copyto!(fp.wcb, fp.wcb_b)
     for k in eachindex(plan.sub_first)
-        launch_subtrees_sorted!(N, S, nzval, fp.ssub[k], plan.sub_local[k], backend, fp.ldl)
+        launch_subtrees_sorted!(N, S, nzval, fp.ssub[k], plan.sub_local[k], backend, fp.ldl, fp.chol_eps)
     end
     nodes = S.schedule.group_nodes
     fill!(fp.pcb, zero(eltype(fp.pcb)))
@@ -2070,7 +2075,8 @@ function refact_fused!(s, nzval, fp)
                                              S.snparent, S.front_ptr, S.front_nrows, S.front_ncols,
                                              S.cb_ptr, S.child_ptr, S.child_list, S.relind_ptr, S.relind,
                                              Int32(base), fp.wbase, fp.wupd, fp.wpd, fp.wtl, fp.wcb,
-                                             fp.scap, fp.eac, N.d, S.super_ptr, Val(fp.ldl), Val(fp.wgf);
+                                             fp.scap, fp.eac, real(eltype(N.factor))(fp.chol_eps), N.d,
+                                             S.super_ptr, Val(fp.ldl), Val(fp.wgf);
                                              ndrange = fp.wgf * count)
         for v in gwides                              # host-wide fronts at the segment boundary
             if fp.ldl && fp.ldl_potrf
@@ -2079,7 +2085,7 @@ function refact_fused!(s, nzval, fp)
                 _factor_panel_c_ldl!(N, S, fp, v, p, backend)
             else
                 SDS._factor_panel_c!(N.factor, fp.pcb, N.work, N.info, v, S.layout.panel_ptr[v],
-                                     S.schedule.rows[v], S.schedule.width[v], fp.cb2_h[v], p)
+                                     S.schedule.rows[v], S.schedule.width[v], fp.cb2_h[v], p, fp.chol_eps)
             end
             signal_kernel!(backend, 32)(fp.cbd, fp.arrived, S.snparent, v; ndrange = 32)
         end

src/numeric/factorize.jl

@@ -109,11 +109,11 @@ end
 
 # regime-C dense kernels of front `s` (assembled panel + contribution block) of every active batch member,
 # `p` from `_front_impls`
-function _factor_front_c!(N::Numeric{T}, S::Symbolic, s::Int, p::NamedTuple) where {T}
+function _factor_front_c!(N::Numeric{T}, S::Symbolic, s::Int, p::NamedTuple, epsc::Real = 0.0) where {T}
     L, sc = S.layout, S.schedule
     if N.nbatch == 1
         _factor_panel_c!(N.factor, N.stack, N.work, N.info, s, L.panel_ptr[s], sc.rows[s], sc.width[s], L.cb_ptr[s],
-                         p)
+                         p, epsc)
         return nothing
     end
     plan = N.plan
@@ -125,10 +125,19 @@ function _factor_front_c!(N::Numeric{T}, S::Symbolic, s::Int, p::NamedTuple) whe
             for j in j0:j1
                 k = Int(plan.members_host[j])
                 _factor_panel_c!(N.factor, _mview(N.stack, k, nb), N.work, N.info, (s - 1) * nb + k,
-                                 panel_offset(L.panel_ptr, s, k, nb), sc.rows[s], sc.width[s], L.cb_ptr[s], p)
+                                 panel_offset(L.panel_ptr, s, k, nb), sc.rows[s], sc.width[s], L.cb_ptr[s], p, epsc)
             end
         else
-            _factor_panels_c_batched!(N, S, s, j0, j1 - j0 + 1, p)
+            if epsc > 0
+                backend = KernelAbstractions.get_backend(N.factor)
+                for j in j0:j1
+                    k = Int(plan.members_host[j])
+                    _chol_diag_clamp_kernel!(backend, 64)(N.factor, panel_offset(L.panel_ptr, s, k, nb),
+                                                          sc.rows[s], sc.width[s], real(T)(epsc);
+                                                          ndrange = sc.width[s])
+                end
+            end
+            _factor_panels_c_batched!(N, S, s, j0, j1 - j0 + 1, p, epsc)
         end
     end
     return nothing
@@ -138,7 +147,7 @@ end
 # are one `f×w×count` strided batch (interleaved layout), factored by strided-batched `potrf`, `trsm` and `gemm`
 # (`F21 F21ᴴ` into `count` blocks of the workspace), then packed into the members' contribution blocks
 function _factor_panels_c_batched!(N::Numeric{T}, S::Symbolic, s::Int, j0::Int, count::Int,
-                                   p::NamedTuple) where {T}
+                                   p::NamedTuple, epsc::Real = 0.0) where {T}
     L, sc = S.layout, S.schedule
     nb = N.nbatch
     f, w = sc.rows[s], sc.width[s]
@@ -149,12 +158,14 @@ function _factor_panels_c_batched!(N::Numeric{T}, S::Symbolic, s::Int, j0::Int,
     F11 = view(P3, 1:w, 1:w, :)
     ioff = (s - 1) * nb + k0 - 1                              # info of front s, member k0 (interleaved)
     ptrs = N.plan.vendor_ptrs[s]                              # (F11, F21) member pointers, or nothing
+    dmin = sqrt(real(T)(epsc))                                # epsilon floor (#127); 0 keeps today's check
     if p.bpotrf === :vendor && ptrs !== nothing
         vendor_potrf_batched_ptrs!('L', w, view(ptrs[1], k0:(k0 + count - 1)), f,
                                    view(N.info, (ioff + 1):(ioff + count)), count)
-        ka_chol_check_info!(N.info, ioff, F11)
+        ka_chol_check_info!(N.info, ioff, F11, dmin)
     else
         _potrf_batched_impl!(p.bpotrf, F11, N.info, ioff)
+        epsc > 0 && ka_chol_check_info!(N.info, ioff, F11, dmin)
     end
     m > 0 || return nothing
     F21 = view(P3, (w + 1):f, 1:w, :)
@@ -175,13 +186,36 @@ end
 # `potrf`/`trsm`/`syrk` (`herk`) on the `f×w` panel at `factor[p0]`; the `syrk` result goes to the
 # `m×m` workspace `work[1:m^2]` and is added to the packed contribution block at `stack[c0]`
 # (`c0 = 0`: none); status into `info[s]` (also used by bench/front_bins.jl)
+# pivots in [0, epsc] of the still-unfactored panel diagonal are raised to epsc before
+# potrf (the cuDSS pivot_epsilon static perturbation on the vendor/host Cholesky path);
+# nonpositive diagonals are left for potrf to reject, as in the kernel paths
+@kernel function _chol_diag_clamp_kernel!(factor, p0::Int, f::Int, w::Int, epsc)
+    j = @index(Global)
+    @inbounds if j <= w
+        idx = p0 + (j - 1) * f + j - 1
+        d = real(factor[idx])
+        # [0, epsc]: the kernel routes perturb an exact zero too (cuDSS semantics)
+        (d >= zero(d)) & (d <= epsc) && (factor[idx] = epsc)
+    end
+end
+
 function _factor_panel_c!(factor::AbstractVector{T}, stack::AbstractVector{T}, work::AbstractVector{T},
                           info::AbstractVector{Int32}, s::Int, p0::Int, f::Int, w::Int, c0::Int,
-                          p::NamedTuple) where {T}
+                          p::NamedTuple, epsc::Real = 0.0) where {T}
     m = f - w
     P = reshape(view(factor, p0:(p0 + f * w - 1)), f, w)
     F11 = view(P, 1:w, 1:w)
-    _potrf_info_impl!(p.potrf, 'L', F11, info, s)
+    if epsc > 0
+        # same-queue ordering covers every potrf impl (the CPU backend already ran
+        # LAPACK straight after the assembly kernels before this PR): no host sync
+        _chol_diag_clamp_kernel!(KernelAbstractions.get_backend(factor), 64)(factor, p0, f, w, real(T)(epsc);
+                                                                             ndrange = w)
+    end
+    # the clamp handles input-diagonal tininess; a pivot that becomes tiny DURING the
+    # elimination is caught after the fact by the sqrt(epsilon) factor-diagonal floor
+    # inside _potrf_info_impl!, reported through info (the potrf route DETECTS where the
+    # kernel routes perturb -- either way MadNLP gets its regularization signal, #127)
+    _potrf_info_impl!(p.potrf, 'L', F11, info, s, sqrt(real(T)(epsc)))
     m > 0 || return nothing
     F21 = view(P, (w + 1):f, 1:w)
     _trsm_impl!(p.trsm, 'R', 'L', 'C', 'N', one(T), F11, F21)
@@ -255,23 +289,27 @@ function factorize_cholesky!(N::Numeric{T}, S::Symbolic, nzval::AbstractVector;
     plan = N.plan
     nodes = S.schedule.group_nodes
     flag = opts.user_host_interrupt
+    # cuDSS pivot_epsilon semantics on the Cholesky path (issue #127): pivots in
+    # (0, epsc] are statically perturbed to epsc so a near-singular SPD system
+    # factors with bounded growth; nonpositive/NaN pivots still fail the front
+    epsc = resolved_pivot_epsilon(opts, real(T))
     for k in eachindex(plan.sub_first)
         _poll_interrupt(flag)
         a, b = plan.sub_first[k], plan.sub_last[k]
-        factorize_subtrees!(N, S, nzval, a, b - a + 1, plan.sub_local[k])
+        factorize_subtrees!(N, S, nzval, a, b - a + 1, plan.sub_local[k], epsc)
     end
     for k in eachindex(plan.group_first)
         _poll_interrupt(flag)
         a, b = plan.group_first[k], plan.group_last[k]
         W = plan.group_width[k]
         if W > 0
-            factorize_fronts_b!(N, S, nzval, a, b - a + 1, plan.group_maxchild[k], W)
+            factorize_fronts_b!(N, S, nzval, a, b - a + 1, plan.group_maxchild[k], W, epsc)
         else
             zero_fronts!(N, S, a, b - a + 1)
             scatter_A!(N, S, nzval, a, b - a + 1)
             extend_add!(N, S, a, b - a + 1, plan.group_maxchild[k])
             for q in a:b
-                _factor_front_c!(N, S, nodes[q], p)
+                _factor_front_c!(N, S, nodes[q], p, epsc)
             end
         end
     end