[mlir][vector][transform] Register vector dialect patterns

[springerm]

Summary

Depends On: D151983

[springerm]

rebase

[springerm]

This has landed as 4b48063b521bfcc9835269c729de79459d93229e.

mlir/include/mlir/Dialect/Transform/IR/TransformOps.h

 /// `RewritePatternSet`. Registered patterns can be applied with the
 /// "transform.apply_patterns" op.
 class PatternRegistry : public TransformDialectData<PatternRegistry> {
 public:
   PatternRegistry(MLIRContext *ctx) : TransformDialectData(ctx), builder(ctx) {}
 
   /// A function that populates a `RewritePatternSet`.
   using PopulatePatternsFn = std::function<void(RewritePatternSet &)>;
+  /// A function that populates a `RewritePatternSet` with a specified benefit.
+  using PopulatePatternsWithBenefitFn =
+      std::function<void(RewritePatternSet &, PatternBenefit)>;
 
   /// Registers patterns with the specified identifier. The identifier should
   /// be prefixed with the dialect to which the patterns belong.
   void registerPatterns(StringRef identifier, PopulatePatternsFn &&fn);
 
+  /// Registers patterns with the specified identifier. The identifier should
+  /// be prefixed with the dialect to which the patterns belong. The pattern
+  /// benefit is currently ignored.
+  void registerPatterns(StringRef identifier,
+                        PopulatePatternsWithBenefitFn &&fn);
+
 protected:
   friend class ApplyPatternsOp;
 
   /// Returns "true" if patterns are registered with the specified identifier.
   bool hasPatterns(StringAttr identifier) const;
 
   /// Populates the given pattern set with the specified patterns.
   void populatePatterns(StringAttr identifier,

mlir/lib/Dialect/Transform/IR/TransformOps.cpp

 
 void transform::PatternRegistry::registerPatterns(StringRef identifier,
                                                   PopulatePatternsFn &&fn) {
   StringAttr attr = builder.getStringAttr(identifier);
   assert(!patterns.contains(attr) && "patterns identifier is already in use");
   patterns.try_emplace(attr, std::move(fn));
 }
 
+void transform::PatternRegistry::registerPatterns(
+    StringRef identifier, PopulatePatternsWithBenefitFn &&fn) {
+  StringAttr attr = builder.getStringAttr(identifier);
+  assert(!patterns.contains(attr) && "patterns identifier is already in use");
+  patterns.try_emplace(attr, [f = move(fn)](RewritePatternSet &patternSet) {
+    f(patternSet, /*benefit=*/1);
+  });
+}
+
 void transform::PatternRegistry::populatePatterns(
     StringAttr identifier, RewritePatternSet &patternSet) const {
   auto it = patterns.find(identifier);
   assert(it != patterns.end() && "patterns not registered in registry");
   it->second(patternSet);
 }
 
 bool transform::PatternRegistry::hasPatterns(StringAttr identifier) const {

mlir/lib/Dialect/Vector/TransformOps/VectorTransformOps.cpp

 //===- VectorTransformOps.cpp - Implementation of Vector transform ops ----===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Vector/TransformOps/VectorTransformOps.h"
 #include "mlir/Conversion/VectorToSCF/VectorToSCF.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/Transform/IR/TransformDialect.h"
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
+#include "mlir/Dialect/Transform/IR/TransformOps.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/Dialect/Vector/Transforms/LoweringPatterns.h"
 #include "mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h"
 #include "mlir/Dialect/Vector/Transforms/VectorTransforms.h"
 #include "mlir/Dialect/X86Vector/Transforms.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
 using namespace mlir;

 public:
   VectorTransformDialectExtension() {
     declareGeneratedDialect<vector::VectorDialect>();
     declareGeneratedDialect<LLVM::LLVMDialect>();
     registerTransformOps<
 #define GET_OP_LIST
 #include "mlir/Dialect/Vector/TransformOps/VectorTransformOps.cpp.inc"
         >();
+
+    addDialectDataInitializer<transform::PatternRegistry>(
+        [&](transform::PatternRegistry &registry) {
+          registry.registerPatterns("vector.outer_product_lowering",
+                                    populateVectorOuterProductLoweringPatterns);
+          registry.registerPatterns("vector.broadcast_lowering",
+                                    populateVectorBroadcastLoweringPatterns);
+          registry.registerPatterns("vector.mask_op_lowering",
+                                    populateVectorMaskOpLoweringPatterns);
+          registry.registerPatterns("vector.shape_cast_lowering",
+                                    populateVectorShapeCastLoweringPatterns);
+          registry.registerPatterns(
+              "vector.transfer_lowering",
+              [&](RewritePatternSet &set, PatternBenefit benefit) {
+                return populateVectorTransferLoweringPatterns(
+                    set, /*maxTransferRank=*/std::nullopt, benefit);
+              });
+          registry.registerPatterns(
+              "vector.transfer_permutation_map_lowering",
+              populateVectorTransferPermutationMapLoweringPatterns);
+          registry.registerPatterns("vector.scan_lowering",
+                                    populateVectorScanLoweringPatterns);
+          registry.registerPatterns("vector.vector_gather_lowering",
+                                    populateVectorGatherLoweringPatterns);
+          registry.registerPatterns(
+              "vector.mask_lowering_for_side_effecting_ops",
+              populateVectorMaskLoweringPatternsForSideEffectingOps);
+        });
   }
 };
 } // namespace
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Vector/TransformOps/VectorTransformOps.cpp.inc"
 
 void mlir::vector::registerTransformDialectExtension(
     DialectRegistry &registry) {
   registry.addExtensions<VectorTransformDialectExtension>();
 }