Skip to content

Commit 227e1ff

Browse files
ergawyergawy
authored
[flang][fir] Add locality specifiers modeling to fir.do_concurrent.loop (llvm#138506)
1 parent 88e6887 commit 227e1ff

File tree

5 files changed

+185
-26
lines changed

5 files changed

+185
-26
lines changed

‎flang/include/flang/Optimizer/Dialect/FIROps.td‎

Lines changed: 32 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3647,6 +3647,13 @@ def fir_DoConcurrentOp : fir_Op<"do_concurrent",
36473647
let hasVerifier = 1;
36483648
}
36493649

3650+
def fir_LocalSpecifier {
3651+
dag arguments = (ins
3652+
Variadic<AnyType>:$local_vars,
3653+
OptionalAttr<SymbolRefArrayAttr>:$local_syms
3654+
);
3655+
}
3656+
36503657
def fir_DoConcurrentLoopOp : fir_Op<"do_concurrent.loop",
36513658
[AttrSizedOperandSegments, DeclareOpInterfaceMethods<LoopLikeOpInterface,
36523659
["getLoopInductionVars"]>,
@@ -3700,7 +3707,7 @@ def fir_DoConcurrentLoopOp : fir_Op<"do_concurrent.loop",
37003707
LLVM.
37013708
}];
37023709

3703-
let arguments = (ins
3710+
defvar opArgs = (ins
37043711
Variadic<Index>:$lowerBound,
37053712
Variadic<Index>:$upperBound,
37063713
Variadic<Index>:$step,
@@ -3709,16 +3716,40 @@ def fir_DoConcurrentLoopOp : fir_Op<"do_concurrent.loop",
37093716
OptionalAttr<LoopAnnotationAttr>:$loopAnnotation
37103717
);
37113718

3719+
let arguments = !con(opArgs, fir_LocalSpecifier.arguments);
3720+
37123721
let regions = (region SizedRegion<1>:$region);
37133722

37143723
let hasCustomAssemblyFormat = 1;
37153724
let hasVerifier = 1;
37163725

37173726
let extraClassDeclaration = [{
3727+
unsigned getNumInductionVars() { return getLowerBound().size(); }
3728+
3729+
unsigned getNumLocalOperands() { return getLocalVars().size(); }
3730+
3731+
mlir::Block::BlockArgListType getInductionVars() {
3732+
return getBody()->getArguments().slice(0, getNumInductionVars());
3733+
}
3734+
3735+
mlir::Block::BlockArgListType getRegionLocalArgs() {
3736+
return getBody()->getArguments().slice(getNumInductionVars(),
3737+
getNumLocalOperands());
3738+
}
3739+
3740+
/// Number of operands controlling the loop
3741+
unsigned getNumControlOperands() { return getLowerBound().size() * 3; }
3742+
37183743
// Get Number of reduction operands
37193744
unsigned getNumReduceOperands() {
37203745
return getReduceOperands().size();
37213746
}
3747+
3748+
mlir::Operation::operand_range getLocalOperands() {
3749+
return getOperands()
3750+
.slice(getNumControlOperands() + getNumReduceOperands(),
3751+
getNumLocalOperands());
3752+
}
37223753
}];
37233754
}
37243755

‎flang/lib/Lower/Bridge.cpp‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2460,7 +2460,7 @@ class FirConverter : public Fortran::lower::AbstractConverter {
24602460
nestReduceAttrs.empty()
24612461
? nullptr
24622462
: mlir::ArrayAttr::get(builder->getContext(), nestReduceAttrs),
2463-
nullptr);
2463+
nullptr, /*local_vars=*/std::nullopt, /*local_syms=*/nullptr);
24642464

24652465
llvm::SmallVector<mlir::Type> loopBlockArgTypes(
24662466
incrementLoopNestInfo.size(), builder->getIndexType());

‎flang/lib/Optimizer/Dialect/FIROps.cpp‎

Lines changed: 94 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -5033,29 +5033,33 @@ mlir::ParseResult fir::DoConcurrentLoopOp::parse(mlir::OpAsmParser &parser,
50335033
mlir::OperationState &result) {
50345034
auto &builder = parser.getBuilder();
50355035
// Parse an opening `(` followed by induction variables followed by `)`
5036-
llvm::SmallVector<mlir::OpAsmParser::Argument, 4> ivs;
5037-
if (parser.parseArgumentList(ivs, mlir::OpAsmParser::Delimiter::Paren))
5036+
llvm::SmallVector<mlir::OpAsmParser::Argument, 4> regionArgs;
5037+
5038+
if (parser.parseArgumentList(regionArgs, mlir::OpAsmParser::Delimiter::Paren))
50385039
return mlir::failure();
50395040

5041+
llvm::SmallVector<mlir::Type> argTypes(regionArgs.size(),
5042+
builder.getIndexType());
5043+
50405044
// Parse loop bounds.
50415045
llvm::SmallVector<mlir::OpAsmParser::UnresolvedOperand, 4> lower;
50425046
if (parser.parseEqual() ||
5043-
parser.parseOperandList(lower, ivs.size(),
5047+
parser.parseOperandList(lower, regionArgs.size(),
50445048
mlir::OpAsmParser::Delimiter::Paren) ||
50455049
parser.resolveOperands(lower, builder.getIndexType(), result.operands))
50465050
return mlir::failure();
50475051

50485052
llvm::SmallVector<mlir::OpAsmParser::UnresolvedOperand, 4> upper;
50495053
if (parser.parseKeyword("to") ||
5050-
parser.parseOperandList(upper, ivs.size(),
5054+
parser.parseOperandList(upper, regionArgs.size(),
50515055
mlir::OpAsmParser::Delimiter::Paren) ||
50525056
parser.resolveOperands(upper, builder.getIndexType(), result.operands))
50535057
return mlir::failure();
50545058

50555059
// Parse step values.
50565060
llvm::SmallVector<mlir::OpAsmParser::UnresolvedOperand, 4> steps;
50575061
if (parser.parseKeyword("step") ||
5058-
parser.parseOperandList(steps, ivs.size(),
5062+
parser.parseOperandList(steps, regionArgs.size(),
50595063
mlir::OpAsmParser::Delimiter::Paren) ||
50605064
parser.resolveOperands(steps, builder.getIndexType(), result.operands))
50615065
return mlir::failure();
@@ -5086,20 +5090,72 @@ mlir::ParseResult fir::DoConcurrentLoopOp::parse(mlir::OpAsmParser &parser,
50865090
builder.getArrayAttr(arrayAttr));
50875091
}
50885092

5089-
// Now parse the body.
5090-
mlir::Region *body = result.addRegion();
5091-
for (auto &iv : ivs)
5092-
iv.type = builder.getIndexType();
5093-
if (parser.parseRegion(*body, ivs))
5094-
return mlir::failure();
5093+
llvm::SmallVector<mlir::OpAsmParser::UnresolvedOperand> localOperands;
5094+
if (succeeded(parser.parseOptionalKeyword("local"))) {
5095+
std::size_t oldArgTypesSize = argTypes.size();
5096+
if (failed(parser.parseLParen()))
5097+
return mlir::failure();
5098+
5099+
llvm::SmallVector<mlir::SymbolRefAttr> localSymbolVec;
5100+
if (failed(parser.parseCommaSeparatedList([&]() {
5101+
if (failed(parser.parseAttribute(localSymbolVec.emplace_back())))
5102+
return mlir::failure();
5103+
5104+
if (parser.parseOperand(localOperands.emplace_back()) ||
5105+
parser.parseArrow() ||
5106+
parser.parseArgument(regionArgs.emplace_back()))
5107+
return mlir::failure();
5108+
5109+
return mlir::success();
5110+
})))
5111+
return mlir::failure();
5112+
5113+
if (failed(parser.parseColon()))
5114+
return mlir::failure();
5115+
5116+
if (failed(parser.parseCommaSeparatedList([&]() {
5117+
if (failed(parser.parseType(argTypes.emplace_back())))
5118+
return mlir::failure();
5119+
5120+
return mlir::success();
5121+
})))
5122+
return mlir::failure();
5123+
5124+
if (regionArgs.size() != argTypes.size())
5125+
return parser.emitError(parser.getNameLoc(),
5126+
"mismatch in number of local arg and types");
5127+
5128+
if (failed(parser.parseRParen()))
5129+
return mlir::failure();
5130+
5131+
for (auto operandType : llvm::zip_equal(
5132+
localOperands, llvm::drop_begin(argTypes, oldArgTypesSize)))
5133+
if (parser.resolveOperand(std::get<0>(operandType),
5134+
std::get<1>(operandType), result.operands))
5135+
return mlir::failure();
5136+
5137+
llvm::SmallVector<mlir::Attribute> symbolAttrs(localSymbolVec.begin(),
5138+
localSymbolVec.end());
5139+
result.addAttribute(getLocalSymsAttrName(result.name),
5140+
builder.getArrayAttr(symbolAttrs));
5141+
}
50955142

50965143
// Set `operandSegmentSizes` attribute.
50975144
result.addAttribute(DoConcurrentLoopOp::getOperandSegmentSizeAttr(),
50985145
builder.getDenseI32ArrayAttr(
50995146
{static_cast<int32_t>(lower.size()),
51005147
static_cast<int32_t>(upper.size()),
51015148
static_cast<int32_t>(steps.size()),
5102-
static_cast<int32_t>(reduceOperands.size())}));
5149+
static_cast<int32_t>(reduceOperands.size()),
5150+
static_cast<int32_t>(localOperands.size())}));
5151+
5152+
// Now parse the body.
5153+
for (auto [arg, type] : llvm::zip_equal(regionArgs, argTypes))
5154+
arg.type = type;
5155+
5156+
mlir::Region *body = result.addRegion();
5157+
if (parser.parseRegion(*body, regionArgs))
5158+
return mlir::failure();
51035159

51045160
// Parse attributes.
51055161
if (parser.parseOptionalAttrDict(result.attributes))
@@ -5109,8 +5165,9 @@ mlir::ParseResult fir::DoConcurrentLoopOp::parse(mlir::OpAsmParser &parser,
51095165
}
51105166

51115167
void fir::DoConcurrentLoopOp::print(mlir::OpAsmPrinter &p) {
5112-
p << " (" << getBody()->getArguments() << ") = (" << getLowerBound()
5113-
<< ") to (" << getUpperBound() << ") step (" << getStep() << ")";
5168+
p << " (" << getBody()->getArguments().slice(0, getNumInductionVars())
5169+
<< ") = (" << getLowerBound() << ") to (" << getUpperBound() << ") step ("
5170+
<< getStep() << ")";
51145171

51155172
if (!getReduceOperands().empty()) {
51165173
p << " reduce(";
@@ -5123,12 +5180,27 @@ void fir::DoConcurrentLoopOp::print(mlir::OpAsmPrinter &p) {
51235180
p << ')';
51245181
}
51255182

5183+
if (!getLocalVars().empty()) {
5184+
p << " local(";
5185+
llvm::interleaveComma(llvm::zip_equal(getLocalSymsAttr(), getLocalVars(),
5186+
getRegionLocalArgs()),
5187+
p, [&](auto it) {
5188+
p << std::get<0>(it) << " " << std::get<1>(it)
5189+
<< " -> " << std::get<2>(it);
5190+
});
5191+
p << " : ";
5192+
llvm::interleaveComma(getLocalVars(), p,
5193+
[&](auto it) { p << it.getType(); });
5194+
p << ")";
5195+
}
5196+
51265197
p << ' ';
51275198
p.printRegion(getRegion(), /*printEntryBlockArgs=*/false);
51285199
p.printOptionalAttrDict(
51295200
(*this)->getAttrs(),
51305201
/*elidedAttrs=*/{DoConcurrentLoopOp::getOperandSegmentSizeAttr(),
5131-
DoConcurrentLoopOp::getReduceAttrsAttrName()});
5202+
DoConcurrentLoopOp::getReduceAttrsAttrName(),
5203+
DoConcurrentLoopOp::getLocalSymsAttrName()});
51325204
}
51335205

51345206
llvm::SmallVector<mlir::Region *> fir::DoConcurrentLoopOp::getLoopRegions() {
@@ -5139,6 +5211,7 @@ llvm::LogicalResult fir::DoConcurrentLoopOp::verify() {
51395211
mlir::Operation::operand_range lbValues = getLowerBound();
51405212
mlir::Operation::operand_range ubValues = getUpperBound();
51415213
mlir::Operation::operand_range stepValues = getStep();
5214+
mlir::Operation::operand_range localVars = getLocalVars();
51425215

51435216
if (lbValues.empty())
51445217
return emitOpError(
@@ -5152,11 +5225,13 @@ llvm::LogicalResult fir::DoConcurrentLoopOp::verify() {
51525225
// Check that the body defines the same number of block arguments as the
51535226
// number of tuple elements in step.
51545227
mlir::Block *body = getBody();
5155-
if (body->getNumArguments() != stepValues.size())
5228+
unsigned numIndVarArgs = body->getNumArguments() - localVars.size();
5229+
5230+
if (numIndVarArgs != stepValues.size())
51565231
return emitOpError() << "expects the same number of induction variables: "
51575232
<< body->getNumArguments()
51585233
<< " as bound and step values: " << stepValues.size();
5159-
for (auto arg : body->getArguments())
5234+
for (auto arg : body->getArguments().slice(0, numIndVarArgs))
51605235
if (!arg.getType().isIndex())
51615236
return emitOpError(
51625237
"expects arguments for the induction variable to be of index type");
@@ -5171,7 +5246,8 @@ llvm::LogicalResult fir::DoConcurrentLoopOp::verify() {
51715246

51725247
std::optional<llvm::SmallVector<mlir::Value>>
51735248
fir::DoConcurrentLoopOp::getLoopInductionVars() {
5174-
return llvm::SmallVector<mlir::Value>{getBody()->getArguments()};
5249+
return llvm::SmallVector<mlir::Value>{
5250+
getBody()->getArguments().slice(0, getLowerBound().size())};
51755251
}
51765252

51775253
//===----------------------------------------------------------------------===//

‎flang/test/Fir/do_concurrent.fir‎

Lines changed: 53 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -91,7 +91,6 @@ func.func @dc_2d_reduction(%i_lb: index, %i_ub: index, %i_st: index,
9191
// CHECK: }
9292
// CHECK: }
9393

94-
9594
fir.local {type = local} @local_privatizer : i32
9695

9796
// CHECK: fir.local {type = local} @[[LOCAL_PRIV_SYM:local_privatizer]] : i32
@@ -109,3 +108,56 @@ fir.local {type = local_init} @local_init_privatizer : i32 copy {
109108
// CHECK: fir.store %[[ORIG_VAL_LD]] to %[[LOCAL_VAL]] : !fir.ref<i32>
110109
// CHECK: fir.yield(%[[LOCAL_VAL]] : !fir.ref<i32>)
111110
// CHECK: }
111+
112+
func.func @do_concurrent_with_locality_specs() {
113+
%3 = fir.alloca i32 {bindc_name = "local_init_var"}
114+
%4:2 = hlfir.declare %3 {uniq_name = "_QFdo_concurrentElocal_init_var"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
115+
%5 = fir.alloca i32 {bindc_name = "local_var"}
116+
%6:2 = hlfir.declare %5 {uniq_name = "_QFdo_concurrentElocal_var"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
117+
%c1 = arith.constant 1 : index
118+
%c10 = arith.constant 10 : index
119+
120+
fir.do_concurrent {
121+
%9 = fir.alloca i32 {bindc_name = "i"}
122+
%10:2 = hlfir.declare %9 {uniq_name = "_QFdo_concurrentEi"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
123+
124+
fir.do_concurrent.loop (%arg0) = (%c1) to (%c10) step (%c1)
125+
local(@local_privatizer %6#0 -> %arg1, @local_init_privatizer %4#0 -> %arg2 : !fir.ref<i32>, !fir.ref<i32>) {
126+
%11 = fir.convert %arg0 : (index) -> i32
127+
fir.store %11 to %10#0 : !fir.ref<i32>
128+
%13:2 = hlfir.declare %arg1 {uniq_name = "_QFdo_concurrentElocal_var"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
129+
%15:2 = hlfir.declare %arg2 {uniq_name = "_QFdo_concurrentElocal_init_var"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
130+
}
131+
}
132+
return
133+
}
134+
135+
// CHECK-LABEL: func.func @do_concurrent_with_locality_specs() {
136+
// CHECK: %[[LOC_INIT_ALLOC:.*]] = fir.alloca i32 {bindc_name = "local_init_var"}
137+
// CHECK: %[[LOC_INIT_DECL:.*]]:2 = hlfir.declare %[[LOC_INIT_ALLOC]]
138+
139+
// CHECK: %[[LOC_ALLOC:.*]] = fir.alloca i32 {bindc_name = "local_var"}
140+
// CHECK: %[[LOC_DECL:.*]]:2 = hlfir.declare %[[LOC_ALLOC]]
141+
142+
// CHECK: %[[C1:.*]] = arith.constant 1 : index
143+
// CHECK: %[[C10:.*]] = arith.constant 10 : index
144+
145+
// CHECK: fir.do_concurrent {
146+
// CHECK: %[[DC_I_ALLOC:.*]] = fir.alloca i32 {bindc_name = "i"}
147+
// CHECK: %[[DC_I_DECL:.*]]:2 = hlfir.declare %[[DC_I_ALLOC]]
148+
149+
// CHECK: fir.do_concurrent.loop (%[[IV:.*]]) = (%[[C1]]) to
150+
// CHECK-SAME: (%[[C10]]) step (%[[C1]])
151+
// CHECK-SAME: local(@[[LOCAL_PRIV_SYM]] %[[LOC_DECL]]#0 -> %[[LOC_ARG:[^,]*]],
152+
// CHECK-SAME: @[[LOCAL_INIT_PRIV_SYM]] %[[LOC_INIT_DECL]]#0 -> %[[LOC_INIT_ARG:.*]] :
153+
// CHECK-SAME: !fir.ref<i32>, !fir.ref<i32>) {
154+
155+
// CHECK: %[[IV_CVT:.*]] = fir.convert %[[IV]] : (index) -> i32
156+
// CHECK: fir.store %[[IV_CVT]] to %[[DC_I_DECL]]#0 : !fir.ref<i32>
157+
158+
// CHECK: %[[LOC_PRIV_DECL:.*]]:2 = hlfir.declare %[[LOC_ARG]]
159+
// CHECK: %[[LOC_INIT_PRIV_DECL:.*]]:2 = hlfir.declare %[[LOC_INIT_ARG]]
160+
// CHECK: }
161+
// CHECK: }
162+
// CHECK: return
163+
// CHECK: }

‎flang/test/Fir/invalid.fir‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1196,7 +1196,7 @@ func.func @dc_0d() {
11961196

11971197
func.func @dc_invalid_parent(%arg0: index, %arg1: index) {
11981198
// expected-error@+1 {{'fir.do_concurrent.loop' op expects parent op 'fir.do_concurrent'}}
1199-
"fir.do_concurrent.loop"(%arg0, %arg1) <{operandSegmentSizes = array<i32: 1, 1, 0, 0>}> ({
1199+
"fir.do_concurrent.loop"(%arg0, %arg1) <{operandSegmentSizes = array<i32: 1, 1, 0, 0, 0>}> ({
12001200
^bb0(%arg2: index):
12011201
%tmp = "fir.alloca"() <{in_type = i32, operandSegmentSizes = array<i32: 0, 0>}> : () -> !fir.ref<i32>
12021202
}) : (index, index) -> ()
@@ -1208,7 +1208,7 @@ func.func @dc_invalid_parent(%arg0: index, %arg1: index) {
12081208
func.func @dc_invalid_control(%arg0: index, %arg1: index) {
12091209
// expected-error@+2 {{'fir.do_concurrent.loop' op different number of tuple elements for lowerBound, upperBound or step}}
12101210
fir.do_concurrent {
1211-
"fir.do_concurrent.loop"(%arg0, %arg1) <{operandSegmentSizes = array<i32: 1, 1, 0, 0>}> ({
1211+
"fir.do_concurrent.loop"(%arg0, %arg1) <{operandSegmentSizes = array<i32: 1, 1, 0, 0, 0>}> ({
12121212
^bb0(%arg2: index):
12131213
%tmp = "fir.alloca"() <{in_type = i32, operandSegmentSizes = array<i32: 0, 0>}> : () -> !fir.ref<i32>
12141214
}) : (index, index) -> ()
@@ -1221,7 +1221,7 @@ func.func @dc_invalid_control(%arg0: index, %arg1: index) {
12211221
func.func @dc_invalid_ind_var(%arg0: index, %arg1: index) {
12221222
// expected-error@+2 {{'fir.do_concurrent.loop' op expects the same number of induction variables: 2 as bound and step values: 1}}
12231223
fir.do_concurrent {
1224-
"fir.do_concurrent.loop"(%arg0, %arg1, %arg0) <{operandSegmentSizes = array<i32: 1, 1, 1, 0>}> ({
1224+
"fir.do_concurrent.loop"(%arg0, %arg1, %arg0) <{operandSegmentSizes = array<i32: 1, 1, 1, 0, 0>}> ({
12251225
^bb0(%arg3: index, %arg4: index):
12261226
%tmp = "fir.alloca"() <{in_type = i32, operandSegmentSizes = array<i32: 0, 0>}> : () -> !fir.ref<i32>
12271227
}) : (index, index, index) -> ()
@@ -1234,7 +1234,7 @@ func.func @dc_invalid_ind_var(%arg0: index, %arg1: index) {
12341234
func.func @dc_invalid_ind_var_type(%arg0: index, %arg1: index) {
12351235
// expected-error@+2 {{'fir.do_concurrent.loop' op expects arguments for the induction variable to be of index type}}
12361236
fir.do_concurrent {
1237-
"fir.do_concurrent.loop"(%arg0, %arg1, %arg0) <{operandSegmentSizes = array<i32: 1, 1, 1, 0>}> ({
1237+
"fir.do_concurrent.loop"(%arg0, %arg1, %arg0) <{operandSegmentSizes = array<i32: 1, 1, 1, 0, 0>}> ({
12381238
^bb0(%arg3: i32):
12391239
%tmp = "fir.alloca"() <{in_type = i32, operandSegmentSizes = array<i32: 0, 0>}> : () -> !fir.ref<i32>
12401240
}) : (index, index, index) -> ()
@@ -1248,7 +1248,7 @@ func.func @dc_invalid_reduction(%arg0: index, %arg1: index) {
12481248
%sum = fir.alloca i32
12491249
// expected-error@+2 {{'fir.do_concurrent.loop' op mismatch in number of reduction variables and reduction attributes}}
12501250
fir.do_concurrent {
1251-
"fir.do_concurrent.loop"(%arg0, %arg1, %arg0, %sum) <{operandSegmentSizes = array<i32: 1, 1, 1, 1>}> ({
1251+
"fir.do_concurrent.loop"(%arg0, %arg1, %arg0, %sum) <{operandSegmentSizes = array<i32: 1, 1, 1, 1, 0>}> ({
12521252
^bb0(%arg3: index):
12531253
%tmp = "fir.alloca"() <{in_type = i32, operandSegmentSizes = array<i32: 0, 0>}> : () -> !fir.ref<i32>
12541254
}) : (index, index, index, !fir.ref<i32>) -> ()

0 commit comments

Comments
 (0)