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

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-openmp

@llvm/pr-subscribers-flang-openmp

Author: Tom Eccles (tblah)

<details>
<summary>Changes</summary>

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.

---
Full diff: https://github.com/llvm/llvm-project/pull/88289.diff


3 Files Affected:

- (modified) mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp (+11-2) 
- (modified) mlir/test/Target/LLVMIR/openmp-parallel-reduction-cleanup.mlir (+4-2) 
- (modified) mlir/test/Target/LLVMIR/openmp-wsloop-reduction-cleanup.mlir (+4-2) 


``````````diff
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:

``````````

</details>


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