forked from OSchip/llvm-project
Add Allocate Clause to MLIR Parallel Operation Definition
Differential Revision: https://reviews.llvm.org/D87684
This commit is contained in:
Irina Dobrescu
parent
4212533961
commit
65b9b9aa50
|
|
@ -109,7 +109,7 @@ genOMP(Fortran::lower::AbstractConverter &converter,
|
|||
auto parallelOp = firOpBuilder.create<mlir::omp::ParallelOp>(
|
||||
currentLocation, argTy, Value(), numThreads,
|
||||
defaultValue.dyn_cast_or_null<StringAttr>(), ValueRange(), ValueRange(),
|
||||
ValueRange(), ValueRange(),
|
||||
ValueRange(), ValueRange(), ValueRange(), ValueRange(),
|
||||
procBindValue.dyn_cast_or_null<StringAttr>());
|
||||
firOpBuilder.createBlock(¶llelOp.getRegion());
|
||||
auto &block = parallelOp.getRegion().back();
|
||||
|
|
|
|||
|
|
@ -65,6 +65,9 @@ def ParallelOp : OpenMP_Op<"parallel", [AttrSizedOperandSegments]> {
|
|||
are a variadic list of values that specify the data sharing attribute of
|
||||
those values.
|
||||
|
||||
The $allocators_vars and $allocate_vars parameters are a variadic list of values
|
||||
that specify the memory allocator to be used to obtain storage for private values.
|
||||
|
||||
The optional $proc_bind_val attribute controls the thread affinity for the execution
|
||||
of the parallel region.
|
||||
}];
|
||||
|
|
@ -76,12 +79,15 @@ def ParallelOp : OpenMP_Op<"parallel", [AttrSizedOperandSegments]> {
|
|||
Variadic<AnyType>:$firstprivate_vars,
|
||||
Variadic<AnyType>:$shared_vars,
|
||||
Variadic<AnyType>:$copyin_vars,
|
||||
Variadic<AnyType>:$allocate_vars,
|
||||
Variadic<AnyType>:$allocators_vars,
|
||||
OptionalAttr<ProcBindKind>:$proc_bind_val);
|
||||
|
||||
let regions = (region AnyRegion:$region);
|
||||
|
||||
let parser = [{ return parseParallelOp(parser, result); }];
|
||||
let printer = [{ return printParallelOp(p, *this); }];
|
||||
let verifier = [{ return ::verifyParallelOp(*this); }];
|
||||
}
|
||||
|
||||
def TerminatorOp : OpenMP_Op<"terminator", [Terminator]> {
|
||||
|
|
|
|||
|
|
@ -41,7 +41,7 @@ void OpenMPDialect::initialize() {
|
|||
///
|
||||
/// operand-and-type-list ::= `(` ssa-id-and-type-list `)`
|
||||
/// ssa-id-and-type-list ::= ssa-id-and-type |
|
||||
/// ssa-id-and-type ',' ssa-id-and-type-list
|
||||
/// ssa-id-and-type `,` ssa-id-and-type-list
|
||||
/// ssa-id-and-type ::= ssa-id `:` type
|
||||
static ParseResult
|
||||
parseOperandAndTypeList(OpAsmParser &parser,
|
||||
|
|
@ -65,6 +65,52 @@ parseOperandAndTypeList(OpAsmParser &parser,
|
|||
return success();
|
||||
}
|
||||
|
||||
/// Parse an allocate clause with allocators and a list of operands with types.
|
||||
///
|
||||
/// operand-and-type-list ::= `(` allocate-operand-list `)`
|
||||
/// allocate-operand-list :: = allocate-operand |
|
||||
/// allocator-operand `,` allocate-operand-list
|
||||
/// allocate-operand :: = ssa-id-and-type -> ssa-id-and-type
|
||||
/// ssa-id-and-type ::= ssa-id `:` type
|
||||
static ParseResult parseAllocateAndAllocator(
|
||||
OpAsmParser &parser,
|
||||
SmallVectorImpl<OpAsmParser::OperandType> &operandsAllocate,
|
||||
SmallVectorImpl<Type> &typesAllocate,
|
||||
SmallVectorImpl<OpAsmParser::OperandType> &operandsAllocator,
|
||||
SmallVectorImpl<Type> &typesAllocator) {
|
||||
if (parser.parseLParen())
|
||||
return failure();
|
||||
|
||||
do {
|
||||
OpAsmParser::OperandType operand;
|
||||
Type type;
|
||||
|
||||
if (parser.parseOperand(operand) || parser.parseColonType(type))
|
||||
return failure();
|
||||
operandsAllocator.push_back(operand);
|
||||
typesAllocator.push_back(type);
|
||||
if (parser.parseArrow())
|
||||
return failure();
|
||||
if (parser.parseOperand(operand) || parser.parseColonType(type))
|
||||
return failure();
|
||||
|
||||
operandsAllocate.push_back(operand);
|
||||
typesAllocate.push_back(type);
|
||||
} while (succeeded(parser.parseOptionalComma()));
|
||||
|
||||
if (parser.parseRParen())
|
||||
return failure();
|
||||
|
||||
return success();
|
||||
}
|
||||
|
||||
static LogicalResult verifyParallelOp(ParallelOp op) {
|
||||
if (op.allocate_vars().size() != op.allocators_vars().size())
|
||||
return op.emitError(
|
||||
"expected equal sizes for allocate and allocator variables");
|
||||
return success();
|
||||
}
|
||||
|
||||
static void printParallelOp(OpAsmPrinter &p, ParallelOp op) {
|
||||
p << "omp.parallel";
|
||||
|
||||
|
|
@ -84,10 +130,26 @@ static void printParallelOp(OpAsmPrinter &p, ParallelOp op) {
|
|||
}
|
||||
}
|
||||
};
|
||||
|
||||
// Print allocator and allocate parameters
|
||||
auto printAllocateAndAllocator = [&p](OperandRange varsAllocate,
|
||||
OperandRange varsAllocator) {
|
||||
if (varsAllocate.empty())
|
||||
return;
|
||||
|
||||
p << " allocate(";
|
||||
for (unsigned i = 0; i < varsAllocate.size(); ++i) {
|
||||
std::string separator = i == varsAllocate.size() - 1 ? ")" : ", ";
|
||||
p << varsAllocator[i] << " : " << varsAllocator[i].getType() << " -> ";
|
||||
p << varsAllocate[i] << " : " << varsAllocate[i].getType() << separator;
|
||||
}
|
||||
};
|
||||
|
||||
printDataVars("private", op.private_vars());
|
||||
printDataVars("firstprivate", op.firstprivate_vars());
|
||||
printDataVars("shared", op.shared_vars());
|
||||
printDataVars("copyin", op.copyin_vars());
|
||||
printAllocateAndAllocator(op.allocate_vars(), op.allocators_vars());
|
||||
|
||||
if (auto def = op.default_val())
|
||||
p << " default(" << def->drop_front(3) << ")";
|
||||
|
|
@ -118,6 +180,7 @@ static ParseResult allowedOnce(OpAsmParser &parser, llvm::StringRef clause,
|
|||
/// firstprivate ::= `firstprivate` operand-and-type-list
|
||||
/// shared ::= `shared` operand-and-type-list
|
||||
/// copyin ::= `copyin` operand-and-type-list
|
||||
/// allocate ::= `allocate` operand-and-type `->` operand-and-type-list
|
||||
/// default ::= `default` `(` (`private` | `firstprivate` | `shared` | `none`)
|
||||
/// procBind ::= `proc_bind` `(` (`master` | `close` | `spread`) `)`
|
||||
///
|
||||
|
|
@ -134,7 +197,11 @@ static ParseResult parseParallelOp(OpAsmParser &parser,
|
|||
SmallVector<Type, 4> sharedTypes;
|
||||
SmallVector<OpAsmParser::OperandType, 4> copyins;
|
||||
SmallVector<Type, 4> copyinTypes;
|
||||
std::array<int, 6> segments{0, 0, 0, 0, 0, 0};
|
||||
SmallVector<OpAsmParser::OperandType, 4> allocates;
|
||||
SmallVector<Type, 4> allocateTypes;
|
||||
SmallVector<OpAsmParser::OperandType, 4> allocators;
|
||||
SmallVector<Type, 4> allocatorTypes;
|
||||
std::array<int, 8> segments{0, 0, 0, 0, 0, 0, 0, 0};
|
||||
llvm::StringRef keyword;
|
||||
bool defaultVal = false;
|
||||
bool procBind = false;
|
||||
|
|
@ -145,6 +212,8 @@ static ParseResult parseParallelOp(OpAsmParser &parser,
|
|||
const int firstprivateClausePos = 3;
|
||||
const int sharedClausePos = 4;
|
||||
const int copyinClausePos = 5;
|
||||
const int allocateClausePos = 6;
|
||||
const int allocatorPos = 7;
|
||||
const llvm::StringRef opName = result.name.getStringRef();
|
||||
|
||||
while (succeeded(parser.parseOptionalKeyword(&keyword))) {
|
||||
|
|
@ -192,6 +261,15 @@ static ParseResult parseParallelOp(OpAsmParser &parser,
|
|||
if (parseOperandAndTypeList(parser, copyins, copyinTypes))
|
||||
return failure();
|
||||
segments[copyinClausePos] = copyins.size();
|
||||
} else if (keyword == "allocate") {
|
||||
// fail if there was already another allocate clause
|
||||
if (segments[allocateClausePos])
|
||||
return allowedOnce(parser, "allocate", opName);
|
||||
if (parseAllocateAndAllocator(parser, allocates, allocateTypes,
|
||||
allocators, allocatorTypes))
|
||||
return failure();
|
||||
segments[allocateClausePos] = allocates.size();
|
||||
segments[allocatorPos] = allocators.size();
|
||||
} else if (keyword == "default") {
|
||||
// fail if there was already another default clause
|
||||
if (defaultVal)
|
||||
|
|
@ -261,6 +339,18 @@ static ParseResult parseParallelOp(OpAsmParser &parser,
|
|||
result.operands))
|
||||
return failure();
|
||||
|
||||
// Add allocate parameters
|
||||
if (segments[allocateClausePos] &&
|
||||
parser.resolveOperands(allocates, allocateTypes, allocates[0].location,
|
||||
result.operands))
|
||||
return failure();
|
||||
|
||||
// Add allocator parameters
|
||||
if (segments[allocatorPos] &&
|
||||
parser.resolveOperands(allocators, allocatorTypes, allocators[0].location,
|
||||
result.operands))
|
||||
return failure();
|
||||
|
||||
result.addAttribute("operand_segment_sizes",
|
||||
parser.getBuilder().getI32VectorAttr(segments));
|
||||
|
||||
|
|
|
|||
|
|
@ -52,37 +52,43 @@ func @omp_terminator() -> () {
|
|||
}
|
||||
|
||||
func @omp_parallel(%data_var : memref<i32>, %if_cond : i1, %num_threads : si32) -> () {
|
||||
// CHECK: omp.parallel if(%{{.*}}) num_threads(%{{.*}} : si32) private(%{{.*}} : memref<i32>) firstprivate(%{{.*}} : memref<i32>) shared(%{{.*}} : memref<i32>) copyin(%{{.*}} : memref<i32>)
|
||||
"omp.parallel" (%if_cond, %num_threads, %data_var, %data_var, %data_var, %data_var) ({
|
||||
// CHECK: omp.parallel if(%{{.*}}) num_threads(%{{.*}} : si32) private(%{{.*}} : memref<i32>) firstprivate(%{{.*}} : memref<i32>) shared(%{{.*}} : memref<i32>) copyin(%{{.*}} : memref<i32>) allocate(%{{.*}} : memref<i32> -> %{{.*}} : memref<i32>)
|
||||
"omp.parallel" (%if_cond, %num_threads, %data_var, %data_var, %data_var, %data_var, %data_var, %data_var) ({
|
||||
|
||||
// test without if condition
|
||||
// CHECK: omp.parallel num_threads(%{{.*}} : si32) private(%{{.*}} : memref<i32>) firstprivate(%{{.*}} : memref<i32>) shared(%{{.*}} : memref<i32>) copyin(%{{.*}} : memref<i32>)
|
||||
"omp.parallel"(%num_threads, %data_var, %data_var, %data_var, %data_var) ({
|
||||
// CHECK: omp.parallel num_threads(%{{.*}} : si32) private(%{{.*}} : memref<i32>) firstprivate(%{{.*}} : memref<i32>) shared(%{{.*}} : memref<i32>) copyin(%{{.*}} : memref<i32>) allocate(%{{.*}} : memref<i32> -> %{{.*}} : memref<i32>)
|
||||
"omp.parallel"(%num_threads, %data_var, %data_var, %data_var, %data_var, %data_var, %data_var) ({
|
||||
omp.terminator
|
||||
}) {operand_segment_sizes = dense<[0,1,1,1,1,1]>: vector<6xi32>, default_val = "defshared"} : (si32, memref<i32>, memref<i32>, memref<i32>, memref<i32>) -> ()
|
||||
}) {operand_segment_sizes = dense<[0,1,1,1,1,1,1,1]>: vector<8xi32>, default_val = "defshared"} : (si32, memref<i32>, memref<i32>, memref<i32>, memref<i32>, memref<i32>, memref<i32>) -> ()
|
||||
|
||||
// CHECK: omp.barrier
|
||||
omp.barrier
|
||||
|
||||
// test without num_threads
|
||||
// CHECK: omp.parallel if(%{{.*}}) private(%{{.*}} : memref<i32>) firstprivate(%{{.*}} : memref<i32>) shared(%{{.*}} : memref<i32>) copyin(%{{.*}} : memref<i32>)
|
||||
"omp.parallel"(%if_cond, %data_var, %data_var, %data_var, %data_var) ({
|
||||
// CHECK: omp.parallel if(%{{.*}}) private(%{{.*}} : memref<i32>) firstprivate(%{{.*}} : memref<i32>) shared(%{{.*}} : memref<i32>) copyin(%{{.*}} : memref<i32>) allocate(%{{.*}} : memref<i32> -> %{{.*}} : memref<i32>)
|
||||
"omp.parallel"(%if_cond, %data_var, %data_var, %data_var, %data_var, %data_var, %data_var) ({
|
||||
omp.terminator
|
||||
}) {operand_segment_sizes = dense<[1,0,1,1,1,1]> : vector<6xi32>} : (i1, memref<i32>, memref<i32>, memref<i32>, memref<i32>) -> ()
|
||||
}) {operand_segment_sizes = dense<[1,0,1,1,1,1,1,1]> : vector<8xi32>} : (i1, memref<i32>, memref<i32>, memref<i32>, memref<i32>, memref<i32>, memref<i32>) -> ()
|
||||
|
||||
// test without allocate
|
||||
// CHECK: omp.parallel if(%{{.*}}) num_threads(%{{.*}} : si32) private(%{{.*}} : memref<i32>) firstprivate(%{{.*}} : memref<i32>) shared(%{{.*}} : memref<i32>) copyin(%{{.*}} : memref<i32>)
|
||||
"omp.parallel"(%if_cond, %num_threads, %data_var, %data_var, %data_var, %data_var) ({
|
||||
omp.terminator
|
||||
}) {operand_segment_sizes = dense<[1,1,1,1,1,1,0,0]> : vector<8xi32>} : (i1, si32, memref<i32>, memref<i32>, memref<i32>, memref<i32>) -> ()
|
||||
|
||||
omp.terminator
|
||||
}) {operand_segment_sizes = dense<[1,1,1,1,1,1]> : vector<6xi32>, proc_bind_val = "spread"} : (i1, si32, memref<i32>, memref<i32>, memref<i32>, memref<i32>) -> ()
|
||||
}) {operand_segment_sizes = dense<[1,1,1,1,1,1,1,1]> : vector<8xi32>, proc_bind_val = "spread"} : (i1, si32, memref<i32>, memref<i32>, memref<i32>, memref<i32>, memref<i32>, memref<i32>) -> ()
|
||||
|
||||
// test with multiple parameters for single variadic argument
|
||||
// CHECK: omp.parallel private(%{{.*}} : memref<i32>) firstprivate(%{{.*}} : memref<i32>, %{{.*}} : memref<i32>) shared(%{{.*}} : memref<i32>) copyin(%{{.*}} : memref<i32>)
|
||||
"omp.parallel" (%data_var, %data_var, %data_var, %data_var, %data_var) ({
|
||||
// CHECK: omp.parallel private(%{{.*}} : memref<i32>) firstprivate(%{{.*}} : memref<i32>, %{{.*}} : memref<i32>) shared(%{{.*}} : memref<i32>) copyin(%{{.*}} : memref<i32>) allocate(%{{.*}} : memref<i32> -> %{{.*}} : memref<i32>)
|
||||
"omp.parallel" (%data_var, %data_var, %data_var, %data_var, %data_var, %data_var, %data_var) ({
|
||||
omp.terminator
|
||||
}) {operand_segment_sizes = dense<[0,0,1,2,1,1]> : vector<6xi32>} : (memref<i32>, memref<i32>, memref<i32>, memref<i32>, memref<i32>) -> ()
|
||||
}) {operand_segment_sizes = dense<[0,0,1,2,1,1,1,1]> : vector<8xi32>} : (memref<i32>, memref<i32>, memref<i32>, memref<i32>, memref<i32>, memref<i32>, memref<i32>) -> ()
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
func @omp_parallel_pretty(%data_var : memref<i32>, %if_cond : i1, %num_threads : si32) -> () {
|
||||
func @omp_parallel_pretty(%data_var : memref<i32>, %if_cond : i1, %num_threads : si32, %allocator : si32) -> () {
|
||||
// CHECK: omp.parallel
|
||||
omp.parallel {
|
||||
omp.terminator
|
||||
|
|
@ -93,6 +99,11 @@ func @omp_parallel_pretty(%data_var : memref<i32>, %if_cond : i1, %num_threads :
|
|||
omp.terminator
|
||||
}
|
||||
|
||||
// CHECK: omp.parallel allocate(%{{.*}} : memref<i32> -> %{{.*}} : memref<i32>)
|
||||
omp.parallel allocate(%data_var : memref<i32> -> %data_var : memref<i32>) {
|
||||
omp.terminator
|
||||
}
|
||||
|
||||
// CHECK: omp.parallel private(%{{.*}} : memref<i32>, %{{.*}} : memref<i32>) firstprivate(%{{.*}} : memref<i32>)
|
||||
omp.parallel private(%data_var : memref<i32>, %data_var : memref<i32>) firstprivate(%data_var : memref<i32>) {
|
||||
omp.terminator
|
||||
|
|
|
|||
Loading…
Reference in New Issue