[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

1696

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

1696

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
+//===----------------------------------------------------------------------===//
+
+// atomic_default_mem_order clause values not defined here because they can be
+// represented by the OMPC_MemoryOrder enumeration instead.
+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 clauses",
+    [
+      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<
       }],
       /*retTy=*/"void",
       /*methodName=*/"setIsDevice",
       (ins "bool":$isDevice), [{}], [{
         $_op->setAttr(
           mlir::StringAttr::get($_op->getContext(), llvm::Twine{"omp.is_device"}),
             mlir::BoolAttr::get($_op->getContext(), isDevice));
       }]>,
-      InterfaceMethod<
       /*description=*/[{
         Get the attribute on the current module if it exists and
         return its value, if it doesn't exist it returns false by default.
       }],
       /*retTy=*/"bool",
       /*methodName=*/"getIsDevice",
       (ins), [{}], [{
         if (Attribute isDevice = $_op->getAttr("omp.is_device"))
           if (isDevice.isa<mlir::BoolAttr>())
            return isDevice.dyn_cast<BoolAttr>().getValue();
         return false;
       }]>,
-      InterfaceMethod<
       /*description=*/[{
         Get the FlagsAttr attribute on the current module if it exists
         and return the attribute, if it doesn't exit it returns a nullptr
       }],
       /*retTy=*/"mlir::omp::FlagsAttr",
       /*methodName=*/"getFlags",
       (ins), [{}], [{
         if (Attribute flags = $_op->getAttr("omp.flags"))
InterfaceMethod<
            "llvm::StringRef":$targetFeatures), [{}], [{
         if (targetCPU.empty())
           return;
         $_op->setAttr(("omp." + mlir::omp::TargetAttr::getMnemonic()).str(),
                   mlir::omp::TargetAttr::get($_op->getContext(),
                                              targetCPU.str(),
                                              targetFeatures.str()));
       }]>,
-      InterfaceMethod<
       /*description=*/[{
         Set a StringAttr on the current module containing the host IR file path. This
         file path is used in two-phase compilation during the device phase to generate
         device side LLVM IR when lowering MLIR.
       }],
       /*retTy=*/"void",
       /*methodName=*/"setHostIRFilePath",
       (ins "std::string":$hostIRFilePath), [{}], [{
         $_op->setAttr(
           mlir::StringAttr::get($_op->getContext(), llvm::Twine{"omp.host_ir_filepath"}),
             mlir::StringAttr::get($_op->getContext(), hostIRFilePath));
        }]>,
-      InterfaceMethod<
       /*description=*/[{
         Find the host-ir file path StringAttr from the current module if it exists and
         return its contained value, if it doesn't exist it returns an empty string. This
         file path is used in two-phase compilation during the device phase to generate
         device side LLVM IR when lowering MLIR.
       }],
       /*retTy=*/"llvm::StringRef",
       /*methodName=*/"getHostIRFilePath",
       (ins), [{}], [{
         if (Attribute filepath = $_op->getAttr("omp.host_ir_filepath"))
           if (::llvm::isa<mlir::StringAttr>(filepath))
             return ::llvm::dyn_cast<mlir::StringAttr>(filepath).getValue();
         return {};
+      }]>,
+    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));
       }]>
   ];
 }
 
 #endif // OpenMP_OPS_INTERFACES

mlir/test/Dialect/OpenMP/ops.mlir

for (%i, %j) : i32 = (%lb, %ub) to (%ub, %lb) step (%step, %step) {
     // CHECK: omp.terminator
     omp.terminator
   }
 
   // CHECK: return
   return
 }
 
+// CHECK: func.func @omp_requires_one
+// CHECK-SAME: omp.requires = #omp<clause_requires reverse_offload>
+func.func @omp_requires_one() -> ()
+    attributes {omp.requires = #omp<clause_requires reverse_offload>} {
+  return
+}
+
+// CHECK: func.func @omp_requires_multiple
+// CHECK-SAME: omp.requires = #omp<clause_requires unified_address|dynamic_allocators>
+func.func @omp_requires_multiple() -> ()
+    attributes {omp.requires = #omp<clause_requires unified_address|dynamic_allocators>} {
+  return
+}
+
 // -----
 
 // CHECK-LABEL: @opaque_pointers_atomic_rwu
 // CHECK-SAME: (%[[v:.*]]: !llvm.ptr, %[[x:.*]]: !llvm.ptr)
 func.func @opaque_pointers_atomic_rwu(%v: !llvm.ptr, %x: !llvm.ptr) {
   // CHECK: omp.atomic.read %[[v]] = %[[x]] : !llvm.ptr, i32
   // CHECK: %[[VAL:.*]] = llvm.load %[[x]] : !llvm.ptr -> i32
   // CHECK: omp.atomic.write %[[v]] = %[[VAL]] : !llvm.ptr, i32