[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

957

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

981

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.

983

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

957

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

981

Good idea, adjusted the code according to your suggestion.

983

Ack.

[rsuderman]

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

144

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

144–145

Include error message.

151

Use ++i instead of i++.

957

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.

1021

Missing newline.

mlir/test/IR/transpose-fold.mlir

47

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

144

Good point, updated to importmative error messages.

144–145

Updated to importmative error messages.

151

Ack.

957

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.

1021

Ack.

mlir/test/IR/transpose-fold.mlir

47

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, 6> getConstantPerms();
+    mlir::OpFoldResult foldIdentityTranspose();
+    mlir::OpFoldResult foldCancellableTranspose();
+  }];
+
   let hasCanonicalizer = 1;
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // TOSA Spec Section 2.10
 // Operator Class: Scatter/gather Operations.
 //===----------------------------------------------------------------------===//

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

struct NoOpOptimization : public OpRewritePattern<tosa::TransposeOp> {
   using OpRewritePattern::OpRewritePattern;
 
   LogicalResult matchAndRewrite(tosa::TransposeOp op,
                                 PatternRewriter &rewriter) const override {
     auto perm = op.getPerms();
 
     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();
   }
for (Attribute val : getMultiples().getValue()) {
     allOnes = allOnes && val.cast<IntegerAttr>().getValue().getSExtValue() == 1;
   }
 
   if (allOnes && getInput1().getType() == getType())
     return getInput1();
   return {};
 }
 
+llvm::SmallVector<int64_t, 6> TransposeOp::getConstantPerms() {
+  // Perms must be constants.
+  DenseIntElementsAttr perms;
+  if (!matchPattern(getPerms(), m_Constant(&perms)))
+    return {};
+
+  return llvm::to_vector<6>(
+      llvm::map_range(perms.getValues<APInt>(),
+                      [](const APInt &val) { return val.getSExtValue(); }));
+}
+
+OpFoldResult TransposeOp::foldIdentityTranspose() {
+  // Transpose does not change the input type.
+  if (getInput1().getType() != getType())
+    return {};
+
+  // Transpose is not the identity transpose.
+  auto perms = getConstantPerms();
+  if (!llvm::equal(llvm::seq<int64_t>(0, perms.size()), perms))
+    return {};
+
+  return getInput1();
+}
+
+OpFoldResult TransposeOp::foldCancellableTranspose() {
+  // Input is also TransposeOp - transpose(transpose(A)).
+  auto innerTranspose = getInput1().getDefiningOp<TransposeOp>();
+  if (!innerTranspose)
+    return {};
+
+  auto transposePerms = getConstantPerms();
+  auto innerTransposePerms = innerTranspose.getConstantPerms();
+
+  // Number of Perms values is same and positive.
+  if (transposePerms.size() != innerTransposePerms.size() ||
+      transposePerms.empty())
+    return {};
+
+  // Transpose & inner transpose cancel each other.
+  for (unsigned int i = 0; i < transposePerms.size(); ++i) {
+    if (transposePerms[innerTransposePerms[i]] != i)
+      return {};
+  }
+
+  return innerTranspose.getInput1();
+}
+
 OpFoldResult TransposeOp::fold(ArrayRef<Attribute> operands) {
   if (!operands[1])
     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>());
   }
 
-  auto perms = llvm::to_vector<6>(llvm::map_range(
-      operands[1].cast<DenseIntElementsAttr>().getValues<APInt>(),
-      [](const APInt &val) { return val.getSExtValue(); }));
+  if (auto folded = foldIdentityTranspose())
+    return folded;
+  if (auto folded = foldCancellableTranspose())
+    return folded;
 
-  if (llvm::equal(llvm::seq<int64_t>(0, perms.size()), perms) &&
-      getInput1().getType() == getType())
-    return getInput1();
   return {};
 }

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, 1, 0, 2]> : tensor<4xi32>
+// CHECK:           %[[VAL_2:.*]] = arith.constant dense<[1, 2, 0, 3]> : tensor<4xi32>
+// CHECK:           %[[VAL_3:.*]] = "tosa.transpose"(%[[VAL_0]], %[[VAL_2]]) : (tensor<1x2x3x4xi32>, tensor<4xi32>) -> tensor<2x3x1x4xi32>
+// CHECK:           %[[VAL_4:.*]] = "tosa.transpose"(%[[VAL_3]], %[[VAL_1]]) : (tensor<2x3x1x4xi32>, tensor<4xi32>) -> tensor<4x3x2x1xi32>
+// CHECK:           return %[[VAL_4]] : 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