[mlir][tosa] Prefer tosa.transpose composition canonicalization to reshape

[rsuderman]

Summary

It is preferred to merge tosa.transpose operations together rather than convert
one to a tosa.reshape. This is to leverage the tosa.transpose -> tosa.transpose
merging canonicalization.

[AviadCo]

This is good idea!
I am wondering if we can use the "benefit" (https://mlir.llvm.org/docs/PatternRewriter/#benefit) mechnism here of Pattern Rewriting, something like:

enum TransposeRewriteBenefit : uint8_t {
  None = 0,
  ReshapeTranspose,
  ConsolidateTransposes,
} TransposeBenefit;

And then use is in c'tors:

TransposeIsReshape(MLIRContext *context)
    : OpRewritePattern(context, ReshapeTranspose) {}

ConsolidateTransposeOptimization(MLIRContext *context)
    : OpRewritePattern(context, ConsolidateTransposes) {}

More at: https://mlir.llvm.org/docs/DeclarativeRewrites/#adjusting-benefits

[rsuderman]

Moved transpose-fold.mlir to the correct folder

[rsuderman]

In D141434#4042451, @AviadCo wrote:

This is good idea!
I am wondering if we can use the "benefit" (https://mlir.llvm.org/docs/PatternRewriter/#benefit) mechnism here of Pattern Rewriting, something like:

enum TransposeRewriteBenefit : uint8_t {
  None = 0,
  ReshapeTranspose,
  ConsolidateTransposes,
} TransposeBenefit;

And then use is in c'tors:

TransposeIsReshape(MLIRContext *context)
    : OpRewritePattern(context, ReshapeTranspose) {}

ConsolidateTransposeOptimization(MLIRContext *context)
    : OpRewritePattern(context, ConsolidateTransposes) {}

More at: https://mlir.llvm.org/docs/DeclarativeRewrites/#adjusting-benefits

Unfortunately benefit cannot be used to target the beneficial pattern as the transpose-transpose case depends on the interaction of two operations. Depending on the processing order it is possible for the top-most case to canonicalize to a reshape and be unable to merge with the consuming transpose. It may be possible to rearrange whether a consumer/producer is checked for being a transpose allowing the benefit pattern to be used however if a pattern used the canonicalization patterns but processed it in the opposite ordering it would then fail.

[AviadCo]

I see, good point.

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

struct TransposeIsReshape : public OpRewritePattern<tosa::TransposeOp> {
   using OpRewritePattern::OpRewritePattern;
 
   LogicalResult matchAndRewrite(tosa::TransposeOp op,
                                 PatternRewriter &rewriter) const override {
     DenseIntElementsAttr permAttr;
     if (!matchPattern(op.getPerms(), m_Constant(&permAttr)))
       return rewriter.notifyMatchFailure(op, "Non-constant permutation");
 
+    if (op.getInput1().getDefiningOp<tosa::TransposeOp>())
+      return rewriter.notifyMatchFailure(
+          op, "Src is from transpose, can compose transposes");
+
+    Value result = op.getResult();
+    for (Operation *subop : result.getUsers()) {
+      if (dyn_cast_or_null<tosa::TransposeOp>(subop))
+        return rewriter.notifyMatchFailure(
+            op, "Dest is used by transpose, can compose transposes");
+    }
+
     auto input = op.getInput1();
     auto inputTy = input.getType().cast<ShapedType>();
     if (!inputTy.hasRank())
       return rewriter.notifyMatchFailure(op, "Unranked input.");
 
     int64_t numDynDims = 0;
     for (int i = 0; i < inputTy.getRank(); ++i)
       if (inputTy.isDynamicDim(i))

mlir/test/IR/transpose-fold.mlir

 
 func.func @test_do_not_cancel_different_transpose(%arg0: tensor<2x3x4x5xi32>) -> (tensor<5x4x3x2xi32>) {
   %0 = arith.constant dense<[1, 2, 0, 3]> : tensor<4xi32>
   %1 = "tosa.transpose"(%arg0, %0) : (tensor<2x3x4x5xi32>, tensor<4xi32>) -> (tensor<3x4x2x5xi32>)
   %2 = arith.constant dense<[3, 1, 0, 2]> : tensor<4xi32>
   %3 = "tosa.transpose"(%1, %2) : (tensor<3x4x2x5xi32>, tensor<4xi32>) -> tensor<5x4x3x2xi32>
   return %3 : tensor<5x4x3x2xi32>
 }
+
+// -----
+
+// CHECK-LABEL:   func.func @test_prefer_compose_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_prefer_compose_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>
+}