Quantized versions of Jacobi and Identity preconditioners - #2869
Open
pcarruscag wants to merge 26 commits into
Open
Quantized versions of Jacobi and Identity preconditioners#2869pcarruscag wants to merge 26 commits into
pcarruscag wants to merge 26 commits into
Conversation
Extends the Q_LU_SGS quantization scheme to JACOBI and IDENTITY: the matrix-vector product used by the Krylov solver gets int8-quantized off-diagonal storage, while each preconditioner's own apply step stays at full precision (Jacobi's inverse diagonal is always computed from the full-precision diagonal; identity is a no-op either way). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Q_JACOBI and Q_IDENTITY no longer force quantized_mode off under CUDA: their shared matrix-vector product has no sequential dependency (unlike Q_LU_SGS's sweeps, which stay host-only), so it can run on the device. Adds device mirrors of the quantized off-diagonal storage (refreshed by HtDTransfer) and a new QuantizedBlockLDU_SpMV_kernel mirroring the host QuantizedRowProduct. The quantized diagonal is populated by a new QuantizeDiagonalBlocksGPU() kernel that reads directly from gpu.d rather than transferring the host-quantized diagonal, since diagonal quantization only happens in Build(), which runs after HtDTransfer has already uploaded the matrix for the solve. Validated against the host path on a NACA0012 inviscid case (RTX 4070 Ti): GPU quantized-Jacobi results match CPU quantized-Jacobi exactly, and GPU plain-Jacobi is unaffected. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
…h L/U upload Templatize the existing LDU struct on its value type and reuse it for the quantized storage: q_scale/q_blocks (host) and d_q_scale/d_q_blocks (device) replace the twelve loose QuantType* members, cutting the constructor/destructor/Initialize boilerplate down to a couple of generic-lambda calls per group. Also change how the quantized diagonal reaches the device: instead of a GPU kernel that quantized straight from gpu.d, the diagonal is always quantized on the host in QuantizeDiagonalBlocks() and then uploaded. HtDTransfer() kicks off the (larger) L/U quantized transfer via cudaMemcpyAsync before Build() is reached, so it can overlap with that host-side diagonal quantization instead of blocking on it first; same default stream ordering keeps any later kernel launch correct without an extra explicit sync. Re-validated against the pre-refactor build on the NACA0012 case (RTX 4070 Ti): CPU/GPU x plain/quantized Jacobi and GPU quantized-identity all reproduce identical CL/CD to the prior commit. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
CQuantizedLUSGSPreconditioner, CQuantizedJacobiPreconditioner and CQuantizedIdentityPreconditioner each differed from their plain counterpart only by an extra QuantizeDiagonalBlocks() call in Build(). Since that call is already a no-op unless the matrix was actually set up for quantization, push it into the plain wrappers instead: - CLU_SGSPreconditioner/CIdentityPreconditioner now call QuantizeDiagonalBlocks() directly in Build(). - CSysMatrix::BuildJacobiPreconditioner() calls QuantizeDiagonalBlocks() itself (before the jacobi_on_device branch, since diagonal quantization is independent of invM and jacobi_on_device returns early), so CJacobiPreconditioner needs no change at all. This removes three duplicate classes; the factory now maps Q_LU_SGS/ Q_JACOBI/Q_IDENTITY to the same wrapper as their plain counterpart, and CSysSolve's discrete-adjoint transpose-rebuild switch merges Q_JACOBI into the JACOBI/LINELET case. Re-validated on the NACA0012 case: LU_SGS/Q_LU_SGS (host) and CPU/GPU x plain/quantized Jacobi and GPU quantized-identity all reproduce the same CL/CD as before this refactor. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
QuantizeDiagonalBlocksGPU() had shrunk to two cudaMemcpy calls with no kernel involved, so it doesn't need the usual CSysMatrixGPU.cu hook indirection: cudaMemcpy (unlike a kernel launch) is a plain CUDA runtime call, already available in CSysMatrix.cpp transitively via allocation_toolbox.hpp -> GPUComms.cuh, the same way Initialize() already calls cudaMalloc/cudaMemcpy directly for the other GPU buffers. Inlines the upload at the end of QuantizeDiagonalBlocks() and removes the now-empty hook (declaration, .cu implementation, and its INSTANTIATE_MATRIX entry). Re-validated on the NACA0012 case: GPU quantized-Jacobi and quantized- identity reproduce the same CL/CD as before this change. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
…y overlaps cudaMemcpyAsync is only genuinely asynchronous (non-blocking on the host) when the host-side buffer is page-locked; from regular pageable memory (what q_scale/q_blocks were using) the driver silently falls back to a synchronous staged copy, so the async transfer added earlier was never actually overlapping with the host quantizing the diagonal. Adds GPUMemoryAllocation::pinned_alloc/pinned_free (cudaMallocHost/ cudaFreeHost) and switches q_scale/q_blocks to use them instead of aligned_alloc whenever useCuda, matched by conditional frees in the destructor. Also makes the diagonal upload async (now pinned too) and adds an explicit cudaStreamSynchronize after it, matching the pattern every other GPU-touching function in this file already follows, rather than relying solely on implicit default-stream ordering. Verified with nsys (gputrace) that the transfers now show up with SrcMemKd=Pinned. Also chased a small (~0.01%) CL/CD difference this introduced for the GPU quantized-identity case: confirmed via a CPU- only, non-CUDA run of the same case that varying OMP_NUM_THREADS alone produces a difference of the same order, so this is pre-existing floating-point non-associativity sensitivity in that (deliberately unpreconditioned) test case, not a bug in the new transfer code - plain/quantized Jacobi (which is actually preconditioned) stay bit-identical throughout. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Both BlockLDU_SpMV_kernel and QuantizedBlockLDU_SpMV_kernel launched one CUDA block per matrix row with blockDim.x == nVar (~4-6 threads), leaving most of a warp idle and capping occupancy at a few resident, mostly-empty warps per SM - latency-bound well before DRAM bandwidth becomes the limit, which is why the quantized kernel (4x fewer matrix bytes) wasn't measurably faster than the float one: shrinking the matrix doesn't help a kernel that isn't bandwidth-bound in the first place, and int8 element loads cost the same one-load-per-element as the float32 ones did. Rows are now batched into blocks of ~128 threads (threadIdx.x maps to (row-within-block, output variable) via divmod by nVar), the same batching for both kernels since it's the float kernel's launch shape that needs fixing first. The quantized kernel additionally reads each row of quantized mantissas 4 bytes at a time (memcpy into a uint32_t, row bytes aren't generally 4-byte aligned) instead of one byte at a time; true dp4a doesn't apply since only the matrix is quantized, not the vector, so accumulation stays one scalar FMA per element in the same order as before (bit-identical results) - the win is fewer, wider load instructions. Verified bit-identical results against the pre-change build on the NACA0012 case (nVar=4, full vectorized path) and a quick 5-iteration ONERAM6 3D run (nVar=5, exercises the vectorized-load remainder loop), GPU quantized-Jacobi matching the CPU reference exactly in both cases. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
ApplyJacobiPreconditionerKernel already had good occupancy (128 threads/block, one thread per point), so the SpMV row-batching fix doesn't apply here as-is. Its actual problem is coalescing: each thread read a whole nVar x nVar invM block serially, so at any given load instruction consecutive threads (consecutive points) were nVar^2 elements apart - e.g. 200 bytes for nVar=5, double - with essentially no sharing of the 32-byte memory sectors. Switches to the same thread mapping as the SpMV kernels (threadIdx.x -> (point-within-block, output variable) via divmod by nVar), so each thread computes one output row via its own dot-product loop rather than the whole block. This shrinks the inter-thread stride to nVar elements (a real, if partial, coalescing improvement - not the full stride-1 coalescing a one-thread-per-block-entry-plus-reduction version would get) while adding no shared memory or synchronization, since occupancy was never the problem here. Verified bit-identical to the pre-change kernel: NACA0012 (nVar=4) GPU Jacobi/Q_Jacobi runs reproduce their prior CL/CD exactly, and a 3D ONERAM6 (nVar=5) GPU Jacobi run matches its CPU reference exactly. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
For completeness, compare against the thread-per-row version just committed: threadIdx.x now maps to (point-within-block, iVar, jVar), so consecutive threads read consecutive elements of invM (stride 1, fully coalesced across the whole block, not just within one point, since points are stored back to back) instead of landing nVar elements apart. Each thread computes one product into shared memory; the jVar==0 thread of each row then sums its nVar partials in the same order as the scalar loop it replaces (bit-identical result). Points are still batched per block for occupancy. The cost is a shared-memory write/read and one __syncthreads() per call, which the thread-per-row version didn't need - not obviously a win at this problem size without measuring, hence trying both. Verified bit-identical to the prior kernel: NACA0012 (nVar=4) GPU Jacobi/Q_Jacobi reproduce their known CL/CD exactly, and a 3D ONERAM6 (nVar=5) GPU Jacobi run matches its CPU reference exactly. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This reverts commit 56fe054.
InvertDiagonalBlocksKernel launched one block per row with blockDim.x == nVar*nVar (16-36 threads for nVar=4-6), under a full warp, wasting lanes the same way the SpMV/apply kernels did before their fix. Batches ~128 threads worth of rows per block instead (threadIdx.x -> (row-within-block, block entry) via divmod by nVar*nVar), each row getting its own slice of the dynamic shared memory buffer. The inversion itself (SU2_LinAlg::MatrixInverse, shared host/device code) stays serial on one thread per row: it's a row-oriented Gauss-Jordan with a genuine sequential dependency chain of ~nVar^2/2 elimination substeps, so spreading it across threads would trade cheap serial FLOPs (~nVar^3, tiny at this size) for that many __syncthreads() barriers - likely a net loss, and not worth it anyway since this kernel only runs once per Build(), not per Krylov iteration. Batching just fixes the launch-shape waste, same as the other kernels. Verified bit-identical: NACA0012 (nVar=4) GPU Jacobi/Q_Jacobi reproduce their known CL/CD exactly, and a 3D ONERAM6 (nVar=5) GPU Jacobi run matches its CPU reference exactly. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This reverts commit e65bfa8.
…oding logic
gpu.d is already uploaded unconditionally by HtDTransfer() (Jacobi's
own build needs the full precision diagonal regardless of
quantization), so quantize it there directly under CUDA instead of
computing on the host and uploading the int8 result - one less
host/device round trip for the diagonal specifically (L/U still have
to be quantized on the host, during assembly, since only the host ever
touches the matrix as it is being built).
To do this without duplicating the encoding algorithm (as previously
happened with a hand-rolled DecodeQuantScaleDevice), DecodeQuantScale
and EncodeQuantBlock (CSysMatrix.hpp) are now genuinely shared
__host__ __device__ code, following the same SU2_CUDA_HOST_DEVICE
pattern CMatrixInverse.hpp already established for MatrixInverse. That
macro was previously defined locally (with a matching #undef at EOF) in
three separate files - CMatrixInverse.hpp, CSysVector.hpp and
vector_expressions.hpp - which do not compose: whichever of them is
included last silently undefines it for anything appearing after, in
the same translation unit. Centralized the single definition in
code_config.hpp instead.
Getting the shared quantization code to actually behave correctly
under nvcc took two rounds of fixes, found by bisecting with a local,
device-only reimplementation as a known-good reference:
- memcpy-based float/uint32_t reinterpretation, used for the packed
scale exponent, silently misinterprets the bits when compiled for
the device (no diagnostic, just a wrong answer) - now branches on
__CUDA_ARCH__ to use __float_as_uint/__uint_as_float on device.
- std::max/std::min are not actually callable from device code without
--expt-relaxed-constexpr (not used in this build); nvcc only warns
("calling a constexpr __host__ function ... is not allowed") rather
than erroring, and silently emits broken device code. Added
QuantMax/QuantMin, which branch the same way (CUDA's own unqualified
max/min on device, std::max/std::min on host).
Both bugs were fully silent (clean compile, no crash, no NaN) and
produced a wrong-but-plausible-looking answer, so they were only
caught by comparing against the known-correct CPU/pre-refactor GPU
results - a compute-sanitizer racecheck run would have been the more
direct tool but the toolchain in this environment (2022.4.1) could not
attach to this MPI-launched binary.
Verified against CPU/pre-refactor references: NACA0012 (nVar=4)
GPU Jacobi/Q_Jacobi/Q_Identity all reproduce their known CL/CD exactly,
and a 3D ONERAM6 (nVar=5) GPU Q_Jacobi run matches its CPU counterpart
exactly.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
The reverted batching attempt (e65bfa8) floor-divided a ~128 thread target by nVar*nVar to pick rows-per-block, which for nVar values that don't divide 32 evenly (e.g. nVar=5, blockSize=25 -> 125 threads) left the block's last warp partly idle. Round rowsPerBlock up instead to the smallest multiple that makes blockDim.x itself a multiple of 32. This does not touch the actual bottleneck (MatrixInverse is a serial per-row Gauss-Jordan, so only one thread per row ever does real work regardless of alignment) - it only tidies up the parallel load-into- shared-memory phase. Kept for the learning value of confirming that, even warp-aligned, batching does not help here. Verified bit-identical: NACA0012 (nVar=4) GPU Jacobi/Q_Jacobi/Q_Identity reproduce their known CL/CD exactly, and Q_Jacobi on ONERAM6 (nVar=5, which now batches 32 rows/800 threads per block instead of 5 rows/125 threads) matches its CPU reference to 5 decimals. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
…kernel" This reverts commit 0bfd05b.
…ries InvertDiagonalBlocksKernel used nVar*nVar threads per row but only thread 0 ever did real work (SU2_LinAlg::MatrixInverse, a serial forward-elimination + back-substitution chain), the rest just cooperated on the parallel load into shared memory. Replace it with full (both-directions) Gauss-Jordan elimination: for each pivot k, every entry update A(i,j) -= A(i,k)*A(k,j) for i != k is independent of every other entry in that same step, so all nVar*nVar threads do real work every step instead of just one. This drops the serial chain from one thread doing ~nVar^3 scalar ops to nVar barrier- separated steps with the full block active. Uses __syncthreads() rather than __syncwarp() so it stays correct once nVar*nVar exceeds one warp (nVar > ~5), at the cost of a full block-wide barrier even when the whole block already fits in one warp. This is a different algorithm from the shared SU2_LinAlg::MatrixInverse (host path, and IluFactorColorKernel, both left untouched), not a device branch of it, since the elimination order genuinely differs - not just a codegen difference from __CUDA_ARCH__. Verified bit-identical to the previous serial-inversion values: NACA0012 (nVar=4) GPU Jacobi/Q_Jacobi/Q_Identity reproduce their known CL/CD exactly, GPU Q_Jacobi matches the CPU reference exactly, and ONERAM6 (nVar=5) GPU Q_Jacobi matches its CPU reference to 5 decimals (same float/double gap as before). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
QuantizeDiagonalBlocksKernel assigned one thread to a whole point's nVar*nVar block (looping over all nVar rows itself, each with its own max-abs scan and quantize pass), the coarsest-grained kernel in the file - every other kernel here (SpMV, Jacobi apply, diagonal inversion) assigns one thread to at most one block-row. Consecutive threads also landed nVar*nVar elements apart in mat_d, the same de-coalesced pattern ApplyJacobiPreconditionerKernel's row-major thread mapping fixed. Split EncodeQuantBlock's per-row body out into EncodeQuantRow (still SU2_CUDA_HOST_DEVICE, still shared - EncodeQuantBlock now just calls it once per row for the host path), since each block-row's scale and quantization are already fully independent of every other row. The kernel now assigns one thread per (point, row), batched into blocks of ~128 threads the same way as the Jacobi apply/SpMV kernels (threadIdx.x -> (point-within-block, row) via divmod by nVar). Verified bit-identical: NACA0012 (nVar=4) GPU Jacobi/Q_Jacobi/Q_Identity reproduce their known CL/CD exactly, GPU Q_Jacobi matches the CPU reference exactly, and ONERAM6 (nVar=5) GPU Q_Jacobi matches its CPU reference to 5 decimals (same float/double gap as before). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
HtDTransfer issued the async L/U copies on the default stream, the same stream QuantizeDiagonalBlocksGPU (and everything else) launches on. A single stream is strictly ordered, so despite being "async" copies, the diagonal-quantization kernel just queued behind them instead of actually running concurrently on the GPU - the copy engine and the SMs are independent hardware and there is no real data dependency between the two (the kernel only touches gpu.d and d_q_scale/blocks.d, the transfer only touches d_q_scale/blocks.l/u). Move the L/U copies to a dedicated stream (htd_stream, same pattern as the existing ilu_stream) and record an event (htd_event) at the end of them. The one default-stream consumer that actually reads the result - the quantized SpMV kernel in MatrixVectorProductGPU - waits on that event via cudaStreamWaitEvent before launching, a device-side wait that does not block the host thread, unlike the cudaStreamSynchronize calls elsewhere in this file. Verified bit-identical over 5 repeated runs (to catch any nondeterministic cross-stream race): NACA0012 (nVar=4) GPU Jacobi/Q_Jacobi/Q_Identity reproduce their known CL/CD exactly, GPU Q_Jacobi matches the CPU reference exactly, and ONERAM6 (nVar=5, several outer iterations so HtDTransfer/the event fire repeatedly) GPU Q_Jacobi matches its CPU reference to 5 decimals. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
IluFactorColorKernel's diagonal-inversion tail had the exact same shape InvertDiagonalBlocksKernel used to: nVar*nVar threads already active in the block, but only tid==0 doing the actual SU2_LinAlg::MatrixInverse work while the rest sat idle. It also already had two blockSize-sized shared buffers on hand (work, and Lij - free by that point, its last use storing it into Block_ij is already done), the exact inputs ParallelMatrixInverse needs. Factor the parallel algorithm out of InvertDiagonalBlocksKernel into a standalone ParallelMatrixInverse device function and call it from both kernels, instead of duplicating the elimination loop. Verified bit-identical (including against a pre-change build of the same config) over 5 repeated runs of GPU ILU (to catch any nondeterministic issue from reusing Lij's shared memory as the second buffer): NACA0012 (nVar=4) GPU Jacobi/Q_Jacobi/Q_Identity/ILU all reproduce their known/previous values exactly, GPU Q_Jacobi matches the CPU reference exactly. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
ilu_stream and htd_stream never actually need to exist at the same time: quantized_mode (which needs htd_stream for HtDTransfer's async L/U copy) and ILU (which needs ilu_stream for its build/apply CUDA graphs) are alternative preconditioner choices decided once in Initialize() - a given CSysMatrix instance is never both. Merge them into a single aux_stream member instead of keeping two separate streams (and two create/destroy pairs) that are never both live. Verified: GPU Q_Jacobi matches CPU Q_Jacobi exactly and GPU ILU converges normally, both reproducible over 3 repeated runs. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
pcarruscag
commented
Aug 16, 2026
Co-authored-by: Pedro Gomes <38071223+pcarruscag@users.noreply.github.com>
EncodeQuantRow cast F's return straight to float, on the assumption that quantization is compile-time disabled whenever ScalarType could be AD-active. That's only true for reverse-mode AD - quantized_mode (CSysMatrix.cpp's quantized_offdiag_needed) is gated on `#ifndef CODI_REVERSE_TYPE`, so it's still compiled in for forward-mode AD (CODI_FORWARD_TYPE), where ScalarType is codi::ActiveType<...> and a direct static_cast<float> doesn't compile. Branch on __CUDA_ARCH__ instead: on device (only ever instantiated for plain ScalarType, never AD-active) cast directly; on host go through SU2_TYPE::PassiveValue first (host-only, not SU2_CUDA_HOST_DEVICE) to safely extract the passive value before the float cast. Verified: SU2_CFD_DIRECTDIFF (forward-mode AD) now builds, and the GPU regression set (Jacobi/Q_Jacobi/Q_Identity/ILU) is unaffected - bit-identical to before this fix. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Proposed Changes
Also compatible with GPU
PR Checklist
pre-commit run --allto format old commits.