[OpenMP] Add verifier and tests for workdistribute mlir op.

Author: skc7

flang/lib/Lower/OpenMP/OpenMP.cpp

@@ -536,10 +536,10 @@ static void processHostEvalClauses(lower::AbstractConverter &converter,
       break;
 
     case OMPD_teams_workdistribute:
-      cp.processThreadLimit(stmtCtx, hostInfo.ops);
+      cp.processThreadLimit(stmtCtx, hostInfo->ops);
       [[fallthrough]];
     case OMPD_target_teams_workdistribute:
-      cp.processNumTeams(stmtCtx, hostInfo.ops);
+      cp.processNumTeams(stmtCtx, hostInfo->ops);
       processSingleNestedIf([](Directive nestedDir) {
         return topDistributeSet.test(nestedDir) || topLoopSet.test(nestedDir);
       });

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

@@ -2132,6 +2132,7 @@ def WorkdistributeOp : OpenMP_Op<"workdistribute"> {
     ```
   }];
   let regions = (region AnyRegion:$region);
+  let hasVerifier = 1;
   let assemblyFormat = "$region attr-dict";
 }
 

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

@@ -3875,6 +3875,21 @@ LogicalResult AllocateDirOp::verify() {
   return success();
 }
 
+//===----------------------------------------------------------------------===//
+// WorkdistributeOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult WorkdistributeOp::verify() {
+  Region &region = getRegion();
+  if (!region.hasOneBlock())
+    return emitOpError("region must contain exactly one block");
+
+  Operation *parentOp = (*this)->getParentOp();
+  if (!llvm::dyn_cast<TeamsOp>(parentOp))
+    return emitOpError("workdistribute must be nested under teams");
+  return success();
+}
+
 #define GET_ATTRDEF_CLASSES
 #include "mlir/Dialect/OpenMP/OpenMPOpsAttributes.cpp.inc"
 

mlir/test/Dialect/OpenMP/invalid.mlir

@@ -3017,3 +3017,26 @@ func.func @invalid_allocate_allocator(%arg0 : memref<i32>) -> () {
 
   return
 }
+
+// -----
+func.func @invalid_workdistribute() -> () {
+// expected-error @below {{workdistribute must be nested under teams}}
+  omp.workdistribute {
+    omp.terminator
+  }
+  return
+}
+
+// -----
+func.func @invalid_workdistribute_with_multiple_blocks() -> () {
+  omp.teams {
+  // expected-error @below {{region must contain exactly one block}}
+  omp.workdistribute {
+    cf.br ^bb1
+  ^bb1:
+    omp.terminator
+  }
+  omp.terminator
+  }
+  return
+}

mlir/test/Dialect/OpenMP/ops.mlir

@@ -3230,3 +3230,15 @@ func.func @omp_allocate_dir(%arg0 : memref<i32>, %arg1 : memref<i32>) -> () {
   return
 }
 
+// CHECK-LABEL: func.func @omp_workdistribute
+func.func @omp_workdistribute() {
+  // CHECK: omp.teams
+  omp.teams {
+  // CHECK: omp.workdistribute
+  omp.workdistribute {
+    omp.terminator
+  }
+  omp.terminator
+  }
+  return
+}