Merge pull request #427 from Xilinx/matthias.fix_concat_zero_fold

TOSA: concat: fix canonicalization that would result in concat with no operands
Xilinx · Dec 16, 2024 · 14e4586 · 14e4586
2 parents c9c2863 + ae05589
commit 14e4586
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 9 deletions.
diff --git a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
@@ -1349,19 +1349,13 @@ OpFoldResult tosa::AbsOp::fold(FoldAdaptor adaptor) {
   return {};
 }
 
-static bool hasZeroSize(Type ty) {
-  auto ranked = dyn_cast<RankedTensorType>(ty);
-  if (!ranked)
-    return false;
-  return any_of(ranked.getShape(), [](auto d) { return d == 0; });
-}
-
 OpFoldResult ConcatOp::fold(FoldAdaptor adaptor) {
   /// Remove operands that have zero elements.
   bool changed = false;
   for (size_t i = 0; i < getInput1().size(); ) {
-    auto input = getInput1()[i];
-    if (hasZeroSize(input.getType())) {
+    auto input = cast<RankedTensorType>(getInput1()[i].getType());
+    // Ensure that we have at least one operand left.
+    if (input.getDimSize(getAxis()) == 0 && getInput1().size() > 1) {
       getInput1Mutable().erase(i);
       changed = true;
     } else {

diff --git a/mlir/test/Dialect/Tosa/canonicalize.mlir b/mlir/test/Dialect/Tosa/canonicalize.mlir
@@ -204,6 +204,32 @@ func.func @concat_fold_zero(%arg0: tensor<?x0xf32>, %arg1: tensor<?x1xf32>, %arg
   %0 = tosa.concat %arg0, %arg1, %arg2 {axis = 1 : i32}: (tensor<?x0xf32>, tensor<?x1xf32>, tensor<?x2xf32>) -> tensor<?x3xf32>
   return %0 : tensor<?x3xf32>
 }
+// -----
+
+// CHECK-LABEL: @concat_fold_zero
+func.func @concat_fold_zero_all(%arg0: tensor<?x0xf32>, %arg1: tensor<?x0xf32>) -> tensor<?x0xf32> {
+  // CHECK: return %arg1
+  %0 = tosa.concat %arg0, %arg1 {axis = 1 : i32}: (tensor<?x0xf32>, tensor<?x0xf32>) -> tensor<?x0xf32>
+  return %0 : tensor<?x0xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @concat_fold_zero
+func.func @concat_fold_zero_different_axis(%arg0: tensor<0x2xf32>, %arg1: tensor<0x4xf32>) -> tensor<0x6xf32> {
+  // CHECK: tosa.concat %arg0, %arg1
+  %0 = tosa.concat %arg0, %arg1 {axis = 1 : i32}: (tensor<0x2xf32>, tensor<0x4xf32>) -> tensor<0x6xf32>
+  return %0 : tensor<0x6xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @concat_fold_zero_size
+func.func @concat_fold_zero_size(%arg0: tensor<?x0xf32>, %arg1: tensor<?x1xf32>, %arg2: tensor<?x2xf32>) -> tensor<?x3xf32> {
+  // CHECK: tosa.concat %arg1, %arg2 {axis = 1 : i32}
+  %0 = tosa.concat %arg0, %arg1, %arg2 {axis = 1 : i32}: (tensor<?x0xf32>, tensor<?x1xf32>, tensor<?x2xf32>) -> tensor<?x3xf32>
+  return %0 : tensor<?x3xf32>
+}
 
 // -----