Skip to content

Commit 1a76e4d

Browse files
ergawyergawy
committed
[OpenMP][OMPIRBuilder] Refactor reduction initialization logic into one util
This refactors the logic needed to emit init logic for reductions by moving some duplicated code into a shared util. The logic for doing is quite involved and is needed for any construct that has reductions. Moreover, when a construct has both private and reduction clauses, both sets of clauses need to cooperate with each other when emitting the logic needed for allocation and initialization. Therefore, this PR clearly sets the boundaries for the logic needed to initialize reductions.
1 parent 95566af commit 1a76e4d

File tree

4 files changed

+109
-127
lines changed

4 files changed

+109
-127
lines changed

‎mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

Lines changed: 105 additions & 123 deletions
Original file line numberDiff line numberDiff line change
@@ -1030,6 +1030,99 @@ mapInitializationArgs(T loop, LLVM::ModuleTranslation &moduleTranslation,
10301030
}
10311031
}
10321032

1033+
template <typename OP>
1034+
static LogicalResult
1035+
initReductionVars(OP op, ArrayRef<BlockArgument> reductionArgs,
1036+
llvm::IRBuilderBase &builder,
1037+
LLVM::ModuleTranslation &moduleTranslation,
1038+
llvm::BasicBlock *latestAllocaBlock,
1039+
SmallVectorImpl<omp::DeclareReductionOp> &reductionDecls,
1040+
SmallVectorImpl<llvm::Value *> &privateReductionVariables,
1041+
DenseMap<Value, llvm::Value *> &reductionVariableMap,
1042+
llvm::ArrayRef<bool> isByRef,
1043+
SmallVectorImpl<DeferredStore> &deferredStores) {
1044+
if (op.getNumReductionVars() == 0)
1045+
return success();
1046+
1047+
llvm::IRBuilderBase::InsertPointGuard guard(builder);
1048+
1049+
builder.SetInsertPoint(latestAllocaBlock->getTerminator());
1050+
llvm::BasicBlock *initBlock = splitBB(builder, true, "omp.reduction.init");
1051+
auto allocaIP = llvm::IRBuilderBase::InsertPoint(
1052+
latestAllocaBlock, latestAllocaBlock->getTerminator()->getIterator());
1053+
builder.restoreIP(allocaIP);
1054+
SmallVector<llvm::Value *> byRefVars(op.getNumReductionVars());
1055+
1056+
for (unsigned i = 0; i < op.getNumReductionVars(); ++i) {
1057+
if (isByRef[i]) {
1058+
if (!reductionDecls[i].getAllocRegion().empty())
1059+
continue;
1060+
1061+
// TODO: remove after all users of by-ref are updated to use the alloc
1062+
// region: Allocate reduction variable (which is a pointer to the real
1063+
// reduciton variable allocated in the inlined region)
1064+
byRefVars[i] = builder.CreateAlloca(
1065+
moduleTranslation.convertType(reductionDecls[i].getType()));
1066+
}
1067+
}
1068+
1069+
builder.SetInsertPoint(&*initBlock->getFirstNonPHIOrDbgOrAlloca());
1070+
1071+
// store result of the alloc region to the allocated pointer to the real
1072+
// reduction variable
1073+
for (auto [data, addr] : deferredStores)
1074+
builder.CreateStore(data, addr);
1075+
1076+
// Before the loop, store the initial values of reductions into reduction
1077+
// variables. Although this could be done after allocas, we don't want to mess
1078+
// up with the alloca insertion point.
1079+
for (unsigned i = 0; i < op.getNumReductionVars(); ++i) {
1080+
SmallVector<llvm::Value *, 1> phis;
1081+
1082+
// map block argument to initializer region
1083+
mapInitializationArgs(op, moduleTranslation, reductionDecls,
1084+
reductionVariableMap, i);
1085+
1086+
if (failed(inlineConvertOmpRegions(reductionDecls[i].getInitializerRegion(),
1087+
"omp.reduction.neutral", builder,
1088+
moduleTranslation, &phis)))
1089+
return failure();
1090+
1091+
assert(phis.size() == 1 && "expected one value to be yielded from the "
1092+
"reduction neutral element declaration region");
1093+
1094+
builder.SetInsertPoint(builder.GetInsertBlock()->getTerminator());
1095+
1096+
if (isByRef[i]) {
1097+
if (!reductionDecls[i].getAllocRegion().empty())
1098+
// done in allocReductionVars
1099+
continue;
1100+
1101+
// TODO: this path can be removed once all users of by-ref are updated to
1102+
// use an alloc region
1103+
1104+
// Store the result of the inlined region to the allocated reduction var
1105+
// ptr
1106+
builder.CreateStore(phis[0], byRefVars[i]);
1107+
1108+
privateReductionVariables[i] = byRefVars[i];
1109+
moduleTranslation.mapValue(reductionArgs[i], phis[0]);
1110+
reductionVariableMap.try_emplace(op.getReductionVars()[i], phis[0]);
1111+
} else {
1112+
// for by-ref case the store is inside of the reduction region
1113+
builder.CreateStore(phis[0], privateReductionVariables[i]);
1114+
// the rest was handled in allocByValReductionVars
1115+
}
1116+
1117+
// forget the mapping for the initializer region because we might need a
1118+
// different mapping if this reduction declaration is re-used for a
1119+
// different variable
1120+
moduleTranslation.forgetMapping(reductionDecls[i].getInitializerRegion());
1121+
}
1122+
1123+
return success();
1124+
}
1125+
10331126
/// Collect reduction info
10341127
template <typename T>
10351128
static void collectReductionInfo(
@@ -1183,6 +1276,7 @@ static LogicalResult allocAndInitializeReductionVars(
11831276
if (op.getNumReductionVars() == 0)
11841277
return success();
11851278

1279+
llvm::IRBuilderBase::InsertPointGuard guard(builder);
11861280
SmallVector<DeferredStore> deferredStores;
11871281

11881282
if (failed(allocReductionVars(op, reductionArgs, builder, moduleTranslation,
@@ -1191,59 +1285,10 @@ static LogicalResult allocAndInitializeReductionVars(
11911285
deferredStores, isByRef)))
11921286
return failure();
11931287

1194-
// store result of the alloc region to the allocated pointer to the real
1195-
// reduction variable
1196-
for (auto [data, addr] : deferredStores)
1197-
builder.CreateStore(data, addr);
1198-
1199-
// Before the loop, store the initial values of reductions into reduction
1200-
// variables. Although this could be done after allocas, we don't want to mess
1201-
// up with the alloca insertion point.
1202-
for (unsigned i = 0; i < op.getNumReductionVars(); ++i) {
1203-
SmallVector<llvm::Value *, 1> phis;
1204-
1205-
// map block argument to initializer region
1206-
mapInitializationArgs(op, moduleTranslation, reductionDecls,
1207-
reductionVariableMap, i);
1208-
1209-
if (failed(inlineConvertOmpRegions(reductionDecls[i].getInitializerRegion(),
1210-
"omp.reduction.neutral", builder,
1211-
moduleTranslation, &phis)))
1212-
return failure();
1213-
assert(phis.size() == 1 && "expected one value to be yielded from the "
1214-
"reduction neutral element declaration region");
1215-
if (isByRef[i]) {
1216-
if (!reductionDecls[i].getAllocRegion().empty())
1217-
// done in allocReductionVars
1218-
continue;
1219-
1220-
// TODO: this path can be removed once all users of by-ref are updated to
1221-
// use an alloc region
1222-
1223-
// Allocate reduction variable (which is a pointer to the real reduction
1224-
// variable allocated in the inlined region)
1225-
llvm::Value *var = builder.CreateAlloca(
1226-
moduleTranslation.convertType(reductionDecls[i].getType()));
1227-
// Store the result of the inlined region to the allocated reduction var
1228-
// ptr
1229-
builder.CreateStore(phis[0], var);
1230-
1231-
privateReductionVariables[i] = var;
1232-
moduleTranslation.mapValue(reductionArgs[i], phis[0]);
1233-
reductionVariableMap.try_emplace(op.getReductionVars()[i], phis[0]);
1234-
} else {
1235-
// for by-ref case the store is inside of the reduction region
1236-
builder.CreateStore(phis[0], privateReductionVariables[i]);
1237-
// the rest was handled in allocByValReductionVars
1238-
}
1239-
1240-
// forget the mapping for the initializer region because we might need a
1241-
// different mapping if this reduction declaration is re-used for a
1242-
// different variable
1243-
moduleTranslation.forgetMapping(reductionDecls[i].getInitializerRegion());
1244-
}
1245-
1246-
return success();
1288+
return initReductionVars(op, reductionArgs, builder, moduleTranslation,
1289+
allocaIP.getBlock(), reductionDecls,
1290+
privateReductionVariables, reductionVariableMap,
1291+
isByRef, deferredStores);
12471292
}
12481293

12491294
/// Allocate delayed private variables. Returns the basic block which comes
@@ -1929,6 +1974,7 @@ convertOmpParallel(omp::ParallelOp opInst, llvm::IRBuilderBase &builder,
19291974
splitBB(builder, /*CreateBranch=*/true, "omp.private.copy");
19301975
builder.SetInsertPoint(copyBlock->getFirstNonPHIOrDbgOrAlloca());
19311976
}
1977+
19321978
for (auto [decl, mlirVar, llvmVar] :
19331979
llvm::zip_equal(privateDecls, mlirPrivateVars, llvmPrivateVars)) {
19341980
if (decl.getDataSharingType() != omp::DataSharingClauseType::FirstPrivate)
@@ -1960,76 +2006,12 @@ convertOmpParallel(omp::ParallelOp opInst, llvm::IRBuilderBase &builder,
19602006
moduleTranslation.forgetMapping(copyRegion);
19612007
}
19622008

1963-
// Initialize reduction vars
1964-
builder.SetInsertPoint(builder.GetInsertBlock()->getTerminator());
1965-
llvm::BasicBlock *initBlock = splitBB(builder, true, "omp.reduction.init");
1966-
allocaIP =
1967-
InsertPointTy(allocaIP.getBlock(),
1968-
allocaIP.getBlock()->getTerminator()->getIterator());
1969-
1970-
builder.restoreIP(allocaIP);
1971-
SmallVector<llvm::Value *> byRefVars(opInst.getNumReductionVars());
1972-
for (unsigned i = 0; i < opInst.getNumReductionVars(); ++i) {
1973-
if (isByRef[i]) {
1974-
if (!reductionDecls[i].getAllocRegion().empty())
1975-
continue;
1976-
1977-
// TODO: remove after all users of by-ref are updated to use the alloc
1978-
// region: Allocate reduction variable (which is a pointer to the real
1979-
// reduciton variable allocated in the inlined region)
1980-
byRefVars[i] = builder.CreateAlloca(
1981-
moduleTranslation.convertType(reductionDecls[i].getType()));
1982-
}
1983-
}
1984-
1985-
builder.SetInsertPoint(initBlock->getFirstNonPHIOrDbgOrAlloca());
1986-
1987-
// insert stores deferred until after all allocas
1988-
// these store the results of the alloc region into the allocation for the
1989-
// pointer to the reduction variable
1990-
for (auto [data, addr] : deferredStores)
1991-
builder.CreateStore(data, addr);
1992-
1993-
for (unsigned i = 0; i < opInst.getNumReductionVars(); ++i) {
1994-
SmallVector<llvm::Value *> phis;
1995-
1996-
// map the block argument
1997-
mapInitializationArgs(opInst, moduleTranslation, reductionDecls,
1998-
reductionVariableMap, i);
1999-
if (failed(inlineConvertOmpRegions(
2000-
reductionDecls[i].getInitializerRegion(), "omp.reduction.neutral",
2001-
builder, moduleTranslation, &phis)))
2002-
return llvm::createStringError(
2003-
"failed to inline `init` region of `omp.declare_reduction`");
2004-
assert(phis.size() == 1 &&
2005-
"expected one value to be yielded from the "
2006-
"reduction neutral element declaration region");
2007-
2008-
builder.SetInsertPoint(builder.GetInsertBlock()->getTerminator());
2009-
2010-
if (isByRef[i]) {
2011-
if (!reductionDecls[i].getAllocRegion().empty())
2012-
continue;
2013-
2014-
// TODO: remove after all users of by-ref are updated to use the alloc
2015-
2016-
// Store the result of the inlined region to the allocated reduction var
2017-
// ptr
2018-
builder.CreateStore(phis[0], byRefVars[i]);
2019-
2020-
privateReductionVariables[i] = byRefVars[i];
2021-
moduleTranslation.mapValue(reductionArgs[i], phis[0]);
2022-
reductionVariableMap.try_emplace(opInst.getReductionVars()[i], phis[0]);
2023-
} else {
2024-
// for by-ref case the store is inside of the reduction init region
2025-
builder.CreateStore(phis[0], privateReductionVariables[i]);
2026-
// the rest is done in allocByValReductionVars
2027-
}
2028-
2029-
// clear block argument mapping in case it needs to be re-created with a
2030-
// different source for another use of the same reduction decl
2031-
moduleTranslation.forgetMapping(reductionDecls[i].getInitializerRegion());
2032-
}
2009+
if (failed(
2010+
initReductionVars(opInst, reductionArgs, builder, moduleTranslation,
2011+
afterAllocas.get()->getSinglePredecessor(),
2012+
reductionDecls, privateReductionVariables,
2013+
reductionVariableMap, isByRef, deferredStores)))
2014+
return llvm::make_error<PreviouslyReportedError>();
20332015

20342016
// Store the mapping between reduction variables and their private copies on
20352017
// ModuleTranslation stack. It can be then recovered when translating

‎mlir/test/Target/LLVMIR/openmp-reduction-array-sections.mlir

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -91,12 +91,12 @@ llvm.func @sectionsreduction_(%arg0: !llvm.ptr {fir.bindc_name = "x"}) attribute
9191
// CHECK: %[[VAL_14:.*]] = alloca [1 x ptr], align 8
9292
// CHECK: br label %[[VAL_15:.*]]
9393
// CHECK: omp.reduction.init: ; preds = %[[VAL_16:.*]]
94+
// CHECK: store ptr %[[VAL_20]], ptr %[[VAL_21]], align 8
9495
// CHECK: br label %[[VAL_17:.*]]
9596
// CHECK: omp.par.region: ; preds = %[[VAL_15]]
9697
// CHECK: br label %[[VAL_18:.*]]
9798
// CHECK: omp.par.region1: ; preds = %[[VAL_17]]
9899
// CHECK: %[[VAL_19:.*]] = alloca { ptr, i64, i32, i8, i8, i8, i8, [1 x [3 x i64]] }, i64 1, align 8
99-
// CHECK: store ptr %[[VAL_20]], ptr %[[VAL_21]], align 8
100100
// CHECK: br label %[[VAL_22:.*]]
101101
// CHECK: omp_section_loop.preheader: ; preds = %[[VAL_18]]
102102
// CHECK: store i32 0, ptr %[[VAL_7]], align 4

‎mlir/test/Target/LLVMIR/openmp-reduction-sections.mlir

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -51,11 +51,11 @@ llvm.func @sections_(%arg0: !llvm.ptr {fir.bindc_name = "x"}) attributes {fir.in
5151
// CHECK: %[[VAL_21:.*]] = alloca [1 x ptr], align 8
5252
// CHECK: br label %[[VAL_22:.*]]
5353
// CHECK: omp.reduction.init: ; preds = %[[VAL_23:.*]]
54+
// CHECK: store float 0.000000e+00, ptr %[[VAL_20]], align 4
5455
// CHECK: br label %[[VAL_24:.*]]
5556
// CHECK: omp.par.region: ; preds = %[[VAL_22]]
5657
// CHECK: br label %[[VAL_25:.*]]
5758
// CHECK: omp.par.region1: ; preds = %[[VAL_24]]
58-
// CHECK: store float 0.000000e+00, ptr %[[VAL_20]], align 4
5959
// CHECK: br label %[[VAL_26:.*]]
6060
// CHECK: omp_section_loop.preheader: ; preds = %[[VAL_25]]
6161
// CHECK: store i32 0, ptr %[[VAL_13]], align 4

‎mlir/test/Target/LLVMIR/openmp-wsloop-reduction-cleanup.mlir

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -51,11 +51,11 @@
5151
llvm.func @free(%arg0 : !llvm.ptr) -> ()
5252

5353
// Private reduction variable and its initialization.
54-
// CHECK: %[[MALLOC_I:.+]] = call ptr @malloc(i64 4)
5554
// CHECK: %[[PRIV_PTR_I:.+]] = alloca ptr
55+
// CHECK: %[[PRIV_PTR_J:.+]] = alloca ptr
56+
// CHECK: %[[MALLOC_I:.+]] = call ptr @malloc(i64 4)
5657
// CHECK: store ptr %[[MALLOC_I]], ptr %[[PRIV_PTR_I]]
5758
// CHECK: %[[MALLOC_J:.+]] = call ptr @malloc(i64 4)
58-
// CHECK: %[[PRIV_PTR_J:.+]] = alloca ptr
5959
// CHECK: store ptr %[[MALLOC_J]], ptr %[[PRIV_PTR_J]]
6060

6161
// Call to the reduction function.

0 commit comments

Comments
 (0)