[MLIR] Support memrefs with memref element types.

[kernhanda]

Summary

With this change, it is now possible to create a memref that's comprised of other memref types.

[kernhanda]

Perhaps memref should be agnostic to the element type it contains? Tensors and Vector types can remain more constrained, but memref should be able to act on many kinds of types, not just the basic fundamental ones.

[ftynse]

This is largely insufficient. What is the dynamic representation of such memrefs? How do we lower loads/stores? Do we add type some amount of type erasure (e.g. should it matter that the element type is memref<10xf32> or memref<?xf32> ?)

Don't get me wrong, I think it is important to properly nested pointers for, e.g., sparse storage. They key thing is "properly" though. We need to make sure we are working towards that instead of just relaxing the type constraints on standard types without considering the implications.

Perhaps memref should be agnostic to the element type it contains? Tensors and Vector types can remain more constrained, but memref should be able to act on many kinds of types, not just the basic fundamental ones.

Memref points to a buffer in memory. It does not make sense to have memrefs of a type that don't have a well-defined in-memory representation.

[kernhanda]

Is there any guidance on what other changes would be required to land this change? I'll write some tests in the meanwhile to ensure that the standard to llvm conversion works well, at least.

mlir/lib/IR/StandardTypes.cpp

 MemRefType MemRefType::getImpl(ArrayRef<int64_t> shape, Type elementType,
                                ArrayRef<AffineMap> affineMapComposition,
                                unsigned memorySpace,
                                Optional<Location> location) {
   auto *context = elementType.getContext();
 
   // Check that memref is formed from allowed types.
   if (!elementType.isIntOrFloat() && !elementType.isa<VectorType>() &&
-      !elementType.isa<ComplexType>())
+      !elementType.isa<ComplexType>() && !elementType.isa<BaseMemRefType>())
     return emitOptionalError(location, "invalid memref element type"),
            MemRefType();
 
   for (int64_t s : shape) {
     // Negative sizes are not allowed except for `-1` that means dynamic size.
     if (s < -1)
       return emitOptionalError(location, "invalid memref size"), MemRefType();
   }
unsigned UnrankedMemRefType::getMemorySpace() const {
   return getImpl()->memorySpace;
 }
 
 LogicalResult UnrankedMemRefType::verifyConstructionInvariants(
     Optional<Location> loc, MLIRContext *context, Type elementType,
     unsigned memorySpace) {
   // Check that memref is formed from allowed types.
   if (!elementType.isIntOrFloat() && !elementType.isa<VectorType>() &&
-      !elementType.isa<ComplexType>())
+      !elementType.isa<ComplexType>() && !elementType.isa<BaseMemRefType>())
     return emitOptionalError(*loc, "invalid memref element type");
   return success();
 }
 
 /// Given MemRef `sizes` that are either static or dynamic, returns the
 /// canonical "contiguous" strides AffineExpr. Strides are multiplicative and
 /// once a dynamic dimension is encountered, all canonical strides become
 /// dynamic and need to be encoded with a different symbol.

mlir/test/IR/parser.mlir

 func @memref_with_complex_elems(memref<1x?xcomplex<f32>>)
 
 // CHECK: func @memref_with_vector_elems(memref<1x?xvector<10xf32>>)
 func @memref_with_vector_elems(memref<1x?xvector<10xf32>>)
 
 // CHECK: func @unranked_memref_with_complex_elems(memref<*xcomplex<f32>>)
 func @unranked_memref_with_complex_elems(memref<*xcomplex<f32>>)
 
+// CHECK: func @memref_with_memref_elems(memref<1x?xmemref<f32>>)
+func @memref_with_memref_elems(memref<1x?xmemref<f32>>)
+
+// CHECK: func @unranked_memref_with_memref_elems(memref<*xmemref<f32>>)
+func @unranked_memref_with_memref_elems(memref<*xmemref<f32>>)
+
+// CHECK: func @functions_with_memref((memref<1x?x4x?x?xmemref<i32>, #map0>, memref<8xi8>) -> (), () -> ())
+func @functions_with_memref((memref<1x?x4x?x?xmemref<i32>, #map0, 0>, memref<8xi8, #map1, 0>) -> (), ()->())
+
 // CHECK: func @functions((memref<1x?x4x?x?xi32, #map0>, memref<8xi8>) -> (), () -> ())
 func @functions((memref<1x?x4x?x?xi32, #map0, 0>, memref<8xi8, #map1, 0>) -> (), ()->())
 
 // CHECK-LABEL: func @simpleCFG(%{{.*}}: i32, %{{.*}}: f32) -> i1 {
 func @simpleCFG(%arg0: i32, %f: f32) -> i1 {
   // CHECK: %{{.*}} = "foo"() : () -> i64
   %1 = "foo"() : ()->i64
   // CHECK: "bar"(%{{.*}}) : (i64) -> (i1, i1, i1)