Fix powi nan (rust-lang#507)

Author: wsmoses

enzyme/Enzyme/AdjointGenerator.h

@@ -3328,10 +3328,10 @@ class AdjointGenerator
         if (vdiff && !gutils->isConstantValue(orig_ops[0])) {
           Value *op0 = gutils->getNewFromOriginal(orig_ops[0]);
           Value *op1 = gutils->getNewFromOriginal(orig_ops[1]);
+          Value *nop1 = lookup(op1, Builder2);
           SmallVector<Value *, 2> args = {
               lookup(op0, Builder2),
-              Builder2.CreateSub(lookup(op1, Builder2),
-                                 ConstantInt::get(op1->getType(), 1))};
+              Builder2.CreateSub(nop1, ConstantInt::get(op1->getType(), 1))};
           auto &CI = cast<CallInst>(I);
 #if LLVM_VERSION_MAJOR >= 11
           auto *PowF = CI.getCalledOperand();
@@ -3349,6 +3349,9 @@ class AdjointGenerator
               Builder2.CreateFMul(vdiff, cal),
               Builder2.CreateSIToFP(lookup(op1, Builder2),
                                     op0->getType()->getScalarType()));
+          dif0 = Builder2.CreateSelect(
+              Builder2.CreateICmpEQ(ConstantInt::get(nop1->getType(), 0), nop1),
+              Constant::getNullValue(dif0->getType()), dif0);
           addToDiffe(orig_ops[0], dif0, Builder2, I.getType());
         }
         return;
@@ -3808,8 +3811,12 @@ class AdjointGenerator
               Builder2.CreateSIToFP(op1, op0->getType()->getScalarType());
           Value *op = diffe(orig_ops[0], Builder2);
 
+          Value *cmp = Builder2.CreateICmpEQ(
+              ConstantInt::get(args[1]->getType(), 0), op1);
           auto rule = [&](Value *op) {
-            return Builder2.CreateFMul(Builder2.CreateFMul(op, cal), cast);
+            return Builder2.CreateSelect(
+                cmp, Constant::getNullValue(op->getType()),
+                Builder2.CreateFMul(Builder2.CreateFMul(op, cal), cast));
           };
 
           Value *dif0 = applyChainRule(I.getType(), Builder2, rule, op);

enzyme/test/Enzyme/ForwardMode/powi13.ll

@@ -1,4 +1,4 @@
-; RUN: if [ %llvmver -ge 13 ]; then %opt < %s %loadEnzyme -enzyme -enzyme-preopt=false -mem2reg -instsimplify -simplifycfg -S | FileCheck %s; fi
+; RUN: %opt < %s %loadEnzyme -enzyme -enzyme-preopt=false -mem2reg -instsimplify -simplifycfg -S | FileCheck %s
 
 ; Function Attrs: noinline nounwind readnone uwtable
 define double @tester(double %x, i32 %y) {
@@ -22,10 +22,12 @@ declare double @__enzyme_fwddiff(double (double, i32)*, ...)
 ; CHECK: define internal {{(dso_local )?}}double @fwddiffetester(double %x, double %"x'", i32 %y)
 ; CHECK-NEXT: entry:
 ; CHECK-NEXT:   %[[ym1:.+]] = sub i32 %y, 1
-; CHECK-NEXT:   %[[newpow:.+]] = call fast double @llvm.powi.f64.i32(double %x, i32 %[[ym1]])
+; CHECK-NEXT:   %[[newpow:.+]] = call fast double @llvm.powi.f64{{(\.i32)?}}(double %x, i32 %[[ym1]])
 ; CHECK-DAG:    %[[sitofp:.+]] = sitofp i32 %y to double
+; CHECK-DAG:    %[[cmp:.+]] = icmp eq i32 0, %y
 ; CHECK-DAG:    %[[newpowdret:.+]] = fmul fast double %"x'", %[[newpow]]
 ; CHECK-NEXT:   %[[dx:.+]] = fmul fast double %[[newpowdret]], %[[sitofp]]
-; CHECK-NEXT:   ret double %[[dx:.+]]
+; CHECK-NEXT:   %[[res:.+]] = select {{(fast )?}}i1 %[[cmp]], double 0.000000e+00, double %[[dx]]
+; CHECK-NEXT:   ret double %[[res]]
 ; CHECK-NEXT: }
 

enzyme/test/Enzyme/ForwardModeVector/powi13.ll

@@ -1,4 +1,4 @@
-; RUN: if [ %llvmver -ge 13 ]; then %opt < %s %loadEnzyme -enzyme -enzyme-preopt=false -mem2reg -instsimplify -simplifycfg -S | FileCheck %s; fi
+; RUN: %opt < %s %loadEnzyme -enzyme -enzyme-preopt=false -mem2reg -instsimplify -simplifycfg -S | FileCheck %s
 
 %struct.Gradients = type { double, double, double }
 
@@ -22,22 +22,26 @@ entry:
 declare double @llvm.powi.f64.i32(double, i32)
 
 
-; CHECK: define internal [3 x double] @fwddiffe3tester(double %x, [3 x double] %"x'", i32 %y) #1 {
-; CHECK-NEXT entry:
-; CHECK-NEXT   %0 = sub i32 %y, 1
-; CHECK-NEXT   %1 = call fast double @llvm.powi.f64.i32(double %x, i32 %0)
-; CHECK-NEXT   %2 = sitofp i32 %y to double
-; CHECK-NEXT   %3 = extractvalue [3 x double] %"x'", 0
-; CHECK-NEXT   %4 = fmul fast double %3, %1
-; CHECK-NEXT   %5 = fmul fast double %4, %2
-; CHECK-NEXT   %6 = insertvalue [3 x double] undef, double %5, 0
-; CHECK-NEXT   %7 = extractvalue [3 x double] %"x'", 1
-; CHECK-NEXT   %8 = fmul fast double %7, %1
-; CHECK-NEXT   %9 = fmul fast double %8, %2
-; CHECK-NEXT   %10 = insertvalue [3 x double] %6, double %9, 1
-; CHECK-NEXT   %11 = extractvalue [3 x double] %"x'", 2
-; CHECK-NEXT   %12 = fmul fast double %11, %1
-; CHECK-NEXT   %13 = fmul fast double %12, %2
-; CHECK-NEXT   %14 = insertvalue [3 x double] %10, double %13, 2
-; CHECK-NEXT   ret [3 x double] %14
-; CHECK-NEXT }
+; CHECK: define internal [3 x double] @fwddiffe3tester(double %x, [3 x double] %"x'", i32 %y)
+; CHECK-NEXT: entry:
+; CHECK-NEXT:   %0 = sub i32 %y, 1
+; CHECK-NEXT:   %1 = call fast double @llvm.powi.f64{{(\.i32)?}}(double %x, i32 %0)
+; CHECK-NEXT:   %2 = sitofp i32 %y to double
+; CHECK-NEXT:   %3 = icmp eq i32 0, %y
+; CHECK-NEXT:   %4 = extractvalue [3 x double] %"x'", 0
+; CHECK-NEXT:   %5 = fmul fast double %4, %1
+; CHECK-NEXT:   %6 = fmul fast double %5, %2
+; CHECK-NEXT:   %7 = select {{(fast )?}}i1 %3, double 0.000000e+00, double %6
+; CHECK-NEXT:   %8 = insertvalue [3 x double] undef, double %7, 0
+; CHECK-NEXT:   %9 = extractvalue [3 x double] %"x'", 1
+; CHECK-NEXT:   %10 = fmul fast double %9, %1
+; CHECK-NEXT:   %11 = fmul fast double %10, %2
+; CHECK-NEXT:   %12 = select {{(fast )?}}i1 %3, double 0.000000e+00, double %11
+; CHECK-NEXT:   %13 = insertvalue [3 x double] %8, double %12, 1
+; CHECK-NEXT:   %14 = extractvalue [3 x double] %"x'", 2
+; CHECK-NEXT:   %15 = fmul fast double %14, %1
+; CHECK-NEXT:   %16 = fmul fast double %15, %2
+; CHECK-NEXT:   %17 = select {{(fast )?}}i1 %3, double 0.000000e+00, double %16
+; CHECK-NEXT:   %18 = insertvalue [3 x double] %13, double %17, 2
+; CHECK-NEXT:   ret [3 x double] %18
+; CHECK-NEXT }

enzyme/test/Enzyme/ReverseMode/jl_powi.ll

@@ -1,42 +1,43 @@
-; RUN: if [ %llvmver -ge 13 ]; then %opt < %s %loadEnzyme -enzyme -enzyme-preopt=false -mem2reg -instsimplify -simplifycfg -S | FileCheck %s; fi
+; RUN: %opt < %s %loadEnzyme -enzyme -enzyme-preopt=false -mem2reg -instsimplify -simplifycfg -S | FileCheck %s
+
 source_filename = "text"
 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-pc-linux-gnu"
 
-define private fastcc double @julia___2797(double %0, i64 signext %1) unnamed_addr #0 !dbg !7 {
+define private fastcc double @julia___2797(double %x0, i64 signext %x1) unnamed_addr #0 {
 top:
-  switch i64 %1, label %L20 [
+  switch i64 %x1, label %L20 [
     i64 -1, label %L3
     i64 0, label %L7
     i64 1, label %L7.fold.split
     i64 2, label %L13
     i64 3, label %L17
-  ], !dbg !9
+  ]
 
 L3:                                               ; preds = %top
-  %2 = fdiv double 1.000000e+00, %0, !dbg !10
-  ret double %2, !dbg !9
+  %x2 = fdiv double 1.000000e+00, %x0
+  ret double %x2
 
 L7.fold.split:                                    ; preds = %top
-  br label %L7, !dbg !16
+  br label %L7
 
 L7:                                               ; preds = %top, %L7.fold.split
-  %merge = phi double [ 1.000000e+00, %top ], [ %0, %L7.fold.split ]
-  ret double %merge, !dbg !16
+  %merge = phi double [ 1.000000e+00, %top ], [ %x0, %L7.fold.split ]
+  ret double %merge
 
 L13:                                              ; preds = %top
-  %3 = fmul double %0, %0, !dbg !17
-  ret double %3, !dbg !19
+  %x3 = fmul double %x0, %x0
+  ret double %x3
 
 L17:                                              ; preds = %top
-  %4 = fmul double %0, %0, !dbg !20
-  %5 = fmul double %4, %0, !dbg !20
-  ret double %5, !dbg !24
+  %x4 = fmul double %x0, %x0
+  %x5 = fmul double %x4, %x0
+  ret double %x5
 
 L20:                                              ; preds = %top
-  %6 = sitofp i64 %1 to double, !dbg !25
-  %7 = call double @llvm.pow.f64(double %0, double %6), !dbg !27
-  ret double %7, !dbg !27
+  %x6 = sitofp i64 %x1 to double
+  %x7 = call double @llvm.pow.f64(double %x0, double %x6)
+  ret double %x7
 }
 
 ; Function Attrs: nofree nosync nounwind readnone speculatable willreturn
@@ -46,10 +47,10 @@ declare double @llvm.pow.f64(double, double) #1
 declare double @__enzyme_autodiff(double (double, i64)*, ...)
 
 ; Function Attrs: alwaysinline nosync readnone
-define double @julia_f_2794(double %0, i64 signext %1) local_unnamed_addr #2 !dbg !28 {
+define double @julia_f_2794(double %y0, i64 signext %y1) {
 entry:
-  %2 = call fastcc double @julia___2797(double %0, i64 signext %1) #5, !dbg !29
-  ret double %2
+  %y2 = call fastcc double @julia___2797(double %y0, i64 signext %y1) #5
+  ret double %y2
 }
 
 define double @test_derivative(double %x, i64 %y) {
@@ -58,13 +59,15 @@ entry:
   ret double %0
 }
 
-; CHECK: define internal { double } @diffejulia_f_2794(double %0, i64 signext %1, double %differeturn) local_unnamed_addr #5 !dbg !35 {
+; CHECK: define internal { double } @diffejulia_f_2794(double %y0, i64 signext %y1, double %differeturn)
 ; CHECK-NEXT: entry:
-; CHECK-NEXT:   %2 = sub i64 %1, 1
-; CHECK-NEXT:   %3 = call fast fastcc double @julia___2797(double %0, i64 %2), !dbg !36
-; CHECK-NEXT:   %4 = sitofp i64 %1 to double
-; CHECK-NEXT:   %5 = fmul fast double %differeturn, %3
-; CHECK-NEXT:   %6 = fmul fast double %5, %4
+; CHECK-NEXT:   %0 = sub i64 %y1, 1
+; CHECK-NEXT:   %1 = call fast fastcc double @julia___2797(double %y0, i64 %0)
+; CHECK-NEXT:   %2 = sitofp i64 %y1 to double
+; CHECK-NEXT:   %3 = fmul fast double %differeturn, %1
+; CHECK-NEXT:   %4 = fmul fast double %3, %2
+; CHECK-NEXT:   %5 = icmp eq i64 0, %y1
+; CHECK-NEXT:   %6 = select {{(fast )?}}i1 %5, double 0.000000e+00, double %4
 ; CHECK-NEXT:   %7 = insertvalue { double } undef, double %6, 0
 ; CHECK-NEXT:   ret { double } %7
 ; CHECK-NEXT: }
@@ -78,44 +81,10 @@ declare noalias nonnull {} addrspace(10)* @jl_gc_pool_alloc(i8*, i32, i32) #4
 ; Function Attrs: allocsize(1)
 declare noalias nonnull {} addrspace(10)* @jl_gc_big_alloc(i8*, i64) #4
 
-attributes #0 = { noinline nosync readnone "enzyme_math"="powi" "enzyme_shouldrecompute"="powi" "probe-stack"="inline-asm" }
-attributes #1 = { nofree nosync nounwind readnone speculatable willreturn }
-attributes #2 = { alwaysinline nosync readnone "probe-stack"="inline-asm" }
+attributes #0 = { noinline readnone "enzyme_math"="powi" "enzyme_shouldrecompute"="powi"}
+attributes #1 = { nounwind readnone speculatable }
+attributes #2 = { alwaysinline readnone "probe-stack"="inline-asm" }
 attributes #3 = { inaccessiblemem_or_argmemonly }
 attributes #4 = { allocsize(1) }
 attributes #5 = { "probe-stack"="inline-asm" }
 
-!llvm.module.flags = !{!0, !1}
-!llvm.dbg.cu = !{!2, !5}
-
-!0 = !{i32 2, !"Dwarf Version", i32 4}
-!1 = !{i32 2, !"Debug Info Version", i32 3}
-!2 = distinct !DICompileUnit(language: DW_LANG_Julia, file: !3, producer: "julia", isOptimized: true, runtimeVersion: 0, emissionKind: LineTablesOnly, enums: !4, nameTableKind: None)
-!3 = !DIFile(filename: "math.jl", directory: ".")
-!4 = !{}
-!5 = distinct !DICompileUnit(language: DW_LANG_Julia, file: !6, producer: "julia", isOptimized: true, runtimeVersion: 0, emissionKind: LineTablesOnly, enums: !4, nameTableKind: None)
-!6 = !DIFile(filename: "REPL[3]", directory: ".")
-!7 = distinct !DISubprogram(name: "^", linkageName: "julia_^_2797", scope: null, file: !3, line: 922, type: !8, scopeLine: 922, spFlags: DISPFlagDefinition | DISPFlagOptimized, unit: !2, retainedNodes: !4)
-!8 = !DISubroutineType(types: !4)
-!9 = !DILocation(line: 923, scope: !7)
-!10 = !DILocation(line: 408, scope: !11, inlinedAt: !13)
-!11 = distinct !DISubprogram(name: "/;", linkageName: "/", scope: !12, file: !12, type: !8, spFlags: DISPFlagDefinition | DISPFlagOptimized, unit: !2, retainedNodes: !4)
-!12 = !DIFile(filename: "float.jl", directory: ".")
-!13 = !DILocation(line: 243, scope: !14, inlinedAt: !9)
-!14 = distinct !DISubprogram(name: "inv;", linkageName: "inv", scope: !15, file: !15, type: !8, spFlags: DISPFlagDefinition | DISPFlagOptimized, unit: !2, retainedNodes: !4)
-!15 = !DIFile(filename: "number.jl", directory: ".")
-!16 = !DILocation(line: 924, scope: !7)
-!17 = !DILocation(line: 405, scope: !18, inlinedAt: !19)
-!18 = distinct !DISubprogram(name: "*;", linkageName: "*", scope: !12, file: !12, type: !8, spFlags: DISPFlagDefinition | DISPFlagOptimized, unit: !2, retainedNodes: !4)
-!19 = !DILocation(line: 926, scope: !7)
-!20 = !DILocation(line: 405, scope: !18, inlinedAt: !21)
-!21 = !DILocation(line: 655, scope: !22, inlinedAt: !24)
-!22 = distinct !DISubprogram(name: "*;", linkageName: "*", scope: !23, file: !23, type: !8, spFlags: DISPFlagDefinition | DISPFlagOptimized, unit: !2, retainedNodes: !4)
-!23 = !DIFile(filename: "operators.jl", directory: ".")
-!24 = !DILocation(line: 927, scope: !7)
-!25 = !DILocation(line: 146, scope: !26, inlinedAt: !27)
-!26 = distinct !DISubprogram(name: "Float64;", linkageName: "Float64", scope: !12, file: !12, type: !8, spFlags: DISPFlagDefinition | DISPFlagOptimized, unit: !2, retainedNodes: !4)
-!27 = !DILocation(line: 928, scope: !7)
-!28 = distinct !DISubprogram(name: "f", linkageName: "julia_f_2794", scope: null, file: !6, line: 1, type: !8, scopeLine: 1, spFlags: DISPFlagDefinition | DISPFlagOptimized, unit: !5, retainedNodes: !4)
-!29 = !DILocation(line: 1, scope: !28, inlinedAt: !30)
-!30 = distinct !DILocation(line: 0, scope: !28)

enzyme/test/Enzyme/ReverseMode/powi13.ll

@@ -1,4 +1,4 @@
-; RUN: if [ %llvmver -ge 13 ]; then %opt < %s %loadEnzyme -enzyme -enzyme-preopt=false -mem2reg -instsimplify -simplifycfg -S | FileCheck %s; fi
+; RUN: %opt < %s %loadEnzyme -enzyme -enzyme-preopt=false -mem2reg -instsimplify -simplifycfg -S | FileCheck %s
 
 ; Function Attrs: noinline nounwind readnone uwtable
 define double @tester(double %x, i32 %y) {
@@ -22,10 +22,12 @@ declare double @__enzyme_autodiff(double (double, i32)*, ...)
 ; CHECK: define internal {{(dso_local )?}}{ double } @diffetester(double %x, i32 %y, double %differeturn)
 ; CHECK-NEXT: entry:
 ; CHECK-NEXT:   %[[ym1:.+]] = sub i32 %y, 1
-; CHECK-NEXT:   %[[newpow:.+]] = call fast double @llvm.powi.f64.i32(double %x, i32 %[[ym1]])
+; CHECK-NEXT:   %[[newpow:.+]] = call fast double @llvm.powi.f64{{(\.i32)?}}(double %x, i32 %[[ym1]])
 ; CHECK-DAG:    %[[sitofp:.+]] = sitofp i32 %y to double
 ; CHECK-DAG:    %[[newpowdret:.+]] = fmul fast double %differeturn, %[[newpow]]
 ; CHECK-NEXT:   %[[dx:.+]] = fmul fast double %[[newpowdret]], %[[sitofp]]
-; CHECK-NEXT:   %[[interres:.+]] = insertvalue { double } undef, double %[[dx:.+]], 0
+; CHECK-NEXT:   %[[cmp:.+]] = icmp eq i32 0, %y
+; CHECK-NEXT:   %[[res:.+]] = select {{(fast )?}}i1 %[[cmp]], double 0.000000e+00, double %[[dx]]
+; CHECK-NEXT:   %[[interres:.+]] = insertvalue { double } undef, double %[[res]], 0
 ; CHECK-NEXT:   ret { double } %[[interres]]
 ; CHECK-NEXT: }