Add fence inst handling (#922)

EnzymeAD · Oct 31, 2022 · 60aa782 · 60aa782
1 parent 9f102eb
commit 60aa782
Show file tree

Hide file tree

Showing 3 changed files with 70 additions and 0 deletions.
diff --git a/enzyme/Enzyme/ActivityAnalysis.cpp b/enzyme/Enzyme/ActivityAnalysis.cpp
@@ -639,6 +639,13 @@ bool ActivityAnalyzer::isConstantInstruction(TypeResults const &TR,
     return true;
   }
 
+  if (isa<FenceInst>(I)) {
+    if (EnzymePrintActivity)
+      llvm::errs() << " constant fence instruction " << *I << "\n";
+    InsertConstantInstruction(TR, I);
+    return true;
+  }
+
   if (auto CI = dyn_cast<CallInst>(I)) {
     if (CI->hasFnAttr("enzyme_active")) {
       if (EnzymePrintActivity)
@@ -1620,6 +1627,9 @@ bool ActivityAnalyzer::isConstantValue(TypeResults const &TR, Value *Val) {
       if (notForAnalysis.count(I->getParent()))
         return false;
 
+      if (isa<FenceInst>(I))
+        return false;
+
       // If this is a malloc or free, this doesn't impact the activity
       if (auto CI = dyn_cast<CallInst>(I)) {
         if (CI->hasFnAttr("enzyme_inactive"))

diff --git a/enzyme/Enzyme/AdjointGenerator.h b/enzyme/Enzyme/AdjointGenerator.h
@@ -3646,6 +3646,31 @@ class AdjointGenerator
     eraseIfUnused(MTI);
   }
 
+  void visitFenceInst(llvm::FenceInst &FI) {
+    switch (Mode) {
+    default:
+      break;
+    case DerivativeMode::ReverseModeGradient:
+    case DerivativeMode::ReverseModeCombined: {
+      IRBuilder<> Builder2(FI.getParent());
+      getReverseBuilder(Builder2);
+      auto order = FI.getOrdering();
+      switch (order) {
+      case AtomicOrdering::Acquire:
+        order = AtomicOrdering::Release;
+        break;
+      case AtomicOrdering::Release:
+        order = AtomicOrdering::Acquire;
+        break;
+      default:
+        break;
+      }
+      Builder2.CreateFence(order, FI.getSyncScopeID());
+    }
+    }
+    eraseIfUnused(FI);
+  }
+
   void visitIntrinsicInst(llvm::IntrinsicInst &II) {
     if (II.getIntrinsicID() == Intrinsic::stacksave) {
       eraseIfUnused(II, /*erase*/ true, /*check*/ false);

diff --git a/enzyme/test/Enzyme/ReverseMode/fence.ll b/enzyme/test/Enzyme/ReverseMode/fence.ll
@@ -0,0 +1,35 @@
+; RUN: %opt < %s %loadEnzyme -enzyme -enzyme-preopt=false -mem2reg -early-cse -instsimplify -simplifycfg -S | FileCheck %s
+
+; Function Attrs: noinline nounwind readnone uwtable
+define double @tester(double* %xp) {
+entry:
+  fence syncscope("singlethread") seq_cst
+  %x = load double, double* %xp, align 8
+  fence syncscope("singlethread") seq_cst
+  %x2 = fmul double %x, %x
+  ret double %x2
+}
+
+define double @test_derivative(double* %x, double* %y) {
+entry:
+  %0 = tail call double (double (double*)*, ...) @__enzyme_autodiff(double (double*)* nonnull @tester, double* %x, double* %y)
+  ret double %0
+}
+
+; Function Attrs: nounwind
+declare double @__enzyme_autodiff(double (double*)*, ...)
+
+; CHECK: define internal void @diffetester(double* %xp, double* %"xp'", double %differeturn) 
+; CHECK-NEXT: entry:
+; CHECK-NEXT:   fence syncscope("singlethread") seq_cst
+; CHECK-NEXT:   %x = load double, double* %xp, align 8
+; CHECK-NEXT:   fence syncscope("singlethread") seq_cst
+; CHECK-NEXT:   %m0diffex = fmul fast double %differeturn, %x
+; CHECK-NEXT:   %[[i0:.+]] = fadd fast double %m0diffex, %m0diffex
+; CHECK-NEXT:   fence syncscope("singlethread") seq_cst
+; CHECK-NEXT:   %[[i1:.+]] = load double, double* %"xp'", align 8
+; CHECK-NEXT:   %[[i2:.+]] = fadd fast double %[[i1]], %[[i0]]
+; CHECK-NEXT:   store double %[[i2]], double* %"xp'", align 8
+; CHECK-NEXT:   fence syncscope("singlethread") seq_cst
+; CHECK-NEXT:   ret void
+; CHECK-NEXT: }