Handle Freeze instruction

Author: wsmoses

enzyme/Enzyme/AdjointGenerator.h

@@ -232,6 +232,46 @@ class AdjointGenerator
     return B.CreateLoad(alloc);
   }
 
+#if LLVM_VERSION_MAJOR >= 10
+  void visitFreezeInst(llvm::FreezeInst &inst) {
+    eraseIfUnused(inst);
+    if (gutils->isConstantInstruction(&inst))
+      return;
+    Value *orig_op0 = inst.getOperand(0);
+
+    switch (Mode) {
+    case DerivativeMode::ReverseModeCombined:
+    case DerivativeMode::ReverseModeGradient: {
+      IRBuilder<> Builder2(inst.getParent());
+      getReverseBuilder(Builder2);
+
+      Value *idiff = diffe(&inst, Builder2);
+      Value *dif1 = Builder2.CreateFreeze(idiff);
+      setDiffe(&inst, Constant::getNullValue(inst.getType()), Builder2);
+      size_t size = 1;
+      if (inst.getType()->isSized())
+        size = (gutils->newFunc->getParent()->getDataLayout().getTypeSizeInBits(
+                    orig_op0->getType()) +
+                7) /
+               8;
+      addToDiffe(orig_op0, dif1, Builder2, TR.addingType(size, orig_op0));
+      return;
+    }
+    case DerivativeMode::ForwardMode: {
+      IRBuilder<> BuilderZ(&inst);
+      getForwardBuilder(BuilderZ);
+
+      Value *idiff = diffe(orig_op0, BuilderZ);
+      Value *dif1 = BuilderZ.CreateFreeze(idiff);
+      setDiffe(&inst, dif1, BuilderZ);
+      return;
+    }
+    case DerivativeMode::ReverseModePrimal:
+      return;
+    }
+  }
+#endif
+
   void visitInstruction(llvm::Instruction &inst) {
     // TODO explicitly handle all instructions rather than using the catch all
     // below

enzyme/Enzyme/DifferentialUseAnalysis.h

@@ -70,6 +70,9 @@ static inline bool is_use_directly_needed_in_reverse(
        cast<InsertElementInst>(user)->getOperand(2) != val) ||
       (isa<ExtractElementInst>(user) &&
        cast<ExtractElementInst>(user)->getIndexOperand() != val)
+#if LLVM_VERSION_MAJOR >= 10
+      || isa<FreezeInst>(user)
+#endif
       // isa<ExtractElement>(use) ||
       // isa<InsertElementInst>(use) || isa<ShuffleVectorInst>(use) ||
       // isa<ExtractValueInst>(use) || isa<AllocaInst>(use)

enzyme/Enzyme/GradientUtils.cpp

@@ -252,6 +252,21 @@ Value *GradientUtils::unwrapM(Value *const val, IRBuilder<> &BuilderM,
     return val;
   } else if (isa<AllocaInst>(val)) {
     return val;
+#if LLVM_VERSION_MAJOR >= 10
+  } else if (auto op = dyn_cast<FreezeInst>(val)) {
+    auto op0 = getOp(op->getOperand(0));
+    if (op0 == nullptr)
+      goto endCheck;
+    auto toreturn = BuilderM.CreateFreeze(op0, op->getName() + "_unwrap");
+    if (auto newi = dyn_cast<Instruction>(toreturn)) {
+      newi->copyIRFlags(op);
+      unwrappedLoads[newi] = val;
+    }
+    if (permitCache)
+      unwrap_cache[BuilderM.GetInsertBlock()][idx] = toreturn;
+    assert(val->getType() == toreturn->getType());
+    return toreturn;
+#endif
   } else if (auto op = dyn_cast<CastInst>(val)) {
     auto op0 = getOp(op->getOperand(0));
     if (op0 == nullptr)

enzyme/Enzyme/TypeAnalysis/TypeAnalysis.cpp

@@ -1514,6 +1514,15 @@ void TypeAnalyzer::visitIntToPtrInst(IntToPtrInst &I) {
     updateAnalysis(I.getOperand(0), getAnalysis(&I), &I);
 }
 
+#if LLVM_VERSION_MAJOR >= 10
+void TypeAnalyzer::visitFreezeInst(FreezeInst &I) {
+  if (direction & DOWN)
+    updateAnalysis(&I, getAnalysis(I.getOperand(0)), &I);
+  if (direction & UP)
+    updateAnalysis(I.getOperand(0), getAnalysis(&I), &I);
+}
+#endif
+
 void TypeAnalyzer::visitBitCastInst(BitCastInst &I) {
   if (I.getType()->isIntOrIntVectorTy() || I.getType()->isFPOrFPVectorTy()) {
     if (direction & DOWN)

enzyme/Enzyme/TypeAnalysis/TypeAnalysis.h

@@ -282,6 +282,10 @@ class TypeAnalyzer : public llvm::InstVisitor<TypeAnalyzer> {
 
   void visitBitCastInst(llvm::BitCastInst &I);
 
+#if LLVM_VERSION_MAJOR >= 10
+  void visitFreezeInst(llvm::FreezeInst &I);
+#endif
+
   void visitSelectInst(llvm::SelectInst &I);
 
   void visitExtractElementInst(llvm::ExtractElementInst &I);

enzyme/test/Enzyme/ReverseMode/freeze.ll

@@ -0,0 +1,26 @@
+; RUN: if [ %llvmver -ge 10 ]; then %opt < %s %loadEnzyme -enzyme -enzyme-preopt=false -mem2reg -instsimplify -simplifycfg -S | FileCheck %s; fi
+
+; Function Attrs: nounwind readnone uwtable
+define double @tester(double %x) {
+entry:
+  %out = freeze double %x
+  ret double %out
+}
+
+define double @test_derivative(double %x) {
+entry:
+  %0 = tail call double (double (double)*, ...) @__enzyme_autodiff(double (double)* nonnull @tester, double %x)
+  ret double %0
+}
+
+declare double @cbrt(double)
+
+; Function Attrs: nounwind
+declare double @__enzyme_autodiff(double (double)*, ...)
+
+; CHECK: define internal { double } @diffetester(double %x, double %differeturn) {
+; CHECK-NEXT: entry:
+; CHECK-NEXT:   %0 = freeze double %differeturn
+; CHECK-NEXT:   %1 = insertvalue { double } undef, double %0, 0
+; CHECK-NEXT:   ret { double } %1
+; CHECK-NEXT: }