[mlir][gpu] Guard negative workgroup_attributions in GPU ops

[meshtag]

Prevent crashes from negative workgroup_attributions by validating the attribute in gpu.launch and gpu.func verification. We emit an error with a helpful message instead of an outright crash in such cases now.

Fixes #159674

[llvmbot]

@llvm/pr-subscribers-mlir-gpu

@llvm/pr-subscribers-mlir

Author: Prathamesh Tagore (meshtag)

Changes

Prevent crashes from negative workgroup_attributions by validating the attribute in gpu.launch and gpu.func verification. We emit an error with a helpful message instead of an outright crash in such cases now.

Fixes #159674

---

Full diff: https://github.com/llvm/llvm-project/pull/174409.diff

2 Files Affected:

• (modified) mlir/lib/Dialect/GPU/IR/GPUDialect.cpp (+17-2)
• (modified) mlir/test/Dialect/GPU/invalid.mlir (+24)

diff --git a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
index 21c0d369b8d1c..7e174f1d21adb 100644
--- a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
+++ b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
@@ -886,12 +886,20 @@ LogicalResult LaunchOp::verify() {
 }
 
 LogicalResult LaunchOp::verifyRegions() {
+  int64_t numWorkgroupAttributions = 0;
+  if (auto attr = (*this)->getAttrOfType<IntegerAttr>(
+          getNumWorkgroupAttributionsAttrName())) {
+    numWorkgroupAttributions = attr.getInt();
+    if (numWorkgroupAttributions < 0)
+      return emitOpError() << "expected non-negative workgroup_attributions";
+  }
+
   // Kernel launch takes kNumConfigOperands leading operands for grid/block
   // sizes and transforms them into kNumConfigRegionAttributes region arguments
   // for block/thread identifiers and grid/block sizes.
   if (!getBody().empty()) {
     if (getBody().getNumArguments() <
-        kNumConfigRegionAttributes + getNumWorkgroupAttributions())
+        kNumConfigRegionAttributes + numWorkgroupAttributions)
       return emitOpError("unexpected number of region arguments");
   }
 
@@ -1814,7 +1822,14 @@ LogicalResult GPUFuncOp::verifyBody() {
   if (empty())
     return emitOpError() << "expected body with at least one block";
   unsigned numFuncArguments = getNumArguments();
-  unsigned numWorkgroupAttributions = getNumWorkgroupAttributions();
+  int64_t numWorkgroupAttributions = 0;
+  if (auto attr = (*this)->getAttrOfType<IntegerAttr>(
+          getNumWorkgroupAttributionsAttrName())) {
+    numWorkgroupAttributions = attr.getInt();
+    if (numWorkgroupAttributions < 0)
+      return emitOpError() << "expected non-negative workgroup_attributions";
+  }
+
   unsigned numBlockArguments = front().getNumArguments();
   if (numBlockArguments < numFuncArguments + numWorkgroupAttributions)
     return emitOpError() << "expected at least "
diff --git a/mlir/test/Dialect/GPU/invalid.mlir b/mlir/test/Dialect/GPU/invalid.mlir
index 26bcf948bc85d..bc06e78186b19 100644
--- a/mlir/test/Dialect/GPU/invalid.mlir
+++ b/mlir/test/Dialect/GPU/invalid.mlir
@@ -1021,3 +1021,27 @@ func.func @warp_mismatch_rank(%laneid: index) {
   }
   return
 }
+
+// -----
+
+func.func @launch_negative_workgroup_attributions(%sz : index) {
+  // expected-error@+1 {{expected non-negative workgroup_attributions}}
+  "gpu.launch"(%sz, %sz, %sz, %sz, %sz, %sz) ({})
+      {operandSegmentSizes = array<i32: 0, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0>,
+       workgroup_attributions = -1 : i64}
+      : (index, index, index, index, index, index) -> ()
+  return
+}
+
+// -----
+
+module {
+  gpu.module @gpu_funcs {
+    // expected-error@+1 {{expected non-negative workgroup_attributions}}
+    "gpu.func"() ({
+    ^bb0:
+      gpu.return
+    }) {function_type = () -> (), sym_name = "gpu_func_negative_workgroup_attributions",
+         workgroup_attributions = -1 : i64} : () -> ()
+  }
+}