[mlir][gpu] Fix logic error in D79508 computing number of private attributions.

[whchung]

Summary

Fix logic error in D79508. The old logic would make the first check in
GPUFuncOp::verifyBody always pass.

[herhut]

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

693

I don't think this should be in a verifier. getWorkgroupAttributions() is implemented by means of getNumWorkgroupAttributions, so this is not an invariant we need to check for each operation.

If we want to prevent this from regressing, we would need to write a unit test that shows the wrong behavior. You could specify the illegal case by using the generic syntax, which does not go through the parser but has to explicitly specify attributes. That should allow you to construct also illegal cases.

You can use mlir-opt -mlir-print-op-generic to look at the generic syntax.

[whchung]

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

693

Thanks for the hint. Will revise the patch.

[whchung]

Revise the fix and added positive and negative tests.

[herhut]

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

680–681

This produces error messages that are not overly helpful. I would verify this is follows. We need at least getNumArguments() + getNumWorkgroupAttributions() many operands, as otherwise there are not enough. The getNumPrivateAttributions() does not need to be verified, as it is computed as the remaining operands that are left over.

So the only error case is when we do not have enough operands to satisfy function arguments and workgroup attributions.

[whchung]

Address code review comments.

[herhut]

Thanks! Please excuse the long delay due to me travelling.

[whchung]

This patch was landed in commit 603b974cf7103766a0e5e4a0320fedb7c4b570f9.

mlir/include/mlir/Dialect/GPU/GPUOps.td

let extraClassDeclaration = [{
     }
 
     /// Adds a new block argument that corresponds to buffers located in
     /// workgroup memory.
     BlockArgument addWorkgroupAttribution(Type type);
 
     /// Returns the number of buffers located in the private memory.
     unsigned getNumPrivateAttributions() {
-      return getOperation()->getNumOperands() - getType().getNumInputs() -
+      return getBody().front().getNumArguments() - getType().getNumInputs() -
           getNumWorkgroupAttributions();
     }
  
     /// Returns a list of block arguments that correspond to buffers located in
     /// the private memory.
     ArrayRef<BlockArgument> getPrivateAttributions() {
       // Buffers on the private memory always come after buffers on the workgroup
       // memory.

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

static LogicalResult verifyAttributions(Operation *op,
   }
   return success();
 }
 
 /// Verifies the body of the function.
 LogicalResult GPUFuncOp::verifyBody() {
   unsigned numFuncArguments = getNumArguments();
   unsigned numWorkgroupAttributions = getNumWorkgroupAttributions();
-  unsigned numPrivateAttributions = getNumPrivateAttributions();
   unsigned numBlockArguments = front().getNumArguments();
-  if (numBlockArguments <
-      numFuncArguments + numWorkgroupAttributions + numPrivateAttributions)
+  if (numBlockArguments < numFuncArguments + numWorkgroupAttributions)
     return emitOpError() << "expected at least "
-                         << numFuncArguments + numWorkgroupAttributions +
-                                numPrivateAttributions
+                         << numFuncArguments + numWorkgroupAttributions
                          << " arguments to body region";
 
   ArrayRef<Type> funcArgTypes = getType().getInputs();
   for (unsigned i = 0; i < numFuncArguments; ++i) {
     Type blockArgType = front().getArgument(i).getType();
     if (funcArgTypes[i] != blockArgType)
       return emitOpError() << "expected body region argument #" << i
                            << " to be of type " << funcArgTypes[i] << ", got "
                            << blockArgType;
   }
 
   if (failed(verifyAttributions(getOperation(), getWorkgroupAttributions(),
                                 GPUDialect::getWorkgroupAddressSpace())) ||
       failed(verifyAttributions(getOperation(), getPrivateAttributions(),
                                 GPUDialect::getPrivateAddressSpace())))
     return failure();
 
   return success();
 }
 

mlir/test/Dialect/GPU/invalid.mlir

module {
   gpu.module @gpu_funcs {
     // expected-error @+1 {{'gpu.func' op expected void return type for kernel function}}
     gpu.func @kernel() -> index kernel {
       %0 = constant 0 : index
       gpu.return
     }
   }
 }
+
+// -----
+
+module {
+  gpu.module @gpu_funcs {
+    // expected-error @+1 {{'gpu.func' op expected at least 5 arguments to body region}}
+    "gpu.func"() ( {
+    ^bb0(%arg0: f32, %arg1: memref<?xf32>, %arg2: memref<5xf32, 3>, %arg3: memref<5xf32, 5>):
+      "gpu.return"() : () -> ()
+    } ) {gpu.kernel, sym_name = "kernel_1", type = (f32, memref<?xf32>) -> (), workgroup_attributions = 3: i64} : () -> ()
+  }
+}

mlir/test/Dialect/GPU/ops.mlir

gpu.module @gpu_funcs {
 
     // CHECK-LABEL: @empty_attribution
     // CHECK:       {
     gpu.func @empty_attribution(%arg0: f32) workgroup() private() {
       gpu.return
     }
   }
 
+  gpu.module @explicit_attributions {
+    // CHECK-LABEL: gpu.func @kernel_1({{.*}}: f32, {{.*}}: memref<?xf32>) workgroup({{.*}}: memref<5xf32, 3>) private({{.*}}: memref<5xf32, 5>)
+    "gpu.func"() ( {
+    ^bb0(%arg0: f32, %arg1: memref<?xf32>, %arg2: memref<5xf32, 3>, %arg3: memref<5xf32, 5>):
+      "gpu.return"() : () -> ()
+    } ) {gpu.kernel, sym_name = "kernel_1", type = (f32, memref<?xf32>) -> (), workgroup_attributions = 1: i64} : () -> ()
+  }
 }