[MLIR][Shape] Realize `shape` to `std` lowering with declarative patterns

[frgossen]

Summary

Setup declarative rewrite rules to lower the shape dialect to the std
dialect with two exemplary rules for from/to_extent_tensor.

[pifon2a]

thanks!

mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp

22

you can also move it right before populate...() function

[bondhugula]

mlir/lib/Conversion/ShapeToStandard/ShapeToStandardPatterns.td

5

Comment here and below please.

[frgossen]

Address comments

mlir/lib/Conversion/ShapeToStandard/CMakeLists.txt

+set(LLVM_TARGET_DEFINITIONS ShapeToStandardPatterns.td)
+mlir_tablegen(ShapeToStandardPatterns.inc -gen-rewriters)
+add_public_tablegen_target(ShapeToStandardPatternsIncGen)
+
 add_mlir_conversion_library(MLIRShapeToStandard
   ShapeToStandard.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Conversion/ShapeToStandard
 
   DEPENDS
   MLIRConversionPassIncGen
+  ShapeToStandardPatternsIncGen
 
   LINK_COMPONENTS
   Core
 
   LINK_LIBS PUBLIC
   MLIREDSC
   MLIRIR
   MLIRShape
   MLIRPass
   MLIRSCF
   MLIRTransforms
   )

mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp

 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/Transforms/DialectConversion.h"
 
 using namespace mlir;
 using namespace mlir::shape;
 
 namespace {
 
+/// Generated conversion patterns.
+#include "ShapeToStandardPatterns.inc"
+
 /// Conversion patterns.
 template <typename SrcOpTy, typename DstOpTy>
 class BinaryOpConversion : public OpConversionPattern<SrcOpTy> {
 public:
   using OpConversionPattern<SrcOpTy>::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(SrcOpTy op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     typename SrcOpTy::Adaptor adaptor(operands);
     rewriter.replaceOpWithNewOp<DstOpTy>(op.getOperation(), adaptor.lhs(),
                                          adaptor.rhs());
     return success();
   }
 };
 
-class FromExtentTensorOpConversion
-    : public OpConversionPattern<FromExtentTensorOp> {
-public:
-  using OpConversionPattern<FromExtentTensorOp>::OpConversionPattern;
-
-  LogicalResult
-  matchAndRewrite(FromExtentTensorOp op, ArrayRef<Value> operands,
-                  ConversionPatternRewriter &rewriter) const override {
-    FromExtentTensorOp::Adaptor transformed(operands);
-    rewriter.replaceOp(op.getOperation(), transformed.input());
-    return success();
-  }
-};
-
 class IndexToSizeOpConversion : public OpConversionPattern<IndexToSizeOp> {
 public:
   using OpConversionPattern<IndexToSizeOp>::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(IndexToSizeOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     IndexToSizeOp::Adaptor transformed(operands);
public:
   matchAndRewrite(SizeToIndexOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     SizeToIndexOp::Adaptor transformed(operands);
     rewriter.replaceOp(op.getOperation(), transformed.arg());
     return success();
   }
 };
 
-class ToExtentTensorOpConversion
-    : public OpConversionPattern<ToExtentTensorOp> {
-public:
-  using OpConversionPattern<ToExtentTensorOp>::OpConversionPattern;
-
-  LogicalResult
-  matchAndRewrite(ToExtentTensorOp op, ArrayRef<Value> operands,
-                  ConversionPatternRewriter &rewriter) const override {
-    ToExtentTensorOp::Adaptor transformed(operands);
-    rewriter.replaceOp(op.getOperation(), transformed.input());
-    return success();
-  }
-};
-
 class ConstSizeOpConverter : public OpConversionPattern<ConstSizeOp> {
 public:
   using OpConversionPattern<ConstSizeOp>::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(ConstSizeOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     rewriter.replaceOpWithNewOp<ConstantIndexOp>(op.getOperation(),
addConversion([ctx](ShapeType type) {
                                    IndexType::get(ctx));
     });
   }
 };
 
 /// Conversion pass.
 class ConvertShapeToStandardPass
     : public ConvertShapeToStandardBase<ConvertShapeToStandardPass> {
+
   void runOnOperation() override {
     // Setup type conversion.
     MLIRContext &ctx = getContext();
     ShapeTypeConverter typeConverter(&ctx);
 
     // Setup target legality.
     ConversionTarget target(ctx);
     target.addLegalDialect<scf::SCFDialect, StandardOpsDialect>();
if (failed(applyFullConversion(module, target, patterns, &typeConverter)))
       signalPassFailure();
   }
 };
 
 } // namespace
 
 void mlir::populateShapeToStandardConversionPatterns(
     OwningRewritePatternList &patterns, MLIRContext *ctx) {
+  populateWithGenerated(ctx, &patterns);
   // clang-format off
   patterns.insert<
       BinaryOpConversion<AddOp, AddIOp>,
       BinaryOpConversion<MulOp, MulIOp>,
       ConstSizeOpConverter,
-      FromExtentTensorOpConversion,
       IndexToSizeOpConversion,
-      SizeToIndexOpConversion,
-      ToExtentTensorOpConversion>(ctx);
+      SizeToIndexOpConversion>(ctx);
   // clang-format on
 }
 
 std::unique_ptr<OperationPass<ModuleOp>>
 mlir::createConvertShapeToStandardPass() {
   return std::make_unique<ConvertShapeToStandardPass>();
 }

mlir/lib/Conversion/ShapeToStandard/ShapeToStandardPatterns.td

+include "mlir/Dialect/Shape/IR/ShapeOps.td"
+include "mlir/Dialect/StandardOps/IR/Ops.td"
+
+// Convert `from_extent_tensor` and `to_extent_tensor` to no-ops as shapes will
+// be represented as extent tensors.
+def FromExtentTensorOpConversion : Pat<
+    (Shape_FromExtentTensorOp $input),
+    (replaceWithValue $input)>;
+def ToExtentTensorOpConversion : Pat<
+    (Shape_ToExtentTensorOp $input),
+    (replaceWithValue $input)>;
+