[Mlir-commits] [flang] [llvm] [mlir] [MLIR][OpenMP] Handle privatization for global values in MLIR->LLVM translation (PR #104407)
Kareem Ergawy
llvmlistbot at llvm.org
Fri Aug 16 01:47:46 PDT 2024
https://github.com/ergawy updated https://github.com/llvm/llvm-project/pull/104407
>From 27639c8472e55af9e7a394d1af40a73ed1bfef33 Mon Sep 17 00:00:00 2001
From: ergawy <kareem.ergawy at amd.com>
Date: Wed, 14 Aug 2024 03:40:59 -0500
Subject: [PATCH 1/2] [MLIR][OpenMP] Fix MLIR->LLVM value matching in
privatization logic
Fixes #102935
Updates matching logic for finding the LLVM value that corresponds to
an MLIR value. We need that matching to find the delayed privatizer for
an LLVM value being privatized.
The issue occures when there is an "indirect" correspondence between
MLIR and LLVM values: in some cases the values we are trying to match
stem from a pair of load/store ops that point to the same memref. This
PR adds such matching logic.
---
.../OpenMP/OpenMPToLLVMIRTranslation.cpp | 129 ++++++++++++++----
.../Target/LLVMIR/openmp-firstprivate.mlir | 42 ++++++
2 files changed, 143 insertions(+), 28 deletions(-)
diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index 458d05d5059db7..82a1e322b1b2b7 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -690,14 +690,14 @@ inlineOmpRegionCleanup(llvm::SmallVectorImpl<Region *> &cleanupRegions,
: &builder.GetInsertBlock()->back();
if (potentialTerminator && potentialTerminator->isTerminator())
builder.SetInsertPoint(potentialTerminator);
- llvm::Value *prviateVarValue =
+ llvm::Value *privateVarValue =
shouldLoadCleanupRegionArg
? builder.CreateLoad(
moduleTranslation.convertType(entry.getArgument(0).getType()),
privateVariables[i])
: privateVariables[i];
- moduleTranslation.mapValue(entry.getArgument(0), prviateVarValue);
+ moduleTranslation.mapValue(entry.getArgument(0), privateVarValue);
if (failed(inlineConvertOmpRegions(*cleanupRegion, regionName, builder,
moduleTranslation)))
@@ -1424,35 +1424,107 @@ convertOmpParallel(omp::ParallelOp opInst, llvm::IRBuilderBase &builder,
}
};
- SmallVector<omp::PrivateClauseOp> privatizerClones;
- SmallVector<llvm::Value *> privateVariables;
+ SmallVector<omp::PrivateClauseOp> mlirPrivatizerClones;
+ SmallVector<llvm::Value *> llvmPrivateVars;
// TODO: Perform appropriate actions according to the data-sharing
// attribute (shared, private, firstprivate, ...) of variables.
// Currently shared and private are supported.
auto privCB = [&](InsertPointTy allocaIP, InsertPointTy codeGenIP,
- llvm::Value &, llvm::Value &vPtr,
- llvm::Value *&replacementValue) -> InsertPointTy {
- replacementValue = &vPtr;
+ llvm::Value &, llvm::Value &llvmOmpRegionInputPtr,
+ llvm::Value *&llvmReplacementValue) -> InsertPointTy {
+ llvmReplacementValue = &llvmOmpRegionInputPtr;
// If this is a private value, this lambda will return the corresponding
// mlir value and its `PrivateClauseOp`. Otherwise, empty values are
// returned.
- auto [privVar, privatizerClone] =
+ auto [mlirPrivVar, mlirPrivatizerClone] =
[&]() -> std::pair<mlir::Value, omp::PrivateClauseOp> {
if (!opInst.getPrivateVars().empty()) {
- auto privateVars = opInst.getPrivateVars();
- auto privateSyms = opInst.getPrivateSyms();
+ auto mlirPrivVars = opInst.getPrivateVars();
+ auto mlirPrivSyms = opInst.getPrivateSyms();
- for (auto [privVar, privatizerAttr] :
- llvm::zip_equal(privateVars, *privateSyms)) {
+ // Try to find a privatizer that corresponds to the LLVM value being
+ // prvatized.
+ for (auto [mlirPrivVar, mlirPrivatizerAttr] :
+ llvm::zip_equal(mlirPrivVars, *mlirPrivSyms)) {
// Find the MLIR private variable corresponding to the LLVM value
// being privatized.
- llvm::Value *llvmPrivVar = moduleTranslation.lookupValue(privVar);
- if (llvmPrivVar != &vPtr)
+ llvm::Value *mlirToLLVMPrivVar =
+ moduleTranslation.lookupValue(mlirPrivVar);
+
+ // Check if the LLVM value being privatized matches the LLVM value
+ // mapped to privVar. In some cases, this is not trivial ...
+ auto isMatch = [&]() {
+ if (mlirToLLVMPrivVar == nullptr)
+ return false;
+
+ // If both values are trivially equal, we found a match.
+ if (mlirToLLVMPrivVar == &llvmOmpRegionInputPtr)
+ return true;
+
+ // Otherwise, we check if both llvmOmpRegionInputPtr and
+ // mlirToLLVMPrivVar refer to the same memory (through a load/store
+ // pair). This happens if a struct (i.e. multi-field value) is being
+ // privatized.
+ //
+ // For example, if the privatized value is defined by:
+ // ```
+ // %priv_val = alloca { ptr, i64 }, align 8
+ // ```
+ //
+ // Initialized this value (outside the omp region) will be something
+ // like this:
+ //
+ // clang-format off
+ // ```
+ // %partially_init_priv_val = insertvalue { ptr, i64 } undef,
+ // ptr %some_ptr, 0
+ // %fully_init_priv_val = insertvalue { ptr, i64 } %partially_init_priv_val,
+ // i64 %some_i64, 1
+ // ...
+ // store { ptr, i64 } %fully_init_priv_val, ptr %priv_val, align 8
+ // ```
+ // clang-format on
+ //
+ // Now, we enter the OMP region, in order to access this privatized
+ // value, we need to load from the allocated memory:
+ // ```
+ // omp.par.entry:
+ // %priv_val_load = load { ptr, i64 }, ptr %priv_val, align 8
+ // ```
+ //
+ // The 2 LLVM values tracked here map as follows:
+ // - `mlirToLLVMPrivVar` -> `%fully_init_priv_val`
+ // - `llvmOmpRegionInputPtr` -> `priv_val_load`
+ //
+ // Even though they eventually refer to the same memory reference
+ // (the memory being privatized), they are 2 distinct LLVM values.
+ // Therefore, we need to discover their correspondence by finding
+ // out if the store into and load from the same mem ref.
+ auto *llvmOmpRegionInputPtrLoad =
+ llvm::dyn_cast_if_present<llvm::LoadInst>(
+ &llvmOmpRegionInputPtr);
+
+ if (llvmOmpRegionInputPtrLoad == nullptr)
+ return false;
+
+ for (auto &use : mlirToLLVMPrivVar->uses()) {
+ auto *mlirToLLVMPrivVarStore =
+ llvm::dyn_cast_if_present<llvm::StoreInst>(use.getUser());
+ if (mlirToLLVMPrivVarStore &&
+ (llvmOmpRegionInputPtrLoad->getPointerOperand() ==
+ mlirToLLVMPrivVarStore->getPointerOperand()))
+ return true;
+ }
+
+ return false;
+ };
+
+ if (!isMatch())
continue;
- SymbolRefAttr privSym = llvm::cast<SymbolRefAttr>(privatizerAttr);
+ SymbolRefAttr privSym = llvm::cast<SymbolRefAttr>(mlirPrivatizerAttr);
omp::PrivateClauseOp privatizer =
SymbolTable::lookupNearestSymbolFrom<omp::PrivateClauseOp>(
opInst, privSym);
@@ -1480,24 +1552,24 @@ convertOmpParallel(omp::ParallelOp opInst, llvm::IRBuilderBase &builder,
counter);
clone.setSymName(cloneName);
- return {privVar, clone};
+ return {mlirPrivVar, clone};
}
}
return {mlir::Value(), omp::PrivateClauseOp()};
}();
- if (privVar) {
- Region &allocRegion = privatizerClone.getAllocRegion();
+ if (mlirPrivVar) {
+ Region &allocRegion = mlirPrivatizerClone.getAllocRegion();
// If this is a `firstprivate` clause, prepare the `omp.private` op by:
- if (privatizerClone.getDataSharingType() ==
+ if (mlirPrivatizerClone.getDataSharingType() ==
omp::DataSharingClauseType::FirstPrivate) {
auto oldAllocBackBlock = std::prev(allocRegion.end());
omp::YieldOp oldAllocYieldOp =
llvm::cast<omp::YieldOp>(oldAllocBackBlock->getTerminator());
- Region ©Region = privatizerClone.getCopyRegion();
+ Region ©Region = mlirPrivatizerClone.getCopyRegion();
mlir::IRRewriter copyCloneBuilder(&moduleTranslation.getContext());
// 1. Cloning the `copy` region to the end of the `alloc` region.
@@ -1524,7 +1596,7 @@ convertOmpParallel(omp::ParallelOp opInst, llvm::IRBuilderBase &builder,
// This way, the body of the privatizer will be changed from using the
// region/block argument to the value being privatized.
auto allocRegionArg = allocRegion.getArgument(0);
- replaceAllUsesInRegionWith(allocRegionArg, privVar, allocRegion);
+ replaceAllUsesInRegionWith(allocRegionArg, mlirPrivVar, allocRegion);
auto oldIP = builder.saveIP();
builder.restoreIP(allocaIP);
@@ -1535,15 +1607,15 @@ convertOmpParallel(omp::ParallelOp opInst, llvm::IRBuilderBase &builder,
opInst.emitError("failed to inline `alloc` region of an `omp.private` "
"op in the parallel region");
bodyGenStatus = failure();
- privatizerClone.erase();
+ mlirPrivatizerClone.erase();
} else {
assert(yieldedValues.size() == 1);
- replacementValue = yieldedValues.front();
+ llvmReplacementValue = yieldedValues.front();
// Keep the LLVM replacement value and the op clone in case we need to
// emit cleanup (i.e. deallocation) logic.
- privateVariables.push_back(replacementValue);
- privatizerClones.push_back(privatizerClone);
+ llvmPrivateVars.push_back(llvmReplacementValue);
+ mlirPrivatizerClones.push_back(mlirPrivatizerClone);
}
builder.restoreIP(oldIP);
@@ -1571,13 +1643,14 @@ convertOmpParallel(omp::ParallelOp opInst, llvm::IRBuilderBase &builder,
bodyGenStatus = failure();
SmallVector<Region *> privateCleanupRegions;
- llvm::transform(privatizerClones, std::back_inserter(privateCleanupRegions),
+ llvm::transform(mlirPrivatizerClones,
+ std::back_inserter(privateCleanupRegions),
[](omp::PrivateClauseOp privatizer) {
return &privatizer.getDeallocRegion();
});
if (failed(inlineOmpRegionCleanup(
- privateCleanupRegions, privateVariables, moduleTranslation, builder,
+ privateCleanupRegions, llvmPrivateVars, moduleTranslation, builder,
"omp.private.dealloc", /*shouldLoadCleanupRegionArg=*/false)))
bodyGenStatus = failure();
@@ -1604,7 +1677,7 @@ convertOmpParallel(omp::ParallelOp opInst, llvm::IRBuilderBase &builder,
ompBuilder->createParallel(ompLoc, allocaIP, bodyGenCB, privCB, finiCB,
ifCond, numThreads, pbKind, isCancellable));
- for (mlir::omp::PrivateClauseOp privatizerClone : privatizerClones)
+ for (mlir::omp::PrivateClauseOp privatizerClone : mlirPrivatizerClones)
privatizerClone.erase();
return bodyGenStatus;
diff --git a/mlir/test/Target/LLVMIR/openmp-firstprivate.mlir b/mlir/test/Target/LLVMIR/openmp-firstprivate.mlir
index 65ae98b2a74c6e..b06ad96f4592c5 100644
--- a/mlir/test/Target/LLVMIR/openmp-firstprivate.mlir
+++ b/mlir/test/Target/LLVMIR/openmp-firstprivate.mlir
@@ -114,3 +114,45 @@ omp.private {type = firstprivate} @multi_block.privatizer : !llvm.ptr alloc {
llvm.store %arg2, %arg3 : f32, !llvm.ptr
omp.yield(%arg3 : !llvm.ptr)
}
+
+// -----
+
+// Verifies fix for https://github.com/llvm/llvm-project/issues/102935.
+//
+// The issue happens since we previously failed to match MLIR values to their
+// corresponding LLVM values in some cases (e.g. char strings with non-const
+// length).
+llvm.func @non_const_len_char_test(%n: !llvm.ptr {fir.bindc_name = "n"}) {
+ %n_val = llvm.load %n : !llvm.ptr -> i64
+ %orig_alloc = llvm.alloca %n_val x i8 {bindc_name = "str"} : (i64) -> !llvm.ptr
+ %orig_val = llvm.mlir.undef : !llvm.struct<(ptr, i64)>
+ %orig_val_init = llvm.insertvalue %orig_alloc, %orig_val[0] : !llvm.struct<(ptr, i64)>
+ omp.parallel private(@non_const_len_char %orig_val_init -> %priv_arg : !llvm.struct<(ptr, i64)>) {
+ %dummy = llvm.extractvalue %priv_arg[0] : !llvm.struct<(ptr, i64)>
+ omp.terminator
+ }
+ llvm.return
+}
+
+omp.private {type = firstprivate} @non_const_len_char : !llvm.struct<(ptr, i64)> alloc {
+^bb0(%orig_val: !llvm.struct<(ptr, i64)>):
+ %str_len = llvm.extractvalue %orig_val[1] : !llvm.struct<(ptr, i64)>
+ %priv_alloc = llvm.alloca %str_len x i8 {bindc_name = "str", pinned} : (i64) -> !llvm.ptr
+ %priv_val = llvm.mlir.undef : !llvm.struct<(ptr, i64)>
+ %priv_val_init = llvm.insertvalue %priv_alloc, %priv_val[0] : !llvm.struct<(ptr, i64)>
+ omp.yield(%priv_val_init : !llvm.struct<(ptr, i64)>)
+} copy {
+^bb0(%orig_val: !llvm.struct<(ptr, i64)>, %priv_val: !llvm.struct<(ptr, i64)>):
+ llvm.call @foo() : () -> ()
+ omp.yield(%priv_val : !llvm.struct<(ptr, i64)>)
+}
+
+llvm.func @foo()
+
+// CHECK-LABEL: @non_const_len_char_test..omp_par({{.*}})
+// CHECK-NEXT: omp.par.entry:
+// Verify that we found the privatizer by checking that we properly inlined the
+// bodies of the alloc and copy regions.
+// CHECK: %[[STR_LEN:.*]] = extractvalue { ptr, i64 } %{{.*}}, 1
+// CHECK: %{{.*}} = alloca i8, i64 %[[STR_LEN]], align 1
+// CHECK: call void @foo()
>From 5d93023a8d87d91d3e2e1afeacac62e00a5a2586 Mon Sep 17 00:00:00 2001
From: ergawy <kareem.ergawy at amd.com>
Date: Thu, 15 Aug 2024 01:16:14 -0500
Subject: [PATCH 2/2] [MLIR][OpenMP] Handle privatization for global values in
MLIR->LLVM translation
Fix for https://github.com/llvm/llvm-project/issues/102939.
The issues occurs because the CodeExtractor component only collect inputs
(to the parallel regions) that are defined in the same function in which the
parallel regions is present. Howerver, this is problematic because if we are
privatizing a global value (e.g. a `target` variable which is emitted as a
global), then we miss finding that input and we do not privatize the
variable.
This commit attempts to fix the issue by adding a flag to the
CodeExtractor so that we can collect global inputs.
---
...privatization-allocatable-firstprivate.f90 | 5 +-
.../llvm/Transforms/Utils/CodeExtractor.h | 3 +-
llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp | 4 +-
llvm/lib/Transforms/Utils/CodeExtractor.cpp | 7 ++-
.../Target/LLVMIR/openmp-firstprivate.mlir | 46 +++++++++++++++++++
5 files changed, 58 insertions(+), 7 deletions(-)
diff --git a/flang/test/Lower/OpenMP/delayed-privatization-allocatable-firstprivate.f90 b/flang/test/Lower/OpenMP/delayed-privatization-allocatable-firstprivate.f90
index 833976ff284a86..5f09371bbaba2e 100644
--- a/flang/test/Lower/OpenMP/delayed-privatization-allocatable-firstprivate.f90
+++ b/flang/test/Lower/OpenMP/delayed-privatization-allocatable-firstprivate.f90
@@ -57,6 +57,5 @@ end program compilation_to_obj
! LLVM: @[[GLOB_VAR:[^[:space:]]+]]t = internal global
! LLVM: define internal void @_QQmain..omp_par
-! LLVM: %[[LOCAL_VAR:.*]] = alloca { ptr, i64, i32, i8, i8, i8, i8, [1 x [3 x i64]] }, align 8
-! LLVM-NEXT: %[[GLOB_VAL:.*]] = load { ptr, i64, i32, i8, i8, i8, i8, [1 x [3 x i64]] }, ptr @[[GLOB_VAR]]t, align 8
-! LLVM-NEXT: store { ptr, i64, i32, i8, i8, i8, i8, [1 x [3 x i64]] } %[[GLOB_VAL]], ptr %[[LOCAL_VAR]], align 8
+! LLVM: %[[GLOB_VAL:.*]] = load { ptr, i64, i32, i8, i8, i8, i8, [1 x [3 x i64]] }, ptr @[[GLOB_VAR]]t, align 8
+! LLVM-NEXT: store { ptr, i64, i32, i8, i8, i8, i8, [1 x [3 x i64]] } %[[GLOB_VAL]], ptr %{{.*}}, align 8
diff --git a/llvm/include/llvm/Transforms/Utils/CodeExtractor.h b/llvm/include/llvm/Transforms/Utils/CodeExtractor.h
index 68eb00a50fe030..826347e79f7195 100644
--- a/llvm/include/llvm/Transforms/Utils/CodeExtractor.h
+++ b/llvm/include/llvm/Transforms/Utils/CodeExtractor.h
@@ -187,7 +187,8 @@ class CodeExtractorAnalysisCache {
/// sets, before extraction occurs. These modifications won't have any
/// significant impact on the cost however.
void findInputsOutputs(ValueSet &Inputs, ValueSet &Outputs,
- const ValueSet &Allocas) const;
+ const ValueSet &Allocas,
+ bool CollectGlobalInputs = false) const;
/// Check if life time marker nodes can be hoisted/sunk into the outline
/// region.
diff --git a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
index f9b070e6f1eae4..2ea37f33dab45f 100644
--- a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
+++ b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
@@ -1542,7 +1542,9 @@ IRBuilder<>::InsertPoint OpenMPIRBuilder::createParallel(
BasicBlock *CommonExit = nullptr;
SetVector<Value *> Inputs, Outputs, SinkingCands, HoistingCands;
Extractor.findAllocas(CEAC, SinkingCands, HoistingCands, CommonExit);
- Extractor.findInputsOutputs(Inputs, Outputs, SinkingCands);
+
+ Extractor.findInputsOutputs(Inputs, Outputs, SinkingCands,
+ /*CollectGlobalInputs=*/true);
LLVM_DEBUG(dbgs() << "Before privatization: " << *OuterFn << "\n");
diff --git a/llvm/lib/Transforms/Utils/CodeExtractor.cpp b/llvm/lib/Transforms/Utils/CodeExtractor.cpp
index 81d3243c887fce..d378c6c3a4b01c 100644
--- a/llvm/lib/Transforms/Utils/CodeExtractor.cpp
+++ b/llvm/lib/Transforms/Utils/CodeExtractor.cpp
@@ -632,14 +632,17 @@ bool CodeExtractor::isEligible() const {
}
void CodeExtractor::findInputsOutputs(ValueSet &Inputs, ValueSet &Outputs,
- const ValueSet &SinkCands) const {
+ const ValueSet &SinkCands,
+ bool CollectGlobalInputs) const {
for (BasicBlock *BB : Blocks) {
// If a used value is defined outside the region, it's an input. If an
// instruction is used outside the region, it's an output.
for (Instruction &II : *BB) {
for (auto &OI : II.operands()) {
Value *V = OI;
- if (!SinkCands.count(V) && definedInCaller(Blocks, V))
+ if (!SinkCands.count(V) &&
+ (definedInCaller(Blocks, V) ||
+ (CollectGlobalInputs && llvm::isa<llvm::GlobalVariable>(V))))
Inputs.insert(V);
}
diff --git a/mlir/test/Target/LLVMIR/openmp-firstprivate.mlir b/mlir/test/Target/LLVMIR/openmp-firstprivate.mlir
index b06ad96f4592c5..02ce6b5b19ceaf 100644
--- a/mlir/test/Target/LLVMIR/openmp-firstprivate.mlir
+++ b/mlir/test/Target/LLVMIR/openmp-firstprivate.mlir
@@ -156,3 +156,49 @@ llvm.func @foo()
// CHECK: %[[STR_LEN:.*]] = extractvalue { ptr, i64 } %{{.*}}, 1
// CHECK: %{{.*}} = alloca i8, i64 %[[STR_LEN]], align 1
// CHECK: call void @foo()
+
+// -----
+
+// Verifies fix for https://github.com/llvm/llvm-project/issues/102939.
+//
+// The issues occurs because the CodeExtractor component only collect inputs
+// (to the parallel regions) that are defined in the same function in which the
+// parallel regions is present. Howerver, this is problematic because if we are
+// privatizing a global value (e.g. a `target` variable which is emitted as a
+// global), then we miss finding that input and we do not privatize the
+// variable.
+
+omp.private {type = firstprivate} @global_privatizer : !llvm.ptr alloc {
+^bb0(%arg0: !llvm.ptr):
+ %0 = llvm.mlir.constant(1 : i64) : i64
+ %1 = llvm.alloca %0 x f32 {bindc_name = "global", pinned} : (i64) -> !llvm.ptr
+ omp.yield(%1 : !llvm.ptr)
+} copy {
+^bb0(%arg0: !llvm.ptr, %arg1: !llvm.ptr):
+ %0 = llvm.load %arg0 : !llvm.ptr -> f32
+ llvm.store %0, %arg1 : f32, !llvm.ptr
+ omp.yield(%arg1 : !llvm.ptr)
+}
+
+llvm.func @global_accessor() {
+ %global_addr = llvm.mlir.addressof @global : !llvm.ptr
+ omp.parallel private(@global_privatizer %global_addr -> %arg0 : !llvm.ptr) {
+ %1 = llvm.mlir.constant(3.140000e+00 : f32) : f32
+ llvm.store %1, %arg0 : f32, !llvm.ptr
+ omp.terminator
+ }
+ llvm.return
+}
+
+llvm.mlir.global internal @global() {addr_space = 0 : i32} : f32 {
+ %0 = llvm.mlir.zero : f32
+ llvm.return %0 : f32
+}
+
+// CHECK-LABEL: @global_accessor..omp_par({{.*}})
+// CHECK-NEXT: omp.par.entry:
+// Verify that we found the privatizer by checking that we properly inlined the
+// bodies of the alloc and copy regions.
+// CHECK: %[[PRIV_ALLOC:.*]] = alloca float, i64 1, align 4
+// CHECK: %[[GLOB_VAL:.*]] = load float, ptr @global, align 4
+// CHECK: store float %[[GLOB_VAL]], ptr %[[PRIV_ALLOC]], align 4
More information about the Mlir-commits
mailing list