feat: hoist chain of ops

Author: avik-pal

src/enzyme_ad/jax/Implementations/WhileLoopInfo.cpp

@@ -214,7 +214,8 @@ bool WhileLoopInfo::isConstantValue(Value v, llvm::APInt &constVal) {
     return true;
 
   Value outerValue;
-  if (isConstantAcrossIterations(v, outerValue) &&
+  SmallVector<Operation *> canBeHoisted;
+  if (isConstantAcrossIterations(v, outerValue, canBeHoisted, false) &&
       matchPattern(outerValue, m_ConstantInt(&constVal)))
     return true;
   return false;
@@ -295,11 +296,13 @@ void WhileLoopInfo::propagateAffineIndexInfo() {
 
 bool WhileLoopInfo::isConstantAcrossIterations(Value v, bool checkOperands) {
   Value outerValue;
-  return isConstantAcrossIterations(v, outerValue, checkOperands);
+  SmallVector<Operation *> canBeHoisted;
+  return isConstantAcrossIterations(v, outerValue, canBeHoisted, checkOperands);
 }
 
-bool WhileLoopInfo::isConstantAcrossIterations(Value v, Value &outerValue,
-                                               bool checkOperands) {
+bool WhileLoopInfo::isConstantAcrossIterations(
+    Value v, Value &outerValue, SmallVector<Operation *> &canBeHoisted,
+    bool checkOperands) {
   if (definedOutside(v, op)) {
     outerValue = v;
     return true;
@@ -326,12 +329,15 @@ bool WhileLoopInfo::isConstantAcrossIterations(Value v, Value &outerValue,
 
   // all operands of the defining op are constant across iterations
   // don't populate the outerValue in this case
-  return llvm::all_of(defOp->getOperands(), [&](Value operand) {
-    // TODO: we should do `isConstantAcrossIterations` but for now we do a more
-    // conservative check
-    // return isConstantAcrossIterations(operand);
-    return definedOutside(operand, op);
-  });
+  if (llvm::all_of(defOp->getOperands(), [&](Value operand) {
+        return isConstantAcrossIterations(operand, outerValue, canBeHoisted,
+                                          true);
+      })) {
+    outerValue = nullptr;
+    canBeHoisted.push_back(defOp);
+    return true;
+  }
+  return false;
 }
 
 template <typename OpTy>
@@ -423,8 +429,6 @@ bool WhileLoopInfo::hoistOperationFromLoop(
   if (!canHoistOperationFromLoop(sliceOp, dimensions))
     return false;
 
-  auto totalIterCount = getConstantNumIters();
-
   auto depIndex = sliceOp.getStartIndices()[sliceIndex];
   auto indexTy = depIndex.getType();
 
@@ -439,7 +443,6 @@ bool WhileLoopInfo::hoistOperationFromLoop(
   auto step = getConstantStep().value();
   auto lb = getConstantStart().value();
   auto ub = getConstantLimit().value();
-  int64_t N = ub - lb;
 
   auto rawMin = scale * lb + offset;
   auto rawMax = scale * (ub - 1) + offset;

src/enzyme_ad/jax/Implementations/WhileLoopInfo.h

@@ -65,6 +65,7 @@ struct WhileLoopInfo {
 
   bool isConstantAcrossIterations(Value v, bool checkOperands = true);
   bool isConstantAcrossIterations(Value v, Value &outerValue,
+                                  SmallVector<Operation *> &canBeHoisted,
                                   bool checkOperands = true);
 
   bool canHoistOperationFromLoop(mlir::stablehlo::DynamicSliceOp sliceOp,

src/enzyme_ad/jax/Passes/AutoBatching.cpp

@@ -1,6 +1,7 @@
 #include "src/enzyme_ad/jax/Passes/AutoBatching.h"
 
 #include "Enzyme/MLIR/Passes/EnzymeBatchPass.h"
+#include "mlir/Analysis/TopologicalSortUtils.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Matchers.h"
@@ -925,11 +926,14 @@ bool GreedyWhileLoopBatchFission::liftOperationByBatching(
   SmallVector<SmallVector<int64_t>> sliceDims(op->getNumOperands());
   SmallVector<int64_t> hoistedDims(op->getNumOperands());
   SmallVector<DynamicSliceInfo> mappedSliceInfos(op->getNumOperands());
+  DenseMap<Value, SmallVector<Operation *>> hoistMap;
+
   for (int i = 0; i < op->getNumOperands(); i++) {
     auto operand = op->getOperand(i);
 
     Value outerValue;
-    if (info.isConstantAcrossIterations(operand, outerValue)) {
+    SmallVector<Operation *> canBeHoisted;
+    if (info.isConstantAcrossIterations(operand, outerValue, canBeHoisted)) {
       if (outerValue) {
         SplatElementsAttr splat;
         if (matchPattern(operand, m_Constant(&splat))) {
@@ -939,6 +943,7 @@ bool GreedyWhileLoopBatchFission::liftOperationByBatching(
         }
         batchOperands[i] = outerValue;
       } else {
+        hoistMap[operand] = canBeHoisted;
         hoistedDims[i] = cast<mlir::OpResult>(operand).getResultNumber();
         batchLiftingModes[i] = BatchLiftingMode::NEEDS_HOISTING_OUTSIDE_WHILE;
         batchOperands[i] = operand;
@@ -1040,6 +1045,34 @@ bool GreedyWhileLoopBatchFission::liftOperationByBatching(
 
   rewriter.setInsertionPoint(whileOp);
 
+  // hoist any operations that can be hoisted
+  DenseMap<Value, Value> hoistedValues;
+  for (auto &[val, ops] : hoistMap) {
+    llvm::SetVector<Operation *> toHoist(ops.begin(), ops.end());
+    auto sorted = mlir::topologicalSort(toHoist);
+    IRMapping mapper;
+
+    for (auto &op : sorted) {
+      for (auto operand : op->getOperands()) {
+        if (!definedOutside(operand, whileOp)) {
+          Value outerValue;
+          SmallVector<Operation *> canBeHoisted;
+          if (info.isConstantAcrossIterations(operand, outerValue, canBeHoisted,
+                                              false)) {
+            mapper.map(operand, outerValue);
+          }
+        }
+      }
+      auto hoisted = rewriter.clone(*op, mapper);
+      for (auto [origRes, newRes] :
+           llvm::zip(op->getResults(), hoisted->getResults())) {
+        mapper.map(origRes, newRes);
+      }
+    }
+
+    hoistedValues[val] = mapper.lookup(val);
+  }
+
   SmallVector<Value> newOperands;
   for (auto [consType, baseOp, sliceDim, sliceInfo, hoistDim] :
        llvm::zip(batchLiftingModes, batchOperands, sliceDims, mappedSliceInfos,
@@ -1097,9 +1130,8 @@ bool GreedyWhileLoopBatchFission::liftOperationByBatching(
       break;
     }
     case BatchLiftingMode::NEEDS_HOISTING_OUTSIDE_WHILE: {
-      auto hoisted = rewriter.clone(*baseOp.getDefiningOp());
-      baseOp = hoisted->getResult(hoistDim);
-      // intentionally fallthrough
+      baseOp = hoistedValues[baseOp];
+      LLVM_FALLTHROUGH;
     }
     case BatchLiftingMode::DEFINED_OUTSIDE_WHILE: {
       auto operandShape = operandType.getShape();

src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp

@@ -2210,7 +2210,6 @@ struct DUSPad final
 
     SmallVector<int64_t> newDusStartIndexValues;
     SmallVector<Value> newDusStartIndices;
-    SmallVector<int64_t> newDusStartIndexValues;
     Location loc = dus.getLoc();
     auto indexElementType =
         cast<ShapedType>(startIndices[0].getType()).getElementType();

test/lit_tests/autobatching/higher_order_post_diff.mlir

@@ -1,4 +1,4 @@
-// RUN: enzymexlamlir-opt --enzyme-hlo-generate-td="patterns=dot_general_licm(0);elementwise_licm(0);greedy_while_loop_batch_fission;while_is_copy_simplify;remove_no_ops_from_while_loop;dynamic_slice_reshape_dynamic_slice" --transform-interpreter --enzyme-hlo-remove-transform --inline --enzyme-hlo-opt %s | FileCheck %s
+// RUN: enzymexlamlir-opt --enzyme-hlo-generate-td="patterns=greedy_while_loop_batch_fission;while_is_copy_simplify;remove_no_ops_from_while_loop;dynamic_slice_reshape_dynamic_slice" --transform-interpreter --enzyme-hlo-remove-transform --enzyme-hlo-opt %s | FileCheck %s
 
 func.func @main(%arg0: tensor<5x5xf32>, %arg1: tensor<5xf32>, %arg2: tensor<3x5xf32>) -> tensor<3x5xf32> {
     %cst = stablehlo.constant dense<"0x000000400000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000004000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000040000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000400000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000004000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000040000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000400000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000004000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000040000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000400000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000004000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000040000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000400000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000004000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000040"> : tensor<15x3x5xf32>
@@ -83,29 +83,29 @@ func.func @main(%arg0: tensor<5x5xf32>, %arg1: tensor<5xf32>, %arg2: tensor<3x5x
 // CHECK-NEXT:     %4 = stablehlo.dot_general %0, %cst_0, batching_dims = [0] x [0], contracting_dims = [1] x [2] : (tensor<15x5x5xf32>, tensor<15x3x5xf32>) -> tensor<15x5x3xf32>
 // CHECK-NEXT:     %5 = stablehlo.dot_general %arg0, %arg2, contracting_dims = [0] x [1] : (tensor<5x5xf32>, tensor<3x5xf32>) -> tensor<5x3xf32>
 // CHECK-NEXT:     %6 = stablehlo.add %5, %3 : tensor<5x3xf32>
-// CHECK-NEXT:     %7 = stablehlo.broadcast_in_dim %6, dims = [1, 2] : (tensor<5x3xf32>) -> tensor<15x5x3xf32>
-// CHECK-NEXT:     %8 = stablehlo.multiply %7, %4 : tensor<15x5x3xf32>
-// CHECK-NEXT:     %9 = stablehlo.multiply %8, %cst : tensor<15x5x3xf32>
-// CHECK-NEXT:     %10 = stablehlo.multiply %cst_4, %6 : tensor<5x3xf32>
-// CHECK-NEXT:     %11 = stablehlo.broadcast_in_dim %10, dims = [1, 2] : (tensor<5x3xf32>) -> tensor<15x5x3xf32>
-// CHECK-NEXT:     %12 = stablehlo.multiply %9, %11 : tensor<15x5x3xf32>
-// CHECK-NEXT:     %13 = stablehlo.multiply %6, %6 : tensor<5x3xf32>
-// CHECK-NEXT:     %14 = stablehlo.multiply %13, %cst_3 : tensor<5x3xf32>
-// CHECK-NEXT:     %15 = stablehlo.add %14, %cst_2 : tensor<5x3xf32>
-// CHECK-NEXT:     %16 = stablehlo.broadcast_in_dim %15, dims = [1, 2] : (tensor<5x3xf32>) -> tensor<15x5x3xf32>
-// CHECK-NEXT:     %17 = stablehlo.multiply %2, %16 : tensor<15x5x3xf32>
-// CHECK-NEXT:     %18 = stablehlo.add %17, %12 : tensor<15x5x3xf32>
-// CHECK-NEXT:     %19 = stablehlo.multiply %10, %15 : tensor<5x3xf32>
-// CHECK-NEXT:     %20 = stablehlo.logistic %19 : tensor<5x3xf32>
-// CHECK-NEXT:     %21 = stablehlo.broadcast_in_dim %20, dims = [1, 2] : (tensor<5x3xf32>) -> tensor<15x5x3xf32>
-// CHECK-NEXT:     %22 = stablehlo.multiply %1, %21 : tensor<15x5x3xf32>
-// CHECK-NEXT:     %23 = stablehlo.logistic %19 : tensor<5x3xf32>
+// CHECK-NEXT:     %7 = stablehlo.multiply %6, %6 : tensor<5x3xf32>
+// CHECK-NEXT:     %8 = stablehlo.multiply %7, %cst_3 : tensor<5x3xf32>
+// CHECK-NEXT:     %9 = stablehlo.add %8, %cst_2 : tensor<5x3xf32>
+// CHECK-NEXT:     %10 = stablehlo.broadcast_in_dim %9, dims = [1, 2] : (tensor<5x3xf32>) -> tensor<15x5x3xf32>
+// CHECK-NEXT:     %11 = stablehlo.multiply %2, %10 : tensor<15x5x3xf32>
+// CHECK-NEXT:     %12 = stablehlo.broadcast_in_dim %6, dims = [1, 2] : (tensor<5x3xf32>) -> tensor<15x5x3xf32>
+// CHECK-NEXT:     %13 = stablehlo.multiply %12, %4 : tensor<15x5x3xf32>
+// CHECK-NEXT:     %14 = stablehlo.multiply %cst_4, %6 : tensor<5x3xf32>
+// CHECK-NEXT:     %15 = stablehlo.multiply %14, %9 : tensor<5x3xf32>
+// CHECK-NEXT:     %16 = stablehlo.logistic %15 : tensor<5x3xf32>
+// CHECK-NEXT:     %17 = stablehlo.broadcast_in_dim %16, dims = [1, 2] : (tensor<5x3xf32>) -> tensor<15x5x3xf32>
+// CHECK-NEXT:     %18 = stablehlo.multiply %1, %17 : tensor<15x5x3xf32>
+// CHECK-NEXT:     %19 = stablehlo.multiply %13, %cst : tensor<15x5x3xf32>
+// CHECK-NEXT:     %20 = stablehlo.broadcast_in_dim %14, dims = [1, 2] : (tensor<5x3xf32>) -> tensor<15x5x3xf32>
+// CHECK-NEXT:     %21 = stablehlo.multiply %19, %20 : tensor<15x5x3xf32>
+// CHECK-NEXT:     %22 = stablehlo.add %11, %21 : tensor<15x5x3xf32>
+// CHECK-NEXT:     %23 = stablehlo.logistic %15 : tensor<5x3xf32>
 // CHECK-NEXT:     %24 = stablehlo.subtract %cst_2, %23 : tensor<5x3xf32>
 // CHECK-NEXT:     %25 = stablehlo.multiply %23, %24 : tensor<5x3xf32>
 // CHECK-NEXT:     %26 = stablehlo.broadcast_in_dim %25, dims = [1, 2] : (tensor<5x3xf32>) -> tensor<15x5x3xf32>
-// CHECK-NEXT:     %27 = stablehlo.multiply %18, %26 : tensor<15x5x3xf32>
-// CHECK-NEXT:     %28 = stablehlo.multiply %27, %7 : tensor<15x5x3xf32>
-// CHECK-NEXT:     %29 = stablehlo.add %22, %28 : tensor<15x5x3xf32>
+// CHECK-NEXT:     %27 = stablehlo.multiply %22, %26 : tensor<15x5x3xf32>
+// CHECK-NEXT:     %28 = stablehlo.multiply %27, %12 : tensor<15x5x3xf32>
+// CHECK-NEXT:     %29 = stablehlo.add %18, %28 : tensor<15x5x3xf32>
 // CHECK-NEXT:     %30:2 = stablehlo.while(%iterArg = %c_7, %iterArg_12 = %cst_9) : tensor<i64>, tensor<3x5xf32>
 // CHECK-NEXT:     cond {
 // CHECK-NEXT:       %31 = stablehlo.compare  LT, %iterArg, %c_11 : (tensor<i64>, tensor<i64>) -> tensor<i1>