[MLIR] Add `to/from_extent_tensor` lowering to the standard dialect

[frgossen]

Summary

The operations to_extent_tensor and from_extent_tensor become no-ops when
lowered to the standard dialect.
This is possible with a lowering from shape.shape to tensor<?xindex>.

Depends On D81161

mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp

 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/Transforms/DialectConversion.h"
 
 using namespace mlir;
 
 namespace {
 
 /// Conversion patterns.
-class SizeToIndexOpConversion
-    : public OpConversionPattern<shape::SizeToIndexOp> {
+class FromExtentTensorOpConversion
+    : public OpConversionPattern<shape::FromExtentTensorOp> {
 public:
-  using OpConversionPattern<shape::SizeToIndexOp>::OpConversionPattern;
+  using OpConversionPattern<shape::FromExtentTensorOp>::OpConversionPattern;
 
   LogicalResult
-  matchAndRewrite(shape::SizeToIndexOp op, ArrayRef<Value> operands,
+  matchAndRewrite(shape::FromExtentTensorOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
-    shape::SizeToIndexOpOperandAdaptor transformed(operands);
-    rewriter.replaceOp(op.getOperation(), transformed.arg());
+    shape::FromExtentTensorOpOperandAdaptor transformed(operands);
+    rewriter.replaceOp(op.getOperation(), transformed.input());
     return success();
   }
 };
 
 class IndexToSizeOpConversion
     : public OpConversionPattern<shape::IndexToSizeOp> {
 public:
   using OpConversionPattern<shape::IndexToSizeOp>::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(shape::IndexToSizeOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     shape::IndexToSizeOpOperandAdaptor transformed(operands);
     rewriter.replaceOp(op.getOperation(), transformed.arg());
     return success();
   }
 };
 
+class SizeToIndexOpConversion
+    : public OpConversionPattern<shape::SizeToIndexOp> {
+public:
+  using OpConversionPattern<shape::SizeToIndexOp>::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(shape::SizeToIndexOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    shape::SizeToIndexOpOperandAdaptor transformed(operands);
+    rewriter.replaceOp(op.getOperation(), transformed.arg());
+    return success();
+  }
+};
+
+class ToExtentTensorOpConversion
+    : public OpConversionPattern<shape::ToExtentTensorOp> {
+public:
+  using OpConversionPattern<shape::ToExtentTensorOp>::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(shape::ToExtentTensorOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    shape::ToExtentTensorOpOperandAdaptor transformed(operands);
+    rewriter.replaceOp(op.getOperation(), transformed.input());
+    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));
     });
   }
 };
 

 };
 
 } // namespace
 
 void mlir::populateShapeToStandardConversionPatterns(
     OwningRewritePatternList &patterns, MLIRContext *ctx) {
   // clang-format off
   patterns.insert<
+      FromExtentTensorOpConversion,
       IndexToSizeOpConversion,
-      SizeToIndexOpConversion>(ctx);
+      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

 
 // 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 `to_extent_tensor` operation to no-op.
+// CHECK-LABEL: @to_extent_tensor
+// CHECK-SAME: (%[[SHAPE:.*]]: tensor<?xindex>) -> tensor<?xindex>
+func @to_extent_tensor(%shape : !shape.shape) -> tensor<?xindex> {
+  // CHECK-NEXT: return %[[SHAPE]] : tensor<?xindex>
+  %tensor = "shape.to_extent_tensor"(%shape) : (!shape.shape) -> tensor<?xindex>
+  return %tensor : tensor<?xindex>
+}
+
+// -----
+
+// Lower `from_extent_tensor` operation to no-op.
+// CHECK-LABEL: @from_extent_tensor
+// CHECK-SAME: (%[[TENSOR:.*]]: tensor<?xindex>) -> tensor<?xindex>
+func @from_extent_tensor(%tensor : tensor<?xindex>) -> !shape.shape {
+  // CHECK-NEXT: return %[[TENSOR]] : tensor<?xindex>
+  %shape = "shape.from_extent_tensor"(%tensor)
+      : (tensor<?xindex>) -> !shape.shape
+  return %shape : !shape.shape
+}