refactor: support dynamicsliceop

Author: avik-pal

src/enzyme_ad/jax/Passes/AutoBatching.cpp

@@ -265,6 +265,19 @@ constructAndExtractBatchOperands(PatternRewriter &rewriter,
   return std::make_tuple(operands, operandIndexMap);
 }
 
+std::tuple<bool, bool> allSameBool(const SmallVector<bool> &bools) {
+  return {
+      llvm::all_of(bools, [&](bool b) { return b == bools.front(); }),
+      bools.front(),
+  };
+}
+
+bool allOpsAreUnique(const SmallVector<Operation *> &ops) {
+  SmallPtrSet<Operation *, 8> seen;
+  return llvm::all_of(ops,
+                      [&](Operation *op) { return seen.insert(op).second; });
+}
+
 LogicalResult ConcatInsertDimToBatchBase::matchAndRewriteImpl(
     stablehlo::ConcatenateOp concatOp, PatternRewriter &rewriter) const {
   if (concatOp.getNumOperands() <= 1)
@@ -732,29 +745,8 @@ bool SliceToBatchBase::areSlicesContiguous(
   return true;
 }
 
-std::tuple<bool, bool>
-SliceToBatchBase::allSameBool(const SmallVector<bool> &bools) const {
-  bool first = bools[0];
-  for (auto b : bools) {
-    if (b != first)
-      return {false, first};
-  }
-  return {true, first};
-}
-
-bool SliceToBatchBase::allOpsAreUnique(
-    const SmallVector<Operation *> &ops) const {
-  SmallPtrSet<Operation *, 8> seen;
-  for (auto op : ops) {
-    if (!seen.insert(op).second)
-      return false;
-  }
-  return true;
-}
-
-LogicalResult
-GreedyWhileLoopBatchFission::matchAndRewrite(stablehlo::WhileOp whileOp,
-                                             PatternRewriter &rewriter) const {
+LogicalResult GreedyWhileLoopBatchFission::matchAndRewriteImpl(
+    stablehlo::WhileOp whileOp, PatternRewriter &rewriter) const {
   auto info = WhileLoopInfo(whileOp);
   auto computeInfoSuccess = info.computeInfo();
   if (computeInfoSuccess.failed())

src/enzyme_ad/jax/Passes/AutoBatching.h

@@ -17,6 +17,9 @@ struct WhileLoopInfo;
 }; // namespace enzyme
 }; // namespace mlir
 
+std::tuple<bool, bool> allSameBool(const llvm::SmallVector<bool> &bools);
+bool allOpsAreUnique(const llvm::SmallVector<mlir::Operation *> &ops);
+
 struct BatchOperandConstructionInfo {
   mlir::stablehlo::SliceOp sliceOp;
   int32_t sliceOperandIndex;
@@ -25,6 +28,27 @@ struct BatchOperandConstructionInfo {
   bool intermediateReshape;
 };
 
+// TODO: update old code to use this new slice info
+template <typename OpTy> struct NewSliceInfo {
+  OpTy sliceOp;
+  llvm::SmallVector<mlir::Value> dynamicStartIndices;
+  llvm::SmallVector<int64_t> startIndices;
+  llvm::SmallVector<int64_t> sliceSizes;
+  int64_t sliceDim;
+  int64_t sliceStart;
+  bool supported;
+};
+
+NewSliceInfo<mlir::stablehlo::SliceOp>
+constructNewSliceInfo(mlir::stablehlo::SliceOp sliceOp);
+NewSliceInfo<mlir::stablehlo::DynamicSliceOp>
+constructNewSliceInfo(mlir::stablehlo::DynamicSliceOp sliceOp);
+
+bool areSlicesContiguous(
+    llvm::SmallVector<NewSliceInfo<mlir::stablehlo::SliceOp>> &slices);
+bool areSlicesContiguous(
+    llvm::SmallVector<NewSliceInfo<mlir::stablehlo::DynamicSliceOp>> &slices);
+
 struct ConcatInsertDimToBatchBase
     : public mlir::enzyme::CheckedOpRewritePattern<
           mlir::stablehlo::ConcatenateOp, ConcatInsertDimToBatchBase> {
@@ -108,11 +132,6 @@ struct SliceToBatchBase
 
   SliceInfo extractSliceInfo(mlir::stablehlo::SliceOp slice) const;
   bool areSlicesContiguous(llvm::SmallVector<SliceInfo> &slices) const;
-  std::tuple<bool, bool>
-  allSameBool(const llvm::SmallVector<bool> &bools) const;
-  std::tuple<bool, bool>
-  allSameBool(const llvm::SmallVector<mlir::Operation *> &ops) const;
-  bool allOpsAreUnique(const llvm::SmallVector<mlir::Operation *> &ops) const;
 
 protected:
   std::function<mlir::Operation *(mlir::Operation *)> isValidTargetOp;
@@ -182,13 +201,16 @@ struct SliceToBatchElementwise : public SliceToBatchBase {
 };
 
 struct GreedyWhileLoopBatchFission
-    : public mlir::OpRewritePattern<mlir::stablehlo::WhileOp> {
-  using Base = mlir::OpRewritePattern<mlir::stablehlo::WhileOp>;
+    : public mlir::enzyme::CheckedOpRewritePattern<
+          mlir::stablehlo::WhileOp, GreedyWhileLoopBatchFission> {
+  using Base =
+      mlir::enzyme::CheckedOpRewritePattern<mlir::stablehlo::WhileOp,
+                                            GreedyWhileLoopBatchFission>;
   using Base::Base;
 
   mlir::LogicalResult
-  matchAndRewrite(mlir::stablehlo::WhileOp whileOp,
-                  mlir::PatternRewriter &rewriter) const override;
+  matchAndRewriteImpl(mlir::stablehlo::WhileOp whileOp,
+                      mlir::PatternRewriter &rewriter) const;
 
 private:
   struct DynamicSliceInfo {