Merge branch 'feature/onnx-to-tosa' into chaitany.fix_convtranpose_4x4_mismatch

Author: chaitanyakamarapu

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/Conversion/ONNXToTOSA/Tensor/PaddingOp.cpp

@@ -22,8 +22,8 @@
 #include "src/Conversion/ONNXToTOSA/ONNXToTOSALegalizeUtils.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
 #include "llvm/ADT/APFloat.h"
+#include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/SmallVector.h"
-#include "llvm/Support/Casting.h"
 
 using namespace mlir;
 
@@ -61,11 +61,9 @@ class ONNXPadOpLoweringToTOSA : public OpConversionPattern<ONNXPadOp> {
           op, "only 'constant' mode is supported");
     }
 
-    if (!pads.getDefiningOp<mlir::tosa::ConstOp>() ||
-        !(constValue.getDefiningOp<mlir::tosa::ConstOp>() ||
-            constValue.getDefiningOp<ONNXNoneOp>())) {
+    if (!pads.getDefiningOp<mlir::tosa::ConstOp>()) {
       return rewriter.notifyMatchFailure(
-          op, "only tosa.const operands are supported");
+          op, "only tosa.const 'padding' values are supported");
     }
     // creating the DenseElementsAttr using pads values.
     auto denseAttr = tosa::getValueFromTosaConst<ElementsAttr>(pads);
@@ -90,33 +88,7 @@ class ONNXPadOpLoweringToTOSA : public OpConversionPattern<ONNXPadOp> {
     mlir::Type resultType =
         getTypeConverter()->convertType(op.getResult().getType());
 
-    if (!isa<NoneType>(constValue.getType())) {
-      auto valueAttr = tosa::getValueFromTosaConst<ElementsAttr>(constValue);
-      TosaBuilder tosaBuilder(rewriter, loc);
-
-      Value constTosaTensor;
-      if (isa<FloatType>(valueAttr.getElementType())) {
-        auto valueIt = valueAttr.getValues<FloatAttr>().begin();
-        const float valueFloat = cast<FloatAttr>(*valueIt).getValueAsDouble();
-        constTosaTensor = tosaBuilder.getSplattedConst(
-            valueFloat, valueAttr.getElementType(), 0);
-      } else {
-        assert(isTOSAInt(elementDtype) && "Already validated");
-        auto valueIt = valueAttr.getValues<IntegerAttr>().begin();
-        auto valueAsAPInt = cast<IntegerAttr>(*valueIt).getValue();
-        auto asIntegerTy = cast<IntegerType>(valueAttr.getElementType());
-        if (asIntegerTy.isUnsigned()) {
-          constTosaTensor = tosaBuilder.getSplattedConst(
-              valueAsAPInt.getZExtValue(), asIntegerTy, 0);
-        } else {
-          constTosaTensor = tosaBuilder.getSplattedConst(
-              valueAsAPInt.getSExtValue(), asIntegerTy, 0);
-        }
-      }
-      rewriter.replaceOpWithNewOp<mlir::tosa::PadOp>(
-          op, resultType, data, padsList1, constTosaTensor);
-
-    } else {
+    if (isa<NoneType>(constValue.getType())) {
       auto constType = RankedTensorType::get({}, elementDtype);
 
       DenseElementsAttr constAttr;
@@ -134,8 +106,12 @@ class ONNXPadOpLoweringToTOSA : public OpConversionPattern<ONNXPadOp> {
           padsList1,
           rewriter.create<mlir::tosa::ConstOp>(
               op->getLoc(), constType, constAttr));
+    } else {
+      TosaBuilder tosaBuilder(rewriter, loc);
+      Value reshapeToSplattedConst = tosaBuilder.reshape(constValue, {});
+      rewriter.replaceOpWithNewOp<mlir::tosa::PadOp>(
+          op, resultType, data, padsList1, reshapeToSplattedConst);
     }
-
     return success();
   }
 };

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/conversion/onnx_to_tosa/Tensor/Padding.mlir

@@ -1,4 +1,4 @@
-// RUN: onnx-mlir-opt --shape-inference --convert-onnx-to-tosa --cse %s -split-input-file | FileCheck %s
+// RUN: onnx-mlir-opt --shape-inference --convert-onnx-to-tosa --canonicalize --cse %s -split-input-file | FileCheck %s
 
 func.func @test_pad_f32(%arg0: tensor<20x16x44x32xf32>) ->  tensor<24x22x52x42xf32>     {
     %noval = "onnx.NoValue"() {value} : () -> none
@@ -160,3 +160,59 @@ func.func @test_pad_f16_constant_none(%arg0: tensor<256x1x1x5x1xf16>) -> tensor<
 // CHECK: %[[VAR2:.*]] = tosa.pad %arg0, %[[VAR0]], %[[VAR1]] : (tensor<256x1x1x5x1xf16>, !tosa.shape<10>, tensor<f16>) -> tensor<256x1x1x5x2xf16>
 // CHECK: return %[[VAR2]] : tensor<256x1x1x5x2xf16>
 }
+
+// -----
+
+func.func @test_pad_f32_non_constant_padval(%arg0: tensor<20x16x44x32xf32>, %arg1: tensor<f32>) ->  tensor<24x22x52x42xf32>     {
+    %noval = "onnx.NoValue"() {value} : () -> none
+    %0 = "onnx.Constant"() {value = dense<[0, 1, 2, 3, 4, 5, 6, 7]> : tensor<8xi64>} : () -> tensor<8xi64> 
+    %2 = "onnx.Pad"(%arg0, %0, %arg1, %noval) {mode = "constant"} : (tensor<20x16x44x32xf32>, tensor<8xi64>, tensor<f32>, none) -> tensor<24x22x52x42xf32> 
+    return %2 :   tensor<24x22x52x42xf32> 
+// CHECK-LABEL:  func.func @test_pad_f32_non_constant_padval
+// CHECK-SAME:   ([[PARAM_0_:%.+]]: tensor<20x16x44x32xf32>, [[PARAM_1_:%.+]]: tensor<f32>) -> tensor<24x22x52x42xf32> {
+// CHECK:           [[VAR_0_:%.+]] = tosa.const_shape {value = dense<[0, 4, 1, 5, 2, 6, 3, 7]> : tensor<8xindex>} : () -> !tosa.shape<8>
+// CHECK:           [[VAR_1_:%.+]] = tosa.pad [[PARAM_0_]], [[VAR_0_]], [[PARAM_1_]] : (tensor<20x16x44x32xf32>, !tosa.shape<8>, tensor<f32>) -> tensor<24x22x52x42xf32>
+// CHECK:           return [[VAR_1_]] : tensor<24x22x52x42xf32>
+}
+
+// -----
+
+func.func @test_pad_f32_non_constant_1Dpadval(%arg0: tensor<20x16x44x32xf32>, %arg1: tensor<1xf32>) ->  tensor<24x22x52x42xf32>     {
+    %noval = "onnx.NoValue"() {value} : () -> none
+    %0 = "onnx.Constant"() {value = dense<[0, 1, 2, 3, 4, 5, 6, 7]> : tensor<8xi64>} : () -> tensor<8xi64> 
+    %2 = "onnx.Pad"(%arg0, %0, %arg1, %noval) {mode = "constant"} : (tensor<20x16x44x32xf32>, tensor<8xi64>, tensor<1xf32>, none) -> tensor<24x22x52x42xf32> 
+    return %2 :   tensor<24x22x52x42xf32> 
+// CHECK-LABEL:  func.func @test_pad_f32_non_constant_1Dpadval
+// CHECK-SAME:   ([[PARAM_0_:%.+]]: tensor<20x16x44x32xf32>, [[PARAM_1_:%.+]]: tensor<1xf32>) -> tensor<24x22x52x42xf32> {
+// CHECK-DAG:       [[VAR_0_:%.+]] = tosa.const_shape {value = dense<[0, 4, 1, 5, 2, 6, 3, 7]> : tensor<8xindex>} : () -> !tosa.shape<8>
+// CHECK-DAG:       [[VAL_1_:%.+]] = tosa.reshape [[PARAM_1_]] {new_shape = array<i64>} : (tensor<1xf32>) -> tensor<f32>
+// CHECK:           [[VAR_2_:%.+]] = tosa.pad [[PARAM_0_]], [[VAR_0_]], [[VAL_1_]] : (tensor<20x16x44x32xf32>, !tosa.shape<8>, tensor<f32>) -> tensor<24x22x52x42xf32>
+// CHECK:           return [[VAR_2_]] : tensor<24x22x52x42xf32>
+}
+
+// -----
+
+func.func @test_pad_i64_non_constant_padval(%arg0: tensor<20x16x44x32xi64>, %arg1: tensor<i64>) ->  tensor<24x22x52x42xi64>     {
+    %noval = "onnx.NoValue"() {value} : () -> none
+    %0 = "onnx.Constant"() {value = dense<[0, 1, 2, 3, 4, 5, 6, 7]> : tensor<8xi64>} : () -> tensor<8xi64> 
+    %2 = "onnx.Pad"(%arg0, %0, %arg1, %noval) {mode = "constant"} : (tensor<20x16x44x32xi64>, tensor<8xi64>, tensor<i64>, none) -> tensor<24x22x52x42xi64> 
+    return %2 :   tensor<24x22x52x42xi64> 
+// CHECK-LABEL:  func.func @test_pad_i64_non_constant_padval
+// CHECK-SAME:   ([[PARAM_0_:%.+]]: tensor<20x16x44x32xi64>, [[PARAM_1_:%.+]]: tensor<i64>) -> tensor<24x22x52x42xi64> {
+// CHECK:           [[VAR_0_:%.+]] = tosa.const_shape {value = dense<[0, 4, 1, 5, 2, 6, 3, 7]> : tensor<8xindex>} : () -> !tosa.shape<8>
+// CHECK:           [[VAR_1_:%.+]] = tosa.pad [[PARAM_0_]], [[VAR_0_]], [[PARAM_1_]] : (tensor<20x16x44x32xi64>, !tosa.shape<8>, tensor<i64>) -> tensor<24x22x52x42xi64>
+// CHECK:           return [[VAR_1_]] : tensor<24x22x52x42xi64>
+}
+
+// -----
+func.func @test_pad_f16_non_constant_padval(%arg0: tensor<20x16x44x32xf16>, %arg1: tensor<f16>) ->  tensor<24x22x52x42xf16>     {
+    %noval = "onnx.NoValue"() {value} : () -> none
+    %0 = "onnx.Constant"() {value = dense<[0, 1, 2, 3, 4, 5, 6, 7]> : tensor<8xi64>} : () -> tensor<8xi64> 
+    %2 = "onnx.Pad"(%arg0, %0, %arg1, %noval) {mode = "constant"} : (tensor<20x16x44x32xf16>, tensor<8xi64>, tensor<f16>, none) -> tensor<24x22x52x42xf16> 
+    return %2 :   tensor<24x22x52x42xf16> 
+// CHECK-LABEL:  func.func @test_pad_f16_non_constant_padval
+// CHECK-SAME:   ([[PARAM_0_:%.+]]: tensor<20x16x44x32xf16>, [[PARAM_1_:%.+]]: tensor<f16>) -> tensor<24x22x52x42xf16> {
+// CHECK:           [[VAR_0_:%.+]] = tosa.const_shape {value = dense<[0, 4, 1, 5, 2, 6, 3, 7]> : tensor<8xindex>} : () -> !tosa.shape<8>
+// CHECK:           [[VAR_1_:%.+]] = tosa.pad [[PARAM_0_]], [[VAR_0_]], [[PARAM_1_]] : (tensor<20x16x44x32xf16>, !tosa.shape<8>, tensor<f16>) -> tensor<24x22x52x42xf16>
+// CHECK:           return [[VAR_1_]] : tensor<24x22x52x42xf16>
+}

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>