Fix extract invertpointer ordering (#1437)

EnzymeAD · Sep 18, 2023 · e6e5a03 · e6e5a03
1 parent 86fc287
commit e6e5a03
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 6 deletions.
diff --git a/enzyme/Enzyme/GradientUtils.cpp b/enzyme/Enzyme/GradientUtils.cpp
@@ -5490,12 +5490,13 @@ Value *GradientUtils::invertPointerM(Value *const oval, IRBuilder<> &BuilderM,
     }
     goto end;
   } else if (auto arg = dyn_cast<ExtractValueInst>(oval)) {
-    IRBuilder<> bb(getNewFromOriginal(arg));
+    auto newi = getNewFromOriginal(arg);
+    IRBuilder<> bb(newi->getNextNode());
     auto ip = invertPointerM(arg->getOperand(0), bb, nullShadow);
 
-    auto rule = [&bb, &arg, this](Value *ip) -> llvm::Value * {
+    auto rule = [&bb, &arg, &newi, this](Value *ip) -> llvm::Value * {
       if (ip == getNewFromOriginal(arg->getOperand(0)))
-        return getNewFromOriginal(arg);
+        return newi;
       return bb.CreateExtractValue(ip, arg->getIndices(),
                                    arg->getName() + "'ipev");
     };

diff --git a/enzyme/test/Enzyme/ForwardMode/augmentedreturn.ll b/enzyme/test/Enzyme/ForwardMode/augmentedreturn.ll
@@ -34,8 +34,8 @@ entry:
 ; CHECK-NEXT:   %0 = call { { i8*, double }, { i8*, double } } @fwddiffeaugsquare(double %x, double %"x'")
 ; CHECK-NEXT:   %1 = extractvalue { { i8*, double }, { i8*, double } } %0, 0
 ; CHECK-NEXT:   %2 = extractvalue { { i8*, double }, { i8*, double } } %0, 1
-; CHECK-NEXT:   %[[i3:.+]] = extractvalue { i8*, double } %2, 1
 ; CHECK-NEXT:   %o = extractvalue { i8*, double } %1, 1
+; CHECK-NEXT:   %[[i3:.+]] = extractvalue { i8*, double } %2, 1
 ; CHECK-NEXT:   %[[i4:.+]] = fmul fast double %[[i3]], %o
 ; CHECK-NEXT:   %[[i5:.+]] = fmul fast double %[[i3]], %o
 ; CHECK-NEXT:   %[[i6:.+]] = fadd fast double %[[i4]], %[[i5]]

diff --git a/enzyme/test/Enzyme/ForwardMode/freeuse.ll b/enzyme/test/Enzyme/ForwardMode/freeuse.ll
@@ -31,8 +31,8 @@ declare void @free(i8*)
 
 ; CHECK: define internal void @fwddiffejac_rev(double* nocapture readonly %r, { i8* } %tapeArg, { i8* } %"tapeArg'", i1 %cmp)
 ; CHECK-NEXT: entry:
-; CHECK-NEXT:   %"arg0'ipev" = extractvalue { i8* } %"tapeArg'", 0
 ; CHECK-NEXT:   %arg0 = extractvalue { i8* } %tapeArg, 0
+; CHECK-NEXT:   %"arg0'ipev" = extractvalue { i8* } %"tapeArg'", 0
 ; CHECK-NEXT:   store i8 0, i8* %"arg0'ipev", align 8
 ; CHECK-NEXT:   store i8 0, i8* %arg0, align 8
 ; CHECK-NEXT:   br i1 %cmp, label %invertbaz, label %invertfoo
@@ -52,7 +52,7 @@ declare void @free(i8*)
 ; CHECK-NEXT:   ret void
 
 ; CHECK: invertfoo:                                        ; preds = %entry
-; CHECK-NEXT:   %"arg1'ipev" = extractvalue { i8* } %"tapeArg'", 0
 ; CHECK-NEXT:   %arg1 = extractvalue { i8* } %tapeArg, 0
+; CHECK-NEXT:   %"arg1'ipev" = extractvalue { i8* } %"tapeArg'", 0
 ; CHECK-NEXT:   br label %invertbaz
 ; CHECK-NEXT: }