Kokkos audit: add missing fences after async H2D copies, fix view/index bugs - #263
Merged
Conversation
Fences (each commented with what it fences after): async deep_copys of IS indices/host mirrors were racing ISRestoreIndices or scope exit of the host buffer in VecISCopyLocalk, Gmres_Polyk, Gmres_Poly_Newtonk and PETSc_Helperk (MatCreateSubMatrix paths). Also: SAI_Zk no longer returns early past its MatSeqAIJRestoreKokkosView calls, remove_small_from_sparse_kokkos inserts the lumped diagonal at the correct local column when row/col ownership starts differ, MatDiagDomk fences before destroying the scatter Vecs, the compute_R_from_Z reuse path only applies the identity-skip offset when the identity was included, MatSetAllValues_kokkos fences before exit, plus small cleanups (PetscCallVoid on PetscFree, PetscCallMPIAbort on MPI_Scan, removed deprecated volatile reducer joins, stray semicolon). Co-Authored-By: Claude Fable 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.
Summary
Full review of all Kokkos code (
src/*.kokkos.cxx+include/kokkos_helper.hpp) for bugs, with emphasis on missing fences.Missing fences (use-after-free race class)
Async
Kokkos::deep_copy(exec, device, host)copies were followed by release of the host source buffer —ISRestoreIndices(which frees a temp array for stride/block IS types) or scope exit of a managedcreate_mirror_viewhost mirror — with no fence in between. These pass on CUDA today only because async copies from pageable host memory block the host; they are genuine races on HIP/SYCL or with pinned mirrors. Each new fence carries a single-line comment stating exactly what it fences after:VecISCopyLocalk.kokkos.cxx— fine and coarse IS index copies vsISRestoreIndicesGmres_Polyk.kokkos.cxx/Gmres_Poly_Newtonk.kokkos.cxx— input→submat colmap mapping vs host mirror freed at scope exitPETSc_Helperk.kokkos.cxx—MatCreateSubMatrix_kokkosis_row/is_col copies and theMatCreateSubMatrix_kokkos_viewreuse-path iscol_o copy vsISRestoreIndicesOther fixes
SAI_Zk.kokkos.cxx: the "nothing to do" early return skipped everyMatSeqAIJRestoreKokkosView/RestoreKokkosViewWrite; control now falls through to the shared restore/cleanup epilogue (the existingj_max > 0/count_iter > 0guards already skip the kernels)PETSc_Helperk.kokkos.cxx(remove_small_from_sparse_kokkos): a lumped diagonal inserted into the local block used local row indexias the column; now usesrow_index_global - global_col_start, correct when row/col ownership starts differMatDiagDomk.kokkos.cxx: fence moved before theVecDestroycalls whose device buffers the conversion kernel readsGrid_Transferk.kokkos.cxx(compute_R_from_Z_kokkosreuse path): identity-skip offset now guarded byidentity_intPMISR_Modulek.kokkos.cxx: documented that the ADD-reverse-scatter marker merge relies on the strength matrix being structurally symmetric (S+S^T)MatSetAllValues_kokkos: trailing fence added to match the file conventionPetscCallVoid(PetscFree(...))instead of(void)PetscFree(...),PetscCallMPIAbortaround an uncheckedMPI_Scan, removed the deprecatedvolatile join()reducer overloads inkokkos_helper.hpp, stray semicolonTesting
With
make -j3 build_testswarning-clean,make check+make tests_shortpass:PETSC_OPTIONS="-mat_type aijkokkos -vec_type kokkos -dm_mat_type aijkokkos -dm_vec_type kokkos -on_error_abort"PFLARE_KOKKOS_DEBUG=1(CPU/Kokkos debug-compare)The fence fixes remove latent races that don't reproduce on a CUDA host (pageable-memory async copies block), so the tests validate no regression rather than reproducing the race.
🤖 Generated with Claude Code