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

[llvmlistbot at llvm.org]

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

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.

>From ce28b9c1489c56aa12dbeeedfcdb451c5ccce27b Mon Sep 17 00:00:00 2001
From: Kazuaki Matsumura <kmatsumura at nvidia.com>
Date: Thu, 17 Sep 2026 07:56:12 -0700
Subject: [PATCH] [OpenACC] Honor precomputed active par dims on a predicate
 region

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.

Let a predicate region state which dims run it unpredicated, the way a
privatization already can. 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.
---
 .../Dialect/OpenACC/Transforms/ACCCGToGPU.cpp |  5 +-
 ...g-to-gpu-predicate-region-active-dims.mlir | 58 +++++++++++++++++++
 2 files changed, 61 insertions(+), 2 deletions(-)
 create mode 100644 mlir/test/Dialect/OpenACC/acc-cg-to-gpu-predicate-region-active-dims.mlir

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
+}