Add verification for dimension attribute on GPUDialect index operations.

PiperOrigin-RevId: 266073204

diff --git a/mlir/include/mlir/Dialect/GPU/GPUOps.td b/mlir/include/mlir/Dialect/GPU/GPUOps.td
index b38a597..6c713b7 100644 (file)
--- a/mlir/include/mlir/Dialect/GPU/GPUOps.td
+++ b/mlir/include/mlir/Dialect/GPU/GPUOps.td
@@ -37,7 +37,9 @@ class GPU_Op<string mnemonic, list<OpTrait> traits = []> :
 
 class GPU_IndexOp<string mnemonic, list<OpTrait> traits = []> :
     GPU_Op<mnemonic, !listconcat(traits, [NoSideEffect])>,
-    Arguments<(ins StrAttr:$dimension)>, Results<(outs Index)>;
+    Arguments<(ins StrAttr:$dimension)>, Results<(outs Index)> {
+  let verifier = [{ return ::verifyIndexOp(*this); }];
+}
 
 def gpu_BlockDim : GPU_IndexOp<"block_dim">;
 def gpu_BlockId : GPU_IndexOp<"block_id">;

diff --git a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
index 22d433a..5609725 100644 (file)
--- a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
+++ b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
@@ -47,6 +47,13 @@ GPUDialect::GPUDialect(MLIRContext *context)
                 >();
 }
 
+template <typename T> static LogicalResult verifyIndexOp(T op) {
+  auto dimension = op.dimension();
+  if (dimension != "x" && dimension != "y" && dimension != "z")
+    return op.emitError("dimension \"") << dimension << "\" is invalid";
+  return success();
+}
+
 #define GET_OP_CLASSES
 #include "mlir/Dialect/GPU/GPUOps.cpp.inc"
 

diff --git a/mlir/test/Dialect/GPU/invalid.mlir b/mlir/test/Dialect/GPU/invalid.mlir
index bfc18fa..2a0be68 100644 (file)
--- a/mlir/test/Dialect/GPU/invalid.mlir
+++ b/mlir/test/Dialect/GPU/invalid.mlir
@@ -162,3 +162,39 @@ func @launch_func_kernel_operand_types(%sz : index, %arg : f32) {
       : (index, index, index, index, index, index, f32) -> ()
   return
 }
+
+// -----
+
+func @illegal_dimension() {
+  // expected-error@+1 {{dimension "o" is invalid}}
+  %tIdX = "gpu.thread_id"() {dimension = "o"} : () -> (index)
+
+  return
+}
+
+// -----
+
+func @illegal_dimension() {
+  // expected-error@+1 {{dimension "o" is invalid}}
+  %bDimX = "gpu.block_dim"() {dimension = "o"} : () -> (index)
+
+  return
+}
+
+// -----
+
+func @illegal_dimension() {
+  // expected-error@+1 {{dimension "o" is invalid}}
+  %bIdX = "gpu.block_id"() {dimension = "o"} : () -> (index)
+
+  return
+}
+
+// -----
+
+func @illegal_dimension() {
+  // expected-error@+1 {{dimension "o" is invalid}}
+  %gDimX = "gpu.grid_dim"() {dimension = "o"} : () -> (index)
+
+  return
+}