[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

 };
 
 /// Collect a set of transformation patterns that are related to contracting
 /// or expanding vector operations:
 ///   ContractionOpLowering,
 ///   ShapeCastOp2DDownCastRewritePattern,
 ///   ShapeCastOp2DUpCastRewritePattern
 ///   BroadcastOpLowering,
-///   TransposeOpLowering
 ///   OuterproductOpLowering
 /// These transformation express higher level vector ops in terms of more
 /// elementary extraction, insertion, reduction, product, and broadcast ops.
 void populateVectorContractLoweringPatterns(
     RewritePatternSet &patterns,
     VectorTransformsOptions vectorTransformOptions = VectorTransformsOptions());
 
+/// Insert TransposeLowering patterns into extraction/insertion.
+void populateVectorTransposeLoweringPatterns(
+    RewritePatternSet &patterns,
+    VectorTransformsOptions vectorTransformOptions = VectorTransformsOptions());
+
 /// Returns the integer type required for subscripts in the vector dialect.
 IntegerType getVectorSubscriptType(Builder &builder);
 
 /// Returns an integer array attribute containing the given values using
 /// the integer type required for subscripts in the vector dialect.
 ArrayAttr getVectorSubscriptAttr(Builder &b, ArrayRef<int64_t> values);
 
 namespace impl {

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVMPass.cpp

 void LowerVectorToLLVMPass::runOnOperation() {
   // Perform progressive lowering of operations on slices and
   // all contraction operations. Also applies folding and DCE.
   {
     RewritePatternSet patterns(&getContext());
     populateVectorToVectorCanonicalizationPatterns(patterns);
     populateVectorSlicesLoweringPatterns(patterns);
     populateVectorContractLoweringPatterns(patterns);
+    populateVectorTransposeLoweringPatterns(patterns);
     (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
   }
 
   // Convert to the LLVM IR dialect.
   LLVMTypeConverter converter(&getContext());
   RewritePatternSet patterns(&getContext());
   populateVectorMaskMaterializationPatterns(patterns, enableIndexOptimizations);
   populateVectorToLLVMMatrixConversionPatterns(converter, patterns);

mlir/lib/Dialect/Vector/VectorTransforms.cpp

void mlir::vector::populateVectorContractLoweringPatterns(
   // clang-format off
   patterns.add<BroadcastOpLowering,
                   CreateMaskOpLowering,
                   ConstantMaskOpLowering,
                   OuterProductOpLowering,
                   ShapeCastOp2DDownCastRewritePattern,
                   ShapeCastOp2DUpCastRewritePattern,
                   ShapeCastOpRewritePattern>(patterns.getContext());
-  patterns.add<TransposeOpLowering,
-                  ContractionOpLowering,
+  patterns.add<ContractionOpLowering,
                   ContractionOpToMatmulOpLowering,
                   ContractionOpToOuterProductOpLowering>(parameters, patterns.getContext());
   // clang-format on
 }
 
+void mlir::vector::populateVectorTransposeLoweringPatterns(
+    RewritePatternSet &patterns,
+    VectorTransformsOptions vectorTransformOptions) {
+  patterns.add<TransposeOpLowering>(vectorTransformOptions,
+                                    patterns.getContext());
+}
+
 void mlir::vector::populateVectorTransferLoweringPatterns(
     RewritePatternSet &patterns) {
   patterns
       .add<TransferReadToVectorLoadLowering, TransferWriteToVectorStoreLowering,
            TransferReadPermutationLowering, TransferOpReduceRank>(
           patterns.getContext());
 }

mlir/test/lib/Transforms/TestConvVectorization.cpp

void TestConvVectorization::runOnOperation() {
   pm.addNestedPass<FuncOp>(createConvertLinalgToLoopsPass());
   if (failed(pm.run(module)))
     llvm_unreachable("Unexpected failure in linalg to loops pass.");
 
   // Programmatic controlled lowering of vector.contract only.
   RewritePatternSet vectorContractLoweringPatterns(context);
   populateVectorContractLoweringPatterns(vectorContractLoweringPatterns,
                                          vectorTransformsOptions);
+  populateVectorTransposeLoweringPatterns(vectorContractLoweringPatterns,
+                                          vectorTransformsOptions);
   (void)applyPatternsAndFoldGreedily(module,
                                      std::move(vectorContractLoweringPatterns));
 
   // Programmatic controlled lowering of vector.transfer only.
   RewritePatternSet vectorToLoopsPatterns(context);
   populateVectorToSCFConversionPatterns(vectorToLoopsPatterns,
                                         VectorTransferToSCFOptions());
   (void)applyPatternsAndFoldGreedily(module, std::move(vectorToLoopsPatterns));

mlir/test/lib/Transforms/TestVectorTransforms.cpp

void runOnFunction() override {
     if (lowerToFlatMatrix)
       contractLowering = VectorContractLowering::Matmul;
     VectorTransposeLowering transposeLowering =
         VectorTransposeLowering::EltWise;
     if (lowerToFlatTranspose)
       transposeLowering = VectorTransposeLowering::Flat;
     VectorTransformsOptions options{contractLowering, transposeLowering};
     populateVectorContractLoweringPatterns(patterns, options);
+    populateVectorTransposeLoweringPatterns(patterns, options);
     (void)applyPatternsAndFoldGreedily(getFunction(), std::move(patterns));
   }
 };
 
 struct TestVectorUnrollingPatterns
     : public PassWrapper<TestVectorUnrollingPatterns, FunctionPass> {
   TestVectorUnrollingPatterns() = default;
   TestVectorUnrollingPatterns(const TestVectorUnrollingPatterns &pass) {}