[MLIR] [TOSA]: fix input index of tosa.pad folding. Improve lit tests.

Xilinx · Sep 8, 2023 · dc3e56c · dc3e56c
1 parent cc3af09
commit dc3e56c
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 37 deletions.
diff --git a/mlir/lib/Dialect/Tosa/Transforms/TosaFolders.cpp b/mlir/lib/Dialect/Tosa/Transforms/TosaFolders.cpp
@@ -1318,14 +1318,23 @@ DenseElementsAttr padType(ShapedType inputType, ElementsAttr inputValues,
   for (size_t outIndex = 0, e = outputValues.size(); outIndex < e; ++outIndex) {
     auto indexInTarget = offsetToIndex(outputShape, outIndex);
 
+    llvm::for_each(llvm::enumerate(indexInTarget), [&](const auto &dimInfo) {
+      auto index = dimInfo.index();
+      auto i = dimInfo.value() - paddingVals[index * 2];
+
+      // Update index so it points to the right position
+      // when this is not a padConst value.
+      indexInTarget[index] = i;
+    });
+
     bool isPad =
         llvm::any_of(llvm::enumerate(indexInTarget), [&](const auto &dimInfo) {
           auto index = dimInfo.index();
-          auto i = dimInfo.value() - paddingVals[index * 2];
-          return static_cast<bool>(i < 0 || i >= inputShape[index]);
+          auto value = dimInfo.value();
+          return static_cast<bool>(value < 0 || value >= inputShape[index]);
         });
 
-    auto inputIndexOffset = indexToOffset(outputShape, indexInTarget);
+    auto inputIndexOffset = indexToOffset(inputShape, indexInTarget);
     outputValues[outIndex] = isPad ? padConst : values[inputIndexOffset];
   }
   return DenseElementsAttr::get(outputType,

diff --git a/mlir/test/Dialect/Tosa/constant-pad.mlir b/mlir/test/Dialect/Tosa/constant-pad.mlir
@@ -14,8 +14,8 @@ func.func @pad_bool() -> (tensor<5x5xi1>) {
 
 // CHECK-LABEL: @pad_int8
 func.func @pad_int8() -> (tensor<5x5xi8>) {
-  // CHECK: "tosa.const"() <{value = dense<{{\[\[}}1, 1, 1, 1, 1], [1, 2, 2, 1, 1], [1, 2, 2, 1, 1], [1, 1, 1, 1, 1], [1, 1, 1, 1, 1]]>
-  %0 = "tosa.const"() {value = dense<2> : tensor<2x2xi8>} : () -> tensor<2x2xi8>
+  // CHECK: "tosa.const"() <{value = dense<{{\[\[}}1, 1, 1, 1, 1], [1, 3, 4, 1, 1], [1, 5, 6, 1, 1], [1, 1, 1, 1, 1], [1, 1, 1, 1, 1]]>
+  %0 = "tosa.const"() {value = dense<[[3, 4], [5, 6]]> : tensor<2x2xi8>} : () -> tensor<2x2xi8>
   %1 = "tosa.const"() {value = dense<[[1, 2], [1, 2]]> : tensor<2x2xi64>} : () -> tensor<2x2xi64>
   %2 = "tosa.const"() {value = dense<1> : tensor<i8>} : () -> tensor<i8>
   %3 = "tosa.pad"(%0, %1, %2) : (tensor<2x2xi8>, tensor<2x2xi64>, tensor<i8>) -> tensor<5x5xi8>
@@ -24,8 +24,8 @@ func.func @pad_int8() -> (tensor<5x5xi8>) {
 
 // CHECK-LABEL: @pad_int32
 func.func @pad_int32() -> (tensor<5x5xi32>) {
-  // CHECK: "tosa.const"() <{value = dense<{{\[\[}}1, 1, 1, 1, 1], [1, 2, 2, 1, 1], [1, 2, 2, 1, 1], [1, 1, 1, 1, 1], [1, 1, 1, 1, 1]]>
-  %0 = "tosa.const"() {value = dense<2> : tensor<2x2xi32>} : () -> tensor<2x2xi32>
+  // CHECK: "tosa.const"() <{value = dense<{{\[\[}}1, 1, 1, 1, 1], [1, 3, 4, 1, 1], [1, 5, 6, 1, 1], [1, 1, 1, 1, 1], [1, 1, 1, 1, 1]]>
+  %0 = "tosa.const"() {value = dense<[[3, 4], [5, 6]]> : tensor<2x2xi32>} : () -> tensor<2x2xi32>
   %1 = "tosa.const"() {value = dense<[[1, 2], [1, 2]]> : tensor<2x2xi64>} : () -> tensor<2x2xi64>
   %2 = "tosa.const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
   %3 = "tosa.pad"(%0, %1, %2) : (tensor<2x2xi32>, tensor<2x2xi64>, tensor<i32>) -> tensor<5x5xi32>
@@ -34,8 +34,8 @@ func.func @pad_int32() -> (tensor<5x5xi32>) {
 
 // CHECK-LABEL: @pad_int32_default_value
 func.func @pad_int32_default_value() -> (tensor<5x5xi32>) {
-  // CHECK: "tosa.const"() <{value = dense<{{\[\[}}0, 0, 0, 0, 0], [0, 2, 2, 0, 0], [0, 2, 2, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0]]>
-  %0 = "tosa.const"() {value = dense<2> : tensor<2x2xi32>} : () -> tensor<2x2xi32>
+  // CHECK: "tosa.const"() <{value = dense<{{\[\[}}0, 0, 0, 0, 0], [0, 3, 4, 0, 0], [0, 5, 6, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0]]>
+  %0 = "tosa.const"() {value = dense<[[3, 4], [5, 6]]> : tensor<2x2xi32>} : () -> tensor<2x2xi32>
   %1 = "tosa.const"() {value = dense<[[1, 2], [1, 2]]> : tensor<2x2xi64>} : () -> tensor<2x2xi64>
   %2 = "tosa.pad"(%0, %1) : (tensor<2x2xi32>, tensor<2x2xi64>) -> tensor<5x5xi32>
   return %2 : tensor<5x5xi32>
@@ -44,14 +44,14 @@ func.func @pad_int32_default_value() -> (tensor<5x5xi32>) {
 // CHECK-LABEL: @pad_bfloat16
 func.func @pad_bfloat16() -> (tensor<5x5xbf16>) {
   // CHECK: "tosa.const"()
-  // CHECK-SAME: {{\[\[}}2.000000e+00, 2.000000e+00, 2.000000e+00, 2.000000e+00, 2.000000e+00],
-  // CHECK-SAME: [2.000000e+00, 1.000000e+00, 1.000000e+00, 2.000000e+00, 2.000000e+00],
-  // CHECK-SAME: [2.000000e+00, 1.000000e+00, 1.000000e+00, 2.000000e+00, 2.000000e+00],
-  // CHECK-SAME: [2.000000e+00, 2.000000e+00, 2.000000e+00, 2.000000e+00, 2.000000e+00],
-  // CHECK-SAME: [2.000000e+00, 2.000000e+00, 2.000000e+00, 2.000000e+00, 2.000000e+00]]>
-  %0 = "tosa.const"() {value = dense<1.0> : tensor<2x2xbf16>} : () -> tensor<2x2xbf16>
+  // CHECK-SAME: {{\[\[}}-1.000000e+00, -1.000000e+00, -1.000000e+00, -1.000000e+00, -1.000000e+00],
+  // CHECK-SAME: [-1.000000e+00, 1.000000e+00, 2.000000e+00, -1.000000e+00, -1.000000e+00],
+  // CHECK-SAME: [-1.000000e+00, 3.000000e+00, 4.000000e+00, -1.000000e+00, -1.000000e+00],
+  // CHECK-SAME: [-1.000000e+00, -1.000000e+00, -1.000000e+00, -1.000000e+00, -1.000000e+00],
+  // CHECK-SAME: [-1.000000e+00, -1.000000e+00, -1.000000e+00, -1.000000e+00, -1.000000e+00]]>
+  %0 = "tosa.const"() {value = dense<[[1.0, 2.0], [3.0, 4.0]]> : tensor<2x2xbf16>} : () -> tensor<2x2xbf16>
   %1 = "tosa.const"() {value = dense<[[1, 2], [1, 2]]> : tensor<2x2xi64>} : () -> tensor<2x2xi64>
-  %2 = "tosa.const"() {value = dense<2.0> : tensor<bf16>} : () -> tensor<bf16>
+  %2 = "tosa.const"() {value = dense<-1.0> : tensor<bf16>} : () -> tensor<bf16>
   %3 = "tosa.pad"(%0, %1, %2) : (tensor<2x2xbf16>, tensor<2x2xi64>, tensor<bf16>) -> tensor<5x5xbf16>
   return %3 : tensor<5x5xbf16>
 }
@@ -60,11 +60,11 @@ func.func @pad_bfloat16() -> (tensor<5x5xbf16>) {
 func.func @pad_bfloat16_default_value() -> (tensor<5x5xbf16>) {
   // CHECK: "tosa.const"()
   // CHECK-SAME: {{\[\[}}0.000000e+00, 0.000000e+00, 0.000000e+00, 0.000000e+00, 0.000000e+00],
-  // CHECK-SAME: [0.000000e+00, 1.000000e+00, 1.000000e+00, 0.000000e+00, 0.000000e+00],
-  // CHECK-SAME: [0.000000e+00, 1.000000e+00, 1.000000e+00, 0.000000e+00, 0.000000e+00],
+  // CHECK-SAME: [0.000000e+00, 1.000000e+00, 2.000000e+00, 0.000000e+00, 0.000000e+00],
+  // CHECK-SAME: [0.000000e+00, 3.000000e+00, 4.000000e+00, 0.000000e+00, 0.000000e+00],
   // CHECK-SAME: [0.000000e+00, 0.000000e+00, 0.000000e+00, 0.000000e+00, 0.000000e+00],
-  // CHECK-SAME: [0.000000e+00, 0.000000e+00, 0.000000e+00, 0.000000e+00, 0.000000e+00]]>
-  %0 = "tosa.const"() {value = dense<1.0> : tensor<2x2xbf16>} : () -> tensor<2x2xbf16>
+  // CHECK-SAME: [0.000000e+00, 0.000000e+00, 0.000000e+00, 0.000000e+00, 0.000000e+00]]
+  %0 = "tosa.const"() {value = dense<[[1.0, 2.0], [3.0, 4.0]]> : tensor<2x2xbf16>} : () -> tensor<2x2xbf16>
   %1 = "tosa.const"() {value = dense<[[1, 2], [1, 2]]> : tensor<2x2xi64>} : () -> tensor<2x2xi64>
   %2 = "tosa.pad"(%0, %1) : (tensor<2x2xbf16>, tensor<2x2xi64>) -> tensor<5x5xbf16>
   return %2 : tensor<5x5xbf16>
@@ -73,35 +73,35 @@ func.func @pad_bfloat16_default_value() -> (tensor<5x5xbf16>) {
 // CHECK-LABEL: @pad_f32_3d
 func.func @pad_f32_3d() -> (tensor<3x3x4xf32>) {
   // CHECK: "tosa.const"()
-  // CHECK-SAME: {{\[\[}}[1.000000e+00, 1.000000e+00, 1.000000e+00, 1.000000e+00],
-  // CHECK-SAME: [1.000000e+00, 1.000000e+00, 1.000000e+00, 1.000000e+00],
-  // CHECK-SAME: [1.000000e+00, 1.000000e+00, 1.000000e+00, 1.000000e+00]],
-  // CHECK-SAME: {{\[\[}}1.000000e+00, 2.000000e+00, 2.000000e+00, 1.000000e+00],
-  // CHECK-SAME: [1.000000e+00, 2.000000e+00, 2.000000e+00, 1.000000e+00],
-  // CHECK-SAME: [1.000000e+00, 1.000000e+00, 1.000000e+00, 1.000000e+00]],
-  // CHECK-SAME: {{\[\[}}1.000000e+00, 2.000000e+00, 2.000000e+00, 1.000000e+00],
-  // CHECK-SAME: [1.000000e+00, 2.000000e+00, 2.000000e+00, 1.000000e+00],
-  // CHECK-SAME: [1.000000e+00, 1.000000e+00, 1.000000e+00, 1.000000e+00]]]>
-  %0 = "tosa.const"() {value = dense<2.0> : tensor<2x2x2xf32>} : () -> tensor<2x2x2xf32>
+  // CHECK-SAME: {{\[\[}}-1.000000e+00, -1.000000e+00, -1.000000e+00, -1.000000e+00],
+  // CHECK-SAME: [-1.000000e+00, -1.000000e+00, -1.000000e+00, -1.000000e+00],
+  // CHECK-SAME: [-1.000000e+00, -1.000000e+00, -1.000000e+00, -1.000000e+00]],
+  // CHECK-SAME: {{\[\[}}-1.000000e+00, 1.000000e+00, 2.000000e+00, -1.000000e+00],
+  // CHECK-SAME: [-1.000000e+00, 3.000000e+00, 4.000000e+00, -1.000000e+00],
+  // CHECK-SAME: [-1.000000e+00, -1.000000e+00, -1.000000e+00, -1.000000e+00]],
+  // CHECK-SAME: {{\[\[}}-1.000000e+00, 5.000000e+00, 6.000000e+00, -1.000000e+00],
+  // CHECK-SAME: [-1.000000e+00, 7.000000e+00, 8.000000e+00, -1.000000e+00],
+  // CHECK-SAME: [-1.000000e+00, -1.000000e+00, -1.000000e+00, -1.000000e+00]]]>
+  %0 = "tosa.const"() {value = dense<[[[1.0, 2.0], [3.0, 4.0]],[[5.0, 6.0], [7.0, 8.0]]]> : tensor<2x2x2xf32>} : () -> tensor<2x2x2xf32>
   %1 = "tosa.const"() {value = dense<[[1, 0], [0, 1], [1, 1]]> : tensor<3x2xi64>} : () -> tensor<3x2xi64>
-  %2 = "tosa.const"() {value = dense<1.0> : tensor<f32>} : () -> tensor<f32>
+  %2 = "tosa.const"() {value = dense<-1.0> : tensor<f32>} : () -> tensor<f32>
   %3 = "tosa.pad"(%0, %1, %2) : (tensor<2x2x2xf32>, tensor<3x2xi64>, tensor<f32>) -> tensor<3x3x4xf32>
   return %3 : tensor<3x3x4xf32>
 }
 
 // CHECK-LABEL: @pad_f32_3d_default_value
 func.func @pad_f32_3d_default_value() -> (tensor<3x3x4xf32>) {
   // CHECK: "tosa.const"()
-  // CHECK-SAME: {{\[\[}}[0.000000e+00, 0.000000e+00, 0.000000e+00, 0.000000e+00],
+  // CHECK-SAME: {{\[\[}}0.000000e+00, 0.000000e+00, 0.000000e+00, 0.000000e+00],
   // CHECK-SAME: [0.000000e+00, 0.000000e+00, 0.000000e+00, 0.000000e+00],
   // CHECK-SAME: [0.000000e+00, 0.000000e+00, 0.000000e+00, 0.000000e+00]],
-  // CHECK-SAME: {{\[\[}}0.000000e+00, 2.000000e+00, 2.000000e+00, 0.000000e+00],
-  // CHECK-SAME: [0.000000e+00, 2.000000e+00, 2.000000e+00, 0.000000e+00],
+  // CHECK-SAME: {{\[\[}}0.000000e+00, 1.000000e+00, 2.000000e+00, 0.000000e+00],
+  // CHECK-SAME: [0.000000e+00, 3.000000e+00, 4.000000e+00, 0.000000e+00],
   // CHECK-SAME: [0.000000e+00, 0.000000e+00, 0.000000e+00, 0.000000e+00]],
-  // CHECK-SAME: {{\[\[}}0.000000e+00, 2.000000e+00, 2.000000e+00, 0.000000e+00],
-  // CHECK-SAME: [0.000000e+00, 2.000000e+00, 2.000000e+00, 0.000000e+00],
+  // CHECK-SAME: {{\[\[}}0.000000e+00, 5.000000e+00, 6.000000e+00, 0.000000e+00],
+  // CHECK-SAME: [0.000000e+00, 7.000000e+00, 8.000000e+00, 0.000000e+00],
   // CHECK-SAME: [0.000000e+00, 0.000000e+00, 0.000000e+00, 0.000000e+00]]]>
-  %0 = "tosa.const"() {value = dense<2.0> : tensor<2x2x2xf32>} : () -> tensor<2x2x2xf32>
+  %0 = "tosa.const"() {value = dense<[[[1.0, 2.0], [3.0, 4.0]],[[5.0, 6.0], [7.0, 8.0]]]> : tensor<2x2x2xf32>} : () -> tensor<2x2x2xf32>
   %1 = "tosa.const"() {value = dense<[[1, 0], [0, 1], [1, 1]]> : tensor<3x2xi64>} : () -> tensor<3x2xi64>
   %2 = "tosa.pad"(%0, %1) : (tensor<2x2x2xf32>, tensor<3x2xi64>) -> tensor<3x3x4xf32>
   return %2 : tensor<3x3x4xf32>