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

[Wen-Heng (Jack) Chung via Phabricator via llvm-commits]

whchung updated this revision to Diff 267637.
whchung added a comment.

Revise the fix and added positive and negative tests.


Repository:
  rG LLVM Github Monorepo

CHANGES SINCE LAST ACTION
  https://reviews.llvm.org/D80766/new/

https://reviews.llvm.org/D80766

Files:
  mlir/include/mlir/Dialect/GPU/GPUOps.td
  mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
  mlir/test/Dialect/GPU/invalid.mlir
  mlir/test/Dialect/GPU/ops.mlir


Index: mlir/test/Dialect/GPU/ops.mlir
===================================================================
--- mlir/test/Dialect/GPU/ops.mlir
+++ mlir/test/Dialect/GPU/ops.mlir
@@ -132,4 +132,11 @@
     }
   }
 
+  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} : () -> ()
+  }
 }
Index: mlir/test/Dialect/GPU/invalid.mlir
===================================================================
--- mlir/test/Dialect/GPU/invalid.mlir
+++ mlir/test/Dialect/GPU/invalid.mlir
@@ -423,3 +423,15 @@
     }
   }
 }
+
+// -----
+
+module {
+  gpu.module @gpu_funcs {
+    // expected-error @+1 {{'gpu.func' op expected at least 4294967300 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} : () -> ()
+  }
+}
Index: mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
===================================================================
--- mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
+++ mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
@@ -673,10 +673,12 @@
 
 /// Verifies the body of the function.
 LogicalResult GPUFuncOp::verifyBody() {
-  unsigned numFuncArguments = getNumArguments();
-  unsigned numWorkgroupAttributions = getNumWorkgroupAttributions();
-  unsigned numPrivateAttributions = getNumPrivateAttributions();
-  unsigned numBlockArguments = front().getNumArguments();
+  // Explicitly widen the range to int64_t, to catch the case any of the values
+  // wraps around within unsigned type.
+  int64_t numFuncArguments = getNumArguments();
+  int64_t numWorkgroupAttributions = getNumWorkgroupAttributions();
+  int64_t numPrivateAttributions = getNumPrivateAttributions();
+  int64_t numBlockArguments = front().getNumArguments();
   if (numBlockArguments <
       numFuncArguments + numWorkgroupAttributions + numPrivateAttributions)
     return emitOpError() << "expected at least "
Index: mlir/include/mlir/Dialect/GPU/GPUOps.td
===================================================================
--- mlir/include/mlir/Dialect/GPU/GPUOps.td
+++ mlir/include/mlir/Dialect/GPU/GPUOps.td
@@ -200,7 +200,7 @@
 
     /// Returns the number of buffers located in the private memory.
     unsigned getNumPrivateAttributions() {
-      return getOperation()->getNumOperands() - getType().getNumInputs() -
+      return getBody().front().getNumArguments() - getType().getNumInputs() -
           getNumWorkgroupAttributions();
     }
  


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D80766.267637.patch
Type: text/x-patch
Size: 2947 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20200601/59c65cc7/attachment.bin>