Move boolAttr checks in each analysis to shared check in localGuaranteedWithSetAttr

Author: gaurav-arya

src/enzyme_ad/jax/Utils.cpp

@@ -748,13 +748,6 @@ NoNanResultAnalysis::localGuaranteed(Operation *op,
                                      PatternRewriter &rewriter) {
   assert(op);
 
-  if (auto boolAttr = op->getAttrOfType<BoolAttr>(getAttrName())) {
-    if (boolAttr.getValue())
-      return State::GUARANTEED;
-    else
-      return State::NOTGUARANTEED;
-  }
-
   DenseElementsAttr denseAttr;
   if (matchPattern(op, m_Constant(&denseAttr))) {
     if (guaranteedConstantOp(op, denseAttr, rewriter)) {
@@ -891,13 +884,6 @@ FiniteResultAnalysis::localGuaranteed(Operation *op,
                                       PatternRewriter &rewriter) {
   assert(op);
 
-  if (auto boolAttr = op->getAttrOfType<BoolAttr>(getAttrName())) {
-    if (boolAttr.getValue())
-      return State::GUARANTEED;
-    else
-      return State::NOTGUARANTEED;
-  }
-
   DenseElementsAttr denseAttr;
   if (matchPattern(op, m_Constant(&denseAttr))) {
     if (guaranteedConstantOp(op, denseAttr, rewriter)) {
@@ -1004,13 +990,6 @@ NonNegativeResultAnalysis::State NonNegativeResultAnalysis::localGuaranteed(
     PatternRewriter &rewriter) {
   assert(op);
 
-  if (auto boolAttr = op->getAttrOfType<BoolAttr>(getAttrName())) {
-    if (boolAttr.getValue())
-      return State::GUARANTEED;
-    else
-      return State::NOTGUARANTEED;
-  }
-
   DenseElementsAttr denseAttr;
   if (matchPattern(op, m_Constant(&denseAttr))) {
     if (guaranteedConstantOp(op, denseAttr, rewriter)) {

src/enzyme_ad/jax/Utils.h

@@ -541,8 +541,18 @@ template <typename Child> class GuaranteedResultAnalysisBase {
   State localGuaranteedWithSetAttr(Operation *op,
                                    SmallVectorImpl<Operation *> &localtodo,
                                    PatternRewriter &rewriter) {
-    auto state = ((Child *)this)->localGuaranteed(op, localtodo, rewriter);
+
     auto attrName = ((Child *)this)->getAttrName();
+
+    if (auto boolAttr = op->getAttrOfType<BoolAttr>(attrName)) {
+      if (boolAttr.getValue())
+        return State::GUARANTEED;
+      else
+        return State::NOTGUARANTEED;
+    }
+
+    auto state = ((Child *)this)->localGuaranteed(op, localtodo, rewriter);
+
     switch (state) {
     case State::GUARANTEED:
       rewriter.modifyOpInPlace(op, [&]() {