[MLIR][Shape] Lower `size_to_index` and `index_to_size` with declarative rules

[frgossen]

Summary

Replace implemented rewrite patterns with equivalent declarative rules.

Depends On D82022

[frgossen]

Add comment

mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp

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 IndexToSizeOpConversion : public OpConversionPattern<IndexToSizeOp> {
-public:
-  using OpConversionPattern<IndexToSizeOp>::OpConversionPattern;
-
-  LogicalResult
-  matchAndRewrite(IndexToSizeOp op, ArrayRef<Value> operands,
-                  ConversionPatternRewriter &rewriter) const override {
-    IndexToSizeOp::Adaptor transformed(operands);
-    rewriter.replaceOp(op.getOperation(), transformed.arg());
-    return success();
-  }
-};
-
-class SizeToIndexOpConversion : public OpConversionPattern<SizeToIndexOp> {
-public:
-  using OpConversionPattern<SizeToIndexOp>::OpConversionPattern;
-
-  LogicalResult
-  matchAndRewrite(SizeToIndexOp op, ArrayRef<Value> operands,
-                  ConversionPatternRewriter &rewriter) const override {
-    SizeToIndexOp::Adaptor transformed(operands);
-    rewriter.replaceOp(op.getOperation(), transformed.arg());
-    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(),

 
 void mlir::populateShapeToStandardConversionPatterns(
     OwningRewritePatternList &patterns, MLIRContext *ctx) {
   populateWithGenerated(ctx, &patterns);
   // clang-format off
   patterns.insert<
       BinaryOpConversion<AddOp, AddIOp>,
       BinaryOpConversion<MulOp, MulIOp>,
-      ConstSizeOpConverter,
-      IndexToSizeOpConversion,
-      SizeToIndexOpConversion>(ctx);
+      ConstSizeOpConverter>(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)>;
 
+// Convert `index_to_size` and `size_to_index` to no-ops as sizes will be
+// represented as indices.
+def IndexToSizeOpConversion : Pat<
+    (Shape_IndexToSizeOp $arg),
+    (replaceWithValue $arg)>;
+def SizeToIndexOpConversion : Pat<
+    (Shape_SizeToIndexOp $arg),
+    (replaceWithValue $arg)>;
+