[Mlir-commits] [mlir] [ACCRecipeMaterialization] materialize destroy region for reductions (PR #203935)

[Scott Manley]

https://github.com/rscottmanley created https://github.com/llvm/llvm-project/pull/203935

Materialize the destroy region when the reduction recipe has one

>From 62c1c0f38a4735b954589d76d52534bdc88456d6 Mon Sep 17 00:00:00 2001
From: Scott Manley <scmanley at nvidia.com>
Date: Mon, 15 Jun 2026 09:00:33 -0700
Subject: [PATCH] [ACCRecipeMaterialization] materialize destroy region for
 reductions

Materialize the destroy region when the reduction recipe has one
---
 .../Transforms/ACCRecipeMaterialization.cpp       | 15 +++++++--------
 .../acc-recipe-materialization-reduction.mlir     |  8 +++++++-
 2 files changed, 14 insertions(+), 9 deletions(-)

diff --git a/mlir/lib/Dialect/OpenACC/Transforms/ACCRecipeMaterialization.cpp b/mlir/lib/Dialect/OpenACC/Transforms/ACCRecipeMaterialization.cpp
index b60103e9cd85f..ea4db8b6ab2f8 100644
--- a/mlir/lib/Dialect/OpenACC/Transforms/ACCRecipeMaterialization.cpp
+++ b/mlir/lib/Dialect/OpenACC/Transforms/ACCRecipeMaterialization.cpp
@@ -116,6 +116,7 @@ static void saveVarName(Value src, Value dst) {
 // according to the recipe specifications.
 template <typename RecipeOpTy>
 static void cloneDestroy(RecipeOpTy recipe, mlir::Block *block,
+                         Block::iterator ip,
                          const llvm::SmallVector<mlir::Value> &arguments) {
   IRMapping mapping{};
   Region &destroyRegion = recipe.getDestroyRegion();
@@ -123,8 +124,7 @@ static void cloneDestroy(RecipeOpTy recipe, mlir::Block *block,
              arguments.size() &&
          "unexpected acc recipe destroy block arguments");
   mapping.map(destroyRegion.getBlocks().front().getArguments(), arguments);
-  acc::cloneACCRegionInto(&destroyRegion, block, std::prev(block->end()),
-                          mapping,
+  acc::cloneACCRegionInto(&destroyRegion, block, ip, mapping,
                           /*resultsToReplace=*/{});
 }
 
@@ -263,7 +263,7 @@ ACCRecipeMaterialization::materialize(OpTy op, RecipeOpTy recipe, AccOpTy accOp,
     if (!recipe.getDestroyRegion().empty()) {
       results.insert(results.begin(), origPtr);
       results.append(triples);
-      cloneDestroy(recipe, block, results);
+      cloneDestroy(recipe, block, block->end(), results);
     }
   } else if constexpr (std::is_same_v<OpTy, acc::FirstprivateOp>) {
     // Clone the init region for a firstprivate.
@@ -284,7 +284,7 @@ ACCRecipeMaterialization::materialize(OpTy op, RecipeOpTy recipe, AccOpTy accOp,
                             mapping, {});
     if (!recipe.getDestroyRegion().empty()) {
       // origPtr was already pushed.
-      cloneDestroy(recipe, block, results);
+      cloneDestroy(recipe, block, block->end(), results);
     }
   } else if constexpr (std::is_same_v<OpTy, acc::ReductionOp>) {
     auto cloneRegionIntoAccRegion = [&](Region *src, Region *dest,
@@ -353,10 +353,9 @@ ACCRecipeMaterialization::materialize(OpTy op, RecipeOpTy recipe, AccOpTy accOp,
     setSeqParDimsForRecipeLoops(&combineRegionOp.getRegion());
 
     if (!recipe.getDestroyRegion().empty()) {
-      (void)accSupport.emitNYI(
-          recipe.getLoc(),
-          "OpenACC reduction variable that requires destruction code");
-      return failure();
+      SmallVector<Value> results{origPtr, reductionOp.getResult()};
+      Block::iterator ip = std::next(Block::iterator(combineRegionOp));
+      cloneDestroy(recipe, combineRegionOp->getBlock(), ip, results);
     }
   } else {
     llvm_unreachable("unexpected op type");
diff --git a/mlir/test/Dialect/OpenACC/acc-recipe-materialization-reduction.mlir b/mlir/test/Dialect/OpenACC/acc-recipe-materialization-reduction.mlir
index 9b023eede5bdc..083a6839b7d0a 100644
--- a/mlir/test/Dialect/OpenACC/acc-recipe-materialization-reduction.mlir
+++ b/mlir/test/Dialect/OpenACC/acc-recipe-materialization-reduction.mlir
@@ -13,6 +13,10 @@ acc.reduction.recipe @reduction_add_memref_f64 : memref<f64> reduction_operator
   %2 = arith.addf %0, %1 fastmath<contract> : f64
   memref.store %2, %arg0[] : memref<f64>
   acc.yield %arg0 : memref<f64>
+} destroy {
+^bb0(%arg0: memref<f64>, %arg1: memref<f64>):
+  memref.dealloc %arg1 : memref<f64>
+  acc.terminator
 }
 
 // Verify that the reduction init and combine recipes attached to compute
@@ -27,11 +31,13 @@ acc.reduction.recipe @reduction_add_memref_f64 : memref<f64> reduction_operator
 // CHECK:       } {{.*}}acc.var_name = #acc.var_name<"tmp">
 // CHECK:       memref.load [[PRIVATE]][]
 // CHECK:       memref.store {{.*}}, [[PRIVATE]][]
-// CHECK:       acc.reduction_combine_region [[PRIVATE]] into [[REDUCVAR:%.*]] :
+// 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:  }
+// CHECK-NEXT:  memref.dealloc [[PRIVATE]] : memref<f64>
 // CHECK:       acc.yield
 
 func.func @par_reduction_clause_(%arg0: memref<f64>) {