[mlir][Linalg] Make codegen strategy late transformations opt-in

[nicolasvasilache]

Summary

Making the late transformations opt-in results in less surprising behavior when composing multiple calls to the codegen strategy.

[nicolasvasilache]

Don;t make the hoisting part optional, they are often desired.

[dcaballe]

Thanks for addressing this, Nicolas! LGTM!
I agree that LICM and the other hoisting passes are always beneficial but I wonder if they should run by default as many times as the number of codegen strategies. I think it could be a good idea to disable all the existing (and future) cleanup passes by default to avoid surprises and minimize compile time. WDYT?

[nicolasvasilache]

@dcaballe thanks!

This is by no means in any final shape or form, I am happy to evolve based on data.

To give more context, atm my first order bottleneck is being able to compose transformations easily and run autotuning at scale.
Once that first bottleneck is lifted, we will have have a good to statistical way to measure whether such improvements result in concrete gains.

In such a context, I think the metrics I'll be looking for are:

1. does a particular change increase / decrease the density of "good" strategies in the search space or do we lose important points and degrade the overall result (think of this as an analogy to the final accuracy in an ML model / mixture of models)
2. assuming 1. is not adversely impacted, does the change result in speed improvements and ideally faster time to solution (think of of this as the analogy to training speed in ML models)

Also, keep in mind that we are operating at the Linalg level here and the number of ops is generally quite limited since they have pretty big granularity: basically we are talking about ops that can fuse together with the caveat that combinations map/broadcast/transpose/elementwise can all be fused in a single Linalg op.
As we move towards mixing with finer grained transformations, maybe at the loop and load/store level, I am sure these improvements will be more noticeable than they are today.

mlir/include/mlir/Dialect/Linalg/Transforms/CodegenStrategy.h

 
 private:
   std::string opName;
   linalg::LinalgVectorizationOptions options;
 };
 
 /// Options to control the application of late transformations.
 struct LateCodegenStrategyOptions {
+  /// Hoisting transformations are always deemed beneficial and must disabled
+  /// explicitly.
   bool enableLICM = true;
   bool enableHoistRedundantVectorTransfers = true;
   bool enableHoistRedundantVectorTransfersOnTensor = true;
-  bool enableVectorTransferPartialRewrite = true;
-  bool enableVectorContractLowering = true;
-  bool enableVectorToSCFConversion = true;
+  /// Vector lowering operations may result in surprising behavior when
+  /// composing multiple codegen strategies and must be enabled explicitly.
+  bool enableVectorTransferPartialRewrite = false;
+  bool enableVectorContractLowering = false;
+  bool enableVectorToSCFConversion = false;
 };
 
 /// Codegen strategy controls how a Linalg op is progressively lowered.
 /// The application uses a 3-level staged patterns strategy which allows
 /// ordering transformations by using the Linalg `applyStagedPatterns`
 /// function, where:
 ///   1. The first stage consists of the successive `tile`, `promote` and
 ///   `vectorize` patterns, applied sequentially.

mlir/test/lib/Dialect/Linalg/TestLinalgCodegenStrategy.cpp

strategy.tileIf<LinalgOp>(!tileSizes.empty(), anchorOpName, tilingOptions)
       .tileIf<LinalgOp>(!registerTileSizes.empty(), anchorOpName,
                         registerTilingOptions)
       .promoteIf<LinalgOp>(
           registerPromote, anchorOpName,
           LinalgPromotionOptions()
               .setAlignment(16)
               .setUseFullTileBuffersByDefault(registerPromoteFullTile))
       .vectorizeIf(vectorize, anchorOpName)
+      .setEnableVectorTransferPartialRewrite(true)
+      .setEnableVectorContractLowering(true)
+      .setEnableVectorToSCFConversion(true)
       .setVectorTransformsOptions(
           vector::VectorTransformsOptions()
               .setVectorTransformsOptions(vectorContractLowering)
               .setVectorTransferSplit(vectorTransferSplit))
       .setVectorTransferToSCFOptions(
           VectorTransferToSCFOptions().setUnroll(unrollVectorTransfers));
   strategy.transform(getFunction());
 }
strategy.tileIf<OpType>(!tileSizes.empty(), tilingOptions)
       .template tileIf<OpType>(!registerTileSizes.empty(),
                                registerTilingOptions)
       .template promoteIf<OpType>(
           registerPromote,
           LinalgPromotionOptions()
               .setAlignment(16)
               .setUseFullTileBuffersByDefault(registerPromoteFullTile))
       .template vectorizeIf<OpType>(vectorize)
+      .setEnableVectorTransferPartialRewrite(true)
+      .setEnableVectorContractLowering(true)
+      .setEnableVectorToSCFConversion(true)
       .setVectorTransformsOptions(
           vector::VectorTransformsOptions()
               .setVectorTransformsOptions(vectorContractLowering)
               .setVectorTransferSplit(vectorTransferSplit))
       .setVectorTransferToSCFOptions(
           VectorTransferToSCFOptions().setUnroll(unrollVectorTransfers));
   strategy.transform(getFunction());
 }