[mlr][Linalg] NFC - Add option to hook vector.multi_reduction lowering to strategies.

[gysit]

LGTM!

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

LinalgVectorLoweringOptions &enableTransferPartialRewrite(bool val = true) {
     return *this;
   }
   /// Enable lowering of vector.contract.
   bool contractionLowering = false;
   LinalgVectorLoweringOptions &enableContractionLowering(bool val = true) {
     contractionLowering = val;
     return *this;
   }
+  /// Enable lowering of vector.multi_reduce.
+  bool multiReductionLowering = false;
+  LinalgVectorLoweringOptions &enableMultiReductionLowering(bool val = true) {
+    multiReductionLowering = val;
+    return *this;
+  }
   /// Enable lowering of vector.transfer to scf.
   bool transferToSCFConversion = false;
   LinalgVectorLoweringOptions &enableTransferToSCFConversion(bool val = true) {
     transferToSCFConversion = val;
     return *this;
   }
   /// Configure late vector transformations.
   vector::VectorTransformsOptions vectorTransformOptions;

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

 
 namespace vector {
 class VectorDialect;
 
 namespace detail {
 struct BitmaskEnumStorage;
 } // namespace detail
 
+/// Enum to control the lowering of `vector.contract` operations.
+enum class VectorContractLowering {
+  /// Progressively lower to finer grained `vector.contract` and dot-products.
+  Dot = 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.multi_reduction` operations.
+enum class VectorMultiReductionLowering {
+  /// Lower multi_reduction into outer-reduction and inner-parallel ops.
+  InnerParallel = 0,
+  /// Lower multi_reduction into outer-parallel and inner-reduction ops.
+  InnerReduction = 1,
+};
+/// 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,
+};
+/// 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,
+  /// Split using an in-bounds vector.transfer + linalg.fill + linalg.copy
+  /// operations.
+  LinalgCopy = 2,
+  /// Do not split vector transfer operation but instead mark it as "in-bounds".
+  ForceInBounds = 3
+};
+/// Structure to control the behavior of vector transform patterns.
+struct VectorTransformsOptions {
+  /// Option to control the lowering of vector.contract.
+  VectorContractLowering vectorContractLowering = VectorContractLowering::Dot;
+  VectorTransformsOptions &
+  setVectorTransformsOptions(VectorContractLowering opt) {
+    vectorContractLowering = opt;
+    return *this;
+  }
+  /// Option to control the lowering of vector.multi_reduction.
+  VectorMultiReductionLowering vectorMultiReductionLowering =
+      VectorMultiReductionLowering::InnerParallel;
+  VectorTransformsOptions &
+  setVectorMultiReductionLowering(VectorMultiReductionLowering opt) {
+    vectorMultiReductionLowering = opt;
+    return *this;
+  }
+  /// Option to control the lowering of vector.transpose.
+  VectorTransposeLowering vectorTransposeLowering =
+      VectorTransposeLowering::EltWise;
+  VectorTransformsOptions &
+  setVectorTransposeLowering(VectorTransposeLowering opt) {
+    vectorTransposeLowering = opt;
+    return *this;
+  }
+  /// Option to control the splitting of vector transfers.
+  VectorTransferSplit vectorTransferSplit = VectorTransferSplit::None;
+  VectorTransformsOptions &setVectorTransferSplit(VectorTransferSplit opt) {
+    vectorTransferSplit = opt;
+    return *this;
+  }
+};
+
 /// Return whether `srcType` can be broadcast to `dstVectorType` under the
 /// semantics of the `vector.broadcast` op.
 enum class BroadcastableToResult {
   Success = 0,
   SourceRankHigher = 1,
   DimensionMismatch = 2,
   SourceTypeNotAVector = 3
 };

 /// obtain a sequence of extract + vector.reduction + insert. This can further
 /// lower to horizontal reduction ops.
 /// - OneDimMultiReductionToTwoDim: for cases that reduce to 1-D vector<k>
 /// reduction (and are thus missing either a parallel or a reduction), we lift
 /// them back up to 2-D with a simple vector.shape_cast to vector<1xk> so that
 /// the other patterns can kick in, thus fully exiting out of the
 /// vector.multi_reduction abstraction.
 void populateVectorMultiReductionLoweringPatterns(
-    RewritePatternSet &patterns, bool useInnerDimsForReduction = false);
+    RewritePatternSet &patterns,
+    VectorMultiReductionLowering options =
+        vector::VectorMultiReductionLowering::InnerParallel);
 
 /// Collect a set of patterns to propagate insert_map/extract_map in the ssa
 /// chain.
 void populatePropagateVectorDistributionPatterns(RewritePatternSet &patterns);
 
 /// An attribute that specifies the combining function for `vector.contract`,
 /// and `vector.reduction`.
 class CombiningKindAttr
     : public Attribute::AttrBase<CombiningKindAttr, Attribute,
                                  detail::BitmaskEnumStorage> {
 public:
   using Base::Base;
 
   static CombiningKindAttr get(CombiningKind kind, MLIRContext *context);
 
   CombiningKind getKind() const;
 
   void print(DialectAsmPrinter &p) const;
   static Attribute parse(DialectAsmParser &parser);
 };
 
-/// Enum to control the lowering of `vector.contract` operations.
-enum class VectorContractLowering {
-  /// Progressively lower to finer grained `vector.contract` and dot-products.
-  Dot = 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,
-};
-/// 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,
-  /// Split using an in-bounds vector.transfer + linalg.fill + linalg.copy
-  /// operations.
-  LinalgCopy = 2,
-  /// Do not split vector transfer operation but instead mark it as "in-bounds".
-  ForceInBounds = 3
-};
-/// Structure to control the behavior of vector transform patterns.
-struct VectorTransformsOptions {
-  /// Option to control the lowering of vector.contract.
-  VectorContractLowering vectorContractLowering = VectorContractLowering::Dot;
-  VectorTransformsOptions &
-  setVectorTransformsOptions(VectorContractLowering opt) {
-    vectorContractLowering = opt;
-    return *this;
-  }
-  /// Option to control the lowering of vector.transpose.
-  VectorTransposeLowering vectorTransposeLowering =
-      VectorTransposeLowering::EltWise;
-  VectorTransformsOptions &
-  setVectorTransposeLowering(VectorTransposeLowering opt) {
-    vectorTransposeLowering = opt;
-    return *this;
-  }
-  /// Option to control the splitting of vector transfers.
-  VectorTransferSplit vectorTransferSplit = VectorTransferSplit::None;
-  VectorTransformsOptions &setVectorTransferSplit(VectorTransferSplit opt) {
-    vectorTransferSplit = opt;
-    return *this;
-  }
-};
-
 /// Collects patterns to progressively lower vector.broadcast ops on high-D
 /// vectors to low-D vector ops.
 void populateVectorBroadcastLoweringPatterns(RewritePatternSet &patterns);
 
 /// Collects patterns to progressively lower vector contraction ops on high-D
 /// into low-D reduction and product ops.
 void populateVectorContractLoweringPatterns(
     RewritePatternSet &patterns,

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

struct LinalgStrategyLowerVectorsPass
 
   void runOnFunction() override {
     auto funcOp = getFunction();
     if (!anchorFuncName.empty() && funcOp.getName() != anchorFuncName)
       return;
 
     MLIRContext *context = funcOp.getContext();
     RewritePatternSet patterns(context);
+    vector::populateVectorToVectorCanonicalizationPatterns(patterns);
     if (options.transferLowering) {
       vector::populateVectorTransferLoweringPatterns(patterns,
                                                      options.maxTransferRank);
     }
     if (options.transferPartialRewrite) {
       patterns.add<vector::VectorTransferFullPartialRewriter>(
           context, options.vectorTransformOptions);
     }
     if (options.contractionLowering) {
       patterns.add<ContractionOpToOuterProductOpLowering,
                    ContractionOpToMatmulOpLowering, ContractionOpLowering>(
           options.vectorTransformOptions, context);
       vector::populateVectorTransferPermutationMapLoweringPatterns(patterns);
     }
+    if (options.multiReductionLowering) {
+      vector::populateVectorMultiReductionLoweringPatterns(
+          patterns,
+          options.vectorTransformOptions.vectorMultiReductionLowering);
+    }
     if (options.transferToSCFConversion) {
       populateVectorToSCFConversionPatterns(patterns,
                                             options.vectorTransferToSCFOptions);
     }
     (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
   }
 
   LinalgVectorLoweringOptions options;

mlir/lib/Dialect/Vector/VectorMultiDimReductionTransforms.cpp

 
 /// Converts vector.multi_reduction into inner-most/outer-most reduction form
 /// by using vector.transpose
 class InnerOuterDimReductionConversion
     : public OpRewritePattern<vector::MultiDimReductionOp> {
 public:
   using OpRewritePattern<vector::MultiDimReductionOp>::OpRewritePattern;
 
-  explicit InnerOuterDimReductionConversion(MLIRContext *context,
-                                            bool useInnerDimsForReduction)
+  explicit InnerOuterDimReductionConversion(
+      MLIRContext *context, vector::VectorMultiReductionLowering options)
       : mlir::OpRewritePattern<vector::MultiDimReductionOp>(context),
-        useInnerDimsForReduction(useInnerDimsForReduction) {}
+        useInnerDimsForReduction(
+            options == vector::VectorMultiReductionLowering::InnerReduction) {}
 
   LogicalResult matchAndRewrite(vector::MultiDimReductionOp multiReductionOp,
                                 PatternRewriter &rewriter) const override {
     auto src = multiReductionOp.source();
     auto loc = multiReductionOp.getLoc();
     auto srcRank = multiReductionOp.getSourceVectorType().getRank();
 
     // Separate reduction and parallel dims

 
 /// Reduces the rank of vector.multi_reduction nd -> 2d given all reduction
 /// dimensions are either inner most or outer most.
 class ReduceMultiDimReductionRank
     : public OpRewritePattern<vector::MultiDimReductionOp> {
 public:
   using OpRewritePattern<vector::MultiDimReductionOp>::OpRewritePattern;
 
-  explicit ReduceMultiDimReductionRank(MLIRContext *context,
-                                       bool useInnerDimsForReduction)
+  explicit ReduceMultiDimReductionRank(
+      MLIRContext *context, vector::VectorMultiReductionLowering options)
       : mlir::OpRewritePattern<vector::MultiDimReductionOp>(context),
-        useInnerDimsForReduction(useInnerDimsForReduction) {}
+        useInnerDimsForReduction(
+            options == vector::VectorMultiReductionLowering::InnerReduction) {}
 
   LogicalResult matchAndRewrite(vector::MultiDimReductionOp multiReductionOp,
                                 PatternRewriter &rewriter) const override {
     auto srcRank = multiReductionOp.getSourceVectorType().getRank();
     auto srcShape = multiReductionOp.getSourceVectorType().getShape();
     auto loc = multiReductionOp.getLoc();
 
     // If rank less than 2, nothing to do.
Value reduced = rewriter.create<vector::MultiDimReductionOp>(
         loc, cast, mask, multiReductionOp.kind());
     rewriter.replaceOpWithNewOp<vector::ExtractOp>(multiReductionOp, reduced,
                                                    ArrayRef<int64_t>{0});
     return success();
   }
 };
 
 void mlir::vector::populateVectorMultiReductionLoweringPatterns(
-    RewritePatternSet &patterns, bool useInnerDimsForReduction) {
-  patterns.add<InnerOuterDimReductionConversion, ReduceMultiDimReductionRank,
-               OneDimMultiReductionToTwoDim>(patterns.getContext(),
-                                             useInnerDimsForReduction);
-  if (useInnerDimsForReduction)
+    RewritePatternSet &patterns, VectorMultiReductionLowering options) {
+  patterns.add<InnerOuterDimReductionConversion, ReduceMultiDimReductionRank>(
+      patterns.getContext(), options);
+  patterns.add<OneDimMultiReductionToTwoDim>(patterns.getContext());
+  if (options == VectorMultiReductionLowering ::InnerReduction)
     patterns.add<TwoDimMultiReductionToReduction>(patterns.getContext());
   else
     patterns.add<TwoDimMultiReductionToElementWise>(patterns.getContext());
 }

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

(void)linalg::applyStagedPatterns(module, frozenStage1Patterns,
                                     std::move(stage2Patterns),
                                     stage3Transforms);
 
   //===--------------------------------------------------------------------===//
   // Post staged patterns transforms
   //===--------------------------------------------------------------------===//
 
   VectorTransformsOptions vectorTransformOptions{
-      VectorContractLowering::Dot, VectorTransposeLowering::EltWise};
+      VectorContractLowering::Dot, VectorMultiReductionLowering::InnerParallel,
+      VectorTransposeLowering::EltWise};
 
   RewritePatternSet vectorTransferPatterns(context);
   // Pattern is not applied because rank-reducing vector transfer is not yet
   // supported as can be seen in splitFullAndPartialTransferPrecondition,
   // VectorTransforms.cpp
   vectorTransferPatterns.add<VectorTransferFullPartialRewriter>(
       context, vectorTransformOptions);
   (void)applyPatternsAndFoldGreedily(module, std::move(vectorTransferPatterns));

mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp

if (lowerToFilterOuterProduct) {
       (void)applyPatternsAndFoldGreedily(getFunction(), std::move(patterns));
       return;
     }
 
     // Test on all contract lowering patterns.
     VectorContractLowering contractLowering = VectorContractLowering::Dot;
     if (lowerToFlatMatrix)
       contractLowering = VectorContractLowering::Matmul;
+    VectorMultiReductionLowering vectorMultiReductionLowering =
+        VectorMultiReductionLowering::InnerParallel;
     VectorTransposeLowering transposeLowering =
         VectorTransposeLowering::EltWise;
     if (lowerToFlatTranspose)
       transposeLowering = VectorTransposeLowering::Flat;
-    VectorTransformsOptions options{contractLowering, transposeLowering};
+    VectorTransformsOptions options{
+        contractLowering, vectorMultiReductionLowering, transposeLowering};
     populateVectorBroadcastLoweringPatterns(patterns);
     populateVectorContractLoweringPatterns(patterns, options);
     populateVectorMaskOpLoweringPatterns(patterns);
     populateVectorShapeCastLoweringPatterns(patterns);
     populateVectorTransposeLoweringPatterns(patterns, options);
     (void)applyPatternsAndFoldGreedily(getFunction(), std::move(patterns));
   }
 };
return "Test conversion patterns to lower vector.multi_reduction to other "
            "vector ops";
   }
   Option<bool> useOuterReductions{
       *this, "use-outer-reductions",
       llvm::cl::desc("Move reductions to outer most dimensions"),
       llvm::cl::init(false)};
   void runOnFunction() override {
     RewritePatternSet patterns(&getContext());
-    populateVectorMultiReductionLoweringPatterns(patterns, !useOuterReductions);
+    populateVectorMultiReductionLoweringPatterns(
+        patterns, useOuterReductions
+                      ? vector::VectorMultiReductionLowering::InnerParallel
+                      : vector::VectorMultiReductionLowering::InnerReduction);
     (void)applyPatternsAndFoldGreedily(getFunction(), std::move(patterns));
   }
 };
 
 struct TestVectorTransferCollapseInnerMostContiguousDims
     : public PassWrapper<TestVectorTransferCollapseInnerMostContiguousDims,
                          FunctionPass> {
   TestVectorTransferCollapseInnerMostContiguousDims() = default;