[OpenMP][Flang][MLIR] Implement OffloadModuleInterface for OpenMP Dialect and convert is_device to an Attribute

[agozillon]

Summary

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).

[agozillon]

A patch based on https://reviews.llvm.org/D146721 by @kiranchandramohan an attempt at adding a module interface for OpenMP offload related attributes and functions that apply to modules. Using the existing omp.is_device work as the initial attribute (although this has been tested with the RTLModuleFlag's patch i have up and it works fine as well).

It unfortunately wasn't as simple as just adding the ODS definition, it appears that a concrete model for the interface must also be defined (and attached rather than the interface) for you to get access to the default methods defined in the ODS. But this is my first attempt at using interfaces like this, so perhaps I am missing something, please do bring it up if there is a better method! I found that just using the interface as in the original proposed patch will result in a runtime assertion whenever you try to access the interface functions as I do in this patch (even with pre-defined default ODS implementations). Perhaps when you attach an interface to an operation using the ODS framework as opposed to the attachInterface it does some table gen magic to create the Model for the operation.

[agozillon]

• [MLIR][OpenMP] Newline at end of OpenMPInterfaces.h file

[kiranchandramohan]

Thanks @agozillon for the changes. Looks fine.

It unfortunately wasn't as simple as just adding the ODS definition, it appears that a concrete model for the interface must also be defined (and attached rather than the interface) for you to get access to the default methods defined in the ODS.

Sorry about the difficulty here. I guess this is necessary. As I mentioned, I did not have sufficient time to go through this in detail but wanted to give some direction.

Did you (or @domada) check whether this can be used in the translation phase (where we connect with the OpenMP IRBuilder) as well?

Also, is there a way to test that other top-level builtin modules (if any) are not affected by this change? I guess in the flang flow there is only one.

flang/lib/Frontend/FrontendActions.cpp

227–228

Can this code be shared with bbc?

mlir/include/mlir/Dialect/OpenMP/OpenMPDialect.h

31

Nit: Is this change required?

[agozillon]

In D146850#4223461, @kiranchandramohan wrote:

Thanks @agozillon for the changes. Looks fine.

It unfortunately wasn't as simple as just adding the ODS definition, it appears that a concrete model for the interface must also be defined (and attached rather than the interface) for you to get access to the default methods defined in the ODS.

Sorry about the difficulty here. I guess this is necessary. As I mentioned, I did not have sufficient time to go through this in detail but wanted to give some direction.

No need to apologies! I just wanted to give some context to the changes from your original patch that's all, always happy to learn more.

Did you (or @domada) check whether this can be used in the translation phase (where we connect with the OpenMP IRBuilder) as well?

Yes, I have tested it, I can cast to the interface and call getIsDevice and get the correct results, so it appears to be functioning fine even in the OpenMPToLLVMIRTranslation phase, no unfortunate discarding of the interface thankfully!

Also, is there a way to test that other top-level builtin modules (if any) are not affected by this change? I guess in the flang flow there is only one.

I don't know of an easy way to test this unfortunately, however, yes, in the Flang flow there is only one module at a time (in MLIR in general it seems to assume only one builtin.module, you'll run into errors with existing tools if you have more than one (or none the only reason it won't complain in these cases is that the tools wrap MLIR implicitly in builtin.modules)). I am fairly certain it'll attach to other builtin modules in the given context (which will just be the one). If we do ever end up with two in a single context, it shouldn't be an issue as long as code segments are written based off of the contained attributes on the module, and not the existence of the interface itself I believe.

flang/lib/Frontend/FrontendActions.cpp

227–228

I use the same code in bbc at the moment a little lower in this patch :-) However, if you mean if this code can be shared across both the bbc tool and the Flang frontend then I am not so sure, I believe they do not run through the same paths, they create separate LoweringBridges and Modules in different code locations to generate their output.

mlir/include/mlir/Dialect/OpenMP/OpenMPDialect.h

31

In this case, we need a place to declare/define our C++ implementations of interfaces (models or if someone eventually decides to create an interface in C++) I thought separating it out into a separate file may be prudent. You also can't get away with forward declaring the model inside of OpenMPDialect.h as you need the fully defined class quite early on in compilation (although you could define the full model inside of the OpenMPDialect.h, but that may be messier than just the header, I am happy to do this though if that's what we wish or any other suggestions).

The ordering is also important in this case, the Module Interface (and possibly others in the future) rely on the OMP MLIR Attributes, and the ODS implementation requires them to be defined before use in either the ODS definition or the C++ definition, forward declaring the attributes inside of OpenMPDialect.h doesn't satisfy the compiler in this instance unfortunately.

[kiranchandramohan]

LGTM.

Ideally, it should be an interface whose existence can be checked to see whether it is a device Module. But I guess we can proceed for now with this. Later on we can investigate if an attribute interface would have been a better choice here.

Did you (or @domada) check whether this can be used in the translation phase (where we connect with the OpenMP IRBuilder) as well?

Yes, I have tested it, I can cast to the interface and call getIsDevice and get the correct results, so it appears to be functioning fine even in the OpenMPToLLVMIRTranslation phase, no unfortunate discarding of the interface thankfully!

OK. That is great. Hope it is working fine at the translation phase for both with and without attributes.

Also, is there a way to test that other top-level builtin modules (if any) are not affected by this change? I guess in the flang flow there is only one.

I don't know of an easy way to test this unfortunately, however, yes, in the Flang flow there is only one module at a time (in MLIR in general it seems to assume only one builtin.module, you'll run into errors with existing tools if you have more than one (or none the only reason it won't complain in these cases is that the tools wrap MLIR implicitly in builtin.modules)). I am fairly certain it'll attach to other builtin modules in the given context (which will just be the one). If we do ever end up with two in a single context, it shouldn't be an issue as long as code segments are written based off of the contained attributes on the module, and not the existence of the interface itself I believe.

I hope we can test with mlir-translate at the translation phase with two modules where one has the attributes and the other does not.

flang/lib/Frontend/FrontendActions.cpp

227–228

So, I meant some function that given a module attaches the interface, applies the interface and inserts all the attributes. If that can be placed in a common place (/flang/include/flang/Tools ?) then both bbc and frontendactions can share it.

[awarzynski]

Thanks, LGTM!

[agozillon]

In D146850#4224910, @kiranchandramohan wrote:

LGTM.

Ideally, it should be an interface whose existence can be checked to see whether it is a device Module. But I guess we can proceed for now with this. Later on we can investigate if an attribute interface would have been a better choice here.

True, you can't check if it's a device module via the interface, but you can through the is_device attribute. Which is similar to Clang, just a simple flag check, perhaps not ideal, but it works.

Did you (or @domada) check whether this can be used in the translation phase (where we connect with the OpenMP IRBuilder) as well?

Yes, I have tested it, I can cast to the interface and call getIsDevice and get the correct results, so it appears to be functioning fine even in the OpenMPToLLVMIRTranslation phase, no unfortunate discarding of the interface thankfully!

OK. That is great. Hope it is working fine at the translation phase for both with and without attributes.

I believe it currently does and it breaks no existing tests.

Also, is there a way to test that other top-level builtin modules (if any) are not affected by this change? I guess in the flang flow there is only one.

I don't know of an easy way to test this unfortunately, however, yes, in the Flang flow there is only one module at a time (in MLIR in general it seems to assume only one builtin.module, you'll run into errors with existing tools if you have more than one (or none the only reason it won't complain in these cases is that the tools wrap MLIR implicitly in builtin.modules)). I am fairly certain it'll attach to other builtin modules in the given context (which will just be the one). If we do ever end up with two in a single context, it shouldn't be an issue as long as code segments are written based off of the contained attributes on the module, and not the existence of the interface itself I believe.

I hope we can test with mlir-translate at the translation phase with two modules where one has the attributes and the other does not.

I'm not sure how this would test that the interface is not attached to other modules, but I may be misunderstanding what you mean :-) As far as I am aware the only time this interface exists on a module is during either the bbc tool or Flang-new's execution, it's not present in the IR and will not be parsed in or out at any point. The only thing that will be present are the attributes if they've been applied and the interface being attached doesn't mean that the entire subset of attributes are automatically applied, that will likely depend on the flags provided and whether it's a host or device pass.

This more general MLIR test may answer the question, if I am of course understanding the original question which I may not be (so I'm very sorry if I'm missing the point!): https://github.com/llvm/llvm-project/blob/main/mlir/unittests/IR/InterfaceAttachmentTest.cpp#L332 The interface should not be present on other modules, provided they are utilising different MLIRContexts, if they share the same context then it will. In both the case of Flang and bbc, a new MLIRContext appears to be generated on a per source file basis, and only one module is generated per source file (and as far as mlir-translate is concerned at the moment appears to be legal per file), so modules should not accidentally have this interface applied to them, the only time it will be present is if the -fopenmp flag has been given on the command line (or been generated as an argument by something like split compilation).

Thank you both for the review!

[agozillon]

• [FLANG] Add helper function for applying interface and setting attributes

[agozillon]

As requested by @kiranchandramohan I've updated the patch with a cross tool helper function that allows shared assignment of the attributes and interfaces, so no attriibute gets left behind! In theory at least. As there may be some Flags bbc will not support (so the attribute will be set to a default or not applied).

If you are both still happy with the patch and ready for it to land as is, I would appreciate a final acceptance after the last change please @kiranchandramohan @awarzynski and then I'll land it and close this tomorrow (Tuesday 28th).

Thank you both very much.

flang/lib/Frontend/FrontendActions.cpp

227–228

Ah, thank you for clarifying! I have now made one, hopefully it is in a reasonable location. I made a new file for general purpose cross tool functionality as there didn't appear to be a clear place within the Tools subdirectory.

[agozillon]

• [FLANG] Add helper function for applying interface and setting attributes

[kiranchandramohan]

In D146850#4225348, @agozillon wrote:

In D146850#4224910, @kiranchandramohan wrote:

LGTM.

Ideally, it should be an interface whose existence can be checked to see whether it is a device Module. But I guess we can proceed for now with this. Later on we can investigate if an attribute interface would have been a better choice here.

True, you can't check if it's a device module via the interface, but you can through the is_device attribute. Which is similar to Clang, just a simple flag check, perhaps not ideal, but it works.

Did you (or @domada) check whether this can be used in the translation phase (where we connect with the OpenMP IRBuilder) as well?

Yes, I have tested it, I can cast to the interface and call getIsDevice and get the correct results, so it appears to be functioning fine even in the OpenMPToLLVMIRTranslation phase, no unfortunate discarding of the interface thankfully!

OK. That is great. Hope it is working fine at the translation phase for both with and without attributes.

I believe it currently does and it breaks no existing tests.

Also, is there a way to test that other top-level builtin modules (if any) are not affected by this change? I guess in the flang flow there is only one.

I don't know of an easy way to test this unfortunately, however, yes, in the Flang flow there is only one module at a time (in MLIR in general it seems to assume only one builtin.module, you'll run into errors with existing tools if you have more than one (or none the only reason it won't complain in these cases is that the tools wrap MLIR implicitly in builtin.modules)). I am fairly certain it'll attach to other builtin modules in the given context (which will just be the one). If we do ever end up with two in a single context, it shouldn't be an issue as long as code segments are written based off of the contained attributes on the module, and not the existence of the interface itself I believe.

I hope we can test with mlir-translate at the translation phase with two modules where one has the attributes and the other does not.

I'm not sure how this would test that the interface is not attached to other modules, but I may be misunderstanding what you mean :-) As far as I am aware the only time this interface exists on a module is during either the bbc tool or Flang-new's execution, it's not present in the IR and will not be parsed in or out at any point. The only thing that will be present are the attributes if they've been applied and the interface being attached doesn't mean that the entire subset of attributes are automatically applied, that will likely depend on the flags provided and whether it's a host or device pass.

This more general MLIR test may answer the question, if I am of course understanding the original question which I may not be (so I'm very sorry if I'm missing the point!): https://github.com/llvm/llvm-project/blob/main/mlir/unittests/IR/InterfaceAttachmentTest.cpp#L332 The interface should not be present on other modules, provided they are utilising different MLIRContexts, if they share the same context then it will. In both the case of Flang and bbc, a new MLIRContext appears to be generated on a per source file basis, and only one module is generated per source file (and as far as mlir-translate is concerned at the moment appears to be legal per file), so modules should not accidentally have this interface applied to them, the only time it will be present is if the -fopenmp flag has been given on the command line (or been generated as an argument by something like split compilation).

Thank you both for the review!

I was assuming that you will use the interface to access the attributes during the translation phase as well and hence this can be tested with mlir-translate independent of the flang flow.

[agozillon]

• [FLANG][OpenMP][MLIR] Move offload interface registration to OpenMPDialect

[agozillon]

• [FLANG][OpenMP][MLIR] Move offload interface registration to OpenMPDialect

[agozillon]

I have moved the Model (interface) registration into the OpenMPDialect, so now it should be registered for all mlir::ModuleOp's in tools that initialise the OpenMPDialect (which should include mlir-translate). The mlir-translate test will come in an additional patch when there is an module attribute that can be appropriately tested in the OpenMP -> LLVM-IR lowering phase, likely in the following patch when rebased on this and other changes: https://reviews.llvm.org/D145264 as discussed external to the patch.

If it's possible could I ask for a final acceptance for this diff, before I commit the patch? @kiranchandramohan @awarzynski

[kiranchandramohan]

LGTM. Thanks for the patch and addressing the change requests.

flang/tools/bbc/bbc.cpp

248–249

Nit: Braces are not required.

flang/include/flang/Tools/CrossToolHelpers.h

+//===-- 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

flang/lib/Frontend/FrontendActions.cpp

 #include "flang/Parser/dump-parse-tree.h"
 #include "flang/Parser/parsing.h"
 #include "flang/Parser/provenance.h"
 #include "flang/Parser/source.h"
 #include "flang/Parser/unparse.h"
 #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"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Support/LLVM.h"
 #include "mlir/Target/LLVMIR/Import.h"
 #include "mlir/Target/LLVMIR/ModuleTranslation.h"
 #include "clang/Basic/Diagnostic.h"
lower::LoweringBridge lb = Fortran::lower::LoweringBridge::create(
       kindMap, ci.getInvocation().getLoweringOpts(),
       ci.getInvocation().getFrontendOpts().envDefaults);
 
   // Fetch module from lb, so we can set
   mlirModule = std::make_unique<mlir::ModuleOp>(lb.getModule());
 
   if (ci.getInvocation().getFrontendOpts().features.IsEnabled(
           Fortran::common::LanguageFeature::OpenMP)) {
-    mlir::omp::OpenMPDialect::setIsDevice(
+    setOffloadModuleInterfaceAttributes(
         *mlirModule, ci.getInvocation().getLangOpts().OpenMPIsDevice);
   }
 
   if (!setUpTargetMachine())
     return false;
   const llvm::DataLayout &dl = tm->createDataLayout();
   setMLIRDataLayout(*mlirModule, dl);
 
   // Create a parse tree and lower it to FIR
   Fortran::parser::Program &parseTree{*ci.getParsing().parseTree()};

flang/test/Lower/OpenMP/omp-is-device.f90

 !RUN: %flang_fc1 -emit-fir -fopenmp -fopenmp-is-device %s -o - | FileCheck %s --check-prefix=DEVICE
 !RUN: %flang_fc1 -emit-fir -fopenmp %s -o - | FileCheck %s --check-prefix=HOST
 !RUN: %flang_fc1 -emit-fir -fopenmp-is-device %s -o - | FileCheck %s --check-prefix=DEVICE-FLAG-ONLY
 !RUN: bbc -fopenmp -fopenmp-is-device -emit-fir -o - %s | FileCheck %s --check-prefix=DEVICE
 !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

flang/tools/bbc/bbc.cpp

 #include "flang/Parser/parse-tree.h"
 #include "flang/Parser/parsing.h"
 #include "flang/Parser/provenance.h"
 #include "flang/Parser/unparse.h"
 #include "flang/Semantics/expression.h"
 #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"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/Parser/Parser.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassManager.h"
static mlir::LogicalResult convertFortranSourceToMLIR(
   loweringOptions.setLowerToHighLevelFIR(useHLFIR);
   auto burnside = Fortran::lower::LoweringBridge::create(
       ctx, semanticsContext, defKinds, semanticsContext.intrinsics(),
       semanticsContext.targetCharacteristics(), parsing.allCooked(), "",
       kindMap, loweringOptions, {});
   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())
     outputName = llvm::sys::path::stem(inputFilename).str().append(".mlir");
   llvm::raw_fd_ostream out(outputName, ec);
   if (ec)
     return mlir::emitError(mlir::UnknownLoc::get(&ctx),
                            "could not open output file ")
            << outputName;

mlir/include/mlir/Dialect/OpenMP/OpenMPDialect.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/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"
 
 #endif // MLIR_DIALECT_OPENMP_OPENMPDIALECT_H_

mlir/include/mlir/Dialect/OpenMP/OpenMPInterfaces.h

-//===- OpenMPDialect.h - MLIR Dialect for OpenMP ----------------*- C++ -*-===//
+//===- 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 the OpenMP dialect in MLIR.
+// This file declares OpenMP Interface implementations for the OpenMP dialect.
 //
 //===----------------------------------------------------------------------===//
 
-#ifndef MLIR_DIALECT_OPENMP_OPENMPDIALECT_H_
-#define MLIR_DIALECT_OPENMP_OPENMPDIALECT_H_
+#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/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"
+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
 
-#define GET_OP_CLASSES
-#include "mlir/Dialect/OpenMP/OpenMPOps.h.inc"
-
-#endif // MLIR_DIALECT_OPENMP_OPENMPDIALECT_H_
+#endif // MLIR_DIALECT_OPENMP_OPENMPINTERFACES_H_

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

 include "mlir/Dialect/OpenMP/OpenMPOpsInterfaces.td"
 include "mlir/Dialect/OpenMP/OpenMPTypeInterfaces.td"
 
 def OpenMP_Dialect : Dialect {
   let name = "omp";
   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>;
 
 // Type which can be constraint accepting standard integers and indices.
 def IntLikeType : AnyTypeOf<[AnyInteger, Index]>;
 
 def OpenMP_PointerLikeType : TypeAlias<OpenMP_PointerLikeTypeInterface,
   "OpenMP-compatible variable type">;

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

let methods = [
     InterfaceMethod<
       "Get reduction vars", "::mlir::SmallVector<::mlir::Value>",
       "getAllReductionVars", (ins), [{}], [{
         return $_op.getReductionVars();
       }]>,
   ];
 }
 
+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

mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp

#include "mlir/Dialect/OpenMP/OpenMPOps.cpp.inc"
   addAttributes<
 #define GET_ATTRDEF_LIST
 #include "mlir/Dialect/OpenMP/OpenMPOpsAttributes.cpp.inc"
       >();
 
   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());
 }
 
 //===----------------------------------------------------------------------===//
 // Parser and printer for Allocate Clause
 //===----------------------------------------------------------------------===//
 
 /// Parse an allocate clause with allocators and a list of operands with types.
 ///
if ((cct == ClauseCancellationConstructType::Sections) &&
       !(isa<SectionsOp>(parentOp) || isa<SectionOp>(parentOp))) {
     return emitOpError() << "cancellation point sections must appear "
                          << "inside a sections region";
   }
   // TODO : Add more when we support taskgroup.
   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"
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/OpenMP/OpenMPOps.cpp.inc"