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

[pifon2a]

Remove shape::.

[herhut]

Thanks!

[frgossen]

Thanks

mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp

public:
   matchAndRewrite(shape::ToExtentTensorOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     shape::ToExtentTensorOpOperandAdaptor transformed(operands);
     rewriter.replaceOp(op.getOperation(), transformed.input());
     return success();
   }
 };
 
+class ConstSizeOpConverter : public OpConversionPattern<shape::ConstSizeOp> {
+public:
+  using OpConversionPattern<shape::ConstSizeOp>::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(shape::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) {
       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