[Mlir-commits] [mlir] [mlir][OpenMP] add missing load for reduction cleanup region (PR #88289)

[Tom Eccles]

https://github.com/tblah created https://github.com/llvm/llvm-project/pull/88289

I missed this before. For by-ref reductions, the private reduction variable is a pointer to the pointer to the variable. So an extra load is required to get the right value.

See the "red.private.value.n" loads in the reduction combiner region for reference.

>From e161d1280a5e5db465abd256eee974a524daa4fa Mon Sep 17 00:00:00 2001
From: Tom Eccles <tom.eccles at arm.com>
Date: Wed, 10 Apr 2024 15:09:16 +0000
Subject: [PATCH] [mlir][OpenMP] add missing load for reduction cleanup region

I missed this before. For by-ref reductions, the private reduction
variable is a pointer to the pointer to the variable. So an extra load
is required to get the right value.

See the "red.private.value.n" loads in the reduction combiner region for
reference.
---
 .../Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp    | 13 +++++++++++--
 .../LLVMIR/openmp-parallel-reduction-cleanup.mlir   |  6 ++++--
 .../LLVMIR/openmp-wsloop-reduction-cleanup.mlir     |  6 ++++--
 3 files changed, 19 insertions(+), 6 deletions(-)

diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index a59677c02fc392..300fc8ba56fc50 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -889,8 +889,17 @@ static LogicalResult inlineReductionCleanup(
 
     // map the argument to the cleanup region
     Block &entry = cleanupRegion.front();
-    moduleTranslation.mapValue(entry.getArgument(0),
-                               privateReductionVariables[i]);
+
+    llvm::Instruction *potentialTerminator =
+        builder.GetInsertBlock()->empty() ? nullptr
+                                          : &builder.GetInsertBlock()->back();
+    if (potentialTerminator && potentialTerminator->isTerminator())
+      builder.SetInsertPoint(potentialTerminator);
+    llvm::Value *reductionVar = builder.CreateLoad(
+        moduleTranslation.convertType(entry.getArgument(0).getType()),
+        privateReductionVariables[i]);
+
+    moduleTranslation.mapValue(entry.getArgument(0), reductionVar);
 
     if (failed(inlineConvertOmpRegions(cleanupRegion, "omp.reduction.cleanup",
                                        builder, moduleTranslation)))
diff --git a/mlir/test/Target/LLVMIR/openmp-parallel-reduction-cleanup.mlir b/mlir/test/Target/LLVMIR/openmp-parallel-reduction-cleanup.mlir
index 9ae4c4ad392b13..b7f71f438e56b0 100644
--- a/mlir/test/Target/LLVMIR/openmp-parallel-reduction-cleanup.mlir
+++ b/mlir/test/Target/LLVMIR/openmp-parallel-reduction-cleanup.mlir
@@ -86,8 +86,10 @@
 
 // Cleanup region:
 // CHECK: [[OMP_FINALIZE]]:
-// CHECK: call void @free(ptr %[[PRIV_PTR_I]])
-// CHECK: call void @free(ptr %[[PRIV_PTR_J]])
+// CHECK: %[[PRIV_I:.+]] = load ptr, ptr %[[PRIV_PTR_I]], align 8
+// CHECK: call void @free(ptr %[[PRIV_I]])
+// CHECK: %[[PRIV_J:.+]] = load ptr, ptr %[[PRIV_PTR_J]], align 8
+// CHECK: call void @free(ptr %[[PRIV_J]])
 
 // Reduction function.
 // CHECK: define internal void @[[REDFUNC]]
diff --git a/mlir/test/Target/LLVMIR/openmp-wsloop-reduction-cleanup.mlir b/mlir/test/Target/LLVMIR/openmp-wsloop-reduction-cleanup.mlir
index a1e17afa53e215..3842522934e48e 100644
--- a/mlir/test/Target/LLVMIR/openmp-wsloop-reduction-cleanup.mlir
+++ b/mlir/test/Target/LLVMIR/openmp-wsloop-reduction-cleanup.mlir
@@ -63,8 +63,10 @@
 // Weirdly the finalization block is generated before the reduction blocks:
 // CHECK: [[FINALIZE:.+]]:
 // CHECK: call void @__kmpc_barrier
-// CHECK: call void @free(ptr %[[PRIV_PTR_I]])
-// CHECK: call void @free(ptr %[[PRIV_PTR_J]])
+// CHECK: %[[PRIV_I:.+]] = load ptr, ptr %[[PRIV_PTR_I]], align 8
+// CHECK: call void @free(ptr %[[PRIV_I]])
+// CHECK: %[[PRIV_J:.+]] = load ptr, ptr %[[PRIV_PTR_J]], align 8
+// CHECK: call void @free(ptr %[[PRIV_J]])
 // CHECK: ret void
 
 // Non-atomic reduction: