[mlir] [VectorOps] Implement a simple folder for identity vector.transpose operations.

[grosul1]

Summary

[VectorOps] Implement a simple folder for identity vector.transpose operations.

[nicolasvasilache]

Is there a followup revision planned in which you would canonicalize transpose(transpose()) and have it compose with these newly added folders?

[rriddle]

mlir/lib/Dialect/Vector/VectorOps.cpp

1540

Please run clang-format.

[aartbik]

please label title with mlir as well

[mlir] [VectorOps]

mlir/include/mlir/Dialect/Vector/VectorOps.td

1318

note that the attribute is really called "transp" (perhaps not such a good name looking back), so getTransp or getTranspArray would be slightly more in style; but I am okay with this too and perhaps rename the array later

mlir/lib/Dialect/Vector/VectorOps.cpp

1535

wouldn't it read a little easier as

for (int64_t i = 0, n = permutation.size; i < n; i++)

if (permutation[i] != i) 
...

i.e. less casting?

mlir/test/Dialect/Vector/canonicalize.mlir

132

maybe for completeness of what you have right now also a combination of transpose-ID and transpose-non-ID to make sure only the latter remain?

[grosul1]

In D77088#1951080, @nicolasvasilache wrote:

Is there a followup revision planned in which you would canonicalize transpose(transpose()) and have it compose with these newly added folders?

Thanks, Nicolas, yes, I expect to make an update to eliminate combinations of transposes.

[grosul1]

removeme

[grosul1]

mlir/include/mlir/Dialect/Vector/VectorOps.td

1318

Done! Renamed everything to transp/Transp, no discrepancy for now. Everything can be renamed at once if/when someone gets to do it.

mlir/lib/Dialect/Vector/VectorOps.cpp

1535

Used "e = " as it seems to be a convention; in this file, at least.

[aartbik]

Good to go, but some last comments to consider in the follow up.

mlir/lib/Dialect/Vector/VectorOps.cpp

1536

Given the llvm style of leaving the {} out of trivial for/ifs (something I am not a big fan of), I am wondering if that actually also applies to this case, trivial for loop with trivial if, so no {} on for either; perhaps others will comment, otherwise just leave it

mlir/test/Dialect/Vector/canonicalize.mlir

155

this of course captures that the intermediate tranpose is gone, but to be completely consistent with your other tests, you would expect a CHECK-NOT here (to make sure it is also removed as dead code)

[grosul1]

Update the test

[grosul1]

mlir/lib/Dialect/Vector/VectorOps.cpp

1536

I can turn this loop (back) into a lambda.
Not quite the same case, but similar: https://llvm.org/docs/CodingStandards.html#turn-predicate-loops-into-predicate-functions.

[grosul1]

mlir/lib/Dialect/Vector/VectorOps.cpp

1536

Submitting as is.

[grosul1]

Refresh

mlir/include/mlir/Dialect/Vector/VectorOps.td

def Vector_TransposeOp :
   }];
   let extraClassDeclaration = [{
     VectorType getVectorType() {
       return vector().getType().cast<VectorType>();
     }
     VectorType getResultType() {
       return result().getType().cast<VectorType>();
     }
+    void getTransp(SmallVectorImpl<int64_t> &results);
   }];
   let assemblyFormat = [{
     $vector `,` $transp attr-dict `:` type($vector) `to` type($result)
   }];
+  let hasFolder = 1;
 }
 
 def Vector_TupleGetOp :
   Vector_Op<"tuple_get", [NoSideEffect]>,
     Arguments<(ins TupleOf<[AnyVector]>:$vectors, APIntAttr:$index)>,
     Results<(outs AnyVector)> {
   let summary = "vector tuple get operation";
   let description = [{

mlir/lib/Dialect/Vector/VectorOps.cpp

 //===- VectorOps.cpp - MLIR Super Vectorizer Operations -------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //

 }
 
 static LogicalResult verify(TupleOp op) { return success(); }
 
 //===----------------------------------------------------------------------===//
 // TransposeOp
 //===----------------------------------------------------------------------===//
 
+// Eliminates transpose operations, which produce values identical to their
+// input values. This happens when the dimensions of the input vector remain in
+// their original order after the transpose operation.
+OpFoldResult TransposeOp::fold(ArrayRef<Attribute> operands) {
+  SmallVector<int64_t, 4> transp;
+  getTransp(transp);
+
+  // Check if the permutation of the dimensions contains sequential values:
+  // {0, 1, 2, ...}.
+  for (int64_t i = 0, e = transp.size(); i < e; i++) {
+    if (transp[i] != i)
+      return {};
+  }
+
+  return vector();
+}
+
 static LogicalResult verify(TransposeOp op) {
   VectorType vectorType = op.getVectorType();
   VectorType resultType = op.getResultType();
   int64_t rank = resultType.getRank();
   if (vectorType.getRank() != rank)
     return op.emitOpError("vector result rank mismatch: ") << rank;
   // Verify transposition array.
   auto transpAttr = op.transp().getValue();
if (seen[i])
       return op.emitOpError("duplicate position index: ") << i;
     seen[i] = true;
     if (resultType.getDimSize(ta.index()) != vectorType.getDimSize(i))
       return op.emitOpError("dimension size mismatch at: ") << i;
   }
   return success();
 }
 
+void TransposeOp::getTransp(SmallVectorImpl<int64_t> &results) {
+  populateFromInt64AttrArray(transp(), results);
+}
+
 //===----------------------------------------------------------------------===//
 // TupleGetOp
 //===----------------------------------------------------------------------===//
 
 static ParseResult parseTupleGetOp(OpAsmParser &parser,
                                    OperationState &result) {
   OpAsmParser::OperandType operandInfo;
   IntegerAttr indexAttr;

mlir/test/Dialect/Vector/canonicalize.mlir

 func @strided_slice_of_constant_mask() -> (vector<2x1xi1>) {
   %0 = vector.constant_mask [2, 2] : vector<4x3xi1>
   %1 = vector.strided_slice %0
     {offsets = [1, 1], sizes = [2, 1], strides = [1, 1]}
       : vector<4x3xi1> to vector<2x1xi1>
   // CHECK: vector.constant_mask [1, 1] : vector<2x1xi1>
   return %1 : vector<2x1xi1>
 }
+
+// -----
+
+// CHECK-LABEL: transpose_1D_identity
+// CHECK-SAME: ([[ARG:%.*]]: vector<4xf32>)
+func @transpose_1D_identity(%arg : vector<4xf32>) -> vector<4xf32> {
+  // CHECK-NOT: transpose
+  %0 = vector.transpose %arg, [0] : vector<4xf32> to vector<4xf32>
+  // CHECK-NEXT: return [[ARG]]
+  return %0 : vector<4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: transpose_2D_identity
+// CHECK-SAME: ([[ARG:%.*]]: vector<4x3xf32>)
+func @transpose_2D_identity(%arg : vector<4x3xf32>) -> vector<4x3xf32> {
+  // CHECK-NOT: transpose
+  %0 = vector.transpose %arg, [0, 1] : vector<4x3xf32> to vector<4x3xf32>
+  // CHECK-NEXT: return [[ARG]]
+  return %0 : vector<4x3xf32>
+}
+
+// -----
+
+// CHECK-LABEL: transpose_3D_identity
+// CHECK-SAME: ([[ARG:%.*]]: vector<4x3x2xf32>)
+func @transpose_3D_identity(%arg : vector<4x3x2xf32>) -> vector<4x3x2xf32> {
+  // CHECK-NOT: transpose
+  %0 = vector.transpose %arg, [0, 1, 2] : vector<4x3x2xf32> to vector<4x3x2xf32>
+  // CHECK-NEXT: return [[ARG]]
+  return %0 : vector<4x3x2xf32>
+}
+
+// -----
+
+// CHECK-LABEL: transpose_2D_sequence
+// CHECK-SAME: ([[ARG:%.*]]: vector<4x3xf32>)
+func @transpose_2D_sequence(%arg : vector<4x3xf32>) -> vector<3x4xf32> {
+  // CHECK-NOT: transpose
+  %0 = vector.transpose %arg, [0, 1] : vector<4x3xf32> to vector<4x3xf32>
+  // CHECK: [[T1:%.*]] = vector.transpose [[ARG]], [1, 0]
+  %1 = vector.transpose %0, [1, 0] : vector<4x3xf32> to vector<3x4xf32>
+  // CHECK-NOT: transpose
+  %2 = vector.transpose %1, [0, 1] : vector<3x4xf32> to vector<3x4xf32>
+  // CHECK: [[ADD:%.*]] = addf [[T1]], [[T1]]
+  %4 = addf %1, %2 : vector<3x4xf32>
+  // CHECK-NEXT: return [[ADD]]
+  return %4 : vector<3x4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: transpose_3D_sequence
+// CHECK-SAME: ([[ARG:%.*]]: vector<4x3x2xf32>)
+func @transpose_3D_sequence(%arg : vector<4x3x2xf32>) -> vector<2x3x4xf32> {
+  // CHECK: [[T0:%.*]] = vector.transpose [[ARG]], [1, 2, 0]
+  %0 = vector.transpose %arg, [1, 2, 0] : vector<4x3x2xf32> to vector<3x2x4xf32>
+  // CHECK-NOT: transpose
+  %1 = vector.transpose %0, [0, 1, 2] : vector<3x2x4xf32> to vector<3x2x4xf32>
+  // CHECK: [[T2:%.*]] = vector.transpose [[T0]], [1, 0, 2]
+  %2 = vector.transpose %1, [1, 0, 2] : vector<3x2x4xf32> to vector<2x3x4xf32>
+  // CHECK: [[ADD:%.*]] = addf [[T2]], [[T2]]
+  %3 = addf %2, %2 : vector<2x3x4xf32>
+  // CHECK-NOT: transpose
+  %4 = vector.transpose %3, [0, 1, 2] : vector<2x3x4xf32> to vector<2x3x4xf32>
+  // CHECK-NEXT: return [[ADD]]
+  return %4 : vector<2x3x4xf32>
+}