fix: guaranteed analysis bfs queue was busted (#1538)

* fix: guaranteed analysis bfs queue was busted

* chore: fmt

* fix: try making things deterministic

* chore: bazel format

* fix: address review comments

* fix: revert to SmallPtrSet

* chore: run fmt

* fix: reuse found value

Author: avik-pal

src/enzyme_ad/jax/BUILD

@@ -693,6 +693,7 @@ cc_library(
         "@llvm-project//mlir:Support",
         "@llvm-project//mlir:ViewLikeInterface",
         "@stablehlo//:base",
+        "@stablehlo//:chlo_ops",
         "@stablehlo//:stablehlo_ops",
     ],
 )

src/enzyme_ad/jax/Utils.cpp

@@ -30,6 +30,7 @@
 #include "llvm/ADT/SetVector.h"
 #include "llvm/Support/Debug.h"
 
+#include "stablehlo/dialect/ChloOps.h"
 #include "stablehlo/dialect/StablehloOps.h"
 
 #include <set>
@@ -885,6 +886,10 @@ NonNegativeResultAnalysis::State NonNegativeResultAnalysis::localGuaranteed(
     return State::GUARANTEED;
   }
 
+  if (isa<chlo::ErfInvOp>(op)) {
+    return State::NOTGUARANTEED;
+  }
+
   // Any non-negative operation that produces a non-negative result
   // Here we recur on the rhs, as that is more likely to be a constant.
   if (isa<stablehlo::MaxOp>(op)) {

src/enzyme_ad/jax/Utils.h

@@ -7,6 +7,8 @@
 #include "mlir/IR/Types.h"
 #include "llvm/ADT/APFloat.h"
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/MapVector.h"
+#include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallVector.h"
 
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
@@ -354,7 +356,7 @@ template <typename Child> class GuaranteedResultAnalysisBase {
 
     // Map of operations we have seen before. The target of the map[o] is a list
     // of sub-queries, that if all true prove that `o` is no-nan.
-    DenseMap<Operation *, SmallPtrSet<Operation *, 2>> seen;
+    llvm::MapVector<Operation *, llvm::SmallPtrSet<Operation *, 2>> seen;
 
     // Inverse of seen. A map of operations `p` we still need to prove, to a
     // list of values that require `p` to be proven.
@@ -454,10 +456,12 @@ template <typename Child> class GuaranteedResultAnalysisBase {
       case State::PENDING: {
         assert(localtodo.size());
         assert(seen.find(cur) == seen.end());
-        SmallPtrSet<Operation *, 2> set(localtodo.begin(), localtodo.end());
-        for (auto v : set) {
+        for (auto v : localtodo) {
           reverseSeen[v].push_back(cur);
+          todo.push_back(v);
         }
+        llvm::SmallPtrSet<Operation *, 2> set(localtodo.begin(),
+                                              localtodo.end());
         seen[cur] = std::move(set);
         break;
       }
@@ -475,11 +479,15 @@ template <typename Child> class GuaranteedResultAnalysisBase {
       });
     }
 
-    assert(opCache.find(op) != opCache.end());
-    rewriter.modifyOpInPlace(op, [&]() {
-      op->setAttr(attrName, BoolAttr::get(op->getContext(), true));
-    });
-    return true;
+    auto found = opCache.find(op);
+    if (found != opCache.end()) {
+      bool guaranteed = found->second;
+      rewriter.modifyOpInPlace(op, [&]() {
+        op->setAttr(attrName, BoolAttr::get(op->getContext(), guaranteed));
+      });
+      return guaranteed;
+    }
+    return false;
   }
 
   bool guaranteed(stablehlo::ConstantOp constOp, PatternRewriter &rewriter) {

test/lit_tests/abspositive.mlir

@@ -9,8 +9,8 @@ func.func @test1() -> tensor<18500xi64> {
 }
 
 // CHECK:  func.func @test1() -> tensor<18500xi64> {
-// CHECK-NEXT:    %c = stablehlo.constant dense<1> : tensor<18500xi64>
-// CHECK-NEXT:    %0 = stablehlo.iota dim = 0 : tensor<18500xi64>
+// CHECK-NEXT:    %c = stablehlo.constant {enzymexla.guaranteed_non_negative = true} dense<1> : tensor<18500xi64>
+// CHECK-NEXT:    %0 = stablehlo.iota dim = 0 {enzymexla.guaranteed_non_negative = true} : tensor<18500xi64>
 // CHECK-NEXT:    %1 = stablehlo.add %0, %c {enzymexla.guaranteed_non_negative = true} : tensor<18500xi64>
 // CHECK-NEXT:    return %1 : tensor<18500xi64>
 // CHECK-NEXT:  }
@@ -45,7 +45,27 @@ func.func @test4(%arg0: tensor<12xf64>) -> tensor<4x3xf64> {
 }
 
 // CHECK:  func.func @test4(%arg0: tensor<12xf64>) -> tensor<4x3xf64> {
-// CHECK-NEXT:    %0 = stablehlo.multiply %arg0, %arg0 : tensor<12xf64>
+// CHECK-NEXT:    %0 = stablehlo.multiply %arg0, %arg0 {enzymexla.guaranteed_non_negative = true} : tensor<12xf64>
 // CHECK-NEXT:    %1 = stablehlo.reshape %0 {enzymexla.guaranteed_non_negative = true} : (tensor<12xf64>) -> tensor<4x3xf64>
 // CHECK-NEXT:    return %1 : tensor<4x3xf64>
 // CHECK-NEXT:  }
+
+func.func @test5(%arg0: tensor<12xf64>, %arg1: tensor<12xf64>, %arg2: tensor<12xf64>) -> tensor<12xf64> {
+    // CHECK: chlo.erf_inv %arg0 {enzymexla.guaranteed_non_negative = false}
+    %0 = chlo.erf_inv %arg0 : tensor<12xf64> -> tensor<12xf64>
+    %1 = stablehlo.add %arg1, %arg2 : tensor<12xf64>
+    %2 = stablehlo.add %0, %1 : tensor<12xf64>
+    // CHECK: stablehlo.abs
+    %3 = stablehlo.abs %2 : tensor<12xf64>
+    return %3 : tensor<12xf64>
+}
+
+func.func @test6(%arg0: tensor<12xf64>, %arg1: tensor<12xf64>, %arg2: tensor<12xf64>) -> tensor<12xf64> {
+    %cst = stablehlo.constant dense<-1.000000e+00> : tensor<12xf64>
+    %0 = stablehlo.multiply %arg0, %cst : tensor<12xf64>
+    %1 = stablehlo.add %arg1, %arg2 : tensor<12xf64>
+    %2 = stablehlo.add %0, %1 : tensor<12xf64>
+    // CHECK: stablehlo.abs
+    %3 = stablehlo.abs %2 : tensor<12xf64>
+    return %3 : tensor<12xf64>
+}

test/lit_tests/raising/affine_to_stablehlo6.mlir

@@ -111,7 +111,7 @@ module {
 // CHECK-NEXT:    %c_4 = stablehlo.constant dense<-92> : tensor<185xi64>
 // CHECK-NEXT:    %0 = stablehlo.iota dim = 0 : tensor<185xi64>
 // CHECK-NEXT:    %1 = stablehlo.add %0, %c_4 : tensor<185xi64>
-// CHECK-NEXT:    %2 = stablehlo.convert %1 {enzymexla.guaranteed_finite = true, enzymexla.guaranteed_non_negative = false} : (tensor<185xi64>) -> tensor<185xf64>
+// CHECK-NEXT:    %2 = stablehlo.convert %1 {enzymexla.guaranteed_finite = true, enzymexla.guaranteed_no_nan = true, enzymexla.guaranteed_non_negative = false} : (tensor<185xi64>) -> tensor<185xf64>
 // CHECK-NEXT:    %3 = stablehlo.add %cst_3, %2 {enzymexla.guaranteed_finite = true} : tensor<185xf64>
 // CHECK-NEXT:    %4 = stablehlo.multiply %3, %cst_0 : tensor<185xf64>
 // CHECK-NEXT:    %5 = stablehlo.cosine %4 : tensor<185xf64>
@@ -128,7 +128,7 @@ module {
 // CHECK-NEXT:    %16 = stablehlo.slice %arg3 [0:1] : (tensor<186xf64>) -> tensor<1xf64>
 // CHECK-NEXT:    %17 = stablehlo.concatenate %16, %6, dim = 0 : (tensor<1xf64>, tensor<185xf64>) -> tensor<186xf64>
 // CHECK-NEXT:    %18 = stablehlo.add %0, %c_2 : tensor<185xi64>
-// CHECK-NEXT:    %19 = stablehlo.convert %18 {enzymexla.guaranteed_finite = true} : (tensor<185xi64>) -> tensor<185xf64>
+// CHECK-NEXT:    %19 = stablehlo.convert %18 {enzymexla.guaranteed_finite = true, enzymexla.guaranteed_no_nan = true} : (tensor<185xi64>) -> tensor<185xf64>
 // CHECK-NEXT:    %20 = stablehlo.multiply %19, %cst_0 : tensor<185xf64>
 // CHECK-NEXT:    %21 = stablehlo.sine %20 : tensor<185xf64>
 // CHECK-NEXT:    %22 = stablehlo.sine %9 : tensor<185xf64>
@@ -138,7 +138,7 @@ module {
 // CHECK-NEXT:    %26 = stablehlo.concatenate %25, %24, dim = 0 : (tensor<1xf64>, tensor<185xf64>) -> tensor<186xf64>
 // CHECK-NEXT:    %27 = stablehlo.sine %4 : tensor<185xf64>
 // CHECK-NEXT:    %28 = stablehlo.add %0, %c : tensor<185xi64>
-// CHECK-NEXT:    %29 = stablehlo.convert %28 {enzymexla.guaranteed_finite = true, enzymexla.guaranteed_non_negative = false} : (tensor<185xi64>) -> tensor<185xf64>
+// CHECK-NEXT:    %29 = stablehlo.convert %28 {enzymexla.guaranteed_finite = true, enzymexla.guaranteed_no_nan = true, enzymexla.guaranteed_non_negative = false} : (tensor<185xi64>) -> tensor<185xf64>
 // CHECK-NEXT:    %30 = stablehlo.add %cst_3, %29 {enzymexla.guaranteed_finite = true} : tensor<185xf64>
 // CHECK-NEXT:    %31 = stablehlo.multiply %30, %cst_0 : tensor<185xf64>
 // CHECK-NEXT:    %32 = stablehlo.sine %31 : tensor<185xf64>