From 584a39f7b6eb52b4a68ee6a7ac22ab6d410f2989 Mon Sep 17 00:00:00 2001
From: Valentin Churavy <v.churavy@gmail.com>
Date: Thu, 1 Oct 2026 18:15:37 +0000
Subject: [PATCH] Keep the augmented call when it returns an empty tape
 directly

When a call is differentiated with an augmented forward pass that returns
only its tape (tapeIdx == -1), the tape value is the augmented call itself.
If the tape type is empty ({}), the isEmptyTy() branch erased the tape to
replace it with undef, which deleted the augmented call and the primal side
effects of the callee.

Only erase the tape when it is an extractvalue of the augmented call.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01G9QZewtK4oCErzZphHWd8u
---
 enzyme/Enzyme/AdjointGenerator.h              |  5 +-
 .../ReverseMode/custom-split-empty-tape.ll    | 69 +++++++++++++++++++
 2 files changed, 73 insertions(+), 1 deletion(-)
 create mode 100644 enzyme/test/Enzyme/ReverseMode/custom-split-empty-tape.ll

diff --git a/enzyme/Enzyme/AdjointGenerator.h b/enzyme/Enzyme/AdjointGenerator.h
index cbd80ae4d47..883995ccc7d 100644
--- a/enzyme/Enzyme/AdjointGenerator.h
+++ b/enzyme/Enzyme/AdjointGenerator.h
@@ -5998,7 +5998,10 @@ class AdjointGenerator : public llvm::InstVisitor<AdjointGenerator> {
                                     augmentcall, {(unsigned)tval}, "subcache");
           if (tape->getType()->isEmptyTy()) {
             auto tt = tape->getType();
-            gutils->erase(cast<Instruction>(tape));
+            // If the tape is returned directly, it is the augmented call
+            // itself, which must be kept for its side effects.
+            if (tape != augmentcall)
+              gutils->erase(cast<Instruction>(tape));
             tape = UndefValue::get(tt);
           } else {
             gutils->TapesToPreventRecomputation.insert(cast<Instruction>(tape));
diff --git a/enzyme/test/Enzyme/ReverseMode/custom-split-empty-tape.ll b/enzyme/test/Enzyme/ReverseMode/custom-split-empty-tape.ll
new file mode 100644
index 00000000000..d0062800575
--- /dev/null
+++ b/enzyme/test/Enzyme/ReverseMode/custom-split-empty-tape.ll
@@ -0,0 +1,69 @@
+; RUN: if [ %llvmver -ge 15 ]; then %opt < %s %OPnewLoadEnzyme -passes="enzyme" -enzyme-preopt=false -S | FileCheck %s; fi
+
+; @square has a custom augmented forward pass that returns only an empty tape
+; ({}) and writes to %out. @top reads %out after calling @caller, so Enzyme
+; differentiates @caller in split mode. The augmented forward pass of @caller
+; returns the tape of @square directly, and must keep the call to
+; @augment_square for its store to %out.
+
+define internal {} @augment_square(ptr %x, ptr %dx, ptr %out, ptr %dout) {
+entry:
+  %v = load double, ptr %x
+  %m = fmul double %v, %v
+  store double %m, ptr %out
+  ret {} undef
+}
+
+define internal {} @gradient_square(ptr %x, ptr %dx, ptr %out, ptr %dout, {} %tape) {
+entry:
+  %v = load double, ptr %x
+  %d = load double, ptr %dout
+  store double 0.000000e+00, ptr %dout
+  %t = fmul double %d, %v
+  %t2 = fmul double %t, 2.000000e+00
+  %o = load double, ptr %dx
+  %n = fadd double %o, %t2
+  store double %n, ptr %dx
+  ret {} undef
+}
+
+declare !enzyme_augment !{ptr @augment_square} !enzyme_gradient !{ptr @gradient_square} void @square(ptr %x, ptr %out)
+
+define void @caller(ptr %x, ptr %out) {
+entry:
+  call void @square(ptr %x, ptr %out)
+  ret void
+}
+
+define void @top(ptr %x, ptr %out) {
+entry:
+  call void @caller(ptr %x, ptr %out)
+  %v = load double, ptr %out
+  %m = fmul double %v, %v
+  store double %m, ptr %out
+  ret void
+}
+
+define void @dtop(ptr %x, ptr %dx, ptr %out, ptr %dout) {
+entry:
+  call void (...) @__enzyme_autodiff(ptr @top, ptr %x, ptr %dx, ptr %out, ptr %dout)
+  ret void
+}
+
+declare void @__enzyme_autodiff(...)
+
+; CHECK: define internal {} @augmented_caller(ptr %x, ptr %"x'", ptr %out, ptr %"out'")
+; CHECK-NEXT: entry:
+; CHECK-NEXT:   %0 = alloca {}, align 8
+; CHECK-NEXT:   %_augmented = call {} @augment_square(ptr %x, ptr %"x'", ptr %out, ptr %"out'")
+; CHECK-NEXT:   %1 = load {}, ptr %0, align 1
+; CHECK-NEXT:   ret {} %1
+; CHECK-NEXT: }
+
+; CHECK: define internal void @diffecaller(ptr %x, ptr %"x'", ptr %out, ptr %"out'", {} %tapeArg)
+; CHECK-NEXT: entry:
+; CHECK-NEXT:   br label %invertentry
+; CHECK: invertentry:
+; CHECK-NEXT:   %0 = call {} @gradient_square(ptr %x, ptr %"x'", ptr %out, ptr %"out'", {} undef)
+; CHECK-NEXT:   ret void
+; CHECK-NEXT: }