[Mlir-commits] [mlir] a7e91b1 - [OpenMP][MLIR] Add Flags attribute to OMP OffloadModuleInterface

[Andrew Gozillon]

Author: Andrew Gozillon
Date: 2023-03-29T12:54:38-05:00
New Revision: a7e91b18b73b9971d3af2c6df1811e3da63e61c8

URL: https://github.com/llvm/llvm-project/commit/a7e91b18b73b9971d3af2c6df1811e3da63e61c8
DIFF: https://github.com/llvm/llvm-project/commit/a7e91b18b73b9971d3af2c6df1811e3da63e61c8.diff

LOG: [OpenMP][MLIR] Add Flags attribute to OMP OffloadModuleInterface

The intent of this attribute is for it to be applied to a module and
then hold information on runtime library (RTL) flags given to
Flang (or other OpenMP frontend) that should be lowered down to
LLVM-IR for devices as LLVM globals. The following related
flags are:

-fopenmp-target-debug
-fopenmp-assume-threads-oversubscription
-fopenmp-assume-teams-oversubscription
-fopenmp-assume-no-nested-parallelism
-fopenmp-assume-no-thread-state

These exist within Clang and are lowered into the IR when
offloading for device. This attribute allows this infromation
to be carried down from the Flang frontend to the
LLVM/OpenMP Dialect to LLVM-IR translation phase
and then be lowered to LLVM-IR.

Reviewers: kiranchandramohan

Differential Revision: https://reviews.llvm.org/D144896

Added: 
    mlir/test/Dialect/OpenMP/attr.mlir

Modified: 
    mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
    mlir/include/mlir/Dialect/OpenMP/OpenMPOpsInterfaces.td

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
index f79aa50328b9d..8a50d404bdd4e 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
@@ -52,6 +52,22 @@ def IsDeviceAttr : OpenMP_Attr<"IsDevice", "isdevice"> {
   let assemblyFormat = "`<` struct(params) `>`";
 }
 
+//===----------------------------------------------------------------------===//
+// Runtime library flag's attribute that holds information for lowering to LLVM
+//===----------------------------------------------------------------------===//
+
+def FlagsAttr : OpenMP_Attr<"Flags", "flags"> {
+  let parameters = (ins 
+    DefaultValuedParameter<"uint32_t", "0">:$debug_kind,
+    DefaultValuedParameter<"bool", "false">:$assume_teams_oversubscription,
+    DefaultValuedParameter<"bool", "false">:$assume_threads_oversubscription,
+    DefaultValuedParameter<"bool", "false">:$assume_no_thread_state,
+    DefaultValuedParameter<"bool", "false">:$assume_no_nested_parallelism
+  );
+  
+  let assemblyFormat = "`<` struct(params) `>`";
+}
+
 class OpenMP_Op<string mnemonic, list<Trait> traits = []> :
       Op<OpenMP_Dialect, mnemonic, traits>;
 

diff  --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOpsInterfaces.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOpsInterfaces.td
index 69cb6050b1acc..934698ae5d6fa 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOpsInterfaces.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOpsInterfaces.td
@@ -81,6 +81,35 @@ def OffloadModuleInterface : OpInterface<"OffloadModuleInterface"> {
             return isDevice.dyn_cast<IsDeviceAttr>().getIsDevice();
         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"))
+          return flags.dyn_cast_or_null<mlir::omp::FlagsAttr>();
+        return nullptr;
+      }]>,  
+      InterfaceMethod<
+      /*description=*/[{
+        Apply an omp.FlagsAttr to a module with the specified values 
+        for the flags
+      }],
+      /*retTy=*/"void",
+      /*methodName=*/"setFlags", 
+      (ins "uint32_t":$debugKind,
+            "bool":$assumeTeamsOversubscription,
+            "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));
+      }]>,
   ];
 }
 

diff  --git a/mlir/test/Dialect/OpenMP/attr.mlir b/mlir/test/Dialect/OpenMP/attr.mlir
new file mode 100644
index 0000000000000..af0ca8f5e9546
--- /dev/null
+++ b/mlir/test/Dialect/OpenMP/attr.mlir
@@ -0,0 +1,31 @@
+// RUN: mlir-opt %s | mlir-opt | FileCheck %s
+
+// CHECK: module attributes {omp.flags = #omp.flags<>} {
+module attributes {omp.flags = #omp.flags<debug_kind = 0, assume_teams_oversubscription = false, assume_threads_oversubscription = false, assume_no_thread_state = false, assume_no_nested_parallelism = false>} {}
+
+// CHECK: module attributes {omp.flags = #omp.flags<debug_kind = 20>} {
+module attributes {omp.flags = #omp.flags<debug_kind = 20, assume_teams_oversubscription = false, assume_threads_oversubscription = false, assume_no_thread_state = false, assume_no_nested_parallelism = false>} {}
+
+// CHECK: module attributes {omp.flags = #omp.flags<debug_kind = 100, assume_teams_oversubscription = true>} {
+module attributes {omp.flags = #omp.flags<debug_kind = 100, assume_teams_oversubscription = true, assume_threads_oversubscription = false, assume_no_thread_state = false, assume_no_nested_parallelism = false>} {}
+
+// CHECK: module attributes {omp.flags = #omp.flags<debug_kind = 200, assume_teams_oversubscription = true, assume_threads_oversubscription = true>} {
+module attributes {omp.flags = #omp.flags<debug_kind = 200, assume_teams_oversubscription = true, assume_threads_oversubscription = true, assume_no_thread_state = false, assume_no_nested_parallelism = false>} {}
+
+// CHECK: module attributes {omp.flags = #omp.flags<debug_kind = 300, assume_teams_oversubscription = true, assume_threads_oversubscription = true, assume_no_thread_state = true>} {
+module attributes {omp.flags = #omp.flags<debug_kind = 300, assume_teams_oversubscription = true, assume_threads_oversubscription = true, assume_no_thread_state = true, assume_no_nested_parallelism = false>} {}
+
+// CHECK: module attributes {omp.flags = #omp.flags<debug_kind = 400, assume_teams_oversubscription = true, assume_threads_oversubscription = true, assume_no_thread_state = true, assume_no_nested_parallelism = true>} {
+module attributes {omp.flags = #omp.flags<debug_kind = 400, assume_teams_oversubscription = true, assume_threads_oversubscription = true, assume_no_thread_state = true, assume_no_nested_parallelism = true>} {}
+
+// CHECK: module attributes {omp.flags = #omp.flags<>} {
+module attributes {omp.flags = #omp.flags<>} {}
+
+// CHECK: module attributes {omp.flags = #omp.flags<assume_teams_oversubscription = true>} {
+module attributes {omp.flags = #omp.flags<assume_teams_oversubscription = true>} {}
+
+// CHECK: module attributes {omp.flags = #omp.flags<assume_teams_oversubscription = true, assume_no_thread_state = true>} {
+module attributes {omp.flags = #omp.flags<assume_teams_oversubscription = true, assume_no_thread_state = true>} {}
+
+// CHECK: module attributes {omp.flags = #omp.flags<assume_teams_oversubscription = true, assume_no_thread_state = true>} {
+module attributes {omp.flags = #omp.flags<assume_no_thread_state = true, assume_teams_oversubscription = true>} {}