feat: clear out unwanted no-ops from loop body

Author: avik-pal

src/enzyme_ad/jax/Passes/AutoBatching.cpp

@@ -831,7 +831,8 @@ bool GreedyWhileLoopBatchFission::liftElementwiseOp(
   if (elem->getNumOperands() != 1) {
     // For non-unary elementwise ops, we can
     //   1. lift operands if they are produced by DSs in the list
-    //   2. defined outside the loop body, in which case we simply do a broadcast_in_dim
+    //   2. defined outside the loop body, in which case we simply do a
+    //   broadcast_in_dim
     return false;
   }
 

src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp

@@ -22,11 +22,13 @@
 #include "mlir/IR/IRMapping.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/Visitors.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "shardy/dialect/sdy/ir/utils.h"
 #include "src/enzyme_ad/jax/Dialect/Dialect.h"
 #include "src/enzyme_ad/jax/Dialect/Ops.h"
+#include "src/enzyme_ad/jax/Implementations/WhileLoopInfo.h"
 #include "src/enzyme_ad/jax/Passes/EnzymeHLOPatterns.h"
 #include "src/enzyme_ad/jax/Passes/Passes.h"
 #include "src/enzyme_ad/jax/Passes/StructuredTensors.h"
@@ -23215,6 +23217,55 @@ struct CaseToIf : public CheckedOpRewritePattern<stablehlo::CaseOp, CaseToIf> {
   }
 };
 
+struct RemoveNoOpsFromWhileLoop
+    : public CheckedOpRewritePattern<stablehlo::WhileOp,
+                                     RemoveNoOpsFromWhileLoop> {
+  using CheckedOpRewritePattern<
+      stablehlo::WhileOp, RemoveNoOpsFromWhileLoop>::CheckedOpRewritePattern;
+
+  LogicalResult matchAndRewriteImpl(stablehlo::WhileOp whileOp,
+                                    PatternRewriter &rewriter) const {
+    auto info = WhileLoopInfo(whileOp);
+    auto computeInfoSuccess = info.computeInfo();
+    if (computeInfoSuccess.failed())
+      return computeInfoSuccess;
+
+    if (!info.isValid() || !info.isConstant())
+      return failure();
+
+    auto &whileBody = whileOp.getBody().front();
+    auto inductionVar = whileBody.getArgument(0);
+
+    bool anyNoOpRemoved = false;
+
+    auto limit = info.getConstantLimit().value();
+    auto start = info.getConstantStart().value();
+    auto step = info.getConstantStep().value();
+
+    // currently only removes remainder of induction variable
+    whileBody.walk([&](stablehlo::RemOp remOp) -> WalkResult {
+      if (remOp.getLhs() == inductionVar) {
+        SplatElementsAttr remRhsAttr;
+        if (matchPattern(remOp.getRhs(), m_Constant(&remRhsAttr))) {
+          auto attr = remRhsAttr.getSplatValue<Attribute>();
+          if (auto intAttr = dyn_cast<IntegerAttr>(attr)) {
+            auto rhsRemValue = intAttr.getValue();
+            APInt startAPInt(rhsRemValue.getBitWidth(), start, true);
+            APInt limitAPInt(rhsRemValue.getBitWidth(), limit, true);
+            if (rhsRemValue.sge(limitAPInt) && startAPInt.slt(rhsRemValue)) {
+              rewriter.replaceOp(remOp, remOp.getLhs());
+              anyNoOpRemoved = true;
+            }
+          }
+        }
+      }
+      return WalkResult::advance();
+    });
+
+    return anyNoOpRemoved ? success() : failure();
+  }
+};
+
 ///////////////  End Imported from stablehlo
 
 // clang-format off
@@ -23826,7 +23877,8 @@ struct EnzymeHLOOptPass
         MulReduceSliceFusion,
         MinReduceSliceFusion,
         MaxReduceSliceFusion,
-        CaseToIf
+        CaseToIf,
+        RemoveNoOpsFromWhileLoop
       >(context);
 
     patterns.add<

src/enzyme_ad/jax/TransformOps/TransformOps.td

@@ -2378,3 +2378,8 @@ def ApplyElementwiseSliceToBatchPatterns : EnzymeHLOPatternOp<
 def ApplyCaseToIfPatterns : EnzymeHLOPatternOp<"case_to_if"> {
   let patterns = ["CaseToIf"];
 }
+
+def ApplyRemoveNoOpsFromWhileLoopPatterns : EnzymeHLOPatternOp<
+    "remove_no_ops_from_while_loop"> {
+  let patterns = ["RemoveNoOpsFromWhileLoop"];
+}