Phambinh/wip cached allocs pingpong

[phambinhfin]

Mode	Command Buffer Types	Cached Allocs	Median	vs Baseline
A. Baseline	none	n/a	13.601 ms	--
B. FUSION only	FUSION	ON	14.369 ms	+5.6% (overhead > savings)
C. No collectives	FUSION + CUBLAS + CUBLASLT + CUSTOM_CALL	ON	13.413 ms	-1.4%
D. All types	FUSION + CUBLAS + CUBLASLT + CUSTOM_CALL + COLLECTIVES	ON	10.236 ms	-24.7%
E. All types (broken)	same as D	OFF	287.085 ms	+2010% (re-tracing every step)

[claude]

Memory leak: persistent allocations are never freed.

executor->Allocate() creates device memory that is stored in alloc_address_cache, but there is no cleanup path. The ExecutorCommandBuffer has no destructor that frees these allocations, and the GpuExecutable destructor only frees entries in cached_allocs_ (the executable-level cache), not the thunk-level alloc_address_cache.

This means every persistent buffer allocated here leaks GPU memory until process exit, and the leak grows with the number of stabilized allocations across all executors.

Consider adding a destructor to ExecutorCommandBuffer that iterates alloc_address_cache and frees each allocation via StreamExecutor::Deallocate, or unifying this cache with the executable-level cached_allocs_ so a single cleanup path covers both.

[claude]

Bug: collective_init_done set before initialization completes.

The flag is set to true at line 227 before commands_.Record(...) runs (line 248). If Record fails via TF_RETURN_IF_ERROR, the flag stays true, preventing any future retry of collective initialization. This means a transient failure permanently disables collective init.

Move cmd_buffer->collective_init_done = true; to after the Record call succeeds (e.g., after line 249).

[claude]

Data race: CachedAllocator holds a raw pointer to cached_ptrs used outside the lock.

The CachedAllocator stores a raw pointer to cache.cached_ptrs (a flat_hash_set), which is protected by allocs_cache_mutex_. However, CachedAllocator::Deallocate reads cached_ptrs_ (via contains()) outside the lock — e.g., during TearDown or ExecutionOutput destruction.

Meanwhile, BufferForAllocation (line 1076) inserts into cached_ptrs under the lock. A concurrent Deallocate call reading the same flat_hash_set without holding the lock is undefined behavior.

Options:

1. Have CachedAllocator take its own lock (or share allocs_cache_mutex_) when accessing cached_ptrs_.
2. Use a thread-safe set for cached_ptrs.

[claude]

Removing empty-graph safety checks may cause ROCm crashes.

This PR removes both:

1. The Trace() check that returned an error when a traced HIP graph had zero root nodes
2. The PrepareFinalization() empty-node insertion that prevented crashes from empty graphs

The CUDA path still retains both its empty-graph check in Trace() and its PrepareFinalization() empty-body handling. The PR description provides no justification for why this is safe on ROCm. If any custom call traces an empty graph (no HIP operations launched), this will crash at instantiation time.

Could you add a comment explaining why this is safe, or add a test that verifies empty-graph behavior on ROCm?

[claude]

Destructor silently swallows deallocation errors.

(void)real_allocator->DeallocateUnderlying(device_ordinal, addr) discards the absl::Status. If deallocation fails (e.g., device context already torn down, or wrong device ordinal), the error is silently lost. At minimum, log a warning:

Suggested change

	}
	// Free cached allocations through the real (underlying) allocator.
	if (enable_cached_allocs_) {
	absl::MutexLock lock(&allocs_cache_mutex_);
	for (auto& [device_ordinal, cache] : cached_allocs_) {
	auto it = cached_allocators_.find(device_ordinal);
	if (it == cached_allocators_.end() || !it->second) continue;
	auto* real_allocator =
	static_cast<CachedAllocator*>(it->second.get());
	for (auto& [idx, addr] : cache.entries) {
	if (!addr.is_null()) {
	(void)real_allocator->DeallocateUnderlying(device_ordinal, addr);
	}
	for (auto& [idx, addr] : cache.entries) {
	if (!addr.is_null()) {
	auto status = real_allocator->DeallocateUnderlying(device_ordinal, addr);
	if (!status.ok()) {
	LOG(WARNING) << "CachedAllocs: failed to deallocate buffer "
	<< idx << ": " << status;
	}
	}
	}

Also, the static_cast<CachedAllocator*>(it->second.get()) on line 414 is unnecessary — cached_allocators_ stores std::unique_ptr<CachedAllocator>, so .get() already returns CachedAllocator*.

[claude]

CachedAllocator is an implementation detail exposed in a public header.

CachedAllocator and AllocationsCache are only used internally by GpuExecutable, but they are defined in the public header gpu_executable.h. This pollutes the public API and exposes implementation details (raw pointer members, internal caching strategy).

Consider moving these to an anonymous namespace in gpu_executable.cc or to a separate *_internal.h header.

[claude]

Debug counters computed unconditionally, even when VLOG is disabled.

The counters n_params, n_const, n_thread, n_alloc are incremented for every allocation on every call to GenerateBufferAllocations, but the output is gated behind VLOG(1). Consider gating the counting behind VLOG_IS_ON(1) to avoid the overhead:

if (VLOG_IS_ON(1)) {
  int64_t n_params = 0, n_const = 0, n_thread = 0, n_alloc = 0;
  // ... counting loop ...
  VLOG(1) << "GenBufAllocs ...";
}

[claude]

Nit: std::getenv is called redundantly at two locations.

The env var XLA_COMMAND_BUFFERS_USE_CACHED_ALLOCS is read here (per CommandBufferThunk construction) and also in GpuExecutable's constructor (gpu_executable.cc:393). This creates maintenance risk — if one location is changed, the other could be forgotten. std::getenv is also not thread-safe with respect to setenv/putenv.

Consider reading it once and passing the value through the construction chain, or using a tsl::DebugOptions flag instead of an env var.

[claude]

TOCTOU race: double-locking of allocs_cache_mutex_ allows duplicate allocations.

The mutex is acquired for the cache lookup (line 1049), released, then re-acquired for the cache insertion (line 1072). Between the two lock acquisitions, another thread could insert the same arg_idx, causing:

1. A redundant allocation via memory_allocator->Allocate
2. The first allocation being silently overwritten in cache.entries, leaking device memory

Consider holding the lock across the entire lookup-allocate-insert sequence, or using a per-entry lock/flag to prevent concurrent allocation for the same index.

[claude]

Review Summary

This PR implements a cached allocations optimization for GPU command buffer thunks, gated behind XLA_COMMAND_BUFFERS_USE_CACHED_ALLOCS. When enabled, donated parameter addresses are stabilized via D2D memcpy to/from persistent buffers, avoiding costly command buffer re-tracing. The PR also removes empty-graph safety checks from the ROCm command buffer path.

Key issues found (see inline comments for details):

• Memory leak — Thunk-level alloc_address_cache persistent allocations are never freed (no destructor cleanup path)
• Data race — CachedAllocator reads cached_ptrs without holding allocs_cache_mutex_, while BufferForAllocation concurrently inserts into it
• TOCTOU race — Double-locking in BufferForAllocation allows duplicate allocations and memory leaks between lock acquisitions
• Bug — collective_init_done flag set before Record() completes; a transient failure permanently disables collective init
• Safety concern — ROCm empty-graph checks removed without justification; CUDA path retains equivalent checks
• Minor — Destructor swallows deallocation errors silently; debug counters computed unconditionally; redundant std::getenv calls; CachedAllocator exposed in public header
• No tests for the caching mechanism

The concurrency issues (data race on cached_ptrs, TOCTOU in BufferForAllocation) and the memory leak in alloc_address_cache should be addressed before merging.