[mlir][tensor][transform] Register tensor dialect patterns

[springerm]

Summary

Depends On: D151983

[ftynse]

It would be nice to find a way to document these. Maybe update docs/Dialects/Transform.md to include a newly created docs/Dialects/TransformTensor.md, which itself contains a textual description of these pattern sets followed by include docs/Dialect/<generated_doc>.md.

[springerm]

In D151984#4390589, @ftynse wrote:

It would be nice to find a way to document these. Maybe update docs/Dialects/Transform.md to include a newly created docs/Dialects/TransformTensor.md, which itself contains a textual description of these pattern sets followed by include docs/Dialect/<generated_doc>.md.

I'm wondering if there's a way to place the documentation somewhere next to the implementation and auto-generate the .md. So that docs and code don't start to diverge... The registrations here could also be used for pattern tests. I.e., having a pass such as -test-registered-patterns="tensor.reassociative_reshape_folding" instead of Option<bool> testReassociativeReshapeFolding{ etc. in TestTensorTransforms.cpp.

mlir/include/mlir/Dialect/Tensor/Transforms/Transforms.h

 /// use cases can inherit from this pattern and add necessary controls.
 FailureOr<TilingResult> replaceExtractSliceWithTiledProducer(
     OpBuilder &builder, tensor::ExtractSliceOp sliceOp, OpResult producerOp);
 
 //===----------------------------------------------------------------------===//
 // Populate functions.
 //===----------------------------------------------------------------------===//
 
-/// Collects a set of patterns to rewrite ops within the tensor dialect.
-void populateExpandOpsPatterns(RewritePatternSet &patterns);
-
 /// Appends patterns for folding tensor aliasing ops into consumer load/store
 /// ops into `patterns`.
 void populateFoldTensorSubsetOpPatterns(RewritePatternSet &patterns);
 
 /// Collects patterns to merge consecutive tensor.insert_slice/extract_slice
 /// into one. These patterns are in in this separate entry point because the
 /// bufferization is sensitive over IR structure, particularly those
 /// tensor.extract_slice and tensor.insert_slice ops for creating the slices.

mlir/lib/Dialect/Tensor/TransformOps/TensorTransformOps.cpp

public:
   void init() {
     declareGeneratedDialect<affine::AffineDialect>();
     declareGeneratedDialect<tensor::TensorDialect>();
 
     registerTransformOps<
 #define GET_OP_LIST
 #include "mlir/Dialect/Tensor/TransformOps/TensorTransformOps.cpp.inc"
         >();
+
+    addDialectDataInitializer<transform::PatternRegistry>(
+        [&](transform::PatternRegistry &registry) {
+          registry.registerPatterns("tensor.fold_tensor_subset_ops",
+                                    tensor::populateFoldTensorSubsetOpPatterns);
+          registry.registerPatterns(
+              "tensor.merge_consecutive_insert_extract_slice",
+              tensor::populateMergeConsecutiveInsertExtractSlicePatterns);
+          registry.registerPatterns(
+              "tensor.drop_redundant_insert_slice_rank_expansion",
+              tensor::populateDropRedundantInsertSliceRankExpansionPatterns);
+          registry.registerPatterns(
+              "tensor.reassociative_reshape_folding",
+              tensor::populateReassociativeReshapeFoldingPatterns);
+          registry.registerPatterns("tensor.fold_tensor_empty",
+                                    tensor::populateFoldTensorEmptyPatterns);
+          registry.registerPatterns(
+              "tensor.fold_into_pack_and_unpack",
+              tensor::populateFoldIntoPackAndUnpackPatterns);
+        });
   }
 };
 } // namespace
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Tensor/TransformOps/TensorTransformOps.cpp.inc"
 
 void mlir::tensor::registerTransformDialectExtension(
     DialectRegistry &registry) {
   registry.addExtensions<TensorTransformDialectExtension>();
 }