[Mlir-commits] [mlir] [mlir][gpu] Guard negative workgroup_attributions in GPU ops (PR #174409)

[Prathamesh Tagore]

================
@@ -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)
----------------
meshtag wrote:

Following is the function definition of `getNumWorkgroupAttributions()` in tablegen: 
https://github.com/llvm/llvm-project/blob/4f90ce4e6ec20db592836421ce309aaab134c353/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td#L455

It is returning an unsigned value which will overflow for negative numbers and then the comparison there won't remain meaningful. Hence the attribute was read separately for the verification purpose.

We can choose to return an `int64_t` from the function so that we don't compromise on number of representable values and then use the same function for verification, but then this would mean allocating a larger register entry for each op instance at all times even though we only need it to verify the operation once during its creation.

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