NVPTX reverse-mode cache growth does not release superseded buffers
- Dominant language
- LLVM
- Stars
- 1.7k
- Forks
- 188
- Avg merge
- 2d 4h
- Merged PRs (30d)
- 22
Description
When a reverse-mode cache grows on an NVPTX target, the generated code allocates a larger buffer and copies the previous contents without releasing the previous buffer. Reverse execution releases only the final buffer. Repeated execution therefore retains device-heap memory until allocation fails.
## Reproducer
Tested with Enzyme **v0.0.293** and LLVM **20.1.8**.
Save as `mwe.ll`:
```llvm
target triple = "nvptx64-nvidia-cuda"
define double @f(double %x) {
entry:
br label %loop
loop:
%s = phi double [ 1.0, %entry ], [ %next, %loop ]
%next = fmul double %s, %x
%done = fcmp olt double %next, 0.1
br i1 %done, label %exit, label %loop
exit:
ret double %next
}
declare double @__enzyme_autodiff(ptr, ...)
define double @df(double %x) {
%dx = call double (ptr, ...) @__enzyme_autodiff(ptr @f, double %x)
ret double %dx
}
```
Run with the Enzyme library built for the same LLVM version:
```sh
opt -load-pass-plugin=/path/to/libEnzyme-20.so \
-enzyme-preopt=false \
-passes='enzyme,function(mem2reg,instsimplify)' \
-S mwe.ll -o differentiated.ll
```
## Observed and expected behaviour
Relevant instructions from the cache-growth block in `differentiated.ll` (attributes and tail initialization omitted):
```llvm
%11 = call noalias nonnull ptr @malloc(i64 %7)
call void @llvm.memcpy.p0.p0.i64(ptr %11, ptr %s_cache.0, i64 %10, i1 false)
br label %__enzyme_exponentialallocation.exit
```
The old allocation, `%s_cache.0`, becomes unreachable when the cache pointer is replaced. The reverse pass frees only the final allocation.
Expected: release each previous allocation after copying its contents, excluding the initial null pointer. The final allocation should remain owned by reverse execution and be freed there once.
## Source (I think)
In [`getOrInsertExponentialAllocator`](https://github.com/EnzymeAD/Enzyme/blob/v0.0.293/enzyme/Enzyme/Utils.cpp#L515-L562), NVPTX selects the allocation-and-copy fallback instead of `realloc`. That fallback does not call `CreateDealloc` for the previous buffer.
A patch adds conditional `CreateDealloc` after the copy and updates the result PHI's predecessor. The reproducer then emits the missing cleanup. In a separate GPU runtime test, the patch removes the measured cache leak of 270,064 bytes per pullback and successfully completes 1,000 pullbacks with the default 8 MiB device heap. Gradients continue to match an analytic reference.
## Proposed patch
```diff
diff --git a/enzyme/Enzyme/Utils.cpp b/enzyme/Enzyme/Utils.cpp
index 8884fff..9b345af 100644
--- a/enzyme/Enzyme/Utils.cpp
+++ b/enzyme/Enzyme/Utils.cpp
@@ -556,6 +556,15 @@ Function *getOrInsertExponentialAllocator(Module &M, Function *newFunc,
margs[2]->getType()};
auto memsetF = getIntrinsicDeclaration(&M, Intrinsic::memcpy, tys);
B.CreateCall(memsetF, margs);
+ // Unlike realloc, allocation followed by copying does not release ptr.
+ // The first growth has no previous allocation to deallocate.
+ auto freeOld = BasicBlock::Create(M.getContext(), "freeold", F);
+ auto afterFree = BasicBlock::Create(M.getContext(), "afterfree", F);
+ B.CreateCondBr(isFirstSize, afterFree, freeOld);
+ B.SetInsertPoint(freeOld);
+ CreateDealloc(B, ptr);
+ B.CreateBr(afterFree);
+ B.SetInsertPoint(afterFree);
if (SubZero) {
ZeroInit = false;
IRBuilder<> BB(SubZero);
@@ -582,10 +591,11 @@ Function *getOrInsertExponentialAllocator(Module &M, Function *newFunc,
}
gVal = B.CreatePointerCast(gVal, ptr->getType());
+ auto growthEnd = B.GetInsertBlock();
B.CreateBr(ok);
B.SetInsertPoint(ok);
auto phi = B.CreatePHI(ptr->getType(), 2);
- phi->addIncoming(gVal, grow);
+ phi->addIncoming(gVal, growthEnd);
phi->addIncoming(ptr, entry);
B.CreateRet(phi);
return F;
```
_Note: root cause anlaysis and patch heavily assisted by LLMs_
Contributor guide
Research direction
Start in enzyme/Enzyme/Utils.cpp at getOrInsertExponentialAllocator and reproduce the NVPTX case with the supplied mwe.ll and opt command. Compare the generated differentiated.ll cache-growth block and the separate GPU runtime result; done means superseded buffers are released without freeing the initial null or final buffer twice, while gradients and repeated pullbacks still succeed.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp
- Domain
- compilers
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 68/100