feat: more scatter gather optimization patterns (#1033)

* feat: simplify mul of scatter

* chore: run fmt

* feat: gather constant propagation

* feat: unary elementwise scatter simplify

Author: avik-pal

src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp

@@ -19411,6 +19411,278 @@ struct SplitConvolutionIntoReverseConvolution final
   }
 };
 
+stablehlo::GatherDimensionNumbersAttr
+getGatherDims(mlir::MLIRContext *ctx,
+              stablehlo::ScatterDimensionNumbersAttr scatterDimNumbers) {
+  return stablehlo::GatherDimensionNumbersAttr::get(
+      ctx, scatterDimNumbers.getUpdateWindowDims(),
+      scatterDimNumbers.getInsertedWindowDims(),
+      scatterDimNumbers.getInputBatchingDims(),
+      scatterDimNumbers.getScatterIndicesBatchingDims(),
+      scatterDimNumbers.getScatterDimsToOperandDims(),
+      scatterDimNumbers.getIndexVectorDim());
+}
+
+bool isScatterSetindexOp(stablehlo::ScatterOp &scatterOp) {
+  auto &updateComputation = scatterOp.getUpdateComputation();
+
+  if (!updateComputation.hasOneBlock())
+    return false;
+
+  auto &block = updateComputation.front();
+  if (block.getNumArguments() != 2)
+    return false;
+
+  auto originalValue = block.getArgument(0);
+  auto updateValue = block.getArgument(1);
+
+  // The block should have exactly one operation (the return)
+  if (block.getOperations().size() != 1)
+    return false;
+
+  auto &returnOp = block.front();
+  auto stablehloReturn = dyn_cast<stablehlo::ReturnOp>(returnOp);
+  if (!stablehloReturn)
+    return false;
+
+  if (stablehloReturn.getNumOperands() != 1)
+    return false;
+
+  // The returned value should be the update value (second argument)
+  return stablehloReturn.getOperand(0) == updateValue;
+}
+
+SmallVector<int64_t> computeGatherSliceSizes(stablehlo::ScatterOp &scatterOp) {
+  auto inputType = cast<ShapedType>(scatterOp.getInputs()[0].getType());
+  auto updateType = cast<ShapedType>(scatterOp.getUpdates()[0].getType());
+  auto scatterDimNumbers = scatterOp.getScatterDimensionNumbers();
+
+  auto inputShape = inputType.getShape();
+  auto updateShape = updateType.getShape();
+
+  SmallVector<int64_t> sliceSizes(inputShape.size(), 1);
+
+  auto updateWindowDims = scatterDimNumbers.getUpdateWindowDims();
+  auto scatterIndicesBatchingDims =
+      scatterDimNumbers.getScatterIndicesBatchingDims();
+
+  // Map update window dimensions to their corresponding input dimensions
+  for (int64_t i = 0; i < updateWindowDims.size(); ++i) {
+    int64_t inputDim = updateWindowDims[i];
+
+    // Calculate the corresponding dimension in the update tensor
+    // Update tensor layout: [scatter_indices_batching_dims...,
+    // update_window_dims...]
+    int64_t updateDimIndex = scatterIndicesBatchingDims.size() + i;
+
+    if (updateDimIndex < updateShape.size()) {
+      sliceSizes[inputDim] = updateShape[updateDimIndex];
+    }
+  }
+
+  return sliceSizes;
+}
+
+struct ScatterMultiplySimplify final
+    : public CheckedOpRewritePattern<stablehlo::MulOp,
+                                     ScatterMultiplySimplify> {
+  using CheckedOpRewritePattern<
+      stablehlo::MulOp, ScatterMultiplySimplify>::CheckedOpRewritePattern;
+
+  LogicalResult matchAndRewriteImpl(stablehlo::MulOp op,
+                                    PatternRewriter &rewriter) const {
+    auto lhs = op.getLhs();
+    auto rhs = op.getRhs();
+
+    stablehlo::ScatterOp scatterOp;
+    mlir::Value otherValue;
+
+    auto lhsScatterOp = lhs.getDefiningOp<stablehlo::ScatterOp>();
+    auto rhsScatterOp = rhs.getDefiningOp<stablehlo::ScatterOp>();
+    if (!lhsScatterOp && !rhsScatterOp) {
+      return failure();
+    } else {
+      if (lhsScatterOp) {
+        scatterOp = lhsScatterOp;
+        otherValue = rhs;
+      } else {
+        scatterOp = rhsScatterOp;
+        otherValue = lhs;
+      }
+    }
+
+    if (scatterOp.getInputs().size() != 1)
+      return rewriter.notifyMatchFailure(
+          op, "ScatterOp with more than one input not supported");
+
+    auto input = scatterOp.getInputs()[0];
+    if (!matchPattern(input, m_AnyZeroFloat()) &&
+        !matchPattern(input, m_Zero()))
+      return rewriter.notifyMatchFailure(op, "ScatterOp with non-zero input");
+
+    if (!scatterOp.getResult(0).hasOneUse())
+      return rewriter.notifyMatchFailure(op, "ScatterOp with multiple uses");
+
+    if (!isScatterSetindexOp(scatterOp))
+      return rewriter.notifyMatchFailure(op, "ScatterOp with non-setindex");
+
+    auto scatterDimNumbers = scatterOp.getScatterDimensionNumbers();
+
+    SmallVector<int64_t> sliceSizes = computeGatherSliceSizes(scatterOp);
+
+    auto gatheredValues = rewriter.create<stablehlo::GatherOp>(
+        op.getLoc(), otherValue, scatterOp.getScatterIndices(),
+        getGatherDims(rewriter.getContext(), scatterDimNumbers),
+        rewriter.getDenseI64ArrayAttr(sliceSizes),
+        scatterOp.getIndicesAreSortedAttr());
+
+    auto newUpdates = rewriter.create<stablehlo::MulOp>(
+        op.getLoc(), gatheredValues, scatterOp.getUpdates()[0]);
+
+    auto newScatterOp = rewriter.create<stablehlo::ScatterOp>(
+        op.getLoc(), scatterOp.getResultTypes(), scatterOp.getInputs(),
+        scatterOp.getScatterIndices(), ValueRange(newUpdates),
+        scatterOp.getScatterDimensionNumbersAttr(),
+        scatterOp.getIndicesAreSortedAttr(), scatterOp.getUniqueIndicesAttr());
+    newScatterOp.getUpdateComputation().takeBody(
+        scatterOp.getUpdateComputation());
+    rewriter.replaceOp(op, newScatterOp);
+
+    return success();
+  }
+};
+
+struct GatherConstProp final
+    : public CheckedOpRewritePattern<stablehlo::GatherOp, GatherConstProp> {
+  using CheckedOpRewritePattern<stablehlo::GatherOp,
+                                GatherConstProp>::CheckedOpRewritePattern;
+
+  LogicalResult matchAndRewriteImpl(stablehlo::GatherOp op,
+                                    PatternRewriter &rewriter) const {
+    DenseElementsAttr operandAttr;
+    if (!matchPattern(op.getOperand(), m_Constant(&operandAttr)))
+      return rewriter.notifyMatchFailure(op,
+                                         "GatherOp with non-constant input");
+
+    if (operandAttr.isSplat()) {
+      // In this case the indices don't matter and we can construct a new
+      // splatted result
+      rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
+          op, op.getType(), operandAttr.resizeSplat(op.getType()));
+      return success();
+    }
+
+    DenseElementsAttr startIndicesAttr;
+    if (!matchPattern(op.getStartIndices(), m_Constant(&startIndicesAttr))) {
+      return rewriter.notifyMatchFailure(
+          op, "GatherOp with non-constant start indices and unsplatted input");
+    }
+
+    stablehlo::Tensor operandTensor = stablehlo::constantOp(operandAttr);
+    stablehlo::Tensor startIndicesTensor =
+        stablehlo::constantOp(startIndicesAttr);
+    auto gatherDims = op.getDimensionNumbers();
+
+    auto sliceSizes = op.getSliceSizes();
+    auto elementType = rewriter.getIntegerType(64);
+    auto attrType =
+        RankedTensorType::get({(int64_t)sliceSizes.size()}, elementType);
+    auto sliceSizesAttr = DenseElementsAttr::get(attrType, sliceSizes);
+
+    auto result = stablehlo::gatherOp(
+        operandTensor, startIndicesTensor,
+        stablehlo::Axes(gatherDims.getOffsetDims()),
+        stablehlo::Axes(gatherDims.getCollapsedSliceDims()),
+        stablehlo::Axes(gatherDims.getOperandBatchingDims()),
+        stablehlo::Axes(gatherDims.getStartIndicesBatchingDims()),
+        stablehlo::Axes(gatherDims.getStartIndexMap()),
+        stablehlo::Axis(gatherDims.getIndexVectorDim()),
+        stablehlo::makeSizes(stablehlo::constantOp(sliceSizesAttr)),
+        op.getIndicesAreSorted(), op.getType());
+    rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(op, op.getType(),
+                                                       fromTensor(result));
+    return success();
+  }
+};
+
+struct UnaryElementwiseScatterSimplify final
+    : public CheckedOpTraitRewritePattern<OpTrait::Elementwise,
+                                          UnaryElementwiseScatterSimplify> {
+  using CheckedOpTraitRewritePattern<
+      OpTrait::Elementwise,
+      UnaryElementwiseScatterSimplify>::CheckedOpTraitRewritePattern;
+
+  LogicalResult matchAndRewriteImpl(Operation *op,
+                                    PatternRewriter &rewriter) const {
+    if (op->getNumOperands() != 1)
+      return rewriter.notifyMatchFailure(op, "not a unary elementwise op");
+
+    auto input = op->getOperand(0);
+    auto scatterOp = input.getDefiningOp<stablehlo::ScatterOp>();
+    if (!scatterOp)
+      return rewriter.notifyMatchFailure(op, "not a scatter op");
+
+    if (scatterOp.getInputs().size() != 1)
+      return rewriter.notifyMatchFailure(
+          op, "ScatterOp with more than one input not supported");
+
+    if (!scatterOp.getResult(0).hasOneUse())
+      return rewriter.notifyMatchFailure(op, "ScatterOp with multiple uses");
+
+    if (!isScatterSetindexOp(scatterOp))
+      return rewriter.notifyMatchFailure(op, "ScatterOp with non-setindex");
+
+    auto scatterInput = scatterOp.getInputs()[0];
+    DenseElementsAttr scatterInputAttr;
+    // In this case, we are will definitely increase the compute cost
+    if (!matchPattern(scatterInput, m_Constant(&scatterInputAttr)))
+      return rewriter.notifyMatchFailure(op,
+                                         "ScatterOp with non-constant input");
+
+    auto elemType =
+        cast<RankedTensorType>(op->getResult(0).getType()).getElementType();
+
+    // TODO: support convert op. we need to rewrite the update computation to
+    // take the converted element type
+    if (isa<stablehlo::ConvertOp>(op))
+      return rewriter.notifyMatchFailure(op,
+                                         "ConvertOp not supported for now.");
+
+    // should get constant propagated
+    auto scatterInputElem = rewriter.create(
+        op->getLoc(), op->getName().getIdentifier(), ValueRange(scatterInput),
+        TypeRange{RankedTensorType::get(
+            cast<RankedTensorType>(scatterInput.getType()).getShape(),
+            elemType)},
+        op->getAttrs(), {}, {});
+
+    auto scatterUpdatesElem = rewriter.create(
+        op->getLoc(), op->getName().getIdentifier(),
+        ValueRange(scatterOp.getUpdates()),
+        TypeRange{RankedTensorType::get(
+            cast<RankedTensorType>(scatterOp.getUpdates()[0].getType())
+                .getShape(),
+            elemType)},
+        op->getAttrs(), {}, {});
+
+    auto resultType = RankedTensorType::get(
+        cast<RankedTensorType>(scatterOp.getResultTypes()[0]).getShape(),
+        elemType);
+
+    auto newScatterOp = rewriter.create<stablehlo::ScatterOp>(
+        op->getLoc(), TypeRange(resultType),
+        ValueRange(scatterInputElem->getResult(0)),
+        scatterOp.getScatterIndices(),
+        ValueRange(scatterUpdatesElem->getResult(0)),
+        scatterOp.getScatterDimensionNumbersAttr(),
+        scatterOp.getIndicesAreSortedAttr(), scatterOp.getUniqueIndicesAttr());
+    newScatterOp.getUpdateComputation().takeBody(
+        scatterOp.getUpdateComputation());
+    rewriter.replaceOp(op, newScatterOp->getResult(0));
+    return success();
+  }
+};
+
 ///////////////  End Imported from stablehlo
 
 // clang-format off
@@ -19697,6 +19969,8 @@ struct EnzymeHLOOptPass
                  BinaryConstProp<stablehlo::SubtractOp, stablehlo::subtractOp>,
                  BinaryConstProp<stablehlo::XorOp, stablehlo::xorOp>>(context);
 
+    patterns.add<GatherConstProp>(context);
+
     patterns.add<BinaryOpTransposeSimplify<stablehlo::AddOp>,
                  BinaryOpTransposeSimplify<stablehlo::SubtractOp>,
                  BinaryOpTransposeSimplify<stablehlo::MulOp>,
@@ -19929,7 +20203,9 @@ struct EnzymeHLOOptPass
         InvolutionSimplify<chlo::ConjOp>,
         RealConjSimplify,
         ConjComplexSimplify,
-        SplitConvolutionIntoReverseConvolution
+        SplitConvolutionIntoReverseConvolution,
+        ScatterMultiplySimplify,
+        UnaryElementwiseScatterSimplify
       >(context);
 
     patterns.add<SumToReduceWindow<stablehlo::AddOp>,

src/enzyme_ad/jax/TransformOps/TransformOps.td

@@ -2015,3 +2015,21 @@ def ApplySplitConvolutionIntoReverseConvolution : EnzymeHLOPatternOp<
     "SplitConvolutionIntoReverseConvolution"
   ];
 }
+
+def ApplyScatterMultiplySimplify : EnzymeHLOPatternOp<"scatter_multiply_simplify"> {
+  let patterns = [
+    "ScatterMultiplySimplify"
+  ];
+}
+
+def ApplyGatherConstProp : EnzymeHLOPatternOp<"gather_const_prop"> {
+  let patterns = [
+    "GatherConstProp"
+  ];
+}
+
+def ApplyUnaryElementwiseScatterSimplify : EnzymeHLOPatternOp<"unary_elementwise_scatter_simplify"> {
+  let patterns = [
+    "UnaryElementwiseScatterSimplify"
+  ];
+}

test/lit_tests/gatherconstprop.mlir

@@ -0,0 +1,51 @@
+// RUN: enzymexlamlir-opt %s --enzyme-hlo-opt | FileCheck %s
+
+func.func @gather_constprop1(%arg0: tensor<4xi64>, %arg1: tensor<6xi64>, %arg2: tensor<6x4xf64>) -> tensor<6x4xf64> {
+    %c = stablehlo.constant dense<1> : tensor<24x2xi64>
+    %cst = stablehlo.constant dense<0.000000e+00> : tensor<1024x1024xf64>
+    %0 = stablehlo.concatenate %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, dim = 0 : (tensor<4xi64>, tensor<4xi64>, tensor<4xi64>, tensor<4xi64>, tensor<4xi64>, tensor<4xi64>) -> tensor<24xi64>
+    %1 = stablehlo.broadcast_in_dim %arg1, dims = [0] : (tensor<6xi64>) -> tensor<6x4xi64>
+    %2 = stablehlo.reshape %0 : (tensor<24xi64>) -> tensor<24x1xi64>
+    %3 = stablehlo.reshape %1 : (tensor<6x4xi64>) -> tensor<24x1xi64>
+    %4 = stablehlo.concatenate %2, %3, dim = 1 : (tensor<24x1xi64>, tensor<24x1xi64>) -> tensor<24x2xi64>
+    %5 = stablehlo.subtract %4, %c : tensor<24x2xi64>
+    %6 = "stablehlo.gather"(%cst, %5) <{dimension_numbers = #stablehlo.gather<collapsed_slice_dims = [0, 1], start_index_map = [0, 1], index_vector_dim = 1>, indices_are_sorted = false, slice_sizes = array<i64: 1, 1>}> : (tensor<1024x1024xf64>, tensor<24x2xi64>) -> tensor<24xf64>
+    %7 = stablehlo.reshape %6 : (tensor<24xf64>) -> tensor<6x4xf64>
+    %8 = stablehlo.multiply %7, %arg2 : tensor<6x4xf64>
+    return %8 : tensor<6x4xf64>
+}
+
+// CHECK: func.func @gather_constprop1(%arg0: tensor<4xi64>, %arg1: tensor<6xi64>, %arg2: tensor<6x4xf64>) -> tensor<6x4xf64> {
+// CHECK-NEXT:     %cst = stablehlo.constant dense<0.000000e+00> : tensor<6x4xf64>
+// CHECK-NEXT:     %0 = stablehlo.multiply %cst, %arg2 : tensor<6x4xf64>
+// CHECK-NEXT:     return %0 : tensor<6x4xf64>
+// CHECK-NEXT: }
+
+func.func @gather_constprop2(%arg0: tensor<6x4xf64>) -> tensor<6x4xf64> {
+    %cst = stablehlo.constant dense<1.000000e+00> : tensor<1024x1024xf64>
+    %c = stablehlo.constant dense<[[0, 0], [3, 0], [5, 0], [11, 0], [0, 31], [3, 31], [5, 31], [11, 31], [0, 32], [3, 32], [5, 32], [11, 32], [0, 34], [3, 34], [5, 34], [11, 34], [0, 11], [3, 11], [5, 11], [11, 11], [0, 110], [3, 110], [5, 110], [11, 110]]> : tensor<24x2xi64>
+    %0 = "stablehlo.gather"(%cst, %c) <{dimension_numbers = #stablehlo.gather<collapsed_slice_dims = [0, 1], start_index_map = [0, 1], index_vector_dim = 1>, indices_are_sorted = false, slice_sizes = array<i64: 1, 1>}> : (tensor<1024x1024xf64>, tensor<24x2xi64>) -> tensor<24xf64>
+    %1 = stablehlo.reshape %0 : (tensor<24xf64>) -> tensor<6x4xf64>
+    return %1 : tensor<6x4xf64>
+}
+
+// CHECK: func.func @gather_constprop2(%arg0: tensor<6x4xf64>) -> tensor<6x4xf64> {
+// CHECK-NEXT:     %cst = stablehlo.constant dense<1.000000e+00> : tensor<6x4xf64>
+// CHECK-NEXT:     return %cst : tensor<6x4xf64>
+// CHECK-NEXT: }
+
+func.func @gather_constprop3(%arg0: tensor<6x4xf64>) -> tensor<6x4xf64> {
+    %cst = stablehlo.constant dense<"0x0000803F000050410000C841000014420000444200007442000092420000AA420000C2420000DA420000F2420000054300000040000060410000D041000018420000484200007842000094420000AC420000C4420000DC420000F4420000064300004040000070410000D84100001C4200004C4200007C42000096420000AE420000C6420000DE420000F6420000074300008040000080410000E041000020420000504200008042000098420000B0420000C8420000E0420000F842000008430000A040000088410000E84100002442000054420000824200009A420000B2420000CA420000E2420000FA42000009430000C040000090410000F04100002842000058420000844200009C420000B4420000CC420000E4420000FC4200000A430000E040000098410000F84100002C4200005C420000864200009E420000B6420000CE420000E6420000FE4200000B43000000410000A041000000420000304200006042000088420000A0420000B8420000D0420000E8420000004300000C43000010410000A84100000442000034420000644200008A420000A2420000BA420000D2420000EA420000014300000D43000020410000B04100000842000038420000684200008C420000A4420000BC420000D4420000EC420000024300000E43000030410000B84100000C4200003C4200006C4200008E420000A6420000BE420000D6420000EE420000034300000F43000040410000C041000010420000404200007042000090420000A8420000C0420000D8420000F0420000044300001043"> : tensor<12x12xf32>
+    %c = stablehlo.constant dense<[[0, 0], [3, 0], [5, 0], [11, 0], [0, 1], [3, 1], [5, 1], [11, 1], [0, 2], [3, 2], [5, 2], [11, 2], [0, 4], [3, 4], [5, 4], [11, 4], [0, 11], [3, 11], [5, 11], [11, 11], [0, 10], [3, 10], [5, 10], [11, 10]]> : tensor<24x2xi64>
+    %0 = "stablehlo.gather"(%cst, %c) <{dimension_numbers = #stablehlo.gather<collapsed_slice_dims = [0, 1], start_index_map = [0, 1], index_vector_dim = 1>, indices_are_sorted = false, slice_sizes = array<i64: 1, 1>}> : (tensor<12x12xf32>, tensor<24x2xi64>) -> tensor<24xf32>
+    %1 = stablehlo.convert %0 : (tensor<24xf32>) -> tensor<24xf64>
+    %2 = stablehlo.reshape %1 : (tensor<24xf64>) -> tensor<6x4xf64>
+    %3 = stablehlo.multiply %2, %arg0 : tensor<6x4xf64>
+    return %3 : tensor<6x4xf64>
+}
+
+// CHECK: func.func @gather_constprop3(%arg0: tensor<6x4xf64>) -> tensor<6x4xf64> {
+// CHECK-NEXT{LITERAL}:     %cst = stablehlo.constant dense<[[1.000000e+00, 4.000000e+00, 6.000000e+00, 1.200000e+01], [1.300000e+01, 1.600000e+01, 1.800000e+01, 2.400000e+01], [2.500000e+01, 2.800000e+01, 3.000000e+01, 3.600000e+01], [4.900000e+01, 5.200000e+01, 5.400000e+01, 6.000000e+01], [1.330000e+02, 1.360000e+02, 1.380000e+02, 1.440000e+02], [1.210000e+02, 1.240000e+02, 1.260000e+02, 1.320000e+02]]> : tensor<6x4xf64>
+// CHECK-NEXT:     %0 = stablehlo.multiply %cst, %arg0 : tensor<6x4xf64>
+// CHECK-NEXT:     return %0 : tensor<6x4xf64>
+// CHECK-NEXT: }