[Mlir-commits] [mlir] 7e5d8a3 - [MLIR] Support memrefs with complex element types.

[Alexander Belyaev]

Author: Alexander Belyaev
Date: 2020-02-12T09:07:15+01:00
New Revision: 7e5d8a34e341a06a1b8ce221a36146163a00aaa3

URL: https://github.com/llvm/llvm-project/commit/7e5d8a34e341a06a1b8ce221a36146163a00aaa3
DIFF: https://github.com/llvm/llvm-project/commit/7e5d8a34e341a06a1b8ce221a36146163a00aaa3.diff

LOG: [MLIR] Support memrefs with complex element types.

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

Added: 
    

Modified: 
    mlir/lib/IR/StandardTypes.cpp
    mlir/test/IR/parser.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/IR/StandardTypes.cpp b/mlir/lib/IR/StandardTypes.cpp
index bd12cff65f7e..c4269af9030e 100644
--- a/mlir/lib/IR/StandardTypes.cpp
+++ b/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();
 }

diff  --git a/mlir/test/IR/parser.mlir b/mlir/test/IR/parser.mlir
index b2c93ce15291..71e82b0fe090 100644
--- a/mlir/test/IR/parser.mlir
+++ b/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>) -> (), ()->())