From 7a9258e9bbf0f2a057fc894299d9a5a79d8c321d Mon Sep 17 00:00:00 2001
From: Alexander Belyaev <pifon@google.com>
Date: Tue, 16 Jun 2020 15:07:59 +0200
Subject: [PATCH] [mlir][shape] Add a func to populate ShapeToShape patterns.

Differential Revision: https://reviews.llvm.org/D81933
---
 mlir/include/mlir/Dialect/Shape/Transforms/Passes.h      | 5 +++++
 .../Dialect/Shape/Transforms/ShapeToShapeLowering.cpp    | 9 ++++++++-
 2 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/mlir/include/mlir/Dialect/Shape/Transforms/Passes.h b/mlir/include/mlir/Dialect/Shape/Transforms/Passes.h
index 29cf9d1b6715..7e6065341608 100644
--- a/mlir/include/mlir/Dialect/Shape/Transforms/Passes.h
+++ b/mlir/include/mlir/Dialect/Shape/Transforms/Passes.h
@@ -18,6 +18,8 @@
 
 namespace mlir {
 
+class MLIRContext;
+class OwningRewritePatternList;
 class Pass;
 
 /// Creates an instance of the ShapeToShapeLowering pass that legalizes Shape
@@ -25,6 +27,9 @@ class Pass;
 /// transformed to `shape.reduce`, which can be lowered to SCF and Standard.
 std::unique_ptr<Pass> createShapeToShapeLowering();
 
+/// Collects a set of patterns to rewrite ops within the Shape dialect.
+void populateShapeRewritePatterns(MLIRContext *context,
+                                  OwningRewritePatternList &patterns);
 } // end namespace mlir
 
 #endif // MLIR_DIALECT_SHAPE_TRANSFORMS_PASSES_H_
diff --git a/mlir/lib/Dialect/Shape/Transforms/ShapeToShapeLowering.cpp b/mlir/lib/Dialect/Shape/Transforms/ShapeToShapeLowering.cpp
index 1ba68a0a94ee..467f3d33ce23 100644
--- a/mlir/lib/Dialect/Shape/Transforms/ShapeToShapeLowering.cpp
+++ b/mlir/lib/Dialect/Shape/Transforms/ShapeToShapeLowering.cpp
@@ -54,8 +54,10 @@ struct ShapeToShapeLowering
 } // namespace
 
 void ShapeToShapeLowering::runOnFunction() {
+  MLIRContext &ctx = getContext();
+
   OwningRewritePatternList patterns;
-  patterns.insert<NumElementsOpConverter>(&getContext());
+  populateShapeRewritePatterns(&ctx, patterns);
 
   ConversionTarget target(getContext());
   target.addLegalDialect<ShapeDialect>();
@@ -64,6 +66,11 @@ void ShapeToShapeLowering::runOnFunction() {
     signalPassFailure();
 }
 
+void mlir::populateShapeRewritePatterns(MLIRContext *context,
+                                        OwningRewritePatternList &patterns) {
+  patterns.insert<NumElementsOpConverter>(context);
+}
+
 std::unique_ptr<Pass> mlir::createShapeToShapeLowering() {
   return std::make_unique<ShapeToShapeLowering>();
 }