[mlir] Centralize handling of memref element types.

[silvas]

Summary

This also beefs up the test coverage:

• Make unranked memref testing consistent with ranked memrefs.
• Add testing for the invalid element type cases.

This is not quite NFC: index types are now allowed in unranked memrefs.

mlir/include/mlir/IR/StandardTypes.h

 // BaseMemRefType
 //===----------------------------------------------------------------------===//
 
 /// Base MemRef for Ranked and Unranked variants
 class BaseMemRefType : public ShapedType {
 public:
   using ShapedType::ShapedType;
 
+  /// Return true if the specified element type is ok in a memref.
+  static bool isValidElementType(Type type) {
+    return type.isIntOrIndexOrFloat() || type.isa<VectorType, ComplexType>();
+  }
+
   /// Methods for support type inquiry through isa, cast, and dyn_cast.
   static bool classof(Type type);
 };
 
 //===----------------------------------------------------------------------===//
 // MemRefType
 
 /// MemRef types represent a region of memory that have a shape with a fixed

mlir/lib/IR/StandardTypes.cpp

 /// emit detailed error messages.  To emit errors when the location is unknown,
 /// pass in an instance of UnknownLoc.
 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.isIntOrIndexOrFloat() &&
-      !elementType.isa<VectorType, ComplexType>())
+  if (!BaseMemRefType::isValidElementType(elementType))
     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(Location loc, Type elementType,
                                                  unsigned memorySpace) {
-  // Check that memref is formed from allowed types.
-  if (!elementType.isIntOrFloat() &&
-      !elementType.isa<VectorType, ComplexType>())
+  if (!BaseMemRefType::isValidElementType(elementType))
     return emitError(loc, "invalid memref element type");
   return success();
 }
 
 // Fallback cases for terminal dim/sym/cst that are not part of a binary op (
 // i.e. single term). Accumulate the AffineExpr into the existing one.
 static void extractStridesFromTerm(AffineExpr e,
                                    AffineExpr multiplicativeFactor,

mlir/test/IR/invalid.mlir

 }
 
 // -----
 
 func @nestedtensor(tensor<tensor<i8>>) -> () // expected-error {{invalid tensor element type}}
 
 // -----
 
+func @illegalmemrefelementtype(memref<?xtensor<i8>>) -> () // expected-error {{invalid memref element type}}
+
+// -----
+
+func @illegalunrankedmemrefelementtype(memref<*xtensor<i8>>) -> () // expected-error {{invalid memref element type}}
+
+// -----
+
 func @indexvector(vector<4 x index>) -> () // expected-error {{vector elements must be int or float type}}
 
 // -----
 // Test no map in memref type.
 func @memrefs(memref<2x4xi8, >) // expected-error {{expected list element}}
 
 // -----
 // Test non-existent map in memref type.

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 @unranked_memref_with_index_elems(memref<*xindex>)
+func @unranked_memref_with_index_elems(memref<*xindex>)
+
+// CHECK: func @unranked_memref_with_vector_elems(memref<*xvector<10xf32>>)
+func @unranked_memref_with_vector_elems(memref<*xvector<10xf32>>)
+
 // 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)