[MLIR][OpenMP] Attribute to include WsLoop upperbound

[kiranchandramohan]

Summary

This patch adds an attribute inclusive which if present causes
the upperbound to be included in the loop iteration interval.

mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td

 //===----------------------------------------------------------------------===//
 
 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.
+    parallel 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.
 
     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>
let arguments = (ins Variadic<IntLikeType>:$lowerBound,
              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,
+             UnitAttr:$nowait,
              Confined<OptionalAttr<I64Attr>, [IntMinValue<0>]>:$ordered_val,
-             OptionalAttr<OrderKind>:$order_val);
+             OptionalAttr<OrderKind>:$order_val,
+             UnitAttr:$inclusive);
 
   let builders = [
     OpBuilderDAG<(ins "ValueRange":$lowerBound, "ValueRange":$upperBound,
                   "ValueRange":$step,
                   CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes)>
   ];
 
   let regions = (region AnyRegion:$region);

mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp

void WsLoopOp::build(OpBuilder &builder, OperationState &state,
                      ValueRange lowerBound, ValueRange upperBound,
                      ValueRange step, ArrayRef<NamedAttribute> attributes) {
   build(builder, state, TypeRange(), lowerBound, upperBound, step,
         /*private_vars=*/ValueRange(),
         /*firstprivate_vars=*/ValueRange(), /*lastprivate_vars=*/ValueRange(),
         /*linear_vars=*/ValueRange(), /*linear_step_vars=*/ValueRange(),
         /*schedule_val=*/nullptr, /*schedule_chunk_var=*/nullptr,
         /*collapse_val=*/nullptr,
-        /*nowait=*/nullptr, /*ordered_val=*/nullptr, /*order_val=*/nullptr);
+        /*nowait=*/false, /*ordered_val=*/nullptr, /*order_val=*/nullptr,
+        /*inclusive=*/false);
   state.addAttributes(attributes);
 }
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/OpenMP/OpenMPOps.cpp.inc"

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

auto bodyGen = [&](llvm::OpenMPIRBuilder::InsertPointTy ip, llvm::Value *iv) {
     // Convert the body of the loop.
     convertOmpOpRegions(loop.region(), "omp.wsloop.region", valueMapping,
                         blockMapping, *entryBlock, *exitBlock, builder,
                         bodyGenStatus);
   };
 
   // Delegate actual loop construction to the OpenMP IRBuilder.
   // TODO: this currently assumes WsLoop is semantically similar to SCF loop,
-  // i.e. it has a positive step, uses signed integer semantics, and its upper
-  // bound is not included. Reconsider this code when WsLoop clearly supports
-  // more cases.
+  // i.e. it has a positive step, uses signed integer semantics. Reconsider
+  // this code when WsLoop clearly supports more cases.
   llvm::BasicBlock *insertBlock = builder.GetInsertBlock();
   llvm::CanonicalLoopInfo *loopInfo = ompBuilder->createCanonicalLoop(
       ompLoc, bodyGen, lowerBound, upperBound, step, /*IsSigned=*/true,
-      /*InclusiveStop=*/false);
+      /*InclusiveStop=*/loop.inclusive());
   if (failed(bodyGenStatus))
     return failure();
 
   // TODO: get the alloca insertion point from the parallel operation builder.
   // If we insert the at the top of the current function, they will be passed as
   // extra arguments into the function the parallel operation builder outlines.
   // Put them at the start of the current block for now.
   llvm::OpenMPIRBuilder::InsertPointTy allocaIP(
       insertBlock, insertBlock->getFirstInsertionPt());
-  loopInfo = ompBuilder->createStaticWorkshareLoop(
-      ompLoc, loopInfo, allocaIP,
-      !loop.nowait().hasValue() || loop.nowait().getValue(), chunk);
+  loopInfo = ompBuilder->createStaticWorkshareLoop(ompLoc, loopInfo, allocaIP,
+                                                   !loop.nowait(), chunk);
 
   // Continue building IR after the loop.
   builder.restoreIP(loopInfo->getAfterIP());
   return success();
 }
 
 /// Given an OpenMP MLIR operation, create the corresponding LLVM IR
 /// (including OpenMP runtime calls).

mlir/test/Target/openmp-llvm.mlir

^bb0(%arg1: i64):
       llvm.store %3, %4 : !llvm.ptr<float>
       omp.yield
       // CHECK: call void @__kmpc_for_static_fini(%struct.ident_t* @[[$wsloop_loc_struct]],
     }) {operand_segment_sizes = dense<[1, 1, 1, 0, 0, 0, 0, 0, 0]> : vector<9xi32>} : (i64, i64, i64) -> ()
     omp.terminator
   }
   llvm.return
 }
+
+// CHECK-LABEL: @wsloop_inclusive_1
+llvm.func @wsloop_inclusive_1(%arg0: !llvm.ptr<float>) {
+  %0 = llvm.mlir.constant(42 : index) : !llvm.i64
+  %1 = llvm.mlir.constant(10 : index) : !llvm.i64
+  %2 = llvm.mlir.constant(1 : index) : !llvm.i64
+  // CHECK: store i64 31, i64* %{{.*}}upperbound
+  "omp.wsloop"(%1, %0, %2) ( {
+  ^bb0(%arg1: !llvm.i64):
+    %3 = llvm.mlir.constant(2.000000e+00 : f32) : !llvm.float
+    %4 = llvm.getelementptr %arg0[%arg1] : (!llvm.ptr<float>, !llvm.i64) -> !llvm.ptr<float>
+    llvm.store %3, %4 : !llvm.ptr<float>
+    omp.yield
+  }) {operand_segment_sizes = dense<[1, 1, 1, 0, 0, 0, 0, 0, 0]> : vector<9xi32>} : (!llvm.i64, !llvm.i64, !llvm.i64) -> ()
+  llvm.return
+}
+
+// CHECK-LABEL: @wsloop_inclusive_2
+llvm.func @wsloop_inclusive_2(%arg0: !llvm.ptr<float>) {
+  %0 = llvm.mlir.constant(42 : index) : !llvm.i64
+  %1 = llvm.mlir.constant(10 : index) : !llvm.i64
+  %2 = llvm.mlir.constant(1 : index) : !llvm.i64
+  // CHECK: store i64 32, i64* %{{.*}}upperbound
+  "omp.wsloop"(%1, %0, %2) ( {
+  ^bb0(%arg1: !llvm.i64):
+    %3 = llvm.mlir.constant(2.000000e+00 : f32) : !llvm.float
+    %4 = llvm.getelementptr %arg0[%arg1] : (!llvm.ptr<float>, !llvm.i64) -> !llvm.ptr<float>
+    llvm.store %3, %4 : !llvm.ptr<float>
+    omp.yield
+  }) {inclusive, operand_segment_sizes = dense<[1, 1, 1, 0, 0, 0, 0, 0, 0]> : vector<9xi32>} : (!llvm.i64, !llvm.i64, !llvm.i64) -> ()
+  llvm.return
+}