[mlir][shape] Lower Shape `ConstSizeOp` to Standard `ConstantOp`.

[pifon2a]

Remove shape::.

[herhut]

Thanks!

[frgossen]

Thanks

mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp

matchAndRewrite(SrcOpTy op, ArrayRef<Value> operands,
     typename SrcOpTy::OperandAdaptor adaptor(operands);
     rewriter.replaceOpWithNewOp<DstOpTy>(op.getOperation(), adaptor.lhs(),
                                          adaptor.rhs());
     return success();
   }
 };
 
 class FromExtentTensorOpConversion
-    : public OpConversionPattern<shape::FromExtentTensorOp> {
+    : public OpConversionPattern<FromExtentTensorOp> {
 public:
-  using OpConversionPattern<shape::FromExtentTensorOp>::OpConversionPattern;
+  using OpConversionPattern<FromExtentTensorOp>::OpConversionPattern;
 
   LogicalResult
-  matchAndRewrite(shape::FromExtentTensorOp op, ArrayRef<Value> operands,
+  matchAndRewrite(FromExtentTensorOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
-    shape::FromExtentTensorOpOperandAdaptor transformed(operands);
+    FromExtentTensorOpOperandAdaptor transformed(operands);
     rewriter.replaceOp(op.getOperation(), transformed.input());
     return success();
   }
 };
 
-class IndexToSizeOpConversion
-    : public OpConversionPattern<shape::IndexToSizeOp> {
+class IndexToSizeOpConversion : public OpConversionPattern<IndexToSizeOp> {
 public:
-  using OpConversionPattern<shape::IndexToSizeOp>::OpConversionPattern;
+  using OpConversionPattern<IndexToSizeOp>::OpConversionPattern;
 
   LogicalResult
-  matchAndRewrite(shape::IndexToSizeOp op, ArrayRef<Value> operands,
+  matchAndRewrite(IndexToSizeOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
-    shape::IndexToSizeOpOperandAdaptor transformed(operands);
+    IndexToSizeOpOperandAdaptor transformed(operands);
     rewriter.replaceOp(op.getOperation(), transformed.arg());
     return success();
   }
 };
 
-class SizeToIndexOpConversion
-    : public OpConversionPattern<shape::SizeToIndexOp> {
+class SizeToIndexOpConversion : public OpConversionPattern<SizeToIndexOp> {
 public:
-  using OpConversionPattern<shape::SizeToIndexOp>::OpConversionPattern;
+  using OpConversionPattern<SizeToIndexOp>::OpConversionPattern;
 
   LogicalResult
-  matchAndRewrite(shape::SizeToIndexOp op, ArrayRef<Value> operands,
+  matchAndRewrite(SizeToIndexOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
-    shape::SizeToIndexOpOperandAdaptor transformed(operands);
+    SizeToIndexOpOperandAdaptor transformed(operands);
     rewriter.replaceOp(op.getOperation(), transformed.arg());
     return success();
   }
 };
 
 class ToExtentTensorOpConversion
-    : public OpConversionPattern<shape::ToExtentTensorOp> {
+    : public OpConversionPattern<ToExtentTensorOp> {
 public:
-  using OpConversionPattern<shape::ToExtentTensorOp>::OpConversionPattern;
+  using OpConversionPattern<ToExtentTensorOp>::OpConversionPattern;
 
   LogicalResult
-  matchAndRewrite(shape::ToExtentTensorOp op, ArrayRef<Value> operands,
+  matchAndRewrite(ToExtentTensorOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
-    shape::ToExtentTensorOpOperandAdaptor transformed(operands);
+    ToExtentTensorOpOperandAdaptor 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(),
+                                                 op.value().getSExtValue());
+    return success();
+  }
+};
+
 /// Type conversions.
 class ShapeTypeConverter : public TypeConverter {
 public:
   using TypeConverter::convertType;
 
   ShapeTypeConverter(MLIRContext *ctx) {
     // Add default pass-through conversion.
     addConversion([&](Type type) { return type; });
 
-    addConversion([ctx](shape::SizeType type) { return IndexType::get(ctx); });
-    addConversion([ctx](shape::ShapeType type) {
+    addConversion([ctx](SizeType type) { return IndexType::get(ctx); });
+    addConversion([ctx](ShapeType type) {
       return RankedTensorType::get({ShapedType::kDynamicSize},
                                    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>();
     target.addLegalOp<ModuleOp, ModuleTerminatorOp, ReturnOp>();

 } // namespace
 
 void mlir::populateShapeToStandardConversionPatterns(
     OwningRewritePatternList &patterns, MLIRContext *ctx) {
   // clang-format off
   patterns.insert<
       BinaryOpConversion<AddOp, AddIOp>,
       BinaryOpConversion<MulOp, MulIOp>,
+      ConstSizeOpConverter,
       FromExtentTensorOpConversion,
       IndexToSizeOpConversion,
       SizeToIndexOpConversion,
       ToExtentTensorOpConversion>(ctx);
   // clang-format on
 }
 
 std::unique_ptr<OperationPass<ModuleOp>>
 mlir::createConvertShapeToStandardPass() {
   return std::make_unique<ConvertShapeToStandardPass>();
 }

mlir/test/Conversion/ShapeToStandard/shape-to-standard.mlir

 // CHECK-SAME: (%[[LHS:.*]]: index, %[[RHS:.*]]: index)
 func @binary_ops(%lhs : !shape.size, %rhs : !shape.size) {
   %sum = "shape.add"(%lhs, %rhs) : (!shape.size, !shape.size) -> !shape.size
   // CHECK-NEXT: addi %[[LHS]], %[[RHS]] : index
   %product = shape.mul %lhs, %rhs
   // CHECK-NEXT: muli %[[LHS]], %[[RHS]] : index
   return
 }
+
+// -----
+
+// Convert `const_size` to `constant` op.
+// CHECK-LABEL: @size_const
+func @size_const() -> !shape.size {
+  %c1 = shape.const_size 1
+  return %c1 : !shape.size
+}
+// CHECK: %[[C1:.*]] = constant 1 : index
+// CHECK: return %[[C1]] : index