[mlir][OpenMP] Add custom parser and pretty printer for parallel construct

[DavidTruby]

This is my first pass at a custom parser and pretty printer but I had a couple of questions:

1. Is it necessary to pass and parse the type for each data variable? I couldn't find any other way of finding out what the type of that argument is at the right time, even though the variable hasn't changed types between the two uses.
2. Is there a better way of doing position-independent keywords than my above while loop? I feel like that code is messier than it needs to be.

[DavidTruby]

Added additional tests

[DavidTruby]

Added tests for multiple variables in the same clause

[kiranchandramohan]

@DavidTruby I have a few questions/comments.

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

113

would stringswitch make any difference?

121

Is there no map from keyword to segment number?

138–142

Would having a wrapper for this common pattern make any sense?

ParseResult parseOpenMPClause(parser,segment,operand,operandtype)

183

should we have negative tests?

188

Can this be along with the parsing code inside the loop?

[DavidTruby]

Copyin bug fix.

[DavidTruby]

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

113

As I understand it StringSwitch only works for cases that return a single value, rather than having arbitrary code in each case, so we can't use it here.

138–142

Since we need to return early on failure I'm not sure this is that easy to abstract. I could try it though.

188

The parsing code inside the loop is trying to be position independent (so that the keywords can appear in any order), but the parameters do need to be in the correct order. I believe the only way to ensure that is to add them in the right order down here.

[mehdi_amini]

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

47

Please avoid trivial braces in MLIR in general.

116

In general we just return failure when calling into the parser fails, because it already emits an error.
If you fail based on other condition, you should emit a textual error otherwise it'll fail silently.

e.g. : https://github.com/llvm/llvm-project/blob/master/mlir/lib/Dialect/StandardOps/IR/Ops.cpp#L288
(we also test these kind of diagnostics)

mlir/test/Dialect/OpenMP/ops.mlir

108

This seems a bit light on CHECK lines?

[DavidTruby]

Improved error messages with tests and more thorough CHECK lines on existing tests.

[mehdi_amini]

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

25

This should go with the first include section. I suspect the "inc" is here because it depends on the other includes.

[jdoerfert]

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

121

Agreed. Maybe:

keywords ={"if", "...", ...}
...
idx = 0
if (keyword == keywords[idx++]) {
  ...
}

Or even more structure, e.g., we create these through tablegen ;)

(Nit: segments is used as bool array, correct?)

[DavidTruby]

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

121

Segments isn't a bool array, it's an array of argument segment sizes that is needed for passing of variadic arguments in MLIR. In the case of if and num_threads the value will always be 0 or 1 (as there can only be one argument passed), but in the case of private/firstprivate/shared/copyin the value will be the number of arguments passed to that clause.

[ftynse]

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

42

Please add documentation to this function and the functions below. We usually have some variant of EBNF describing what a parsing function accepts.

76

MLIR mostly uses unsigned in such cases

124

Nit: let's not hardcode more than necessary, the operation name can be fetched as result.name.getStrRef() and cached somewhere.

128

Nit: I'd consider named constants for positions in the segment, e.g. segments[kIfClausePos].

237

Nit: do we really need {}?

[DavidTruby]

Added documentation to the new functions.

[ftynse]

LGTM after clang-format and addressing Mehdi's comment about #include order

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

25

Please fix.

[DavidTruby]

clang-format fixes

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

let arguments = (ins Optional<I1>:$if_expr_var,
              OptionalAttr<ClauseDefault>:$default_val,
              Variadic<AnyType>:$private_vars,
              Variadic<AnyType>:$firstprivate_vars,
              Variadic<AnyType>:$shared_vars,
              Variadic<AnyType>:$copyin_vars,
              OptionalAttr<ClauseProcBind>:$proc_bind_val);
 
   let regions = (region AnyRegion:$region);
+
+  let parser = [{ return parseParallelOp(parser, result); }];
+  let printer = [{ return printParallelOp(p, *this); }];
 }
 
 def TerminatorOp : OpenMP_Op<"terminator", [Terminator]> {
   let summary = "terminator for OpenMP regions";
   let description = [{
     A terminator operation for regions that appear in the body of OpenMP
     operation.  These regions are not expected to return any value so the
     terminator takes no operands. The terminator op returns control to the

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

 //===- OpenMPDialect.cpp - MLIR Dialect for OpenMP implementation ---------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements the OpenMP dialect and its operations.
 //
 //===----------------------------------------------------------------------===//
 
 #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/StringRef.h"
 #include "llvm/ADT/StringSwitch.h"
+#include <cstddef>
 
 #include "mlir/Dialect/OpenMP/OpenMPOpsEnums.cpp.inc"
 
 using namespace mlir;
 using namespace mlir::omp;
 
 OpenMPDialect::OpenMPDialect(MLIRContext *context)
     : Dialect(getDialectNamespace(), context) {
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/OpenMP/OpenMPOps.cpp.inc"
       >();
 }
 
+//===----------------------------------------------------------------------===//
+// ParallelOp
+//===----------------------------------------------------------------------===//
+
+/// Parse a list of operands with types.
+///
+/// operand-and-type-list ::= `(` ssa-id-and-type-list `)`
+/// ssa-id-and-type-list ::= ssa-id-and-type |
+///                          ssa-id-and-type ',' ssa-id-and-type-list
+/// ssa-id-and-type ::= ssa-id `:` type
+static ParseResult
+parseOperandAndTypeList(OpAsmParser &parser,
+                        SmallVectorImpl<OpAsmParser::OperandType> &operands,
+                        SmallVectorImpl<Type> &types) {
+  if (parser.parseLParen())
+    return failure();
+
+  do {
+    OpAsmParser::OperandType operand;
+    Type type;
+    if (parser.parseOperand(operand) || parser.parseColonType(type))
+      return failure();
+    operands.push_back(operand);
+    types.push_back(type);
+  } while (succeeded(parser.parseOptionalComma()));
+
+  if (parser.parseRParen())
+    return failure();
+
+  return success();
+}
+
+static void printParallelOp(OpAsmPrinter &p, ParallelOp op) {
+  p << "omp.parallel";
+
+  if (auto ifCond = op.if_expr_var())
+    p << " if(" << ifCond << ")";
+
+  if (auto threads = op.num_threads_var())
+    p << " num_threads(" << threads << " : " << threads.getType() << ")";
+
+  // Print private, firstprivate, shared and copyin parameters
+  auto printDataVars = [&p](StringRef name, OperandRange vars) {
+    if (vars.size()) {
+      p << " " << name << "(";
+      for (unsigned i = 0; i < vars.size(); ++i) {
+        std::string separator = i == vars.size() - 1 ? ")" : ", ";
+        p << vars[i] << " : " << vars[i].getType() << separator;
+      }
+    }
+  };
+  printDataVars("private", op.private_vars());
+  printDataVars("firstprivate", op.firstprivate_vars());
+  printDataVars("shared", op.shared_vars());
+  printDataVars("copyin", op.copyin_vars());
+
+  if (auto def = op.default_val())
+    p << " default(" << def->drop_front(3) << ")";
+
+  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) {
+  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
+/// clause-list ::= clause | clause clause-list
+/// clause ::= if | numThreads | private | firstprivate | shared | copyin |
+///            default | procBind
+/// if ::= `if` `(` ssa-id `)`
+/// numThreads ::= `num_threads` `(` ssa-id-and-type `)`
+/// private ::= `private` operand-and-type-list
+/// firstprivate ::= `firstprivate` operand-and-type-list
+/// shared ::= `shared` operand-and-type-list
+/// copyin ::= `copyin` operand-and-type-list
+/// default ::= `default` `(` (`private` | `firstprivate` | `shared` | `none`)
+/// procBind ::= `proc_bind` `(` (`master` | `close` | `spread`) `)`
+///
+/// Note that each clause can only appear once in the clase-list.
+static ParseResult parseParallelOp(OpAsmParser &parser,
+                                   OperationState &result) {
+  OpAsmParser::OperandType ifCond;
+  std::pair<OpAsmParser::OperandType, Type> numThreads;
+  llvm::SmallVector<OpAsmParser::OperandType, 4> privates;
+  llvm::SmallVector<Type, 4> privateTypes;
+  llvm::SmallVector<OpAsmParser::OperandType, 4> firstprivates;
+  llvm::SmallVector<Type, 4> firstprivateTypes;
+  llvm::SmallVector<OpAsmParser::OperandType, 4> shareds;
+  llvm::SmallVector<Type, 4> sharedTypes;
+  llvm::SmallVector<OpAsmParser::OperandType, 4> copyins;
+  llvm::SmallVector<Type, 4> copyinTypes;
+  std::array<int, 6> segments{0, 0, 0, 0, 0, 0};
+  llvm::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 llvm::StringRef opName = result.name.getStringRef();
+
+  while (succeeded(parser.parseOptionalKeyword(&keyword))) {
+    if (keyword == "if") {
+      // Fail if there was already another if condition
+      if (segments[ifClausePos])
+        return allowedOnce(parser, "if", opName);
+      if (parser.parseLParen() || parser.parseOperand(ifCond) ||
+          parser.parseRParen())
+        return failure();
+      segments[ifClausePos] = 1;
+    } else if (keyword == "num_threads") {
+      // 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
+      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 == "shared") {
+      // 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
+      if (segments[copyinClausePos])
+        return allowedOnce(parser, "copyin", opName);
+      if (parseOperandAndTypeList(parser, copyins, copyinTypes))
+        return failure();
+      segments[copyinClausePos] = copyins.size();
+    } else if (keyword == "default") {
+      // fail if there was already another default clause
+      if (defaultVal)
+        return allowedOnce(parser, "default", opName);
+      defaultVal = true;
+      llvm::StringRef defval;
+      if (parser.parseLParen() || parser.parseKeyword(&defval) ||
+          parser.parseRParen())
+        return failure();
+      llvm::SmallString<16> attrval;
+      // The def prefix is required for the attribute as "private" is a keyword
+      // 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 default clause
+      if (procBind)
+        return allowedOnce(parser, "proc_bind", opName);
+      procBind = true;
+      llvm::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
+  if (segments[ifClausePos]) {
+    parser.resolveOperand(ifCond, parser.getBuilder().getI1Type(),
+                          result.operands);
+  }
+
+  // Add num_threads parameter
+  if (segments[numThreadsClausePos]) {
+    parser.resolveOperand(numThreads.first, numThreads.second, result.operands);
+  }
+
+  // Add private parameters
+  if (segments[privateClausePos]) {
+    parser.resolveOperands(privates, privateTypes, privates[0].location,
+                           result.operands);
+  }
+
+  // Add firstprivate parameters
+  if (segments[firstprivateClausePos]) {
+    parser.resolveOperands(firstprivates, firstprivateTypes,
+                           firstprivates[0].location, result.operands);
+  }
+
+  // Add shared parameters
+  if (segments[sharedClausePos]) {
+    parser.resolveOperands(shareds, sharedTypes, shareds[0].location,
+                           result.operands);
+  }
+
+  // Add copyin parameters
+  if (segments[copyinClausePos]) {
+    parser.resolveOperands(copyins, copyinTypes, copyins[0].location,
+                           result.operands);
+  }
+
+  result.addAttribute("operand_segment_sizes",
+                      parser.getBuilder().getI32VectorAttr(segments));
+
+  Region *body = result.addRegion();
+  llvm::SmallVector<OpAsmParser::OperandType, 4> regionArgs;
+  llvm::SmallVector<Type, 4> regionArgTypes;
+  if (parser.parseRegion(*body, regionArgs, regionArgTypes))
+    return failure();
+  return success();
+}
+
 namespace mlir {
 namespace omp {
 #define GET_OP_CLASSES
 #include "mlir/Dialect/OpenMP/OpenMPOps.cpp.inc"
 } // namespace omp
 } // namespace mlir

mlir/test/Dialect/OpenMP/invalid.mlir

+// RUN: mlir-opt -split-input-file -verify-diagnostics %s
+
+func @unknown_clause() {
+  // expected-error@+1 {{invalid is not a valid clause for the omp.parallel operation}}
+  omp.parallel invalid {
+  }
+
+  return
+}
+
+// -----
+
+func @if_once(%n : i1) {
+  // expected-error@+1 {{at most one if clause can appear on the omp.parallel operation}}
+  omp.parallel if(%n) if(%n) {
+  }
+
+  return
+}
+
+// -----
+
+func @num_threads_once(%n : si32) {
+  // expected-error@+1 {{at most one num_threads clause can appear on the omp.parallel operation}}
+  omp.parallel num_threads(%n : si32) num_threads(%n : si32) {
+  }
+
+  return
+}
+
+// -----
+
+func @private_once(%n : memref<i32>) {
+  // expected-error@+1 {{at most one private clause can appear on the omp.parallel operation}}
+  omp.parallel private(%n : memref<i32>) private(%n : memref<i32>) {
+  }
+
+  return
+}
+
+// -----
+
+func @firstprivate_once(%n : memref<i32>) {
+  // expected-error@+1 {{at most one firstprivate clause can appear on the omp.parallel operation}}
+  omp.parallel firstprivate(%n : memref<i32>) firstprivate(%n : memref<i32>) {
+  }
+
+  return
+}
+
+// -----
+
+func @shared_once(%n : memref<i32>) {
+  // expected-error@+1 {{at most one shared clause can appear on the omp.parallel operation}}
+  omp.parallel shared(%n : memref<i32>) shared(%n : memref<i32>) {
+  }
+
+  return
+}
+
+// -----
+
+func @copyin_once(%n : memref<i32>) {
+  // expected-error@+1 {{at most one copyin clause can appear on the omp.parallel operation}}
+  omp.parallel copyin(%n : memref<i32>) copyin(%n : memref<i32>) {
+  }
+
+  return
+}
+
+// -----
+ 
+func @default_once() {
+  // expected-error@+1 {{at most one default clause can appear on the omp.parallel operation}}
+  omp.parallel default(private) default(firstprivate) {
+  }
+
+  return
+}
+
+// -----
+
+func @proc_bind_once() {
+  // expected-error@+1 {{at most one proc_bind clause can appear on the omp.parallel operation}}
+  omp.parallel proc_bind(close) proc_bind(spread) {
+  }
+
+  return
+}

mlir/test/Dialect/OpenMP/ops.mlir

-// RUN: mlir-opt -verify-diagnostics %s | FileCheck %s
+// RUN: mlir-opt %s | mlir-opt | FileCheck %s
 
 func @omp_barrier() -> () {
   // CHECK: omp.barrier
   omp.barrier
   return
 }
 
 func @omp_master() -> () {

 }
 
 func @omp_terminator() -> () {
   // CHECK: omp.terminator
   omp.terminator
 }
 
 func @omp_parallel(%data_var : memref<i32>, %if_cond : i1, %num_threads : si32) -> () {
-  // CHECK: omp.parallel
+  // CHECK: omp.parallel if(%{{.*}}) num_threads(%{{.*}} : si32) private(%{{.*}} : memref<i32>) firstprivate(%{{.*}} : memref<i32>) shared(%{{.*}} : memref<i32>) copyin(%{{.*}} : memref<i32>)
   "omp.parallel" (%if_cond, %num_threads, %data_var, %data_var, %data_var, %data_var) ({
 
   // test without if condition
-  // CHECK: omp.parallel
+  // CHECK: omp.parallel num_threads(%{{.*}} : si32) private(%{{.*}} : memref<i32>) firstprivate(%{{.*}} : memref<i32>) shared(%{{.*}} : memref<i32>) copyin(%{{.*}} : memref<i32>)
     "omp.parallel"(%num_threads, %data_var, %data_var, %data_var, %data_var) ({
       omp.terminator
     }) {operand_segment_sizes = dense<[0,1,1,1,1,1]>: vector<6xi32>, default_val = "defshared"} : (si32, memref<i32>, memref<i32>, memref<i32>, memref<i32>) -> ()
 
   // CHECK: omp.barrier
     omp.barrier
 
   // test without num_threads
-  // CHECK: omp.parallel
+  // CHECK: omp.parallel if(%{{.*}}) private(%{{.*}} : memref<i32>) firstprivate(%{{.*}} : memref<i32>) shared(%{{.*}} : memref<i32>) copyin(%{{.*}} : memref<i32>)
     "omp.parallel"(%if_cond, %data_var, %data_var, %data_var, %data_var) ({
       omp.terminator
     }) {operand_segment_sizes = dense<[1,0,1,1,1,1]> : vector<6xi32>} : (i1, memref<i32>, memref<i32>, memref<i32>, memref<i32>) -> ()
 
     omp.terminator
   }) {operand_segment_sizes = dense<[1,1,1,1,1,1]> : vector<6xi32>, proc_bind_val = "spread"} : (i1, si32, memref<i32>, memref<i32>, memref<i32>, memref<i32>) -> ()
 
   // test with multiple parameters for single variadic argument
-  // CHECK: omp.parallel
+  // CHECK: omp.parallel private(%{{.*}} : memref<i32>) firstprivate(%{{.*}} : memref<i32>, %{{.*}} : memref<i32>) shared(%{{.*}} : memref<i32>) copyin(%{{.*}} : memref<i32>)
   "omp.parallel" (%data_var, %data_var, %data_var, %data_var, %data_var) ({
     omp.terminator
   }) {operand_segment_sizes = dense<[0,0,1,2,1,1]> : vector<6xi32>} : (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) -> () {
+  // 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 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) {
+      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) 
+               private(%data_var : memref<i32>) proc_bind(close) {
+    omp.terminator
+  }
+
+  return
+}