[MLIR] Add type conversion for `shape.shape`

[frgossen]

Summary

Convert shape.shape to tensor<?xindex> when lowering the shape to the
std dialect.

[pifon2a]

mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp

59

why is it always dynamic?

[frgossen]

mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp

59

The actual shape is not dynamic but the shape type is really not holding that information at compile time.
I would not know how to concretise this here without extending the shape type itself.

[jpienaar]

mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp

59

Yes this is where the type/value boundary gets spanned (what is part of the type in tensor is part of the value in shape).

We could expand shape's type to hold rank (known, unknown, error) but starting with dynamic and having this be refined along with the lowering of the rest of the shape computations should work too.

mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp

 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 {

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

 // Lower `index_to_size` conversion to no-op.
 // CHECK-LABEL: @index_to_size
 // CHECK-SAME: (%[[INDEX:.*]]: index) -> index
 func @index_to_size(%index : index) -> !shape.size {
   // CHECK-NEXT: return %[[INDEX]] : index
   %size = shape.index_to_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
+}