Fix shadow augmented rematerialization (rust-lang#557)

* Shadow promotable remat

* Preserve forward shadow if non-float

Co-authored-by: Ubuntu <[email protected]>
Co-authored-by: Valentin Churavy <[email protected]>

Author: 3 people

enzyme/Enzyme/AdjointGenerator.h

@@ -10443,10 +10443,12 @@ class AdjointGenerator
 #endif
         tapeIdx = 0;
 
-        if (subretType == DIFFE_TYPE::DUP_ARG ||
-            subretType == DIFFE_TYPE::DUP_NONEED) {
+        if (!orig->getType()->isVoidTy()) {
           returnIdx = 1;
-          differetIdx = 2;
+          if (subretType == DIFFE_TYPE::DUP_ARG ||
+              subretType == DIFFE_TYPE::DUP_NONEED) {
+            differetIdx = 2;
+          }
         }
 
       } else {

enzyme/Enzyme/GradientUtils.h

@@ -611,6 +611,28 @@ class GradientUtils : public CacheUtility {
             todo.push_back(std::make_pair(I, (Value *)cur));
         }
       } else if (auto load = dyn_cast<LoadInst>(cur)) {
+
+        // If loaded value is an int or pointer, may need
+        // to preserve initialization within the primal.
+        auto TT = TR.query(load)[{-1}];
+        if (!TT.isFloat()) {
+          // ok to duplicate in forward /
+          // reverse if it is a stack or GC allocation.
+          // Said memory will still be shadow initialized.
+          StringRef funcName = "";
+          if (auto CI = dyn_cast<CallInst>(V))
+            if (Function *originCall = getFunctionFromCall(CI))
+              funcName = originCall->getName();
+          if (isa<AllocaInst>(V) || hasMetadata(V, "enzyme_fromstack") ||
+              funcName == "jl_alloc_array_1d" ||
+              funcName == "jl_alloc_array_2d" ||
+              funcName == "jl_alloc_array_3d" || funcName == "jl_array_copy" ||
+              funcName == "julia.gc_alloc_obj") {
+            primalInitializationOfShadow = true;
+          } else {
+            shadowpromotable = false;
+          }
+        }
         loads.insert(load);
       } else if (auto store = dyn_cast<StoreInst>(cur)) {
         // TODO only add store to shadow iff non float type
@@ -770,10 +792,7 @@ class GradientUtils : public CacheUtility {
       outer = getAncestor(outer, OrigLI.getLoopFor(S->getParent()));
     }
 
-    if (!shadowpromotable)
-      return;
-
-    if (!isConstantValue(V)) {
+    if (shadowpromotable && !isConstantValue(V)) {
       backwardsOnlyShadows[V] = ShadowRematerializer(
           stores, frees, primalInitializationOfShadow, outer);
     }

enzyme/test/Enzyme/ReverseMode/callshadowar.ll

@@ -0,0 +1,76 @@
+; RUN: %opt < %s %loadEnzyme -enzyme -enzyme-preopt=false -mem2reg -instsimplify -adce -simplifycfg -S | FileCheck %s
+
+; ModuleID = 'ode-unopt.ll'
+source_filename = "ode.cpp"
+target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+define double @inner(double* %data) {
+entry:
+  %d = load double, double* %data, align 8
+  %r = fadd double %d, %d
+  ret double %r
+}
+
+define double @inner2(double* %data) {
+entry:
+  %d = load double, double* %data, align 8
+  %r = fsub double %d, %d
+  ret double %r
+}
+
+define double @sub(double* %in, i64 %v) {
+entry:
+  %a = alloca double (double*)*, i64 2, align 8
+  store double (double*)* @inner, double (double*)** %a, align 8
+  %a1 = getelementptr double (double*)*, double (double*)** %a, i64 1
+  store double (double*)* @inner2, double (double*)** %a1, align 8
+  %fa = getelementptr double (double*)*, double (double*)** %a, i64 %v
+  %f = load double (double*)*, double (double*)** %fa, align 8
+  %r = call double %f(double* %in)
+  ret double %r
+}
+
+define void @outer(double* %in, i64 %v) {
+entry:
+  %r = call double @sub(double* %in, i64 %v)
+  store double %r, double* %in, align 8
+  ret void
+}
+
+define void @caller(double* %in, double* %d_in) {
+entry:
+  call void (...) @__enzyme_autodiff(void (double*, i64)* nonnull @outer, double* %in, double* %d_in, i64 0) 
+  ret void
+}
+
+declare void @__enzyme_autodiff(...)
+
+; CHECK: define internal { i8*, double } @augmented_sub(double* %in, double* %"in'", i64 %v)
+; CHECK-NEXT: entry:
+; CHECK-NEXT:   %0 = alloca { i8*, double }
+; CHECK-NEXT:   %1 = getelementptr inbounds { i8*, double }, { i8*, double }* %0, i32 0, i32 0
+; CHECK-NEXT:   %2 = alloca i8, i64 16, align 8
+; CHECK-NEXT:   %"malloccall'mi" = alloca i8, i64 16, align 8
+; CHECK-NEXT:   call void @llvm.memset.p0i8.i64(i8* nonnull dereferenceable(16) dereferenceable_or_null(16) %"malloccall'mi", i8 0, i64 16, i1 false)
+; CHECK-NEXT:   %"a'ipc" = bitcast i8* %"malloccall'mi" to double (double*)*
+; CHECK-NEXT:   %a = bitcast i8* %2 to double (double*)**
+; CHECK-NEXT:   store double (double*)* bitcast ({ { i8*, double } (double*, double*)*, void (double*, double*, double, i8*)* }* @"_enzyme_reverse_inner'" to double (double*)*), double (double*)** %"a'ipc", align 8
+; CHECK-NEXT:   store double (double*)* @inner, double (double*)** %a, align 8
+; CHECK-NEXT:   %"a1'ipg" = getelementptr double (double*)*, double (double*)** %"a'ipc", i64 1
+; CHECK-NEXT:   %a1 = getelementptr double (double*)*, double (double*)** %a, i64 1
+; CHECK-NEXT:   store double (double*)* bitcast ({ { i8*, double } (double*, double*)*, void (double*, double*, double, i8*)* }* @"_enzyme_reverse_inner2'" to double (double*)*), double (double*)** %"a1'ipg", align 8
+; CHECK-NEXT:   store double (double*)* @inner2, double (double*)** %a1, align 8
+; CHECK-NEXT:   %"fa'ipg" = getelementptr double (double*)*, double (double*)** %"a'ipc", i64 %v
+; CHECK-NEXT:   %"f'ipl" = load double (double*)*, double (double*)** %"fa'ipg", align 8
+; CHECK-NEXT:   %3 = bitcast double (double*)* %"f'ipl" to { i8*, double } (double*, double*)**
+; CHECK-NEXT:   %4 = load { i8*, double } (double*, double*)*, { i8*, double } (double*, double*)** %3
+; CHECK-NEXT:   %r_augmented = call { i8*, double } %4(double* %in, double* %"in'")
+; CHECK-NEXT:   %subcache = extractvalue { i8*, double } %r_augmented, 0
+; CHECK-NEXT:   store i8* %subcache, i8** %1
+; CHECK-NEXT:   %r = extractvalue { i8*, double } %r_augmented, 1
+; CHECK-NEXT:   %5 = getelementptr inbounds { i8*, double }, { i8*, double }* %0, i32 0, i32 1
+; CHECK-NEXT:   store double %r, double* %5
+; CHECK-NEXT:   %6 = load { i8*, double }, { i8*, double }* %0
+; CHECK-NEXT:   ret { i8*, double } %6
+; CHECK-NEXT: }