[mlir][memref][transform] Register memref dialect patterns

mlir/include/mlir/Dialect/MemRef/Transforms/Transforms.h

 //===- Transforms.h - MemRef Dialect transformations ------------*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
-/// This header declares functions that assit transformations in the MemRef
+/// This header declares functions that assist transformations in the MemRef
 /// dialect.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_MEMREF_TRANSFORMS_TRANSFORMS_H
 #define MLIR_DIALECT_MEMREF_TRANSFORMS_TRANSFORMS_H
 
 #include "mlir/Support/LogicalResult.h"

 void populateExpandOpsPatterns(RewritePatternSet &patterns);
 
 /// Appends patterns for folding memref aliasing ops into consumer load/store
 /// ops into `patterns`.
 void populateFoldMemRefAliasOpPatterns(RewritePatternSet &patterns);
 
 /// Appends patterns that resolve `memref.dim` operations with values that are
 /// defined by operations that implement the
-/// `ReifyRankedShapeTypeShapeOpInterface`, in terms of shapes of its input
+/// `ReifyRankedShapedTypeOpInterface`, in terms of shapes of its input
 /// operands.
-void populateResolveRankedShapeTypeResultDimsPatterns(
+void populateResolveRankedShapedTypeResultDimsPatterns(
     RewritePatternSet &patterns);
 
 /// Appends patterns that resolve `memref.dim` operations with values that are
 /// defined by operations that implement the `InferShapedTypeOpInterface`, in
 /// terms of shapes of its input operands.
 void populateResolveShapedTypeResultDimsPatterns(RewritePatternSet &patterns);
 
 /// Appends patterns for expanding memref operations that modify the metadata
 /// (sizes, offset, strides) of a memref into easier to analyze constructs.
 void populateExpandStridedMetadataPatterns(RewritePatternSet &patterns);
 
 /// Appends patterns for resolving `memref.extract_strided_metadata` into
 /// `memref.extract_strided_metadata` of its source.
 void populateResolveExtractStridedMetadataPatterns(RewritePatternSet &patterns);
 
 /// Appends patterns for emulating wide integer memref operations with ops over
 /// narrower integer types.
 void populateMemRefWideIntEmulationPatterns(
     arith::WideIntEmulationConverter &typeConverter,
     RewritePatternSet &patterns);
 
-/// Appends type converions for emulating wide integer memref operations with
+/// Appends type conversions for emulating wide integer memref operations with
 /// ops over narrowe integer types.
 void populateMemRefWideIntEmulationConversions(
     arith::WideIntEmulationConverter &typeConverter);
 
 /// Transformation to do multi-buffering/array expansion to remove dependencies
 /// on the temporary allocation between consecutive loop iterations.
 /// It returns the new allocation if the original allocation was multi-buffered
 /// and returns failure() otherwise.

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

patterns.add<FoldUnitDimLoops, RankReducedExtractSliceOp,
                RankReducedInsertSliceOp<tensor::InsertSliceOp>,
                RankReducedInsertSliceOp<tensor::ParallelInsertSliceOp>>(
       context);
   linalg::FillOp::getCanonicalizationPatterns(patterns, context);
   tensor::CollapseShapeOp::getCanonicalizationPatterns(patterns, context);
   tensor::EmptyOp::getCanonicalizationPatterns(patterns, context);
   tensor::ExpandShapeOp::getCanonicalizationPatterns(patterns, context);
   tensor::populateFoldTensorEmptyPatterns(patterns);
-  memref::populateResolveRankedShapeTypeResultDimsPatterns(patterns);
+  memref::populateResolveRankedShapedTypeResultDimsPatterns(patterns);
   memref::populateResolveShapedTypeResultDimsPatterns(patterns);
 }
 
 void mlir::linalg::populateFoldUnitExtentDimsViaSlicesPatterns(
     RewritePatternSet &patterns) {
   auto *context = patterns.getContext();
   patterns.add<ReplaceUnitExtents>(context,
                                    RankReductionStrategy::ExtractInsertSlice);
   patterns.add<FoldUnitDimLoops>(context);
   // TODO: Patterns unrelated to unit dim folding should be factored out.
   linalg::FillOp::getCanonicalizationPatterns(patterns, context);
   tensor::EmptyOp::getCanonicalizationPatterns(patterns, context);
   tensor::populateFoldTensorEmptyPatterns(patterns);
-  memref::populateResolveRankedShapeTypeResultDimsPatterns(patterns);
+  memref::populateResolveRankedShapedTypeResultDimsPatterns(patterns);
   memref::populateResolveShapedTypeResultDimsPatterns(patterns);
 }
 
 void mlir::linalg::populateMoveInitOperandsToInputPattern(
     RewritePatternSet &patterns) {
   patterns.add<MoveInitOperandsToInput>(patterns.getContext());
 }
 

mlir/lib/Dialect/MemRef/TransformOps/MemRefTransformOps.cpp

 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/MemRef/Transforms/Passes.h"
 #include "mlir/Dialect/MemRef/Transforms/Transforms.h"
 #include "mlir/Dialect/NVGPU/IR/NVGPUDialect.h"
 #include "mlir/Dialect/SCF/IR/SCF.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/Interfaces/LoopLikeInterface.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "llvm/Support/Debug.h"
 
 using namespace mlir;
 
 #define DEBUG_TYPE "memref-transforms"
void init() {
     declareGeneratedDialect<memref::MemRefDialect>();
     declareGeneratedDialect<nvgpu::NVGPUDialect>();
     declareGeneratedDialect<vector::VectorDialect>();
 
     registerTransformOps<
 #define GET_OP_LIST
 #include "mlir/Dialect/MemRef/TransformOps/MemRefTransformOps.cpp.inc"
         >();
+
+    addDialectDataInitializer<transform::PatternRegistry>(
+        [&](transform::PatternRegistry &registry) {
+          registry.registerPatterns("memref.expand_ops",
+                                    memref::populateExpandOpsPatterns);
+          registry.registerPatterns("memref.fold_memref_alias_ops",
+                                    memref::populateFoldMemRefAliasOpPatterns);
+          registry.registerPatterns(
+              "memref.resolve_ranked_shaped_type_result_dims",
+              memref::populateResolveRankedShapedTypeResultDimsPatterns);
+          registry.registerPatterns(
+              "memref.expand_strided_metadata",
+              memref::populateExpandStridedMetadataPatterns);
+          registry.registerPatterns(
+              "memref.extract_address_computations",
+              memref::populateExtractAddressComputationsPatterns);
+        });
   }
 };
 } // namespace
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/MemRef/TransformOps/MemRefTransformOps.cpp.inc"
 
 void mlir::memref::registerTransformDialectExtension(
     DialectRegistry &registry) {
   registry.addExtensions<MemRefTransformDialectExtension>();
 }

mlir/lib/Dialect/MemRef/Transforms/ResolveShapedTypeResultDims.cpp

 struct ResolveShapedTypeResultDimsPass final
     : public memref::impl::ResolveShapedTypeResultDimsBase<
           ResolveShapedTypeResultDimsPass> {
   void runOnOperation() override;
 };
 
 } // namespace
 
-void memref::populateResolveRankedShapeTypeResultDimsPatterns(
+void memref::populateResolveRankedShapedTypeResultDimsPatterns(
     RewritePatternSet &patterns) {
   patterns.add<DimOfReifyRankedShapedTypeOpInterface<memref::DimOp>,
                DimOfReifyRankedShapedTypeOpInterface<tensor::DimOp>>(
       patterns.getContext());
 }
 
 void memref::populateResolveShapedTypeResultDimsPatterns(
     RewritePatternSet &patterns) {
   // TODO: Move tensor::DimOp pattern to the Tensor dialect.
   patterns.add<DimOfShapedTypeOpInterface<memref::DimOp>,
                DimOfShapedTypeOpInterface<tensor::DimOp>>(
       patterns.getContext());
 }
 
 void ResolveRankedShapeTypeResultDimsPass::runOnOperation() {
   RewritePatternSet patterns(&getContext());
-  memref::populateResolveRankedShapeTypeResultDimsPatterns(patterns);
+  memref::populateResolveRankedShapedTypeResultDimsPatterns(patterns);
   if (failed(applyPatternsAndFoldGreedily(getOperation(), std::move(patterns))))
     return signalPassFailure();
 }
 
 void ResolveShapedTypeResultDimsPass::runOnOperation() {
   RewritePatternSet patterns(&getContext());
-  memref::populateResolveRankedShapeTypeResultDimsPatterns(patterns);
+  memref::populateResolveRankedShapedTypeResultDimsPatterns(patterns);
   memref::populateResolveShapedTypeResultDimsPatterns(patterns);
   if (failed(applyPatternsAndFoldGreedily(getOperation(), std::move(patterns))))
     return signalPassFailure();
 }
 
 std::unique_ptr<Pass> memref::createResolveShapedTypeResultDimsPass() {
   return std::make_unique<ResolveShapedTypeResultDimsPass>();
 }
 
 std::unique_ptr<Pass> memref::createResolveRankedShapeTypeResultDimsPass() {
   return std::make_unique<ResolveRankedShapeTypeResultDimsPass>();
 }