feat: dynamic_gather to gather opt pass (#193)

* feat: dynamic_gather to gather opt pass * test: test conversion * feat: support non-int64 slice sizes * feat: add the pass to td and py
EnzymeAD · Dec 16, 2024 · 51687b0 · 51687b0
1 parent 30861b8
commit 51687b0
Show file tree

Hide file tree

Showing 4 changed files with 99 additions and 10 deletions.
diff --git a/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp b/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp
@@ -6602,6 +6602,54 @@ struct IfToSelect final : public OpRewritePattern<mlir::stablehlo::IfOp> {
   }
 };
 
+struct DynamicGatherOpIsNotDynamic
+    : public OpRewritePattern<stablehlo::DynamicGatherOp> {
+  using OpRewritePattern<stablehlo::DynamicGatherOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(stablehlo::DynamicGatherOp op,
+                                PatternRewriter &rewriter) const override {
+    // Check if slice sizes are constant.
+    DenseIntElementsAttr sliceSizesAttr;
+    if (!matchPattern(op.getSliceSizes(), m_Constant(&sliceSizesAttr))) {
+      return failure();
+    }
+
+    // dynamic_gather allows non-int64 slice sizes, but we need to convert them
+    // to int64 for the gather.
+    if (!sliceSizesAttr.getType().getElementType().isInteger(64)) {
+      SmallVector<APInt> sliceSizes;
+      for (auto size : sliceSizesAttr.getValues<APInt>()) {
+        sliceSizes.push_back(size);
+      }
+      auto newSliceSizesAttr = DenseElementsAttr::get(
+          RankedTensorType::get(sliceSizesAttr.getType().getShape(),
+                                rewriter.getIntegerType(64)),
+          sliceSizes);
+      sliceSizesAttr = newSliceSizesAttr.cast<DenseIntElementsAttr>();
+    }
+
+    SmallVector<int64_t> sliceSizes;
+    for (auto size : sliceSizesAttr.getValues<int64_t>()) {
+      sliceSizes.push_back(size);
+    }
+    auto sliceSizesArrayAttr =
+        DenseI64ArrayAttr::get(op.getContext(), sliceSizes);
+
+    rewriter.replaceOpWithNewOp<stablehlo::GatherOp>(
+        op, op.getType(), op.getOperand(), op.getStartIndices(),
+        stablehlo::GatherDimensionNumbersAttr::get(
+            op.getContext(), op.getDimensionNumbers().getOffsetDims(),
+            op.getDimensionNumbers().getCollapsedSliceDims(),
+            /*operandBatchingDims=*/{},
+            /*startIndicesBatchingDims=*/{},
+            op.getDimensionNumbers().getStartIndexMap(),
+            op.getDimensionNumbers().getIndexVectorDim()),
+        sliceSizesArrayAttr);
+
+    return success();
+  }
+};
+
 /// Check if a `t` is a tensor with zero extents.
 static std::optional<RankedTensorType> isZeroExtent(Type t) {
   auto type = t.dyn_cast<RankedTensorType>();
@@ -6848,16 +6896,17 @@ struct EnzymeHLOOptPass : public EnzymeHLOOptPassBase<EnzymeHLOOptPass> {
     if (no_nan || all_finite)
       patterns.add<NoNan>(context);
 
-    patterns
-        .add<CompareOpCanon, BroadcastInDimOpCanon, ConvertOpCanon,
-             DynamicBroadcastInDimOpNotActuallyDynamic,
-             ChainedDynamicBroadcastInDimCanonicalization,
-             DynamicBroadcastInDimAllDimsNonExpanding, NoopReduceOpCanon,
-             EmptyReduceOpCanon, DynamicReshapeOpCanon, GetTupleElementOpCanon,
-             RealOpCanon, ImagOpCanon, ConjComplexNegate,
-             GetDimensionSizeOpCanon, GatherOpCanon, ReshapeOpCanon,
-             MergeConsecutiveReshapes, TransposeIsReshape, IfInline, IfToSelect,
-             ZeroExtentTensorCanon, ReorderElementwiseAndShapeOp>(context);
+    patterns.add<CompareOpCanon, BroadcastInDimOpCanon, ConvertOpCanon,
+                 DynamicBroadcastInDimOpNotActuallyDynamic,
+                 ChainedDynamicBroadcastInDimCanonicalization,
+                 DynamicBroadcastInDimAllDimsNonExpanding, NoopReduceOpCanon,
+                 EmptyReduceOpCanon, DynamicReshapeOpCanon,
+                 GetTupleElementOpCanon, RealOpCanon, ImagOpCanon,
+                 ConjComplexNegate, GetDimensionSizeOpCanon, GatherOpCanon,
+                 ReshapeOpCanon, MergeConsecutiveReshapes, TransposeIsReshape,
+                 IfInline, IfToSelect, ZeroExtentTensorCanon,
+                 ReorderElementwiseAndShapeOp, DynamicGatherOpIsNotDynamic>(
+        context);
     patterns.add<SelectOpCanon>(max_constant_expansion, context,
                                 PatternBenefit(65000));
     patterns.add<ConcatenateOpCanon>(max_constant_expansion, context,

diff --git a/src/enzyme_ad/jax/TransformOps/TransformOps.td b/src/enzyme_ad/jax/TransformOps/TransformOps.td
@@ -659,6 +659,10 @@ def ApplyConstPropThroughBarrierPatterns : EnzymeHLOPatternOp<
     "const_prop_through_barrier"> {
   let patterns = ["ConstPropThroughBarrier"];
 }
+def DynamicGatherOpIsNotDynamic : EnzymeHLOPatternOp<
+    "dynamic_gather_op_is_not_dynamic"> {
+  let patterns = ["DynamicGatherOpIsNotDynamic"];
+}
 
 // TODO: better naming for parameters requires a static interface for
 // constructing them in search.

diff --git a/src/enzyme_ad/jax/primitives.py b/src/enzyme_ad/jax/primitives.py
@@ -110,6 +110,7 @@ def hlo_opts():
 transpose_is_reshape<16>;
 zero_extent_tensor_canon<16>;
 reorder_elementwise_and_shape_op<16>;
+dynamic_gather_op_is_not_dynamic<16>;
 
 cse_broadcast_in_dim<16>;
 cse_slice<16>;

diff --git a/test/lit_tests/dynamicgathertogather.mlir b/test/lit_tests/dynamicgathertogather.mlir
@@ -0,0 +1,35 @@
+// RUN: enzymexlamlir-opt --enzyme-hlo-opt %s | FileCheck %s
+
+module {
+  func.func @main(%arg0: tensor<4x4xf64>) -> tensor<2xf64> {
+    %c = stablehlo.constant dense<1> : tensor<2xi64>
+    %c_0 = stablehlo.constant dense<[[1, -1], [2, 0]]> : tensor<2x2xi64>
+    %0 = stablehlo.transpose %arg0, dims = [1, 0] : (tensor<4x4xf64>) -> tensor<4x4xf64>
+    %1 = "stablehlo.dynamic_gather"(%0, %c_0, %c) <{dimension_numbers = #stablehlo.gather<collapsed_slice_dims = [0, 1], start_index_map = [0, 1], index_vector_dim = 1>}> : (tensor<4x4xf64>, tensor<2x2xi64>, tensor<2xi64>) -> tensor<2xf64>
+    return %1 : tensor<2xf64>
+  }
+}
+
+// CHECK: func.func @main(%arg0: tensor<4x4xf64>) -> tensor<2xf64> {
+// CHECK-NEXT:    %c = stablehlo.constant dense<{{\[\[}}1, -1{{\]}}, {{\[}}2, 0{{\]\]}}> : tensor<2x2xi64>
+// CHECK-NEXT:    %0 = stablehlo.transpose %arg0, dims = [1, 0] : (tensor<4x4xf64>) -> tensor<4x4xf64>
+// CHECK-NEXT:    %1 = "stablehlo.gather"(%0, %c) <{dimension_numbers = #stablehlo.gather<collapsed_slice_dims = [0, 1], start_index_map = [0, 1], index_vector_dim = 1>, slice_sizes = array<i64: 1, 1>}> : (tensor<4x4xf64>, tensor<2x2xi64>) -> tensor<2xf64>
+// CHECK-NEXT:    return %1 : tensor<2xf64>
+// CHECK-NEXT:  }
+
+module {
+  func.func @main(%arg0: tensor<4x4xf64>) -> tensor<2xf64> {
+    %c = stablehlo.constant dense<1> : tensor<2xi32>
+    %c_0 = stablehlo.constant dense<[[1, -1], [2, 0]]> : tensor<2x2xi64>
+    %0 = stablehlo.transpose %arg0, dims = [1, 0] : (tensor<4x4xf64>) -> tensor<4x4xf64>
+    %1 = "stablehlo.dynamic_gather"(%0, %c_0, %c) <{dimension_numbers = #stablehlo.gather<collapsed_slice_dims = [0, 1], start_index_map = [0, 1], index_vector_dim = 1>}> : (tensor<4x4xf64>, tensor<2x2xi64>, tensor<2xi32>) -> tensor<2xf64>
+    return %1 : tensor<2xf64>
+  }
+}
+
+// CHECK: func.func @main(%arg0: tensor<4x4xf64>) -> tensor<2xf64> {
+// CHECK-NEXT:    %c = stablehlo.constant dense<{{\[\[}}1, -1{{\]}}, {{\[}}2, 0{{\]\]}}> : tensor<2x2xi64>
+// CHECK-NEXT:    %0 = stablehlo.transpose %arg0, dims = [1, 0] : (tensor<4x4xf64>) -> tensor<4x4xf64>
+// CHECK-NEXT:    %1 = "stablehlo.gather"(%0, %c) <{dimension_numbers = #stablehlo.gather<collapsed_slice_dims = [0, 1], start_index_map = [0, 1], index_vector_dim = 1>, slice_sizes = array<i64: 1, 1>}> : (tensor<4x4xf64>, tensor<2x2xi64>) -> tensor<2xf64>
+// CHECK-NEXT:    return %1 : tensor<2xf64>
+// CHECK-NEXT:  }