[Mlir-commits] [mlir] [OpenACC] Honor precomputed active par dims on a predicate region (PR #224390)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-openacc

Author: Matsu (khaki3)

<details>
<summary>Changes</summary>

Example:
```mlir
acc.predicate_region {
  memref.store %updated, %shared[] : memref<f64>
} {acc.active_par_dims = #acc<active_par_dims[]>}
```

Code outside a gang-level loop runs on every thread block, so the block dims
are treated as active and never predicated away. That is right for a redundant
computation, but not for an in-place update of memory the whole launch shares:
every block applies it. The attribute above was ignored here, because only a
privatization could state which dims run it unpredicated.

Fix: honor a precomputed active set on a predicate region too. An empty set
leaves one block and one thread performing the update; listing the thread dims
of the enclosing loops keeps a work-shared update distributing its iterations.
Behavior is unchanged for a region without the attribute.

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


2 Files Affected:

- (modified) mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp (+3-2) 
- (added) mlir/test/Dialect/OpenACC/acc-cg-to-gpu-predicate-region-active-dims.mlir (+58) 


``````````diff
diff --git a/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp b/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
index b9d634e059a01..716662f4d605d 100644
--- a/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
+++ b/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
@@ -1472,9 +1472,10 @@ ACCCGToGPULowering::computeActiveAndInactiveParDims(Operation *op,
   }
 
   // The active set is precomputed separately from the ownership par_dims for
-  // privatizations; prefer that attribute when present. The inactive dims are
+  // privatizations, and a predicate region may likewise state which dims run
+  // it unpredicated; prefer that attribute when present. The inactive dims are
   // the launch dims which are not active.
-  if (isa<acc::PrivateLocalOp, acc::PrivatizeOp>(op)) {
+  if (isa<acc::PrivateLocalOp, acc::PrivatizeOp, acc::PredicateRegionOp>(op)) {
     if (mlir::acc::ActiveParDimsAttr precomputedActiveParDims =
             getActiveParDimsAttr(op)) {
       mlir::acc::GPUParallelDimAttr lowestParDim =
diff --git a/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-predicate-region-active-dims.mlir b/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-predicate-region-active-dims.mlir
new file mode 100644
index 0000000000000..3fda43c53ce56
--- /dev/null
+++ b/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-predicate-region-active-dims.mlir
@@ -0,0 +1,58 @@
+// RUN: mlir-opt %s --pass-pipeline="builtin.module(func.func(acc-cg-to-gpu))" --split-input-file | FileCheck %s
+
+// Without the attribute the update runs on every block: only the thread dim
+// is predicated.
+// CHECK-LABEL: func.func @gang_redundant_update
+// CHECK:       gpu.launch blocks(%[[BID:[a-z0-9_]+]], %{{.*}}) in ({{.*}}) threads(%[[TID:[a-z0-9_]+]], %{{.*}}) in ({{.*}}) {
+// CHECK:       %[[P:.*]] = arith.cmpi eq, %[[TID]], %c0
+// CHECK-NOT:   arith.cmpi eq, %[[BID]]
+// CHECK:       scf.if %[[P]]
+
+func.func @gang_redundant_update(%arg0: memref<f32>) {
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+  %0 = acc.par_width %c2 par_dim(#acc.par_dim<block_x>)
+  %1 = acc.par_width %c2 par_dim(#acc.par_dim<thread_x>)
+  acc.kernel_environment {
+    acc.compute_region launch(%arg1 = %0, %arg2 = %1) ins(%arg10 = %arg0) : (memref<f32>) {
+      %cst = arith.constant 1.000000e+00 : f32
+      %2 = memref.load %arg10[] : memref<f32>
+      %3 = arith.addf %2, %cst : f32
+      acc.predicate_region {
+        memref.store %3, %arg10[] : memref<f32>
+      }
+      acc.yield
+    } <{origin = "acc.kernels"}>
+  }
+  return
+}
+
+// -----
+
+// An empty active set means no dim runs it unpredicated, so the block dim is
+// predicated too and one block applies the update.
+// CHECK-LABEL: func.func @single_block_update
+// CHECK:       gpu.launch blocks(%[[BID:[a-z0-9_]+]], %{{.*}}) in ({{.*}}) threads(%[[TID:[a-z0-9_]+]], %{{.*}}) in ({{.*}}) {
+// CHECK:       %[[PB:.*]] = arith.cmpi eq, %[[BID]], %c0
+// CHECK:       %[[PT:.*]] = arith.cmpi eq, %[[TID]], %c0
+// CHECK:       %[[P:.*]] = arith.andi %[[PT]], %[[PB]]
+// CHECK:       scf.if %[[P]]
+
+func.func @single_block_update(%arg0: memref<f32>) {
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+  %0 = acc.par_width %c2 par_dim(#acc.par_dim<block_x>)
+  %1 = acc.par_width %c2 par_dim(#acc.par_dim<thread_x>)
+  acc.kernel_environment {
+    acc.compute_region launch(%arg1 = %0, %arg2 = %1) ins(%arg10 = %arg0) : (memref<f32>) {
+      %cst = arith.constant 1.000000e+00 : f32
+      %2 = memref.load %arg10[] : memref<f32>
+      %3 = arith.addf %2, %cst : f32
+      acc.predicate_region {
+        memref.store %3, %arg10[] : memref<f32>
+      } {acc.active_par_dims = #acc<active_par_dims[]>}
+      acc.yield
+    } <{origin = "acc.kernels"}>
+  }
+  return
+}

``````````

</details>


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