[mlir][GPU] treat the absence of workgroup attributes correctly

The helper function in GPUFuncOp incorrectly assumed the workgroup
attribution attribute is always present. Instead, treat its absence as
if its value was zero, i.e., no workgroup attributions are specified.

Closes #58045.

Reviewed By: nicolasvasilache

Differential Revision: https://reviews.llvm.org/D134865

diff --git a/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td b/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td
index c781606..673be1b 100644 (file)
--- a/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td
+++ b/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td
@@ -263,8 +263,9 @@ def GPU_GPUFuncOp : GPU_Op<"func", [
 
     /// Returns the number of buffers located in the workgroup memory.
     unsigned getNumWorkgroupAttributions() {
-      return (*this)->getAttrOfType<IntegerAttr>(
-          getNumWorkgroupAttributionsAttrName()).getInt();
+      auto attr = (*this)->getAttrOfType<IntegerAttr>(
+          getNumWorkgroupAttributionsAttrName());
+      return attr ? attr.getInt() : 0;
     }
 
     /// Returns a list of block arguments that correspond to buffers located in

diff --git a/mlir/test/Dialect/GPU/ops.mlir b/mlir/test/Dialect/GPU/ops.mlir
index f882a8e..0ffaf22 100644 (file)
--- a/mlir/test/Dialect/GPU/ops.mlir
+++ b/mlir/test/Dialect/GPU/ops.mlir
@@ -284,3 +284,10 @@ module attributes {gpu.container_module} {
     return
   }
 }
+
+// Just check that this doesn't crash.
+gpu.module @module {
+  "gpu.func"() ({
+    gpu.return
+  }) {function_type = () -> (), sym_name = "func"} : () -> ()
+}