Merge branch 'main' into jm/multirotate_recognize

Author: wsmoses

src/enzyme_ad/jax/Implementations/StableHLOAutoDiffOpInterfaceImpl.cpp

@@ -3939,6 +3939,42 @@ struct SHLOReverseOpBatchInterface
   }
 };
 
+struct SHLOPadOpBatchInterface
+    : public BatchOpInterface::ExternalModel<SHLOPadOpBatchInterface,
+                                             stablehlo::PadOp> {
+  mlir::LogicalResult createBatch(Operation *src, OpBuilder &builder,
+                                  IRMapping &mapper,
+                                  ArrayRef<int64_t> batchSizes) const {
+    auto op = cast<stablehlo::PadOp>(src);
+
+    auto batchedPadValue = mapper.lookup(op.getPaddingValue());
+    auto scalarPadValue =
+        getScalarValue(batchedPadValue.getDefiningOp(), builder);
+    if (!scalarPadValue) {
+      return genericCreateBatch(src, builder, mapper, batchSizes);
+    }
+
+    int64_t nBatches = batchSizes.size();
+    SmallVector<int64_t> newLow(nBatches, 0);
+    newLow.append(op.getEdgePaddingLow().begin(), op.getEdgePaddingLow().end());
+    SmallVector<int64_t> newHigh(nBatches, 0);
+    newHigh.append(op.getEdgePaddingHigh().begin(),
+                   op.getEdgePaddingHigh().end());
+    SmallVector<int64_t> newInterior(nBatches, 0);
+    newInterior.append(op.getInteriorPadding().begin(),
+                       op.getInteriorPadding().end());
+
+    auto newPadOp = stablehlo::PadOp::create(
+        builder, op.getLoc(), mapper.lookup(op.getOperand()), scalarPadValue,
+        builder.getDenseI64ArrayAttr(newLow),
+        builder.getDenseI64ArrayAttr(newHigh),
+        builder.getDenseI64ArrayAttr(newInterior));
+
+    mapper.map(src->getResult(0), newPadOp.getResult());
+    return success();
+  }
+};
+
 // https://github.com/jax-ml/jax/blob/2a8cb54b82f1b0d17181d43f9be78d2b349df333/jax/_src/lax/convolution.py#L613-L629
 struct SHLOConvolutionOpBatchInterface
     : public BatchOpInterface::ExternalModel<SHLOConvolutionOpBatchInterface,
@@ -4110,6 +4146,7 @@ void mlir::enzyme::registerStableHLODialectAutoDiffInterface(
         *context);
     ReverseOp::attachInterface<SHLOReverseOpBatchInterface>(*context);
     ConvolutionOp::attachInterface<SHLOConvolutionOpBatchInterface>(*context);
+    PadOp::attachInterface<SHLOPadOpBatchInterface>(*context);
 
     ScatterOp::attachInterface<SHLOGenericBatchOpInterface<ScatterOp>>(
         *context); // TODO: simpler version with newly named dims

src/enzyme_ad/jax/Passes/AutoBatching.cpp

@@ -976,8 +976,7 @@ LogicalResult GreedyWhileLoopBatchFission::matchAndRewriteImpl(
   for (auto &[op, slices] : userOpToSlicesMap) {
     bool avoidBatching =
         llvm::TypeSwitch<Operation *, bool>(op)
-            .Case<stablehlo::DynamicSliceOp, stablehlo::ReshapeOp,
-                  stablehlo::SliceOp,
+            .Case<stablehlo::ReshapeOp, stablehlo::SliceOp,
                   // TODO: avoid scatter since that lowers to loop right now
                   stablehlo::ScatterOp>([=](auto op) { return true; })
             .Case<stablehlo::BroadcastInDimOp, stablehlo::TransposeOp>(
@@ -987,9 +986,13 @@ LogicalResult GreedyWhileLoopBatchFission::matchAndRewriteImpl(
       continue;
     }
 
-    if ((dyn_cast<BatchOpInterface>(op) ||
-         stablehlo::hasTraitElementwise(op)) &&
-        op->getNumResults() == 1) {
+    if (auto dsOp = dyn_cast<stablehlo::DynamicSliceOp>(op)) {
+      if (raiseDynamicSliceToGather(rewriter, whileOp, slices, dsOp, info)) {
+        anyOpRewritten = true;
+      }
+    } else if ((dyn_cast<BatchOpInterface>(op) ||
+                stablehlo::hasTraitElementwise(op)) &&
+               op->getNumResults() == 1) {
       if (liftOperationByBatching(rewriter, whileOp, slices, op, info)) {
         anyOpRewritten = true;
       } else if (liftReduceLikeOperation(rewriter, whileOp, slices, op, info)) {
@@ -1440,6 +1443,247 @@ bool liftReduceLikeOperation(
   return true;
 }
 
+bool raiseDynamicSliceToGather(
+    PatternRewriter &rewriter, stablehlo::WhileOp whileOp,
+    ArrayRef<SliceInfo<stablehlo::DynamicSliceOp>> slices,
+    stablehlo::DynamicSliceOp dsOp, WhileLoopInfo info) {
+  // Pattern: x[..., y[idx], z[idx], ...] where idx is an affine function of
+  // the loop induction variable. We need to:
+  // 1. Hoist the computation of y[idx], z[idx], etc. for all loop iterations
+  // 2. Create a gather operation from x using those indices
+  // 3. Replace dsOp uses with a dynamic slice into the gather result
+
+  // Find all start indices that are dependent on the loop and come from inner
+  // dynamic slices (through possible reshape)
+  SmallVector<int64_t> dependentDims;
+  SmallVector<Value> innerSliceOperands;
+  SmallVector<SliceInfo<stablehlo::DynamicSliceOp>> innerSliceInfos;
+
+  for (auto [i, startIndex] : llvm::enumerate(dsOp.getStartIndices())) {
+    // Use traverseOperandsForHoisting to classify this operand
+    SmallVector<BatchLiftingMode> modes;
+    SmallVector<Value> operands;
+    SmallVector<SmallVector<int64_t>> dims;
+    SmallVector<int64_t> hoisted;
+    SmallVector<SliceInfo<stablehlo::DynamicSliceOp>> mapped;
+    DenseMap<Value, SmallVector<Operation *>> hoistMap;
+
+    SmallVector<Value> singleOperand = {startIndex};
+    if (!traverseOperandsForHoisting(singleOperand, whileOp, slices, info,
+                                     modes, operands, dims, hoisted, mapped,
+                                     hoistMap)) {
+      return false;
+    }
+
+    if (modes[0] == BatchLiftingMode::DYNAMIC_SLICE) {
+      dependentDims.push_back(i);
+      innerSliceOperands.push_back(operands[0]);
+      innerSliceInfos.push_back(mapped[0]);
+    }
+  }
+
+  if (dependentDims.empty()) {
+    return false;
+  }
+
+  // Get outer operand - it must be constant across iterations (loop invariant)
+  Value outerOperand;
+  Value dsOperand = dsOp.getOperand();
+  SmallVector<Operation *> canBeHoisted;
+  if (!info.isConstantAcrossIterations(dsOperand, outerOperand, canBeHoisted,
+                                       true)) {
+    return false;
+  }
+  if (!outerOperand) {
+    // The operand is defined inside the loop but is hoistable - hoist it
+    DenseMap<Value, SmallVector<Operation *>> hoistMap;
+    hoistMap[dsOperand] = canBeHoisted;
+    DenseMap<Value, Value> hoistedValues;
+    hoistChainOfOps(hoistMap, rewriter, whileOp, info, hoistedValues);
+    outerOperand = hoistedValues[dsOperand];
+  }
+
+  // Verify all non-dependent start indices are constant across iterations
+  for (auto [i, startIndex] : llvm::enumerate(dsOp.getStartIndices())) {
+    if (llvm::is_contained(dependentDims, i)) {
+      continue;
+    }
+    if (!info.isConstantAcrossIterations(startIndex, true)) {
+      return false;
+    }
+  }
+
+  int64_t numIters = info.getConstantNumIters();
+  Location loc = dsOp.getLoc();
+
+  rewriter.setInsertionPoint(whileOp);
+
+  // Step 1: Hoist each inner slice operand and construct the gather indices
+  // We need to gather all indices for all loop iterations and concatenate them.
+  SmallVector<Value> hoistedIndicesList;
+  Type hoistedIndicesElemTy;
+
+  for (size_t idx = 0; idx < dependentDims.size(); idx++) {
+    Value hoistedIndices;
+    if (!info.hoistOperationFromLoop(
+            rewriter, innerSliceOperands[idx], innerSliceInfos[idx].sliceOp,
+            innerSliceInfos[idx].dimensions, hoistedIndices)) {
+      return false;
+    }
+
+    auto hoistedTy = cast<RankedTensorType>(hoistedIndices.getType());
+    if (idx == 0) {
+      hoistedIndicesElemTy = hoistedTy.getElementType();
+    }
+
+    // Reshape to [numIters, 1] for use as gather indices
+    SmallVector<int64_t> reshapeShape = {numIters, 1};
+    auto reshapeTy = RankedTensorType::get(reshapeShape, hoistedIndicesElemTy);
+
+    // Convert type if needed
+    if (hoistedTy.getElementType() != hoistedIndicesElemTy) {
+      hoistedIndices = stablehlo::ConvertOp::create(
+          rewriter, loc,
+          RankedTensorType::get(hoistedTy.getShape(), hoistedIndicesElemTy),
+          hoistedIndices);
+    }
+
+    Value reshaped =
+        stablehlo::ReshapeOp::create(rewriter, loc, reshapeTy, hoistedIndices);
+    hoistedIndicesList.push_back(reshaped);
+  }
+
+  // Concatenate all hoisted indices along the last dimension
+  Value gatherIndices;
+  if (hoistedIndicesList.size() == 1) {
+    gatherIndices = hoistedIndicesList[0];
+  } else {
+    // Result shape: [numIters, numDependentDims]
+    SmallVector<int64_t> concatShape = {numIters,
+                                        (int64_t)dependentDims.size()};
+    auto concatTy = RankedTensorType::get(concatShape, hoistedIndicesElemTy);
+    gatherIndices = stablehlo::ConcatenateOp::create(
+        rewriter, loc, concatTy, hoistedIndicesList, /*dimension=*/1);
+  }
+
+  // Step 2: Create the gather operation from the outer operand
+  auto outerOperandTy = cast<RankedTensorType>(outerOperand.getType());
+  auto dsSliceSizes = dsOp.getSliceSizes();
+
+  // The gather slice sizes: dependent dimensions get 1, others get original
+  SmallVector<int64_t> gatherSliceSizes;
+  for (size_t i = 0; i < dsSliceSizes.size(); i++) {
+    if (llvm::is_contained(dependentDims, i)) {
+      gatherSliceSizes.push_back(1);
+    } else {
+      gatherSliceSizes.push_back(dsSliceSizes[i]);
+    }
+  }
+
+  // offsetDims: output dimensions corresponding to non-collapsed slice dims
+  // Start at 1 since batch dim is at position 0, then consecutive for each
+  // non-collapsed dimension
+  SmallVector<int64_t> offsetDims;
+  int64_t offsetDimIdx = 1; // Start after the batch dimension
+  for (size_t i = 0; i < outerOperandTy.getRank(); i++) {
+    if (!llvm::is_contained(dependentDims, i)) {
+      offsetDims.push_back(offsetDimIdx);
+      offsetDimIdx++;
+    }
+  }
+
+  // collapsedSliceDims: the dimensions we're indexing into
+  SmallVector<int64_t> collapsedSliceDims = llvm::to_vector(dependentDims);
+
+  // startIndexMap: maps index vector dimensions to operand dimensions
+  SmallVector<int64_t> startIndexMap = llvm::to_vector(dependentDims);
+
+  // Calculate output shape: [numIters, ...sliceSizes for non-dependent dims...]
+  SmallVector<int64_t> gatherOutputShape;
+  gatherOutputShape.push_back(numIters);
+  for (size_t i = 0; i < dsSliceSizes.size(); i++) {
+    if (!llvm::is_contained(dependentDims, i)) {
+      gatherOutputShape.push_back(dsSliceSizes[i]);
+    }
+  }
+
+  auto gatherResultTy =
+      RankedTensorType::get(gatherOutputShape, outerOperandTy.getElementType());
+
+  auto gatherOp = stablehlo::GatherOp::create(
+      rewriter, loc, gatherResultTy, outerOperand, gatherIndices,
+      stablehlo::GatherDimensionNumbersAttr::get(
+          rewriter.getContext(),
+          /*offsetDims=*/offsetDims,
+          /*collapsedSliceDims=*/collapsedSliceDims,
+          /*operandBatchingDims=*/{},
+          /*startIndicesBatchingDims=*/{},
+          /*startIndexMap=*/startIndexMap,
+          /*indexVectorDim=*/1),
+      gatherSliceSizes);
+
+  // Step 3: Replace the dsOp with a dynamic slice into the gather result
+  // The dynamic slice will index using the loop induction variable
+  rewriter.setInsertionPointAfter(dsOp);
+
+  auto inductionVar = info.getInductionVariable();
+  auto inductionVarType = cast<RankedTensorType>(inductionVar.getType());
+
+  // Compute the index for the dynamic slice
+  Value sliceIndex;
+  if (info.isConstantStart() && info.getConstantStart() == 0) {
+    sliceIndex = inductionVar;
+  } else {
+    sliceIndex = stablehlo::SubtractOp::create(rewriter, loc, inductionVar,
+                                               info.getStart());
+  }
+  if (!info.isStepOne()) {
+    sliceIndex = stablehlo::DivOp::create(rewriter, loc, sliceIndex,
+                                          info.getStep(rewriter));
+  }
+
+  // Convert sliceIndex to the same type as gather indices if needed
+  if (inductionVarType.getElementType() != hoistedIndicesElemTy) {
+    auto newIndexTy = RankedTensorType::get({}, hoistedIndicesElemTy);
+    sliceIndex =
+        stablehlo::ConvertOp::create(rewriter, loc, newIndexTy, sliceIndex);
+  }
+
+  // Create constZero with the same type as sliceIndex (after conversion)
+  auto sliceIndexTy = cast<RankedTensorType>(sliceIndex.getType());
+  auto constZero = stablehlo::ConstantOp::create(
+      rewriter, loc, sliceIndexTy,
+      cast<ElementsAttr>(makeAttr(sliceIndexTy, 0)));
+  // Build the start indices for dynamic slice
+  SmallVector<Value> dynSliceStarts;
+  dynSliceStarts.push_back(sliceIndex);
+  for (size_t i = 0; i < dsSliceSizes.size(); i++) {
+    if (!llvm::is_contained(dependentDims, i)) {
+      dynSliceStarts.push_back(constZero);
+    }
+  }
+
+  // Build the slice sizes (1 for the batch dim, original sizes for others)
+  SmallVector<int64_t> dynSliceSizes;
+  dynSliceSizes.push_back(1);
+  for (size_t i = 0; i < dsSliceSizes.size(); i++) {
+    if (!llvm::is_contained(dependentDims, i)) {
+      dynSliceSizes.push_back(dsSliceSizes[i]);
+    }
+  }
+
+  auto dynSlice = stablehlo::DynamicSliceOp::create(
+      rewriter, loc, gatherOp.getResult(), dynSliceStarts, dynSliceSizes);
+
+  // Reshape to match the original dsOp output type
+  auto replacement =
+      stablehlo::ReshapeOp::create(rewriter, loc, dsOp.getType(), dynSlice);
+
+  rewriter.replaceOp(dsOp, replacement.getResult());
+
+  return true;
+}
+
 bool liftOperationByBatching(
     PatternRewriter &rewriter, stablehlo::WhileOp whileOp,
     ArrayRef<SliceInfo<stablehlo::DynamicSliceOp>> slices, Operation *op,

src/enzyme_ad/jax/Passes/AutoBatching.h

@@ -144,6 +144,11 @@ struct SliceToBatchElementwise : public SliceToBatchBase {
       : SliceToBatchBase(CheckElementwise, ctx, benefit) {}
 };
 
+bool raiseDynamicSliceToGather(
+    mlir::PatternRewriter &rewriter, mlir::stablehlo::WhileOp whileOp,
+    llvm::ArrayRef<SliceInfo<mlir::stablehlo::DynamicSliceOp>> slices,
+    mlir::stablehlo::DynamicSliceOp dsOp, mlir::enzyme::WhileLoopInfo info);
+
 bool liftOperationByBatching(
     mlir::PatternRewriter &rewriter, mlir::stablehlo::WhileOp whileOp,
     llvm::ArrayRef<SliceInfo<mlir::stablehlo::DynamicSliceOp>> slices,

src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp

@@ -1316,6 +1316,9 @@ struct LowerWrap
 
 stablehlo::ConcatenateOp lowerExtend(enzymexla::ExtendOp extend,
                                      PatternRewriter &rewriter, bool replace) {
+  OpBuilder::InsertionGuard guard(rewriter);
+  rewriter.setInsertionPoint(extend);
+
   auto loc = extend.getLoc();
   auto operand = extend.getOperand();
 
@@ -1376,6 +1379,9 @@ struct LowerExtend
 
 stablehlo::ConcatenateOp lowerRotate(enzymexla::RotateOp rotate,
                                      PatternRewriter &rewriter, bool replace) {
+  OpBuilder::InsertionGuard guard(rewriter);
+  rewriter.setInsertionPoint(rotate);
+
   // sl0[A:end], sl1[0:A]
   auto shard = sdy::getShardingPerValue(rotate);
   SmallVector<int64_t> strides(rotate.getType().getShape().size(), 1);
@@ -28963,6 +28969,12 @@ LogicalResult DUSDSSimplifyWithSomeUpdateOverlapHelper(
           dusStartIsConstant) {
         int64_t dsStartInt = dsStartAP.getSExtValue();
         int64_t dusStartInt = dusStartAP.getSExtValue();
+
+        if (dusStartInt + updateShape[i] <= dsStartInt) { // no overlap
+          canReplace = false;
+          break;
+        }
+
         sliceStarts[i] = dsStartInt - dusStartInt;
         continue;
       }
@@ -29039,6 +29051,7 @@ LogicalResult DUSDSSimplifyWithSomeUpdateOverlapHelper(
       }
     }
 
+    rewriter.setInsertionPoint(dsOp);
     auto newDS = stablehlo::DynamicSliceOpCreate(
         rewriter, dsOp.getLoc(), dusOp.getUpdate(), dynamicSliceStarts,
         dsOp.getSliceSizes());
@@ -29051,7 +29064,7 @@ LogicalResult DUSDSSimplifyWithSomeUpdateOverlapHelper(
 
   // simple case
   if (allOffsetsZero && updateShape == dsSliceSizes) {
-    rewriter.replaceAllUsesWith(dsOp.getResult(), dusOp.getUpdate());
+    rewriter.replaceOp(dsOp, dusOp.getUpdate());
     return success();
   }
 
@@ -29081,6 +29094,7 @@ LogicalResult DUSDSSimplifyWithSomeUpdateOverlapHelper(
     return failure();
   }
 
+  rewriter.setInsertionPoint(dsOp);
   Value result =
       stablehlo::SliceOp::create(rewriter, dsOp.getLoc(), dusOp.getUpdate(),
                                  sliceStarts, sliceLimits, sliceStrides);