1 //===- OpenMPDialect.cpp - MLIR Dialect for OpenMP implementation ---------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 //
9 // This file implements the OpenMP dialect and its operations.
10 //
11 //===----------------------------------------------------------------------===//
12 
13 #include "mlir/Dialect/OpenMP/OpenMPDialect.h"
14 #include "mlir/Dialect/StandardOps/IR/Ops.h"
15 #include "mlir/IR/Attributes.h"
16 #include "mlir/IR/OpImplementation.h"
17 #include "mlir/IR/OperationSupport.h"
18 
19 #include "llvm/ADT/SmallString.h"
20 #include "llvm/ADT/StringExtras.h"
21 #include "llvm/ADT/StringRef.h"
22 #include "llvm/ADT/StringSwitch.h"
23 #include <cstddef>
24 
25 #include "mlir/Dialect/OpenMP/OpenMPOpsDialect.cpp.inc"
26 #include "mlir/Dialect/OpenMP/OpenMPOpsEnums.cpp.inc"
27 
28 using namespace mlir;
29 using namespace mlir::omp;
30 
31 void OpenMPDialect::initialize() {
32   addOperations<
33 #define GET_OP_LIST
34 #include "mlir/Dialect/OpenMP/OpenMPOps.cpp.inc"
35       >();
36 }
37 
38 //===----------------------------------------------------------------------===//
39 // ParallelOp
40 //===----------------------------------------------------------------------===//
41 
42 void ParallelOp::build(OpBuilder &builder, OperationState &state,
43                        ArrayRef<NamedAttribute> attributes) {
44   ParallelOp::build(
45       builder, state, /*if_expr_var=*/nullptr, /*num_threads_var=*/nullptr,
46       /*default_val=*/nullptr, /*private_vars=*/ValueRange(),
47       /*firstprivate_vars=*/ValueRange(), /*shared_vars=*/ValueRange(),
48       /*copyin_vars=*/ValueRange(), /*allocate_vars=*/ValueRange(),
49       /*allocators_vars=*/ValueRange(), /*proc_bind_val=*/nullptr);
50   state.addAttributes(attributes);
51 }
52 
53 /// Parse a list of operands with types.
54 ///
55 /// operand-and-type-list ::= `(` ssa-id-and-type-list `)`
56 /// ssa-id-and-type-list ::= ssa-id-and-type |
57 ///                          ssa-id-and-type `,` ssa-id-and-type-list
58 /// ssa-id-and-type ::= ssa-id `:` type
59 static ParseResult
60 parseOperandAndTypeList(OpAsmParser &parser,
61                         SmallVectorImpl<OpAsmParser::OperandType> &operands,
62                         SmallVectorImpl<Type> &types) {
63   if (parser.parseLParen())
64     return failure();
65 
66   do {
67     OpAsmParser::OperandType operand;
68     Type type;
69     if (parser.parseOperand(operand) || parser.parseColonType(type))
70       return failure();
71     operands.push_back(operand);
72     types.push_back(type);
73   } while (succeeded(parser.parseOptionalComma()));
74 
75   if (parser.parseRParen())
76     return failure();
77 
78   return success();
79 }
80 
81 /// Parse an allocate clause with allocators and a list of operands with types.
82 ///
83 /// operand-and-type-list ::= `(` allocate-operand-list `)`
84 /// allocate-operand-list :: = allocate-operand |
85 ///                            allocator-operand `,` allocate-operand-list
86 /// allocate-operand :: = ssa-id-and-type -> ssa-id-and-type
87 /// ssa-id-and-type ::= ssa-id `:` type
88 static ParseResult parseAllocateAndAllocator(
89     OpAsmParser &parser,
90     SmallVectorImpl<OpAsmParser::OperandType> &operandsAllocate,
91     SmallVectorImpl<Type> &typesAllocate,
92     SmallVectorImpl<OpAsmParser::OperandType> &operandsAllocator,
93     SmallVectorImpl<Type> &typesAllocator) {
94   if (parser.parseLParen())
95     return failure();
96 
97   do {
98     OpAsmParser::OperandType operand;
99     Type type;
100 
101     if (parser.parseOperand(operand) || parser.parseColonType(type))
102       return failure();
103     operandsAllocator.push_back(operand);
104     typesAllocator.push_back(type);
105     if (parser.parseArrow())
106       return failure();
107     if (parser.parseOperand(operand) || parser.parseColonType(type))
108       return failure();
109 
110     operandsAllocate.push_back(operand);
111     typesAllocate.push_back(type);
112   } while (succeeded(parser.parseOptionalComma()));
113 
114   if (parser.parseRParen())
115     return failure();
116 
117   return success();
118 }
119 
120 static LogicalResult verifyParallelOp(ParallelOp op) {
121   if (op.allocate_vars().size() != op.allocators_vars().size())
122     return op.emitError(
123         "expected equal sizes for allocate and allocator variables");
124   return success();
125 }
126 
127 static void printParallelOp(OpAsmPrinter &p, ParallelOp op) {
128   p << "omp.parallel";
129 
130   if (auto ifCond = op.if_expr_var())
131     p << " if(" << ifCond << " : " << ifCond.getType() << ")";
132 
133   if (auto threads = op.num_threads_var())
134     p << " num_threads(" << threads << " : " << threads.getType() << ")";
135 
136   // Print private, firstprivate, shared and copyin parameters
137   auto printDataVars = [&p](StringRef name, OperandRange vars) {
138     if (vars.size()) {
139       p << " " << name << "(";
140       for (unsigned i = 0; i < vars.size(); ++i) {
141         std::string separator = i == vars.size() - 1 ? ")" : ", ";
142         p << vars[i] << " : " << vars[i].getType() << separator;
143       }
144     }
145   };
146 
147   // Print allocator and allocate parameters
148   auto printAllocateAndAllocator = [&p](OperandRange varsAllocate,
149                                         OperandRange varsAllocator) {
150     if (varsAllocate.empty())
151       return;
152 
153     p << " allocate(";
154     for (unsigned i = 0; i < varsAllocate.size(); ++i) {
155       std::string separator = i == varsAllocate.size() - 1 ? ")" : ", ";
156       p << varsAllocator[i] << " : " << varsAllocator[i].getType() << " -> ";
157       p << varsAllocate[i] << " : " << varsAllocate[i].getType() << separator;
158     }
159   };
160 
161   printDataVars("private", op.private_vars());
162   printDataVars("firstprivate", op.firstprivate_vars());
163   printDataVars("shared", op.shared_vars());
164   printDataVars("copyin", op.copyin_vars());
165   printAllocateAndAllocator(op.allocate_vars(), op.allocators_vars());
166 
167   if (auto def = op.default_val())
168     p << " default(" << def->drop_front(3) << ")";
169 
170   if (auto bind = op.proc_bind_val())
171     p << " proc_bind(" << bind << ")";
172 
173   p.printRegion(op.getRegion());
174 }
175 
176 /// Emit an error if the same clause is present more than once on an operation.
177 static ParseResult allowedOnce(OpAsmParser &parser, StringRef clause,
178                                StringRef operation) {
179   return parser.emitError(parser.getNameLoc())
180          << " at most one " << clause << " clause can appear on the "
181          << operation << " operation";
182 }
183 
184 /// Parses a parallel operation.
185 ///
186 /// operation ::= `omp.parallel` clause-list
187 /// clause-list ::= clause | clause clause-list
188 /// clause ::= if | numThreads | private | firstprivate | shared | copyin |
189 ///            default | procBind
190 /// if ::= `if` `(` ssa-id `)`
191 /// numThreads ::= `num_threads` `(` ssa-id-and-type `)`
192 /// private ::= `private` operand-and-type-list
193 /// firstprivate ::= `firstprivate` operand-and-type-list
194 /// shared ::= `shared` operand-and-type-list
195 /// copyin ::= `copyin` operand-and-type-list
196 /// allocate ::= `allocate` operand-and-type `->` operand-and-type-list
197 /// default ::= `default` `(` (`private` | `firstprivate` | `shared` | `none`)
198 /// procBind ::= `proc_bind` `(` (`master` | `close` | `spread`) `)`
199 ///
200 /// Note that each clause can only appear once in the clase-list.
201 static ParseResult parseParallelOp(OpAsmParser &parser,
202                                    OperationState &result) {
203   std::pair<OpAsmParser::OperandType, Type> ifCond;
204   std::pair<OpAsmParser::OperandType, Type> numThreads;
205   SmallVector<OpAsmParser::OperandType, 4> privates;
206   SmallVector<Type, 4> privateTypes;
207   SmallVector<OpAsmParser::OperandType, 4> firstprivates;
208   SmallVector<Type, 4> firstprivateTypes;
209   SmallVector<OpAsmParser::OperandType, 4> shareds;
210   SmallVector<Type, 4> sharedTypes;
211   SmallVector<OpAsmParser::OperandType, 4> copyins;
212   SmallVector<Type, 4> copyinTypes;
213   SmallVector<OpAsmParser::OperandType, 4> allocates;
214   SmallVector<Type, 4> allocateTypes;
215   SmallVector<OpAsmParser::OperandType, 4> allocators;
216   SmallVector<Type, 4> allocatorTypes;
217   std::array<int, 8> segments{0, 0, 0, 0, 0, 0, 0, 0};
218   StringRef keyword;
219   bool defaultVal = false;
220   bool procBind = false;
221 
222   const int ifClausePos = 0;
223   const int numThreadsClausePos = 1;
224   const int privateClausePos = 2;
225   const int firstprivateClausePos = 3;
226   const int sharedClausePos = 4;
227   const int copyinClausePos = 5;
228   const int allocateClausePos = 6;
229   const int allocatorPos = 7;
230   const StringRef opName = result.name.getStringRef();
231 
232   while (succeeded(parser.parseOptionalKeyword(&keyword))) {
233     if (keyword == "if") {
234       // Fail if there was already another if condition.
235       if (segments[ifClausePos])
236         return allowedOnce(parser, "if", opName);
237       if (parser.parseLParen() || parser.parseOperand(ifCond.first) ||
238           parser.parseColonType(ifCond.second) || parser.parseRParen())
239         return failure();
240       segments[ifClausePos] = 1;
241     } else if (keyword == "num_threads") {
242       // Fail if there was already another num_threads clause.
243       if (segments[numThreadsClausePos])
244         return allowedOnce(parser, "num_threads", opName);
245       if (parser.parseLParen() || parser.parseOperand(numThreads.first) ||
246           parser.parseColonType(numThreads.second) || parser.parseRParen())
247         return failure();
248       segments[numThreadsClausePos] = 1;
249     } else if (keyword == "private") {
250       // Fail if there was already another private clause.
251       if (segments[privateClausePos])
252         return allowedOnce(parser, "private", opName);
253       if (parseOperandAndTypeList(parser, privates, privateTypes))
254         return failure();
255       segments[privateClausePos] = privates.size();
256     } else if (keyword == "firstprivate") {
257       // Fail if there was already another firstprivate clause.
258       if (segments[firstprivateClausePos])
259         return allowedOnce(parser, "firstprivate", opName);
260       if (parseOperandAndTypeList(parser, firstprivates, firstprivateTypes))
261         return failure();
262       segments[firstprivateClausePos] = firstprivates.size();
263     } else if (keyword == "shared") {
264       // Fail if there was already another shared clause.
265       if (segments[sharedClausePos])
266         return allowedOnce(parser, "shared", opName);
267       if (parseOperandAndTypeList(parser, shareds, sharedTypes))
268         return failure();
269       segments[sharedClausePos] = shareds.size();
270     } else if (keyword == "copyin") {
271       // Fail if there was already another copyin clause.
272       if (segments[copyinClausePos])
273         return allowedOnce(parser, "copyin", opName);
274       if (parseOperandAndTypeList(parser, copyins, copyinTypes))
275         return failure();
276       segments[copyinClausePos] = copyins.size();
277     } else if (keyword == "allocate") {
278       // Fail if there was already another allocate clause.
279       if (segments[allocateClausePos])
280         return allowedOnce(parser, "allocate", opName);
281       if (parseAllocateAndAllocator(parser, allocates, allocateTypes,
282                                     allocators, allocatorTypes))
283         return failure();
284       segments[allocateClausePos] = allocates.size();
285       segments[allocatorPos] = allocators.size();
286     } else if (keyword == "default") {
287       // Fail if there was already another default clause.
288       if (defaultVal)
289         return allowedOnce(parser, "default", opName);
290       defaultVal = true;
291       StringRef defval;
292       if (parser.parseLParen() || parser.parseKeyword(&defval) ||
293           parser.parseRParen())
294         return failure();
295       // The def prefix is required for the attribute as "private" is a keyword
296       // in C++.
297       auto attr = parser.getBuilder().getStringAttr("def" + defval);
298       result.addAttribute("default_val", attr);
299     } else if (keyword == "proc_bind") {
300       // Fail if there was already another proc_bind clause.
301       if (procBind)
302         return allowedOnce(parser, "proc_bind", opName);
303       procBind = true;
304       StringRef bind;
305       if (parser.parseLParen() || parser.parseKeyword(&bind) ||
306           parser.parseRParen())
307         return failure();
308       auto attr = parser.getBuilder().getStringAttr(bind);
309       result.addAttribute("proc_bind_val", attr);
310     } else {
311       return parser.emitError(parser.getNameLoc())
312              << keyword << " is not a valid clause for the " << opName
313              << " operation";
314     }
315   }
316 
317   // Add if parameter.
318   if (segments[ifClausePos] &&
319       parser.resolveOperand(ifCond.first, ifCond.second, result.operands))
320     return failure();
321 
322   // Add num_threads parameter.
323   if (segments[numThreadsClausePos] &&
324       parser.resolveOperand(numThreads.first, numThreads.second,
325                             result.operands))
326     return failure();
327 
328   // Add private parameters.
329   if (segments[privateClausePos] &&
330       parser.resolveOperands(privates, privateTypes, privates[0].location,
331                              result.operands))
332     return failure();
333 
334   // Add firstprivate parameters.
335   if (segments[firstprivateClausePos] &&
336       parser.resolveOperands(firstprivates, firstprivateTypes,
337                              firstprivates[0].location, result.operands))
338     return failure();
339 
340   // Add shared parameters.
341   if (segments[sharedClausePos] &&
342       parser.resolveOperands(shareds, sharedTypes, shareds[0].location,
343                              result.operands))
344     return failure();
345 
346   // Add copyin parameters.
347   if (segments[copyinClausePos] &&
348       parser.resolveOperands(copyins, copyinTypes, copyins[0].location,
349                              result.operands))
350     return failure();
351 
352   // Add allocate parameters.
353   if (segments[allocateClausePos] &&
354       parser.resolveOperands(allocates, allocateTypes, allocates[0].location,
355                              result.operands))
356     return failure();
357 
358   // Add allocator parameters.
359   if (segments[allocatorPos] &&
360       parser.resolveOperands(allocators, allocatorTypes, allocators[0].location,
361                              result.operands))
362     return failure();
363 
364   result.addAttribute("operand_segment_sizes",
365                       parser.getBuilder().getI32VectorAttr(segments));
366 
367   Region *body = result.addRegion();
368   SmallVector<OpAsmParser::OperandType, 4> regionArgs;
369   SmallVector<Type, 4> regionArgTypes;
370   if (parser.parseRegion(*body, regionArgs, regionArgTypes))
371     return failure();
372   return success();
373 }
374 
375 /// linear ::= `linear` `(` linear-list `)`
376 /// linear-list := linear-val | linear-val linear-list
377 /// linear-val := ssa-id-and-type `=` ssa-id-and-type
378 static ParseResult
379 parseLinearClause(OpAsmParser &parser,
380                   SmallVectorImpl<OpAsmParser::OperandType> &vars,
381                   SmallVectorImpl<Type> &types,
382                   SmallVectorImpl<OpAsmParser::OperandType> &stepVars) {
383   if (parser.parseLParen())
384     return failure();
385 
386   do {
387     OpAsmParser::OperandType var;
388     Type type;
389     OpAsmParser::OperandType stepVar;
390     if (parser.parseOperand(var) || parser.parseEqual() ||
391         parser.parseOperand(stepVar) || parser.parseColonType(type))
392       return failure();
393 
394     vars.push_back(var);
395     types.push_back(type);
396     stepVars.push_back(stepVar);
397   } while (succeeded(parser.parseOptionalComma()));
398 
399   if (parser.parseRParen())
400     return failure();
401 
402   return success();
403 }
404 
405 /// schedule ::= `schedule` `(` sched-list `)`
406 /// sched-list ::= sched-val | sched-val sched-list
407 /// sched-val ::= sched-with-chunk | sched-wo-chunk
408 /// sched-with-chunk ::= sched-with-chunk-types (`=` ssa-id-and-type)?
409 /// sched-with-chunk-types ::= `static` | `dynamic` | `guided`
410 /// sched-wo-chunk ::=  `auto` | `runtime`
411 static ParseResult
412 parseScheduleClause(OpAsmParser &parser, SmallString<8> &schedule,
413                     Optional<OpAsmParser::OperandType> &chunkSize) {
414   if (parser.parseLParen())
415     return failure();
416 
417   StringRef keyword;
418   if (parser.parseKeyword(&keyword))
419     return failure();
420 
421   schedule = keyword;
422   if (keyword == "static" || keyword == "dynamic" || keyword == "guided") {
423     if (succeeded(parser.parseOptionalEqual())) {
424       chunkSize = OpAsmParser::OperandType{};
425       if (parser.parseOperand(*chunkSize))
426         return failure();
427     } else {
428       chunkSize = llvm::NoneType::None;
429     }
430   } else if (keyword == "auto" || keyword == "runtime") {
431     chunkSize = llvm::NoneType::None;
432   } else {
433     return parser.emitError(parser.getNameLoc()) << " expected schedule kind";
434   }
435 
436   if (parser.parseRParen())
437     return failure();
438 
439   return success();
440 }
441 
442 /// Parses an OpenMP Workshare Loop operation
443 ///
444 /// operation ::= `omp.wsloop` loop-control clause-list
445 /// loop-control ::= `(` ssa-id-list `)` `:` type `=`  loop-bounds
446 /// loop-bounds := `(` ssa-id-list `)` to `(` ssa-id-list `)` steps
447 /// steps := `step` `(`ssa-id-list`)`
448 /// clause-list ::= clause | empty | clause-list
449 /// clause ::= private | firstprivate | lastprivate | linear | schedule |
450 //             collapse | nowait | ordered | order | inclusive
451 /// private ::= `private` `(` ssa-id-and-type-list `)`
452 /// firstprivate ::= `firstprivate` `(` ssa-id-and-type-list `)`
453 /// lastprivate ::= `lastprivate` `(` ssa-id-and-type-list `)`
454 /// linear ::= `linear` `(` linear-list `)`
455 /// schedule ::= `schedule` `(` sched-list `)`
456 /// collapse ::= `collapse` `(` ssa-id-and-type `)`
457 /// nowait ::= `nowait`
458 /// ordered ::= `ordered` `(` ssa-id-and-type `)`
459 /// order ::= `order` `(` `concurrent` `)`
460 /// inclusive ::= `inclusive`
461 ///
462 static ParseResult parseWsLoopOp(OpAsmParser &parser, OperationState &result) {
463   Type loopVarType;
464   int numIVs;
465 
466   // Parse an opening `(` followed by induction variables followed by `)`
467   SmallVector<OpAsmParser::OperandType> ivs;
468   if (parser.parseRegionArgumentList(ivs, /*requiredOperandCount=*/-1,
469                                      OpAsmParser::Delimiter::Paren))
470     return failure();
471 
472   numIVs = static_cast<int>(ivs.size());
473 
474   if (parser.parseColonType(loopVarType))
475     return failure();
476 
477   // Parse loop bounds.
478   SmallVector<OpAsmParser::OperandType> lower;
479   if (parser.parseEqual() ||
480       parser.parseOperandList(lower, numIVs, OpAsmParser::Delimiter::Paren) ||
481       parser.resolveOperands(lower, loopVarType, result.operands))
482     return failure();
483 
484   SmallVector<OpAsmParser::OperandType> upper;
485   if (parser.parseKeyword("to") ||
486       parser.parseOperandList(upper, numIVs, OpAsmParser::Delimiter::Paren) ||
487       parser.resolveOperands(upper, loopVarType, result.operands))
488     return failure();
489 
490   // Parse step values.
491   SmallVector<OpAsmParser::OperandType> steps;
492   if (parser.parseKeyword("step") ||
493       parser.parseOperandList(steps, numIVs, OpAsmParser::Delimiter::Paren) ||
494       parser.resolveOperands(steps, loopVarType, result.operands))
495     return failure();
496 
497   SmallVector<OpAsmParser::OperandType> privates;
498   SmallVector<Type> privateTypes;
499   SmallVector<OpAsmParser::OperandType> firstprivates;
500   SmallVector<Type> firstprivateTypes;
501   SmallVector<OpAsmParser::OperandType> lastprivates;
502   SmallVector<Type> lastprivateTypes;
503   SmallVector<OpAsmParser::OperandType> linears;
504   SmallVector<Type> linearTypes;
505   SmallVector<OpAsmParser::OperandType> linearSteps;
506   SmallString<8> schedule;
507   Optional<OpAsmParser::OperandType> scheduleChunkSize;
508   std::array<int, 9> segments{numIVs, numIVs, numIVs, 0, 0, 0, 0, 0, 0};
509 
510   const StringRef opName = result.name.getStringRef();
511   StringRef keyword;
512 
513   enum SegmentPos {
514     lbPos = 0,
515     ubPos,
516     stepPos,
517     privateClausePos,
518     firstprivateClausePos,
519     lastprivateClausePos,
520     linearClausePos,
521     linearStepPos,
522     scheduleClausePos,
523   };
524 
525   while (succeeded(parser.parseOptionalKeyword(&keyword))) {
526     if (keyword == "private") {
527       if (segments[privateClausePos])
528         return allowedOnce(parser, "private", opName);
529       if (parseOperandAndTypeList(parser, privates, privateTypes))
530         return failure();
531       segments[privateClausePos] = privates.size();
532     } else if (keyword == "firstprivate") {
533       // fail if there was already another firstprivate clause
534       if (segments[firstprivateClausePos])
535         return allowedOnce(parser, "firstprivate", opName);
536       if (parseOperandAndTypeList(parser, firstprivates, firstprivateTypes))
537         return failure();
538       segments[firstprivateClausePos] = firstprivates.size();
539     } else if (keyword == "lastprivate") {
540       // fail if there was already another shared clause
541       if (segments[lastprivateClausePos])
542         return allowedOnce(parser, "lastprivate", opName);
543       if (parseOperandAndTypeList(parser, lastprivates, lastprivateTypes))
544         return failure();
545       segments[lastprivateClausePos] = lastprivates.size();
546     } else if (keyword == "linear") {
547       // fail if there was already another linear clause
548       if (segments[linearClausePos])
549         return allowedOnce(parser, "linear", opName);
550       if (parseLinearClause(parser, linears, linearTypes, linearSteps))
551         return failure();
552       segments[linearClausePos] = linears.size();
553       segments[linearStepPos] = linearSteps.size();
554     } else if (keyword == "schedule") {
555       if (!schedule.empty())
556         return allowedOnce(parser, "schedule", opName);
557       if (parseScheduleClause(parser, schedule, scheduleChunkSize))
558         return failure();
559       if (scheduleChunkSize) {
560         segments[scheduleClausePos] = 1;
561       }
562     } else if (keyword == "collapse") {
563       auto type = parser.getBuilder().getI64Type();
564       mlir::IntegerAttr attr;
565       if (parser.parseLParen() || parser.parseAttribute(attr, type) ||
566           parser.parseRParen())
567         return failure();
568       result.addAttribute("collapse_val", attr);
569     } else if (keyword == "nowait") {
570       auto attr = UnitAttr::get(parser.getBuilder().getContext());
571       result.addAttribute("nowait", attr);
572     } else if (keyword == "ordered") {
573       mlir::IntegerAttr attr;
574       if (succeeded(parser.parseOptionalLParen())) {
575         auto type = parser.getBuilder().getI64Type();
576         if (parser.parseAttribute(attr, type))
577           return failure();
578         if (parser.parseRParen())
579           return failure();
580       } else {
581         // Use 0 to represent no ordered parameter was specified
582         attr = parser.getBuilder().getI64IntegerAttr(0);
583       }
584       result.addAttribute("ordered_val", attr);
585     } else if (keyword == "order") {
586       StringRef order;
587       if (parser.parseLParen() || parser.parseKeyword(&order) ||
588           parser.parseRParen())
589         return failure();
590       auto attr = parser.getBuilder().getStringAttr(order);
591       result.addAttribute("order", attr);
592     } else if (keyword == "inclusive") {
593       auto attr = UnitAttr::get(parser.getBuilder().getContext());
594       result.addAttribute("inclusive", attr);
595     }
596   }
597 
598   if (segments[privateClausePos]) {
599     parser.resolveOperands(privates, privateTypes, privates[0].location,
600                            result.operands);
601   }
602 
603   if (segments[firstprivateClausePos]) {
604     parser.resolveOperands(firstprivates, firstprivateTypes,
605                            firstprivates[0].location, result.operands);
606   }
607 
608   if (segments[lastprivateClausePos]) {
609     parser.resolveOperands(lastprivates, lastprivateTypes,
610                            lastprivates[0].location, result.operands);
611   }
612 
613   if (segments[linearClausePos]) {
614     parser.resolveOperands(linears, linearTypes, linears[0].location,
615                            result.operands);
616     auto linearStepType = parser.getBuilder().getI32Type();
617     SmallVector<Type> linearStepTypes(linearSteps.size(), linearStepType);
618     parser.resolveOperands(linearSteps, linearStepTypes,
619                            linearSteps[0].location, result.operands);
620   }
621 
622   if (!schedule.empty()) {
623     schedule[0] = llvm::toUpper(schedule[0]);
624     auto attr = parser.getBuilder().getStringAttr(schedule);
625     result.addAttribute("schedule_val", attr);
626     if (scheduleChunkSize) {
627       auto chunkSizeType = parser.getBuilder().getI32Type();
628       parser.resolveOperand(*scheduleChunkSize, chunkSizeType, result.operands);
629     }
630   }
631 
632   result.addAttribute("operand_segment_sizes",
633                       parser.getBuilder().getI32VectorAttr(segments));
634 
635   // Now parse the body.
636   Region *body = result.addRegion();
637   SmallVector<Type> ivTypes(numIVs, loopVarType);
638   if (parser.parseRegion(*body, ivs, ivTypes))
639     return failure();
640   return success();
641 }
642 
643 static void printWsLoopOp(OpAsmPrinter &p, WsLoopOp op) {
644   auto args = op.getRegion().front().getArguments();
645   p << op.getOperationName() << " (" << args << ") : " << args[0].getType()
646     << " = (" << op.lowerBound() << ") to (" << op.upperBound() << ") step ("
647     << op.step() << ")";
648 
649   // Print private, firstprivate, shared and copyin parameters
650   auto printDataVars = [&p](StringRef name, OperandRange vars) {
651     if (vars.empty())
652       return;
653 
654     p << " " << name << "(";
655     llvm::interleaveComma(
656         vars, p, [&](const Value &v) { p << v << " : " << v.getType(); });
657     p << ")";
658   };
659   printDataVars("private", op.private_vars());
660   printDataVars("firstprivate", op.firstprivate_vars());
661   printDataVars("lastprivate", op.lastprivate_vars());
662 
663   auto linearVars = op.linear_vars();
664   auto linearVarsSize = linearVars.size();
665   if (linearVarsSize) {
666     p << " "
667       << "linear"
668       << "(";
669     for (unsigned i = 0; i < linearVarsSize; ++i) {
670       std::string separator = i == linearVarsSize - 1 ? ")" : ", ";
671       p << linearVars[i];
672       if (op.linear_step_vars().size() > i)
673         p << " = " << op.linear_step_vars()[i];
674       p << " : " << linearVars[i].getType() << separator;
675     }
676   }
677 
678   if (auto sched = op.schedule_val()) {
679     auto schedLower = sched->lower();
680     p << " schedule(" << schedLower;
681     if (auto chunk = op.schedule_chunk_var()) {
682       p << " = " << chunk;
683     }
684     p << ")";
685   }
686 
687   if (auto collapse = op.collapse_val())
688     p << " collapse(" << collapse << ")";
689 
690   if (op.nowait())
691     p << " nowait";
692 
693   if (auto ordered = op.ordered_val()) {
694     p << " ordered(" << ordered << ")";
695   }
696 
697   if (op.inclusive()) {
698     p << " inclusive";
699   }
700 
701   p.printRegion(op.region(), /*printEntryBlockArgs=*/false);
702 }
703 
704 //===----------------------------------------------------------------------===//
705 // WsLoopOp
706 //===----------------------------------------------------------------------===//
707 
708 void WsLoopOp::build(OpBuilder &builder, OperationState &state,
709                      ValueRange lowerBound, ValueRange upperBound,
710                      ValueRange step, ArrayRef<NamedAttribute> attributes) {
711   build(builder, state, TypeRange(), lowerBound, upperBound, step,
712         /*private_vars=*/ValueRange(),
713         /*firstprivate_vars=*/ValueRange(), /*lastprivate_vars=*/ValueRange(),
714         /*linear_vars=*/ValueRange(), /*linear_step_vars=*/ValueRange(),
715         /*schedule_val=*/nullptr, /*schedule_chunk_var=*/nullptr,
716         /*collapse_val=*/nullptr,
717         /*nowait=*/nullptr, /*ordered_val=*/nullptr, /*order_val=*/nullptr,
718         /*inclusive=*/nullptr, /*buildBody=*/false);
719   state.addAttributes(attributes);
720 }
721 
722 void WsLoopOp::build(OpBuilder &, OperationState &state, TypeRange resultTypes,
723                      ValueRange operands, ArrayRef<NamedAttribute> attributes) {
724   state.addOperands(operands);
725   state.addAttributes(attributes);
726   (void)state.addRegion();
727   assert(resultTypes.size() == 0u && "mismatched number of return types");
728   state.addTypes(resultTypes);
729 }
730 
731 void WsLoopOp::build(OpBuilder &builder, OperationState &result,
732                      TypeRange typeRange, ValueRange lowerBounds,
733                      ValueRange upperBounds, ValueRange steps,
734                      ValueRange privateVars, ValueRange firstprivateVars,
735                      ValueRange lastprivateVars, ValueRange linearVars,
736                      ValueRange linearStepVars, StringAttr scheduleVal,
737                      Value scheduleChunkVar, IntegerAttr collapseVal,
738                      UnitAttr nowait, IntegerAttr orderedVal,
739                      StringAttr orderVal, UnitAttr inclusive, bool buildBody) {
740   result.addOperands(lowerBounds);
741   result.addOperands(upperBounds);
742   result.addOperands(steps);
743   result.addOperands(privateVars);
744   result.addOperands(firstprivateVars);
745   result.addOperands(linearVars);
746   result.addOperands(linearStepVars);
747   if (scheduleChunkVar)
748     result.addOperands(scheduleChunkVar);
749 
750   if (scheduleVal)
751     result.addAttribute("schedule_val", scheduleVal);
752   if (collapseVal)
753     result.addAttribute("collapse_val", collapseVal);
754   if (nowait)
755     result.addAttribute("nowait", nowait);
756   if (orderedVal)
757     result.addAttribute("ordered_val", orderedVal);
758   if (orderVal)
759     result.addAttribute("order", orderVal);
760   if (inclusive)
761     result.addAttribute("inclusive", inclusive);
762   result.addAttribute(
763       WsLoopOp::getOperandSegmentSizeAttr(),
764       builder.getI32VectorAttr(
765           {static_cast<int32_t>(lowerBounds.size()),
766            static_cast<int32_t>(upperBounds.size()),
767            static_cast<int32_t>(steps.size()),
768            static_cast<int32_t>(privateVars.size()),
769            static_cast<int32_t>(firstprivateVars.size()),
770            static_cast<int32_t>(lastprivateVars.size()),
771            static_cast<int32_t>(linearVars.size()),
772            static_cast<int32_t>(linearStepVars.size()),
773            static_cast<int32_t>(scheduleChunkVar != nullptr ? 1 : 0)}));
774 
775   Region *bodyRegion = result.addRegion();
776   if (buildBody) {
777     OpBuilder::InsertionGuard guard(builder);
778     unsigned numIVs = steps.size();
779     SmallVector<Type, 8> argTypes(numIVs, steps.getType().front());
780     builder.createBlock(bodyRegion, {}, argTypes);
781   }
782 }
783 
784 #define GET_OP_CLASSES
785 #include "mlir/Dialect/OpenMP/OpenMPOps.cpp.inc"
786