Skip to content

Commit 4f3924f

Browse files
ergawyergawy
committed
[NFC][flang][OpenMP] Extract target region utils to map or clone outside values
Following up on #154483, this PR introduces further refactoring to extract some shared utils between OpenMP lowering and `do concurrent` conversion pass. In particular, this PR extracts 2 utils that handle mapping or cloning values used inside target regions but defined outside. Later `do concurrent` PR(s) will use these utils.
1 parent b2a7369 commit 4f3924f

File tree

4 files changed

+152
-100
lines changed

4 files changed

+152
-100
lines changed

‎flang/include/flang/Utils/OpenMP.h‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,10 @@
1111

1212
#include "mlir/Dialect/OpenMP/OpenMPDialect.h"
1313

14+
namespace fir {
15+
class FirOpBuilder;
16+
} // namespace fir
17+
1418
namespace Fortran::utils::openmp {
1519
// TODO We can probably move the stuff inside `Support/OpenMP-utils.h/.cpp` here
1620
// as well.
@@ -28,6 +32,31 @@ mlir::omp::MapInfoOp createMapInfoOp(mlir::OpBuilder &builder,
2832
uint64_t mapType, mlir::omp::VariableCaptureKind mapCaptureType,
2933
mlir::Type retTy, bool partialMap = false,
3034
mlir::FlatSymbolRefAttr mapperId = mlir::FlatSymbolRefAttr());
35+
36+
/// For an mlir value that does not have storage, allocate temporary storage
37+
/// (outside the target region), store the value in that storage, and map the
38+
/// storage to the target region.
39+
///
40+
/// \param firOpBuilder - Operation builder.
41+
/// \param targetOp - Target op to which the temporary value is mapped.
42+
/// \param val - Temp value that should be mapped to the target region.
43+
/// \param name - A string used to identify the created `omp.map.info`
44+
/// op.
45+
///
46+
/// \returns The loaded mapped value inside the target region.
47+
mlir::Value mapTemporaryValue(fir::FirOpBuilder &firOpBuilder,
48+
mlir::omp::TargetOp targetOp, mlir::Value val, llvm::StringRef name);
49+
50+
/// For values used inside a target region but defined outside, either clone
51+
/// these value inside the target region or map them to the region. This
52+
/// function tries firts to clone values (if they are defined by
53+
/// memory-effect-free ops, otherwise, the values are mapped.
54+
///
55+
/// \param firOpBuilder - Operation builder.
56+
/// \param targetOp - The target that needs to be extended by clones and/or
57+
/// maps.
58+
void cloneOrMapRegionOutsiders(
59+
fir::FirOpBuilder &firOpBuilder, mlir::omp::TargetOp targetOp);
3160
} // namespace Fortran::utils::openmp
3261

3362
#endif // FORTRAN_UTILS_OPENMP_H_

‎flang/lib/Lower/OpenMP/OpenMP.cpp‎

Lines changed: 2 additions & 99 deletions
Original file line numberDiff line numberDiff line change
@@ -1422,7 +1422,7 @@ static void genBodyOfTargetOp(
14221422
auto argIface = llvm::cast<mlir::omp::BlockArgOpenMPOpInterface>(*targetOp);
14231423

14241424
mlir::Region &region = targetOp.getRegion();
1425-
mlir::Block *entryBlock = genEntryBlock(firOpBuilder, args, region);
1425+
genEntryBlock(firOpBuilder, args, region);
14261426
bindEntryBlockArgs(converter, targetOp, args);
14271427
if (HostEvalInfo *hostEvalInfo = getHostEvalInfoStackTop(converter))
14281428
hostEvalInfo->bindOperands(argIface.getHostEvalBlockArgs());
@@ -1431,104 +1431,7 @@ static void genBodyOfTargetOp(
14311431
// If so, then either clone them as well if they are MemoryEffectFree, or else
14321432
// copy them to a new temporary and add them to the map and block_argument
14331433
// lists and replace their uses with the new temporary.
1434-
llvm::SetVector<mlir::Value> valuesDefinedAbove;
1435-
mlir::getUsedValuesDefinedAbove(region, valuesDefinedAbove);
1436-
while (!valuesDefinedAbove.empty()) {
1437-
for (mlir::Value val : valuesDefinedAbove) {
1438-
mlir::Operation *valOp = val.getDefiningOp();
1439-
1440-
// NOTE: We skip BoxDimsOp's as the lesser of two evils is to map the
1441-
// indices separately, as the alternative is to eventually map the Box,
1442-
// which comes with a fairly large overhead comparatively. We could be
1443-
// more robust about this and check using a BackwardsSlice to see if we
1444-
// run the risk of mapping a box.
1445-
if (valOp && mlir::isMemoryEffectFree(valOp) &&
1446-
!mlir::isa<fir::BoxDimsOp>(valOp)) {
1447-
mlir::Operation *clonedOp = valOp->clone();
1448-
entryBlock->push_front(clonedOp);
1449-
1450-
auto replace = [entryBlock](mlir::OpOperand &use) {
1451-
return use.getOwner()->getBlock() == entryBlock;
1452-
};
1453-
1454-
valOp->getResults().replaceUsesWithIf(clonedOp->getResults(), replace);
1455-
valOp->replaceUsesWithIf(clonedOp, replace);
1456-
} else {
1457-
auto savedIP = firOpBuilder.getInsertionPoint();
1458-
1459-
if (valOp)
1460-
firOpBuilder.setInsertionPointAfter(valOp);
1461-
else
1462-
// This means val is a block argument
1463-
firOpBuilder.setInsertionPoint(targetOp);
1464-
1465-
auto copyVal =
1466-
firOpBuilder.createTemporary(val.getLoc(), val.getType());
1467-
firOpBuilder.createStoreWithConvert(copyVal.getLoc(), val, copyVal);
1468-
1469-
fir::factory::AddrAndBoundsInfo info =
1470-
fir::factory::getDataOperandBaseAddr(
1471-
firOpBuilder, val, /*isOptional=*/false, val.getLoc());
1472-
llvm::SmallVector<mlir::Value> bounds =
1473-
fir::factory::genImplicitBoundsOps<mlir::omp::MapBoundsOp,
1474-
mlir::omp::MapBoundsType>(
1475-
firOpBuilder, info,
1476-
hlfir::translateToExtendedValue(val.getLoc(), firOpBuilder,
1477-
hlfir::Entity{val})
1478-
.first,
1479-
/*dataExvIsAssumedSize=*/false, val.getLoc());
1480-
1481-
std::stringstream name;
1482-
firOpBuilder.setInsertionPoint(targetOp);
1483-
1484-
llvm::omp::OpenMPOffloadMappingFlags mapFlag =
1485-
llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_IMPLICIT;
1486-
mlir::omp::VariableCaptureKind captureKind =
1487-
mlir::omp::VariableCaptureKind::ByRef;
1488-
1489-
mlir::Type eleType = copyVal.getType();
1490-
if (auto refType =
1491-
mlir::dyn_cast<fir::ReferenceType>(copyVal.getType()))
1492-
eleType = refType.getElementType();
1493-
1494-
if (fir::isa_trivial(eleType) || fir::isa_char(eleType)) {
1495-
captureKind = mlir::omp::VariableCaptureKind::ByCopy;
1496-
} else if (!fir::isa_builtin_cptr_type(eleType)) {
1497-
mapFlag |= llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_TO;
1498-
}
1499-
1500-
mlir::Value mapOp = createMapInfoOp(
1501-
firOpBuilder, copyVal.getLoc(), copyVal,
1502-
/*varPtrPtr=*/mlir::Value{}, name.str(), bounds,
1503-
/*members=*/llvm::SmallVector<mlir::Value>{},
1504-
/*membersIndex=*/mlir::ArrayAttr{},
1505-
static_cast<
1506-
std::underlying_type_t<llvm::omp::OpenMPOffloadMappingFlags>>(
1507-
mapFlag),
1508-
captureKind, copyVal.getType());
1509-
1510-
// Get the index of the first non-map argument before modifying mapVars,
1511-
// then append an element to mapVars and an associated entry block
1512-
// argument at that index.
1513-
unsigned insertIndex =
1514-
argIface.getMapBlockArgsStart() + argIface.numMapBlockArgs();
1515-
targetOp.getMapVarsMutable().append(mapOp);
1516-
mlir::Value clonedValArg = region.insertArgument(
1517-
insertIndex, copyVal.getType(), copyVal.getLoc());
1518-
1519-
firOpBuilder.setInsertionPointToStart(entryBlock);
1520-
auto loadOp = fir::LoadOp::create(firOpBuilder, clonedValArg.getLoc(),
1521-
clonedValArg);
1522-
val.replaceUsesWithIf(loadOp->getResult(0),
1523-
[entryBlock](mlir::OpOperand &use) {
1524-
return use.getOwner()->getBlock() == entryBlock;
1525-
});
1526-
firOpBuilder.setInsertionPoint(entryBlock, savedIP);
1527-
}
1528-
}
1529-
valuesDefinedAbove.clear();
1530-
mlir::getUsedValuesDefinedAbove(region, valuesDefinedAbove);
1531-
}
1434+
cloneOrMapRegionOutsiders(firOpBuilder, targetOp);
15321435

15331436
// Insert dummy instruction to remember the insertion position. The
15341437
// marker will be deleted since there are not uses.

‎flang/lib/Utils/CMakeLists.txt‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,10 +11,17 @@ add_flang_library(FortranUtils
1111

1212
DEPENDS
1313
FIRDialect
14+
FIRBuilder
15+
HLFIRDialect
1416

1517
LINK_LIBS
1618
FIRDialect
17-
19+
FIRBuilder
20+
HLFIRDialect
21+
1822
MLIR_LIBS
23+
MLIRSupport
1924
MLIROpenMPDialect
25+
MLIRTransformUtils
26+
MLIRArithDialect
2027
)

‎flang/lib/Utils/OpenMP.cpp‎

Lines changed: 113 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,10 +8,14 @@
88

99
#include "flang/Utils/OpenMP.h"
1010

11+
#include "flang/Lower/ConvertExprToHLFIR.h"
12+
#include "flang/Optimizer/Builder/DirectivesCommon.h"
13+
#include "flang/Optimizer/Builder/FIRBuilder.h"
1114
#include "flang/Optimizer/Dialect/FIROps.h"
1215
#include "flang/Optimizer/Dialect/FIRType.h"
1316

1417
#include "mlir/Dialect/OpenMP/OpenMPDialect.h"
18+
#include "mlir/Transforms/RegionUtils.h"
1519

1620
namespace Fortran::utils::openmp {
1721
mlir::omp::MapInfoOp createMapInfoOp(mlir::OpBuilder &builder,
@@ -44,4 +48,113 @@ mlir::omp::MapInfoOp createMapInfoOp(mlir::OpBuilder &builder,
4448
builder.getStringAttr(name), builder.getBoolAttr(partialMap));
4549
return op;
4650
}
51+
52+
mlir::Value mapTemporaryValue(fir::FirOpBuilder &firOpBuilder,
53+
mlir::omp::TargetOp targetOp, mlir::Value val, llvm::StringRef name) {
54+
mlir::OpBuilder::InsertionGuard guard(firOpBuilder);
55+
mlir::Operation *valOp = val.getDefiningOp();
56+
57+
if (valOp)
58+
firOpBuilder.setInsertionPointAfter(valOp);
59+
else
60+
// This means val is a block argument
61+
firOpBuilder.setInsertionPoint(targetOp);
62+
63+
auto copyVal = firOpBuilder.createTemporary(val.getLoc(), val.getType());
64+
firOpBuilder.createStoreWithConvert(copyVal.getLoc(), val, copyVal);
65+
66+
fir::factory::AddrAndBoundsInfo info = fir::factory::getDataOperandBaseAddr(
67+
firOpBuilder, val, /*isOptional=*/false, val.getLoc());
68+
llvm::SmallVector<mlir::Value> bounds =
69+
fir::factory::genImplicitBoundsOps<mlir::omp::MapBoundsOp,
70+
mlir::omp::MapBoundsType>(firOpBuilder, info,
71+
hlfir::translateToExtendedValue(
72+
val.getLoc(), firOpBuilder, hlfir::Entity{val})
73+
.first,
74+
/*dataExvIsAssumedSize=*/false, val.getLoc());
75+
76+
firOpBuilder.setInsertionPoint(targetOp);
77+
78+
llvm::omp::OpenMPOffloadMappingFlags mapFlag =
79+
llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_IMPLICIT;
80+
mlir::omp::VariableCaptureKind captureKind =
81+
mlir::omp::VariableCaptureKind::ByRef;
82+
83+
mlir::Type eleType = copyVal.getType();
84+
if (auto refType = mlir::dyn_cast<fir::ReferenceType>(copyVal.getType())) {
85+
eleType = refType.getElementType();
86+
}
87+
88+
if (fir::isa_trivial(eleType) || fir::isa_char(eleType)) {
89+
captureKind = mlir::omp::VariableCaptureKind::ByCopy;
90+
} else if (!fir::isa_builtin_cptr_type(eleType)) {
91+
mapFlag |= llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_TO;
92+
}
93+
94+
mlir::Value mapOp = createMapInfoOp(firOpBuilder, copyVal.getLoc(), copyVal,
95+
/*varPtrPtr=*/mlir::Value{}, name.str(), bounds,
96+
/*members=*/llvm::SmallVector<mlir::Value>{},
97+
/*membersIndex=*/mlir::ArrayAttr{},
98+
static_cast<std::underlying_type_t<llvm::omp::OpenMPOffloadMappingFlags>>(
99+
mapFlag),
100+
captureKind, copyVal.getType());
101+
102+
auto argIface = llvm::cast<mlir::omp::BlockArgOpenMPOpInterface>(*targetOp);
103+
mlir::Region &region = targetOp.getRegion();
104+
105+
// Get the index of the first non-map argument before modifying mapVars,
106+
// then append an element to mapVars and an associated entry block
107+
// argument at that index.
108+
unsigned insertIndex =
109+
argIface.getMapBlockArgsStart() + argIface.numMapBlockArgs();
110+
targetOp.getMapVarsMutable().append(mapOp);
111+
mlir::Value clonedValArg =
112+
region.insertArgument(insertIndex, copyVal.getType(), copyVal.getLoc());
113+
114+
mlir::Block *entryBlock = &region.getBlocks().front();
115+
firOpBuilder.setInsertionPointToStart(entryBlock);
116+
auto loadOp =
117+
firOpBuilder.create<fir::LoadOp>(clonedValArg.getLoc(), clonedValArg);
118+
return loadOp.getResult();
119+
}
120+
121+
void cloneOrMapRegionOutsiders(
122+
fir::FirOpBuilder &firOpBuilder, mlir::omp::TargetOp targetOp) {
123+
mlir::Region &region = targetOp.getRegion();
124+
mlir::Block *entryBlock = &region.getBlocks().front();
125+
126+
llvm::SetVector<mlir::Value> valuesDefinedAbove;
127+
mlir::getUsedValuesDefinedAbove(region, valuesDefinedAbove);
128+
while (!valuesDefinedAbove.empty()) {
129+
for (mlir::Value val : valuesDefinedAbove) {
130+
mlir::Operation *valOp = val.getDefiningOp();
131+
132+
// NOTE: We skip BoxDimsOp's as the lesser of two evils is to map the
133+
// indices separately, as the alternative is to eventually map the Box,
134+
// which comes with a fairly large overhead comparatively. We could be
135+
// more robust about this and check using a BackwardsSlice to see if we
136+
// run the risk of mapping a box.
137+
if (valOp && mlir::isMemoryEffectFree(valOp) &&
138+
!mlir::isa<fir::BoxDimsOp>(valOp)) {
139+
mlir::Operation *clonedOp = valOp->clone();
140+
entryBlock->push_front(clonedOp);
141+
142+
auto replace = [entryBlock](mlir::OpOperand &use) {
143+
return use.getOwner()->getBlock() == entryBlock;
144+
};
145+
146+
valOp->getResults().replaceUsesWithIf(clonedOp->getResults(), replace);
147+
valOp->replaceUsesWithIf(clonedOp, replace);
148+
} else {
149+
mlir::Value mappedTemp = mapTemporaryValue(firOpBuilder, targetOp, val,
150+
/*name=*/{});
151+
val.replaceUsesWithIf(mappedTemp, [entryBlock](mlir::OpOperand &use) {
152+
return use.getOwner()->getBlock() == entryBlock;
153+
});
154+
}
155+
}
156+
valuesDefinedAbove.clear();
157+
mlir::getUsedValuesDefinedAbove(region, valuesDefinedAbove);
158+
}
159+
}
47160
} // namespace Fortran::utils::openmp

0 commit comments

Comments
 (0)