https://github.com/TIFitis created 
https://github.com/llvm/llvm-project/pull/102341

This patch sets the omp.composite unit attr for composite wrapper ops and also 
add appropriate checks to the verifiers of supported ops for the 
presence/absence of the attribute.

This is patch 2/2 in a series of patches.


>From 9e43d5e608935add740e078e74e1884270775862 Mon Sep 17 00:00:00 2001
From: Akash Banerjee <akash.baner...@amd.com>
Date: Mon, 5 Aug 2024 15:31:25 +0100
Subject: [PATCH 1/3] [MLIR][OpenMP] Add a new ComposableLoopWrapperInterface
 to check and set a composite unitAttr.

---
 mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td |  9 +++--
 .../Dialect/OpenMP/OpenMPOpsInterfaces.td     | 34 +++++++++++++++++++
 2 files changed, 41 insertions(+), 2 deletions(-)

diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td 
b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
index 68f92e6952694b..44210025dfe18b 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
@@ -129,6 +129,7 @@ def PrivateClauseOp : OpenMP_Op<"private", 
[IsolatedFromAbove, RecipeInterface]>
 def ParallelOp : OpenMP_Op<"parallel", traits = [
     AttrSizedOperandSegments, AutomaticAllocationScope,
     DeclareOpInterfaceMethods<LoopWrapperInterface>,
+    DeclareOpInterfaceMethods<ComposableLoopWrapperInterface>,
     DeclareOpInterfaceMethods<OutlineableOpenMPOpInterface>,
     RecursiveMemoryEffects
   ], clauses = [
@@ -357,6 +358,7 @@ def LoopNestOp : OpenMP_Op<"loop_nest", traits = [
 
 def WsloopOp : OpenMP_Op<"wsloop", traits = [
     AttrSizedOperandSegments, DeclareOpInterfaceMethods<LoopWrapperInterface>,
+    DeclareOpInterfaceMethods<ComposableLoopWrapperInterface>,
     RecursiveMemoryEffects, SingleBlock
   ], clauses = [
     OpenMP_AllocateClauseSkip<assemblyFormat = true>,
@@ -433,6 +435,7 @@ def WsloopOp : OpenMP_Op<"wsloop", traits = [
 
 def SimdOp : OpenMP_Op<"simd", traits = [
     AttrSizedOperandSegments, DeclareOpInterfaceMethods<LoopWrapperInterface>,
+    DeclareOpInterfaceMethods<ComposableLoopWrapperInterface>,
     RecursiveMemoryEffects, SingleBlock
   ], clauses = [
     OpenMP_AlignedClause, OpenMP_IfClause, OpenMP_LinearClause,
@@ -500,6 +503,7 @@ def YieldOp : OpenMP_Op<"yield",
 
//===----------------------------------------------------------------------===//
 def DistributeOp : OpenMP_Op<"distribute", traits = [
     AttrSizedOperandSegments, DeclareOpInterfaceMethods<LoopWrapperInterface>,
+    DeclareOpInterfaceMethods<ComposableLoopWrapperInterface>,
     RecursiveMemoryEffects, SingleBlock
   ], clauses = [
     OpenMP_AllocateClause, OpenMP_DistScheduleClause, OpenMP_OrderClause,
@@ -587,8 +591,9 @@ def TaskOp : OpenMP_Op<"task", traits = [
 
 def TaskloopOp : OpenMP_Op<"taskloop", traits = [
     AttrSizedOperandSegments, AutomaticAllocationScope,
-    DeclareOpInterfaceMethods<LoopWrapperInterface>, RecursiveMemoryEffects,
-    SingleBlock
+    DeclareOpInterfaceMethods<LoopWrapperInterface>,
+    DeclareOpInterfaceMethods<ComposableLoopWrapperInterface>,
+    RecursiveMemoryEffects, SingleBlock
   ], clauses = [
     OpenMP_AllocateClause, OpenMP_FinalClause, OpenMP_GrainsizeClause,
     OpenMP_IfClause, OpenMP_InReductionClauseSkip<extraClassDeclaration = 
true>,
diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOpsInterfaces.td 
b/mlir/include/mlir/Dialect/OpenMP/OpenMPOpsInterfaces.td
index 2d1de37239c82a..50394eaf236f94 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOpsInterfaces.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOpsInterfaces.td
@@ -142,6 +142,40 @@ def LoopWrapperInterface : 
OpInterface<"LoopWrapperInterface"> {
   ];
 }
 
+def ComposableLoopWrapperInterface : 
OpInterface<"ComposableLoopWrapperInterface"> {
+  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`.
+  }];
+
+  let cppNamespace = "::mlir::omp";
+
+  let methods = [
+    InterfaceMethod<
+      /*description=*/[{
+        Check whether the operation is composable.
+      }],
+      /*retTy=*/"bool",
+      /*methodName=*/"isComposite",
+      (ins ), [{}], [{
+        return $_op->hasAttr("omp.composite");
+      }]
+    >,
+    InterfaceMethod<
+      /*description=*/[{
+        Set the CompositeAttr for the op.
+      }],
+      /*retTy=*/"void",
+      /*methodName=*/"setComposite",
+      (ins), [{}], [{
+        $_op->setDiscardableAttr("omp.composite", 
mlir::UnitAttr::get($_op->getContext()));
+      }]
+    >
+  ];
+}
+
 def DeclareTargetInterface : OpInterface<"DeclareTargetInterface"> {
   let description = [{
     OpenMP operations that support declare target have this interface.

>From 1ce8a0500e95b0bcd5d2e64b1cbcfe3f216a048f Mon Sep 17 00:00:00 2001
From: Akash Banerjee <akash.baner...@amd.com>
Date: Tue, 6 Aug 2024 13:25:06 +0100
Subject: [PATCH 2/3] Address reviewer comments.

---
 mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td | 19 ++++++++++--------
 .../Dialect/OpenMP/OpenMPOpsInterfaces.td     | 20 ++++++++++---------
 2 files changed, 22 insertions(+), 17 deletions(-)

diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td 
b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
index 44210025dfe18b..d63fdd88f79104 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
@@ -128,8 +128,8 @@ def PrivateClauseOp : OpenMP_Op<"private", 
[IsolatedFromAbove, RecipeInterface]>
 
 def ParallelOp : OpenMP_Op<"parallel", traits = [
     AttrSizedOperandSegments, AutomaticAllocationScope,
+    DeclareOpInterfaceMethods<ComposableOpInterface>,
     DeclareOpInterfaceMethods<LoopWrapperInterface>,
-    DeclareOpInterfaceMethods<ComposableLoopWrapperInterface>,
     DeclareOpInterfaceMethods<OutlineableOpenMPOpInterface>,
     RecursiveMemoryEffects
   ], clauses = [
@@ -357,8 +357,9 @@ def LoopNestOp : OpenMP_Op<"loop_nest", traits = [
 
//===----------------------------------------------------------------------===//
 
 def WsloopOp : OpenMP_Op<"wsloop", traits = [
-    AttrSizedOperandSegments, DeclareOpInterfaceMethods<LoopWrapperInterface>,
-    DeclareOpInterfaceMethods<ComposableLoopWrapperInterface>,
+    AttrSizedOperandSegments,
+    DeclareOpInterfaceMethods<ComposableOpInterface>,
+    DeclareOpInterfaceMethods<LoopWrapperInterface>,
     RecursiveMemoryEffects, SingleBlock
   ], clauses = [
     OpenMP_AllocateClauseSkip<assemblyFormat = true>,
@@ -434,8 +435,9 @@ def WsloopOp : OpenMP_Op<"wsloop", traits = [
 
//===----------------------------------------------------------------------===//
 
 def SimdOp : OpenMP_Op<"simd", traits = [
-    AttrSizedOperandSegments, DeclareOpInterfaceMethods<LoopWrapperInterface>,
-    DeclareOpInterfaceMethods<ComposableLoopWrapperInterface>,
+    AttrSizedOperandSegments,
+    DeclareOpInterfaceMethods<ComposableOpInterface>,
+    DeclareOpInterfaceMethods<LoopWrapperInterface>,
     RecursiveMemoryEffects, SingleBlock
   ], clauses = [
     OpenMP_AlignedClause, OpenMP_IfClause, OpenMP_LinearClause,
@@ -502,8 +504,9 @@ def YieldOp : OpenMP_Op<"yield",
 // Distribute construct [2.9.4.1]
 
//===----------------------------------------------------------------------===//
 def DistributeOp : OpenMP_Op<"distribute", traits = [
-    AttrSizedOperandSegments, DeclareOpInterfaceMethods<LoopWrapperInterface>,
-    DeclareOpInterfaceMethods<ComposableLoopWrapperInterface>,
+    AttrSizedOperandSegments,
+    DeclareOpInterfaceMethods<ComposableOpInterface>,
+    DeclareOpInterfaceMethods<LoopWrapperInterface>,
     RecursiveMemoryEffects, SingleBlock
   ], clauses = [
     OpenMP_AllocateClause, OpenMP_DistScheduleClause, OpenMP_OrderClause,
@@ -591,8 +594,8 @@ def TaskOp : OpenMP_Op<"task", traits = [
 
 def TaskloopOp : OpenMP_Op<"taskloop", traits = [
     AttrSizedOperandSegments, AutomaticAllocationScope,
+    DeclareOpInterfaceMethods<ComposableOpInterface>,
     DeclareOpInterfaceMethods<LoopWrapperInterface>,
-    DeclareOpInterfaceMethods<ComposableLoopWrapperInterface>,
     RecursiveMemoryEffects, SingleBlock
   ], clauses = [
     OpenMP_AllocateClause, OpenMP_FinalClause, OpenMP_GrainsizeClause,
diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOpsInterfaces.td 
b/mlir/include/mlir/Dialect/OpenMP/OpenMPOpsInterfaces.td
index 50394eaf236f94..4cf847f55be215 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOpsInterfaces.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOpsInterfaces.td
@@ -142,12 +142,10 @@ def LoopWrapperInterface : 
OpInterface<"LoopWrapperInterface"> {
   ];
 }
 
-def ComposableLoopWrapperInterface : 
OpInterface<"ComposableLoopWrapperInterface"> {
+def ComposableOpInterface : OpInterface<"ComposableOpInterface"> {
   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 can represent a single leaf of a composite OpenMP
+    construct.
   }];
 
   let cppNamespace = "::mlir::omp";
@@ -155,7 +153,8 @@ def ComposableLoopWrapperInterface : 
OpInterface<"ComposableLoopWrapperInterface
   let methods = [
     InterfaceMethod<
       /*description=*/[{
-        Check whether the operation is composable.
+        Check whether the operation is representing a leaf of a composite 
OpenMP
+        construct.
       }],
       /*retTy=*/"bool",
       /*methodName=*/"isComposite",
@@ -165,12 +164,15 @@ def ComposableLoopWrapperInterface : 
OpInterface<"ComposableLoopWrapperInterface
     >,
     InterfaceMethod<
       /*description=*/[{
-        Set the CompositeAttr for the op.
+        Mark the operation as part of an OpenMP composite construct.
       }],
       /*retTy=*/"void",
       /*methodName=*/"setComposite",
-      (ins), [{}], [{
-        $_op->setDiscardableAttr("omp.composite", 
mlir::UnitAttr::get($_op->getContext()));
+      (ins "bool":$val), [{}], [{
+        if(val)
+          $_op->setDiscardableAttr("omp.composite", 
mlir::UnitAttr::get($_op->getContext()));
+        else
+          $_op->removeDiscardableAttr("omp.composite");
       }]
     >
   ];

>From 8d656e464eea5b4c65836589c4e785b10490ce2c Mon Sep 17 00:00:00 2001
From: Akash Banerjee <akash.baner...@amd.com>
Date: Wed, 7 Aug 2024 18:31:08 +0100
Subject: [PATCH 3/3] [OpenMP][MLIR] Set omp.composite attr for composite loop
 wrappers and add verifier checks

This patch sets the omp.composite unit attr for composite wrapper ops and also 
add appropriate checks to the verifiers of supported ops for the 
presence/absence of the attribute.
---
 flang/lib/Lower/OpenMP/OpenMP.cpp            |  8 +++++
 mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp | 32 ++++++++++++++++++++
 2 files changed, 40 insertions(+)

diff --git a/flang/lib/Lower/OpenMP/OpenMP.cpp 
b/flang/lib/Lower/OpenMP/OpenMP.cpp
index bbde77c14f36a1..3b18e7b3ecf80e 100644
--- a/flang/lib/Lower/OpenMP/OpenMP.cpp
+++ b/flang/lib/Lower/OpenMP/OpenMP.cpp
@@ -2063,10 +2063,14 @@ static void genCompositeDistributeSimd(
   // TODO: Populate entry block arguments with private variables.
   auto distributeOp = genWrapperOp<mlir::omp::DistributeOp>(
       converter, loc, distributeClauseOps, /*blockArgTypes=*/{});
+  llvm::cast<mlir::omp::ComposableOpInterface>(distributeOp.getOperation())
+      .setComposite(/*val=*/true);
 
   // TODO: Populate entry block arguments with reduction and private variables.
   auto simdOp = genWrapperOp<mlir::omp::SimdOp>(converter, loc, simdClauseOps,
                                                 /*blockArgTypes=*/{});
+  llvm::cast<mlir::omp::ComposableOpInterface>(simdOp.getOperation())
+      .setComposite(/*val=*/true);
 
   // Construct wrapper entry block list and associated symbols. It is important
   // that the symbol order and the block argument order match, so that the
@@ -2111,10 +2115,14 @@ static void genCompositeDoSimd(lower::AbstractConverter 
&converter,
   // TODO: Add private variables to entry block arguments.
   auto wsloopOp = genWrapperOp<mlir::omp::WsloopOp>(
       converter, loc, wsloopClauseOps, wsloopReductionTypes);
+  llvm::cast<mlir::omp::ComposableOpInterface>(wsloopOp.getOperation())
+      .setComposite(/*val=*/true);
 
   // TODO: Populate entry block arguments with reduction and private variables.
   auto simdOp = genWrapperOp<mlir::omp::SimdOp>(converter, loc, simdClauseOps,
                                                 /*blockArgTypes=*/{});
+  llvm::cast<mlir::omp::ComposableOpInterface>(simdOp.getOperation())
+      .setComposite(/*val=*/true);
 
   // Construct wrapper entry block list and associated symbols. It is important
   // that the symbol and block argument order match, so that the symbol-value
diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp 
b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
index 11780f84697b15..641fbb5a1418f6 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -1546,6 +1546,9 @@ LogicalResult ParallelOp::verify() {
     if (!isWrapper())
       return emitOpError() << "must take a loop wrapper role if nested inside "
                               "of 'omp.distribute'";
+    if (!llvm::cast<ComposableOpInterface>(getOperation()).isComposite())
+      return emitError()
+             << "'omp.composite' attribute missing from composite wrapper";
 
     if (LoopWrapperInterface nested = getNestedWrapper()) {
       // Check for the allowed leaf constructs that may appear in a composite
@@ -1555,6 +1558,9 @@ LogicalResult ParallelOp::verify() {
     } else {
       return emitOpError() << "must not wrap an 'omp.loop_nest' directly";
     }
+  } else if (llvm::cast<ComposableOpInterface>(getOperation()).isComposite()) {
+    return emitError()
+           << "'omp.composite' attribute present in non-composite wrapper";
   }
 
   if (getAllocateVars().size() != getAllocatorVars().size())
@@ -1749,10 +1755,18 @@ LogicalResult WsloopOp::verify() {
     return emitOpError() << "must be a loop wrapper";
 
   if (LoopWrapperInterface nested = getNestedWrapper()) {
+    if (!llvm::cast<ComposableOpInterface>(getOperation()).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 DO/FOR.
     if (!isa<SimdOp>(nested))
       return emitError() << "only supported nested wrapper is 'omp.simd'";
+
+  } else if (llvm::cast<ComposableOpInterface>(getOperation()).isComposite()) {
+    return emitError()
+           << "'omp.composite' attribute present in non-composite wrapper";
   }
 
   return verifyReductionVarList(*this, getReductionSyms(), getReductionVars(),
@@ -1796,6 +1810,10 @@ LogicalResult SimdOp::verify() {
   if (getNestedWrapper())
     return emitOpError() << "must wrap an 'omp.loop_nest' directly";
 
+  if (llvm::cast<ComposableOpInterface>(getOperation()).isComposite())
+    return emitError()
+           << "'omp.composite' attribute present in non-composite wrapper";
+
   return success();
 }
 
@@ -1825,11 +1843,17 @@ LogicalResult DistributeOp::verify() {
     return emitOpError() << "must be a loop wrapper";
 
   if (LoopWrapperInterface nested = getNestedWrapper()) {
+    if (!llvm::cast<ComposableOpInterface>(getOperation()).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'";
+  } else if (llvm::cast<ComposableOpInterface>(getOperation()).isComposite()) {
+    return emitError()
+           << "'omp.composite' attribute present in non-composite wrapper";
   }
 
   return success();
@@ -2031,11 +2055,19 @@ LogicalResult TaskloopOp::verify() {
     return emitOpError() << "must be a loop wrapper";
 
   if (LoopWrapperInterface nested = getNestedWrapper()) {
+    if (!llvm::cast<ComposableOpInterface>(getOperation()).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 TASKLOOP.
     if (!isa<SimdOp>(nested))
       return emitError() << "only supported nested wrapper is 'omp.simd'";
+  } else if (llvm::cast<ComposableOpInterface>(getOperation()).isComposite()) {
+    return emitError()
+           << "'omp.composite' attribute present in non-composite wrapper";
   }
+
   return success();
 }
 

_______________________________________________
llvm-branch-commits mailing list
llvm-branch-commits@lists.llvm.org
https://lists.llvm.org/cgi-bin/mailman/listinfo/llvm-branch-commits

Reply via email to