Fold consecutive tosa.tiles

Author: jorickert

mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp

@@ -1362,6 +1362,21 @@ OpFoldResult TileOp::fold(FoldAdaptor adaptor) {
   bool allOnes = llvm::all_of(getMultiples(), [](int64_t v) { return v == 1; });
   if (allOnes && getInput1().getType() == getType())
     return getInput1();
+
+  if (auto inputTile = getInput1().getDefiningOp<TileOp>()) {
+    if (!inputTile->hasOneUse()) {
+      return {};
+    }
+    llvm::SmallVector<int64_t> newMultiplies{getMultiples()};
+    for (auto [idx, multiplier] : llvm::enumerate(inputTile.getMultiples())) {
+      newMultiplies[idx] *= multiplier;
+    }
+    setMultiples(newMultiplies);
+    setOperand(inputTile->getOperand(0));
+    getOperation()->setLoc(
+        FusedLoc::get(getContext(), {inputTile->getLoc(), getLoc()}));
+    return getResult();
+  }
   return {};
 }
 

mlir/test/Dialect/Tosa/canonicalize.mlir

@@ -691,6 +691,31 @@ func.func @tile_fold(%arg0: tensor<3x4xf32>) -> tensor<3x4xf32> {
 
 // -----
 
+// CHECK-LABEL:  func.func @tile_fold_consecutive
+func.func @tile_fold_consecutive(%arg0: tensor<3x4xf32>) -> tensor<6x16xf32> { 
+  // CHECK-SAME:   ([[PARAM_0_:%.+]]: tensor<3x4xf32>) -> tensor<6x16xf32> {
+  // CHECK:           [[VAR_0_:%.+]] = tosa.tile [[PARAM_0_]] {multiples = array<i64: 2, 4>} : (tensor<3x4xf32>) -> tensor<6x16xf32>
+  // CHECK:           return [[VAR_0_]] : tensor<6x16xf32>
+  %0 = tosa.tile %arg0 { multiples = array<i64: 1, 2> }: (tensor<3x4xf32>) -> tensor<3x8xf32>
+  %1 = tosa.tile %0 { multiples = array<i64: 2, 2> }: (tensor<3x8xf32>) -> tensor<6x16xf32>
+  return %1 : tensor<6x16xf32>
+}
+
+// -----
+
+// CHECK-LABEL:  func.func @tile_no_fold_consecutive_multi_use
+func.func @tile_no_fold_consecutive_multi_use(%arg0: tensor<3x4xf32>) -> (tensor<3x8xf32>, tensor<6x16xf32>) {
+  // CHECK-SAME:   ([[PARAM_0_:%.+]]: tensor<3x4xf32>) -> (tensor<3x8xf32>, tensor<6x16xf32>) {
+  // CHECK:           [[VAR_0_:%.+]] = tosa.tile [[PARAM_0_]] {multiples = array<i64: 1, 2>} : (tensor<3x4xf32>) -> tensor<3x8xf32>
+  // CHECK:           [[VAR_1_:%.+]] = tosa.tile [[VAR_0_]] {multiples = array<i64: 2, 2>} : (tensor<3x8xf32>) -> tensor<6x16xf32>
+  // CHECK:           return [[VAR_0_]], [[VAR_1_]] : tensor<3x8xf32>, tensor<6x16xf32>
+  %0 = tosa.tile %arg0 { multiples = array<i64: 1, 2> }: (tensor<3x4xf32>) -> tensor<3x8xf32>
+  %1 = tosa.tile %0 { multiples = array<i64: 2, 2> }: (tensor<3x8xf32>) -> tensor<6x16xf32>
+  return %0, %1 : tensor<3x8xf32>, tensor<6x16xf32>
+}
+
+// -----
+
 // CHECK-LABEL: @tile_nofold
 func.func @tile_nofold(%arg0: tensor<3x4xf32>) -> tensor<3x8xf32> {
   // CHECK: tosa.tile

mlir/test/Dialect/Tosa/fold_concats.mlir

@@ -21,9 +21,8 @@ func.func @concat_different_axis(%arg0: tensor<1x1x7x7xf32>) -> tensor<2x2x7x7xf
 
 // CHECK-LABEL:  func.func @concat_different_axis
 // CHECK-SAME:   ([[PARAM_0_:%.+]]: tensor<1x1x7x7xf32>) -> tensor<2x2x7x7xf32> {
-// CHECK:           [[VAR_0_:%.+]] = tosa.tile [[PARAM_0_]] {multiples = array<i64: 1, 2, 1, 1>} : (tensor<1x1x7x7xf32>) -> tensor<1x2x7x7xf32>
-// CHECK:           [[VAR_1_:%.+]] = tosa.tile [[VAR_0_]] {multiples = array<i64: 2, 1, 1, 1>} : (tensor<1x2x7x7xf32>) -> tensor<2x2x7x7xf32>
-// CHECK:           return [[VAR_1_]] : tensor<2x2x7x7xf32>
+// CHECK:           [[VAR_0_:%.+]] = tosa.tile [[PARAM_0_]] {multiples = array<i64: 2, 2, 1, 1>} : (tensor<1x1x7x7xf32>) -> tensor<2x2x7x7xf32>
+// CHECK:           return [[VAR_0_]] : tensor<2x2x7x7xf32>
 // CHECK:         }
 
 // -----