[mlir][tosa] Remove redundant "tosa.transpose" operations

[AviadCo]

Summary

We can fold redundant Tosa::TransposeOp actions like identity tranpose/transpose(traspose).

[AviadCo]

Seperated change into NFC part and FC part.

[AviadCo]

Seperate change into to commits - NFC part and FC part.

[AviadCo]

Seperate change into 2 to commits - NFC part and FC part.

[AviadCo]

ping

[rsuderman]

Overall move the double tosa.transpose pattern to a TOSA canonicalization pattern, then replace with a
new tosa.transpose op. We can leverage the existing tosa.transpose folder to remove the extra operation
if it is a standard version.

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

970

Breaking this out to a separate function is unnecessary as the folder should focus on the simple case.

994

This is better handled by adding a TransposeOp + TransposeOp -> TransposeOp canonicalization pattern. This allows multiple transposes to merge together, then reuse the existing folder to remove the noop version.

996

We avoid putting multi-op folders together as folding patterns should avoid becoming cumbersome. These kind of operations should stick to canonicalizers.

[AviadCo]

[mlir][tosa] Remove redundant "tosa.transpose" operations

[AviadCo]

I appreciate your code review! Can you please review my new revision?

Thanks,
Aviad

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

970

This function is still usefull both in TransposeOp::fold & in ConsolidateTransposeOptimization::matchAndRewrite , I think can usefull, can you check the new revision?

994

Good idea, adjusted the code according to your suggestion.

996

Ack.

[rsuderman]

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

145

return rewriter.notifyMatchFailure(op, <error message>);

154

Include error message.

158

Use ++i instead of i++.

970

This function is still usefull both in TransposeOp::fold & in ConsolidateTransposeOptimization::matchAndRewrite , I think can usefull, can you check the new revision?

The problem is you have no way to discern if the perms is a constant value of size-0 or a non-constant value. Typically we avoid cases like returning an empty array as a failure sign to avoid ambiguous returns.

993

Missing newline.

mlir/test/IR/transpose-fold.mlir

45

Missing newline.

[AviadCo]

[mlir][tosa] Remove redundant "tosa.transpose" operations

[AviadCo]

Hi Rob,

I fixed the adidiotnal comments.

Thanks,
Aviad Cohen

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

145

Good point, updated to importmative error messages.

154

Updated to importmative error messages.

158

Ack.

970

Right. I updated the function to return a LogicalResult when failed.
In addition, I moved the function implementation to TosaOps.cpp as it seems better place for it.

993

Ack.

mlir/test/IR/transpose-fold.mlir

45

Ack.

mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td

let arguments = (ins
     Tosa_Tensor1Dto6D:$input1,
     Tosa_Int32Or64Tensor:$perms
   );
 
   let results = (
     outs Tosa_Tensor1Dto6D:$output
   );
 
+  let extraClassDeclaration = [{
+    llvm::SmallVector<int64_t> getConstantPerms();
+  }];
+
   let hasCanonicalizer = 1;
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // TOSA Spec Section 2.10
 // Operator Class: Scatter/gather Operations.
 //===----------------------------------------------------------------------===//

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

if (!notOp)
     return failure();
   rewriter.updateRootInPlace(op, [&]() {
     op.getOperation()->setOperands(
         {notOp.getInput1(), op.getOnFalse(), op.getOnTrue()});
   });
   return success();
 }
 
-struct NoOpOptimization : public OpRewritePattern<tosa::TransposeOp> {
+struct ConsolidateTransposeOptimization
+    : public OpRewritePattern<tosa::TransposeOp> {
   using OpRewritePattern::OpRewritePattern;
 
-  LogicalResult matchAndRewrite(tosa::TransposeOp op,
+  LogicalResult matchAndRewrite(tosa::TransposeOp transposeOp,
                                 PatternRewriter &rewriter) const override {
-    auto perm = op.getPerms();
+    // Input is also TransposeOp - transpose(transpose(A)).
+    auto innerTranspose =
+        transposeOp.getInput1().getDefiningOp<tosa::TransposeOp>();
+    if (!innerTranspose)
+      return failure();
+
+    SmallVector<int64_t> transposePerms = transposeOp.getConstantPerms();
+    SmallVector<int64_t> innerTransposePerms =
+        innerTranspose.getConstantPerms();
+
+    // Number of Perms values is same and positive.
+    if (transposePerms.size() != innerTransposePerms.size() ||
+        transposePerms.empty())
+      return failure();
+
+    // Consolidate transposes into one transpose.
+    SmallVector<int32_t> perms(transposePerms.size());
+    for (int i = 0, s = transposePerms.size(); i < s; i++)
+      perms[i] = innerTransposePerms[transposePerms[i]];
+
+    auto permsTy =
+        RankedTensorType::get(transposePerms.size(), rewriter.getI32Type());
+    auto permsAttr = DenseIntElementsAttr::get(permsTy, perms);
+    Value permsValue =
+        rewriter.create<arith::ConstantOp>(transposeOp.getLoc(), permsAttr);
+
+    rewriter.replaceOpWithNewOp<tosa::TransposeOp>(
+        transposeOp, transposeOp.getResult().getType(),
+        innerTranspose.getInput1(), permsValue);
 
-    DenseIntElementsAttr permAttr;
-    if (!matchPattern(perm, m_Constant(&permAttr))) {
-      return failure();
-    }
-
-    SmallVector<int64_t> permValues = llvm::to_vector<6>(
-        llvm::map_range(permAttr.getValues<APInt>(),
-                        [](const APInt &val) { return val.getSExtValue(); }));
-
-    for (int i = 0, s = permValues.size(); i < s; i++) {
-      if (i != permValues[i]) {
-        return failure();
-      }
-    }
-
-    rewriter.replaceOp(op, op.getInput1());
     return success();
   }
 };
 
 void TransposeOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                               MLIRContext *context) {
-  results.add<NoOpOptimization>(context);
+  results.add<ConsolidateTransposeOptimization>(context);
 }
 
 struct AddZeroOptimization : public OpRewritePattern<tosa::AddOp> {
   using OpRewritePattern::OpRewritePattern;
 
   LogicalResult matchAndRewrite(tosa::AddOp op,
                                 PatternRewriter &rewriter) const override {
     auto input1 = op.getInput1();
for (Attribute val : getMultiples().getValue()) {
     allOnes = allOnes && val.cast<IntegerAttr>().getValue().getSExtValue() == 1;
   }
 
   if (allOnes && getInput1().getType() == getType())
     return getInput1();
   return {};
 }
 
-OpFoldResult TransposeOp::fold(ArrayRef<Attribute> operands) {
-  if (!operands[1])
+llvm::SmallVector<int64_t> TransposeOp::getConstantPerms() {
+  // Perms must be constants.
+  DenseIntElementsAttr perms;
+  if (!matchPattern(getPerms(), m_Constant(&perms)))
     return {};
 
-  // Transposing splat values just means reshaping.
-  if (auto input = operands[0].dyn_cast_or_null<DenseElementsAttr>()) {
-    if (input.isSplat())
-      return input.reshape(getType().cast<ShapedType>());
-  }
+  return llvm::to_vector(
+      llvm::map_range(perms.getValues<APInt>(),
+                      [](const APInt &val) { return val.getSExtValue(); }));
 
-  auto perms = llvm::to_vector<6>(llvm::map_range(
-      operands[1].cast<DenseIntElementsAttr>().getValues<APInt>(),
-      [](const APInt &val) { return val.getSExtValue(); }));
+OpFoldResult TransposeOp::fold(ArrayRef<Attribute> operands) {
+  // Transpose does not change the input type.
+  if (getInput1().getType() != getType())
+    return {};
 
-  if (llvm::equal(llvm::seq<int64_t>(0, perms.size()), perms) &&
-      getInput1().getType() == getType())
-    return getInput1();
+  // Transpose is not the identity transpose.
+  auto perms = getConstantPerms();
+  if (!llvm::equal(llvm::seq<int64_t>(0, perms.size()), perms))
     return {};
-}
+
+  return getInput1();
+}
 No newline at end of file

mlir/test/IR/transpose-fold.mlir

+// RUN: mlir-opt %s --canonicalize -split-input-file | FileCheck %s
+
+// CHECK-LABEL:   func.func @test_cancel_transpose_transpose(
+// CHECK-SAME:                                               %[[VAL_0:.*]]: tensor<1x2x3xi32>) -> tensor<1x2x3xi32> {
+// CHECK:           return %[[VAL_0]] : tensor<1x2x3xi32>
+// CHECK:         }
+
+func.func @test_cancel_transpose_transpose(%arg0: tensor<1x2x3xi32>) -> (tensor<1x2x3xi32>) {
+  %0 = arith.constant dense<[1, 2, 0]> : tensor<3xi32>
+  %1 = "tosa.transpose"(%arg0, %0) : (tensor<1x2x3xi32>, tensor<3xi32>) -> (tensor<2x3x1xi32>)
+  %2 = arith.constant dense<[2, 0, 1]> : tensor<3xi32>
+  %3 = "tosa.transpose"(%1, %2) : (tensor<2x3x1xi32>, tensor<3xi32>) -> tensor<1x2x3xi32>
+  return %3 : tensor<1x2x3xi32>
+}
+
+// -----
+
+// CHECK-LABEL:   func.func @test_remove_identity_transpose(
+// CHECK-SAME:                                              %[[VAL_0:.*]]: tensor<1x2x3xi32>) -> tensor<1x2x3xi32> {
+// CHECK:           return %[[VAL_0]] : tensor<1x2x3xi32>
+// CHECK:         }
+
+func.func @test_remove_identity_transpose(%arg0: tensor<1x2x3xi32>) -> (tensor<1x2x3xi32>) {
+  %0 = arith.constant dense<[0, 1, 2]> : tensor<3xi32>
+  %1 = "tosa.transpose"(%arg0, %0) : (tensor<1x2x3xi32>, tensor<3xi32>) -> (tensor<1x2x3xi32>)
+  return %1 : tensor<1x2x3xi32>
+}
+
+// -----
+
+// CHECK-LABEL:   func.func @test_do_not_cancel_different_transpose(
+// CHECK-SAME:                                                      %[[VAL_0:.*]]: tensor<1x2x3x4xi32>) -> tensor<4x3x2x1xi32> {
+// CHECK:           %[[VAL_1:.*]] = arith.constant dense<[3, 2, 1, 0]> : tensor<4xi32>
+// CHECK:           %[[VAL_2:.*]] = "tosa.transpose"(%[[VAL_0]], %[[VAL_1]]) : (tensor<1x2x3x4xi32>, tensor<4xi32>) -> tensor<4x3x2x1xi32>
+// CHECK:           return %[[VAL_2]] : tensor<4x3x2x1xi32>
+// CHECK:         }
+
+func.func @test_do_not_cancel_different_transpose(%arg0: tensor<1x2x3x4xi32>) -> (tensor<4x3x2x1xi32>) {
+  %0 = arith.constant dense<[1, 2, 0, 3]> : tensor<4xi32>
+  %1 = "tosa.transpose"(%arg0, %0) : (tensor<1x2x3x4xi32>, tensor<4xi32>) -> (tensor<2x3x1x4xi32>)
+  %2 = arith.constant dense<[3, 1, 0, 2]> : tensor<4xi32>
+  %3 = "tosa.transpose"(%1, %2) : (tensor<2x3x1x4xi32>, tensor<4xi32>) -> tensor<4x3x2x1xi32>
+  return %3 : tensor<4x3x2x1xi32>
+}
 No newline at end of file