[OpenMP][Flang][MLIR] Implement OffloadModuleInterface for OpenMP Dia…

…lect and convert is_device to an Attribute This commit adds the OffloadModuleInterface to the OpenMP dialect, which will implement future module attribute get/set's for offloading. Currently it implements set and get's for the omp.is_device attribute, which is promoted to a real attribute in this commit as well (primarily to allow switch cases to work nicely with it for future work and to keep consistency with future module attributes). This interface is attached to mlir::ModuleOp's on registration of the OpenMPDialect and should be accessible anywhere the OpenMP dialect is registered and initialized. Reviewers: kiranchandramohan, awarzynski Differential Revision: https://reviews.llvm.org/D146850
llvm · Mar 28, 2023 · 132feb7 · 132feb7
1 parent 48cd8b5
commit 132feb7
Show file tree

Hide file tree

Showing 9 changed files with 133 additions and 35 deletions.
diff --git a/flang/include/flang/Tools/CrossToolHelpers.h b/flang/include/flang/Tools/CrossToolHelpers.h
@@ -0,0 +1,30 @@
+//===-- Tools/CrossToolHelpers.h --------------------------------- *-C++-*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+// A header file for containing functionallity that is used across Flang tools,
+// such as helper functions which apply or generate information needed accross
+// tools like bbc and flang-new.
+//===----------------------------------------------------------------------===//
+
+#ifndef FORTRAN_TOOLS_CROSS_TOOL_HELPERS_H
+#define FORTRAN_TOOLS_CROSS_TOOL_HELPERS_H
+
+#include "mlir/Dialect/OpenMP/OpenMPDialect.h"
+#include "mlir/IR/BuiltinOps.h"
+
+//  Shares assinging of the OpenMP OffloadModuleInterface and its assorted
+//  attributes accross Flang tools (bbc/flang)
+void setOffloadModuleInterfaceAttributes(
+    mlir::ModuleOp &module, bool isDevice) {
+  // Should be registered by the OpenMPDialect
+  if (auto offloadMod = llvm::dyn_cast<mlir::omp::OffloadModuleInterface>(
+          module.getOperation())) {
+    offloadMod.setIsDevice(isDevice);
+  }
+}
+
+#endif // FORTRAN_TOOLS_CROSS_TOOL_HELPERS_H
diff --git a/flang/lib/Frontend/FrontendActions.cpp b/flang/lib/Frontend/FrontendActions.cpp
@@ -31,6 +31,7 @@
 #include "flang/Semantics/runtime-type-info.h"
 #include "flang/Semantics/semantics.h"
 #include "flang/Semantics/unparse-with-symbols.h"
+#include "flang/Tools/CrossToolHelpers.h"
 
 #include "mlir/IR/Dialect.h"
 #include "mlir/Parser/Parser.h"
@@ -222,7 +223,7 @@ bool CodeGenAction::beginSourceFileAction() {
 
   if (ci.getInvocation().getFrontendOpts().features.IsEnabled(
           Fortran::common::LanguageFeature::OpenMP)) {
-    mlir::omp::OpenMPDialect::setIsDevice(
+    setOffloadModuleInterfaceAttributes(
         *mlirModule, ci.getInvocation().getLangOpts().OpenMPIsDevice);
   }
 

diff --git a/flang/test/Lower/OpenMP/omp-is-device.f90 b/flang/test/Lower/OpenMP/omp-is-device.f90
@@ -5,10 +5,10 @@
 !RUN: bbc -fopenmp -emit-fir -o - %s | FileCheck %s --check-prefix=HOST
 !RUN: bbc -fopenmp-is-device -emit-fir -o - %s | FileCheck %s --check-prefix=DEVICE-FLAG-ONLY
 
-!DEVICE: module attributes {{{.*}}, omp.is_device = true{{.*}}}
-!HOST: module attributes {{{.*}}, omp.is_device = false{{.*}}}
+!DEVICE: module attributes {{{.*}}, omp.is_device = #omp.isdevice<is_device = true>{{.*}}}
+!HOST: module attributes {{{.*}}, omp.is_device = #omp.isdevice<is_device = false>{{.*}}}
 !DEVICE-FLAG-ONLY: module attributes {{{.*}}"
-!DEVICE-FLAG-ONLY-NOT: , omp.is_device = {{.*}}
+!DEVICE-FLAG-ONLY-NOT: , omp.is_device = #omp.isdevice<{{.*}}> 
 !DEVICE-FLAG-ONLY-SAME: }
 subroutine omp_subroutine()
 end subroutine omp_subroutine
diff --git a/flang/tools/bbc/bbc.cpp b/flang/tools/bbc/bbc.cpp
@@ -37,6 +37,7 @@
 #include "flang/Semantics/runtime-type-info.h"
 #include "flang/Semantics/semantics.h"
 #include "flang/Semantics/unparse-with-symbols.h"
+#include "flang/Tools/CrossToolHelpers.h"
 #include "flang/Version.inc"
 #include "mlir/Dialect/OpenMP/OpenMPDialect.h"
 #include "mlir/IR/AsmState.h"
@@ -244,7 +245,7 @@ static mlir::LogicalResult convertFortranSourceToMLIR(
   burnside.lower(parseTree, semanticsContext);
   mlir::ModuleOp mlirModule = burnside.getModule();
   if (enableOpenMP)
-    mlir::omp::OpenMPDialect::setIsDevice(mlirModule, enableOpenMPDevice);
+    setOffloadModuleInterfaceAttributes(mlirModule, enableOpenMPDevice);
   std::error_code ec;
   std::string outputName = outputFilename;
   if (!outputName.size())

diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPDialect.h b/mlir/include/mlir/Dialect/OpenMP/OpenMPDialect.h
@@ -23,12 +23,13 @@
 
 #include "mlir/Dialect/OpenMP/OpenMPOpsDialect.h.inc"
 #include "mlir/Dialect/OpenMP/OpenMPOpsEnums.h.inc"
-#include "mlir/Dialect/OpenMP/OpenMPOpsInterfaces.h.inc"
 #include "mlir/Dialect/OpenMP/OpenMPTypeInterfaces.h.inc"
 
 #define GET_ATTRDEF_CLASSES
 #include "mlir/Dialect/OpenMP/OpenMPOpsAttributes.h.inc"
 
+#include "mlir/Dialect/OpenMP/OpenMPInterfaces.h"
+
 #define GET_OP_CLASSES
 #include "mlir/Dialect/OpenMP/OpenMPOps.h.inc"
 

diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPInterfaces.h b/mlir/include/mlir/Dialect/OpenMP/OpenMPInterfaces.h
@@ -0,0 +1,35 @@
+//===- OpenMPInterfaces.h - MLIR Interfaces for OpenMP ----------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file declares OpenMP Interface implementations for the OpenMP dialect.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_OPENMP_OPENMPINTERFACES_H_
+#define MLIR_DIALECT_OPENMP_OPENMPINTERFACES_H_
+
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/IR/Dialect.h"
+#include "mlir/IR/OpDefinition.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/SymbolTable.h"
+#include "mlir/Interfaces/ControlFlowInterfaces.h"
+#include "mlir/Interfaces/SideEffectInterfaces.h"
+
+#include "mlir/Dialect/OpenMP/OpenMPOpsInterfaces.h.inc"
+
+namespace mlir::omp {
+// You can override defaults here or implement more complex implementations of
+// functions. Or define a completely seperate external model implementation,
+// to override the existing implementation.
+struct OffloadModuleDefaultModel
+    : public OffloadModuleInterface::ExternalModel<OffloadModuleDefaultModel,
+                                                   mlir::ModuleOp> {};
+} // namespace mlir::omp
+
+#endif // MLIR_DIALECT_OPENMP_OPENMPINTERFACES_H_
diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
@@ -28,20 +28,30 @@ def OpenMP_Dialect : Dialect {
   let cppNamespace = "::mlir::omp";
   let dependentDialects = ["::mlir::LLVM::LLVMDialect"];
   let useDefaultAttributePrinterParser = 1;
-
-  let extraClassDeclaration = [{
-    // Set the omp.is_device attribute on the module with the specified boolean
-    static void setIsDevice(Operation* module, bool isDevice);
-
-    // Return the value of the omp.is_device attribute stored in the module if it
-    // exists, otherwise return false by default
-    static bool getIsDevice(Operation* module);
-  }];
 }
 
 // OmpCommon requires definition of OpenACC_Dialect.
 include "mlir/Dialect/OpenMP/OmpCommon.td"
 
+//===----------------------------------------------------------------------===//
+//  OpenMP Attributes
+//===----------------------------------------------------------------------===//
+
+class OpenMP_Attr<string name, string attrMnemonic,
+                list<Trait> traits = [],
+                string baseCppClass = "::mlir::Attribute">
+    : AttrDef<OpenMP_Dialect, name, traits, baseCppClass> {
+  let mnemonic = attrMnemonic;
+}
+
+def IsDeviceAttr : OpenMP_Attr<"IsDevice", "isdevice"> {
+  let parameters = (ins 
+    "bool":$is_device
+  );
+
+  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
@@ -47,4 +47,40 @@ def ReductionClauseInterface : OpInterface<"ReductionClauseInterface"> {
   ];
 }
 
+def OffloadModuleInterface : OpInterface<"OffloadModuleInterface"> {
+  let description = [{
+    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. 
+      }],
+      /*retTy=*/"void",
+      /*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", 
+      (ins), [{}], [{
+        if (Attribute isDevice = $_op->getAttr("omp.is_device"))
+          if (isDevice.isa<mlir::omp::IsDeviceAttr>())
+            return isDevice.dyn_cast<IsDeviceAttr>().getIsDevice();
+      }]>,
+  ];
+}
+
 #endif // OpenMP_OPS_INTERFACES
diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -58,6 +58,10 @@ void OpenMPDialect::initialize() {
   LLVM::LLVMPointerType::attachInterface<
       PointerLikeModel<LLVM::LLVMPointerType>>(*getContext());
   MemRefType::attachInterface<PointerLikeModel<MemRefType>>(*getContext());
+  LLVM::LLVMPointerType::attachInterface<
+      PointerLikeModel<LLVM::LLVMPointerType>>(*getContext());
+  mlir::ModuleOp::attachInterface<mlir::omp::OffloadModuleDefaultModel>(
+      *getContext());
 }
 
 //===----------------------------------------------------------------------===//
@@ -1417,26 +1421,6 @@ LogicalResult CancellationPointOp::verify() {
   return success();
 }
 
-//===----------------------------------------------------------------------===//
-// OpenMPDialect helper functions
-//===----------------------------------------------------------------------===//
-
-// Set the omp.is_device attribute on the module with the specified boolean
-void OpenMPDialect::setIsDevice(Operation* module, bool isDevice) {
-  module->setAttr(
-      mlir::StringAttr::get(module->getContext(), llvm::Twine{"omp.is_device"}),
-      mlir::BoolAttr::get(module->getContext(), isDevice));
-}
-
-// Return the value of the omp.is_device attribute stored in the module if it
-// exists, otherwise return false by default
-bool OpenMPDialect::getIsDevice(Operation* module) {
-  if (Attribute isDevice = module->getAttr("omp.is_device"))
-    if (isDevice.isa<mlir::BoolAttr>())
-      return isDevice.dyn_cast<BoolAttr>().getValue();
-  return false;
-}
-
 #define GET_ATTRDEF_CLASSES
 #include "mlir/Dialect/OpenMP/OpenMPOpsAttributes.cpp.inc"