Update src/enzyme_ad/jax/Passes/WhileLoopOutsideValuesAddToArgumentListPass.cpp

Co-authored-by: Paul Berg <naydex.mc+github@gmail.com>

Author: avik-pal

src/enzyme_ad/jax/Passes/WhileLoopOutsideValuesAddToArgumentListPass.cpp

@@ -27,26 +27,10 @@ struct SHLOWhileOpUpdateArgumentListPattern final
 
   LogicalResult matchAndRewrite(stablehlo::WhileOp whileOp,
                                 PatternRewriter &rewriter) const override {
-    // Collect values used inside cond/body that are defined outside the WhileOp
-    SmallVector<Value, 4> extraValues;
-    SmallPtrSet<Value, 8> seen;
+    SetVector<Value> extraValuesSet;
+    getUsedValuesDefinedAbove(whileOp->getRegions(), extraValuesSet);
+    SmallVector<Value> extraValues = extraValuesSet.takeVector();
 
-    auto collectExternal = [&](Region &region) {
-      region.walk([&](Operation *op) {
-        for (OpOperand &operand : op->getOpOperands()) {
-          Value v = operand.get();
-          if (!v)
-            continue;
-          if (definedOutside(v, whileOp) && !seen.contains(v)) {
-            seen.insert(v);
-            extraValues.push_back(v);
-          }
-        }
-      });
-    };
-
-    collectExternal(whileOp.getCond());
-    collectExternal(whileOp.getBody());
 
     if (extraValues.empty())
       return failure();