[MLIR][Shape] Remove type conversion from lowering to standard

[frgossen]

Summary

Operating on indices and extent tensors directly, the type conversion is no
longer needed for the supported cases.

Depends On D84441

mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp

 RankOpConverter::matchAndRewrite(shape::RankOp op, ArrayRef<Value> operands,
                                  ConversionPatternRewriter &rewriter) const {
   shape::RankOp::Adaptor transformed(operands);
   rewriter.replaceOpWithNewOp<DimOp>(op, transformed.shape(), 0);
   return success();
 }
 
 namespace {
-/// 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](SizeType type) { return IndexType::get(ctx); });
-    addConversion([ctx](ShapeType type) {
-      return RankedTensorType::get({ShapedType::kDynamicSize},
-                                   IndexType::get(ctx));
-    });
-  }
-};
-} // namespace
-
-namespace {
 /// Conversion pass.
 class ConvertShapeToStandardPass
     : public ConvertShapeToStandardBase<ConvertShapeToStandardPass> {
 
   void runOnOperation() override;
 };
 } // namespace
 
 void ConvertShapeToStandardPass::runOnOperation() {
-  // Setup type conversion.
-  MLIRContext &ctx = getContext();
-  ShapeTypeConverter typeConverter(&ctx);
-
   // Setup target legality.
+  MLIRContext &ctx = getContext();
   ConversionTarget target(ctx);
-  target.addLegalDialect<scf::SCFDialect, StandardOpsDialect>();
-  target.addLegalOp<ModuleOp, ModuleTerminatorOp, ReturnOp>();
-  target.addDynamicallyLegalOp<FuncOp>([&](FuncOp op) {
-    return typeConverter.isSignatureLegal(op.getType()) &&
-           typeConverter.isLegal(&op.getBody());
-  });
+  target.addLegalDialect<StandardOpsDialect>();
+  target.addLegalOp<FuncOp, ModuleOp, ModuleTerminatorOp>();
 
   // Setup conversion patterns.
   OwningRewritePatternList patterns;
   populateShapeToStandardConversionPatterns(patterns, &ctx);
-  populateFuncOpTypeConversionPattern(patterns, &ctx, typeConverter);
 
   // Apply conversion.
   auto module = getOperation();
   if (failed(applyFullConversion(module, target, patterns)))
     signalPassFailure();
 }
 
 void mlir::populateShapeToStandardConversionPatterns(

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

 // RUN: mlir-opt --split-input-file --convert-shape-to-std --verify-diagnostics %s | FileCheck %s
 
-// Convert `size` to `index` type.
-// CHECK-LABEL: @size_id
-// CHECK-SAME: (%[[SIZE:.*]]: index)
-func @size_id(%size : !shape.size) -> !shape.size {
-  // CHECK: return %[[SIZE]] : index
-  return %size : !shape.size
-}
-
-// -----
-
-// Convert `shape` to `tensor<?xindex>` type.
-// CHECK-LABEL: @shape_id
-// CHECK-SAME: (%[[SHAPE:.*]]: tensor<?xindex>)
-func @shape_id(%shape : !shape.shape) -> !shape.shape {
-  // CHECK: return %[[SHAPE]] : tensor<?xindex>
-  return %shape : !shape.shape
-}
-
-// -----
-
-// Lower binary ops.
-// CHECK-LABEL: @binary_ops
-// CHECK-SAME: (%[[LHS:.*]]: index, %[[RHS:.*]]: index)
-func @binary_ops(%lhs : !shape.size, %rhs : !shape.size) {
-  // CHECK: addi %[[LHS]], %[[RHS]] : index
-  %sum = "shape.add"(%lhs, %rhs) : (!shape.size, !shape.size) -> !shape.size
-  return
-}
-
-// -----
-
 // Lower binary ops.
 // CHECK-LABEL: @binary_ops
 // CHECK-SAME: (%[[LHS:.*]]: index, %[[RHS:.*]]: index)
 func @binary_ops(%lhs : index, %rhs : index) {
+  // CHECK: addi %[[LHS]], %[[RHS]] : index
+  %sum = shape.add %lhs, %rhs : index, index -> index
   // CHECK: muli %[[LHS]], %[[RHS]] : index
   %product = shape.mul %lhs, %rhs : index, index -> index
   return
 }
 
 // -----
 
 // Lower `shape_of` for statically shaped tensor.