[mlir][OpenMP] Added assemblyFormat for SectionsOp

[shraiysh]

Summary

This patch adds assemblyFormat for omp.sections operation.

Some existing functions have been altered to fit the custom directive
in assemblyFormat. This has led to their callsites to get modified too,
but those will be removed in later patches, when other operations get
their assemblyFormat. All operations were not changed in one patch for
ease of review.

[shraiysh]

Corrected comments.

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

def SectionsOp : OpenMP_Op<"sections", [AttrSizedOperandSegments]> {
   let arguments = (ins Variadic<OpenMP_PointerLikeType>:$reduction_vars,
                        OptionalAttr<SymbolRefArrayAttr>:$reductions,
                        Variadic<AnyType>:$allocate_vars,
                        Variadic<AnyType>:$allocators_vars,
                        UnitAttr:$nowait);
 
   let regions = (region SizedRegion<1>:$region);
 
-  let hasCustomAssemblyFormat = 1;
+  let assemblyFormat = [{
+    oilist( `reduction` `(`
+              custom<ReductionVarList>(
+                $reduction_vars, type($reduction_vars), $reductions
+              ) `)`
+          | `allocate` `(`
+              custom<AllocateAndAllocator>(
+                $allocate_vars, type($allocate_vars),
+                $allocators_vars, type($allocators_vars)
+              ) `)`
+          | `nowait`
+    ) $region attr-dict
+  }];
+
   let hasVerifier = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // 2.9.2 Workshare Loop Construct
 //===----------------------------------------------------------------------===//
 
 def WsLoopOp : OpenMP_Op<"wsloop", [AttrSizedOperandSegments,

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

 }
 
 //===----------------------------------------------------------------------===//
 // Parser and printer for Allocate Clause
 //===----------------------------------------------------------------------===//
 
 /// Parse an allocate clause with allocators and a list of operands with types.
 ///
-/// allocate ::= `allocate` `(` allocate-operand-list `)`
 /// allocate-operand-list :: = allocate-operand |
 ///                            allocator-operand `,` allocate-operand-list
 /// allocate-operand :: = ssa-id-and-type -> ssa-id-and-type
 /// ssa-id-and-type ::= ssa-id `:` type
 static ParseResult parseAllocateAndAllocator(
     OpAsmParser &parser,
     SmallVectorImpl<OpAsmParser::OperandType> &operandsAllocate,
     SmallVectorImpl<Type> &typesAllocate,
if (simd)
     p << ", simd";
   p << ") ";
 }
 
 //===----------------------------------------------------------------------===//
 // Parser, printer and verifier for ReductionVarList
 //===----------------------------------------------------------------------===//
 
-/// reduction ::= `reduction` `(` reduction-entry-list `)`
 /// reduction-entry-list ::= reduction-entry
 ///                        | reduction-entry-list `,` reduction-entry
 /// reduction-entry ::= symbol-ref `->` ssa-id `:` type
-static ParseResult
-parseReductionVarList(OpAsmParser &parser,
-                      SmallVectorImpl<SymbolRefAttr> &symbols,
-                      SmallVectorImpl<OpAsmParser::OperandType> &operands,
-                      SmallVectorImpl<Type> &types) {
-  if (failed(parser.parseLParen()))
-    return failure();
-
+static ParseResult parseReductionVarList(
+    OpAsmParser &parser, SmallVectorImpl<OpAsmParser::OperandType> &operands,
+    SmallVectorImpl<Type> &types, ArrayAttr &redcuctionSymbols) {
+  SmallVector<SymbolRefAttr> reductionVec;
   do {
-    if (parser.parseAttribute(symbols.emplace_back()) || parser.parseArrow() ||
-        parser.parseOperand(operands.emplace_back()) ||
+    if (parser.parseAttribute(reductionVec.emplace_back()) ||
+        parser.parseArrow() || parser.parseOperand(operands.emplace_back()) ||
         parser.parseColonType(types.emplace_back()))
       return failure();
   } while (succeeded(parser.parseOptionalComma()));
-  return parser.parseRParen();
+  SmallVector<Attribute> reductions(reductionVec.begin(), reductionVec.end());
+  redcuctionSymbols = ArrayAttr::get(parser.getContext(), reductions);
+  return success();
 }
 
 /// Print Reduction clause
-static void printReductionVarList(OpAsmPrinter &p,
-                                  Optional<ArrayAttr> reductions,
-                                  OperandRange reductionVars) {
-  p << "reduction(";
+static void printReductionVarList(OpAsmPrinter &p, Operation *op,
+                                  OperandRange reductionVars,
+                                  TypeRange reductionTypes,
+                                  Optional<ArrayAttr> reductions) {
   for (unsigned i = 0, e = reductions->size(); i < e; ++i) {
     if (i != 0)
       p << ", ";
     p << (*reductions)[i] << " -> " << reductionVars[i] << " : "
       << reductionVars[i].getType();
   }
-  p << ") ";
 }
 
 /// Verifies Reduction Clause
 static LogicalResult verifyReductionVarList(Operation *op,
                                             Optional<ArrayAttr> reductions,
                                             OperandRange reductionVars) {
   if (!reductionVars.empty()) {
     if (!reductions || reductions->size() != reductionVars.size())
static ParseResult parseClauses(OpAsmParser &parser, OperationState &result,
   std::pair<OpAsmParser::OperandType, Type> ifCond;
   std::pair<OpAsmParser::OperandType, Type> numThreads;
   std::pair<OpAsmParser::OperandType, Type> device;
   std::pair<OpAsmParser::OperandType, Type> threadLimit;
 
   SmallVector<OpAsmParser::OperandType> allocates, allocators;
   SmallVector<Type> allocateTypes, allocatorTypes;
 
-  SmallVector<SymbolRefAttr> reductionSymbols;
+  ArrayAttr reductions;
   SmallVector<OpAsmParser::OperandType> reductionVars;
   SmallVector<Type> reductionVarTypes;
 
   SmallVector<OpAsmParser::OperandType> linears;
   SmallVector<Type> linearTypes;
   SmallVector<OpAsmParser::OperandType> linearSteps;
 
   SmallString<8> schedule;
if (clauseKeyword == "if") {
       clauseSegments[pos[allocateClause]] = allocates.size();
       clauseSegments[pos[allocateClause] + 1] = allocators.size();
     } else if (clauseKeyword == "proc_bind") {
       if (checkAllowed(procBindClause) ||
           parseClauseAttr<ClauseProcBindKindAttr>(parser, result,
                                                   "proc_bind_val", "proc bind"))
         return failure();
     } else if (clauseKeyword == "reduction") {
-      if (checkAllowed(reductionClause) ||
-          parseReductionVarList(parser, reductionSymbols, reductionVars,
-                                reductionVarTypes))
+      if (checkAllowed(reductionClause) || parser.parseLParen() ||
+          parseReductionVarList(parser, reductionVars, reductionVarTypes,
+                                reductions) ||
+          parser.parseRParen())
         return failure();
       clauseSegments[pos[reductionClause]] = reductionVars.size();
     } else if (clauseKeyword == "nowait") {
       if (checkAllowed(nowaitClause))
         return failure();
       auto attr = UnitAttr::get(parser.getBuilder().getContext());
       result.addAttribute("nowait", attr);
     } else if (clauseKeyword == "linear") {
if (done[allocateClause] && clauseSegments[pos[allocateClause] + 1] &&
                                     allocators[0].location, result.operands)))
     return failure();
 
   // Add reduction parameters and symbols
   if (done[reductionClause] && clauseSegments[pos[reductionClause]]) {
     if (failed(parser.resolveOperands(reductionVars, reductionVarTypes,
                                       parser.getNameLoc(), result.operands)))
       return failure();
-
-    SmallVector<Attribute> reductions(reductionSymbols.begin(),
-                                      reductionSymbols.end());
-    result.addAttribute("reductions",
-                        parser.getBuilder().getArrayAttr(reductions));
+    result.addAttribute("reductions", reductions);
   }
 
   // Add linear parameters
   if (done[linearClause] && clauseSegments[pos[linearClause]]) {
     auto linearStepType = parser.getBuilder().getI32Type();
     SmallVector<Type> linearStepTypes(linearSteps.size(), linearStepType);
     if (failed(parser.resolveOperands(linears, linearTypes, linears[0].location,
                                       result.operands)) ||
static ParseResult parseClauses(OpAsmParser &parser, OperationState &result,
   }
 
   segments.insert(segments.end(), clauseSegments.begin(), clauseSegments.end());
 
   return success();
 }
 
 //===----------------------------------------------------------------------===//
-// Parser, printer and verifier for SectionsOp
+// Verifier for SectionsOp
 //===----------------------------------------------------------------------===//
 
-/// Parses an OpenMP Sections operation
-///
-/// sections ::= `omp.sections` clause-list
-/// clause-list ::= clause clause-list | empty
-/// clause ::= reduction | allocate | nowait
-ParseResult SectionsOp::parse(OpAsmParser &parser, OperationState &result) {
-  SmallVector<ClauseType> clauses = {reductionClause, allocateClause,
-                                     nowaitClause};
-
-  SmallVector<int> segments;
-
-  if (failed(parseClauses(parser, result, clauses, segments)))
-    return failure();
-
-  result.addAttribute("operand_segment_sizes",
-                      parser.getBuilder().getI32VectorAttr(segments));
-
-  // Now parse the body.
-  Region *body = result.addRegion();
-  if (parser.parseRegion(*body))
-    return failure();
-  return success();
-}
-
-void SectionsOp::print(OpAsmPrinter &p) {
-  p << " ";
-
-  if (!reduction_vars().empty())
-    printReductionVarList(p, reductions(), reduction_vars());
-
-  if (!allocate_vars().empty()) {
-    printAllocateAndAllocator(p << "allocate(", *this, allocate_vars(),
-                              allocate_vars().getTypes(), allocators_vars(),
-                              allocators_vars().getTypes());
-    p << ")";
-  }
-
-  if (nowait())
-    p << "nowait";
-
-  p << ' ';
-  p.printRegion(region());
-}
-
 LogicalResult SectionsOp::verify() {
   if (allocate_vars().size() != allocators_vars().size())
     return emitError(
         "expected equal sizes for allocate and allocator variables");
 
   for (auto &inst : *region().begin()) {
     if (!(isa<SectionOp>(inst) || isa<TerminatorOp>(inst))) {
       return emitOpError()
if (nowait())
     p << "nowait ";
 
   if (auto ordered = ordered_val())
     p << "ordered(" << ordered << ") ";
 
   if (auto order = order_val())
     p << "order(" << stringifyClauseOrderKind(*order) << ") ";
 
-  if (!reduction_vars().empty())
-    printReductionVarList(p, reductions(), reduction_vars());
+  if (!reduction_vars().empty()) {
+    printReductionVarList(p << "reduction(", *this, reduction_vars(),
+                          reduction_vars().getTypes(), reductions());
+    p << ")";
+  }
 
   p << ' ';
   p.printRegion(region(), /*printEntryBlockArgs=*/false);
 }
 
 //===----------------------------------------------------------------------===//
 // ReductionOp
 //===----------------------------------------------------------------------===//

mlir/test/Dialect/OpenMP/invalid.mlir

omp.sections {
     "test.payload" () : () -> ()
   }
   return
 }
 
 // -----
 
 func @omp_sections(%cond : i1) {
-  // expected-error @below {{if is not a valid clause for the omp.sections operation}}
+  // expected-error @below {{expected '{' to begin a region}}
   omp.sections if(%cond) {
     omp.terminator
   }
   return
 }
 
 // -----
 
 func @omp_sections() {
-  // expected-error @below {{num_threads is not a valid clause for the omp.sections operation}}
+  // expected-error @below {{expected '{' to begin a region}}
   omp.sections num_threads(10) {
     omp.terminator
   }
   return
 }
 
 // -----
 
 func @omp_sections() {
-  // expected-error @below {{proc_bind is not a valid clause for the omp.sections operation}}
+  // expected-error @below {{expected '{' to begin a region}}
   omp.sections proc_bind(close) {
     omp.terminator
   }
   return
 }
 
 // -----
 
 func @omp_sections(%data_var : memref<i32>, %linear_var : i32) {
-  // expected-error @below {{linear is not a valid clause for the omp.sections operation}}
+  // expected-error @below {{expected '{' to begin a region}}
   omp.sections linear(%data_var = %linear_var : memref<i32>) {
     omp.terminator
   }
   return
 }
 
 // -----
 
 func @omp_sections() {
-  // expected-error @below {{schedule is not a valid clause for the omp.sections operation}}
+  // expected-error @below {{expected '{' to begin a region}}
   omp.sections schedule(static, none) {
     omp.terminator
   }
   return
 }
 
 // -----
 
 func @omp_sections() {
-  // expected-error @below {{collapse is not a valid clause for the omp.sections operation}}
+  // expected-error @below {{expected '{' to begin a region}}
   omp.sections collapse(3) {
     omp.terminator
   }
   return
 }
 
 // -----
 
 func @omp_sections() {
-  // expected-error @below {{ordered is not a valid clause for the omp.sections operation}}
+  // expected-error @below {{expected '{' to begin a region}}
   omp.sections ordered(2) {
     omp.terminator
   }
   return
 }
 
 // -----
 
 func @omp_sections() {
-  // expected-error @below {{order is not a valid clause for the omp.sections operation}}
+  // expected-error @below {{expected '{' to begin a region}}
   omp.sections order(concurrent) {
     omp.terminator
   }
   return
 }
 
 // -----
 

mlir/test/Dialect/OpenMP/ops.mlir

 
 // CHECK-LABEL: omp_sectionsop
 func @omp_sectionsop(%data_var1 : memref<i32>, %data_var2 : memref<i32>,
                      %data_var3 : memref<i32>, %redn_var : !llvm.ptr<f32>) {
   // CHECK: omp.sections allocate(%{{.*}} : memref<i32> -> %{{.*}} : memref<i32>)
   "omp.sections" (%data_var1, %data_var1) ({
     // CHECK: omp.terminator
     omp.terminator
-  }) {operand_segment_sizes = dense<[0,1,1]> : vector<3xi32>} : (memref<i32>, memref<i32>) -> ()
+  }) {allocate, operand_segment_sizes = dense<[0,1,1]> : vector<3xi32>} : (memref<i32>, memref<i32>) -> ()
 
     // CHECK: omp.sections reduction(@add_f32 -> %{{.*}} : !llvm.ptr<f32>)
   "omp.sections" (%redn_var) ({
     // CHECK: omp.terminator
     omp.terminator
-  }) {operand_segment_sizes = dense<[1,0,0]> : vector<3xi32>, reductions=[@add_f32]} : (!llvm.ptr<f32>) -> ()
+  }) {reduction, operand_segment_sizes = dense<[1,0,0]> : vector<3xi32>, reductions=[@add_f32]} : (!llvm.ptr<f32>) -> ()
 
   // CHECK: omp.sections nowait {
   omp.sections nowait {
     // CHECK: omp.terminator
     omp.terminator
   }
 
   // CHECK: omp.sections reduction(@add_f32 -> %{{.*}} : !llvm.ptr<f32>) {