Forward mode cast inst (#194)

* implemented forward mode visitCastInst

Author: tgymnich

enzyme/Enzyme/AdjointGenerator.h

@@ -608,51 +608,75 @@ class AdjointGenerator
         I.getOpcode() == CastInst::CastOps::PtrToInt)
       return;
 
-    if (Mode == DerivativeMode::ReverseModePrimal)
+    switch (Mode) {
+    case DerivativeMode::ReverseModePrimal: {
       return;
+    }
+    case DerivativeMode::ReverseModeGradient:
+    case DerivativeMode::ReverseModeCombined: {
+      Value *orig_op0 = I.getOperand(0);
+      Value *op0 = gutils->getNewFromOriginal(orig_op0);
 
-    Value *orig_op0 = I.getOperand(0);
-    Value *op0 = gutils->getNewFromOriginal(orig_op0);
+      IRBuilder<> Builder2(I.getParent());
+      getReverseBuilder(Builder2);
 
-    IRBuilder<> Builder2(I.getParent());
-    getReverseBuilder(Builder2);
+      if (!gutils->isConstantValue(orig_op0)) {
+        Value *dif = diffe(&I, Builder2);
 
-    if (!gutils->isConstantValue(orig_op0)) {
-      Value *dif = diffe(&I, Builder2);
-
-      size_t size = 1;
-      if (orig_op0->getType()->isSized())
-        size = (gutils->newFunc->getParent()->getDataLayout().getTypeSizeInBits(
-                    orig_op0->getType()) +
-                7) /
-               8;
-      Type *FT = TR.addingType(size, orig_op0);
-      if (!FT) {
-        llvm::errs() << " " << *gutils->oldFunc << "\n";
-        TR.dump();
-        llvm::errs() << " " << *orig_op0 << "\n";
+        size_t size = 1;
+        if (orig_op0->getType()->isSized())
+          size =
+              (gutils->newFunc->getParent()->getDataLayout().getTypeSizeInBits(
+                   orig_op0->getType()) +
+               7) /
+              8;
+        Type *FT = TR.addingType(size, orig_op0);
+        if (!FT) {
+          llvm::errs() << " " << *gutils->oldFunc << "\n";
+          TR.dump();
+          llvm::errs() << " " << *orig_op0 << "\n";
+        }
+        assert(FT);
+        if (I.getOpcode() == CastInst::CastOps::FPTrunc ||
+            I.getOpcode() == CastInst::CastOps::FPExt) {
+          addToDiffe(orig_op0, Builder2.CreateFPCast(dif, op0->getType()),
+                     Builder2, FT);
+        } else if (I.getOpcode() == CastInst::CastOps::BitCast) {
+          addToDiffe(orig_op0, Builder2.CreateBitCast(dif, op0->getType()),
+                     Builder2, FT);
+        } else if (I.getOpcode() == CastInst::CastOps::Trunc) {
+          // TODO CHECK THIS
+          auto trunced = Builder2.CreateZExt(dif, op0->getType());
+          addToDiffe(orig_op0, trunced, Builder2, FT);
+        } else {
+          TR.dump();
+          llvm::errs() << *I.getParent()->getParent() << "\n"
+                       << *I.getParent() << "\n";
+          llvm::errs() << "cannot handle above cast " << I << "\n";
+          report_fatal_error("unknown instruction");
+        }
       }
-      assert(FT);
-      if (I.getOpcode() == CastInst::CastOps::FPTrunc ||
-          I.getOpcode() == CastInst::CastOps::FPExt) {
-        addToDiffe(orig_op0, Builder2.CreateFPCast(dif, op0->getType()),
-                   Builder2, FT);
-      } else if (I.getOpcode() == CastInst::CastOps::BitCast) {
-        addToDiffe(orig_op0, Builder2.CreateBitCast(dif, op0->getType()),
-                   Builder2, FT);
-      } else if (I.getOpcode() == CastInst::CastOps::Trunc) {
-        // TODO CHECK THIS
-        auto trunced = Builder2.CreateZExt(dif, op0->getType());
-        addToDiffe(orig_op0, trunced, Builder2, FT);
+      setDiffe(&I, Constant::getNullValue(I.getType()), Builder2);
+
+      break;
+    }
+    case DerivativeMode::ForwardMode: {
+      Value *orig_op0 = I.getOperand(0);
+
+      IRBuilder<> Builder2(&I);
+      getForwardBuilder(Builder2);
+
+      if (!gutils->isConstantValue(orig_op0)) {
+        Value *dif = diffe(orig_op0, Builder2);
+        setDiffe(&I, Builder2.CreateCast(I.getOpcode(), dif, I.getType()),
+                 Builder2);
       } else {
-        TR.dump();
-        llvm::errs() << *I.getParent()->getParent() << "\n"
-                     << *I.getParent() << "\n";
-        llvm::errs() << "cannot handle above cast " << I << "\n";
-        report_fatal_error("unknown instruction");
+        setDiffe(&I, Constant::getNullValue(I.getType()), Builder2);
       }
+
+      break;
+    }
     }
-    setDiffe(&I, Constant::getNullValue(I.getType()), Builder2);
   }
 
   void visitSelectInst(llvm::SelectInst &SI) {

enzyme/test/Enzyme/ForwardMode/bitcast.ll

@@ -0,0 +1,22 @@
+; RUN: %opt < %s %loadEnzyme -enzyme -enzyme-preopt=false -mem2reg -simplifycfg -adce -instsimplify -S | FileCheck %s
+
+define double @tester(double %x) {
+entry:
+  %y = bitcast double %x to i64
+  %z = bitcast i64 %y to double
+  ret double %z
+}
+
+define double @test_derivative(double %x) {
+entry:
+  %0 = tail call double (double (double)*, ...) @__enzyme_fwddiff(double (double)* nonnull @tester, double %x, double 1.0)
+  ret double %0
+}
+
+declare double @__enzyme_fwddiff(double (double)*, ...)
+
+; CHECK: define internal {{(dso_local )?}}{ double } @diffetester(double %x, double %"x'") {
+; CHECK-NEXT: entry:
+; CHECK-NEXT:   %0 = insertvalue { double } undef, double %"x'", 0
+; CHECK-NEXT:   ret { double } %0
+; CHECK-NEXT: }

enzyme/test/Enzyme/ForwardMode/fpext.ll

@@ -0,0 +1,22 @@
+; RUN: %opt < %s %loadEnzyme -enzyme -enzyme-preopt=false -mem2reg -simplifycfg -adce -instsimplify -S | FileCheck %s
+
+define double @tester(float %x) {
+entry:
+  %y = fpext float %x to double
+  ret double %y
+}
+
+define double @test_derivative(float %x) {
+entry:
+  %0 = tail call double (double (float)*, ...) @__enzyme_fwddiff(double (float)* nonnull @tester, float %x, float 1.0)
+  ret double %0
+}
+
+declare double @__enzyme_fwddiff(double (float)*, ...)
+
+; CHECK: define internal {{(dso_local )?}}{ double } @diffetester(float %x, float %"x'") {
+; CHECK-NEXT: entry:
+; CHECK-NEXT:   %0 = fpext float %"x'" to double
+; CHECK-NEXT:   %1 = insertvalue { double } undef, double %0, 0
+; CHECK-NEXT:   ret { double } %1
+; CHECK-NEXT: }