diff --git a/mlir/include/mlir/Dialect/GPU/GPUOps.td b/mlir/include/mlir/Dialect/GPU/GPUOps.td --- a/mlir/include/mlir/Dialect/GPU/GPUOps.td +++ b/mlir/include/mlir/Dialect/GPU/GPUOps.td @@ -882,6 +882,7 @@ `(` $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 --- a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp +++ b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp @@ -1165,6 +1165,25 @@ //===----------------------------------------------------------------------===// // GPU_AllocOp //===----------------------------------------------------------------------===// + +static LogicalResult verify(AllocOp op) { + auto memRefType = op.memref().getType().cast(); + + if (static_cast(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 --- a/mlir/test/Dialect/GPU/invalid.mlir +++ b/mlir/test/Dialect/GPU/invalid.mlir @@ -574,3 +574,42 @@ %D = gpu.subgroup_mma_compute %A, %B, %C : !gpu.mma_matrix<16x32xf16, "AOp">, !gpu.mma_matrix<16x16xf16, "BOp"> -> !gpu.mma_matrix<16x16xf16, "COp"> 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 +}