EnzymeAD / EnzymeAD/Enzyme

NVPTX reverse-mode cache growth does not release superseded buffers

Open
#3,219 0 comments 0 reactions 0 assignees View on GitHub
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.