convert negative indices correctly

Author: p-lanza

src/Dialect/ONNX/Transforms/SimplifyShapeRelatedOps.cpp

@@ -247,14 +247,16 @@ class PassThroughGatherPattern : public OpRewritePattern<ONNXGatherOp> {
 
     // Rewrite
     MultiDialectBuilder<OnnxBuilder> create(rewriter, loc);
-    int64_t inputRank = getRank(input.getType());
+    ShapedType inputType = llvm::dyn_cast<ShapedType>(input.getType());
+    if (!inputType || !inputType.hasStaticShape())
+      return failure();
 
     // 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);
+      int64_t index = element.getInt();
+      index = (index < 0) ? (index + inputType.getShape()[axis]) : index;
+      indicesI64.emplace_back(index);
     }
 
     // Replace GatherOp by ConcatOp of specific dimensions.

test/mlir/onnx/onnx_simplify_shape_related_ops.mlir

@@ -103,7 +103,7 @@ func.func @test_pass_dims_through_concat(%arg0: tensor<?x256xi64>) -> (tensor<4x
 
 // -----
 
-func.func @test_pass_dims_through_cast_2(%arg0: tensor<?x?x200xf32>) -> tensor<2xi64> {
+func.func @test_pass_dims_through_gather(%arg0: tensor<?x?x200xf32>) -> tensor<2xi64> {
   %0 = onnx.Constant dense<[0, 1]> : 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>
@@ -113,7 +113,28 @@ func.func @test_pass_dims_through_cast_2(%arg0: tensor<?x?x200xf32>) -> tensor<2
   onnx.Return %5 : tensor<2xi64>
 
 // mlir2FileCheck.py
-// CHECK-LABEL:  func.func @test_pass_dims_through_cast_2
+// CHECK-LABEL:  func.func @test_pass_dims_through_gather
+// 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:           onnx.Return [[VAR_2_]] : tensor<2xi64>
+// CHECK:         }
+}
+
+// -----
+
+func.func @test_pass_dims_through_gather_2(%arg0: tensor<?x?x200xf32>) -> tensor<2xi64> {
+  %0 = onnx.Constant dense<[-3, -2]> : 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 dense<200> : 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>
+  onnx.Return %5 : tensor<2xi64>
+
+// mlir2FileCheck.py
+// CHECK-LABEL:  func.func @test_pass_dims_through_gather_2
 // 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>