Add verification for AllocOp.

PiperOrigin-RevId: 213829386

mlir/lib/IR/StandardOps.cpp

@@ -229,31 +229,47 @@ bool AllocOp::parse(OpAsmParser *parser, OperationState *result) {
     return true;
 
   // Check numDynamicDims against number of question marks in memref type.
+  // Note: this check remains here (instead of in verify()), because the
+  // partition between dim operands and symbol operands is lost after parsing.
+  // Verification still checks that the total number of operands matches
+  // the number of symbols in the affine map, plus the number of dynamic
+  // dimensions in the memref.
   if (numDimOperands != type->getNumDynamicDims()) {
     return parser->emitError(parser->getNameLoc(),
                              "dimension operand count does not equal memref "
                              "dynamic dimension count");
   }
-
-  // Check that the number of symbol operands matches the number of symbols in
-  // the first affinemap of the memref's affine map composition.
-  unsigned numSymbols = 0;
-  if (!type->getAffineMaps().empty())
-    numSymbols = type->getAffineMaps()[0]->getNumSymbols();
-
-  if (result->operands.size() - numDimOperands != numSymbols) {
-    return parser->emitError(
-        parser->getNameLoc(),
-        "affine map symbol operand count does not equal memref affine map "
-        "symbol count");
-  }
-
   result->types.push_back(type);
   return false;
 }
 
 bool AllocOp::verify() const {
-  // TODO(andydavis): Verify alloc.
+  auto *memRefType = dyn_cast<MemRefType>(getMemRef()->getType());
+  if (!memRefType)
+    return emitOpError("result must be a memref");
+
+  unsigned numSymbols = 0;
+  if (!memRefType->getAffineMaps().empty()) {
+    AffineMap *affineMap = memRefType->getAffineMaps()[0];
+    // Store number of symbols used in affine map (used in subsequent check).
+    numSymbols = affineMap->getNumSymbols();
+    // Verify that the layout affine map matches the rank of the memref.
+    if (affineMap->getNumDims() != memRefType->getRank())
+      return emitOpError("affine map dimension count must equal memref rank");
+  }
+  unsigned numDynamicDims = memRefType->getNumDynamicDims();
+  // Check that the total number of operands matches the number of symbols in
+  // the affine map, plus the number of dynamic dimensions specified in the
+  // memref type.
+  if (getOperation()->getNumOperands() != numDynamicDims + numSymbols) {
+    return emitOpError(
+        "operand count does not equal dimension plus symbol operand count");
+  }
+  // Verify that all operands are of type AffineInt.
+  for (auto *operand : getOperands()) {
+    if (!operand->getType()->isAffineInt())
+      return emitOpError("requires operands to be of type AffineInt");
+  }
   return false;
 }
 

mlir/test/IR/invalid-ops.mlir

@@ -82,7 +82,7 @@ cfgfunc @bad_alloc_wrong_symbol_count() {
 bb0:
   %0 = "constant"() {value: 7} : () -> affineint
   // Test alloc with wrong number of symbols
-  %1 = alloc(%0) : memref<2x?xf32, (d0, d1)[s0] -> ((d0 + s0), d1), 1> // expected-error {{custom op 'alloc' affine map symbol operand count does not equal memref affine map symbol count}}
+  %1 = alloc(%0) : memref<2x?xf32, (d0, d1)[s0] -> ((d0 + s0), d1), 1> // expected-error {{operand count does not equal dimension plus symbol operand count}}
   return
 }
 
@@ -101,6 +101,14 @@ bb0:
 
 // -----
 
+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}}
+  return
+}
+
+// -----
+
 cfgfunc @intlimit2() {
 bb:
   %0 = "constant"() {value: 0} : () -> i4096