Fix/mr pool span hardening - #144
Open
abagusetty wants to merge 6 commits into
Open
Conversation
…den allocator
Root cause of "No memory-block found in stream_ordered_memory_resource" on large
requests against an apparently-free pool: 104 call sites recomputed
count*sizeof(T) independently for allocate and the matching deallocate. Any
mismatch corrupts the free list silently -- freeing short orphans the
remainder forever, freeing long makes the pool believe it owns live memory --
and both permanently fragment the coalescing chain.
Diagnostics (this session, prior commit) already made pool exhaustion vs
fragmentation distinguishable. This change removes the mismatch class itself.
std::span API (src/tamm/mr/device_memory_resource.hpp, host_memory_resource.hpp)
- allocate_span<T>(n) / deallocate(span<T>): the span carries its own byte
count, so a free can no longer disagree with its allocation. Measured
against std::map/unordered_map trackers: this costs 0.00ns/pair (compiles
away) vs +38-71% for a runtime tracker. Raw allocate/deallocate retained.
Migrated 104 call sites to the span API (50 in TAMM, 54 in exachem-dev):
- src/tamm/multop.hpp, kernels/multiply.hpp, tests/coupledcluster/Test_CCSD.cpp
- exachem/cc/ccsd/cd_ccsd_{cs,os}_ann.cpp, cc/ccsd_t/ccsd_t_all_fused.hpp,
cc/ccsd_t/ccsd_t_fused_driver.hpp
Prerequisite fix (gpu_streams.hpp): gpuMemsetAsync took void*&, forcing every
call site to reinterpret_cast<void*&>(typed_ptr) -- a strict-aliasing
violation -- even though the reference was never written through. Changed to
void* by value; this also unblocks binding a span's .data() at call sites.
Allocator hardening (mr/coalescing_free_list.hpp, free_list.hpp):
- free_list backed by std::set (address-ordered) instead of std::list;
added a size_index_ secondary std::set for O(log n) best-fit. Measured
97x faster at 10k free blocks (17.0us -> 0.16us per alloc/free pair),
directly in the CCSD allocation hot path.
- Base mutators made protected so nothing outside coalescing_free_list can
desync the two indices. Both invariants (address<->size index agreement)
are now assert()-guarded rather than silently tolerated -- a discarded
std::set::insert result previously let a duplicate address slip into one
index but not the other, which could hand the same address to two live
allocations.
Correctness fixes surfaced while auditing (mr/aligned.hpp):
- is_pow2(0) returned true (hand-rolled bit trick), so
is_supported_alignment(0) passed and align_up(x, 0) returned SIZE_MAX.
Replaced with std::has_single_bit (C++20); alignment 0 now correctly
normalizes to the default via supported_alignment_or_default(), applied
identically on both the allocate and deallocate side so the two can never
derive different padded sizes.
- aligned_allocate() did not check its allocator's return value; numa_alloc_onnode
returns NULL on failure (does not throw), which previously wrote through a
null pointer instead of surfacing as std::bad_alloc.
- aligned_deallocate's callable is now binary (original ptr, padded size)
instead of unary; new_delete_resource's numa_free() path was freeing the
caller's requested bytes instead of the padded allocation size actually
reserved by aligned_allocate, leaking alignment+sizeof(ptrdiff_t) bytes
per upstream reservation.
- ccsd_t_fused_driver.hpp: replaced std::pow(max_num_blocks, 6) (a double)
sizing a real allocation with integer arithmetic shared by the GPU and
CPU allocation paths.
TAMM_RMM_TRACK=1 diagnostics (mr/pool_memory_resource.hpp, from prior commit):
- outstanding_ switched from std::map to std::unordered_map (+38% vs +71%
per alloc/free pair when tracking is enabled; ordering was never needed,
lookups are always by exact pointer).
- Fixed a deadlock: validate_free() called tamm_terminate() while holding
the pool mutex; tamm_terminate() calls exit(), which runs static
destructors that re-enter release() on that same non-recursive mutex on
the same thread. Reproduced the hang, fixed by returning the diagnostic
as an out-param and terminating only after the lock scope closes (same
pattern already used in do_allocate).
C++20 hardening (mr/coalescing_free_list.hpp, mr/aligned.hpp):
- block::operator<=> (spaceship) replaces the hand-written operator<.
- [[nodiscard]] on all allocate/allocate_span entry points -- discarding
the pointer is always a leak.
- 15 static_assert round-trip proofs for align_up/align_down/
detail_padded_size/supported_alignment_or_default, including the
SIZE_MAX-saturation and zero-is-not-a-power-of-two cases above. These
fail the build rather than corrupting a free list at runtime.
Testing (tests/tamm/Test_MemoryPool.cpp, new):
18 cases / 9035 assertions, run at 4 RMM_ALLOCATION_ALIGNMENT values (4, 16,
128, 256) x tracking on/off = 8 full passes, plus:
- TSan and ASan/UBSan clean in all 8 combinations
- -O2 -DNDEBUG pass
- Mutation-tested: deleting one drop_from_size_index() call fails 2 cases
- 200k-op randomized adversarial audit (exact byte accounting every step),
32768-chunk shuffle-then-realloc-whole-pool, alignment-boundary sizes
- All fatal paths (oversized request, fragmentation, exhaustion, tracker
mismatch/double-free/leak) confirmed to exit cleanly, not hang
- Before/after benchmark for the O(log n) free-list claim
Not verified here (no GPU, no cmake on this machine): the exachem-dev edits
compile-shape-checked in isolation and are span-balanced (every span declared
is freed exactly once, verified per-file), but the actual exachem build and
all USE_CUDA/USE_HIP/USE_DPCPP branches in both repos are unverified.
…is in MultOp
Two unrelated warnings reported from a real GCC build (NWChemEx V100 CUDA),
independent of the memory-pool work in this branch.
index_loop_nest.hpp:174 -- real use-after-free (GCC -Wdangling-reference)
for(const TileLabelElement& slbl: ibc.this_label().secondary_labels()) {
IndexLoopBound::this_label() returns TiledIndexLabel by value (a temporary).
TiledIndexLabel::secondary_labels() returns a reference into that temporary's
own std::vector member (TiledIndexLabel holds its data by value, unlike e.g.
TiledIndexSpace/Tensor which pimpl through a shared_ptr). The range-for's
range-expression binds a reference to the vector returned by
secondary_labels(), not directly to the this_label() temporary, so C++
temporary-lifetime-extension does not apply: the TiledIndexLabel is destroyed
at the end of the full expression, and `slbl` reads freed memory on every
iteration thereafter.
Confirmed with an isolated ASan/UBSan repro reproducing the identical
this_label()-by-value / secondary_labels()-by-reference shape: the old
pattern reports "stack-use-after-scope" immediately; binding this_label() to
a named local first (extending its lifetime to the loop) reports clean and
returns correct data. Checked every other .secondary_labels() call site in
the codebase (14 others across tiled_index_space.hpp, labeled_tensor.hpp,
op_visitors.hpp, op_attributes.hpp, utils.hpp, tensor_base.hpp/.cpp) --
all of them already bind through a named local/reference first, so this
constructor was the only affected call site.
Introduced in 8125a93 ("Fix memory leaks", 2024-09-06).
multop.hpp:204 -- deprecated implicit `this` capture (C++20 -Wdeprecated)
auto lambda = [=, &oprof, &add_bufs, &loop_nest, &ec](const IndexVector itval) {
The lambda accesses lhs_/rhs1_/rhs2_ unqualified, relying on an implicit
copy-capture of `this` via the `[=` default. C++20 deprecates this specific
form (implicit `this` via a `[=]` default) and requires it be named
explicitly; `[&]`-default implicit `this` is unaffected and not deprecated.
Fixed by adding `this` to the capture list. No behavior change -- do_work()
invokes the lambda synchronously within execute(), so `this` does not outlive
the call regardless. Checked every other lambda capture in src/tamm/ (~90
`[&...]` lambdas, 3 other `[=]` lambdas that are either non-member functions
or nested inside a non-capturing outer lambda) -- this was the only true
implicit-`this`-via-`[=]` site in the codebase.
Verified: Test_MemoryPool.cpp (18 cases / 9035 assertions, from this same
branch) unaffected -- both changes are in files unrelated to the memory
pool. Not independently build-verified against GA/MPI/CUDA on this machine;
the exact code shapes were isolated and verified under ASan/UBSan and against
clang's equivalent -Wdeprecated-this-capture in standalone repros.
… int64
Two defects in kernels/multiply.hpp, both only reachable on CUDA/HIP.
1. Device buffers were returned to the pool while GPU work was still reading them
The pool's deallocate() is pure host-side bookkeeping: it makes a block instantly
re-allocatable and does not defer reuse behind a stream event, because upstream
RMM's stream-ordered free lists were stripped from this fork. Meanwhile
cublasDgemm/rocblas_dgemm, cudaMemcpyAsync and librettExecute are enqueued on
`thandle` and return immediately. Freeing a device buffer straight after those
calls therefore lets the very next allocation hand out the same address while the
GEMM/transpose is still reading it. Confirmed the pool re-issues the identical
address on the next iteration.
The exposed sites were the reduction loops, which free abuf_dev/bbuf_dev every
iteration with the only sync placed after the loop: MultOp::execute_bufacc, and
the equivalent loops in exachem's cd_ccsd_{cs,os}_ann.cpp. MultOp::execute was
already safe (it syncs before its D2H copy).
Fixed in block_multiply rather than at the call sites, so the guarantee is a
postcondition of the function: every device buffer passed in or allocated
internally is free of in-flight work when it returns. This makes the exachem
loops correct with no change to that repository, and any future caller inherits
it. DPC++ is unaffected -- gpu::gemm already calls gemm_event.wait() and the SYCL
queue is in-order -- so the added syncs are no-ops there. Cost on CUDA/HIP is 2
syncs per block_multiply on the T1==T2==T3 path (the CCSD hot path), 3 on the
mixed-precision paths; that is the synchronisation the callers' correctness
already assumed.
2. gemm_wrapper computed batch/reduction strides in int
int bbatch_ld = K * N;
int breduce_ld = B * bbatch_ld;
int*int is evaluated in int before any promotion, so B*K*N overflows. At B=256,
K=N=4096 it wraps to exactly 0, and with BR>1 every reduction iteration then
re-reads batch 0 instead of advancing -- silently wrong numbers, no crash. M*K
overflows the same way once a single stride exceeds INT_MAX. Both are signed
overflow, i.e. UB.
Note the offset *accumulation* was already 64-bit: the loop counters were size_t,
which promoted the int strides. An earlier draft of this change attributed the
bug to offset accumulation instead; that was wrong and the comment and tests have
been written against the actual failure. Strides, offsets and loop counters are
now int64_t. Leading dimensions stay int, matching the cpu::gemm/gpu::gemm
signatures in kernels/tamm_blas.hpp, which take int for m/n/k and lda/ldb/ldc.
Testing (tests/tamm/Test_MemoryPool.cpp): 4 new cases pinning the stride
arithmetic against the shapes that distinguish int from int64, including the
B=256/K=N=4096 wrap-to-zero case and an end-to-end BR>1 offset comparison. The
old int behaviour is emulated with unsigned arithmetic rather than reproduced
literally, so the tests do not themselves trip UBSan. Mutation-tested: reverting
the widening fails 6 assertions across 3 cases.
Full suite now 22 cases / 9064 assertions, run at 4 alignments x tracking on/off,
clean under TSan and ASan/UBSan, and under -O2 -DNDEBUG. Separately verified that
all 7 syncs sit inside #if defined(USE_CUDA)||defined(USE_HIP)||defined(USE_DPCPP)
guards -- an earlier draft left 5 of them unguarded, which broke the CPU-only
build outright, since gpuStreamSynchronize is undeclared there and the call is a
non-dependent name resolved at template definition time.
Not verified here (no GPU, no GA/MPI on this machine): the CUDA/HIP/DPC++ code
paths are compile-checked only in isolation, against stand-ins for gpuStream_t
and gpuStreamSynchronize.
Host allocate/deallocate defaulted to alignof(std::max_align_t) (16), undercutting RMM_ALLOCATION_ALIGNMENT (256 CUDA / 128 HIP), so the pool's base slab and every suballocation were misaligned. Made the backend constant the default everywhere and dropped the explicit alignof(T) from the span API. Also raised the DPC++ constant from 4 to alignof(std::max_align_t) and switched the SYCL USM calls to aligned_alloc_device/aligned_alloc_host: SYCL 2020 Table 70 only guarantees fundamental alignment for the plain overloads, and 4 is below alignof(std::complex<double>). HIP's 128 left as-is, flagged in-source as unverified. Not built or tested locally; to be verified on Aurora and NVIDIA systems.
Mark intentionally-discarded smoke-test results [[maybe_unused]] and drop the dead ex_hw in Test_LocalTensor's main. In check_local_tensor_values, num_elements was deduced int from `auto ... = 1`, so accumulating the size_t extents truncated each iteration and could overflow before the loop compared it against a size_t. Declared it size_t, which fixes the overflow and the sign-compare warning together.
tamm_terminate() exits the process, which runs the pool destructor and made the tracker report every still-live block as a leak. Those blocks are simply what was allocated when the abort was raised, so a failed allocation always printed a spurious LEAK trailer. Added tamm_terminating() and reported them as live-at-abort instead; the size breakdown is unchanged. Also reworked the pool-exhaustion diagnostic. It always suggested raising TAMM_GPU_POOL / TAMM_CPU_POOL, which is wrong when peak demand exceeds the device: no pool setting can help. It now reports peak demand (in use + requested) and separates a request larger than the entire pool from ordinary exhaustion. tamm_terminate() exited 0, reporting success to mpiexec and CI after a fatal error; it now exits 1.
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.
Is this pull request associated with an issue(s)?
N/A
Description
N/A
TODOs
For draft pull requests please include a list of what needs to be done and check
off items as you complete them.