[MLIR] Support memrefs with complex element types.

Differential Revision: https://reviews.llvm.org/D74307

Author: pifon2a

mlir/lib/IR/StandardTypes.cpp

@@ -333,7 +333,8 @@ MemRefType MemRefType::getImpl(ArrayRef<int64_t> shape, Type elementType,
   auto *context = elementType.getContext();
 
   // Check that memref is formed from allowed types.
-  if (!elementType.isIntOrFloat() && !elementType.isa<VectorType>())
+  if (!elementType.isIntOrFloat() && !elementType.isa<VectorType>() &&
+      !elementType.isa<ComplexType>())
     return emitOptionalError(location, "invalid memref element type"),
            MemRefType();
 
@@ -411,7 +412,8 @@ 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>())
+  if (!elementType.isIntOrFloat() && !elementType.isa<VectorType>() &&
+      !elementType.isa<ComplexType>())
     return emitOptionalError(*loc, "invalid memref element type");
   return success();
 }

mlir/test/IR/parser.mlir

@@ -133,6 +133,16 @@ func @memrefs_compose_with_id(memref<2x2xi8, affine_map<(d0, d1) -> (d0, d1)>,
 // CHECK: func @complex_types(complex<i1>) -> complex<f32>
 func @complex_types(complex<i1>) -> complex<f32>
 
+
+// CHECK: func @memref_with_complex_elems(memref<1x?xcomplex<f32>>)
+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 @functions((memref<1x?x4x?x?xi32, #map0>, memref<8xi8>) -> (), () -> ())
 func @functions((memref<1x?x4x?x?xi32, #map0, 0>, memref<8xi8, #map1, 0>) -> (), ()->())