handle indirect enzyme_const loads (rust-lang#753)

* add testcase for enzyme loads behind phis * handle load trough phis * adjust test for older llvms * adjust test for older llvms * adjust test for even older llvms
EnzymeAD · Jul 29, 2022 · f10b38d · f10b38d
1 parent 8abdb9f
commit f10b38d
Show file tree

Hide file tree

Showing 2 changed files with 72 additions and 0 deletions.
diff --git a/enzyme/Enzyme/Enzyme.cpp b/enzyme/Enzyme/Enzyme.cpp
@@ -498,6 +498,38 @@ castToDiffeFunctionArgType(IRBuilder<> &Builder, llvm::CallInst *CI,
   return Builder.CreateBitCast(value, destType);
 }
 
+static Optional<StringRef> getMetadataName(llvm::Value *res);
+
+// if all phi arms are (recursively) based on the same metaString, use that
+static Optional<StringRef> recursePhiReads(PHINode *val) {
+  Optional<StringRef> finalMetadata;
+  SmallVector<PHINode *, 1> todo = {val};
+  SmallSet<PHINode *, 1> done;
+  while (todo.size()) {
+    auto phiInst = todo.back();
+    todo.pop_back();
+    if (done.count(phiInst))
+      continue;
+    done.insert(phiInst);
+    for (unsigned j = 0; j < phiInst->getNumIncomingValues(); ++j) {
+      auto newVal = phiInst->getIncomingValue(j);
+      if (auto phi = dyn_cast<PHINode>(newVal)) {
+        todo.push_back(phi);
+      } else {
+        Optional<StringRef> metaString = getMetadataName(newVal);
+        if (metaString) {
+          if (!finalMetadata) {
+            finalMetadata = metaString;
+          } else if (finalMetadata != metaString) {
+            return None;
+          }
+        }
+      }
+    }
+  }
+  return finalMetadata;
+}
+
 static Optional<StringRef> getMetadataName(llvm::Value *res) {
   if (auto av = dyn_cast<MetadataAsValue>(res)) {
     return cast<MDString>(av->getMetadata())->getString();
@@ -528,6 +560,9 @@ static Optional<StringRef> getMetadataName(llvm::Value *res) {
   } else if (auto gv = dyn_cast<AllocaInst>(res)) {
     return gv->getName();
   } else {
+    if (isa<PHINode>(res)) {
+      return recursePhiReads(cast<PHINode>(res));
+    }
     return Optional<StringRef>();
   }
 }

diff --git a/enzyme/test/Enzyme/ReverseMode/constarg.ll b/enzyme/test/Enzyme/ReverseMode/constarg.ll
@@ -0,0 +1,37 @@
+; RUN: %opt < %s %loadEnzyme -enzyme -enzyme-preopt=false -mem2reg -sroa -instsimplify -simplifycfg -S | FileCheck %s
+
+; ModuleID = 'ld-temp.o'
+source_filename = "ld-temp.o"
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+@enzyme_const = internal global i32 0, align 4
+
+define internal double @callable(double* %x) {
+  ret double 1.000000e+00
+}
+
+define internal void @_Z19testSensitivitiesADv(double* %primal, double* %grad, double* %a) {
+  br label %bb1
+
+bb1:
+  %a1 = load i32, i32* @enzyme_const, align 4
+  br label %bb3
+
+bb2:
+  %a2 = load i32, i32* @enzyme_const, align 4
+  br label %bb3
+
+bb3:
+  %a3 = phi i32 [ %a1, %bb1 ], [ %a2, %bb2 ]
+  ;%a3 = load i32, i32* @enzyme_const, align 4
+
+  %c1 = call double (...) @__enzyme_autodiff(i8* bitcast (double (double*)* @callable to i8*), i32 %a3, double* %a)
+  ret void
+}
+declare dso_local double @__enzyme_autodiff(...)
+
+; CHECK: define internal void @diffecallable(double* %x, double %differeturn)
+; CHECK-NEXT: invert:
+; CHECK-NEXT:  ret void
+; CHECK-NEXT: }