[mlir][vector] Add an option to skip transpose lowering during contractLowering

[ThomasRaoux]

Summary

This is useful as it is sometimes better to try to canonicalize the transpose away after contract lowering rather than lowering it to a chain of scalar insert/extract.

[nicolasvasilache]

shouldn't we instead omit the transpose lowering pattern in places where don't want to apply it?

[ThomasRaoux]

In D101647#2729714, @nicolasvasilache wrote:

shouldn't we instead omit the transpose lowering pattern in places where don't want to apply it?

Do you mean adding a if() addPattern<TransposeLowering> in or split the populateVectorContractLoweringPatterns? Splitting it would make sense, I just need to make sure every user gets updated.

[ThomasRaoux]

Changed it to a NFC change moving transpose lowering in a separate populate function.

[ThomasRaoux]

In D101647#2729726, @ThomasRaoux wrote:

In D101647#2729714, @nicolasvasilache wrote:

shouldn't we instead omit the transpose lowering pattern in places where don't want to apply it?

Do you mean adding a if() addPattern<TransposeLowering> in or split the populateVectorContractLoweringPatterns? Splitting it would make sense, I just need to make sure every user gets updated.

I moved TransposeLowering in a different populate function.

[nicolasvasilache]

yes that is what I meant, thanks!

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

 };
 /// 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,
+  /// Do not lower transpose.
+  None = 2,
 };
 /// Enum to control the splitting of `vector.transfer` operations into
 /// in-bounds and out-of-bounds variants.
 enum class VectorTransferSplit {
   /// Do not split vector transfer operations.
   None = 0,
   /// Split using in-bounds + out-of-bounds vector.transfer operations.
   VectorTransfer = 1,

mlir/lib/Dialect/Vector/VectorTransforms.cpp

public:
 
   TransposeOpLowering(vector::VectorTransformsOptions vectorTransformsOptions,
                       MLIRContext *context)
       : OpRewritePattern<vector::TransposeOp>(context),
         vectorTransformsOptions(vectorTransformsOptions) {}
 
   LogicalResult matchAndRewrite(vector::TransposeOp op,
                                 PatternRewriter &rewriter) const override {
+    if (vectorTransformsOptions.vectorTransposeLowering ==
+        vector::VectorTransposeLowering::None)
+      return failure();
     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());