[Flang][MLIR][OpenMP][OMPIRBuilder] Use target triple to initialize `IsGPU` flag

[skatrak]

Summary

This patch modifies the construction of the OpenMPIRBuilder to initialize the IsGPU flag using target triple information passed down from the Flang frontend. If not present, it will default to false.

This replicates the behavior currently implemented in Clang, where the CodeGenModule::createOpenMPRuntime() method creates a different CGOpenMPRuntime instance depending on the target triple, which in turn has an effect on the IsGPU flag of the OpenMPIRBuilderConfig object.

[kiranchandramohan]

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

1302–1314

Can this logic be moved into the OpenMP IRBuilder?

[skatrak]

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

1302–1314

I think it would be easy to pass in the ArchType enum to the OpenMPIRBuilderConfig object and make this check there. I only did it here because it replicates what it's done in Clang: In CodeGenModule::createOpenMPRuntime() and then the constructors for CGOpenMPRuntime[GPU]. By moving this check into the OpenMPIRBuilderConfig, it would be done twice in the Clang flow (the check cannot really be removed from CodeGenModule::createOpenMPRuntime()), but I'm happy to move the check if that's preferred. Just let me know and I'll work on it, thank you.

[kiranchandramohan]

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

1302–1314

I think we should either implement this in Flang's target handling or in the OpenMPIRBuilder.

[skatrak]

Move initialization of IsTargetCodegen flag to the frontend and pass it down to the OpenMPIRBuilder through an MLIR attribute like it is done for IsDevice.

[skatrak]

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

1302–1314

Thanks for the suggestion. It seems like it makes more sense then to move it earlier in the flow rather than to the OpenMPIRBuilder at the very end, since that's already what it's done for IsDevice, so I think the logic is easier to follow and it also doesn't conflict with Clang.

[skatrak]

Ping for review! Thank you.

[kiranchandramohan]

Will there be a change in the LLVM IR generated if the option is set? If so, it might be good to have a test.

Have a couple of Nit comments regarding name of the flag. LG.

flang/include/flang/Frontend/LangOptions.def

41–42

Nit: Do you mean only for Device?

flang/tools/bbc/bbc.cpp

138–140

Nit: is this specific for GPU? If so, should we have the name GPU in the name?

[skatrak]

Rename compiler option and MLIR attribute to match new name for the flag. Create unit test for a case where that flag impacts the generated IR.

[skatrak]

Thank you for your feedback, Kiran.

In D151903#4467865, @kiranchandramohan wrote:

Will there be a change in the LLVM IR generated if the option is set? If so, it might be good to have a test.

I added an OpenMPIRBuilder test to exercise a code path that is dependent on that flag.

Have a couple of Nit comments regarding name of the flag. LG.

I created D154591 to address the confusion with flag names, because the ones we were using didn't describe very well what they represented. I just updated option / MLIR attribute names here to match what was decided there, which should also address your concerns.

[skatrak]

Fix formatting.

[skatrak]

Ping! Can this patch be landed now or should I wait for another review?

[kiranchandramohan]

LG.

flang/include/flang/Frontend/LangOptions.def

 /// Allow reassociation transformations for floating-point instructions
 LANGOPT(AssociativeMath, 1, false)
 /// Allow division operations to be reassociated
 LANGOPT(ReciprocalMath, 1, false)
 /// Specify version of OpenMP standard
 LANGOPT(OpenMPVersion, 32, 0)
 /// Generate code only for OpenMP target device
 LANGOPT(OpenMPIsTargetDevice, 1, false)
+/// Generate OpenMP target code only for GPUs
+LANGOPT(OpenMPIsGPU, 1, false)
 /// Enable debugging in the OpenMP offloading device RTL
 LANGOPT(OpenMPTargetDebug, 32, 0)
 /// Assume work-shared loops do not have more iterations than participating
 /// threads.
 LANGOPT(OpenMPThreadSubscription, 1, 0)
 /// Assume distributed loops do not have more iterations than participating
 /// teams.
 LANGOPT(OpenMPTeamSubscription, 1, 0)
 /// Assume that no thread in a parallel region will modify an ICV.
 LANGOPT(OpenMPNoThreadState, 1, 0)
 /// Assume that no thread in a parallel region will encounter a parallel region
 LANGOPT(OpenMPNoNestedParallelism, 1, 0)
 
 #undef LANGOPT
 #undef ENUM_LANGOPT

flang/include/flang/Tools/CrossToolHelpers.h

 #include "mlir/Dialect/OpenMP/OpenMPDialect.h"
 #include "mlir/IR/BuiltinOps.h"
 
 struct OffloadModuleOpts {
   OffloadModuleOpts() {}
   OffloadModuleOpts(uint32_t OpenMPTargetDebug, bool OpenMPTeamSubscription,
       bool OpenMPThreadSubscription, bool OpenMPNoThreadState,
       bool OpenMPNoNestedParallelism, bool OpenMPIsTargetDevice,
-      uint32_t OpenMPVersion, std::string OMPHostIRFile = {})
+      bool OpenMPIsGPU, uint32_t OpenMPVersion, std::string OMPHostIRFile = {})
       : OpenMPTargetDebug(OpenMPTargetDebug),
         OpenMPTeamSubscription(OpenMPTeamSubscription),
         OpenMPThreadSubscription(OpenMPThreadSubscription),
         OpenMPNoThreadState(OpenMPNoThreadState),
         OpenMPNoNestedParallelism(OpenMPNoNestedParallelism),
-        OpenMPIsTargetDevice(OpenMPIsTargetDevice),
+        OpenMPIsTargetDevice(OpenMPIsTargetDevice), OpenMPIsGPU(OpenMPIsGPU),
         OpenMPVersion(OpenMPVersion), OMPHostIRFile(OMPHostIRFile) {}
 
   OffloadModuleOpts(Fortran::frontend::LangOptions &Opts)
       : OpenMPTargetDebug(Opts.OpenMPTargetDebug),
         OpenMPTeamSubscription(Opts.OpenMPTeamSubscription),
         OpenMPThreadSubscription(Opts.OpenMPThreadSubscription),
         OpenMPNoThreadState(Opts.OpenMPNoThreadState),
         OpenMPNoNestedParallelism(Opts.OpenMPNoNestedParallelism),
         OpenMPIsTargetDevice(Opts.OpenMPIsTargetDevice),
-        OpenMPVersion(Opts.OpenMPVersion), OMPHostIRFile(Opts.OMPHostIRFile) {}
+        OpenMPIsGPU(Opts.OpenMPIsGPU), OpenMPVersion(Opts.OpenMPVersion),
+        OMPHostIRFile(Opts.OMPHostIRFile) {}
 
   uint32_t OpenMPTargetDebug = 0;
   bool OpenMPTeamSubscription = false;
   bool OpenMPThreadSubscription = false;
   bool OpenMPNoThreadState = false;
   bool OpenMPNoNestedParallelism = false;
   bool OpenMPIsTargetDevice = false;
+  bool OpenMPIsGPU = false;
   uint32_t OpenMPVersion = 11;
   std::string OMPHostIRFile = {};
 };
 
 //  Shares assinging of the OpenMP OffloadModuleInterface and its assorted
 //  attributes accross Flang tools (bbc/flang)
 void setOffloadModuleInterfaceAttributes(
     mlir::ModuleOp &module, OffloadModuleOpts Opts) {
   // Should be registered by the OpenMPDialect
   if (auto offloadMod = llvm::dyn_cast<mlir::omp::OffloadModuleInterface>(
           module.getOperation())) {
     offloadMod.setIsTargetDevice(Opts.OpenMPIsTargetDevice);
+    offloadMod.setIsGPU(Opts.OpenMPIsGPU);
     if (Opts.OpenMPIsTargetDevice) {
       offloadMod.setFlags(Opts.OpenMPTargetDebug, Opts.OpenMPTeamSubscription,
           Opts.OpenMPThreadSubscription, Opts.OpenMPNoThreadState,
           Opts.OpenMPNoNestedParallelism, Opts.OpenMPVersion);
 
       if (!Opts.OMPHostIRFile.empty())
         offloadMod.setHostIRFilePath(Opts.OMPHostIRFile);
     }

flang/lib/Frontend/CompilerInvocation.cpp

if (args.hasArg(clang::driver::options::OPT_fopenmp_is_target_device)) {
             args, clang::driver::options::OPT_fopenmp_target_debug_EQ,
             res.getLangOpts().OpenMPTargetDebug, diags);
 
         if (!res.getLangOpts().OpenMPTargetDebug &&
             args.hasArg(clang::driver::options::OPT_fopenmp_target_debug))
           res.getLangOpts().OpenMPTargetDebug = 1;
       }
     }
+
+    switch (llvm::Triple(res.getTargetOpts().triple).getArch()) {
+    case llvm::Triple::nvptx:
+    case llvm::Triple::nvptx64:
+    case llvm::Triple::amdgcn:
+      if (!res.getLangOpts().OpenMPIsTargetDevice) {
+        const unsigned diagID = diags.getCustomDiagID(
+            clang::DiagnosticsEngine::Error,
+            "OpenMP AMDGPU/NVPTX is only prepared to deal with device code.");
+        diags.Report(diagID);
+      }
+      res.getLangOpts().OpenMPIsGPU = 1;
+      break;
+    default:
+      res.getLangOpts().OpenMPIsGPU = 0;
+      break;
+    }
   }
 
   // -pedantic
   if (args.hasArg(clang::driver::options::OPT_pedantic)) {
     res.setEnableConformanceChecks();
     res.setEnableUsageChecks();
   }
   // -std=f2018

flang/test/Lower/OpenMP/host-ir-flag.f90

 !RUN: %flang_fc1 -emit-llvm-bc -fopenmp -o %t.bc %s 2>&1
 !RUN: %flang_fc1 -emit-mlir -fopenmp -fopenmp-is-target-device -fopenmp-host-ir-file-path %t.bc -o - %s 2>&1 | FileCheck %s
 
-!CHECK: module attributes {{{.*}}, omp.host_ir_filepath = "{{.*}}.bc", omp.is_target_device = true{{.*}}}
+!CHECK: module attributes {{{.*}}, omp.host_ir_filepath = "{{.*}}.bc", omp.is_gpu = false, omp.is_target_device = true{{.*}}}
 subroutine omp_subroutine()
 end subroutine omp_subroutine

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

+!RUN: %flang_fc1 -triple amdgcn-amd-amdhsa -emit-fir -fopenmp -fopenmp-is-target-device %s -o - | FileCheck %s
+!RUN: bbc -fopenmp -fopenmp-is-target-device -fopenmp-is-gpu -emit-fir -o - %s | FileCheck %s
+
+!RUN: not %flang_fc1 -triple amdgcn-amd-amdhsa -emit-fir -fopenmp %s -o - 2>&1 | FileCheck %s --check-prefix=FLANG-ERROR
+!RUN: not bbc -fopenmp -fopenmp-is-gpu -emit-fir %s -o - 2>&1 | FileCheck %s --check-prefix=BBC-ERROR
+
+!CHECK: module attributes {{{.*}}omp.is_gpu = true
+subroutine omp_subroutine()
+end subroutine omp_subroutine
+
+!FLANG-ERROR: error: OpenMP AMDGPU/NVPTX is only prepared to deal with device code.
+!BBC-ERROR: FATAL: -fopenmp-is-gpu can only be set if -fopenmp-is-target-device is also set

flang/test/Lower/OpenMP/rtl-flags.f90

 !RUN: bbc -emit-fir -fopenmp -fopenmp-version=45 -o - %s | FileCheck %s --check-prefix=DEFAULT-HOST-FIR-VERSION
 !RUN: bbc -emit-fir -fopenmp -fopenmp-target-debug=111 -fopenmp-is-target-device -o - %s | FileCheck %s --check-prefix=DBG-EQ-DEVICE-FIR
 !RUN: bbc -emit-fir -fopenmp -fopenmp-assume-teams-oversubscription -fopenmp-is-target-device -o - %s | FileCheck %s --check-prefix=TEAMS-OSUB-DEVICE-FIR
 !RUN: bbc -emit-fir -fopenmp -fopenmp-assume-threads-oversubscription -fopenmp-is-target-device -o - %s | FileCheck %s --check-prefix=THREAD-OSUB-DEVICE-FIR
 !RUN: bbc -emit-fir -fopenmp -fopenmp-assume-no-thread-state -fopenmp-is-target-device -o - %s | FileCheck %s --check-prefix=THREAD-STATE-DEVICE-FIR
 !RUN: bbc -emit-fir -fopenmp -fopenmp-assume-no-nested-parallelism -fopenmp-is-target-device -o - %s | FileCheck %s --check-prefix=NEST-PAR-DEVICE-FIR
 !RUN: bbc -emit-fir -fopenmp -fopenmp-target-debug=1 -fopenmp-assume-teams-oversubscription -fopenmp-assume-no-nested-parallelism -fopenmp-assume-threads-oversubscription -fopenmp-assume-no-thread-state -fopenmp-is-target-device -o - %s | FileCheck %s --check-prefix=ALL-DEVICE-FIR
 
-!DEFAULT-DEVICE-FIR: module attributes {{{.*}}, omp.flags = #omp.flags<openmp_device_version = 11>, omp.is_target_device = true{{.*}}}
-!DEFAULT-DEVICE-FIR-VERSION: module attributes {{{.*}}, omp.flags = #omp.flags<openmp_device_version = 45>, omp.is_target_device = true, omp.version = #omp.version<version = 45>{{.*}}
-!DEFAULT-HOST-FIR: module attributes {{{.*}},  omp.is_target_device = false{{.*}}
-!DEFAULT-HOST-FIR-VERSION: module attributes {{{.*}},  omp.is_target_device = false, omp.version = #omp.version<version = 45>{{.*}}
-!DBG-DEVICE-FIR: module attributes {{{.*}}, omp.flags = #omp.flags<debug_kind = 1, openmp_device_version = 11>{{.*}}}
-!DBG-EQ-DEVICE-FIR: module attributes {{{.*}}, omp.flags = #omp.flags<debug_kind = 111, openmp_device_version = 11>{{.*}}}
-!TEAMS-OSUB-DEVICE-FIR: module attributes {{{.*}}, omp.flags = #omp.flags<assume_teams_oversubscription = true, openmp_device_version = 11>{{.*}}}
-!THREAD-OSUB-DEVICE-FIR: module attributes {{{.*}}, omp.flags = #omp.flags<assume_threads_oversubscription = true, openmp_device_version = 11>{{.*}}}
-!THREAD-STATE-DEVICE-FIR: module attributes {{{.*}}, omp.flags = #omp.flags<assume_no_thread_state = true, openmp_device_version = 11>{{.*}}}
-!NEST-PAR-DEVICE-FIR: module attributes {{{.*}}, omp.flags = #omp.flags<assume_no_nested_parallelism = true, openmp_device_version = 11>{{.*}}}
-!ALL-DEVICE-FIR: module attributes {{{.*}}, omp.flags = #omp.flags<debug_kind = 1, assume_teams_oversubscription = true, assume_threads_oversubscription = true, assume_no_thread_state = true, assume_no_nested_parallelism = true, openmp_device_version = 11>{{.*}}}
+!DEFAULT-DEVICE-FIR: module attributes {{{.*}}omp.flags = #omp.flags<openmp_device_version = 11>
+!DEFAULT-DEVICE-FIR-SAME: omp.is_target_device = true
+!DEFAULT-DEVICE-FIR-VERSION: module attributes {{{.*}}omp.flags = #omp.flags<openmp_device_version = 45>
+!DEFAULT-DEVICE-FIR-VERSION-SAME: omp.is_target_device = true
+!DEFAULT-DEVICE-FIR-VERSION-SAME: omp.version = #omp.version<version = 45>
+!DEFAULT-HOST-FIR: module attributes {{{.*}}omp.is_target_device = false{{.*}}
+!DEFAULT-HOST-FIR-VERSION: module attributes {{{.*}}omp.is_target_device = false
+!DEFAULT-HOST-FIR-VERSION-SAME: omp.version = #omp.version<version = 45>
+!DBG-DEVICE-FIR: module attributes {{{.*}}omp.flags = #omp.flags<debug_kind = 1, openmp_device_version = 11>
+!DBG-EQ-DEVICE-FIR: module attributes {{{.*}}omp.flags = #omp.flags<debug_kind = 111, openmp_device_version = 11>
+!TEAMS-OSUB-DEVICE-FIR: module attributes {{{.*}}omp.flags = #omp.flags<assume_teams_oversubscription = true, openmp_device_version = 11>
+!THREAD-OSUB-DEVICE-FIR: module attributes {{{.*}}omp.flags = #omp.flags<assume_threads_oversubscription = true, openmp_device_version = 11>
+!THREAD-STATE-DEVICE-FIR: module attributes {{{.*}}omp.flags = #omp.flags<assume_no_thread_state = true, openmp_device_version = 11>
+!NEST-PAR-DEVICE-FIR: module attributes {{{.*}}omp.flags = #omp.flags<assume_no_nested_parallelism = true, openmp_device_version = 11>
+!ALL-DEVICE-FIR: module attributes {{{.*}}omp.flags = #omp.flags<debug_kind = 1, assume_teams_oversubscription = true, assume_threads_oversubscription = true, assume_no_thread_state = true, assume_no_nested_parallelism = true, openmp_device_version = 11>
 subroutine omp_subroutine()
 end subroutine omp_subroutine

flang/test/Lower/OpenMP/target_cpu_features.f90

 !REQUIRES: amdgpu-registered-target
-!RUN: %flang_fc1 -emit-fir -triple amdgcn-amd-amdhsa -target-cpu gfx908 -fopenmp %s -o - | FileCheck %s
+!RUN: %flang_fc1 -emit-fir -triple amdgcn-amd-amdhsa -target-cpu gfx908 -fopenmp -fopenmp-is-target-device %s -o - | FileCheck %s
 
 !===============================================================================
 ! Target_Enter Simple
 !===============================================================================
 
 !CHECK: omp.target = #omp.target<target_cpu = "gfx908",
 !CHECK-SAME: target_features = "+dot3-insts,+dot4-insts,+s-memtime-inst,
 !CHECK-SAME: +16-bit-insts,+s-memrealtime,+dot6-insts,+dl-insts,+image-insts,+wavefrontsize64,
 !CHECK-SAME: +gfx9-insts,+gfx8-insts,+ci-insts,+dot10-insts,+dot7-insts,
 !CHECK-SAME: +dot1-insts,+dot5-insts,+mai-insts,+dpp,+dot2-insts">
-!CHECK-LABEL: func.func @_QPomp_target_simple() {
+!CHECK-LABEL: func.func @_QPomp_target_simple()
 subroutine omp_target_simple
+  ! Directive needed to prevent subroutine from being filtered out when
+  ! compiling for the device.
+  !$omp declare target
 end subroutine omp_target_simple
 

flang/tools/bbc/bbc.cpp

static llvm::cl::opt<bool> enableOpenMP("fopenmp",
                                         llvm::cl::desc("enable openmp"),
                                         llvm::cl::init(false));
 
 static llvm::cl::opt<bool>
     enableOpenMPDevice("fopenmp-is-target-device",
                        llvm::cl::desc("enable openmp device compilation"),
                        llvm::cl::init(false));
 
+static llvm::cl::opt<bool>
+    enableOpenMPGPU("fopenmp-is-gpu",
+                    llvm::cl::desc("enable openmp GPU target codegen"),
+                    llvm::cl::init(false));
+
 // A simplified subset of the OpenMP RTL Flags from Flang, only the primary
 // positive options are available, no negative options e.g. fopen_assume* vs
 // fno_open_assume*
 static llvm::cl::opt<uint32_t>
     setOpenMPVersion("fopenmp-version",
                      llvm::cl::desc("OpenMP standard version"),
                      llvm::cl::init(11));
 
static mlir::LogicalResult convertFortranSourceToMLIR(
   loweringOptions.setLowerToHighLevelFIR(useHLFIR || emitHLFIR);
   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) {
-    auto offloadModuleOpts = OffloadModuleOpts(
-        setOpenMPTargetDebug, setOpenMPTeamSubscription,
+    if (enableOpenMPGPU && !enableOpenMPDevice) {
+      llvm::errs() << "FATAL: -fopenmp-is-gpu can only be set if "
+                      "-fopenmp-is-target-device is also set";
+      return mlir::failure();
+    }
+    auto offloadModuleOpts =
+        OffloadModuleOpts(setOpenMPTargetDebug, setOpenMPTeamSubscription,
                           setOpenMPThreadSubscription, setOpenMPNoThreadState,
-        setOpenMPNoNestedParallelism, enableOpenMPDevice, setOpenMPVersion);
+                          setOpenMPNoNestedParallelism, enableOpenMPDevice,
+                          enableOpenMPGPU, setOpenMPVersion);
     setOffloadModuleInterfaceAttributes(mlirModule, offloadModuleOpts);
     setOpenMPVersionAttribute(mlirModule, setOpenMPVersion);
   }
   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);

llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp

EXPECT_EQ(OffloadMetadata->getOperand(0)->getOperand(1).equalsStr(
                   "test_data_int_0"),
               true);
     EXPECT_EQ(OffloadMetadata->getOperand(1)->getOperand(1).equalsStr(
                   "test_data_int_1_decl_tgt_ref_ptr"),
               true);
   }
 }
 
+TEST_F(OpenMPIRBuilderTest, createGPUOffloadEntry) {
+  OpenMPIRBuilder OMPBuilder(*M);
+  OMPBuilder.initialize();
+  OpenMPIRBuilderConfig Config(/* IsTargetDevice = */ true,
+                               /* IsGPU = */ true,
+                               /* HasRequiresUnifiedSharedMemory = */ false,
+                               /* OpenMPOffloadMandatory = */ false);
+  OMPBuilder.setConfig(Config);
+
+  FunctionCallee FnTypeAndCallee =
+      M->getOrInsertFunction("test_kernel", Type::getVoidTy(Ctx));
+
+  auto *Fn = cast<Function>(FnTypeAndCallee.getCallee());
+  OMPBuilder.createOffloadEntry(/* ID = */ nullptr, Fn,
+                                /* Size = */ 0,
+                                /* Flags = */ 0, GlobalValue::WeakAnyLinkage);
+
+  // Check nvvm.annotations only created for GPU kernels
+  NamedMDNode *MD = M->getNamedMetadata("nvvm.annotations");
+  EXPECT_NE(MD, nullptr);
+  EXPECT_EQ(MD->getNumOperands(), 1u);
+
+  MDNode *Annotations = MD->getOperand(0);
+  EXPECT_EQ(Annotations->getNumOperands(), 3u);
+
+  Constant *ConstVal =
+      dyn_cast<ConstantAsMetadata>(Annotations->getOperand(0))->getValue();
+  EXPECT_TRUE(isa<Function>(Fn));
+  EXPECT_EQ(ConstVal, cast<Function>(Fn));
+
+  EXPECT_TRUE(Annotations->getOperand(1).equalsStr("kernel"));
+
+  EXPECT_TRUE(mdconst::hasa<ConstantInt>(Annotations->getOperand(2)));
+  APInt IntVal =
+      mdconst::extract<ConstantInt>(Annotations->getOperand(2))->getValue();
+  EXPECT_EQ(IntVal, 1);
+
+  // Check kernel attributes
+  EXPECT_TRUE(Fn->hasFnAttribute("kernel"));
+  EXPECT_TRUE(Fn->hasFnAttribute(Attribute::MustProgress));
+}
+
 } // namespace

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

InterfaceMethod<
       (ins), [{}], [{
         if (Attribute isTargetDevice = $_op->getAttr("omp.is_target_device"))
           if (isTargetDevice.isa<mlir::BoolAttr>())
            return isTargetDevice.dyn_cast<BoolAttr>().getValue();
         return false;
       }]>,
     InterfaceMethod<
       /*description=*/[{
+      Set the attribute on the current module with the specified boolean
+      argument.
+      }],
+      /*retTy=*/"void",
+      /*methodName=*/"setIsGPU",
+      (ins "bool":$isGPU), [{}], [{
+        $_op->setAttr(
+          mlir::StringAttr::get($_op->getContext(), "omp.is_gpu"),
+            mlir::BoolAttr::get($_op->getContext(), isGPU));
+      }]>,
+    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=*/"getIsGPU",
+      (ins), [{}], [{
+        if (Attribute isTargetCGAttr = $_op->getAttr("omp.is_gpu"))
+          if (auto isTargetCGVal = isTargetCGAttr.dyn_cast<BoolAttr>())
+           return isTargetCGVal.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"))
           return ::llvm::dyn_cast_or_null<mlir::omp::FlagsAttr>(flags);

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

for (Value v : values)
     remapped.push_back(lookupValue(v));
   return remapped;
 }
 
 llvm::OpenMPIRBuilder *ModuleTranslation::getOpenMPBuilder() {
   if (!ompBuilder) {
     ompBuilder = std::make_unique<llvm::OpenMPIRBuilder>(*llvmModule);
 
-    bool isTargetDevice = false;
+    bool isTargetDevice = false, isGPU = false;
     llvm::StringRef hostIRFilePath = "";
 
-    if (Attribute deviceAttr = mlirModule->getAttr("omp.is_target_device"))
-      if (::llvm::isa<mlir::BoolAttr>(deviceAttr))
-        isTargetDevice =
-            ::llvm::dyn_cast<mlir::BoolAttr>(deviceAttr).getValue();
-
-    if (Attribute filepath = mlirModule->getAttr("omp.host_ir_filepath"))
-      if (::llvm::isa<mlir::StringAttr>(filepath))
-        hostIRFilePath =
-            ::llvm::dyn_cast<mlir::StringAttr>(filepath).getValue();
+    if (auto deviceAttr =
+            mlirModule->getAttrOfType<mlir::BoolAttr>("omp.is_target_device"))
+      isTargetDevice = deviceAttr.getValue();
+
+    if (auto gpuAttr = mlirModule->getAttrOfType<mlir::BoolAttr>("omp.is_gpu"))
+      isGPU = gpuAttr.getValue();
+
+    if (auto filepathAttr =
+            mlirModule->getAttrOfType<mlir::StringAttr>("omp.host_ir_filepath"))
+      hostIRFilePath = filepathAttr.getValue();
 
     ompBuilder->initialize(hostIRFilePath);
 
     // TODO: set the flags when available
     llvm::OpenMPIRBuilderConfig config(
-        isTargetDevice, /* IsGPU */ false,
+        isTargetDevice, isGPU,
         /* HasRequiresUnifiedSharedMemory */ false,
         /* OpenMPOffloadMandatory */ false);
     ompBuilder->setConfig(config);
   }
   return ompBuilder.get();
 }
 
 llvm::DILocation *ModuleTranslation::translateLoc(Location loc,
                                                   llvm::DILocalScope *scope) {
   return debugTranslation->translateLoc(loc, scope);
 }
 
 llvm::Metadata *ModuleTranslation::translateDebugInfo(LLVM::DINodeAttr attr) {
   return debugTranslation->translate(attr);
 }
 
 llvm::NamedMDNode *
 ModuleTranslation::getOrInsertNamedModuleMetadata(StringRef name) {