[MLIR][OpenMP]Refactored parser and printer function of grainsize and numtasks clause.

Author: kaviya2510

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

@@ -472,55 +472,76 @@ static void printOrderClause(OpAsmPrinter &p, Operation *op,
     p << stringifyClauseOrderKind(order.getValue());
 }
 
-//===----------------------------------------------------------------------===//
-// Parser and printer for grainsize Clause
-//===----------------------------------------------------------------------===//
-
-// grainsize ::= `grainsize` `(` [strict ':'] grain-size `)`
+template <typename ClauseTypeAttr, typename ClauseType>
 static ParseResult
-parseGrainsizeClause(OpAsmParser &parser, ClauseGrainsizeTypeAttr &grainsizeMod,
-                     std::optional<OpAsmParser::UnresolvedOperand> &grainsize,
-                     Type &grainsizeType) {
-  SMLoc loc = parser.getCurrentLocation();
+parseGranularityClause(OpAsmParser &parser, ClauseTypeAttr &prescriptiveness,
+                       std::optional<OpAsmParser::UnresolvedOperand> &operand,
+                       Type &operandType,
+                       std::optional<ClauseType> (*symbolizeClause)(StringRef),
+                       StringRef clauseName) {
   StringRef enumStr;
-
   if (succeeded(parser.parseOptionalKeyword(&enumStr))) {
-    if (std::optional<ClauseGrainsizeType> enumValue =
-            symbolizeClauseGrainsizeType(enumStr)) {
-      grainsizeMod =
-          ClauseGrainsizeTypeAttr::get(parser.getContext(), *enumValue);
+    if (std::optional<ClauseType> enumValue = symbolizeClause(enumStr)) {
+      prescriptiveness = ClauseTypeAttr::get(parser.getContext(), *enumValue);
       if (parser.parseColon())
         return failure();
     } else {
-      return parser.emitError(loc, "invalid grainsize modifier : '")
-             << enumStr << "'";
+      return parser.emitError(parser.getCurrentLocation())
+             << "invalid " << clauseName << " modifier : '" << enumStr << "'";
+      ;
     }
   }
 
-  OpAsmParser::UnresolvedOperand operand;
-  if (succeeded(parser.parseOperand(operand))) {
-    grainsize = operand;
+  OpAsmParser::UnresolvedOperand var;
+  if (succeeded(parser.parseOperand(var))) {
+    operand = var;
   } else {
     return parser.emitError(parser.getCurrentLocation())
-           << "expected grainsize operand";
+           << "expected " << clauseName << " operand";
   }
 
-  if (grainsize.has_value()) {
-    if (parser.parseColonType(grainsizeType))
+  if (operand.has_value()) {
+    if (parser.parseColonType(operandType))
       return failure();
   }
 
   return success();
 }
 
+template <typename ClauseTypeAttr, typename ClauseType>
+static void
+printGranularityClause(OpAsmPrinter &p, Operation *op,
+                       ClauseTypeAttr prescriptiveness, Value operand,
+                       mlir::Type operandType,
+                       StringRef (*stringifyClauseType)(ClauseType)) {
+
+  if (prescriptiveness)
+    p << stringifyClauseType(prescriptiveness.getValue()) << ": ";
+
+  if (operand)
+    p << operand << ": " << operandType;
+}
+
+//===----------------------------------------------------------------------===//
+// Parser and printer for grainsize Clause
+//===----------------------------------------------------------------------===//
+
+// grainsize ::= `grainsize` `(` [strict ':'] grain-size `)`
+static ParseResult
+parseGrainsizeClause(OpAsmParser &parser, ClauseGrainsizeTypeAttr &grainsizeMod,
+                     std::optional<OpAsmParser::UnresolvedOperand> &grainsize,
+                     Type &grainsizeType) {
+  return parseGranularityClause<ClauseGrainsizeTypeAttr, ClauseGrainsizeType>(
+      parser, grainsizeMod, grainsize, grainsizeType,
+      &symbolizeClauseGrainsizeType, "grainsize");
+}
+
 static void printGrainsizeClause(OpAsmPrinter &p, Operation *op,
                                  ClauseGrainsizeTypeAttr grainsizeMod,
                                  Value grainsize, mlir::Type grainsizeType) {
-  if (grainsizeMod)
-    p << stringifyClauseGrainsizeType(grainsizeMod.getValue()) << ": ";
-
-  if (grainsize)
-    p << grainsize << ": " << grainsizeType;
+  printGranularityClause<ClauseGrainsizeTypeAttr, ClauseGrainsizeType>(
+      p, op, grainsizeMod, grainsize, grainsizeType,
+      &stringifyClauseGrainsizeType);
 }
 
 //===----------------------------------------------------------------------===//
@@ -532,46 +553,16 @@ static ParseResult
 parseNumTasksClause(OpAsmParser &parser, ClauseNumTasksTypeAttr &numTasksMod,
                     std::optional<OpAsmParser::UnresolvedOperand> &numTasks,
                     Type &numTasksType) {
-  SMLoc loc = parser.getCurrentLocation();
-  StringRef enumStr;
-
-  if (succeeded(parser.parseOptionalKeyword(&enumStr))) {
-    if (std::optional<ClauseNumTasksType> enumValue =
-            symbolizeClauseNumTasksType(enumStr)) {
-      numTasksMod =
-          ClauseNumTasksTypeAttr::get(parser.getContext(), *enumValue);
-      if (parser.parseColon())
-        return failure();
-    } else {
-      return parser.emitError(loc, "invalid numTasks modifier : '")
-             << enumStr << "'";
-    }
-  }
-
-  OpAsmParser::UnresolvedOperand operand;
-  if (succeeded(parser.parseOperand(operand))) {
-    numTasks = operand;
-  } else {
-    return parser.emitError(parser.getCurrentLocation())
-           << "expected num_tasks operand";
-  }
-
-  if (numTasks.has_value()) {
-    if (parser.parseColonType(numTasksType))
-      return failure();
-  }
-
-  return success();
+  return parseGranularityClause<ClauseNumTasksTypeAttr, ClauseNumTasksType>(
+      parser, numTasksMod, numTasks, numTasksType, &symbolizeClauseNumTasksType,
+      "num_tasks");
 }
 
 static void printNumTasksClause(OpAsmPrinter &p, Operation *op,
                                 ClauseNumTasksTypeAttr numTasksMod,
                                 Value numTasks, mlir::Type numTasksType) {
-  if (numTasksMod)
-    p << stringifyClauseNumTasksType(numTasksMod.getValue()) << ": ";
-
-  if (numTasks)
-    p << numTasks << ": " << numTasksType;
+  printGranularityClause<ClauseNumTasksTypeAttr, ClauseNumTasksType>(
+      p, op, numTasksMod, numTasks, numTasksType, &stringifyClauseNumTasksType);
 }
 
 //===----------------------------------------------------------------------===//

mlir/test/Dialect/OpenMP/invalid.mlir

@@ -2078,7 +2078,7 @@ func.func @taskloop(%lb: i32, %ub: i32, %step: i32) {
 
 func.func @taskloop(%lb: i32, %ub: i32, %step: i32) {
   %testi64 = "test.i64"() : () -> (i64)
-  // expected-error @below {{invalid numTasks modifier : 'default'}}
+  // expected-error @below {{invalid num_tasks modifier : 'default'}}
   omp.taskloop num_tasks(default: %testi64: i64) {
     omp.loop_nest (%i, %j) : i32 = (%lb, %ub) to (%ub, %lb) step (%step, %step) {
       omp.yield