[MLIR][OpenMP] Pretty printer and parser for omp.wsloop

[kiranchandramohan]

Summary

Co-authored-by: David Truby <[e-mail]>

[ftynse]

I find the linear(%data_var : memref<i32> -> %linear_var) a bit confusing, we are using an equivalent of linear(%linear_var = %data_var : memref<i32>) in SCF and GPU dialects, and generally the %within_region = %outside_region : type pattern applies even more operations. -> seems to be predominantly used for results, which makes it confusing for me here.

[kiranchandramohan]

Fixed failure in SCFToOpenMP.
Converted "->" to "=" for linear clause.

[kiranchandramohan]

In D92327#2436879, @ftynse wrote:

I find the linear(%data_var : memref<i32> -> %linear_var) a bit confusing, we are using an equivalent of linear(%linear_var = %data_var : memref<i32>) in SCF and GPU dialects, and generally the %within_region = %outside_region : type pattern applies even more operations. -> seems to be predominantly used for results, which makes it confusing for me here.

I have used "=" now.

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

457–461

If we want to support loops with different integer types (like index, i32, i64 etc), can we have a syntax like the following?

omp.wsloop (%i) = (%lb) to (%ub) step (%c)

or should it be like the following?

omp.wsloop (%i : i32) = (%lb : i32) to (%ub : i32) step (%c : i32)

And will the code given above for resolveOperands work by trying to parse with different types?

[kiranchandramohan]

Added multi-block tests.

[ftynse]

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

197

Nit: 80cols plz

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

378

Please add a comment with pseudo-BNF of what this parses.

437

s/omp.do/omp.wsloop

443–444

There's no difference between integer types and llvm integer types anymore

457–461

If the op only accepts operands of _identical_ types (which sounds sensible here), it suffices to list the type once, something like

omp.wsloop (%i) : (i32) = (%lb) to %(ub) step (%c) or omp.wsloop (%i : i32) = (%lb) to %(ub) step (%c).

The try-and-fail resolution won't work. The value will be null in case of type mismatch https://github.com/llvm/llvm-project/blob/f6524b4ada823a9766f7cbdfc16e052971e005f6/mlir/lib/Parser/Parser.cpp#L541. Even if it did, it would complain about every mismatching type. You need to parse a specific type instead and use it in resolution. I also suggest _not_ to check which type it is and delegate it to the verifier, just using AnyTypeOf<[Index, I32, I64]> for arguments in ODS combined with AllTypesMatch trait should suffice.

498

The code above supports parsing multiple lbs, ubs and steps, but this hardcodes 1 as their number.

626

Nit: we don't need an explicit number of stack values for SmallVector anymore, unless there is a strong reason to put 4 specifically

633–637

This likely won't work. In any case, at this point you should already know the actual types of region entry arguments - they are the same as lbs/ubs

649

Nit: we prefer early return:

if (vars.empty())
  return;
// code

651

Nit: don't call .size() on each iteration

651–654

llvm::interleaveComma is your friend here

674

Please fix the linter

[kiranchandramohan]

Address review comments.
Updated syntax to specify the type along with the index.
Using BNF for parsing function documentation
Minor fixes, using llvm::interleaveComma.

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

457–461

Thanks for your help here.

I have used the following style.

omp.wsloop (%i) : i32 = (%lb) to %(ub) step (%c)

651–654

Done.
For the linearVars case below i have not used since it refers to two containers linearVars, op.linear_step_vars().

[ftynse]

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

410

Nit: this reads like one can just omit the =, use ('=' ssa-id-and-type)? for a group

411

ultra-nit: delimit literals with backticks

509

Please drop llvm:: from Optional and StringRef, they are re-exported into the mlir namespace.

535

Nit: capitalize the first word and terminate sentence with a period.

728

Why do we need this trivial loop?

[kiranchandramohan]

Addressed review comments.

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

def TerminatorOp : OpenMP_Op<"terminator", [Terminator]> {
 
   let assemblyFormat = "attr-dict";
 }
 
 //===----------------------------------------------------------------------===//
 // 2.9.2 Workshare Loop Construct
 //===----------------------------------------------------------------------===//
 
-def WsLoopOp : OpenMP_Op<"wsloop", [AttrSizedOperandSegments]> {
+def WsLoopOp : OpenMP_Op<"wsloop", [AttrSizedOperandSegments,
+                         AllTypesMatch<["lowerBound", "upperBound", "step"]>]> {
   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 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>
-        %sum = addf %a, %b : f32
-        store %sum, %arrC[%i1, %i2] : memref<?x?xf32>
-        omp.yield
-      }
+    omp.wsloop (%i1, %i2) : index = (%c0, %c0) to (%c10, %c10) step (%c1, %c1) {
     ```
 
     `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.
let arguments = (ins Variadic<IntLikeType>:$lowerBound,
              OptionalAttr<ScheduleKind>:$schedule_val,
              Optional<AnyType>:$schedule_chunk_var,
              Confined<OptionalAttr<I64Attr>, [IntMinValue<0>]>:$collapse_val,
              UnitAttr:$nowait,
              Confined<OptionalAttr<I64Attr>, [IntMinValue<0>]>:$ordered_val,
              OptionalAttr<OrderKind>:$order_val,
              UnitAttr:$inclusive);
 
+  let skipDefaultBuilders = 1;
+
   let builders = [
     OpBuilder<(ins "ValueRange":$lowerBound, "ValueRange":$upperBound,
-                  "ValueRange":$step,
-                  CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes)>
+               CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes)>,
+    OpBuilder<(ins "TypeRange":$resultTypes, "ValueRange":$lowerBound,
+               "ValueRange":$upperBound, "ValueRange":$step,
+               "ValueRange":$privateVars, "ValueRange":$firstprivateVars,
+               "ValueRange":$lastprivate_vars, "ValueRange":$linear_vars,
+               "ValueRange":$linear_step_vars, "StringAttr":$schedule_val,
+               "Value":$schedule_chunk_var, "IntegerAttr":$collapse_val,
+               "UnitAttr":$nowait, "IntegerAttr":$ordered_val,
+               "StringAttr":$order_val, "UnitAttr":$inclusive, CArg<"bool",
+               "true">:$buildBody)>,
+    OpBuilder<(ins "TypeRange":$resultTypes, "ValueRange":$operands,
   ];
 
   let regions = (region AnyRegion:$region);
 
   let extraClassDeclaration = [{
     /// Returns the number of loops in the workshape loop nest.
     unsigned getNumLoops() { return lowerBound().size(); }
   }];
+  let parser = [{ return parseWsLoopOp(parser, result); }];
+  let printer = [{ return printWsLoopOp(p, *this); }];
 }
 
 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

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

 
 #include "mlir/Dialect/OpenMP/OpenMPDialect.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/OperationSupport.h"
 
 #include "llvm/ADT/SmallString.h"
+#include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/StringSwitch.h"
 #include <cstddef>
 
 #include "mlir/Dialect/OpenMP/OpenMPOpsEnums.cpp.inc"
 
 using namespace mlir;
 using namespace mlir::omp;
static void printParallelOp(OpAsmPrinter &p, ParallelOp op) {
 
   if (auto bind = op.proc_bind_val())
     p << " proc_bind(" << bind << ")";
 
   p.printRegion(op.getRegion());
 }
 
 /// Emit an error if the same clause is present more than once on an operation.
-static ParseResult allowedOnce(OpAsmParser &parser, llvm::StringRef clause,
-                               llvm::StringRef operation) {
+static ParseResult allowedOnce(OpAsmParser &parser, StringRef clause,
+                               StringRef operation) {
   return parser.emitError(parser.getNameLoc())
          << " at most one " << clause << " clause can appear on the "
          << operation << " operation";
 }
 
 /// Parses a parallel operation.
 ///
 /// operation ::= `omp.parallel` clause-list
static ParseResult parseParallelOp(OpAsmParser &parser,
   SmallVector<Type, 4> sharedTypes;
   SmallVector<OpAsmParser::OperandType, 4> copyins;
   SmallVector<Type, 4> copyinTypes;
   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;
+  StringRef keyword;
   bool defaultVal = false;
   bool procBind = false;
 
   const int ifClausePos = 0;
   const int numThreadsClausePos = 1;
   const int privateClausePos = 2;
   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();
+  const StringRef opName = result.name.getStringRef();
 
   while (succeeded(parser.parseOptionalKeyword(&keyword))) {
     if (keyword == "if") {
-      // Fail if there was already another if condition
+      // Fail if there was already another if condition.
       if (segments[ifClausePos])
         return allowedOnce(parser, "if", opName);
       if (parser.parseLParen() || parser.parseOperand(ifCond.first) ||
           parser.parseColonType(ifCond.second) || parser.parseRParen())
         return failure();
       segments[ifClausePos] = 1;
     } else if (keyword == "num_threads") {
-      // fail if there was already another num_threads clause
+      // Fail if there was already another num_threads clause.
       if (segments[numThreadsClausePos])
         return allowedOnce(parser, "num_threads", opName);
       if (parser.parseLParen() || parser.parseOperand(numThreads.first) ||
           parser.parseColonType(numThreads.second) || parser.parseRParen())
         return failure();
       segments[numThreadsClausePos] = 1;
     } else if (keyword == "private") {
-      // fail if there was already another private clause
+      // Fail if there was already another private clause.
       if (segments[privateClausePos])
         return allowedOnce(parser, "private", opName);
       if (parseOperandAndTypeList(parser, privates, privateTypes))
         return failure();
       segments[privateClausePos] = privates.size();
     } else if (keyword == "firstprivate") {
-      // fail if there was already another firstprivate clause
+      // Fail if there was already another firstprivate clause.
       if (segments[firstprivateClausePos])
         return allowedOnce(parser, "firstprivate", opName);
       if (parseOperandAndTypeList(parser, firstprivates, firstprivateTypes))
         return failure();
       segments[firstprivateClausePos] = firstprivates.size();
     } else if (keyword == "shared") {
-      // fail if there was already another shared clause
+      // Fail if there was already another shared clause.
       if (segments[sharedClausePos])
         return allowedOnce(parser, "shared", opName);
       if (parseOperandAndTypeList(parser, shareds, sharedTypes))
         return failure();
       segments[sharedClausePos] = shareds.size();
     } else if (keyword == "copyin") {
-      // fail if there was already another copyin clause
+      // Fail if there was already another copyin clause.
       if (segments[copyinClausePos])
         return allowedOnce(parser, "copyin", opName);
       if (parseOperandAndTypeList(parser, copyins, copyinTypes))
         return failure();
       segments[copyinClausePos] = copyins.size();
     } else if (keyword == "allocate") {
-      // fail if there was already another allocate clause
+      // 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
+      // Fail if there was already another default clause.
       if (defaultVal)
         return allowedOnce(parser, "default", opName);
       defaultVal = true;
-      llvm::StringRef defval;
+      StringRef defval;
       if (parser.parseLParen() || parser.parseKeyword(&defval) ||
           parser.parseRParen())
         return failure();
-      llvm::SmallString<16> attrval;
+      SmallString<16> attrval;
       // The def prefix is required for the attribute as "private" is a keyword
-      // in C++
+      // in C++.
       attrval += "def";
       attrval += defval;
       auto attr = parser.getBuilder().getStringAttr(attrval);
       result.addAttribute("default_val", attr);
     } else if (keyword == "proc_bind") {
-      // fail if there was already another proc_bind clause
+      // Fail if there was already another proc_bind clause.
       if (procBind)
         return allowedOnce(parser, "proc_bind", opName);
       procBind = true;
-      llvm::StringRef bind;
+      StringRef bind;
       if (parser.parseLParen() || parser.parseKeyword(&bind) ||
           parser.parseRParen())
         return failure();
       auto attr = parser.getBuilder().getStringAttr(bind);
       result.addAttribute("proc_bind_val", attr);
     } else {
       return parser.emitError(parser.getNameLoc())
              << keyword << " is not a valid clause for the " << opName
              << " operation";
     }
   }
 
-  // Add if parameter
+  // Add if parameter.
   if (segments[ifClausePos] &&
       parser.resolveOperand(ifCond.first, ifCond.second, result.operands))
     return failure();
 
-  // Add num_threads parameter
+  // Add num_threads parameter.
   if (segments[numThreadsClausePos] &&
       parser.resolveOperand(numThreads.first, numThreads.second,
                             result.operands))
     return failure();
 
-  // Add private parameters
+  // Add private parameters.
   if (segments[privateClausePos] &&
       parser.resolveOperands(privates, privateTypes, privates[0].location,
                              result.operands))
     return failure();
 
-  // Add firstprivate parameters
+  // Add firstprivate parameters.
   if (segments[firstprivateClausePos] &&
       parser.resolveOperands(firstprivates, firstprivateTypes,
                              firstprivates[0].location, result.operands))
     return failure();
 
-  // Add shared parameters
+  // Add shared parameters.
   if (segments[sharedClausePos] &&
       parser.resolveOperands(shareds, sharedTypes, shareds[0].location,
                              result.operands))
     return failure();
 
-  // Add copyin parameters
+  // Add copyin parameters.
   if (segments[copyinClausePos] &&
       parser.resolveOperands(copyins, copyinTypes, copyins[0].location,
                              result.operands))
     return failure();
 
-  // Add allocate parameters
+  // Add allocate parameters.
   if (segments[allocateClausePos] &&
       parser.resolveOperands(allocates, allocateTypes, allocates[0].location,
                              result.operands))
     return failure();
 
-  // Add allocator parameters
+  // 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));
 
   Region *body = result.addRegion();
   SmallVector<OpAsmParser::OperandType, 4> regionArgs;
   SmallVector<Type, 4> regionArgTypes;
   if (parser.parseRegion(*body, regionArgs, regionArgTypes))
     return failure();
   return success();
 }
 
+/// linear ::= `linear` `(` linear-list `)`
+/// linear-list := linear-val | linear-val linear-list
+/// linear-val := ssa-id-and-type `=` ssa-id-and-type
+static ParseResult
+parseLinearClause(OpAsmParser &parser,
+                  SmallVectorImpl<OpAsmParser::OperandType> &vars,
+                  SmallVectorImpl<Type> &types,
+                  SmallVectorImpl<OpAsmParser::OperandType> &stepVars) {
+  if (parser.parseLParen())
+    return failure();
+
+  do {
+    OpAsmParser::OperandType var;
+    Type type;
+    OpAsmParser::OperandType stepVar;
+    if (parser.parseOperand(var) || parser.parseEqual() ||
+        parser.parseOperand(stepVar) || parser.parseColonType(type))
+      return failure();
+
+    vars.push_back(var);
+    types.push_back(type);
+    stepVars.push_back(stepVar);
+  } while (succeeded(parser.parseOptionalComma()));
+
+  if (parser.parseRParen())
+    return failure();
+
+  return success();
+}
+
+/// schedule ::= `schedule` `(` sched-list `)`
+/// sched-list ::= sched-val | sched-val sched-list
+/// sched-val ::= sched-with-chunk | sched-wo-chunk
+/// sched-with-chunk ::= sched-with-chunk-types (`=` ssa-id-and-type)?
+/// sched-with-chunk-types ::= `static` | `dynamic` | `guided`
+/// sched-wo-chunk ::=  `auto` | `runtime`
+static ParseResult
+parseScheduleClause(OpAsmParser &parser, SmallString<8> &schedule,
+                    Optional<OpAsmParser::OperandType> &chunkSize) {
+  if (parser.parseLParen())
+    return failure();
+
+  StringRef keyword;
+  if (parser.parseKeyword(&keyword))
+    return failure();
+
+  schedule = keyword;
+  if (keyword == "static" || keyword == "dynamic" || keyword == "guided") {
+    if (succeeded(parser.parseOptionalEqual())) {
+      chunkSize = OpAsmParser::OperandType{};
+      if (parser.parseOperand(*chunkSize))
+        return failure();
+    } else {
+      chunkSize = llvm::NoneType::None;
+    }
+  } else if (keyword == "auto" || keyword == "runtime") {
+    chunkSize = llvm::NoneType::None;
+  } else {
+    return parser.emitError(parser.getNameLoc()) << " expected schedule kind";
+  }
+
+  if (parser.parseRParen())
+    return failure();
+
+  return success();
+}
+
+/// Parses an OpenMP Workshare Loop operation
+///
+/// operation ::= `omp.wsloop` loop-control clause-list
+/// loop-control ::= `(` ssa-id-list `)` `:` type `=`  loop-bounds
+/// loop-bounds := `(` ssa-id-list `)` to `(` ssa-id-list `)` steps
+/// steps := `step` `(`ssa-id-list`)`
+/// clause-list ::= clause | empty | clause-list
+/// clause ::= private | firstprivate | lastprivate | linear | schedule |
+//             collapse | nowait | ordered | order | inclusive
+/// private ::= `private` `(` ssa-id-and-type-list `)`
+/// firstprivate ::= `firstprivate` `(` ssa-id-and-type-list `)`
+/// lastprivate ::= `lastprivate` `(` ssa-id-and-type-list `)`
+/// linear ::= `linear` `(` linear-list `)`
+/// schedule ::= `schedule` `(` sched-list `)`
+/// collapse ::= `collapse` `(` ssa-id-and-type `)`
+/// nowait ::= `nowait`
+/// ordered ::= `ordered` `(` ssa-id-and-type `)`
+/// order ::= `order` `(` `concurrent` `)`
+/// inclusive ::= `inclusive`
+///
+static ParseResult parseWsLoopOp(OpAsmParser &parser, OperationState &result) {
+  Type loopVarType;
+  int numIVs;
+
+  // Parse an opening `(` followed by induction variables followed by `)`
+  SmallVector<OpAsmParser::OperandType> ivs;
+  if (parser.parseRegionArgumentList(ivs, /*requiredOperandCount=*/-1,
+                                     OpAsmParser::Delimiter::Paren))
+    return failure();
+
+  numIVs = static_cast<int>(ivs.size());
+
+  if (parser.parseColonType(loopVarType))
+    return failure();
+
+  // Parse loop bounds.
+  SmallVector<OpAsmParser::OperandType> lower;
+  if (parser.parseEqual() ||
+      parser.parseOperandList(lower, numIVs, OpAsmParser::Delimiter::Paren) ||
+      parser.resolveOperands(lower, loopVarType, result.operands))
+    return failure();
+
+  SmallVector<OpAsmParser::OperandType> upper;
+  if (parser.parseKeyword("to") ||
+      parser.parseOperandList(upper, numIVs, OpAsmParser::Delimiter::Paren) ||
+      parser.resolveOperands(upper, loopVarType, result.operands))
+    return failure();
+
+  // Parse step values.
+  SmallVector<OpAsmParser::OperandType> steps;
+  if (parser.parseKeyword("step") ||
+      parser.parseOperandList(steps, numIVs, OpAsmParser::Delimiter::Paren) ||
+      parser.resolveOperands(steps, loopVarType, result.operands))
+    return failure();
+
+  SmallVector<OpAsmParser::OperandType> privates;
+  SmallVector<Type> privateTypes;
+  SmallVector<OpAsmParser::OperandType> firstprivates;
+  SmallVector<Type> firstprivateTypes;
+  SmallVector<OpAsmParser::OperandType> lastprivates;
+  SmallVector<Type> lastprivateTypes;
+  SmallVector<OpAsmParser::OperandType> linears;
+  SmallVector<Type> linearTypes;
+  SmallVector<OpAsmParser::OperandType> linearSteps;
+  SmallString<8> schedule;
+  Optional<OpAsmParser::OperandType> scheduleChunkSize;
+  std::array<int, 9> segments{numIVs, numIVs, numIVs, 0, 0, 0, 0, 0, 0};
+
+  const StringRef opName = result.name.getStringRef();
+  StringRef keyword;
+
+  enum SegmentPos {
+    lbPos = 0,
+    ubPos,
+    stepPos,
+    privateClausePos,
+    firstprivateClausePos,
+    lastprivateClausePos,
+    linearClausePos,
+    linearStepPos,
+    scheduleClausePos,
+  };
+
+  while (succeeded(parser.parseOptionalKeyword(&keyword))) {
+    if (keyword == "private") {
+      if (segments[privateClausePos])
+        return allowedOnce(parser, "private", opName);
+      if (parseOperandAndTypeList(parser, privates, privateTypes))
+        return failure();
+      segments[privateClausePos] = privates.size();
+    } else if (keyword == "firstprivate") {
+      // fail if there was already another firstprivate clause
+      if (segments[firstprivateClausePos])
+        return allowedOnce(parser, "firstprivate", opName);
+      if (parseOperandAndTypeList(parser, firstprivates, firstprivateTypes))
+        return failure();
+      segments[firstprivateClausePos] = firstprivates.size();
+    } else if (keyword == "lastprivate") {
+      // fail if there was already another shared clause
+      if (segments[lastprivateClausePos])
+        return allowedOnce(parser, "lastprivate", opName);
+      if (parseOperandAndTypeList(parser, lastprivates, lastprivateTypes))
+        return failure();
+      segments[lastprivateClausePos] = lastprivates.size();
+    } else if (keyword == "linear") {
+      // fail if there was already another linear clause
+      if (segments[linearClausePos])
+        return allowedOnce(parser, "linear", opName);
+      if (parseLinearClause(parser, linears, linearTypes, linearSteps))
+        return failure();
+      segments[linearClausePos] = linears.size();
+      segments[linearStepPos] = linearSteps.size();
+    } else if (keyword == "schedule") {
+      if (!schedule.empty())
+        return allowedOnce(parser, "schedule", opName);
+      if (parseScheduleClause(parser, schedule, scheduleChunkSize))
+        return failure();
+      if (scheduleChunkSize) {
+        segments[scheduleClausePos] = 1;
+      }
+    } else if (keyword == "collapse") {
+      auto type = parser.getBuilder().getI64Type();
+      mlir::IntegerAttr attr;
+      if (parser.parseLParen() || parser.parseAttribute(attr, type) ||
+          parser.parseRParen())
+        return failure();
+      result.addAttribute("collapse_val", attr);
+    } else if (keyword == "nowait") {
+      auto attr = UnitAttr::get(parser.getBuilder().getContext());
+      result.addAttribute("nowait", attr);
+    } else if (keyword == "ordered") {
+      mlir::IntegerAttr attr;
+      if (succeeded(parser.parseOptionalLParen())) {
+        auto type = parser.getBuilder().getI64Type();
+        if (parser.parseAttribute(attr, type))
+          return failure();
+        if (parser.parseRParen())
+          return failure();
+      } else {
+        // Use 0 to represent no ordered parameter was specified
+        attr = parser.getBuilder().getI64IntegerAttr(0);
+      }
+      result.addAttribute("ordered_val", attr);
+    } else if (keyword == "order") {
+      StringRef order;
+      if (parser.parseLParen() || parser.parseKeyword(&order) ||
+          parser.parseRParen())
+        return failure();
+      auto attr = parser.getBuilder().getStringAttr(order);
+      result.addAttribute("order", attr);
+    } else if (keyword == "inclusive") {
+      auto attr = UnitAttr::get(parser.getBuilder().getContext());
+      result.addAttribute("inclusive", attr);
+    }
+  }
+
+  if (segments[privateClausePos]) {
+    parser.resolveOperands(privates, privateTypes, privates[0].location,
+                           result.operands);
+  }
+
+  if (segments[firstprivateClausePos]) {
+    parser.resolveOperands(firstprivates, firstprivateTypes,
+                           firstprivates[0].location, result.operands);
+  }
+
+  if (segments[lastprivateClausePos]) {
+    parser.resolveOperands(lastprivates, lastprivateTypes,
+                           lastprivates[0].location, result.operands);
+  }
+
+  if (segments[linearClausePos]) {
+    parser.resolveOperands(linears, linearTypes, linears[0].location,
+                           result.operands);
+    auto linearStepType = parser.getBuilder().getI32Type();
+    SmallVector<Type> linearStepTypes(linearSteps.size(), linearStepType);
+    parser.resolveOperands(linearSteps, linearStepTypes,
+                           linearSteps[0].location, result.operands);
+  }
+
+  if (!schedule.empty()) {
+    schedule[0] = llvm::toUpper(schedule[0]);
+    auto attr = parser.getBuilder().getStringAttr(schedule);
+    result.addAttribute("schedule_val", attr);
+    if (scheduleChunkSize) {
+      auto chunkSizeType = parser.getBuilder().getI32Type();
+      parser.resolveOperand(*scheduleChunkSize, chunkSizeType, result.operands);
+    }
+  }
+
+  result.addAttribute("operand_segment_sizes",
+                      parser.getBuilder().getI32VectorAttr(segments));
+
+  // Now parse the body.
+  Region *body = result.addRegion();
+  SmallVector<Type> ivTypes(numIVs, loopVarType);
+  if (parser.parseRegion(*body, ivs, ivTypes))
+    return failure();
+  return success();
+}
+
+static void printWsLoopOp(OpAsmPrinter &p, WsLoopOp op) {
+  auto args = op.getRegion().front().getArguments();
+  p << op.getOperationName() << " (" << args << ") : " << args[0].getType()
+    << " = (" << op.lowerBound() << ") to (" << op.upperBound() << ") step ("
+    << op.step() << ")";
+
+  // Print private, firstprivate, shared and copyin parameters
+  auto printDataVars = [&p](StringRef name, OperandRange vars) {
+    if (vars.empty())
+      return;
+
+    p << " " << name << "(";
+    llvm::interleaveComma(
+        vars, p, [&](const Value &v) { p << v << " : " << v.getType(); });
+    p << ")";
+  };
+  printDataVars("private", op.private_vars());
+  printDataVars("firstprivate", op.firstprivate_vars());
+  printDataVars("lastprivate", op.lastprivate_vars());
+
+  auto linearVars = op.linear_vars();
+  auto linearVarsSize = linearVars.size();
+  if (linearVarsSize) {
+    p << " "
+      << "linear"
+      << "(";
+    for (unsigned i = 0; i < linearVarsSize; ++i) {
+      std::string separator = i == linearVarsSize - 1 ? ")" : ", ";
+      p << linearVars[i];
+      if (op.linear_step_vars().size() > i)
+        p << " = " << op.linear_step_vars()[i];
+      p << " : " << linearVars[i].getType() << separator;
+    }
+  }
+
+  if (auto sched = op.schedule_val()) {
+    auto schedLower = sched->lower();
+    p << " schedule(" << schedLower;
+    if (auto chunk = op.schedule_chunk_var()) {
+      p << " = " << chunk;
+    }
+    p << ")";
+  }
+
+  if (auto collapse = op.collapse_val())
+    p << " collapse(" << collapse << ")";
+
+  if (op.nowait())
+    p << " nowait";
+
+  if (auto ordered = op.ordered_val()) {
+    p << " ordered(" << ordered << ")";
+  }
+
+  if (op.inclusive()) {
+    p << " inclusive";
+  }
+
+  p.printRegion(op.region(), /*printEntryBlockArgs=*/false);
+}
+
 //===----------------------------------------------------------------------===//
 // WsLoopOp
 //===----------------------------------------------------------------------===//
 
 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=*/false, /*ordered_val=*/nullptr, /*order_val=*/nullptr,
-        /*inclusive=*/false);
+        /*nowait=*/nullptr, /*ordered_val=*/nullptr, /*order_val=*/nullptr,
+        /*inclusive=*/nullptr, /*buildBody=*/false);
   state.addAttributes(attributes);
 }
 
+void WsLoopOp::build(OpBuilder &, OperationState &state, TypeRange resultTypes,
+                     ValueRange operands, ArrayRef<NamedAttribute> attributes) {
+  state.addOperands(operands);
+  state.addAttributes(attributes);
+  (void)state.addRegion();
+  assert(resultTypes.size() == 0u && "mismatched number of return types");
+  state.addTypes(resultTypes);
+}
+
+void WsLoopOp::build(OpBuilder &builder, OperationState &result,
+                     TypeRange typeRange, ValueRange lowerBounds,
+                     ValueRange upperBounds, ValueRange steps,
+                     ValueRange privateVars, ValueRange firstprivateVars,
+                     ValueRange lastprivateVars, ValueRange linearVars,
+                     ValueRange linearStepVars, StringAttr scheduleVal,
+                     Value scheduleChunkVar, IntegerAttr collapseVal,
+                     UnitAttr nowait, IntegerAttr orderedVal,
+                     StringAttr orderVal, UnitAttr inclusive, bool buildBody) {
+  result.addOperands(lowerBounds);
+  result.addOperands(upperBounds);
+  result.addOperands(steps);
+  result.addOperands(privateVars);
+  result.addOperands(firstprivateVars);
+  result.addOperands(linearVars);
+  result.addOperands(linearStepVars);
+  if (scheduleChunkVar)
+    result.addOperands(scheduleChunkVar);
+
+  if (scheduleVal)
+    result.addAttribute("schedule_val", scheduleVal);
+  if (collapseVal)
+    result.addAttribute("collapse_val", collapseVal);
+  if (nowait)
+    result.addAttribute("nowait", nowait);
+  if (orderedVal)
+    result.addAttribute("ordered_val", orderedVal);
+  if (orderVal)
+    result.addAttribute("order", orderVal);
+  if (inclusive)
+    result.addAttribute("inclusive", inclusive);
+  result.addAttribute(
+      WsLoopOp::getOperandSegmentSizeAttr(),
+      builder.getI32VectorAttr(
+          {static_cast<int32_t>(lowerBounds.size()),
+           static_cast<int32_t>(upperBounds.size()),
+           static_cast<int32_t>(steps.size()),
+           static_cast<int32_t>(privateVars.size()),
+           static_cast<int32_t>(firstprivateVars.size()),
+           static_cast<int32_t>(lastprivateVars.size()),
+           static_cast<int32_t>(linearVars.size()),
+           static_cast<int32_t>(linearStepVars.size()),
+           static_cast<int32_t>(scheduleChunkVar != nullptr ? 1 : 0)}));
+
+  Region *bodyRegion = result.addRegion();
+  if (buildBody) {
+    OpBuilder::InsertionGuard guard(builder);
+    unsigned numIVs = steps.size();
+    SmallVector<Type, 8> argTypes(numIVs, steps.getType().front());
+    builder.createBlock(bodyRegion, {}, argTypes);
+  }
+}
+
 #define GET_OP_CLASSES
 #include "mlir/Dialect/OpenMP/OpenMPOps.cpp.inc"

mlir/test/Conversion/OpenMPToLLVM/convert-to-llvmir.mlir

func @branch_loop() {
   return
 }
 
 // CHECK-LABEL: @wsloop
 // CHECK: (%[[ARG0:.*]]: i64, %[[ARG1:.*]]: i64, %[[ARG2:.*]]: i64, %[[ARG3:.*]]: i64, %[[ARG4:.*]]: i64, %[[ARG5:.*]]: i64)
 func @wsloop(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index) {
   // CHECK: omp.parallel
   omp.parallel {
-    // CHECK: omp.wsloop
-    // CHECK: (%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]], %[[ARG5]])
+    // CHECK: omp.wsloop (%[[ARG6:.*]], %[[ARG7:.*]]) : i64 = (%[[ARG0]], %[[ARG1]]) to (%[[ARG2]], %[[ARG3]]) step (%[[ARG4]], %[[ARG5]]) {
     "omp.wsloop"(%arg0, %arg1, %arg2, %arg3, %arg4, %arg5) ( {
-    // CHECK: ^{{.*}}(%[[ARG6:.*]]: i64, %[[ARG7:.*]]: i64):
     ^bb0(%arg6: index, %arg7: index):  // no predecessors
       // CHECK: "test.payload"(%[[ARG6]], %[[ARG7]]) : (i64, i64) -> ()
       "test.payload"(%arg6, %arg7) : (index, index) -> ()
       omp.yield
     }) {operand_segment_sizes = dense<[2, 2, 2, 0, 0, 0, 0, 0, 0]> : vector<9xi32>} : (index, index, index, index, index, index) -> ()
     omp.terminator
   }
   return
 }

mlir/test/Conversion/SCFToOpenMP/scf-to-openmp.mlir

 // RUN: mlir-opt -convert-scf-to-openmp %s | FileCheck %s
 
 // CHECK-LABEL: @parallel
 func @parallel(%arg0: index, %arg1: index, %arg2: index,
           %arg3: index, %arg4: index, %arg5: index) {
   // CHECK: omp.parallel {
-  // CHECK:  "omp.wsloop"({{.*}}) ( {
+  // CHECK: omp.wsloop (%[[LVAR1:.*]], %[[LVAR2:.*]]) : index = (%arg0, %arg1) to (%arg2, %arg3) step (%arg4, %arg5) {
   scf.parallel (%i, %j) = (%arg0, %arg1) to (%arg2, %arg3) step (%arg4, %arg5) {
-    // CHECK:   test.payload
+    // CHECK: "test.payload"(%[[LVAR1]], %[[LVAR2]]) : (index, index) -> ()
     "test.payload"(%i, %j) : (index, index) -> ()
     // CHECK:   omp.yield
     // CHECK: }
   }
   // CHECK:   omp.terminator
   // CHECK: }
   return
 }
 
 // CHECK-LABEL: @nested_loops
 func @nested_loops(%arg0: index, %arg1: index, %arg2: index,
                    %arg3: index, %arg4: index, %arg5: index) {
   // CHECK: omp.parallel {
-  // CHECK:  "omp.wsloop"({{.*}}) ( {
+  // CHECK: omp.wsloop (%[[LVAR_OUT1:.*]]) : index = (%arg0) to (%arg2) step (%arg4) {
   // CHECK-NOT: omp.parallel
   scf.parallel (%i) = (%arg0) to (%arg2) step (%arg4) {
-    // CHECK:   "omp.wsloop"({{.*}}) ( {
+    // CHECK: omp.wsloop (%[[LVAR_IN1:.*]]) : index = (%arg1) to (%arg3) step (%arg5) {
     scf.parallel (%j) = (%arg1) to (%arg3) step (%arg5) {
-      // CHECK:   test.payload
+      // CHECK: "test.payload"(%[[LVAR_OUT1]], %[[LVAR_IN1]]) : (index, index) -> ()
       "test.payload"(%i, %j) : (index, index) -> ()
       // CHECK:   omp.yield
       // CHECK: }
     }
     // CHECK:   omp.yield
     // CHECK: }
   }
   // CHECK:   omp.terminator
   // CHECK: }
   return
 }
 
 func @adjacent_loops(%arg0: index, %arg1: index, %arg2: index,
                      %arg3: index, %arg4: index, %arg5: index) {
   // CHECK: omp.parallel {
-  // CHECK:  "omp.wsloop"({{.*}}) ( {
+  // CHECK: omp.wsloop (%[[LVAR_AL1:.*]]) : index = (%arg0) to (%arg2) step (%arg4) {
   scf.parallel (%i) = (%arg0) to (%arg2) step (%arg4) {
-    // CHECK:   test.payload1
+    // CHECK: "test.payload1"(%[[LVAR_AL1]]) : (index) -> ()
     "test.payload1"(%i) : (index) -> ()
     // CHECK:   omp.yield
     // CHECK: }
   }
   // CHECK:   omp.terminator
   // CHECK: }
 
   // CHECK: omp.parallel {
-  // CHECK:  "omp.wsloop"({{.*}}) ( {
+  // CHECK: omp.wsloop (%[[LVAR_AL2:.*]]) : index = (%arg1) to (%arg3) step (%arg5) {
   scf.parallel (%j) = (%arg1) to (%arg3) step (%arg5) {
-    // CHECK:   test.payload2
+    // CHECK: "test.payload2"(%[[LVAR_AL2]]) : (index) -> ()
     "test.payload2"(%j) : (index) -> ()
     // CHECK:   omp.yield
     // CHECK: }
   }
   // CHECK:   omp.terminator
   // CHECK: }
   return
 }

mlir/test/Dialect/OpenMP/ops.mlir

func @omp_parallel(%data_var : memref<i32>, %if_cond : i1, %num_threads : si32) -> () {
   "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,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, %allocator : si32) -> () {
-  // CHECK: omp.parallel
-  omp.parallel {
-    omp.terminator
-  }
 
-  // CHECK: omp.parallel num_threads(%{{.*}} : si32)
-  omp.parallel num_threads(%num_threads : si32) {
-    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
-  }
 
-  // CHECK omp.parallel shared(%{{.*}} : memref<i32>) copyin(%{{.*}} : memref<i32>, %{{.*}} : memref<i32>)
-  omp.parallel shared(%data_var : memref<i32>) copyin(%data_var : memref<i32>, %data_var : memref<i32>) {
-    omp.parallel if(%if_cond: i1) {
-      omp.terminator
-    }
-    omp.terminator
-  }
 
-  // CHECK omp.parallel if(%{{.*}}) num_threads(%{{.*}} : si32) private(%{{.*}} : memref<i32>) proc_bind(close)
-  omp.parallel num_threads(%num_threads : si32) if(%if_cond: i1)
-               private(%data_var : memref<i32>) proc_bind(close) {
-    omp.terminator
-  }
 
-  return
 }
 
-func @omp_wsloop(%lb : index, %ub : index, %step : index,
-                 %data_var : memref<i32>, %linear_var : si32, %chunk_var : si32) -> () {
+// CHECK-LABEL: omp_wsloop
+func @omp_wsloop(%lb : index, %ub : index, %step : index, %data_var : memref<i32>, %linear_var : i32, %chunk_var : i32) -> () {
 
-  // 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) ({
+  // CHECK: omp.wsloop (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) private(%{{.*}} : memref<i32>, %{{.*}} : memref<i32>) collapse(2) ordered(1)
+  "omp.wsloop" (%lb, %ub, %step, %data_var, %data_var) ({
+    ^bb0(%iv: index):
       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>) -> ()
+  }) {operand_segment_sizes = dense<[1,1,1,2,0,0,0,0,0]> : vector<9xi32>, collapse_val = 2, ordered_val = 1} :
+    (index, index, index, memref<i32>, memref<i32>) -> ()
 
-
-  // CHECK: "omp.wsloop"(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}})
+  // CHECK: omp.wsloop (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) linear(%{{.*}} = %{{.*}} : memref<i32>) schedule(static)
   "omp.wsloop" (%lb, %ub, %step, %data_var, %linear_var) ({
+    ^bb0(%iv: index):
       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) -> ()
+    (index, index, index, memref<i32>, i32) -> ()
+
+  // CHECK: omp.wsloop (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) linear(%{{.*}} = %{{.*}} : memref<i32>, %{{.*}} = %{{.*}} : memref<i32>) schedule(static)
+  "omp.wsloop" (%lb, %ub, %step, %data_var, %data_var, %linear_var, %linear_var) ({
+    ^bb0(%iv: index):
+      omp.yield
+  }) {operand_segment_sizes = dense<[1,1,1,0,0,0,2,2,0]> : vector<9xi32>, schedule_val = "Static"} :
+    (index, index, index, memref<i32>, memref<i32>, i32, i32) -> ()
 
-  // CHECK: "omp.wsloop"(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}})
+  // CHECK: omp.wsloop (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) private(%{{.*}} : memref<i32>) firstprivate(%{{.*}} : memref<i32>) lastprivate(%{{.*}} : memref<i32>) linear(%{{.*}} = %{{.*}} : memref<i32>) schedule(dynamic = %{{.*}}) collapse(3) ordered(2)
   "omp.wsloop" (%lb, %ub, %step, %data_var, %data_var, %data_var, %data_var, %linear_var, %chunk_var) ({
+    ^bb0(%iv: index):
       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) -> ()
+    (index, index, index, memref<i32>, memref<i32>, memref<i32>, memref<i32>, i32, i32) -> ()
 
-  // CHECK: "omp.wsloop"(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}})
+  // CHECK: omp.wsloop (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) private(%{{.*}} : memref<i32>) schedule(auto) nowait
   "omp.wsloop" (%lb, %ub, %step, %data_var) ({
+    ^bb0(%iv: index):
       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
+}
+
+// CHECK-LABEL: omp_wsloop_pretty
+func @omp_wsloop_pretty(%lb : index, %ub : index, %step : index,
+                 %data_var : memref<i32>, %linear_var : i32, %chunk_var : i32) -> () {
+
+  // CHECK: omp.wsloop (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) private(%{{.*}} : memref<i32>)
+  omp.wsloop (%iv) : index = (%lb) to (%ub) step (%step) private(%data_var : memref<i32>) collapse(2) ordered(2) {
+    omp.yield
+  }
+
+  // CHECK: omp.wsloop (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) linear(%{{.*}} = %{{.*}} : memref<i32>) schedule(static)
+  omp.wsloop (%iv) : index = (%lb) to (%ub) step (%step) schedule(static) lastprivate(%data_var : memref<i32>) linear(%data_var = %linear_var : memref<i32>) {
+    omp.yield
+  }
+
+  // CHECK: omp.wsloop (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) private(%{{.*}} : memref<i32>) firstprivate(%{{.*}} : memref<i32>) lastprivate(%{{.*}} : memref<i32>) linear(%{{.*}} = %{{.*}} : memref<i32>) schedule(static = %{{.*}}) collapse(3) ordered(2)
+  omp.wsloop (%iv) : index = (%lb) to (%ub) step (%step) ordered(2) private(%data_var : memref<i32>)
+     firstprivate(%data_var : memref<i32>) lastprivate(%data_var : memref<i32>) linear(%data_var = %linear_var : memref<i32>)
+     schedule(static = %chunk_var) collapse(3) {
+    omp.yield
+  }
+
+  // CHECK: omp.wsloop (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) private({{.*}} : memref<i32>)
+  omp.wsloop (%iv) : index = (%lb) to (%ub) step (%step) private(%data_var : memref<i32>) {
+    omp.yield
+  }
+
+  return
+}
+
+// CHECK-LABEL: omp_wsloop_pretty_multi_block
+func @omp_wsloop_pretty_multi_block(%lb : index, %ub : index, %step : index, %data1 : memref<?xi32>, %data2 : memref<?xi32>) -> () {
+
+  // CHECK: omp.wsloop (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}})
+  omp.wsloop (%iv) : index = (%lb) to (%ub) step (%step) {
+    %1 = "test.payload"(%iv) : (index) -> (i32)
+    br ^bb1(%1: i32)
+  ^bb1(%arg: i32):
+    memref.store %arg, %data1[%iv] : memref<?xi32>
+    omp.yield
+  }
+
+  // CHECK: omp.wsloop (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}})
+  omp.wsloop (%iv) : index = (%lb) to (%ub) step (%step) {
+    %c = "test.condition"(%iv) : (index) -> (i1)
+    %v1 = "test.payload"(%iv) : (index) -> (i32)
+    cond_br %c, ^bb1(%v1: i32), ^bb2(%v1: i32)
+  ^bb1(%arg0: i32):
+    memref.store %arg0, %data1[%iv] : memref<?xi32>
+    br ^bb3
+  ^bb2(%arg1: i32):
+    memref.store %arg1, %data2[%iv] : memref<?xi32>
+    br ^bb3
+  ^bb3:
+    omp.yield
+  }
+
+  // CHECK: omp.wsloop (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}})
+  omp.wsloop (%iv) : index = (%lb) to (%ub) step (%step) {
+    %c = "test.condition"(%iv) : (index) -> (i1)
+    %v1 = "test.payload"(%iv) : (index) -> (i32)
+    cond_br %c, ^bb1(%v1: i32), ^bb2(%v1: i32)
+  ^bb1(%arg0: i32):
+    memref.store %arg0, %data1[%iv] : memref<?xi32>
+    omp.yield
+  ^bb2(%arg1: i32):
+    memref.store %arg1, %data2[%iv] : memref<?xi32>
+    omp.yield
+  }
+
+  return
+}
+
+// CHECK-LABEL: omp_wsloop_pretty_non_index
+func @omp_wsloop_pretty_non_index(%lb1 : i32, %ub1 : i32, %step1 : i32, %lb2 : i64, %ub2 : i64, %step2 : i64,
+                           %data1 : memref<?xi32>, %data2 : memref<?xi64>) -> () {
+
+  // CHECK: omp.wsloop (%{{.*}}) : i32 = (%{{.*}}) to (%{{.*}}) step (%{{.*}})
+  omp.wsloop (%iv1) : i32 = (%lb1) to (%ub1) step (%step1) {
+    %1 = "test.payload"(%iv1) : (i32) -> (index)
+    br ^bb1(%1: index)
+  ^bb1(%arg1: index):
+    memref.store %iv1, %data1[%arg1] : memref<?xi32>
+    omp.yield
+  }
+
+  // CHECK: omp.wsloop (%{{.*}}) : i64 = (%{{.*}}) to (%{{.*}}) step (%{{.*}})
+  omp.wsloop (%iv2) : i64 = (%lb2) to (%ub2) step (%step2) {
+    %2 = "test.payload"(%iv2) : (i64) -> (index)
+    br ^bb1(%2: index)
+  ^bb1(%arg2: index):
+    memref.store %iv2, %data2[%arg2] : memref<?xi64>
+    omp.yield
+  }
+
+  return
+}
+
+// CHECK-LABEL: omp_wsloop_pretty_multiple
+func @omp_wsloop_pretty_multiple(%lb1 : i32, %ub1 : i32, %step1 : i32, %lb2 : i32, %ub2 : i32, %step2 : i32, %data1 : memref<?xi32>) -> () {
+
+  // CHECK: omp.wsloop (%{{.*}}, %{{.*}}) : i32 = (%{{.*}}, %{{.*}}) to (%{{.*}}, %{{.*}}) step (%{{.*}}, %{{.*}})
+  omp.wsloop (%iv1, %iv2) : i32 = (%lb1, %lb2) to (%ub1, %ub2) step (%step1, %step2) {
+    %1 = "test.payload"(%iv1) : (i32) -> (index)
+    %2 = "test.payload"(%iv2) : (i32) -> (index)
+    memref.store %iv1, %data1[%1] : memref<?xi32>
+    memref.store %iv2, %data1[%2] : memref<?xi32>
+    omp.yield
+  }
 
   return
 }