[Mlir-commits] [mlir] [mlir][acc] Preserve section bounds on acc.reduction_init (PR #210443)

[Razvan Lupusoru]

https://github.com/razvanlupusoru updated https://github.com/llvm/llvm-project/pull/210443

>From 01f966dab9321dbc120fe9e4f12c73d5e059b266 Mon Sep 17 00:00:00 2001
From: Razvan Lupusoru <rlupusoru at nvidia.com>
Date: Fri, 17 Jul 2026 14:30:43 -0700
Subject: [PATCH 1/2] [mlir][acc] Preserve section bounds on acc.reduction_init

When `acc.reduction_init` is created from `acc.reduction`, its
bounds are lost because they only live as values folded into the
init region's body. Those bounds describe the legal indexing of
the section: the produced view of a (possibly base-adjusted,
compact) reduction temporary must be indexed with the same bounds
as the original array.

Carry the acc.bounds as optional operands on acc.reduction_init
so this indexing information is directly available. The operands
are optional, so scalar/whole-array reductions and existing
IR are unaffected.
---
 .../mlir/Dialect/OpenACC/OpenACCCGOps.td      | 21 +++++++++++++++++--
 .../Transforms/ACCRecipeMaterialization.cpp   |  4 +++-
 mlir/test/Dialect/OpenACC/ops-cg.mlir         | 14 +++++++++++++
 3 files changed, 36 insertions(+), 3 deletions(-)

diff --git a/mlir/include/mlir/Dialect/OpenACC/OpenACCCGOps.td b/mlir/include/mlir/Dialect/OpenACC/OpenACCCGOps.td
index 53adb91fc2983..730ea1f73a3ac 100644
--- a/mlir/include/mlir/Dialect/OpenACC/OpenACCCGOps.td
+++ b/mlir/include/mlir/Dialect/OpenACC/OpenACCCGOps.td
@@ -30,7 +30,7 @@ include "mlir/IR/CommonAttrConstraints.td"
 
 def OpenACC_ReductionInitOp
     : OpenACC_Op<"reduction_init",
-                 [SameOperandsAndResultType, RecursiveMemoryEffects,
+                 [AllTypesMatch<["var", "result"]>, RecursiveMemoryEffects,
                   DeclareOpInterfaceMethods<RegionBranchOpInterface,
                                             ["getRegionInvocationBounds",
                                              "getSuccessorInputs"]>,
@@ -44,14 +44,31 @@ def OpenACC_ReductionInitOp
 
     The `var` operand is the original/shared reduction variable. The
     `reduction_operator` specifies the reduction kind (e.g. add, mul).
+
+    The optional `bounds` operands describe the element range of the reduction
+    variable (as `acc.bounds` ops) when it refers to an array section. They
+    carry the section's bound information.
   }];
   let arguments = (ins OpenACC_AnyPointerOrMappableType:$var,
+                       Variadic<OpenACC_DataBoundsType>:$bounds,
                        OpenACC_ReductionOperatorAttr:$reductionOperator);
   let results = (outs OpenACC_AnyPointerOrMappableType:$result);
   let regions = (region AnyRegion:$region);
   let assemblyFormat = [{
-    $var $reductionOperator `:` type($result) $region attr-dict
+    $var (`bounds` `(` $bounds^ `)`)? $reductionOperator `:` type($result)
+    $region attr-dict
   }];
+  let builders = [
+    // Convenience builder for reductions without array-section bounds.
+    OpBuilder<(ins "::mlir::Value":$var,
+                   "::mlir::acc::ReductionOperatorAttr":$reductionOperator), [{
+      build($_builder, $_state, var, ::mlir::ValueRange{}, reductionOperator);
+    }]>,
+    OpBuilder<(ins "::mlir::Value":$var,
+                   "::mlir::acc::ReductionOperator":$reductionOperator), [{
+      build($_builder, $_state, var, ::mlir::ValueRange{}, reductionOperator);
+    }]>,
+  ];
   let hasVerifier = 1;
 }
 
diff --git a/mlir/lib/Dialect/OpenACC/Transforms/ACCRecipeMaterialization.cpp b/mlir/lib/Dialect/OpenACC/Transforms/ACCRecipeMaterialization.cpp
index e666bcd9a6242..856ecba953b7a 100644
--- a/mlir/lib/Dialect/OpenACC/Transforms/ACCRecipeMaterialization.cpp
+++ b/mlir/lib/Dialect/OpenACC/Transforms/ACCRecipeMaterialization.cpp
@@ -344,8 +344,10 @@ LogicalResult ACCRecipeMaterialization::materialize(
     else
       llvm_unreachable("unexpected acc op with reduction recipe");
 
+    SmallVector<Value> reductionBounds(acc::getBounds(op));
     auto reductionOp = acc::ReductionInitOp::create(
-        b, op.getLoc(), origPtr, recipe.getReductionOperatorAttr());
+        b, op.getLoc(), origPtr, reductionBounds,
+        recipe.getReductionOperatorAttr());
     saveVarName(op.getAccVar(), reductionOp.getResult());
     cloneRegionIntoAccRegion(&initRegion, &reductionOp.getRegion(),
                              /*hasResult=*/true);
diff --git a/mlir/test/Dialect/OpenACC/ops-cg.mlir b/mlir/test/Dialect/OpenACC/ops-cg.mlir
index 765552e8577a7..cf1c3340870a5 100644
--- a/mlir/test/Dialect/OpenACC/ops-cg.mlir
+++ b/mlir/test/Dialect/OpenACC/ops-cg.mlir
@@ -160,6 +160,20 @@ func.func @compute_region_two_dims(%data: memref<8xi32>,
 
 // -----
 
+// CHECK-LABEL: func @reduction_init_with_bounds
+func.func @reduction_init_with_bounds(%arg0: memref<?xi32>, %lb: index, %ext: index) {
+  %c1 = arith.constant 1 : index
+  %bnd = acc.bounds lowerbound(%lb : index) extent(%ext : index) stride(%c1 : index)
+  %0 = acc.reduction_init %arg0 bounds(%bnd) <add> : memref<?xi32> {
+    acc.yield %arg0 : memref<?xi32>
+  }
+  return
+}
+// CHECK: %[[BND:.*]] = acc.bounds
+// CHECK: acc.reduction_init %{{.*}} bounds(%[[BND]]) <add> : memref<?xi32>
+
+// -----
+
 // CHECK-LABEL: func @compute_region_unknown_width
 func.func @compute_region_unknown_width(%data: memref<100xf32>) {
   %copyin = acc.copyin varPtr(%data : memref<100xf32>) -> memref<100xf32>

>From f7b59b656b807f23c47246cdf5c4e9631019758f Mon Sep 17 00:00:00 2001
From: Razvan Lupusoru <rlupusoru at nvidia.com>
Date: Fri, 17 Jul 2026 14:50:47 -0700
Subject: [PATCH 2/2] Fix formatting

---
 .../Dialect/OpenACC/Transforms/ACCRecipeMaterialization.cpp | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/mlir/lib/Dialect/OpenACC/Transforms/ACCRecipeMaterialization.cpp b/mlir/lib/Dialect/OpenACC/Transforms/ACCRecipeMaterialization.cpp
index 856ecba953b7a..70aacd251c329 100644
--- a/mlir/lib/Dialect/OpenACC/Transforms/ACCRecipeMaterialization.cpp
+++ b/mlir/lib/Dialect/OpenACC/Transforms/ACCRecipeMaterialization.cpp
@@ -345,9 +345,9 @@ LogicalResult ACCRecipeMaterialization::materialize(
       llvm_unreachable("unexpected acc op with reduction recipe");
 
     SmallVector<Value> reductionBounds(acc::getBounds(op));
-    auto reductionOp = acc::ReductionInitOp::create(
-        b, op.getLoc(), origPtr, reductionBounds,
-        recipe.getReductionOperatorAttr());
+    auto reductionOp =
+        acc::ReductionInitOp::create(b, op.getLoc(), origPtr, reductionBounds,
+                                     recipe.getReductionOperatorAttr());
     saveVarName(op.getAccVar(), reductionOp.getResult());
     cloneRegionIntoAccRegion(&initRegion, &reductionOp.getRegion(),
                              /*hasResult=*/true);