refactor: generalize existing passes to operate on traits (#1603)

* fix: bug fix and generalize transpose_unary_transpose patterns

* feat: generalize commutative associative reordering based on traits

* refactor: rewrite associative bin op

* feat: generalize BinaryOpTransposeSimplify

* ci: use specific reactant commit

* chore: run fmt

* fix: induction var in remove_no_ops

* test: use atol and rtol

* fix: tolerance

Author: avik-pal

.github/workflows/test-gb-25.yml

@@ -53,8 +53,8 @@ jobs:
           - 'main'
           # - '0123456789abcdef0123456789abcdef01234567'
         reactant_commit:
-          - 'main'
-          # - 'regenerate-mlir-bindings'
+          # - 'main'
+          - 'ap/pass_generalization'
 
     steps:
       - name: Check GPUs

src/enzyme_ad/jax/Implementations/WhileLoopInfo.h

@@ -39,7 +39,17 @@ struct WhileLoopInfo {
   std::optional<int64_t> getConstantStart();
   std::optional<int64_t> getConstantLimit();
 
-  Value getInductionVariable() { return op.getBody().front().getArgument(0); }
+  // assumes computeInfo() has been called and was successful
+  // returns the induction variable in the body of the while op
+  Value getInductionVariable() {
+    auto &condBlk = op.getCond().front();
+    auto condTerm = cast<stablehlo::ReturnOp>(condBlk.getTerminator());
+    auto condV = condTerm->getOperand(0);
+    auto cond = condV.getDefiningOp<stablehlo::CompareOp>();
+    auto induct = dyn_cast<BlockArgument>(cond.getOperand(0));
+    auto blockArgNum = induct.getArgNumber();
+    return op.getBody().front().getArgument(blockArgNum);
+  }
 
   int64_t getConstantNumIters();
   Value getNumIters(OpBuilder &builder);