Simplify GatherOp when its inputs are dimensions

src/Transform/ONNX/SimplifyShapeRelatedOps.cpp

@@ -60,6 +60,7 @@ Now, it's straighforward to update the output shape of Reshape from
 
 #include "src/Dialect/ONNX/DialectBuilder.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
+#include "src/Dialect/ONNX/ONNXOpsHelper.hpp"
 #include "src/Dialect/ONNX/ShapeInference/ONNXShapeHelper.hpp"
 #include "src/Pass/Passes.hpp"
 #include "src/Support/TypeUtilities.hpp"
@@ -255,6 +256,61 @@ class PassThroughCastPattern : public OpRewritePattern<ONNXCastOp> {
   }
 };
 
+/// Simplify ONNXGatherOp into ONNXConcatOp.
+class PassThroughGatherPattern : public OpRewritePattern<ONNXGatherOp> {
+public:
+  using OpRewritePattern<ONNXGatherOp>::OpRewritePattern;
+  LogicalResult matchAndRewrite(
+      ONNXGatherOp gatherOp, PatternRewriter &rewriter) const override {
+    // Get basic op info.
+    Location loc = gatherOp.getLoc();
+    Value input = gatherOp.data();
+    Value indices = gatherOp.indices();
+    int64_t axis = gatherOp.axis();
+
+    // Match
+    // Gather on axis 0.
+    if (axis != 0)
+      return failure();
+    // Input is defined by Concat of dims, so it has rank of 1.
+    if (!areDimsFromConcat(input))
+      return failure();
+
+    // Indices are constants.
+    if (!definedBy<ONNXConstantOp>(indices))
+      return failure();
+    DenseElementsAttr indicesAttr =
+        getDenseElementAttributeFromONNXValue(indices);
+    if (!indicesAttr)
+      return failure();
+
+    // Rewrite
+    MultiDialectBuilder<OnnxBuilder> create(rewriter, loc);
+    int64_t inputRank = getRank(input.getType());
+
+    // Compute integer indices.
+    SmallVector<int64_t, 4> indicesI64;
+    for (auto element : indicesAttr.getValues<IntegerAttr>()) {
+      int64_t axis = element.getInt();
+      axis = (axis < 0) ? (axis + inputRank) : axis;
+      indicesI64.emplace_back(axis);
+    }
+
+    // Replace GatherOp by ConcatOp of specific dimensions.
+    SmallVector<Value, 4> dims;
+    getDims(input, dims);
+    SmallVector<Value> castedDims;
+    SmallVector<Value> gatherDims;
+    for (int64_t i : indicesI64)
+      gatherDims.emplace_back(dims[i]);
+
+    Value replacedValue = emitConcatOpForDims(create, gatherDims);
+
+    rewriter.replaceOp(gatherOp, replacedValue);
+    return success();
+  }
+};
+
 /// Simplify ONNXSliceOp into ONNXConcatOp.
 class PassThroughSlicePattern : public OpRewritePattern<ONNXSliceOp> {
 public:
@@ -472,6 +528,7 @@ void SimplifyShapeRelatedOpsPass::topDownShapeSimplification(
   // Pass the dimensions through operations of interest.
   patterns.insert<onnx_mlir::PassThroughCastPattern>(context);
   patterns.insert<onnx_mlir::PassThroughConcatPattern>(context);
+  patterns.insert<onnx_mlir::PassThroughGatherPattern>(context);
   patterns.insert<onnx_mlir::PassThroughSlicePattern>(context);
 
   // Update Reshape's output shape using inferred dimensions.

test/mlir/onnx/onnx_simplify_shape_related_ops.mlir

@@ -83,6 +83,27 @@ func.func @test_pass_dims_through_concat(%arg0: tensor<?x256xi64>) -> (tensor<4x
 
 // -----
 
+func.func @test_pass_dims_through_cast(%arg0: tensor<?x?x200xf32>) -> tensor<2xi64> {
+  %0 = "onnx.Constant"() {value = dense<[0, 1]> : tensor<2xi64>} : () -> tensor<2xi64>
+  %1 = "onnx.Dim"(%arg0) {axis = 0 : si64} : (tensor<?x?x200xf32>) -> tensor<1xi64>
+  %2 = "onnx.Dim"(%arg0) {axis = 1 : si64} : (tensor<?x?x200xf32>) -> tensor<1xi64>
+  %3 = "onnx.Constant"() {value = dense<200> : tensor<1xi64>} : () -> tensor<1xi64>
+  %4 = "onnx.Concat"(%1, %2, %3) {axis = 0 : si64} : (tensor<1xi64>, tensor<1xi64>, tensor<1xi64>) -> tensor<3xi64>
+  %5 = "onnx.Gather"(%4, %0) {axis = 0 : si64} : (tensor<3xi64>, tensor<2xi64>) -> tensor<2xi64>
+  return %5 : tensor<2xi64>
+
+// mlir2FileCheck.py
+// CHECK-LABEL:  func.func @test_pass_dims_through_cast
+// CHECK-SAME:   ([[PARAM_0_:%.+]]: tensor<?x?x200xf32>) -> tensor<2xi64> {
+// CHECK-DAG:       [[VAR_0_:%.+]] = "onnx.Dim"([[PARAM_0_]]) {axis = 0 : si64} : (tensor<?x?x200xf32>) -> tensor<1xi64>
+// CHECK-DAG:       [[VAR_1_:%.+]] = "onnx.Dim"([[PARAM_0_]]) {axis = 1 : si64} : (tensor<?x?x200xf32>) -> tensor<1xi64>
+// CHECK:           [[VAR_2_:%.+]] = "onnx.Concat"([[VAR_0_]], [[VAR_1_]]) {axis = 0 : si64} : (tensor<1xi64>, tensor<1xi64>) -> tensor<2xi64>
+// CHECK:           return [[VAR_2_]] : tensor<2xi64>
+// CHECK:         }
+}
+
+// -----
+
 func.func @test_pass_dims_through_slice(%arg0: tensor<?x256xi64>) -> (tensor<1xi64>) {
   %0 = "onnx.Dim"(%arg0) {axis = 0 : si64} : (tensor<?x256xi64>) -> tensor<1xi64>
   %1 = "onnx.Constant"() {value = dense<256> : tensor<1xi64>} : () -> tensor<1xi64>