diff --git a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp index f776129c77405..31c49a41d8146 100644 --- a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp +++ b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp @@ -839,9 +839,13 @@ LogicalResult LaunchOp::verifyRegions() { if (getBody().empty()) { return emitOpError("body region is empty"); } - if (getBody().getNumArguments() < - kNumConfigRegionAttributes + getNumWorkgroupAttributions()) { - return emitOpError("unexpected number of region arguments"); + unsigned actualNumRegionArgs = getBody().getNumArguments(); + unsigned expectedNumRegionArgs = + getNumConfigRegionAttributes() + getNumWorkgroupAttributions(); + if (actualNumRegionArgs < expectedNumRegionArgs) { + return emitOpError("expected at least ") + << expectedNumRegionArgs << " region arguments, but got " + << actualNumRegionArgs; } // Verify Attributions Address Spaces. diff --git a/mlir/test/Dialect/GPU/invalid.mlir b/mlir/test/Dialect/GPU/invalid.mlir index fb01bf0bcce9c..e3a981c1c7afe 100644 --- a/mlir/test/Dialect/GPU/invalid.mlir +++ b/mlir/test/Dialect/GPU/invalid.mlir @@ -20,7 +20,7 @@ func.func @not_enough_sizes(%sz : index) { // ----- func.func @no_region_attrs(%sz : index) { - // expected-error@+1 {{unexpected number of region arguments}} + // expected-error@+1 {{expected at least 12 region arguments, but got 6}} "gpu.launch"(%sz, %sz, %sz, %sz, %sz, %sz) ({ ^bb1(%bx: index, %by: index, %bz: index, %tx: index, %ty: index, %tz: index): @@ -31,6 +31,20 @@ func.func @no_region_attrs(%sz : index) { // ----- +func.func @not_enough_cluster_region_attrs(%sz : index) { + // expected-error@+1 {{expected at least 18 region arguments, but got 12}} + "gpu.launch"(%sz, %sz, %sz, %sz, %sz, %sz, %sz, %sz, %sz) ({ + ^bb1(%bx: index, %by: index, %bz: index, + %tx: index, %ty: index, %tz: index, + %sbx: index, %sby: index, %sbz: index, + %stx: index, %sty: index, %stz: index): + gpu.terminator + }) {operandSegmentSizes = array} : (index, index, index, index, index, index, index, index, index) -> () + return +} + +// ----- + func.func @launch_requires_gpu_return(%sz : index) { // @expected-note@+1 {{in 'gpu.launch' body region}} gpu.launch blocks(%bx, %by, %bz) in (%sbx = %sz, %sby = %sz, %sbz = %sz)