Fix fmuladd fwd inactive (rust-lang#705)

EnzymeAD · Jun 24, 2022 · fb22d2f · fb22d2f
1 parent 3d28f83
commit fb22d2f
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 2 deletions.
diff --git a/enzyme/Enzyme/AdjointGenerator.h b/enzyme/Enzyme/AdjointGenerator.h
@@ -3917,8 +3917,13 @@ class AdjointGenerator
                           : diffe(orig_ops[2], Builder2);
 
         auto rule = [&](Value *dif0, Value *dif1, Value *dif2) {
-          Value *dif = Builder2.CreateFAdd(Builder2.CreateFMul(op0, dif1),
-                                           Builder2.CreateFMul(op1, dif0));
+          Value *dif =
+              Builder2.CreateFAdd(gutils->isConstantValue(orig_ops[1])
+                                      ? Constant::getNullValue(opType1)
+                                      : Builder2.CreateFMul(op0, dif1),
+                                  gutils->isConstantValue(orig_ops[2])
+                                      ? Constant::getNullValue(opType2)
+                                      : Builder2.CreateFMul(op1, dif0));
           return Builder2.CreateFAdd(dif, dif2);
         };
 

diff --git a/enzyme/test/Enzyme/ForwardMode/fmuladd2.ll b/enzyme/test/Enzyme/ForwardMode/fmuladd2.ll
@@ -0,0 +1,26 @@
+; RUN: %opt < %s %loadEnzyme -enzyme -enzyme-preopt=false -mem2reg -gvn -simplifycfg -instcombine -S | FileCheck %s
+
+; Function Attrs: nounwind readnone uwtable
+define double @tester(double* %yp) {
+entry:
+  %y = load double, double* %yp, align 8
+  %0 = tail call fast double @llvm.fmuladd.f64(double 1.000000e+00, double %y, double 0.000000e+00)
+  ret double %0
+}
+
+define double @test_derivative(double* %yp, double* %dyp) {
+entry:
+  %0 = tail call double (double (double*)*, ...) @__enzyme_fwddiff(double (double*)* nonnull @tester, double* %yp, double* %dyp)
+  ret double %0
+}
+
+declare double @llvm.fmuladd.f64(double %a, double %b, double %c)
+
+; Function Attrs: nounwind
+declare double @__enzyme_fwddiff(double (double*)*, ...)
+
+; CHECK: define internal double @fwddiffetester(double* %yp, double* %"yp'")
+; CHECK-NEXT: entry:
+; CHECK-NEXT:   %"y'ipl" = load double, double* %"yp'", align 8
+; CHECK-NEXT:   ret double %"y'ipl"
+; CHECK-NEXT: }