Skip to content

Commit 9e6aa2f

Browse files
DavidTrubykiranchandramohan
DavidTruby
authored and
kiranchandramohan
committed
[mlir][flang][openmp] Rework parallel reduction operations
This patch reworks the way that parallel reduction operations function to better match the expected semantics from the OpenMP specification. Previously specific omp.reduction operations were used inside the region, meaning that the reduction only applied when the correct operation was used, whereas the specification states that any change to the variable inside the region should be taken into account for the reduction. The new semantics create a private reduction variable as a block argument which should be used normally for all operations on that variable in the region; this private variable is then combined with the others into the shared variable. This way no special omp.reduction operations are needed inside the region. This patch only makes the change for the `parallel` operation, the change for the `wsloop` operation will be in a separate patch.
1 parent ffab5a0 commit 9e6aa2f

File tree

8 files changed

+211
-63
lines changed

8 files changed

+211
-63
lines changed

‎flang/lib/Lower/OpenMP.cpp‎

Lines changed: 60 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -621,10 +621,12 @@ class ClauseProcessor {
621621
llvm::SmallVectorImpl<mlir::Location> *mapSymLocs = nullptr,
622622
llvm::SmallVectorImpl<const Fortran::semantics::Symbol *>
623623
*mapSymbols = nullptr) const;
624-
bool processReduction(
625-
mlir::Location currentLocation,
626-
llvm::SmallVectorImpl<mlir::Value> &reductionVars,
627-
llvm::SmallVectorImpl<mlir::Attribute> &reductionDeclSymbols) const;
624+
bool
625+
processReduction(mlir::Location currentLocation,
626+
llvm::SmallVectorImpl<mlir::Value> &reductionVars,
627+
llvm::SmallVectorImpl<mlir::Attribute> &reductionDeclSymbols,
628+
llvm::SmallVectorImpl<const Fortran::semantics::Symbol *>
629+
*reductionSymbols = nullptr) const;
628630
bool processSectionsReduction(mlir::Location currentLocation) const;
629631
bool processTo(llvm::SmallVectorImpl<DeclareTargetCapturePair> &result) const;
630632
bool
@@ -1079,12 +1081,14 @@ class ReductionProcessor {
10791081

10801082
/// Creates a reduction declaration and associates it with an OpenMP block
10811083
/// directive.
1082-
static void addReductionDecl(
1083-
mlir::Location currentLocation,
1084-
Fortran::lower::AbstractConverter &converter,
1085-
const Fortran::parser::OmpReductionClause &reduction,
1086-
llvm::SmallVectorImpl<mlir::Value> &reductionVars,
1087-
llvm::SmallVectorImpl<mlir::Attribute> &reductionDeclSymbols) {
1084+
static void
1085+
addReductionDecl(mlir::Location currentLocation,
1086+
Fortran::lower::AbstractConverter &converter,
1087+
const Fortran::parser::OmpReductionClause &reduction,
1088+
llvm::SmallVectorImpl<mlir::Value> &reductionVars,
1089+
llvm::SmallVectorImpl<mlir::Attribute> &reductionDeclSymbols,
1090+
llvm::SmallVectorImpl<const Fortran::semantics::Symbol *>
1091+
*reductionSymbols = nullptr) {
10881092
fir::FirOpBuilder &firOpBuilder = converter.getFirOpBuilder();
10891093
mlir::omp::ReductionDeclareOp decl;
10901094
const auto &redOperator{
@@ -1114,6 +1118,8 @@ class ReductionProcessor {
11141118
if (const auto *name{
11151119
Fortran::parser::Unwrap<Fortran::parser::Name>(ompObject)}) {
11161120
if (const Fortran::semantics::Symbol * symbol{name->symbol}) {
1121+
if (reductionSymbols)
1122+
reductionSymbols->push_back(symbol);
11171123
mlir::Value symVal = converter.getSymbolAddress(*symbol);
11181124
if (auto declOp = symVal.getDefiningOp<hlfir::DeclareOp>())
11191125
symVal = declOp.getBase();
@@ -1148,6 +1154,8 @@ class ReductionProcessor {
11481154
if (const auto *name{
11491155
Fortran::parser::Unwrap<Fortran::parser::Name>(ompObject)}) {
11501156
if (const Fortran::semantics::Symbol * symbol{name->symbol}) {
1157+
if (reductionSymbols)
1158+
reductionSymbols->push_back(symbol);
11511159
mlir::Value symVal = converter.getSymbolAddress(*symbol);
11521160
if (auto declOp = symVal.getDefiningOp<hlfir::DeclareOp>())
11531161
symVal = declOp.getBase();
@@ -1948,13 +1956,16 @@ bool ClauseProcessor::processMap(
19481956
bool ClauseProcessor::processReduction(
19491957
mlir::Location currentLocation,
19501958
llvm::SmallVectorImpl<mlir::Value> &reductionVars,
1951-
llvm::SmallVectorImpl<mlir::Attribute> &reductionDeclSymbols) const {
1959+
llvm::SmallVectorImpl<mlir::Attribute> &reductionDeclSymbols,
1960+
llvm::SmallVectorImpl<const Fortran::semantics::Symbol *> *reductionSymbols)
1961+
const {
19521962
return findRepeatableClause<ClauseTy::Reduction>(
19531963
[&](const ClauseTy::Reduction *reductionClause,
19541964
const Fortran::parser::CharBlock &) {
19551965
ReductionProcessor rp;
19561966
rp.addReductionDecl(currentLocation, converter, reductionClause->v,
1957-
reductionVars, reductionDeclSymbols);
1967+
reductionVars, reductionDeclSymbols,
1968+
reductionSymbols);
19581969
});
19591970
}
19601971

@@ -2304,6 +2315,14 @@ struct OpWithBodyGenInfo {
23042315
return *this;
23052316
}
23062317

2318+
OpWithBodyGenInfo &
2319+
setReductions(llvm::SmallVector<const Fortran::semantics::Symbol *> *value1,
2320+
llvm::SmallVector<mlir::Type> *value2) {
2321+
reductionSymbols = value1;
2322+
reductionTypes = value2;
2323+
return *this;
2324+
}
2325+
23072326
OpWithBodyGenInfo &setGenRegionEntryCb(GenOMPRegionEntryCBFn value) {
23082327
genRegionEntryCB = value;
23092328
return *this;
@@ -2323,6 +2342,11 @@ struct OpWithBodyGenInfo {
23232342
const Fortran::parser::OmpClauseList *clauses = nullptr;
23242343
/// [in] if provided, processes the construct's data-sharing attributes.
23252344
DataSharingProcessor *dsp = nullptr;
2345+
/// [in] if provided, list of reduction symbols
2346+
llvm::SmallVector<const Fortran::semantics::Symbol *> *reductionSymbols =
2347+
nullptr;
2348+
/// [in] if provided, list of reduction types
2349+
llvm::SmallVector<mlir::Type> *reductionTypes = nullptr;
23262350
/// [in] if provided, emits the op's region entry. Otherwise, an emtpy block
23272351
/// is created in the region.
23282352
GenOMPRegionEntryCBFn genRegionEntryCB = nullptr;
@@ -2567,6 +2591,7 @@ genParallelOp(Fortran::lower::AbstractConverter &converter,
25672591
llvm::SmallVector<mlir::Value> allocateOperands, allocatorOperands,
25682592
reductionVars;
25692593
llvm::SmallVector<mlir::Attribute> reductionDeclSymbols;
2594+
llvm::SmallVector<const Fortran::semantics::Symbol *> reductionSymbols;
25702595

25712596
ClauseProcessor cp(converter, clauseList);
25722597
cp.processIf(Fortran::parser::OmpIfClause::DirectiveNameModifier::Parallel,
@@ -2576,13 +2601,33 @@ genParallelOp(Fortran::lower::AbstractConverter &converter,
25762601
cp.processDefault();
25772602
cp.processAllocate(allocatorOperands, allocateOperands);
25782603
if (!outerCombined)
2579-
cp.processReduction(currentLocation, reductionVars, reductionDeclSymbols);
2604+
cp.processReduction(currentLocation, reductionVars, reductionDeclSymbols,
2605+
&reductionSymbols);
2606+
2607+
llvm::SmallVector<mlir::Type> reductionTypes;
2608+
reductionTypes.reserve(reductionVars.size());
2609+
llvm::transform(reductionVars, std::back_inserter(reductionTypes),
2610+
[](mlir::Value v) { return v.getType(); });
2611+
2612+
auto reductionCallback = [&](mlir::Operation *op) {
2613+
llvm::SmallVector<mlir::Location> locs(reductionVars.size(),
2614+
currentLocation);
2615+
auto block = converter.getFirOpBuilder().createBlock(&op->getRegion(0), {},
2616+
reductionTypes, locs);
2617+
for (auto [arg, prv] :
2618+
llvm::zip_equal(reductionSymbols, block->getArguments())) {
2619+
converter.bindSymbol(*arg, prv);
2620+
}
2621+
return reductionSymbols;
2622+
};
25802623

25812624
return genOpWithBody<mlir::omp::ParallelOp>(
25822625
OpWithBodyGenInfo(converter, currentLocation, eval)
25832626
.setGenNested(genNested)
25842627
.setOuterCombined(outerCombined)
2585-
.setClauses(&clauseList),
2628+
.setClauses(&clauseList)
2629+
.setReductions(&reductionSymbols, &reductionTypes)
2630+
.setGenRegionEntryCb(reductionCallback),
25862631
/*resultTypes=*/mlir::TypeRange(), ifClauseOperand,
25872632
numThreadsClauseOperand, allocateOperands, allocatorOperands,
25882633
reductionVars,
@@ -3634,10 +3679,8 @@ genOMP(Fortran::lower::AbstractConverter &converter,
36343679
break;
36353680
}
36363681

3637-
if (singleDirective) {
3638-
genOpenMPReduction(converter, beginClauseList);
3682+
if (singleDirective)
36393683
return;
3640-
}
36413684

36423685
// Codegen for combined directives
36433686
bool combinedDirective = false;
@@ -3673,7 +3716,6 @@ genOMP(Fortran::lower::AbstractConverter &converter,
36733716
")");
36743717

36753718
genNestedEvaluations(converter, eval);
3676-
genOpenMPReduction(converter, beginClauseList);
36773719
}
36783720

36793721
static void

‎flang/test/Lower/OpenMP/FIR/parallel-reduction-add.f90‎

Lines changed: 17 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -27,9 +27,11 @@
2727
!CHECK: %[[IREF:.*]] = fir.alloca i32 {bindc_name = "i", uniq_name = "_QFsimple_int_addEi"}
2828
!CHECK: %[[I_START:.*]] = arith.constant 0 : i32
2929
!CHECK: fir.store %[[I_START]] to %[[IREF]] : !fir.ref<i32>
30-
!CHECK: omp.parallel reduction(@[[RED_I32_NAME]] -> %[[IREF]] : !fir.ref<i32>) {
31-
!CHECK: %[[I_INCR:.*]] = arith.constant 1 : i32
32-
!CHECK: omp.reduction %[[I_INCR]], %[[IREF]] : i32, !fir.ref<i32>
30+
!CHECK: omp.parallel reduction(@[[RED_I32_NAME]] %[[IREF]] -> %[[PRV:.+]] : !fir.ref<i32>) {
31+
!CHECK: %[[LPRV:.+]] = fir.load %[[PRV]] : !fir.ref<i32>
32+
!CHECK: %[[I_INCR:.+]] = arith.constant 1 : i32
33+
!CHECK: %[[RES:.+]] = arith.addi %[[LPRV]], %[[I_INCR]]
34+
!CHECK: fir.store %[[RES]] to %[[PRV]] : !fir.ref<i32>
3335
!CHECK: omp.terminator
3436
!CHECK: }
3537
!CHECK: return
@@ -48,9 +50,11 @@ subroutine simple_int_add
4850
!CHECK: %[[RREF:.*]] = fir.alloca f32 {bindc_name = "r", uniq_name = "_QFsimple_real_addEr"}
4951
!CHECK: %[[R_START:.*]] = arith.constant 0.000000e+00 : f32
5052
!CHECK: fir.store %[[R_START]] to %[[RREF]] : !fir.ref<f32>
51-
!CHECK: omp.parallel reduction(@[[RED_F32_NAME]] -> %[[RREF]] : !fir.ref<f32>) {
52-
!CHECK: %[[R_INCR:.*]] = arith.constant 1.500000e+00 : f32
53-
!CHECK: omp.reduction %[[R_INCR]], %[[RREF]] : f32, !fir.ref<f32>
53+
!CHECK: omp.parallel reduction(@[[RED_F32_NAME]] %[[RREF]] -> %[[PRV:.+]] : !fir.ref<f32>) {
54+
!CHECK: %[[LPRV:.+]] = fir.load %[[PRV]] : !fir.ref<f32>
55+
!CHECK: %[[R_INCR:.+]] = arith.constant 1.500000e+00 : f32
56+
!CHECK: %[[RES]] = arith.addf %[[LPRV]], %[[R_INCR]] {{.*}} : f32
57+
!CHECK: fir.store %[[RES]] to %[[PRV]] : !fir.ref<f32>
5458
!CHECK: omp.terminator
5559
!CHECK: }
5660
!CHECK: return
@@ -72,11 +76,15 @@ subroutine simple_real_add
7276
!CHECK: fir.store %[[R_START]] to %[[RREF]] : !fir.ref<f32>
7377
!CHECK: %[[I_START:.*]] = arith.constant 0 : i32
7478
!CHECK: fir.store %[[I_START]] to %[[IREF]] : !fir.ref<i32>
75-
!CHECK: omp.parallel reduction(@[[RED_I32_NAME]] -> %[[IREF]] : !fir.ref<i32>, @[[RED_F32_NAME]] -> %[[RREF]] : !fir.ref<f32>) {
79+
!CHECK: omp.parallel reduction(@[[RED_I32_NAME]] %[[IREF]] -> %[[PRV0:.+]] : !fir.ref<i32>, @[[RED_F32_NAME]] %[[RREF]] -> %[[PRV1:.+]] : !fir.ref<f32>) {
7680
!CHECK: %[[R_INCR:.*]] = arith.constant 1.500000e+00 : f32
77-
!CHECK: omp.reduction %[[R_INCR]], %[[RREF]] : f32, !fir.ref<f32>
81+
!CHECK: %[[LPRV1:.+]] = fir.load %[[PRV1]] : !fir.ref<f32>
82+
!CHECK: %[[RES1:.+]] = arith.addf %[[R_INCR]], %[[LPRV1]] {{.*}} : f32
83+
!CHECK: fir.store %[[RES1]] to %[[PRV1]]
84+
!CHECK: %[[LPRV0:.+]] = fir.load %[[PRV0]] : !fir.ref<i32>
7885
!CHECK: %[[I_INCR:.*]] = arith.constant 3 : i32
79-
!CHECK: omp.reduction %[[I_INCR]], %[[IREF]] : i32, !fir.ref<i32>
86+
!CHECK: %[[RES0:.+]] = arith.addi %[[LPRV0]], %[[I_INCR]]
87+
!CHECK: fir.store %[[RES0]] to %[[PRV0]]
8088
!CHECK: omp.terminator
8189
!CHECK: }
8290
!CHECK: return

‎flang/test/Lower/OpenMP/parallel-reduction-add.f90‎

Lines changed: 19 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -28,9 +28,12 @@
2828
!CHECK: %[[I_DECL:.*]]:2 = hlfir.declare %[[IREF]] {uniq_name = "_QFsimple_int_addEi"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
2929
!CHECK: %[[I_START:.*]] = arith.constant 0 : i32
3030
!CHECK: hlfir.assign %[[I_START]] to %[[I_DECL]]#0 : i32, !fir.ref<i32>
31-
!CHECK: omp.parallel reduction(@[[RED_I32_NAME]] -> %[[I_DECL]]#0 : !fir.ref<i32>) {
31+
!CHECK: omp.parallel reduction(@[[RED_I32_NAME]] %[[I_DECL]]#0 -> %[[PRV:.+]] : !fir.ref<i32>) {
32+
!CHECK: %[[P_DECL:.+]]:2 = hlfir.declare %[[PRV]] {{.*}} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
33+
!CHECK: %[[LPRV:.+]] = fir.load %[[P_DECL]]#0 : !fir.ref<i32>
3234
!CHECK: %[[I_INCR:.*]] = arith.constant 1 : i32
33-
!CHECK: omp.reduction %[[I_INCR]], %[[I_DECL]]#0 : i32, !fir.ref<i32>
35+
!CHECK: %[[RES:.+]] = arith.addi %[[LPRV]], %[[I_INCR]] : i32
36+
!CHECK: hlfir.assign %[[RES]] to %[[P_DECL]]#0 : i32, !fir.ref<i32>
3437
!CHECK: omp.terminator
3538
!CHECK: }
3639
!CHECK: return
@@ -50,9 +53,12 @@ subroutine simple_int_add
5053
!CHECK: %[[R_DECL:.*]]:2 = hlfir.declare %[[RREF]] {uniq_name = "_QFsimple_real_addEr"} : (!fir.ref<f32>) -> (!fir.ref<f32>, !fir.ref<f32>)
5154
!CHECK: %[[R_START:.*]] = arith.constant 0.000000e+00 : f32
5255
!CHECK: hlfir.assign %[[R_START]] to %[[R_DECL]]#0 : f32, !fir.ref<f32>
53-
!CHECK: omp.parallel reduction(@[[RED_F32_NAME]] -> %[[R_DECL]]#0 : !fir.ref<f32>) {
56+
!CHECK: omp.parallel reduction(@[[RED_F32_NAME]] %[[R_DECL]]#0 -> %[[PRV:.+]] : !fir.ref<f32>) {
57+
!CHECK: %[[P_DECL:.+]]:2 = hlfir.declare %[[PRV]] {{.*}} : (!fir.ref<f32>) -> (!fir.ref<f32>, !fir.ref<f32>)
58+
!CHECK: %[[LPRV:.+]] = fir.load %[[P_DECL]]#0 : !fir.ref<f32>
5459
!CHECK: %[[R_INCR:.*]] = arith.constant 1.500000e+00 : f32
55-
!CHECK: omp.reduction %[[R_INCR]], %[[R_DECL]]#0 : f32, !fir.ref<f32>
60+
!CHECK: %[[RES:.+]] = arith.addf %[[LPRV]], %[[R_INCR]] {{.*}} : f32
61+
!CHECK: hlfir.assign %[[RES]] to %[[P_DECL]]#0 : f32, !fir.ref<f32>
5662
!CHECK: omp.terminator
5763
!CHECK: }
5864
!CHECK: return
@@ -76,11 +82,17 @@ subroutine simple_real_add
7682
!CHECK: hlfir.assign %[[R_START]] to %[[R_DECL]]#0 : f32, !fir.ref<f32>
7783
!CHECK: %[[I_START:.*]] = arith.constant 0 : i32
7884
!CHECK: hlfir.assign %[[I_START]] to %[[I_DECL]]#0 : i32, !fir.ref<i32>
79-
!CHECK: omp.parallel reduction(@[[RED_I32_NAME]] -> %[[I_DECL]]#0 : !fir.ref<i32>, @[[RED_F32_NAME]] -> %[[R_DECL]]#0 : !fir.ref<f32>) {
85+
!CHECK: omp.parallel reduction(@[[RED_I32_NAME]] %[[I_DECL]]#0 -> %[[IPRV:.+]] : !fir.ref<i32>, @[[RED_F32_NAME]] %[[R_DECL]]#0 -> %[[RPRV:.+]] : !fir.ref<f32>) {
86+
!CHECK: %[[IP_DECL:.+]]:2 = hlfir.declare %[[IPRV]] {{.*}} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
87+
!CHECK: %[[RP_DECL:.+]]:2 = hlfir.declare %[[RPRV]] {{.*}} : (!fir.ref<f32>) -> (!fir.ref<f32>, !fir.ref<f32>)
8088
!CHECK: %[[R_INCR:.*]] = arith.constant 1.500000e+00 : f32
81-
!CHECK: omp.reduction %[[R_INCR]], %[[R_DECL]]#0 : f32, !fir.ref<f32>
89+
!CHECK: %[[R_LPRV:.+]] = fir.load %[[RP_DECL]]#0 : !fir.ref<f32>
90+
!CHECK: %[[RES1:.+]] = arith.addf %[[R_INCR]], %[[R_LPRV]] {{.*}} : f32
91+
!CHECK: hlfir.assign %[[RES1]] to %[[RP_DECL]]#0 : f32, !fir.ref<f32>
92+
!CHECK: %[[I_LPRV:.+]] = fir.load %[[IP_DECL]]#0 : !fir.ref<i32>
8293
!CHECK: %[[I_INCR:.*]] = arith.constant 3 : i32
83-
!CHECK: omp.reduction %[[I_INCR]], %[[I_DECL]]#0 : i32, !fir.ref<i32>
94+
!CHECK: %[[RES0:.+]] = arith.addi %[[I_LPRV]], %[[I_INCR]] : i32
95+
!CHECK: hlfir.assign %[[RES0]] to %[[IP_DECL]]#0 : i32, !fir.ref<i32>
8496
!CHECK: omp.terminator
8597
!CHECK: }
8698
!CHECK: return

‎mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td‎

Lines changed: 3 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -191,19 +191,16 @@ def ParallelOp : OpenMP_Op<"parallel", [
191191
unsigned getNumReductionVars() { return getReductionVars().size(); }
192192
}];
193193
let assemblyFormat = [{
194-
oilist( `reduction` `(`
195-
custom<ReductionVarList>(
196-
$reduction_vars, type($reduction_vars), $reductions
197-
) `)`
198-
| `if` `(` $if_expr_var `:` type($if_expr_var) `)`
194+
oilist(
195+
`if` `(` $if_expr_var `:` type($if_expr_var) `)`
199196
| `num_threads` `(` $num_threads_var `:` type($num_threads_var) `)`
200197
| `allocate` `(`
201198
custom<AllocateAndAllocator>(
202199
$allocate_vars, type($allocate_vars),
203200
$allocators_vars, type($allocators_vars)
204201
) `)`
205202
| `proc_bind` `(` custom<ClauseAttr>($proc_bind_val) `)`
206-
) $region attr-dict
203+
) custom<ParallelRegion>($region, $reduction_vars, type($reduction_vars), $reductions) attr-dict
207204
}];
208205
let hasVerifier = 1;
209206
}

‎mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp‎

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
#include "mlir/Interfaces/FoldInterfaces.h"
2222

2323
#include "llvm/ADT/BitVector.h"
24+
#include "llvm/ADT/STLExtras.h"
2425
#include "llvm/ADT/STLForwardCompat.h"
2526
#include "llvm/ADT/SmallString.h"
2627
#include "llvm/ADT/StringExtras.h"
@@ -34,6 +35,7 @@
3435
#include "mlir/Dialect/OpenMP/OpenMPOpsEnums.cpp.inc"
3536
#include "mlir/Dialect/OpenMP/OpenMPOpsInterfaces.cpp.inc"
3637
#include "mlir/Dialect/OpenMP/OpenMPTypeInterfaces.cpp.inc"
38+
#include "mlir/Support/LogicalResult.h"
3739

3840
using namespace mlir;
3941
using namespace mlir::omp;
@@ -427,6 +429,71 @@ static void printScheduleClause(OpAsmPrinter &p, Operation *op,
427429
// Parser, printer and verifier for ReductionVarList
428430
//===----------------------------------------------------------------------===//
429431

432+
ParseResult
433+
parseReductionClause(OpAsmParser &parser, Region &region,
434+
SmallVectorImpl<OpAsmParser::UnresolvedOperand> &operands,
435+
SmallVectorImpl<Type> &types, ArrayAttr &reductionSymbols,
436+
SmallVectorImpl<OpAsmParser::Argument> &privates) {
437+
if (failed(parser.parseOptionalKeyword("reduction")))
438+
return failure();
439+
440+
SmallVector<SymbolRefAttr> reductionVec;
441+
442+
if (failed(
443+
parser.parseCommaSeparatedList(OpAsmParser::Delimiter::Paren, [&]() {
444+
if (parser.parseAttribute(reductionVec.emplace_back()) ||
445+
parser.parseOperand(operands.emplace_back()) ||
446+
parser.parseArrow() ||
447+
parser.parseArgument(privates.emplace_back()) ||
448+
parser.parseColonType(types.emplace_back()))
449+
return failure();
450+
return success();
451+
})))
452+
return failure();
453+
454+
for (auto [prv, type] : llvm::zip_equal(privates, types)) {
455+
prv.type = type;
456+
}
457+
SmallVector<Attribute> reductions(reductionVec.begin(), reductionVec.end());
458+
reductionSymbols = ArrayAttr::get(parser.getContext(), reductions);
459+
return success();
460+
}
461+
462+
static void printReductionClause(OpAsmPrinter &p, Operation *op, Region &region,
463+
ValueRange operands, TypeRange types,
464+
ArrayAttr reductionSymbols) {
465+
p << "reduction(";
466+
llvm::interleaveComma(llvm::zip_equal(reductionSymbols, operands,
467+
region.front().getArguments(), types),
468+
p, [&p](auto t) {
469+
auto [sym, op, arg, type] = t;
470+
p << sym << " " << op << " -> " << arg << " : "
471+
<< type;
472+
});
473+
p << ") ";
474+
}
475+
476+
static ParseResult
477+
parseParallelRegion(OpAsmParser &parser, Region &region,
478+
SmallVectorImpl<OpAsmParser::UnresolvedOperand> &operands,
479+
SmallVectorImpl<Type> &types, ArrayAttr &reductionSymbols) {
480+
481+
llvm::SmallVector<OpAsmParser::Argument> privates;
482+
if (succeeded(parseReductionClause(parser, region, operands, types,
483+
reductionSymbols, privates)))
484+
return parser.parseRegion(region, privates);
485+
486+
return parser.parseRegion(region);
487+
}
488+
489+
static void printParallelRegion(OpAsmPrinter &p, Operation *op, Region &region,
490+
ValueRange operands, TypeRange types,
491+
ArrayAttr reductionSymbols) {
492+
if (reductionSymbols)
493+
printReductionClause(p, op, region, operands, types, reductionSymbols);
494+
p.printRegion(region, /*printEntryBlockArgs=*/false);
495+
}
496+
430497
/// reduction-entry-list ::= reduction-entry
431498
/// | reduction-entry-list `,` reduction-entry
432499
/// reduction-entry ::= symbol-ref `->` ssa-id `:` type
@@ -1114,6 +1181,7 @@ parseLoopControl(OpAsmParser &parser, Region &region,
11141181
loopVarTypes = SmallVector<Type>(ivs.size(), loopVarType);
11151182
for (auto &iv : ivs)
11161183
iv.type = loopVarType;
1184+
11171185
return parser.parseRegion(region, ivs);
11181186
}
11191187

0 commit comments

Comments
 (0)