[mlir][Linalg] NFC - Reorganize options nesting.

[nicolasvasilache]

Summary

This removes duplication and makes nesting more clear.
It also reduces the amount of changes necessary for exposing future options.

[mehdi_amini]

Looks like a nice improvement!

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

881

What about naming this enableVectorTransferPartialRewrite(bool val = true)?
The enable verb and the default value makes it nice to read at call sites I think:

.enableVectorTransferPartialRewrite()
.enableVectorContractLowering()
.enableVectorToSCFConversion();

Also since this class is also about VectorLowering options, could these be shortened?

.enableTransferPartialRewrite()
.enableContraction()
.enableSCFConversion();

[nicolasvasilache]

Address review.

[nicolasvasilache]

More renames

[nicolasvasilache]

Didn't mean to push yet, reopening.

[nicolasvasilache]

Create a vectorLowering transform, applied on-demand.

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

struct CodegenStrategy {
   /// Conditionally append a pattern to rewrite `LinalgOpType` as a vector
   /// operation.
   CodegenStrategy &
   vectorizeIf(bool b, StringRef opName,
               LinalgTransformationFilter::FilterFunction f = nullptr) {
     return b ? vectorize(opName, f) : *this;
     return *this;
   }
-  /// Configure the post staged-patterns late vector transformations.
+  /// Configure the post staged-patterns late vector lowering options.
   CodegenStrategy &
-  setVectorTransformsOptions(vector::VectorTransformsOptions options) {
-    vectorTransformOptions = options;
+  setLinalgVectorLoweringOptions(LinalgVectorLoweringOptions options) {
+    lateVectorLoweringOptions = options;
     return *this;
   }
-  /// Configure the post staged-patterns late vector.transfer to scf
-  /// conversion.
+  /// Configure the post staged-patterns global enabling passes options.
   CodegenStrategy &
-  setVectorTransferToSCFOptions(VectorTransferToSCFOptions options) {
-    vectorToSCFOptions = options;
-    return *this;
-  }
-  ///
-  /// Configure the application of late transformations.
-  ///
-  CodegenStrategy &setEnableLICM(bool val) {
-    this->lateCodegenStrategyOptions.enableLICM = val;
-    return *this;
-  }
-  CodegenStrategy &setEnableHoistRedundantVectorTransfers(bool val) {
-    this->lateCodegenStrategyOptions.enableHoistRedundantVectorTransfers = val;
-    return *this;
-  }
-  CodegenStrategy &setEnableHoistRedundantVectorTransfersOnTensor(bool val) {
-    this->lateCodegenStrategyOptions
-        .enableHoistRedundantVectorTransfersOnTensor = val;
-    return *this;
-  }
-  CodegenStrategy &setMaxTransferRank(int64_t val) {
-    this->lateCodegenStrategyOptions.maxTransferRank = val;
-    return *this;
-  }
-  CodegenStrategy &setEnableVectorTransferLowering(bool val) {
-    this->lateCodegenStrategyOptions.enableVectorTransferLowering = val;
-    return *this;
-  }
-  CodegenStrategy &setEnableVectorTransferPartialRewrite(bool val) {
-    this->lateCodegenStrategyOptions.enableVectorTransferPartialRewrite = val;
-    return *this;
-  }
-  CodegenStrategy &setEnableVectorContractLowering(bool val) {
-    this->lateCodegenStrategyOptions.enableVectorContractLowering = val;
-    return *this;
-  }
-  CodegenStrategy &setEnableVectorToSCFConversion(bool val) {
-    this->lateCodegenStrategyOptions.enableVectorToSCFConversion = val;
+  setVectorTransferToSCFOptions(LinalgEnablingOptions options) {
+    linalgEnablingOptions = options;
     return *this;
   }
 
   /// Apply the transformation patterns in sequence with cleanup
   /// transformations interleaved.
   void configurePassPipeline(OpPassManager &pm, MLIRContext *context) const;
 
 private:
   LogicalResult postPatternTransforms(Operation *func) const;
 
-  vector::VectorTransformsOptions vectorTransformOptions;
-  VectorTransferToSCFOptions vectorToSCFOptions;
+  LinalgEnablingOptions linalgEnablingOptions;
+  LinalgVectorLoweringOptions lateVectorLoweringOptions;
   SmallVector<std::unique_ptr<Transformation>, 4> transformationSequence;
-  LateCodegenStrategyOptions lateCodegenStrategyOptions;
 };
 
 } // namespace linalg
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LINALG_TRANSFORMS_CODEGENSTRATEGY_H_

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

struct LinalgVectorizationPattern : public LinalgBaseVectorizationPattern {
   LinalgVectorizationPattern(
       StringRef opName, MLIRContext *context,
       LinalgVectorizationOptions options = LinalgVectorizationOptions(),
       LinalgTransformationFilter filter = LinalgTransformationFilter(),
       PatternBenefit benefit = 1)
       : LinalgBaseVectorizationPattern(opName, context, filter, benefit) {}
 };
 
-/// 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;
-  /// Vector lowering operations may result in surprising behavior when
-  /// composing multiple codegen strategies and must be enabled explicitly.
-  int64_t maxTransferRank = 1;
-  bool enableVectorTransferLowering = true;
-  bool enableVectorTransferPartialRewrite = false;
-  bool enableVectorContractLowering = false;
-  bool enableVectorToSCFConversion = false;
-};
-
 /// Options to control the application of enabling transformations.
 /// Hoisting transformations are always deemed beneficial and must be disabled
 /// explicitly.
 struct LinalgEnablingOptions {
-  bool enableLICM = true;
-  bool enableHoistRedundantVectorTransfers = true;
-  bool enableHoistRedundantVectorTransfersOnTensor = true;
+  /// Enable LICM.
+  bool licm = true;
+  LinalgEnablingOptions &enableLICM(bool val) {
+    licm = val;
+    return *this;
+  }
+  /// Enable hoisting of redundant vector transfer ops.
+  bool hoistRedundantVectorTransfers = true;
+  LinalgEnablingOptions &enableHoistRedundantVectorTransfers(bool val) {
+    hoistRedundantVectorTransfers = val;
+    return *this;
+  }
+  /// Enable hoisting of redundant vector transfer ops on tensor.
+  bool hoistRedundantVectorTransfersOnTensor = true;
+  LinalgEnablingOptions &enableHoistRedundantVectorTransfersOnTensor(bool val) {
+    hoistRedundantVectorTransfersOnTensor = val;
+    return *this;
+  }
 };
 
 /// Vector lowering options control how ops are lowered down to 1-D and scf.for
 /// form.
 struct LinalgVectorLoweringOptions {
+  /// Maximal transfer rank under which we do not lower further.
   int64_t maxTransferRank = 1;
-  bool enableVectorTransferLowering = true;
-  bool enableVectorTransferPartialRewrite = false;
-  bool enableVectorContractLowering = false;
-  bool enableVectorToSCFConversion = false;
+  LinalgVectorLoweringOptions &setMaxTransferRank(int64_t val) {
+    maxTransferRank = val;
+    return *this;
+  }
+  /// Vector lowering operations may result in surprising behavior when
+  /// composing multiple codegen strategies and must be enabled explicitly.
+  bool transferLowering = true;
+  LinalgVectorLoweringOptions &enableTransferLowering(bool val) {
+    transferLowering = val;
+    return *this;
+  }
+  /// Trigger full / partial vector.transfer splits.
+  bool transferPartialRewrite = false;
+  LinalgVectorLoweringOptions &enableTransferPartialRewrite(bool val) {
+    transferPartialRewrite = val;
+    return *this;
+  }
+  /// Enable lowering of vector.contract.
+  bool contractionLowering = false;
+  LinalgVectorLoweringOptions &enableContractionLowering(bool val) {
+    contractionLowering = val;
+    return *this;
+  }
+  /// Enable lowering of vector.transfer to scf.
+  bool transferToSCFConversion = false;
+  LinalgVectorLoweringOptions &enableTransferToSCFConversion(bool val) {
+    transferToSCFConversion = val;
+    return *this;
+  }
+  /// Configure late vector transformations.
   vector::VectorTransformsOptions vectorTransformOptions;
+  LinalgVectorLoweringOptions &
+  setVectorTransformsOptions(vector::VectorTransformsOptions options) {
+    vectorTransformOptions = options;
+    return *this;
+  }
+  /// Configure the post staged-patterns late vector.transfer to scf
+  /// conversion.
   VectorTransferToSCFOptions vectorTransferToSCFOptions;
+  LinalgVectorLoweringOptions &
+  setVectorTransferToSCFOptions(VectorTransferToSCFOptions options) {
+    vectorTransferToSCFOptions = options;
+    return *this;
+  }
 };
 
 /// Trait to check if T provides a `getOperationName` method.
 template <typename T, typename... Args>
 using has_get_operation_name = decltype(T::getOperationName());
 template <typename T>
 using detect_has_get_operation_name =
     llvm::is_detected<has_get_operation_name, T>;

mlir/lib/Dialect/Linalg/Transforms/CodegenStrategy.cpp

for (unsigned stepCount = 0, e = transformationSequence.size(); stepCount < e;
     auto filter = (currentState.str() == std::to_string(0))
                       ? linalg::LinalgTransformationFilter(
                             t->filter, ArrayRef<Identifier>{}, nextState)
                       : linalg::LinalgTransformationFilter(
                             t->filter, currentState, nextState);
     t->addToPassPipeline(pm, filter);
     pm.addPass(createLinalgStrategyEnablePass());
   }
-  LinalgVectorLoweringOptions vectorLoweringOptions;
-  vectorLoweringOptions.maxTransferRank =
-      lateCodegenStrategyOptions.maxTransferRank;
-  vectorLoweringOptions.enableVectorTransferLowering =
-      lateCodegenStrategyOptions.enableVectorTransferLowering;
-  vectorLoweringOptions.enableVectorTransferPartialRewrite =
-      lateCodegenStrategyOptions.enableVectorTransferPartialRewrite;
-  vectorLoweringOptions.enableVectorContractLowering =
-      lateCodegenStrategyOptions.enableVectorContractLowering;
-  vectorLoweringOptions.enableVectorToSCFConversion =
-      lateCodegenStrategyOptions.enableVectorToSCFConversion;
-  vectorLoweringOptions.vectorTransformOptions = vectorTransformOptions;
-  vectorLoweringOptions.vectorTransferToSCFOptions = vectorToSCFOptions;
-  pm.addPass(createLinalgStrategyLowerVectorsPass(vectorLoweringOptions));
+  pm.addPass(createLinalgStrategyLowerVectorsPass(lateVectorLoweringOptions));
   pm.addPass(createLinalgStrategyRemoveMarkersPass());
 }

mlir/lib/Dialect/Linalg/Transforms/LinalgStrategyPasses.cpp

void runOnFunction() override {
 
     MLIRContext *context = funcOp.getContext();
     RewritePatternSet patterns =
         linalg::getLinalgTilingCanonicalizationPatterns(context);
     scf::populateSCFForLoopCanonicalizationPatterns(patterns);
     if (failed(applyPatternsAndFoldGreedily(funcOp, std::move(patterns))))
       return signalPassFailure();
 
-    if (options.enableLICM) {
+    if (options.licm) {
       if (funcOp
               ->walk([&](LoopLikeOpInterface loopLike) {
                 if (failed(moveLoopInvariantCode(loopLike)))
                   return WalkResult::interrupt();
                 return WalkResult::advance();
               })
               .wasInterrupted())
         return signalPassFailure();
     }
 
     promoteSingleIterationLoops(funcOp);
-    if (options.enableHoistRedundantVectorTransfers)
+    if (options.hoistRedundantVectorTransfers)
       hoistRedundantVectorTransfers(funcOp);
 
-    if (options.enableHoistRedundantVectorTransfersOnTensor)
+    if (options.hoistRedundantVectorTransfersOnTensor)
       hoistRedundantVectorTransfersOnTensor(funcOp);
   }
 
   LinalgEnablingOptions options;
   LinalgTransformationFilter filter;
 };
 
 /// Configurable pass to lower vector operations.
 struct LinalgStrategyLowerVectorsPass
     : public LinalgStrategyLowerVectorsPassBase<
           LinalgStrategyLowerVectorsPass> {
 
   LinalgStrategyLowerVectorsPass(LinalgVectorLoweringOptions opt,
                                  LinalgTransformationFilter filt)
       : options(opt), filter(filt) {}
 
   void runOnFunction() override {
     auto funcOp = getFunction();
     if (!anchorFuncName.empty() && funcOp.getName() != anchorFuncName)
       return;
 
     MLIRContext *context = funcOp.getContext();
     RewritePatternSet patterns(context);
-    if (options.enableVectorTransferLowering) {
+    if (options.transferLowering) {
       vector::populateVectorTransferLoweringPatterns(patterns,
                                                      options.maxTransferRank);
     }
-    if (options.enableVectorTransferPartialRewrite) {
+    if (options.transferPartialRewrite) {
       patterns.add<vector::VectorTransferFullPartialRewriter>(
           context, options.vectorTransformOptions);
     }
-    if (options.enableVectorContractLowering) {
+    if (options.contractionLowering) {
       patterns.add<ContractionOpToOuterProductOpLowering,
                    ContractionOpToMatmulOpLowering, ContractionOpLowering>(
           options.vectorTransformOptions, context);
       vector::populateVectorTransferPermutationMapLoweringPatterns(patterns);
     }
-    if (options.enableVectorToSCFConversion) {
+    if (options.transferToSCFConversion) {
       populateVectorToSCFConversionPatterns(patterns,
                                             options.vectorTransferToSCFOptions);
     }
     (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
   }
 
   LinalgVectorLoweringOptions options;
   LinalgTransformationFilter filter;
mlir::createLinalgStrategyLowerVectorsPass(LinalgVectorLoweringOptions opt,
   return std::make_unique<LinalgStrategyLowerVectorsPass>(opt, filter);
 }
 
 /// Create a LinalgStrategyRemoveMarkersPass.
 std::unique_ptr<OperationPass<FuncOp>>
 mlir::createLinalgStrategyRemoveMarkersPass() {
   return std::make_unique<LinalgStrategyRemoveMarkersPass>();
 }
 No newline at end of file

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

strategy.tileIf(!tileSizes.empty(), anchorOpName, tilingOptions)
       .tileIf(!registerTileSizes.empty(), anchorOpName, registerTilingOptions)
       .promoteIf(registerPromote, anchorOpName,
                  LinalgPromotionOptions()
                      .setAlignment(16)
                      .setUseFullTileBuffersByDefault(registerPromoteFullTile))
       .generalizeIf(generalize, anchorOpName)
       .interchangeIf(!iteratorInterchange.empty(), iteratorInterchange)
       .vectorizeIf(vectorize, generalize ? genericOpName : anchorOpName)
-      .setEnableVectorTransferPartialRewrite(true)
-      .setEnableVectorContractLowering(true)
-      .setEnableVectorToSCFConversion(true)
+      .setLinalgVectorLoweringOptions(
+          LinalgVectorLoweringOptions()
               .setVectorTransformsOptions(
                   vector::VectorTransformsOptions()
                       .setVectorTransformsOptions(vectorContractLowering)
                       .setVectorTransferSplit(vectorTransferSplit))
               .setVectorTransferToSCFOptions(
-          VectorTransferToSCFOptions().setUnroll(unrollVectorTransfers));
+                  VectorTransferToSCFOptions().setUnroll(unrollVectorTransfers))
+              .enableTransferPartialRewrite(true)
+              .enableContractionLowering(true)
+              .enableTransferToSCFConversion(true));
 
   // Created a nested OpPassManager and run.
   FuncOp funcOp = getFunction();
   OpPassManager dynamicPM("builtin.func");
   strategy.configurePassPipeline(dynamicPM, funcOp.getContext());
   if (failed(runPipeline(dynamicPM, funcOp)))
     return signalPassFailure();
 }