feat: gather elementwise simplify (#1038)

Author: avik-pal

src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp

@@ -19683,6 +19683,40 @@ struct UnaryElementwiseScatterSimplify final
   }
 };
 
+struct GatherElementwise
+    : public CheckedOpRewritePattern<stablehlo::GatherOp, GatherElementwise> {
+  using CheckedOpRewritePattern<stablehlo::GatherOp,
+                                GatherElementwise>::CheckedOpRewritePattern;
+
+  LogicalResult matchAndRewriteImpl(stablehlo::GatherOp op,
+                                    PatternRewriter &rewriter) const {
+    auto gatherInput = op.getOperand();
+    auto defOp = gatherInput.getDefiningOp();
+    if (!defOp || !defOp->hasTrait<mlir::OpTrait::Elementwise>())
+      return rewriter.notifyMatchFailure(op,
+                                         "GatherOp with non-elementwise input");
+
+    if (!isOnlyUsedInOperation(defOp, op))
+      return failure();
+
+    SmallVector<Value> newElementwiseInputs;
+    for (auto input : defOp->getOperands()) {
+      auto neeGatherOp = rewriter.create<stablehlo::GatherOp>(
+          op.getLoc(), input, op.getStartIndices(),
+          op.getDimensionNumbersAttr(), op.getSliceSizesAttr(),
+          op.getIndicesAreSortedAttr());
+      newElementwiseInputs.push_back(neeGatherOp->getResult(0));
+    }
+
+    auto newElemOp = rewriter.create(
+        op.getLoc(), defOp->getName().getIdentifier(),
+        ValueRange(newElementwiseInputs), TypeRange{op.getResult().getType()},
+        defOp->getAttrs(), {}, {});
+    rewriter.replaceOp(op, newElemOp->getResult(0));
+    return success();
+  }
+};
+
 ///////////////  End Imported from stablehlo
 
 // clang-format off
@@ -20205,7 +20239,8 @@ struct EnzymeHLOOptPass
         ConjComplexSimplify,
         SplitConvolutionIntoReverseConvolution,
         ScatterMultiplySimplify,
-        UnaryElementwiseScatterSimplify
+        UnaryElementwiseScatterSimplify,
+        GatherElementwise
       >(context);
 
     patterns.add<SumToReduceWindow<stablehlo::AddOp>,

src/enzyme_ad/jax/TransformOps/TransformOps.td

@@ -2033,3 +2033,7 @@ def ApplyUnaryElementwiseScatterSimplify : EnzymeHLOPatternOp<"unary_elementwise
     "UnaryElementwiseScatterSimplify"
   ];
 }
+
+def ApplyGatherElementwise : EnzymeHLOPatternOp<"gather_elementwise"> {
+  let patterns = ["GatherElementwise"];
+}

test/lit_tests/gather_elementwise.mlir

@@ -0,0 +1,38 @@
+// RUN: enzymexlamlir-opt --enzyme-hlo-opt %s | FileCheck %s
+
+func.func @gather_elementwise1(%arg0: tensor<32x1024xf32>, %arg1: tensor<32x1024xf32>, %arg2: tensor<5xi64>) -> tensor<32x5xf32> {
+    %c = stablehlo.constant dense<1> : tensor<1x5xi64>
+    %0 = stablehlo.multiply %arg0, %arg1 : tensor<32x1024xf32>
+    %1 = stablehlo.reshape %arg2 : (tensor<5xi64>) -> tensor<1x5xi64>
+    %2 = stablehlo.subtract %1, %c : tensor<1x5xi64>
+    %3 = "stablehlo.gather"(%0, %2) <{dimension_numbers = #stablehlo.gather<offset_dims = [0], collapsed_slice_dims = [1], start_index_map = [1]>, indices_are_sorted = false, slice_sizes = array<i64: 32, 1>}> : (tensor<32x1024xf32>, tensor<1x5xi64>) -> tensor<32x5xf32>
+    return %3 : tensor<32x5xf32>
+}
+
+// CHECK: func.func @gather_elementwise1(%arg0: tensor<32x1024xf32>, %arg1: tensor<32x1024xf32>, %arg2: tensor<5xi64>) -> tensor<32x5xf32> {
+// CHECK-NEXT:     %c = stablehlo.constant dense<1> : tensor<1x5xi64>
+// CHECK-NEXT:     %0 = stablehlo.reshape %arg2 : (tensor<5xi64>) -> tensor<1x5xi64>
+// CHECK-NEXT:     %1 = stablehlo.subtract %0, %c : tensor<1x5xi64>
+// CHECK-NEXT:     %2 = "stablehlo.gather"(%arg0, %1) <{dimension_numbers = #stablehlo.gather<offset_dims = [0], collapsed_slice_dims = [1], start_index_map = [1]>, indices_are_sorted = false, slice_sizes = array<i64: 32, 1>}> : (tensor<32x1024xf32>, tensor<1x5xi64>) -> tensor<32x5xf32>
+// CHECK-NEXT:     %3 = "stablehlo.gather"(%arg1, %1) <{dimension_numbers = #stablehlo.gather<offset_dims = [0], collapsed_slice_dims = [1], start_index_map = [1]>, indices_are_sorted = false, slice_sizes = array<i64: 32, 1>}> : (tensor<32x1024xf32>, tensor<1x5xi64>) -> tensor<32x5xf32>
+// CHECK-NEXT:     %4 = stablehlo.multiply %2, %3 : tensor<32x5xf32>
+// CHECK-NEXT:     return %4 : tensor<32x5xf32>
+// CHECK-NEXT: }
+
+func.func @gather_elementwise2(%arg0: tensor<32x1024xf32>, %arg1: tensor<5xi64>) -> tensor<32x5xf16> {
+    %c = stablehlo.constant dense<1> : tensor<5xi64>
+    %0 = stablehlo.convert %arg0 : (tensor<32x1024xf32>) -> tensor<32x1024xf16>
+    %1 = stablehlo.subtract %arg1, %c : tensor<5xi64>
+    %2 = stablehlo.reshape %1 : (tensor<5xi64>) -> tensor<1x5xi64>
+    %3 = "stablehlo.gather"(%0, %2) <{dimension_numbers = #stablehlo.gather<offset_dims = [0], collapsed_slice_dims = [1], start_index_map = [1]>, indices_are_sorted = false, slice_sizes = array<i64: 32, 1>}> : (tensor<32x1024xf16>, tensor<1x5xi64>) -> tensor<32x5xf16>
+    return %3 : tensor<32x5xf16>
+}
+
+// CHECK: func.func @gather_elementwise2(%arg0: tensor<32x1024xf32>, %arg1: tensor<5xi64>) -> tensor<32x5xf16> {
+// CHECK-NEXT:     %c = stablehlo.constant dense<1> : tensor<5xi64>
+// CHECK-NEXT:     %0 = stablehlo.subtract %arg1, %c : tensor<5xi64>
+// CHECK-NEXT:     %1 = stablehlo.reshape %0 : (tensor<5xi64>) -> tensor<1x5xi64>
+// CHECK-NEXT:     %2 = "stablehlo.gather"(%arg0, %1) <{dimension_numbers = #stablehlo.gather<offset_dims = [0], collapsed_slice_dims = [1], start_index_map = [1]>, indices_are_sorted = false, slice_sizes = array<i64: 32, 1>}> : (tensor<32x1024xf32>, tensor<1x5xi64>) -> tensor<32x5xf32>
+// CHECK-NEXT:     %3 = stablehlo.convert %2 : (tensor<32x5xf32>) -> tensor<32x5xf16>
+// CHECK-NEXT:     return %3 : tensor<32x5xf16>
+// CHECK-NEXT: }