[PATCH] D105584: [MLIR][OpenMP] Distribute Construct Operation

[Abid via Phabricator via cfe-commits]

abidmalikwaterloo updated this revision to Diff 473706.
abidmalikwaterloo marked 2 inline comments as done.
abidmalikwaterloo added a comment.

corrected the syntex of the code


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
@@ -121,6 +121,19 @@
   return
 }
 
+// CHECK-LABEL: omp_DistributeOp
+func.func @omp_DistributeOp(%lb : index, %ub : index, %step : index, %data_var : memref<i32>, %chunk_var : i32) -> () {
+  // CHECK: omp.distribute collapse(2)
+  // CHECK-SAME: for (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}})
+  "omp.distribute" (%lb, %ub, %step) ({
+    ^bb0(%iv: index):
+     omp.yield 
+  }) {operand_segment_sizes = dense<[1,1,1,0,0,0]> : vector<6xi32>, collapse_val = 2} :
+    (index, index, index) -> ()
+ 
+ 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) -> () {
 
Index: mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
===================================================================
--- mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -207,6 +207,20 @@
   return success();
 }
 
+//===---------------------------------------------------------------------===//
+// Verifier for Dristribute Op
+//===---------------------------------------------------------------------===//
+
+LogicalResult DistributeOp::verify(){
+  if (this->lowerBound().empty()) 
+    return emitOpError() << "empty lowerbound for distribute loop operation";
+  
+  if (this->upperBound().empty()) 
+    return emitOpError() << "empty upperbound for distribute loop operation";
+  
+  return success();	
+}
+
 /// schedule ::= `schedule` `(` sched-list `)`
 /// sched-list ::= sched-val | sched-val sched-list |
 ///                sched-val `,` sched-modifier
@@ -510,14 +524,14 @@
 }
 
 //===----------------------------------------------------------------------===//
-// WsLoopOp
+// LoopControl
 //===----------------------------------------------------------------------===//
 
 /// loop-control ::= `(` ssa-id-list `)` `:` type `=`  loop-bounds
 /// loop-bounds := `(` ssa-id-list `)` to `(` ssa-id-list `)` inclusive? steps
 /// steps := `step` `(`ssa-id-list`)`
 ParseResult
-parseWsLoopControl(OpAsmParser &parser, Region &region,
+parseLoopControl(OpAsmParser &parser, Region &region,
                    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &lowerBound,
                    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &upperBound,
                    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &steps,
@@ -560,7 +574,7 @@
   return success();
 }
 
-void printWsLoopControl(OpAsmPrinter &p, Operation *op, Region &region,
+void printLoopControl(OpAsmPrinter &p, Operation *op, Region &region,
                         ValueRange lowerBound, ValueRange upperBound,
                         ValueRange steps, TypeRange loopVarTypes,
                         UnitAttr inclusive) {
Index: mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
===================================================================
--- mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
+++ mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
@@ -394,7 +394,7 @@
               custom<ReductionVarList>(
                 $reduction_vars, type($reduction_vars), $reductions
               ) `)`
-    ) `for` custom<WsLoopControl>($region, $lowerBound, $upperBound, $step,
+    ) `for` custom<LoopControl>($region, $lowerBound, $upperBound, $step,
                                   type($step), $inclusive) attr-dict
   }];
   let hasVerifier = 1;
@@ -446,7 +446,7 @@
 def YieldOp : OpenMP_Op<"yield",
     [NoSideEffect, 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
@@ -463,6 +463,68 @@
   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 optional `dist_schedule_var` 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.
+
+    The optional `collapse` attribute specifies the number of loops which
+    are collapsed to form the distribute loop.
+  }];
+
+  let arguments = (ins Variadic<IntLikeType>:$lowerBound,
+             Variadic<IntLikeType>:$upperBound,
+             Variadic<IntLikeType>:$step,
+             UnitAttr:$static_dist_schedule,
+             Optional<IntLikeType>:$schedule_chunk,
+             Confined<OptionalAttr<I64Attr>, [IntMinValue<1>]>:$collapse_val,
+             Variadic<AnyType>:$allocate_vars,
+             Variadic<AnyType>:$allocators_vars,
+             UnitAttr:$inclusive);
+           
+  let assemblyFormat = [{
+    oilist(`dist_schedule` `(` `static`
+    		($static_dist_schedule^)? 
+    		( $schedule_chunk^ `:` type($schedule_chunk))? `)`
+          |`collapse` `(` $collapse_val `)`
+          | `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