[Mlir-commits] [mlir] f1efac7 - Add verifier for gpu.alloc op

[Uday Bondhugula]

Author: Akshay Baviskar
Date: 2022-02-15T15:57:58+05:30
New Revision: f1efac7f081365a937f289c6619adf213a5c58b3

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

LOG: 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

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/GPU/GPUOps.td
    mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
    mlir/test/Dialect/GPU/invalid.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/GPU/GPUOps.td b/mlir/include/mlir/Dialect/GPU/GPUOps.td
index f6a5d4ca7eb8c..d8236eafa9cf1 100644
--- a/mlir/include/mlir/Dialect/GPU/GPUOps.td
+++ b/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;
 }
 

diff  --git a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
index 3825aec5d30e2..24deaad63221c 100644
--- a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
+++ b/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

diff  --git a/mlir/test/Dialect/GPU/invalid.mlir b/mlir/test/Dialect/GPU/invalid.mlir
index 3b2b94b203962..006e1188de3b3 100644
--- a/mlir/test/Dialect/GPU/invalid.mlir
+++ b/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
-}
\ No newline at end of file
+}
+
+// -----
+
+// Number of symbol operand count less than memref symbol count.
+func @alloc() {
+   // expected-error at +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 at +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 at +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 at +1 {{dimension operand count does not equal memref dynamic dimension count}}
+   %1 = gpu.alloc(%0) : memref<2x?x?xf32, 1>
+   return
+}