abidmalikwaterloo updated this revision to Diff 486732.
abidmalikwaterloo added a comment.
Added a pretty test for omp.distrubute in ops.mlir
Repository:
rG LLVM Github Monorepo
CHANGES SINCE LAST ACTION
https://reviews.llvm.org/D105584/new/
https://reviews.llvm.org/D105584
Files:
mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
mlir/test/Dialect/OpenMP/ops.mlir
Index: mlir/test/Dialect/OpenMP/ops.mlir
===================================================================
--- mlir/test/Dialect/OpenMP/ops.mlir
+++ mlir/test/Dialect/OpenMP/ops.mlir
@@ -133,6 +133,30 @@
return
}
+// CHECK-LABEL: omp_DistributeOp
+func.func @omp_DistributeOp(%lb : index, %ub : index, %step : index, %data_var : memref<i32>, %chunk_var : i32) -> () {
+// CHECK: for (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}})
+ "omp.distribute" (%lb, %ub, %step) ({
+ ^bb0(%iv: index):
+ omp.yield
+ }) {operand_segment_sizes = array<i32: 1,1,1,0,0,0>} :
+ (index, index, index) -> ()
+
+ return
+ }
+
+// CHECK-LABEL: omp_distribute_pretty
+ func.func @omp_distribute_pretty(%lb : index, %ub : index, %step : index, %data_var : memref<i32>, %chunk_var : i32) -> () {
+// CHECK: omp.distribute
+// CHECK-SAME: for (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}})
+ omp.distribute
+ for (%iv) : index = (%lb) to (%ub) step (%step) {
+ omp.yield
+ }
+
+ return
+ }
+
// CHECK-LABEL: omp_wsloop
func.func @omp_wsloop(%lb : index, %ub : index, %step : index, %data_var : memref<i32>, %linear_var : i32, %chunk_var : i32) -> () {
@@ -145,7 +169,7 @@
(index, index, index) -> ()
// CHECK: omp.wsloop linear(%{{.*}} = %{{.*}} : memref<i32>) schedule(static)
- // CHECK-SAMe: for (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}})
+ // CHECK-SAME: for (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}})
"omp.wsloop" (%lb, %ub, %step, %data_var, %linear_var) ({
^bb0(%iv: index):
omp.yield
Index: mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
===================================================================
--- mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -281,6 +281,26 @@
return success();
}
+//===---------------------------------------------------------------------===//
+// Verifier for Dristribute Op
+//===---------------------------------------------------------------------===//
+
+LogicalResult DistributeOp::verify(){
+ if (this->getLowerBound().empty())
+ return emitOpError() << "empty lowerbound for distribute loop operation";
+
+ if (this->getUpperBound().empty())
+ return emitOpError() << "empty upperbound for distribute loop operation";
+
+ if (this->getLowerBound().size() != this->getUpperBound().size())
+ return emitOpError() << "upperbound and lowerbound sizes are not equal";
+
+ if (this->getLowerBound().size() != this->getStep().size())
+ return emitOpError() << "upperbound, lowerbound and step sizes are not equal";
+
+ return success();
+}
+
/// schedule ::= `schedule` `(` sched-list `)`
/// sched-list ::= sched-val | sched-val sched-list |
/// sched-val `,` sched-modifier
@@ -590,7 +610,7 @@
}
//===----------------------------------------------------------------------===//
-// WsLoopOp
+// LoopControl
//===----------------------------------------------------------------------===//
/// loop-control ::= `(` ssa-id-list `)` `:` type `=` loop-bounds
Index: mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
===================================================================
--- mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
+++ mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
@@ -452,7 +452,7 @@
def YieldOp : OpenMP_Op<"yield",
[Pure, ReturnLike, Terminator,
ParentOneOf<["WsLoopOp", "ReductionDeclareOp",
- "AtomicUpdateOp", "SimdLoopOp"]>]> {
+ "AtomicUpdateOp", "SimdLoopOp", "DistributeOp"]>]> {
let summary = "loop yield and termination operation";
let description = [{
"omp.yield" yields SSA values from the OpenMP dialect op region and
@@ -469,6 +469,65 @@
let assemblyFormat = [{ ( `(` $results^ `:` type($results) `)` )? attr-dict}];
}
+//===----------------------------------------------------------------------===//
+// 2.9.4.1 distribute Construct
+//===----------------------------------------------------------------------===//
+
+def DistributeOp : OpenMP_Op<"distribute", [AttrSizedOperandSegments,
+ AllTypesMatch<["lowerBound", "upperBound", "step"]>]> {
+ let summary = "distribute loop construct";
+ let description = [{
+ The distribute construct specifies that the iterations of one or more loop
+ will be executed by the initial teams in the context of their implicit
+ tasks. The iterations are distributed across the initial threads of all
+ initial teams that execute the teams region to which the distribute region
+ binds.
+
+ The distribute loop construct specifies that the iterations of the loop(s)
+ will be executed in parallel by threads in the current context. These
+ iterations are spread across threads that already exist in the enclosing
+ region. The lower and upper bounds specify a half-open range: the
+ range includes the lower bound but does not include the upper bound. If the
+ `inclusive` attribute is specified then the upper bound is also included.
+
+ // TODO: private_var, firstprivate_var and lastprivate_var variables
+ // will be defined later
+
+ The `static_dist_schedule` attribute specifies the schedule for this
+ loop, determining how the loop is distributed across the parallel threads.
+ The optional `schedule_chunk` associated with this determines further
+ controls this distribution.
+
+ // TODO: collapse
+ }];
+
+ let arguments = (ins Variadic<IntLikeType>:$lowerBound,
+ Variadic<IntLikeType>:$upperBound,
+ Variadic<IntLikeType>:$step,
+ UnitAttr:$static_dist_schedule,
+ Optional<IntLikeType>:$schedule_chunk,
+ Variadic<AnyType>:$allocate_vars,
+ Variadic<AnyType>:$allocators_vars,
+ UnitAttr:$inclusive);
+
+ let assemblyFormat = [{
+ oilist(`dist_schedule` `(` `static`
+ ($static_dist_schedule^)?
+ ( $schedule_chunk^ `:` type($schedule_chunk))? `)`
+ | `allocate` `(`
+ custom<AllocateAndAllocator>(
+ $allocate_vars, type($allocate_vars),
+ $allocators_vars, type($allocators_vars)
+ ) `)`
+
+ ) `for` custom<LoopControl>($region, $lowerBound, $upperBound, $step,
+ type($step), $inclusive) attr-dict
+ }];
+
+ let regions = (region AnyRegion:$region);
+ let hasVerifier = 1;
+}
+
//===----------------------------------------------------------------------===//
// 2.10.1 task Construct
//===----------------------------------------------------------------------===//
_______________________________________________
cfe-commits mailing list
cfe-commits@lists.llvm.org
https://lists.llvm.org/cgi-bin/mailman/listinfo/cfe-commits
