[mlir] [VectorOps] Use 'vector.flat_transpose' for 2-D 'vector.tranpose'

[aartbik]

Summary

Progressive lowering of vector.transpose into an operation that
is closer to an intrinsic, and thus the hardware ISA. Currently
under the common vector transform testing flag, as we prepare
deploying this transformation in the LLVM lowering pipeline.

[aartbik]

mlir/test/Dialect/Vector/vector-flat-transforms.mlir

6

sending this out for some early discussion; this hits a phase ordering issue

I propose to move the shapecast lowering a bit later, so we can fold them first and only lower them when they cannot be eliminated

[nicolasvasilache]

mlir/include/mlir/Dialect/Vector/VectorOps.h

67

Can we use a struct VectorTransposeLowering to keep it consistent with the one above (even if for now there are really only 2 options atm) ?

mlir/test/Dialect/Vector/vector-flat-transforms.mlir

6

I think this should be resolved in a separate revision and this is fine for now.

Note that this is more an order of visitation problem. Given:

%a = shape_cast %0
%b = shape_cast %a

where the shape casts can fold, if %a is visited before %b then it will be expanded.

I have seen this type of behavior a bunch of times in different place (albeit not involving folding + canonicalization + lowering IIRC).

Seems like ShapeCast should have a canonicalizer / canonicalization pattern (hasCanonicalizer=1) with a separate match and rewrite.
Then ShapeCastLowering could query that on all its uses and fail to lower if any foldable use is left.

In other words, this type of pattern ordering can be resolved by finer-grained case disjunction.
However this seems like it can still "miss folding at a distance": consider a chain of transposes that are lowered in some arbitrary order introducing reshapes.
Folding opportunities would only appear if consecutive transposes are lowered before any newly introduced shape_cast is visited.
This seems like the worklist-based algorithm would handle this ordering naturally but I imagine we can construct more intricate cases where this would not be true?

Pinging @rriddle to see if there are more idiomatic ways of doing this, if this should be integrated in the rewriter itself (i.e. delay pattern application if any operand has folding opportunities), or something else.

[aartbik]

made option an enum

[aartbik]

mlir/include/mlir/Dialect/Vector/VectorOps.h

67

made this an enum

mlir/test/Dialect/Vector/vector-flat-transforms.mlir

6

yes, sounds good fixing this later somehow; the TODO reminds us to follow up :-)

[ftynse]

mlir/test/Dialect/Vector/vector-flat-transforms.mlir

6

Canonicalizer doesn't run in any pattern rewriting (it can in the Nicolas's multi-level driver if you configure it that way), folding is run in both greedy rewriter and dialect converter. However, folding does not recurse upwards on operands of the given op, which it seems what you need here. In dialect conversion, there's a TODO comment related to potentially folding any operation it visits even if it is considered legal. Maybe that will help.

[aartbik]

rebased

mlir/include/mlir/Dialect/Vector/VectorOps.h

 enum class VectorContractLowering {
   /// Progressively lower to finer grained `vector.contract` and `vector.fma`.
   FMA = 0,
   /// Lower to `vector.matrix_multiply`, maps 1-1 to LLVM matrix intrinsics.
   Matmul = 1,
   /// Lower to `vector.outerproduct`.
   OuterProduct = 2,
 };
+/// Enum to control the lowering of `vector.transpose` operations.
+enum class VectorTransposeLowering {
+  // Lower transpose into element-wise extract and inserts.
+  EltWise = 0,
+  /// Lower 2-D transpose to `vector.flat_transpose`, maps 1-1 to LLVM matrix
+  /// intrinsics.
+  Flat = 1,
+};
 /// Structure to control the behavior of vector transform patterns.
 struct VectorTransformsOptions {
   VectorContractLowering vectorContractLowering = VectorContractLowering::FMA;
+  VectorTransposeLowering vectorTransposeLowering =
+      VectorTransposeLowering::EltWise;
   VectorTransformsOptions &
   setVectorTransformsOptions(VectorContractLowering opt) {
     vectorContractLowering = opt;
     return *this;
   }
 };
 
 /// Collect a set of transformation patterns that are related to contracting

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

let extraClassDeclaration = [{
     VectorType getSourceVectorType() {
       return source().getType().cast<VectorType>();
     }
     VectorType getResultVectorType() {
       return getResult().getType().cast<VectorType>();
     }
   }];
   let assemblyFormat = "$source attr-dict `:` type($source) `to` type($result)";
+  let hasFolder = 1;
 }
 
 def Vector_TypeCastOp :
   Vector_Op<"type_cast", [NoSideEffect]>,
     Arguments<(ins StaticShapeMemRefOf<[AnyType]>:$memref)>,
     Results<(outs AnyMemRef:$result)> {
   let summary = "type_cast op converts a scalar memref to a vector memref";
   let description = [{

mlir/lib/Dialect/Vector/VectorOps.cpp

for (unsigned i = 0, e = sourceTupleType.size(); i < e; ++i)
     if (failed(verifyVectorShapeCast(
             op, sourceTupleType.getType(i).cast<VectorType>(),
             resultTupleType.getType(i).cast<VectorType>())))
       return failure();
 
   return success();
 }
 
+OpFoldResult ShapeCastOp::fold(ArrayRef<Attribute> operands) {
+  // Nop shape cast.
+  if (source().getType() == result().getType())
+    return source();
+
+  // Canceling shape casts.
+  if (auto otherOp = source().getDefiningOp<ShapeCastOp>())
+    if (result().getType() == otherOp.source().getType())
+      return otherOp.source();
+
+  return {};
+}
+
 //===----------------------------------------------------------------------===//
 // TypeCastOp
 //===----------------------------------------------------------------------===//
 
 static SmallVector<int64_t, 8> extractShape(MemRefType memRefType) {
   auto vectorType = memRefType.getElementType().dyn_cast<VectorType>();
   SmallVector<int64_t, 8> res(memRefType.getShape().begin(),
                               memRefType.getShape().end());

mlir/lib/Dialect/Vector/VectorTransforms.cpp

 ///   %0 = vector.extract %y[0, 0]
 ///   %1 = vector.insert %0, %z [0, 0]
 ///   ..
 ///   %x = vector.insert .., .. [.., ..]
 class TransposeOpLowering : public OpRewritePattern<vector::TransposeOp> {
 public:
   using OpRewritePattern<vector::TransposeOp>::OpRewritePattern;
 
+  TransposeOpLowering(vector::VectorTransformsOptions vectorTransformsOptions,
+                      MLIRContext *context)
+      : OpRewritePattern<vector::TransposeOp>(context),
+        vectorTransformsOptions(vectorTransformsOptions) {}
+
   LogicalResult matchAndRewrite(vector::TransposeOp op,
                                 PatternRewriter &rewriter) const override {
     auto loc = op.getLoc();
 
     VectorType resType = op.getResultType();
 
     // Set up convenience transposition table.
     SmallVector<int64_t, 4> transp;
     for (auto attr : op.transp())
       transp.push_back(attr.cast<IntegerAttr>().getInt());
 
+    // Handle a true 2-D matrix transpose differently when requested.
+    if (vectorTransformsOptions.vectorTransposeLowering ==
+            vector::VectorTransposeLowering::Flat &&
+        resType.getRank() == 2 && transp[0] == 1 && transp[1] == 0) {
+      Type flattenedType =
+          VectorType::get(resType.getNumElements(), resType.getElementType());
+      auto matrix =
+          rewriter.create<vector::ShapeCastOp>(loc, flattenedType, op.vector());
+      auto rows = rewriter.getI32IntegerAttr(resType.getShape()[0]);
+      auto columns = rewriter.getI32IntegerAttr(resType.getShape()[1]);
+      Value trans = rewriter.create<vector::FlatTransposeOp>(
+          loc, flattenedType, matrix, rows, columns);
+      rewriter.replaceOpWithNewOp<vector::ShapeCastOp>(op, resType, trans);
+      return success();
+    }
+
     // Generate fully unrolled extract/insert ops.
     Value result = rewriter.create<ConstantOp>(loc, resType,
                                                rewriter.getZeroAttr(resType));
     SmallVector<int64_t, 4> lhs(transp.size(), 0);
     SmallVector<int64_t, 4> rhs(transp.size(), 0);
     rewriter.replaceOp(op, expandIndices(loc, resType, 0, transp, lhs, rhs,
                                          op.vector(), result, rewriter));
     return success();
Value expandIndices(Location loc, VectorType resType, int64_t pos,
     for (int64_t d = 0, e = resType.getDimSize(pos); d < e; ++d) {
       lhs[pos] = d;
       rhs[transp[pos]] = d;
       result = expandIndices(loc, resType, pos + 1, transp, lhs, rhs, input,
                              result, rewriter);
     }
     return result;
   }
+
+  /// Options to control the vector patterns.
+  vector::VectorTransformsOptions vectorTransformsOptions;
 };
 
 /// Progressive lowering of OuterProductOp.
 /// One:
 ///   %x = vector.outerproduct %lhs, %rhs, %acc
 /// is replaced by:
 ///   %z = zero-result
 ///   %0 = vector.extract %lhs[0]
void mlir::vector::populateVectorContractLoweringPatterns(
     OwningRewritePatternList &patterns, MLIRContext *context,
     VectorTransformsOptions parameters) {
   // clang-format off
   patterns.insert<BroadcastOpLowering,
                   CreateMaskOpLowering,
                   ConstantMaskOpLowering,
                   OuterProductOpLowering,
                   ShapeCastOp2DDownCastRewritePattern,
-                  ShapeCastOp2DUpCastRewritePattern,
-                  TransposeOpLowering>(context);
-  patterns.insert<ContractionOpLowering,
+                  ShapeCastOp2DUpCastRewritePattern>(context);
+  patterns.insert<TransposeOpLowering,
+                  ContractionOpLowering,
                   ContractionOpToMatmulOpLowering,
                   ContractionOpToOuterProductOpLowering>(parameters, context);
   // clang-format on
 }

mlir/test/Dialect/Vector/vector-contract-transforms.mlir

 // CHECK:      %[[T11:.*]] = vector.insert %[[T10]], %[[T9]] [2, 1] : f32 into vector<3x2xf32>
 // CHECK:      return %[[T11]] : vector<3x2xf32>
 
 func @transpose23(%arg0: vector<2x3xf32>) -> vector<3x2xf32> {
   %0 = vector.transpose %arg0, [1, 0] : vector<2x3xf32> to vector<3x2xf32>
   return %0 : vector<3x2xf32>
 }
 
+
+// CHECK-LABEL: func @nop_shape_cast
+// CHECK-SAME: %[[A:.*]]: vector<16xf32>
+// CHECK:      return %[[A]] : vector<16xf32>
+
+func @nop_shape_cast(%arg0: vector<16xf32>) -> vector<16xf32> {
+  %0 = vector.shape_cast %arg0 : vector<16xf32> to vector<16xf32>
+  return %0 : vector<16xf32>
+}
+
+// CHECK-LABEL: func @cancel_shape_cast
+// CHECK-SAME: %[[A:.*]]: vector<16xf32>
+// CHECK:      return %[[A]] : vector<16xf32>
+
+func @cancel_shape_cast(%arg0: vector<16xf32>) -> vector<16xf32> {
+  %0 = vector.shape_cast %arg0 : vector<16xf32> to vector<4x4xf32>
+  %1 = vector.shape_cast %0 : vector<4x4xf32> to vector<16xf32>
+  return %1 : vector<16xf32>
+}
+
 // Shape up and downcasts for 2-D vectors, for supporting conversion to
 // llvm.matrix operations
 // CHECK-LABEL: func @shape_casts
 func @shape_casts(%a: vector<2x2xf32>) -> (vector<4xf32>, vector<2x2xf32>) {
   // CHECK: %[[cst:.*]] = constant dense<0.000000e+00> : vector<4xf32>
   // CHECK: %[[cst22:.*]] = constant dense<0.000000e+00> : vector<2x2xf32>
   // CHECK: %[[ex0:.*]] = vector.extract %{{.*}}[0] : vector<2x2xf32>
   //

mlir/test/Dialect/Vector/vector-flat-transforms.mlir

+// RUN: mlir-opt %s -test-vector-contraction-conversion=vector-flat-transpose=1 | FileCheck %s --dump-input-on-failure
+
+// Tests for lowering 2-D vector.transpose into vector.flat_transpose.
+//
+// TODO(ajcbik,ntv): having ShapeCastOp2DDownCastRewritePattern and
+//                   ShapeCastOp2DUpCastRewritePattern too early in
+//                   the greedy rewriting patterns misses opportunities
+//                   to fold shape casts!
+
+// No shape cast folding expected.
+//
+// CHECK-LABEL: func @transpose44_44(
+// CHECK-SAME:  %[[A:.*]]: vector<4x4xf32>
+// CHECK:       %[[T0:.*]] = vector.extract %[[A]][0] : vector<4x4xf32>
+// CHECK:       %[[T8:.*]] = vector.flat_transpose %{{.*}} {columns = 4 : i32, rows = 4 : i32} : vector<16xf32> -> vector<16xf32>
+// CHECK:       %[[T9:.*]] = vector.extract_strided_slice %[[T8]] {offsets = [0], sizes = [4], strides = [1]} : vector<16xf32> to vector<4xf32>
+//
+func @transpose44_44(%arg0: vector<4x4xf32>) -> vector<4x4xf32> {
+  %0 = vector.transpose %arg0, [1, 0] : vector<4x4xf32> to vector<4x4xf32>
+  return %0 : vector<4x4xf32>
+}
+
+// Folds preceding shape cast as expected,
+// no following shape cast folding expected.
+//
+// CHECK-LABEL: func @transpose16_44(
+// CHECK-SAME:  %[[A:.*]]: vector<16xf32>
+// CHECK:       %[[T0:.*]] = vector.flat_transpose %[[A]] {columns = 4 : i32, rows = 4 : i32} : vector<16xf32> -> vector<16xf32>
+// CHECK:       %[[T1:.*]] = vector.extract_strided_slice %[[T0]] {offsets = [0], sizes = [4], strides = [1]} : vector<16xf32> to vector<4xf32>
+//
+func @transpose16_44(%arg0: vector<16xf32>) -> vector<4x4xf32> {
+  %0 = vector.shape_cast %arg0 : vector<16xf32> to vector<4x4xf32>
+  %1 = vector.transpose %0, [1, 0] : vector<4x4xf32> to vector<4x4xf32>
+  return %1 : vector<4x4xf32>
+}
+
+// No preceding shape cast folding expected,
+// but FAILS to fold following cast.
+//
+// CHECK-LABEL: func @transpose44_16(
+// CHECK-SAME:  %[[A:.*]]: vector<4x4xf32>
+// CHECK:       %[[T0:.*]] = vector.extract %[[A]][0] : vector<4x4xf32>
+// CHECK:       %[[T8:.*]] = vector.flat_transpose %{{.*}} {columns = 4 : i32, rows = 4 : i32} : vector<16xf32> -> vector<16xf32>
+func @transpose44_16(%arg0: vector<4x4xf32>) -> vector<16xf32> {
+  %0 = vector.transpose %arg0, [1, 0] : vector<4x4xf32> to vector<4x4xf32>
+  %1 = vector.shape_cast %0 : vector<4x4xf32> to vector<16xf32>
+  return %1 : vector<16xf32>
+}
+
+// Folds preceding shape cast as expected,
+// but FAILS to fold following cast.
+//
+// CHECK-LABEL: func @transpose16_16(
+// CHECK-SAME:  %[[A:.*]]: vector<16xf32>
+// CHECK:       %[[T0:.*]] = vector.flat_transpose %[[A]] {columns = 4 : i32, rows = 4 : i32} : vector<16xf32> -> vector<16xf32>
+//
+func @transpose16_16(%arg0: vector<16xf32>) -> vector<16xf32> {
+  %0 = vector.shape_cast %arg0 : vector<16xf32> to vector<4x4xf32>
+  %1 = vector.transpose %0, [1, 0] : vector<4x4xf32> to vector<4x4xf32>
+  %2 = vector.shape_cast %1 : vector<4x4xf32> to vector<16xf32>
+  return %2 : vector<16xf32>
+}

mlir/test/lib/Transforms/TestVectorTransforms.cpp

 };
 
 struct TestVectorContractionConversion
     : public PassWrapper<TestVectorContractionConversion, FunctionPass> {
   TestVectorContractionConversion() = default;
   TestVectorContractionConversion(const TestVectorContractionConversion &pass) {
   }
 
-  Option<bool> lowerToLLVMMatrixIntrinsics{
+  Option<bool> lowerToFlatMatrix{
       *this, "vector-lower-matrix-intrinsics",
       llvm::cl::desc("Lower vector.contract to llvm.intr.matrix.multiply"),
       llvm::cl::init(false)};
+  Option<bool> lowerToFlatTranspose{
+      *this, "vector-flat-transpose",
+      llvm::cl::desc("Lower 2-D vector.transpose to vector.flat_transpose"),
+      llvm::cl::init(false)};
   Option<bool> lowerToOuterProduct{
       *this, "vector-outerproduct",
       llvm::cl::desc("Lower vector.contract to vector.outerproduct"),
       llvm::cl::init(false)};
 
   void runOnFunction() override {
     OwningRewritePatternList patterns;
     if (lowerToOuterProduct) {
       VectorContractLowering lowering = VectorContractLowering::OuterProduct;
       VectorTransformsOptions options{lowering};
       patterns.insert<ContractionOpToOuterProductOpLowering>(options,
                                                              &getContext());
       applyPatternsAndFoldGreedily(getFunction(), patterns);
       return;
     }
 
-    VectorContractLowering lowering = VectorContractLowering::FMA;
-    if (lowerToLLVMMatrixIntrinsics)
-      lowering = VectorContractLowering::Matmul;
-    VectorTransformsOptions options{lowering};
+    VectorContractLowering contractLowering = VectorContractLowering::FMA;
+    if (lowerToFlatMatrix)
+      contractLowering = VectorContractLowering::Matmul;
+    VectorTransposeLowering transposeLowering =
+        VectorTransposeLowering::EltWise;
+    if (lowerToFlatTranspose)
+      transposeLowering = VectorTransposeLowering::Flat;
+    VectorTransformsOptions options{contractLowering, transposeLowering};
     populateVectorContractLoweringPatterns(patterns, &getContext(), options);
     applyPatternsAndFoldGreedily(getFunction(), patterns);
   }
 };
 
 } // end anonymous namespace
 
 namespace mlir {