Fix crash on no return use (#1575)

EnzymeAD · Dec 9, 2023 · 7094012 · 7094012
1 parent 5701321
commit 7094012
Show file tree

Hide file tree

Showing 3 changed files with 59 additions and 13 deletions.
diff --git a/enzyme/Enzyme/AdjointGenerator.h b/enzyme/Enzyme/AdjointGenerator.h
@@ -6173,7 +6173,15 @@ class AdjointGenerator
       return;
 
     if (gutils->isConstantInstruction(&call) &&
-        gutils->isConstantValue(&call)) {
+        (gutils->isConstantValue(&call) || !shadowReturnUsed)) {
+      if (!gutils->isConstantValue(&call)) {
+        auto found = gutils->invertedPointers.find(&call);
+        if (found != gutils->invertedPointers.end()) {
+          PHINode *placeholder = cast<PHINode>(&*found->second);
+          gutils->invertedPointers.erase(found);
+          gutils->erase(placeholder);
+        }
+      }
       bool noFree = Mode == DerivativeMode::ForwardMode;
       noFree |= call.hasFnAttr(Attribute::NoFree);
       if (!noFree && called) {

diff --git a/enzyme/test/Enzyme/ReverseMode/cvalnoti.ll b/enzyme/test/Enzyme/ReverseMode/cvalnoti.ll
@@ -0,0 +1,49 @@
+; RUN: if [ %llvmver -lt 16 ]; then %opt < %s %loadEnzyme -enzyme -enzyme-preopt=false -mem2reg -instsimplify -simplifycfg -S | FileCheck %s; fi
+; RUN: %opt < %s %newLoadEnzyme -passes="enzyme,function(mem2reg,instsimplify,%simplifycfg)" -enzyme-preopt=false -S | FileCheck %s
+
+declare { i8*, double } @__enzyme_augmentfwd(...) 
+
+define double* @c(double* %x) readnone {
+  ret double* %x
+}
+
+define double* @d(double* %x) readnone {
+  ret double* %x
+}
+
+define double @use(double %x, i1 %z) {
+  ret double %x
+}
+
+; differential use analysis crash requires
+;   %r2 is constant instruction, but not constant value
+;   shadow of r2 is not needed
+define double @square(double %x, double* %v) {
+entry:
+  %r1 = call double* @c(double* %v)
+  %r2 = call double* @d(double* %r1)
+  %z = icmp eq double* %r2, null
+  %q = call double @use(double %x, i1 %z)
+  ret double %q
+}
+
+define void @dsquare(double %x) {
+entry:
+  %0 = call { i8*, double } (...) @__enzyme_augmentfwd(double (double, double*)* @square, double %x, metadata !"enzyme_dup", double* null, double* null)
+  ret void
+}
+
+; CHECK: define internal { i8*, double } @augmented_square(double %x, double* %v, double* %"v'")
+; CHECK-NEXT: entry:
+; CHECK-NEXT:   %0 = alloca { i8*, double }, align 8
+; CHECK-NEXT:   %1 = getelementptr inbounds { i8*, double }, { i8*, double }* %0, i32 0, i32 0
+; CHECK-NEXT:   store i8* null, i8** %1, align 8
+; CHECK-NEXT:   %r1 = call double* @nofree_c(double* %v)
+; CHECK-NEXT:   %r2 = call double* @nofree_d(double* %r1)
+; CHECK-NEXT:   %z = icmp eq double* %r2, null
+; CHECK-NEXT:   %q = call fast double @augmented_use(double %x, i1 %z)
+; CHECK-NEXT:   %2 = getelementptr inbounds { i8*, double }, { i8*, double }* %0, i32 0, i32 1
+; CHECK-NEXT:   store double %q, double* %2, align 8
+; CHECK-NEXT:   %3 = load { i8*, double }, { i8*, double }* %0, align 8
+; CHECK-NEXT:   ret { i8*, double } %3
+; CHECK-NEXT: }
diff --git a/enzyme/test/Enzyme/ReverseMode/mincachechain3.ll b/enzyme/test/Enzyme/ReverseMode/mincachechain3.ll
@@ -58,11 +58,6 @@ entry:
 ; CHECK-NEXT:   ret void
 ; CHECK-NEXT: }
 
-; CHECK: define internal i64* @augmented__ZNKSt5arrayIlLm4EEixEm(i64* %ptr)
-; CHECK-NEXT: entry:
-; CHECK-NEXT:   ret i64* %ptr
-; CHECK-NEXT: }
-
 ; CHECK: define internal void @augmented_usesize(double* %ptr, double* %"ptr'", i64 %off)
 ; CHECK-NEXT: entry:
 ; CHECK-NEXT:   ret void
@@ -78,7 +73,7 @@ entry:
 ; CHECK-NEXT:   %iv = phi i64 [ %iv.next, %for.body ], [ 0, %entry ]
 ; CHECK-NEXT:   %size.0 = phi i64 [ 1, %entry ], [ %a2, %for.body ]
 ; CHECK-NEXT:   %iv.next = add nuw nsw i64 %iv, 1
-; CHECK-NEXT:   %call3 = call i64* @augmented__ZNKSt5arrayIlLm4EEixEm(i64* %v)
+; CHECK-NEXT:   %call3 = call i64* @_ZNKSt5arrayIlLm4EEixEm(i64* %v)
 ; CHECK-NEXT:   %a2 = load i64, i64* %call3, align 8, !tbaa !0
 ; CHECK-NEXT:   %0 = getelementptr inbounds i64, i64* %a2_malloccache, i64 %iv
 ; CHECK-NEXT:   store i64 %a2, i64* %0, align 8, !tbaa !0, !invariant.group 
@@ -110,7 +105,6 @@ entry:
 
 ; CHECK: invertfor.body:                                   ; preds = %invertfor.end, %incinvertfor.body
 ; CHECK-NEXT:   %"iv'ac.0" = phi i64 [ 3, %invertfor.end ], [ %[[dec:.+]], %incinvertfor.body ]
-; CHECK-NEXT:   call void @diffe_ZNKSt5arrayIlLm4EEixEm(i64* %v)
 ; CHECK-NEXT:   %[[rcmp:.+]] = icmp eq i64 %"iv'ac.0", 0
 ; CHECK-NEXT:   br i1 %[[rcmp]], label %invertentry, label %incinvertfor.body
 
@@ -123,11 +117,6 @@ entry:
 ; CHECK-NEXT:   br label %invertfor.body
 ; CHECK-NEXT: }
 
-; CHECK: define internal void @diffe_ZNKSt5arrayIlLm4EEixEm(i64* %ptr)
-; CHECK-NEXT: entry:
-; CHECK-NEXT:   ret void
-; CHECK-NEXT: }
-
 ; CHECK: define internal void @diffeusesize(double* %ptr, double* %"ptr'", i64 %off, double %differeturn)
 ; CHECK-NEXT: entry:
 ; CHECK-NEXT:   %"p2'ipg" = getelementptr inbounds double, double* %"ptr'", i64 %off