[Mlir-commits] [mlir] [OpenACC] apply seq par dims to reduction variables (PR #209261)

[Scott Manley]

https://github.com/rscottmanley updated https://github.com/llvm/llvm-project/pull/209261

>From 3d19a90e554316bd98233abc0602c911a944b12f Mon Sep 17 00:00:00 2001
From: Scott Manley <scmanley at nvidia.com>
Date: Mon, 13 Jul 2026 07:00:57 -0700
Subject: [PATCH] [OpenACC] apply seq par dims to reduction variables

When the acc.parallel op is "effectively sequential" apply the correct
par dims during recipe materialization of reduction variables
---
 .../mlir/Dialect/OpenACC/OpenACCUtilsCG.h     | 13 ++++++++
 .../Transforms/ACCRecipeMaterialization.cpp   | 16 ++++++----
 .../acc-recipe-materialization-reduction.mlir | 32 +++++++++++++++++++
 3 files changed, 55 insertions(+), 6 deletions(-)

diff --git a/mlir/include/mlir/Dialect/OpenACC/OpenACCUtilsCG.h b/mlir/include/mlir/Dialect/OpenACC/OpenACCUtilsCG.h
index 8cfc4dac4ab36..9ab5365405d5d 100644
--- a/mlir/include/mlir/Dialect/OpenACC/OpenACCUtilsCG.h
+++ b/mlir/include/mlir/Dialect/OpenACC/OpenACCUtilsCG.h
@@ -95,6 +95,19 @@ void updateParDimsAttr(Operation *op, GPUParallelDimsAttr attr);
 /// Copy parallel dimensions from \p from to \p to.
 void copyParDimsAttr(Operation *from, Operation *to);
 
+/// Create a gang dim 1 GPUParallelDimsAttr based on the mapping policy.
+inline GPUParallelDimsAttr
+getGangDim1ParDimsAttr(MLIRContext *ctx, ACCToGPUMappingPolicy &policy) {
+  return GPUParallelDimsAttr::get(
+      ctx, {policy.gangDim(ctx, acc::ParLevel::gang_dim1)});
+}
+
+/// Create a sequential GPUParallelDimsAttr based on the mapping policy.
+inline GPUParallelDimsAttr getSeqParDimsAttr(MLIRContext *ctx,
+                                             ACCToGPUMappingPolicy &policy) {
+  return GPUParallelDimsAttr::get(ctx, {policy.seqDim(ctx)});
+}
+
 /// Tracks aligned byte consumption against a configurable shared memory cap.
 class SharedMemoryBudget {
 public:
diff --git a/mlir/lib/Dialect/OpenACC/Transforms/ACCRecipeMaterialization.cpp b/mlir/lib/Dialect/OpenACC/Transforms/ACCRecipeMaterialization.cpp
index d2e8d77c168d1..e666bcd9a6242 100644
--- a/mlir/lib/Dialect/OpenACC/Transforms/ACCRecipeMaterialization.cpp
+++ b/mlir/lib/Dialect/OpenACC/Transforms/ACCRecipeMaterialization.cpp
@@ -383,25 +383,29 @@ LogicalResult ACCRecipeMaterialization::materialize(
     cloneRegionIntoAccRegion(&combinerRegion, &combineRegionOp.getRegion(),
                              /*hasResult=*/false);
 
-    auto ctx = b.getContext();
+    auto *ctx = b.getContext();
 
     // For reductions that come from parallel constructs, explicitly set the
     // GPU parallel dimensions attribute to blockXDim since they will always be
     // gang private. GPU parallel dimensions cannot be determined for acc.loop
     // at this point.
     if constexpr (std::is_same_v<AccOpTy, acc::ParallelOp>) {
-      auto parDimsAttr = acc::GPUParallelDimsAttr::get(
-          ctx, {policy.gangDim(ctx, acc::ParLevel::gang_dim1)});
+      acc::GPUParallelDimsAttr parDimsAttr;
+      if (accOp.isEffectivelySerial()) {
+        // If acc.serial has been lowered to a parallel op that is effectively
+        // sequential
+        parDimsAttr = acc::getSeqParDimsAttr(ctx, policy);
+      } else {
+        parDimsAttr = acc::getGangDim1ParDimsAttr(ctx, policy);
+      }
       acc::setParDimsAttr(reductionOp, parDimsAttr);
       acc::setParDimsAttr(combineRegionOp, parDimsAttr);
     }
 
     // Set sequential parallel dimensions attribute for loops in the recipe.
-    auto seqParDimsAttr =
-        acc::GPUParallelDimsAttr::get(ctx, {policy.seqDim(ctx)});
     auto setSeqParDimsForRecipeLoops = [&](Region *r) {
       r->walk([&](LoopLikeOpInterface loopLike) {
-        acc::setParDimsAttr(loopLike, seqParDimsAttr);
+        acc::setParDimsAttr(loopLike, acc::getSeqParDimsAttr(ctx, policy));
       });
     };
     setSeqParDimsForRecipeLoops(&reductionOp.getRegion());
diff --git a/mlir/test/Dialect/OpenACC/acc-recipe-materialization-reduction.mlir b/mlir/test/Dialect/OpenACC/acc-recipe-materialization-reduction.mlir
index 6f2bafb2b6165..e5dc884d37d20 100644
--- a/mlir/test/Dialect/OpenACC/acc-recipe-materialization-reduction.mlir
+++ b/mlir/test/Dialect/OpenACC/acc-recipe-materialization-reduction.mlir
@@ -51,3 +51,35 @@ func.func @par_reduction_clause_(%arg0: memref<f64>) {
   }
   return
 }
+
+// CHECK-LABEL: func.func @par_reduction_clause_serial
+// CHECK:       acc.parallel {{.*}} {
+// CHECK:       [[PRIVATE:%.*]] = acc.reduction_init {{.*}} <add>
+// CHECK-NEXT:  [[ZERO:%.*]] = arith.constant 0.000000e+00 : f64
+// CHECK-NEXT:  [[ALLOCA:%.*]] = memref.alloca() : memref<f64>
+// CHECK-NEXT:  memref.store [[ZERO]], [[ALLOCA]][]
+// CHECK-NEXT:  acc.yield {{.*}}
+// CHECK:       } {{.*}}acc.par_dims = #acc<par_dims[sequential]>, acc.var_name = #acc.var_name<"tmp">
+// CHECK:       memref.load [[PRIVATE]][]
+// CHECK:       memref.store {{.*}}, [[PRIVATE]][]
+// CHECK:       acc.reduction_combine_region [[PRIVATE]] into [[REDUCVAR:%.*]] : memref<f64> {
+// CHECK:       [[LOADVAR:%.*]] = memref.load [[REDUCVAR]][]
+// CHECK-NEXT:  [[LOADPRIV:%.*]] = memref.load [[PRIVATE]][]
+// CHECK-NEXT:  [[COMBINE:%.*]] = arith.addf [[LOADVAR]], [[LOADPRIV]]
+// CHECK-NEXT:  memref.store [[COMBINE]], [[REDUCVAR]][]
+// CHECK-NEXT:  } {acc.par_dims = #acc<par_dims[sequential]>}
+// CHECK-NEXT:  memref.dealloc [[PRIVATE]] : memref<f64>
+// CHECK:       acc.yield
+
+func.func @par_reduction_clause_serial(%arg0: memref<f64>) {
+  %c1_i32 = arith.constant 1 : i32
+  %cst = arith.constant 1.000000e+00 : f64
+  %red = acc.reduction varPtr(%arg0 : memref<f64>) recipe(@reduction_add_memref_f64) -> memref<f64> {name = "tmp"}
+  acc.parallel num_gangs({%c1_i32 : i32}) num_workers(%c1_i32 : i32) vector_length(%c1_i32 : i32) reduction(%red : memref<f64>) {
+    %3 = memref.load %red[] : memref<f64>
+    %4 = arith.addf %3, %cst fastmath<contract> : f64
+    memref.store %4, %red[] : memref<f64>
+    acc.yield
+  }
+  return
+}