Uniformize MemRefType well-formedness checks.

Introduce a new public static member function, MemRefType::getChecked, intended
for the users that want detailed error messages to be emitted during MemRefType
construction and can gracefully handle these errors.  This function takes a
Location of the "MemRef" token if known.  The parser is one user of getChecked
that has location information, it outputs errors as compiler diagnostics.
Other users may pass in an instance of UnknownLoc and still have error messages
emitted.  Compiler-internal users not expecting the MemRefType construction to
fail should call MemRefType::get, which now aborts on failure with a generic
message.

Both "getChecked" and "get" call to a static free function that does actual
construction with well-formedness checks, optionally emits errors and returns
nullptr on failure.

The location information passed to getChecked has voluntarily coarse precision.
The error messages are intended for compiler engineers and do not justify
heavier API than a single location.  The text of the messages can be written so
that it pinpoints the actual location of the error within a MemRef declaration.

PiperOrigin-RevId: 219765902

mlir/include/mlir/IR/Types.h

@@ -24,9 +24,10 @@
 
 namespace mlir {
 class AffineMap;
-class MLIRContext;
-class IntegerType;
 class FloatType;
+class IntegerType;
+class Location;
+class MLIRContext;
 class OtherType;
 
 namespace detail {
@@ -413,11 +414,23 @@ public:
   /* implicit */ MemRefType(Type::ImplType *ptr);
 
   /// Get or create a new MemRefType based on shape, element type, affine
-  /// map composition, and memory space.
+  /// map composition, and memory space.  Assumes the arguments define a
+  /// well-formed MemRef type.  Use getChecked to gracefully handle MemRefType
+  /// construction failures.
   static MemRefType get(ArrayRef<int> shape, Type elementType,
                         ArrayRef<AffineMap> affineMapComposition,
                         unsigned memorySpace);
 
+  /// Get or create a new MemRefType based on shape, element type, affine
+  /// map composition, and memory space declared at the given location.
+  /// If the location is unknown, the last argument should be an instance of
+  /// UnknownLoc.  If the MemRefType defined by the arguments would be
+  /// ill-formed, emits errors (to the handler registered with the context or to
+  /// the error stream) and returns nullptr.
+  static MemRefType getChecked(ArrayRef<int> shape, Type elementType,
+                               ArrayRef<AffineMap> affineMapComposition,
+                               unsigned memorySpace, Location *location);
+
   unsigned getRank() const { return getShape().size(); }
 
   /// Returns an array of memref shape dimension sizes.
@@ -440,6 +453,11 @@ public:
   unsigned getNumDynamicDims() const;
 
   static bool kindof(Kind kind) { return kind == Kind::MemRef; }
+
+private:
+  static MemRefType getSafe(ArrayRef<int> shape, Type elementType,
+                            ArrayRef<AffineMap> affineMapComposition,
+                            unsigned memorySpace, Location *location);
 };
 
 // Make Type hashable.

mlir/lib/IR/MLIRContext.cpp

@@ -33,7 +33,6 @@
 #include "mlir/IR/Types.h"
 #include "mlir/Support/MathExtras.h"
 #include "mlir/Support/STLExtras.h"
-#include "third_party/llvm/llvm/include/llvm/ADT/STLExtras.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/Support/Allocator.h"
@@ -432,10 +431,9 @@ void MLIRContext::emitDiagnostic(Location *location, const llvm::Twine &message,
 
   os << "error: ";
 
-  // The default behavior for errors is to emit them to stderr and exit.
+  // The default behavior for errors is to emit them to stderr.
   os << message.str() << '\n';
   os.flush();
-  exit(1);
 }
 
 //===----------------------------------------------------------------------===//
@@ -725,12 +723,47 @@ UnrankedTensorType UnrankedTensorType::get(Type elementType) {
   return result;
 }
 
-MemRefType MemRefType::get(ArrayRef<int> shape, Type elementType,
-                           ArrayRef<AffineMap> affineMapComposition,
-                           unsigned memorySpace) {
+/// Get or create a new MemRefType defined by the arguments.  If the resulting
+/// type would be ill-formed, return nullptr.  If the location is provided,
+/// i.e. is not nullptr, emit detailed error messages.  To emit errors when
+/// the location is unknown, pass in an instance of UnknownLoc.
+static MemRefType getMemRefType(ArrayRef<int> shape, Type elementType,
+                                ArrayRef<AffineMap> affineMapComposition,
+                                unsigned memorySpace, Location *location) {
   auto *context = elementType.getContext();
   auto &impl = context->getImpl();
 
+  // Check that memref is formed from allowed types.
+  if (!elementType.isa<IntegerType>() && !elementType.isa<FloatType>() &&
+      !elementType.isa<VectorType>()) {
+    if (location)
+      context->emitDiagnostic(location, "invalid memref element type",
+                              MLIRContext::DiagnosticKind::Error);
+    return nullptr;
+  }
+
+  // Check that the structure of the composition is valid, i.e. that each
+  // subsequent affine map has as many inputs as the previous map has results.
+  // Take the dimensionality of the MemRef for the first map.
+  auto dim = shape.size();
+  unsigned i = 0;
+  for (const auto &affineMap : affineMapComposition) {
+    if (affineMap.getNumDims() != dim) {
+      if (location)
+        context->emitDiagnostic(
+            location,
+            "memref affine map dimension mismatch between " +
+                (i == 0 ? Twine("memref rank") : "affine map " + Twine(i)) +
+                " and affine map" + Twine(i + 1) + ": " + Twine(dim) +
+                " != " + Twine(affineMap.getNumDims()),
+            MLIRContext::DiagnosticKind::Error);
+      return nullptr;
+    }
+
+    dim = affineMap.getNumResults();
+    ++i;
+  }
+
   // Drop the unbounded identity maps from the composition.
   // This may lead to the composition becoming empty, which is interpreted as an
   // implicit identity.
@@ -758,13 +791,12 @@ MemRefType MemRefType::get(ArrayRef<int> shape, Type elementType,
   shape = impl.copyInto(shape);
 
   // Copy the affine map composition into the bump pointer.
-  // TODO(andydavis) Assert that the structure of the composition is valid.
   affineMapComposition =
       impl.copyInto(ArrayRef<AffineMap>(affineMapComposition));
 
   // Initialize the memory using placement new.
   new (result) MemRefTypeStorage{
-      {Kind::MemRef, context, static_cast<unsigned int>(shape.size())},
+      {Type::Kind::MemRef, context, static_cast<unsigned int>(shape.size())},
       elementType,
       shape.data(),
       static_cast<unsigned int>(affineMapComposition.size()),
@@ -774,6 +806,26 @@ MemRefType MemRefType::get(ArrayRef<int> shape, Type elementType,
   return *existing.first = result;
 }
 
+// Try constructing a MemRefType, report errors and return a nullptr on failure.
+MemRefType MemRefType::getChecked(ArrayRef<int> shape, Type elementType,
+                                  ArrayRef<AffineMap> affineMapComposition,
+                                  unsigned memorySpace, Location *location) {
+  assert(location &&
+         "location cannot be null, use UnknownLoc if it is unknown");
+  return getMemRefType(shape, elementType, affineMapComposition, memorySpace,
+                       location);
+}
+
+// Try constructing a MemRefType, supressing error messages, abort on failure.
+MemRefType MemRefType::get(ArrayRef<int> shape, Type elementType,
+                           ArrayRef<AffineMap> affineMapComposition,
+                           unsigned memorySpace) {
+  auto type = getMemRefType(shape, elementType, affineMapComposition,
+                            memorySpace, nullptr);
+  assert(type && "failed to construct a MemRef type");
+  return type;
+}
+
 //===----------------------------------------------------------------------===//
 // Attribute uniquing
 //===----------------------------------------------------------------------===//

mlir/lib/Parser/Parser.cpp

@@ -521,15 +521,10 @@ Type Parser::parseMemRefType() {
   if (!elementType)
     return nullptr;
 
-  if (!elementType.isa<IntegerType>() && !elementType.isa<FloatType>() &&
-      !elementType.isa<VectorType>())
-    return (emitError(typeLoc, "invalid memref element type"), nullptr);
-
   // Parse semi-affine-map-composition.
   SmallVector<AffineMap, 2> affineMapComposition;
   unsigned memorySpace = 0;
   bool parsedMemorySpace = false;
-  auto compositionLoc = getToken().getLoc();
 
   auto parseElt = [&]() -> ParseResult {
     if (getToken().is(Token::integer)) {
@@ -566,15 +561,8 @@ Type Parser::parseMemRefType() {
       return nullptr;
   }
 
-  if (!affineMapComposition.empty() &&
-      affineMapComposition.front().getNumDims() != dimensions.size()) {
-    emitError(compositionLoc,
-              "affine map dimension count must equal memref rank");
-    return nullptr;
-  }
-
-  return MemRefType::get(dimensions, elementType, affineMapComposition,
-                         memorySpace);
+  return MemRefType::getChecked(dimensions, elementType, affineMapComposition,
+                                memorySpace, getEncodedSourceLocation(typeLoc));
 }
 
 /// Parse a function type.

mlir/test/IR/invalid-ops.mlir

@@ -110,7 +110,7 @@ mlfunc @test_store_zero_results2(%x: i32, %p: memref<i32>) {
 
 cfgfunc @test_alloc_memref_map_rank_mismatch() {
 bb0:
-  %0 = alloc() : memref<1024x64xf32, (d0) -> (d0), 1> // expected-error {{affine map dimension count must equal memref rank}}
+  %0 = alloc() : memref<1024x64xf32, (d0) -> (d0), 1> // expected-error {{memref affine map dimension mismatch}}
   return
 }
 
@@ -164,4 +164,4 @@ bb0(%a : f32):
 cfgfunc @cfgfunc_with_ops(i32) {
 bb0(%a : i32):
   %sf = addf %a, %a : i32  // expected-error {{'addf' op requires a floating point type}}
-}
+}

mlir/test/IR/invalid.mlir

@@ -44,7 +44,13 @@ extfunc @memrefs(memref<2x4xi8, #map0, 1, #map1>) // expected-error {{affine map
 // The error must be emitted even for the trivial identity layout maps that are
 // dropped in type creation.
 #map0 = (d0, d1) -> (d0, d1)
-extfunc @memrefs(memref<42xi8, #map0>) // expected-error {{affine map dimension count must equal memref rank}}
+extfunc @memrefs(memref<42xi8, #map0>) // expected-error {{memref affine map dimension mismatch}}
+
+// -----
+
+#map0 = (d0, d1) -> (d0, d1)
+#map1 = (d0) -> (d0)
+extfunc @memrefs(memref<42x42xi8, #map0, #map1>) // expected-error {{memref affine map dimension mismatch}}
 
 // -----