review comments - cleanup handling of constant pad values

Author: sayeddla

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;
 
@@ -106,36 +106,12 @@ class ONNXPadOpLoweringToTOSA : public OpConversionPattern<ONNXPadOp> {
           padsList1,
           rewriter.create<mlir::tosa::ConstOp>(
               op->getLoc(), constType, constAttr));
-    } else if (!constValue.getDefiningOp<mlir::tosa::ConstOp>()) {
-      rewriter.replaceOpWithNewOp<mlir::tosa::PadOp>(
-          op, resultType, data, padsList1, constValue);
     } else {
-      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);
-        }
-      }
+      Value reshapeToSplattedConst = tosaBuilder.reshape(constValue, {});
       rewriter.replaceOpWithNewOp<mlir::tosa::PadOp>(
-          op, resultType, data, padsList1, constTosaTensor);
+          op, resultType, data, padsList1, reshapeToSplattedConst);
     }
-
     return success();
   }
 };

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
@@ -170,21 +170,36 @@ func.func @test_pad_f32_non_constant_padval(%arg0: tensor<20x16x44x32xf32>, %arg
     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-DAG:       [[VAR_0_:%.+]] = "tosa.const"() <{value = dense<[0, 4, 1, 5, 2, 6, 3, 7]> : tensor<8xi64>}> : () -> tensor<8xi64>
+// CHECK:           [[VAR_0_:%.+]] = "tosa.const"() <{value = dense<[0, 4, 1, 5, 2, 6, 3, 7]> : tensor<8xi64>}> : () -> tensor<8xi64>
 // CHECK:           [[VAR_1_:%.+]] = tosa.pad [[PARAM_0_]], [[VAR_0_]], [[PARAM_1_]] : (tensor<20x16x44x32xf32>, tensor<8xi64>, 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"() <{value = dense<[0, 4, 1, 5, 2, 6, 3, 7]> : tensor<8xi64>}> : () -> tensor<8xi64>
+// 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>, tensor<8xi64>, 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-DAG:       [[VAR_0_:%.+]] = "tosa.const"() <{value = dense<[0, 4, 1, 5, 2, 6, 3, 7]> : tensor<8xi64>}> : () -> tensor<8xi64>
+// CHECK:           [[VAR_0_:%.+]] = "tosa.const"() <{value = dense<[0, 4, 1, 5, 2, 6, 3, 7]> : tensor<8xi64>}> : () -> tensor<8xi64>
 // CHECK:           [[VAR_1_:%.+]] = tosa.pad [[PARAM_0_]], [[VAR_0_]], [[PARAM_1_]] : (tensor<20x16x44x32xi64>, tensor<8xi64>, tensor<i64>) -> tensor<24x22x52x42xi64>
 // CHECK:           return [[VAR_1_]] : tensor<24x22x52x42xi64>
 }
@@ -197,7 +212,7 @@ func.func @test_pad_f16_non_constant_padval(%arg0: tensor<20x16x44x32xf16>, %arg
     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-DAG:       [[VAR_0_:%.+]] = "tosa.const"() <{value = dense<[0, 4, 1, 5, 2, 6, 3, 7]> : tensor<8xi64>}> : () -> tensor<8xi64>
+// CHECK:           [[VAR_0_:%.+]] = "tosa.const"() <{value = dense<[0, 4, 1, 5, 2, 6, 3, 7]> : tensor<8xi64>}> : () -> tensor<8xi64>
 // CHECK:           [[VAR_1_:%.+]] = tosa.pad [[PARAM_0_]], [[VAR_0_]], [[PARAM_1_]] : (tensor<20x16x44x32xf16>, tensor<8xi64>, tensor<f16>) -> tensor<24x22x52x42xf16>
 // CHECK:           return [[VAR_1_]] : tensor<24x22x52x42xf16>
 }