[MLIR][Flang][OpenMP] Remove omp.parallel from loop wrapper ops

This patch updates the `omp.parallel` operation according to the results of
the discussion in [this RFC](https://discourse.llvm.org/t/rfc-disambiguation-between-loop-and-block-associated-omp-parallelop/79972).
It is removed from the set of loop wrapper operations, changing the expected
MLIR representation for composite `distribute parallel do/for` into the
following:

```mlir
omp.parallel {
  ...
  omp.distribute {
    omp.wsloop {
      omp.loop_nest ... { ... }
      omp.terminator
    }
    omp.terminator
  }
  ...
  omp.terminator
}
```

MLIR verifiers for operations impacted by this representation change are
updated, as well as related tests. The `LoopWrapperInterface` is also updated,
since it's no longer representing an optional "role" of an operation but a
mandatory set of restrictions instead.

Author: skatrak

flang/lib/Lower/OpenMP/DataSharingProcessor.cpp

@@ -231,9 +231,7 @@ void DataSharingProcessor::insertBarrier() {
 void DataSharingProcessor::insertLastPrivateCompare(mlir::Operation *op) {
   mlir::omp::LoopNestOp loopOp;
   if (auto wrapper = mlir::dyn_cast<mlir::omp::LoopWrapperInterface>(op))
-    loopOp = wrapper.isWrapper()
-                 ? mlir::cast<mlir::omp::LoopNestOp>(wrapper.getWrappedLoop())
-                 : nullptr;
+    loopOp = mlir::cast<mlir::omp::LoopNestOp>(wrapper.getWrappedLoop());
 
   bool cmpCreated = false;
   mlir::OpBuilder::InsertionGuard guard(firOpBuilder);

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

@@ -129,7 +129,6 @@ def PrivateClauseOp : OpenMP_Op<"private", [IsolatedFromAbove, RecipeInterface]>
 def ParallelOp : OpenMP_Op<"parallel", traits = [
     AttrSizedOperandSegments, AutomaticAllocationScope,
     DeclareOpInterfaceMethods<ComposableOpInterface>,
-    DeclareOpInterfaceMethods<LoopWrapperInterface>,
     DeclareOpInterfaceMethods<OutlineableOpenMPOpInterface>,
     RecursiveMemoryEffects
   ], clauses = [

mlir/include/mlir/Dialect/OpenMP/OpenMPOpsInterfaces.td

@@ -71,24 +71,23 @@ def ReductionClauseInterface : OpInterface<"ReductionClauseInterface"> {
 
 def LoopWrapperInterface : OpInterface<"LoopWrapperInterface"> {
   let description = [{
-    OpenMP operations that can wrap a single loop nest. When taking a wrapper
-    role, these operations must only contain a single region with a single block
-    in which there's a single operation and a terminator. That nested operation
-    must be another loop wrapper or an `omp.loop_nest`.
+    OpenMP operations that wrap a single loop nest. They must only contain a
+    single region with a single block in which there's a single operation and a
+    terminator. That nested operation must be another loop wrapper or an
+    `omp.loop_nest`.
   }];
 
   let cppNamespace = "::mlir::omp";
 
   let methods = [
     InterfaceMethod<
       /*description=*/[{
-        Tell whether the operation could be taking the role of a loop wrapper.
-        That is, it has a single region with a single block in which there are
-        two operations: another wrapper (also taking a loop wrapper role) or
-        `omp.loop_nest` operation and a terminator.
+        Check whether the operation is a valid loop wrapper. That is, it has a
+        single region with a single block in which there are two operations:
+        another loop wrapper or `omp.loop_nest` operation and a terminator.
       }],
       /*retTy=*/"bool",
-      /*methodName=*/"isWrapper",
+      /*methodName=*/"isValidWrapper",
       (ins ), [{}], [{
         if ($_op->getNumRegions() != 1)
           return false;
@@ -106,34 +105,31 @@ def LoopWrapperInterface : OpInterface<"LoopWrapperInterface"> {
         if (!secondOp.hasTrait<OpTrait::IsTerminator>())
           return false;
 
-        if (auto wrapper = ::llvm::dyn_cast<LoopWrapperInterface>(firstOp))
-          return wrapper.isWrapper();
-
-        return ::llvm::isa<LoopNestOp>(firstOp);
+        return ::llvm::isa<LoopNestOp, LoopWrapperInterface>(firstOp);
       }]
     >,
     InterfaceMethod<
       /*description=*/[{
         If there is another loop wrapper immediately nested inside, return that
-        operation. Assumes this operation is taking a loop wrapper role.
+        operation. Assumes this operation is a valid loop wrapper.
       }],
       /*retTy=*/"::mlir::omp::LoopWrapperInterface",
       /*methodName=*/"getNestedWrapper",
       (ins), [{}], [{
-        assert($_op.isWrapper() && "Unexpected non-wrapper op");
+        assert($_op.isValidWrapper() && "Unexpected non-wrapper op");
         Operation *nested = &*$_op->getRegion(0).op_begin();
         return ::llvm::dyn_cast<LoopWrapperInterface>(nested);
       }]
     >,
     InterfaceMethod<
       /*description=*/[{
         Return the loop nest nested directly or indirectly inside of this loop
-        wrapper. Assumes this operation is taking a loop wrapper role.
+        wrapper. Assumes this operation is a valid loop wrapper.
       }],
       /*retTy=*/"::mlir::Operation *",
       /*methodName=*/"getWrappedLoop",
       (ins), [{}], [{
-        assert($_op.isWrapper() && "Unexpected non-wrapper op");
+        assert($_op.isValidWrapper() && "Unexpected non-wrapper op");
         if (LoopWrapperInterface nested = $_op.getNestedWrapper())
           return nested.getWrappedLoop();
         return &*$_op->getRegion(0).op_begin();

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

@@ -1541,26 +1541,25 @@ static LogicalResult verifyPrivateVarList(OpType &op) {
 }
 
 LogicalResult ParallelOp::verify() {
-  // Check that it is a valid loop wrapper if it's taking that role.
-  if (isa<DistributeOp>((*this)->getParentOp())) {
-    if (!isWrapper())
-      return emitOpError() << "must take a loop wrapper role if nested inside "
-                              "of 'omp.distribute'";
+  auto distributeChildOps = getOps<DistributeOp>();
+  if (!distributeChildOps.empty()) {
     if (!isComposite())
       return emitError()
-             << "'omp.composite' attribute missing from composite wrapper";
+             << "'omp.composite' attribute missing from composite operation";
 
-    if (LoopWrapperInterface nested = getNestedWrapper()) {
-      // Check for the allowed leaf constructs that may appear in a composite
-      // construct directly after PARALLEL.
-      if (!isa<WsloopOp>(nested))
-        return emitError() << "only supported nested wrapper is 'omp.wsloop'";
-    } else {
-      return emitOpError() << "must not wrap an 'omp.loop_nest' directly";
+    auto *ompDialect = getContext()->getLoadedDialect<OpenMPDialect>();
+    Operation &distributeOp = **distributeChildOps.begin();
+    for (Operation &childOp : getOps()) {
+      if (&childOp == &distributeOp || ompDialect != childOp.getDialect())
+        continue;
+
+      if (!childOp.hasTrait<OpTrait::IsTerminator>())
+        return emitError() << "unexpected OpenMP operation inside of composite "
+                              "'omp.parallel'";
     }
   } else if (isComposite()) {
     return emitError()
-           << "'omp.composite' attribute present in non-composite wrapper";
+           << "'omp.composite' attribute present in non-composite operation";
   }
 
   if (getAllocateVars().size() != getAllocatorVars().size())
@@ -1751,15 +1750,12 @@ void WsloopOp::build(OpBuilder &builder, OperationState &state,
 }
 
 LogicalResult WsloopOp::verify() {
-  if (!isWrapper())
-    return emitOpError() << "must be a loop wrapper";
+  if (!isValidWrapper())
+    return emitOpError() << "must be a valid loop wrapper";
 
-  auto wrapper =
-      llvm::dyn_cast_if_present<LoopWrapperInterface>((*this)->getParentOp());
   bool isCompositeChildLeaf =
-      wrapper && wrapper.isWrapper() &&
-      (!llvm::isa<ParallelOp>(wrapper) ||
-       llvm::isa_and_present<DistributeOp>(wrapper->getParentOp()));
+      llvm::dyn_cast_if_present<LoopWrapperInterface>((*this)->getParentOp());
+
   if (LoopWrapperInterface nested = getNestedWrapper()) {
     if (!isComposite())
       return emitError()
@@ -1813,18 +1809,14 @@ LogicalResult SimdOp::verify() {
   if (verifyNontemporalClause(*this, getNontemporalVars()).failed())
     return failure();
 
-  if (!isWrapper())
-    return emitOpError() << "must be a loop wrapper";
+  if (!isValidWrapper())
+    return emitOpError() << "must be a valid loop wrapper";
 
   if (getNestedWrapper())
     return emitOpError() << "must wrap an 'omp.loop_nest' directly";
 
-  auto wrapper =
-      llvm::dyn_cast_if_present<LoopWrapperInterface>((*this)->getParentOp());
   bool isCompositeChildLeaf =
-      wrapper && wrapper.isWrapper() &&
-      (!llvm::isa<ParallelOp>(wrapper) ||
-       llvm::isa_and_present<DistributeOp>(wrapper->getParentOp()));
+      llvm::dyn_cast_if_present<LoopWrapperInterface>((*this)->getParentOp());
 
   if (!isComposite() && isCompositeChildLeaf)
     return emitError()
@@ -1859,18 +1851,22 @@ LogicalResult DistributeOp::verify() {
     return emitError(
         "expected equal sizes for allocate and allocator variables");
 
-  if (!isWrapper())
-    return emitOpError() << "must be a loop wrapper";
+  if (!isValidWrapper())
+    return emitOpError() << "must be a valid loop wrapper";
 
   if (LoopWrapperInterface nested = getNestedWrapper()) {
     if (!isComposite())
       return emitError()
              << "'omp.composite' attribute missing from composite wrapper";
     // Check for the allowed leaf constructs that may appear in a composite
     // construct directly after DISTRIBUTE.
-    if (!isa<ParallelOp, SimdOp>(nested))
-      return emitError() << "only supported nested wrappers are 'omp.parallel' "
-                            "and 'omp.simd'";
+    if (isa<WsloopOp>(nested)) {
+      if (!llvm::dyn_cast_if_present<ParallelOp>((*this)->getParentOp()))
+        return emitError() << "an 'omp.wsloop' nested wrapper is only allowed "
+                              "when 'omp.parallel' is the direct parent";
+    } else if (!isa<SimdOp>(nested))
+      return emitError() << "only supported nested wrappers are 'omp.simd' and "
+                            "'omp.wsloop'";
   } else if (isComposite()) {
     return emitError()
            << "'omp.composite' attribute present in non-composite wrapper";
@@ -2063,8 +2059,8 @@ LogicalResult TaskloopOp::verify() {
         "may not appear on the same taskloop directive");
   }
 
-  if (!isWrapper())
-    return emitOpError() << "must be a loop wrapper";
+  if (!isValidWrapper())
+    return emitOpError() << "must be a valid loop wrapper";
 
   if (LoopWrapperInterface nested = getNestedWrapper()) {
     if (!isComposite())
@@ -2161,11 +2157,8 @@ LogicalResult LoopNestOp::verify() {
              << "range argument type does not match corresponding IV type";
   }
 
-  auto wrapper =
-      llvm::dyn_cast_if_present<LoopWrapperInterface>((*this)->getParentOp());
-
-  if (!wrapper || !wrapper.isWrapper())
-    return emitOpError() << "expects parent op to be a valid loop wrapper";
+  if (!llvm::dyn_cast_if_present<LoopWrapperInterface>((*this)->getParentOp()))
+    return emitOpError() << "expects parent op to be a loop wrapper";
 
   return success();
 }
@@ -2175,8 +2168,6 @@ void LoopNestOp::gatherWrappers(
   Operation *parent = (*this)->getParentOp();
   while (auto wrapper =
              llvm::dyn_cast_if_present<LoopWrapperInterface>(parent)) {
-    if (!wrapper.isWrapper())
-      break;
     wrappers.push_back(wrapper);
     parent = parent->getParentOp();
   }