[OpenMP][Flang][MLIR] Add MLIR support for OpenMP requires directive

[skatrak]

Summary

This patch introduces an MLIR attribute to the OpenMP dialect
representing the clauses that a 'requires' directive can define, with
the exception of the parameterized atomic_default_mem_order
clause. In that case, the existing omp.memory_order enumeration
can be used for representing its information instead.

The OffloadModuleInterface is also updated to provide methods to get
and set two new dialect attributes (omp.requires,
omp.atomic_default_mem_order) to store and use this information during
the lowering stages to LLVM IR.

[skatrak]

Rebase and update.

[skatrak]

Rebase

[skatrak]

Remove omp.atomic_default_mem_order. This clause is now handled in semantics.

[skatrak]

Rebase.

[TIFitis]

Please add tests to mlir/test/Dialect/OpenMP/ops.mlir

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

1633

Add a TODO for unsupported clause

[skatrak]

Add unit test and comment.

[skatrak]

Thank you @TIFitis for the comments. Let me know if these changes address your concerns.

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

1633

I added a comment explaining that there's no need to add flags for the atomic_default_mem_order clause here, since the OMPC_MemoryOrder enum in OMP.td already defines the values it can have.

The currently proposed implementation in Flang, which can be seen in dependent patches to this one, doesn't even store the atomic default memory order clause in MLIR because it's handled early in a semantics pass. Even if later some other frontend needed this information in MLIR, an attribute with an OMPC_MemoryOrder field could be used to represent it.

[TIFitis]

LGTM

[skatrak]

Thank you, Akash. I'll wait for one more day before landing this patch in case somebody else has any further comments.

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

def ReductionOp : OpenMP_Op<"reduction"> {
 
   let arguments= (ins AnyType:$operand, OpenMP_PointerLikeType:$accumulator);
   let assemblyFormat = [{
     $operand `,` $accumulator attr-dict `:` type($operand) `,` type($accumulator)
   }];
   let hasVerifier = 1;
 }
 
+//===----------------------------------------------------------------------===//
+// 8.2 requires directive
+//===----------------------------------------------------------------------===//
+
+def ClauseRequiresNone : I32BitEnumAttrCaseNone<"none">;
+def ClauseRequiresReverseOffload : I32BitEnumAttrCaseBit<"reverse_offload", 0>;
+def ClauseRequiresUnifiedAddress : I32BitEnumAttrCaseBit<"unified_address", 1>;
+def ClauseRequiresUnifiedSharedMemory
+    : I32BitEnumAttrCaseBit<"unified_shared_memory", 2>;
+def ClauseRequiresDynamicAllocators
+    : I32BitEnumAttrCaseBit<"dynamic_allocators", 3>;
+
+def ClauseRequires : I32BitEnumAttr<
+    "ClauseRequires",
+    "requires clause",
+    [
+      ClauseRequiresNone,
+      ClauseRequiresReverseOffload,
+      ClauseRequiresUnifiedAddress,
+      ClauseRequiresUnifiedSharedMemory,
+      ClauseRequiresDynamicAllocators
+    ]> {
+  let genSpecializedAttr = 0;
+  let cppNamespace = "::mlir::omp";
+}
+def ClauseRequiresAttr :
+  EnumAttr<OpenMP_Dialect, ClauseRequires, "clause_requires"> {
+}
+
 #endif // OPENMP_OPS

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

InterfaceMethod<
       "getAllReductionVars", (ins), [{}], [{
         return $_op.getReductionVars();
       }]>,
   ];
 }
 
 def OffloadModuleInterface : OpInterface<"OffloadModuleInterface"> {
   let description = [{
-    Operations that represent a module for offloading (host or device) 
+    Operations that represent a module for offloading (host or device)
     should have this interface.
   }];
 
   let cppNamespace = "::mlir::omp";
 
   let methods = [
     InterfaceMethod<
       /*description=*/[{
-        Set the attribute IsDeviceAttr on the current module with the 
-        specified boolean argument. 
+        Set the attribute IsDeviceAttr on the current module with the
+        specified boolean argument.
       }],
       /*retTy=*/"void",
-      /*methodName=*/"setIsDevice", 
+      /*methodName=*/"setIsDevice",
       (ins "bool":$isDevice), [{}], [{
         $_op->setAttr(
           mlir::StringAttr::get($_op->getContext(), llvm::Twine{"omp.is_device"}),
             mlir::omp::IsDeviceAttr::get($_op->getContext(), isDevice));
       }]>,
-      InterfaceMethod<
       /*description=*/[{
         Get the IsDeviceAttr attribute on the current module if it exists and return
         its value, if it doesn't exit it returns false by default.
       }],
       /*retTy=*/"bool",
-      /*methodName=*/"getIsDevice", 
+      /*methodName=*/"getIsDevice",
       (ins), [{}], [{
         if (Attribute isDevice = $_op->getAttr("omp.is_device"))
           if (isDevice.isa<mlir::omp::IsDeviceAttr>())
             return isDevice.dyn_cast<IsDeviceAttr>().getIsDevice();
         return false;
       }]>,
       InterfaceMethod<
       /*description=*/[{
InterfaceMethod<
             "bool":$assumeThreadsOversubscription,
             "bool":$assumeNoThreadState,
             "bool":$assumeNoNestedParallelism), [{}], [{
         $_op->setAttr(("omp." + mlir::omp::FlagsAttr::getMnemonic()).str(),
                   mlir::omp::FlagsAttr::get($_op->getContext(), debugKind,
                       assumeTeamsOversubscription, assumeThreadsOversubscription, 
                       assumeNoThreadState, assumeNoNestedParallelism));
       }]>,
+    InterfaceMethod<
+      /*description=*/[{
+        Get the omp.requires attribute on the operator if it's present and
+        return its value. If it doesn't exist, return `ClauseRequires::none` by
+        default.
+      }],
+      /*retTy=*/"::mlir::omp::ClauseRequires",
+      /*methodName=*/"getRequires",
+      (ins), [{}], [{
+        if (Attribute requiresAttr = $_op->getAttr("omp.requires"))
+          if (auto requiresVal = requiresAttr.dyn_cast<mlir::omp::ClauseRequiresAttr>())
+            return requiresVal.getValue();
+        return mlir::omp::ClauseRequires::none;
+      }]>,
+    InterfaceMethod<
+      /*description=*/[{
+        Set the omp.requires attribute on the operator to the specified clauses.
+      }],
+      /*retTy=*/"void",
+      /*methodName=*/"setRequires",
+      (ins "::mlir::omp::ClauseRequires":$clauses), [{}], [{
+        $_op->setAttr(mlir::StringAttr::get($_op->getContext(), "omp.requires"),
+          mlir::omp::ClauseRequiresAttr::get($_op->getContext(), clauses));
+      }]>,
+    InterfaceMethod<
+      /*description=*/[{
+        Get the omp.atomic_default_mem_order attribute on the operator and
+        return its value if it's present.
+      }],
+      /*retTy=*/"std::optional<::mlir::omp::ClauseMemoryOrderKind>",
+      /*methodName=*/"getAtomicDefaultMemOrder",
+      (ins), [{}], [{
+        if (mlir::Attribute atomicDefaultAttr = $_op->getAttr("omp.atomic_default_mem_order"))
+          if (auto atomicDefault = atomicDefaultAttr.dyn_cast<mlir::omp::ClauseMemoryOrderKindAttr>())
+            return atomicDefault.getValue();
+        return std::nullopt;
+      }]>,
+    InterfaceMethod<
+      /*description=*/[{
+        Set the omp.atomic_default_mem_order attribute on the operator to the
+        specified memory order kind.
+      }],
+      /*retTy=*/"void",
+      /*methodName=*/"setAtomicDefaultMemOrder",
+      (ins "::mlir::omp::ClauseMemoryOrderKind":$order), [{}], [{
+        $_op->setAttr(mlir::StringAttr::get($_op->getContext(), "omp.atomic_default_mem_order"),
+          mlir::omp::ClauseMemoryOrderKindAttr::get($_op->getContext(), order));
+      }]>
   ];
 }
 
 #endif // OpenMP_OPS_INTERFACES