Merge branch 'feature/onnx-to-tosa' into koshrai2.backport.pad.lowering.changes

Author: kosh-rai

src/Conversion/ONNXToTOSA/Tensor/Gather.cpp

@@ -17,6 +17,7 @@
 #include "src/Conversion/ONNXToTOSA/ONNXToTOSACommon.hpp"
 #include "src/Conversion/ONNXToTOSA/ONNXToTOSALegalizeUtils.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
+#include "src/Support/TypeUtilities.hpp"
 #include "llvm/ADT/SmallVector.h"
 
 using namespace mlir;
@@ -46,7 +47,8 @@ class ONNXGatherLoweringToTOSA : public OpConversionPattern<ONNXGatherOp> {
     if (!onnx_mlir::isRankedShapedType(inputType))
       return rewriter.notifyMatchFailure(op, "input is not a ranked tensor");
 
-    if (!hasStaticShape(result.getType()))
+    if (!hasStaticShape(inputType) || !hasStaticShape(indices.getType()) ||
+        !hasStaticShape(result.getType()))
       return rewriter.notifyMatchFailure(op, "dynamic shapes not supported");
 
     auto resultTy = dyn_cast<TensorType>(op.getType());

src/Dialect/ONNX/Transforms/SimplifyShapeRelatedOps.cpp

@@ -65,7 +65,6 @@ Now, it's straightforward to update the output shape of Reshape from
 #include "src/Dialect/ONNX/ONNXOps/OpHelper.hpp"
 #include "src/Dialect/ONNX/ONNXOps/ShapeHelper.hpp"
 #include "src/Pass/Passes.hpp"
-#include "src/Support/TypeUtilities.hpp"
 
 #define DEBUG_TYPE "simplify_shape_related_ops"
 
@@ -247,14 +246,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/conversion/onnx_to_tosa/Tensor/Gather.mlir

@@ -175,3 +175,21 @@ func.func @test_gather_dynamic_shape_indices_i32(%arg0 : tensor<?x4xf32>, %indic
 // CHECK-LABEL: test_gather_dynamic_shape_indices_i32
 // CHECK: onnx.Gather
 }
+
+// -----
+
+func.func @test_gather_dynamic_input_static_output(%arg0 : tensor<?x2xf32>, %indices: tensor<?xi64>) -> tensor<1x2xf32> {
+  %0 = "onnx.Gather"(%arg0, %indices) {axis = 0 : si64, onnx_node_name = "/Gather_16"} : (tensor<?x2xf32>, tensor<?xi64>) -> tensor<1x2xf32>
+  "func.return"(%0) : (tensor<1x2xf32>) -> ()
+// CHECK-LABEL: test_gather_dynamic_input_static_output
+// CHECK: onnx.Gather
+}
+
+// -----
+
+func.func @test_gather_dynamic_indices(%arg0 : tensor<1x2xf32>, %indices: tensor<?xi64>) -> tensor<1x2xf32> {
+  %0 = "onnx.Gather"(%arg0, %indices) {axis = 0 : si64, onnx_node_name = "/Gather_16"} : (tensor<1x2xf32>, tensor<?xi64>) -> tensor<1x2xf32>
+  "func.return"(%0) : (tensor<1x2xf32>) -> ()
+// CHECK-LABEL: test_gather_dynamic_indices
+// CHECK: onnx.Gather
+}

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>