Add verifier for gpu.alloc op

Add verifier for gpu.alloc op to verify if the dimension operand counts
and symbol operand counts are same as their memref counterparts.

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

Author: akshaybaviskar

mlir/include/mlir/Dialect/GPU/GPUOps.td

@@ -901,6 +901,7 @@ def GPU_AllocOp : GPU_Op<"alloc", [
     `(` $dynamicSizes `)` (`` `[` $symbolOperands^ `]`)? attr-dict `:` type($memref)
   }];
 
+  let verifier = [{ return ::verify(*this); }];
   let hasCanonicalizer = 1;
 }
 

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

@@ -1173,6 +1173,25 @@ LogicalResult MemsetOp::fold(ArrayRef<Attribute> operands,
 //===----------------------------------------------------------------------===//
 // GPU_AllocOp
 //===----------------------------------------------------------------------===//
+
+static LogicalResult verify(AllocOp op) {
+  auto memRefType = op.memref().getType().cast<MemRefType>();
+
+  if (static_cast<int64_t>(op.dynamicSizes().size()) !=
+      memRefType.getNumDynamicDims())
+    return op.emitOpError("dimension operand count does not equal memref "
+                          "dynamic dimension count");
+
+  unsigned numSymbols = 0;
+  if (!memRefType.getLayout().isIdentity())
+    numSymbols = memRefType.getLayout().getAffineMap().getNumSymbols();
+  if (op.symbolOperands().size() != numSymbols)
+    return op.emitOpError("symbol operand count does not equal memref symbol "
+                          "count");
+
+  return success();
+}
+
 namespace {
 
 /// Folding of memref.dim(gpu.alloc(%size), %idx) -> %size similar to

mlir/test/Dialect/GPU/invalid.mlir

@@ -610,4 +610,43 @@ func @async_cp_num_dst_stride(
   gpu.device_async_copy %src[%i, %i], %dst[%i, %i], 16 :
     memref<200x100xf32> to memref<200x100xf32, affine_map<(d0, d1) -> (200*d0 + 2*d1)>, 3>
   return
-}
+}
+
+// -----
+
+// Number of symbol operand count less than memref symbol count.
+func @alloc() {
+   // expected-error@+1 {{symbol operand count does not equal memref symbol count}}
+   %1 = gpu.alloc() : memref<2x4xf32, affine_map<(d0, d1)[s0] -> ((d0 + s0), d1)>, 1>
+   return
+}
+
+// -----
+
+// Number of symbol operand count greater than memref symbol count.
+func @alloc() {
+   %0 = arith.constant 7 : index
+   // expected-error@+1 {{symbol operand count does not equal memref symbol count}}
+   %1 = gpu.alloc()[%0] : memref<2x4xf32, 1>
+   return
+}
+
+// -----
+
+// Number of dynamic dimension operand count greater than memref dynamic dimension count.
+func @alloc() {
+   %0 = arith.constant 7 : index
+   // expected-error@+1 {{dimension operand count does not equal memref dynamic dimension count}}
+   %1 = gpu.alloc(%0, %0) : memref<2x?xf32, 1>
+   return
+}
+
+// -----
+
+// Number of dynamic dimension operand count less than memref dynamic dimension count.
+func @alloc() {
+   %0 = arith.constant 7 : index
+   // expected-error@+1 {{dimension operand count does not equal memref dynamic dimension count}}
+   %1 = gpu.alloc(%0) : memref<2x?x?xf32, 1>
+   return
+}