[mlir][tosa] Added tosa.reverse folder

[rsuderman]

Summary

Fold cases where a tosa.reverse is a splat or reversing a dim
of length-1.

[jpienaar]

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

813

Nit: MLIR convention is to elide braces for single line conditionals.

817

Could you add a comment for this case?

mlir/test/Dialect/Tosa/constant-op-fold.mlir

507

Please fix.

[rsuderman]

Updated for jpienaar comments

[rsuderman]

missed newline

[rsuderman]

git-clang-format

[rsuderman]

Address jpienaar comments

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

def Tosa_ReverseOp: Tosa_Op<"reverse", [
   let arguments = (ins
     Tosa_Tensor1Dto4D:$input,
     I64Attr:$axis
   );
 
   let results = (outs
     Tosa_Tensor1Dto4D:$output
   );
+
+  let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: slice
 //===----------------------------------------------------------------------===//
 def Tosa_SliceOp: Tosa_Op<"slice", [
       DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
                               ["inferReturnTypeComponents"]>, NoSideEffect]> {

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

if (operands[1]) {
     if (densePad.isSplat() && densePad.getSplatValue<APInt>().isZero()) {
       return getInput1();
     }
   }
 
   return {};
 }
 
+OpFoldResult ReverseOp::fold(ArrayRef<Attribute> operands) {
+  auto operand = getInput();
+  auto operandTy = operand.getType().cast<ShapedType>();
+  auto axis = getAxis();
+  auto operandAttr = operands[0].dyn_cast_or_null<SplatElementsAttr>();
+  if (operandAttr) return operandAttr;
+
+  // If the dim-length is 1, tosa.reverse is a no-op.
+  if (operandTy.hasRank() && operandTy.getDimSize(axis) == 1) return operand;
+
+  return {};
+}
+
 OpFoldResult SliceOp::fold(ArrayRef<Attribute> operands) {
   auto inputTy = getInput().getType().dyn_cast<RankedTensorType>();
   auto outputTy = getType().dyn_cast<RankedTensorType>();
 
   if (!inputTy || !outputTy)
     return {};
 
   if (inputTy == outputTy && inputTy.hasStaticShape())

mlir/test/Dialect/Tosa/constant-op-fold.mlir

 // CHECK: func.func @cast_int_to_int_sign
 func.func @cast_int_to_int_sign() -> tensor<i32> {
   %splat = "tosa.const"() {value = dense<-1> : tensor<i16>} : () -> tensor<i16>
   // CHECK: %[[SPLAT:.+]] = "tosa.const"() {value = dense<-1> : tensor<i32>} : () -> tensor<i32>
   %cast = "tosa.cast"(%splat) : (tensor<i16>) -> tensor<i32>
   // CHECK: return %[[SPLAT]]
   return %cast : tensor<i32>
 }
+
+// -----
+
+// CHECK-LABEL: @reverse_splat
+func.func @reverse_splat() -> tensor<10xi32> {
+  // CHECK: %[[SPLAT:.+]] = "tosa.const"() {value = dense<42> : tensor<10xi32>}
+  %splat = "tosa.const"() {value = dense<42> : tensor<10xi32>} : () -> tensor<10xi32>
+  %reverse = "tosa.reverse"(%splat) { axis = 0 : i64 } : (tensor<10xi32>) -> tensor<10xi32>
+  // CHECK: return %[[SPLAT]]
+  return %reverse : tensor<10xi32>
+}
+
+// -----
+
+// CHECK-LABEL: @reverse_length_one
+func.func @reverse_length_one(%arg0 : tensor<10x1xi32>) -> (tensor<10x1xi32>, tensor<10x1xi32>) {
+  %nofold = "tosa.reverse"(%arg0) { axis = 0 : i64 } : (tensor<10x1xi32>) -> tensor<10x1xi32>
+  %fold = "tosa.reverse"(%arg0) { axis = 1 : i64 } : (tensor<10x1xi32>) -> tensor<10x1xi32>
+  // CHECK: %[[NOFOLD:.+]] = "tosa.reverse"(%arg0) {axis = 0 : i64}
+  // CHECK: return %[[NOFOLD]], %arg0
+  return %nofold, %fold : tensor<10x1xi32>, tensor<10x1xi32>
+}