[llvm] 8435250 - [MLIR][OpenMP] Add omp.wsloop operation

Kiran Chandramohan via llvm-commits llvm-commits at lists.llvm.org
Mon Nov 16 07:25:30 PST 2020


Author: David Truby
Date: 2020-11-16T15:24:57Z
New Revision: 843525075b87b367f3d18754ef0a3e35dc2ed162

URL: https://github.com/llvm/llvm-project/commit/843525075b87b367f3d18754ef0a3e35dc2ed162
DIFF: https://github.com/llvm/llvm-project/commit/843525075b87b367f3d18754ef0a3e35dc2ed162.diff

LOG: [MLIR][OpenMP] Add omp.wsloop operation

This adds a simple definition of a "workshare loop" operation for
the OpenMP MLIR dialect, excluding the "reduction" and "allocate"
clauses and without a custom parser and pretty printer.

The schedule clause also does not yet accept the modifiers that are
permitted in OpenMP 5.0.

Co-authored-by: Kiran Chandramohan <kiran.chandramohan at arm.com>

Reviewed By: ftynse, clementval

Differential Revision: https://reviews.llvm.org/D86071

Added: 
    

Modified: 
    llvm/include/llvm/Frontend/OpenMP/OMP.td
    mlir/include/mlir/Dialect/OpenMP/OpenMPDialect.h
    mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
    mlir/test/Dialect/OpenMP/ops.mlir

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/Frontend/OpenMP/OMP.td b/llvm/include/llvm/Frontend/OpenMP/OMP.td
index 7abfdd77bea9..1ecf075c38a9 100644
--- a/llvm/include/llvm/Frontend/OpenMP/OMP.td
+++ b/llvm/include/llvm/Frontend/OpenMP/OMP.td
@@ -116,10 +116,29 @@ def OMPC_ProcBind : Clause<"proc_bind"> {
     OMP_PROC_BIND_unknown
   ];
 }
+
+// static and auto are C++ keywords so need a capital to disambiguate.
+def OMP_SCHEDULE_Static : ClauseVal<"Static", 2, 1> {}
+def OMP_SCHEDULE_Dynamic : ClauseVal<"Dynamic", 3, 1> {}
+def OMP_SCHEDULE_Guided : ClauseVal<"Guided", 4, 1> {}
+def OMP_SCHEDULE_Auto : ClauseVal<"Auto", 5, 1> {}
+def OMP_SCHEDULE_Runtime : ClauseVal<"Runtime", 6, 1> {}
+def OMP_SCHEDULE_Default : ClauseVal<"Default", 7, 0> { let isDefault = 1; }
+
 def OMPC_Schedule : Clause<"schedule"> {
   let clangClass = "OMPScheduleClause";
   let flangClass = "OmpScheduleClause";
+  let enumClauseValue = "ScheduleKind";
+  let allowedClauseValues = [
+    OMP_SCHEDULE_Static,
+    OMP_SCHEDULE_Dynamic,
+    OMP_SCHEDULE_Guided,
+    OMP_SCHEDULE_Auto,
+    OMP_SCHEDULE_Runtime,
+    OMP_SCHEDULE_Default
+  ];
 }
+
 def OMPC_Ordered : Clause<"ordered"> {
   let clangClass = "OMPOrderedClause";
   let flangClassValue = "ScalarIntConstantExpr";
@@ -239,8 +258,14 @@ def OMPC_Allocate : Clause<"allocate"> {
 def OMPC_NonTemporal : Clause<"nontemporal"> {
   let clangClass = "OMPNontemporalClause";
 }
+
+def OMP_ORDER_concurrent : ClauseVal<"default",2,0> { let isDefault = 1; }
 def OMPC_Order : Clause<"order"> {
   let clangClass = "OMPOrderClause";
+  let enumClauseValue = "OrderKind";
+  let allowedClauseValues = [
+    OMP_ORDER_concurrent
+  ];
 }
 def OMPC_Destroy : Clause<"destroy"> {
   let clangClass = "OMPDestroyClause";

diff  --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPDialect.h b/mlir/include/mlir/Dialect/OpenMP/OpenMPDialect.h
index 0715b9ddd394..8f79c4af1ad8 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPDialect.h
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPDialect.h
@@ -13,8 +13,11 @@
 #ifndef MLIR_DIALECT_OPENMP_OPENMPDIALECT_H_
 #define MLIR_DIALECT_OPENMP_OPENMPDIALECT_H_
 
+#include "mlir/Dialect/LLVMIR/LLVMTypes.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/OpDefinition.h"
+#include "mlir/Interfaces/ControlFlowInterfaces.h"
+#include "mlir/Interfaces/SideEffectInterfaces.h"
 
 #include "mlir/Dialect/OpenMP/OpenMPOpsDialect.h.inc"
 #include "mlir/Dialect/OpenMP/OpenMPOpsEnums.h.inc"

diff  --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
index 17845a598c31..d42466af9e17 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
@@ -15,6 +15,9 @@
 #define OPENMP_OPS
 
 include "mlir/IR/OpBase.td"
+include "mlir/Interfaces/SideEffectInterfaces.td"
+include "mlir/Interfaces/ControlFlowInterfaces.td"
+include "mlir/Dialect/LLVMIR/LLVMOpBase.td"
 include "mlir/Dialect/OpenMP/OmpCommon.td"
 
 def OpenMP_Dialect : Dialect {
@@ -25,6 +28,10 @@ def OpenMP_Dialect : Dialect {
 class OpenMP_Op<string mnemonic, list<OpTrait> traits = []> :
       Op<OpenMP_Dialect, mnemonic, traits>;
 
+// Type which can be constraint accepting standard integers, indices and
+// LLVM integer types.
+def IntLikeType : AnyTypeOf<[AnyInteger, Index, LLVM_AnyInteger]>;
+
 //===----------------------------------------------------------------------===//
 // 2.6 parallel Construct
 //===----------------------------------------------------------------------===//
@@ -102,6 +109,91 @@ def TerminatorOp : OpenMP_Op<"terminator", [Terminator]> {
   let assemblyFormat = "attr-dict";
 }
 
+//===----------------------------------------------------------------------===//
+// 2.9.2 Workshare Loop Construct
+//===----------------------------------------------------------------------===//
+
+def WsLoopOp : OpenMP_Op<"wsloop", [AttrSizedOperandSegments]> {
+  let summary = "workshare loop construct";
+  let description = [{
+    The workshare 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
+    parallel region.
+
+    The body region can contain any number of blocks. The region is terminated
+    by "omp.yield" instruction without operands.
+
+    ```
+      omp.wsloop (%i1, %i2) = (%c0, %c0) to (%c10, %c10) step (%c1, %c1) {
+        %a = load %arrA[%i1, %i2] : memref<?x?xf32>
+        %b = load %arrB[%i1, %i2] : memref<?x?xf32>
+        %sum = addf %a, %b : f32
+        store %sum, %arrC[%i1, %i2] : memref<?x?xf32>
+        omp.yield
+      }
+    ```
+
+    `private_vars`, `firstprivate_vars`, `lastprivate_vars` and `linear_vars`
+    arguments are variadic list of operands that specify the data sharing
+    attributes of the list of values. The `linear_step_vars` operand
+    additionally specifies the step for each associated linear operand. Note
+    that the `linear_vars` and `linear_step_vars` variadic lists should contain
+    the same number of elements.
+
+    The optional `schedule_val` attribute specifies the loop schedule for this
+    loop, determining how the loop is distributed across the parallel threads.
+    The optional `schedule_chunk_var` associated with this determines further
+    controls this distribution.
+
+    The optional `collapse_val` attribute specifies the number of loops which
+    are collapsed to form the worksharing loop.
+
+    The `nowait` attribute, when present, signifies that there should be no
+    implicit barrier at the end of the loop.
+
+    The optional `ordered_val` attribute specifies how many loops are associated
+    with the do loop construct.
+
+    The optional `order` attribute specifies which order the iterations of the
+    associate loops are executed in. Currently the only option for this
+    attribute is "concurrent".
+  }];
+
+  let arguments = (ins Variadic<IntLikeType>:$lowerBound,
+             Variadic<IntLikeType>:$upperBound,
+             Variadic<IntLikeType>:$step,
+             Variadic<AnyType>:$private_vars,
+             Variadic<AnyType>:$firstprivate_vars,
+             Variadic<AnyType>:$lastprivate_vars,
+             Variadic<AnyType>:$linear_vars,
+             Variadic<AnyType>:$linear_step_vars,
+             OptionalAttr<ScheduleKind>:$schedule_val,
+             Optional<AnyType>:$schedule_chunk_var,
+             Confined<OptionalAttr<I64Attr>, [IntMinValue<0>]>:$collapse_val,
+             OptionalAttr<UnitAttr>:$nowait,
+             Confined<OptionalAttr<I64Attr>, [IntMinValue<0>]>:$ordered_val,
+             OptionalAttr<OrderKind>:$order_val);
+
+  let regions = (region AnyRegion:$region);
+}
+
+def YieldOp : OpenMP_Op<"yield", [NoSideEffect, ReturnLike, Terminator,
+                                  HasParent<"WsLoopOp">]> {
+  let summary = "loop yield and termination operation";
+  let description = [{
+    "omp.yield" yields SSA values from the OpenMP dialect op region and
+    terminates the region. The semantics of how the values are yielded is
+    defined by the parent operation.
+    If "omp.yield" has any operands, the operands must match the parent
+    operation's results.
+  }];
+
+  let arguments = (ins Variadic<AnyType>:$results);
+
+  let assemblyFormat = [{ ( `(` $results^ `:` type($results) `)` )? attr-dict}];
+}
+
 //===----------------------------------------------------------------------===//
 // 2.10.4 taskyield Construct
 //===----------------------------------------------------------------------===//

diff  --git a/mlir/test/Dialect/OpenMP/ops.mlir b/mlir/test/Dialect/OpenMP/ops.mlir
index add50414d71c..c74b8bc3327d 100644
--- a/mlir/test/Dialect/OpenMP/ops.mlir
+++ b/mlir/test/Dialect/OpenMP/ops.mlir
@@ -125,3 +125,41 @@ func @omp_parallel_pretty(%data_var : memref<i32>, %if_cond : i1, %num_threads :
 
   return
 }
+
+func @omp_wsloop(%lb : index, %ub : index, %step : index,
+                 %data_var : memref<i32>, %linear_var : si32, %chunk_var : si32) -> () {
+
+  // CHECK: "omp.wsloop"(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}})
+  "omp.wsloop" (%lb, %ub, %step, %data_var) ({
+    omp.yield
+  }) {operand_segment_sizes = dense<[1,1,1,1,0,0,0,0,0]> : vector<9xi32>, collapse_val = 2, ordered_val = 1} :
+    (index, index, index, memref<i32>) -> ()
+
+  // CHECK: "omp.wsloop"(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}})
+  "omp.wsloop" (%lb, %lb, %ub, %ub, %step, %step, %data_var) ({
+    omp.yield
+  }) {operand_segment_sizes = dense<[2,2,2,1,0,0,0,0,0]> : vector<9xi32>, collapse_val = 2, ordered_val = 1} :
+    (index, index, index, index, index, index, memref<i32>) -> ()
+
+
+  // CHECK: "omp.wsloop"(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}})
+  "omp.wsloop" (%lb, %ub, %step, %data_var, %linear_var) ({
+    omp.yield
+  }) {operand_segment_sizes = dense<[1,1,1,0,0,0,1,1,0]> : vector<9xi32>, schedule_val = "Static"} :
+    (index, index, index, memref<i32>, si32) -> ()
+
+  // CHECK: "omp.wsloop"(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}})
+  "omp.wsloop" (%lb, %ub, %step, %data_var, %data_var, %data_var, %data_var, %linear_var, %chunk_var) ({
+    omp.yield
+  }) {operand_segment_sizes = dense<[1,1,1,1,1,1,1,1,1]> : vector<9xi32>, schedule_val = "Dynamic", collapse_val = 3, ordered_val = 2} :
+    (index, index, index, memref<i32>, memref<i32>, memref<i32>, memref<i32>, si32, si32) -> ()
+
+  // CHECK: "omp.wsloop"(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}})
+  "omp.wsloop" (%lb, %ub, %step, %data_var) ({
+    omp.yield
+  }) {operand_segment_sizes = dense<[1,1,1,1,0,0,0,0,0]> : vector<9xi32>, nowait, schedule_val = "Auto"} :
+    (index, index, index, memref<i32>) -> ()
+
+
+  return
+}


        


More information about the llvm-commits mailing list