[llvm] [mlir] [Flang][OpenMP][OpenMPIRBuilder] Implement module scope declare target use rewrite mechanism (PR #212920)

via llvm-commits llvm-commits at lists.llvm.org
Thu Jul 30 20:47:53 PDT 2026


https://github.com/agozillon updated https://github.com/llvm/llvm-project/pull/212920

>From c60b1a86138663dd0b2b9c13686f5e9357510944 Mon Sep 17 00:00:00 2001
From: agozillon <Andrew.Gozillon at amd.com>
Date: Wed, 29 Jul 2026 11:31:32 -0500
Subject: [PATCH] [Flang][OpenMP][OpenMPIRBuilder] Implement module scope
 declare target use rewrite mechanism

During lowering of declare target'd variables we generate new global variables for device that
replace the use of the pre-existing global variable. In Flang we currently rewrite this for
each target region, but that's not enough to cover indirect use cases inside of declare target
functions which can be imported into the module and utilised inside of a target region. This
PR tries to extend the scope of the rewriting to the module than a per target region rewrite.

It does so by creating a mechanism where we can register globals for replacement which will
trigger on finalization of the OMPIRBuilder. This is required as due to the ordering of
lowering for MLIR, where we generate the replacement global at the beginning of the module
before any uses have been generated, effectively meaning we cannot replace the uses at that
point. So, we defer the replacement to the OMPIRBuilder as there is no deferral mechanism
directly in the OpenMP MLIR lowering.

The alternative might be to rebind the global maps in ModuleTranslation (which requires
extending ModuleTranslation a bit and might not be looked apon as a great alteration from the
MLIR community) so that the old global points to the new one, but in practice this doesn't
work particularly well for declare target link/usm variables as they neccesitate a load and
the act of rebinding the global doesn't indicate to the lowering that the load is required.
So, the OMPIRBuilder method allows us more flexibility to make this (and other) required
alterations.
---
 .../llvm/Frontend/OpenMP/OMPIRBuilder.h       | 31 +++++++
 llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp     | 89 +++++++++++++++++++
 .../OpenMP/OpenMPToLLVMIRTranslation.cpp      | 88 +++---------------
 ...are-target-module-rewrite-link-device.mlir | 48 ++++++++++
 ...lare-target-module-rewrite-phi-device.mlir | 39 ++++++++
 ...rget-module-rewrite-phi-nested-device.mlir | 60 +++++++++++++
 ...clare-target-module-rewrite-to-device.mlir | 60 +++++++++++++
 ...lare-target-module-rewrite-usm-device.mlir | 62 +++++++++++++
 8 files changed, 403 insertions(+), 74 deletions(-)
 create mode 100644 mlir/test/Target/LLVMIR/omptarget-declare-target-module-rewrite-link-device.mlir
 create mode 100644 mlir/test/Target/LLVMIR/omptarget-declare-target-module-rewrite-phi-device.mlir
 create mode 100644 mlir/test/Target/LLVMIR/omptarget-declare-target-module-rewrite-phi-nested-device.mlir
 create mode 100644 mlir/test/Target/LLVMIR/omptarget-declare-target-module-rewrite-to-device.mlir
 create mode 100644 mlir/test/Target/LLVMIR/omptarget-declare-target-module-rewrite-usm-device.mlir

diff --git a/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h b/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h
index 05307414c94d0..a47795aaa7b4e 100644
--- a/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h
+++ b/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h
@@ -1109,6 +1109,21 @@ class OpenMPIRBuilder {
       std::function<GlobalValue::LinkageTypes()> VariableLinkage,
       Type *LlvmPtrTy, Constant *Addr);
 
+  /// Register a module-scope replacement of a declare target global variable.
+  /// During lowering new globals are generated for certain combinations of
+  /// declare target input, and these new globals require substitution with
+  /// the originals. This replacement occurs during finalization where uses
+  /// of \p Original are rewritten to reference \p Replacement. This is
+  /// predominantly required for Flang where the lowering pattern to LLVM
+  /// prevents immediate use rewrites.
+  /// \param Original - The original global variable that will be replaced.
+  /// \param Replacement - The replacement reference pointer generated by the
+  /// declare target infrastructure (declare target link or unified shared
+  /// memory globals).
+  LLVM_ABI void
+  registerDeclareTargetGlobalReplacement(GlobalValue *Original,
+                                         GlobalValue *Replacement);
+
   /// Get the offset of the OMP_MAP_MEMBER_OF field.
   LLVM_ABI unsigned getFlagMemberOffset();
 
@@ -2669,6 +2684,22 @@ class OpenMPIRBuilder {
   /// outline info's have been processed.
   SmallVector<llvm::Function *, 16> ConstantAllocaRaiseCandidates;
 
+  /// Describes a declare target global variable replacement to be applied
+  /// during finalization.
+  struct DeclareTargetGlobalReplacement {
+    GlobalValue *Original;
+    GlobalValue *Replacement;
+  };
+
+  /// Collection of declare target globals to rewrite uses of during
+  /// finalizaiton.
+  SmallVector<DeclareTargetGlobalReplacement, 8>
+      DeclareTargetGlobalReplacements;
+
+  /// Rewrites uses of registered declare target globals to their
+  /// replacements.
+  void applyDeclareTargetGlobalReplacements();
+
   /// Collection of owned canonical loop objects that eventually need to be
   /// free'd.
   std::forward_list<CanonicalLoopInfo> LoopInfos;
diff --git a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
index fc3812502fb2e..69984f099c38b 100644
--- a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
+++ b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
@@ -1040,6 +1040,9 @@ void OpenMPIRBuilder::finalize(Function *Fn) {
   if (!OffloadInfoManager.empty())
     createOffloadEntriesAndInfoMetadata(ErrorReportFn);
 
+  // Rewrite uses of globals to their replacement declare target globals.
+  applyDeclareTargetGlobalReplacements();
+
   if (Config.EmitLLVMUsedMetaInfo.value_or(false)) {
     std::vector<WeakTrackingVH> LLVMCompilerUsed = {
         M.getGlobalVariable("__openmp_nvptx_data_transfer_temporary_storage")};
@@ -1051,6 +1054,92 @@ void OpenMPIRBuilder::finalize(Function *Fn) {
 
 bool OpenMPIRBuilder::isFinalized() { return IsFinalized; }
 
+void OpenMPIRBuilder::registerDeclareTargetGlobalReplacement(
+    GlobalValue *Original, GlobalValue *Replacement) {
+  assert(Original && Replacement &&
+         "Null values provided to registerDeclareTargetGlobalReplacement");
+  DeclareTargetGlobalReplacements.push_back({Original, Replacement});
+}
+
+void OpenMPIRBuilder::applyDeclareTargetGlobalReplacements() {
+  for (DeclareTargetGlobalReplacement &R : DeclareTargetGlobalReplacements) {
+    GlobalValue *OldGV = R.Original;
+    GlobalValue *NewGV = R.Replacement;
+    if (!OldGV || !NewGV)
+      continue;
+
+    // The replacement global is a reference pointer that holds the
+    // address of the device-resident storage. Every use must load the
+    // reference pointer first and use the loaded address.
+    //
+    // Constant expression users (e.g. a constant GEP embedded in another
+    // global's initializer or in an instruction) cannot have a load inserted
+    // in place, so first expand any constant-expression users that live inside
+    // functions into instructions. Any remaining constant users are handled
+    // via a direct constant rewrite below as we cannot materialize a load
+    // there.
+    //
+    // NOTE: We extend the constant rewrite to module scope, as we replace all
+    // usages.
+    if (auto *OldConst = dyn_cast<Constant>(OldGV))
+      convertUsersOfConstantsToInstructions(OldConst,
+                                            /*RestrictToFunc=*/nullptr,
+                                            /*RemoveDeadConstants=*/false);
+
+    IRBuilderBase::InsertPointGuard Guard(Builder);
+    SmallVector<User *, 16> Users(OldGV->users());
+    for (User *U : Users) {
+      auto *Insn = dyn_cast<Instruction>(U);
+      if (!Insn)
+        continue;
+
+      // A PHI node cannot have a load inserted immediately before it, as PHIs
+      // must remain grouped at the top of their basic block. So we need to
+      // make sure any loads we emit are generated in the preceding edge, a
+      // PHI may reference the global on more than one edge, so every matching
+      // slot must be handled.
+      if (auto *PHI = dyn_cast<PHINode>(Insn)) {
+        for (unsigned I = 0, E = PHI->getNumIncomingValues(); I < E; ++I) {
+          if (PHI->getIncomingValue(I) != OldGV)
+            continue;
+
+          BasicBlock *IncomingBB = PHI->getIncomingBlock(I);
+          Builder.SetInsertPoint(IncomingBB->getTerminator());
+          Builder.SetCurrentDebugLocation(PHI->getDebugLoc());
+          LoadInst *EdgeLoad = Builder.CreateLoad(NewGV->getType(), NewGV);
+          PHI->setIncomingValue(I, EdgeLoad);
+        }
+        continue;
+      }
+
+      Builder.SetInsertPoint(Insn);
+      Builder.SetCurrentDebugLocation(Insn->getDebugLoc());
+      LoadInst *Load = Builder.CreateLoad(NewGV->getType(), NewGV);
+
+      // The replacement declare target global lives in the default address
+      // space, whereas the original global may reside in a non-default
+      // address space. In that case the initial lowering may have
+      // emitted an addrspacecast that is no longer valid.  Replace the
+      // whole addrspacecast with the load and erase it rather than
+      // feeding the load back into the (now pointless) cast.
+      if (auto *ASC = dyn_cast<AddrSpaceCastInst>(Insn)) {
+        unsigned NewGVAS = NewGV->getType()->getPointerAddressSpace();
+        unsigned OldGVAS = OldGV->getType()->getPointerAddressSpace();
+        unsigned DestAS = ASC->getType()->getPointerAddressSpace();
+        if (NewGVAS == 0 && DestAS == 0 && NewGVAS != OldGVAS) {
+          ASC->replaceAllUsesWith(Load);
+          ASC->eraseFromParent();
+          continue;
+        }
+      }
+
+      Insn->replaceUsesOfWith(OldGV, Load);
+    }
+  }
+
+  DeclareTargetGlobalReplacements.clear();
+}
+
 OpenMPIRBuilder::~OpenMPIRBuilder() {
   assert(OutlineInfos.empty() && "There must be no outstanding outlinings");
 }
diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index 1c50ff192c3d5..dcaeec51c5b7a 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -8106,67 +8106,6 @@ static void getTargetEntryUniqueInfo(llvm::TargetRegionEntryInfo &targetInfo,
       ompBuilder.getTargetEntryUniqueInfo(fileInfoCallBack, vfs, parentName);
 }
 
-static void
-handleDeclareTargetMapVar(MapInfoData &mapData,
-                          LLVM::ModuleTranslation &moduleTranslation,
-                          llvm::IRBuilderBase &builder, llvm::Function *func) {
-  assert(moduleTranslation.getOpenMPBuilder()->Config.isTargetDevice() &&
-         "function only supported for target device codegen");
-  llvm::IRBuilderBase::InsertPointGuard guard(builder);
-  for (size_t i = 0; i < mapData.MapClause.size(); ++i) {
-    // In the case of declare target mapped variables, the basePointer is
-    // the reference pointer generated by the convertDeclareTargetAttr
-    // method. Whereas the kernelValue is the original variable, so for
-    // the device we must replace all uses of this original global variable
-    // (stored in kernelValue) with the reference pointer (stored in
-    // basePointer for declare target mapped variables), as for device the
-    // data is mapped into this reference pointer and should be loaded
-    // from it, the original variable is discarded. On host both exist and
-    // metadata is generated (elsewhere in the convertDeclareTargetAttr)
-    // function to link the two variables in the runtime and then both the
-    // reference pointer and the pointer are assigned in the kernel argument
-    // structure for the host.
-    if (!mapData.IsDeclareTarget[i])
-      continue;
-    // If the original map value is a constant, then we have to make sure all
-    // of it's uses within the current kernel/function that we are going to
-    // rewrite are converted to instructions, as we will be altering the old
-    // use (OriginalValue) from a constant to an instruction, which will be
-    // illegal and ICE the compiler if the user is a constant expression of
-    // some kind e.g. a constant GEP.
-    if (auto *constant = dyn_cast<llvm::Constant>(mapData.OriginalValue[i]))
-      convertUsersOfConstantsToInstructions(constant, func, false);
-
-    // The users iterator will get invalidated if we modify an element,
-    // so we populate this vector of uses to alter each user on an
-    // individual basis to emit its own load (rather than one load for
-    // all).
-    llvm::SmallVector<llvm::User *> userVec;
-    for (llvm::User *user : mapData.OriginalValue[i]->users())
-      userVec.push_back(user);
-
-    for (llvm::User *user : userVec) {
-      auto *insn = dyn_cast<llvm::Instruction>(user);
-      if (!insn || insn->getFunction() != func)
-        continue;
-      auto mapOp = cast<omp::MapInfoOp>(mapData.MapClause[i]);
-      llvm::Value *substitute = mapData.BasePointers[i];
-      auto declTarPtr =
-          mapOp.getVarPtrPtr() ? mapOp.getVarPtrPtr() : mapOp.getVarPtr();
-      if (isDeclareTargetLink(declTarPtr) ||
-          (isDeclareTargetTo(declTarPtr) &&
-           moduleTranslation.getOpenMPBuilder()
-               ->Config.hasRequiresUnifiedSharedMemory())) {
-        builder.SetCurrentDebugLocation(insn->getDebugLoc());
-        substitute = builder.CreateLoad(mapData.BasePointers[i]->getType(),
-                                        mapData.BasePointers[i]);
-        cast<llvm::LoadInst>(substitute)->moveBefore(insn->getIterator());
-      }
-      user->replaceUsesOfWith(mapData.OriginalValue[i], substitute);
-    }
-  }
-}
-
 // The createDeviceArgumentAccessor function generates
 // instructions for retrieving (acessing) kernel
 // arguments inside of the device kernel for use by
@@ -9126,12 +9065,6 @@ convertOmpTarget(Operation &opInst, llvm::IRBuilderBase &builder,
   if (dds.DepArray)
     builder.CreateFree(dds.DepArray);
 
-  // Remap access operations to declare target reference pointers for the
-  // device, essentially generating extra loadop's as necessary
-  if (moduleTranslation.getOpenMPBuilder()->Config.isTargetDevice())
-    handleDeclareTargetMapVar(mapData, moduleTranslation, builder,
-                              llvmOutlinedFn);
-
   return success();
 }
 
@@ -9283,12 +9216,11 @@ convertDeclareTargetAttr(Operation *op, mlir::omp::DeclareTargetAttr attribute,
       if (ompBuilder->Config.isTargetDevice() &&
           (captureClause == omp::DeclareTargetCaptureClause::link ||
            requiresUSM)) {
-        llvm::Type *ptrTy = gVal->getType();
-        // For USM the global type becomes a pointer handle, as opposed to the
-        // globals original type.
-        if (requiresUSM)
-          ptrTy = llvm::PointerType::get(llvmModule->getContext(), 0);
-        bool addrGlobalCreated = ompBuilder->getAddrOfDeclareTargetVar(
+        // For USM and link we generate a global reference pointer in the
+        // default address space (e.g address space 0), as opposed to the
+        // globals original type and address space.
+        llvm::Type *ptrTy = llvm::PointerType::get(llvmModule->getContext(), 0);
+        llvm::Constant *refPtr = ompBuilder->getAddrOfDeclareTargetVar(
             captureClauseKind, deviceClause, isDeclaration, isExternallyVisible,
             ompBuilder->getTargetEntryUniqueInfo(fileInfoCallBack, vfs),
             mangledName, generatedRefs, /*OpenMPSimd*/ false, targetTriple,
@@ -9299,8 +9231,16 @@ convertDeclareTargetAttr(Operation *op, mlir::omp::DeclareTargetAttr attribute,
         // linkage to optimize out the unneeded full-variable storage later,
         // since we can't prevent the LLVM dialect from generating globals
         // without also breaking target lowering.
-        if (addrGlobalCreated)
+        if (refPtr) {
           gVar->setLinkage(llvm::GlobalValue::InternalLinkage);
+
+          // Register the (original global, reference pointer) pair so that the
+          // OpenMPIRBuilder can rewrite uses of the original global during
+          // finalization.
+          if (auto *newGV =
+                  dyn_cast<llvm::GlobalValue>(refPtr->stripPointerCasts()))
+            ompBuilder->registerDeclareTargetGlobalReplacement(gVal, newGV);
+        }
       }
 
       // Mark 'device_type(host) enter(...)' variables as external in the device
diff --git a/mlir/test/Target/LLVMIR/omptarget-declare-target-module-rewrite-link-device.mlir b/mlir/test/Target/LLVMIR/omptarget-declare-target-module-rewrite-link-device.mlir
new file mode 100644
index 0000000000000..8d7f7c9d3c3c5
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/omptarget-declare-target-module-rewrite-link-device.mlir
@@ -0,0 +1,48 @@
+// RUN: mlir-translate -mlir-to-llvmir %s | FileCheck %s
+
+// This test verifies the module-scope declare target global use rewrite
+// mechanism for a `declare target link` variable when compiling for device.
+//
+// The declare target global is used both directly inside of a target region
+// and indirectly inside of a declare target function that is invoked from
+// within that target region. Because the rewrite is now applied at module
+// scope during OpenMPIRBuilder finalization (rather than only within the
+// outlined target region), the use of the original global inside of the
+// declare target function should also be rewritten to load from the generated
+// reference pointer.
+
+module attributes {llvm.target_triple = "amdgcn-amd-amdhsa", omp.is_target_device = true} {
+  // CHECK-DAG: @_QMtest_0Esp_decl_tgt_ref_ptr = weak global ptr null, align 8
+  llvm.mlir.global external @_QMtest_0Esp() {addr_space = 0 : i32, omp.declare_target = #omp.declaretarget<device_type = (any), capture_clause = (link)>} : i32 {
+    %0 = llvm.mlir.constant(0 : i32) : i32
+    llvm.return %0 : i32
+  }
+
+  // CHECK-LABEL: define {{.*}} @_QMtest_0Puse_global
+  // CHECK: %[[REF:.*]] = load ptr, ptr @_QMtest_0Esp_decl_tgt_ref_ptr, align 8
+  // CHECK: store i32 2, ptr %[[REF]], align 4
+  llvm.func @_QMtest_0Puse_global() attributes {omp.declare_target = #omp.declaretarget<device_type = (any), capture_clause = (enter)>} {
+    %0 = llvm.mlir.addressof @_QMtest_0Esp : !llvm.ptr
+    %1 = llvm.mlir.constant(2 : i32) : i32
+    llvm.store %1, %0 : i32, !llvm.ptr
+    llvm.return
+  }
+
+  llvm.func @_QQmain() attributes {} {
+    %0 = llvm.mlir.addressof @_QMtest_0Esp : !llvm.ptr
+
+    // CHECK-DAG:   omp.target:
+    // CHECK-DAG: %[[V:.*]] = load ptr, ptr @_QMtest_0Esp_decl_tgt_ref_ptr, align 8
+    // CHECK-DAG: store i32 1, ptr %[[V]], align 4
+    // CHECK-DAG: call void @_QMtest_0Puse_global()
+    %map = omp.map.info var_ptr(%0 : !llvm.ptr, i32) map_clauses(tofrom) capture(ByRef) -> !llvm.ptr {name = ""}
+    omp.target kernel_type(generic) map_entries(%map -> %arg0 : !llvm.ptr) {
+      %1 = llvm.mlir.constant(1 : i32) : i32
+      llvm.store %1, %arg0 : i32, !llvm.ptr
+      llvm.call @_QMtest_0Puse_global() : () -> ()
+      omp.terminator
+    }
+
+    llvm.return
+  }
+}
diff --git a/mlir/test/Target/LLVMIR/omptarget-declare-target-module-rewrite-phi-device.mlir b/mlir/test/Target/LLVMIR/omptarget-declare-target-module-rewrite-phi-device.mlir
new file mode 100644
index 0000000000000..72f8951773853
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/omptarget-declare-target-module-rewrite-phi-device.mlir
@@ -0,0 +1,39 @@
+// RUN: mlir-translate -mlir-to-llvmir %s | FileCheck %s
+
+// This test verifies the declare target global use rewrite mechanism when
+// the original global is consumed by a PHI node. Making sure we rewrite this
+// case correctly.
+
+module attributes {llvm.target_triple = "amdgcn-amd-amdhsa", omp.is_target_device = true} {
+  // CHECK-DAG: @_QMtest_0Esp_decl_tgt_ref_ptr = weak global ptr null, align 8
+  llvm.mlir.global external @_QMtest_0Esp() {addr_space = 0 : i32, omp.declare_target = #omp.declaretarget<device_type = (any), capture_clause = (link)>} : i32 {
+    %0 = llvm.mlir.constant(0 : i32) : i32
+    llvm.return %0 : i32
+  }
+
+  // CHECK-LABEL: define hidden void @_QMtest_0Puse_global
+  // CHECK: %[[LOAD0:.*]] = load ptr, ptr @_QMtest_0Esp_decl_tgt_ref_ptr, align 8
+  // CHECK: %[[LOAD1:.*]] = load ptr, ptr @_QMtest_0Esp_decl_tgt_ref_ptr, align 8
+  // CHECK: br i1 %{{.*}}, label %[[BB_A:.*]], label %[[BB_B:.*]]
+  // CHECK: [[BB_A]]:
+  // CHECK: %[[PHI_A:.*]] = phi ptr [ %[[LOAD1]], %{{.*}} ]
+  // CHECK: br label %[[MERGE:.*]]
+  // CHECK: [[BB_B]]:
+  // CHECK: %[[PHI_B:.*]] = phi ptr [ %[[LOAD0]], %{{.*}} ]
+  // CHECK: br label %[[MERGE]]
+  // CHECK: [[MERGE]]:
+  // CHECK: %[[PHI:.*]] = phi ptr [ %[[PHI_B]], %[[BB_B]] ], [ %[[PHI_A]], %[[BB_A]] ]
+  // CHECK: store i32 2, ptr %[[PHI]], align 4
+  llvm.func @_QMtest_0Puse_global(%cond : i1) attributes {omp.declare_target = #omp.declaretarget<device_type = (any), capture_clause = (enter)>} {
+    %0 = llvm.mlir.addressof @_QMtest_0Esp : !llvm.ptr
+    llvm.cond_br %cond, ^bb1(%0 : !llvm.ptr), ^bb2(%0 : !llvm.ptr)
+  ^bb1(%arg1 : !llvm.ptr):
+    llvm.br ^bb3(%arg1 : !llvm.ptr)
+  ^bb2(%arg2 : !llvm.ptr):
+    llvm.br ^bb3(%arg2 : !llvm.ptr)
+  ^bb3(%arg3 : !llvm.ptr):
+    %1 = llvm.mlir.constant(2 : i32) : i32
+    llvm.store %1, %arg3 : i32, !llvm.ptr
+    llvm.return
+  }
+}
diff --git a/mlir/test/Target/LLVMIR/omptarget-declare-target-module-rewrite-phi-nested-device.mlir b/mlir/test/Target/LLVMIR/omptarget-declare-target-module-rewrite-phi-nested-device.mlir
new file mode 100644
index 0000000000000..ef1e5ffe029f3
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/omptarget-declare-target-module-rewrite-phi-nested-device.mlir
@@ -0,0 +1,60 @@
+// RUN: mlir-translate -mlir-to-llvmir %s | FileCheck %s
+
+// A more complicated exercise of the declare target global use rewrite when
+//  the original global is consumed by multiple PHI nodes, some nested, and
+//  additionally the original global lives in a non-default address space
+//  (address space 2)
+
+module attributes {llvm.target_triple = "amdgcn-amd-amdhsa", omp.is_target_device = true} {
+  // CHECK-DAG: @_QMtest_0Esp_decl_tgt_ref_ptr = weak global ptr null, align 8
+  llvm.mlir.global external @_QMtest_0Esp() {addr_space = 2 : i32, omp.declare_target = #omp.declaretarget<device_type = (any), capture_clause = (link)>} : i32 {
+    %0 = llvm.mlir.constant(0 : i32) : i32
+    llvm.return %0 : i32
+  }
+
+  // CHECK-LABEL: define hidden void @_QMtest_0Puse_global_nested
+  //
+  // CHECK: %[[L0:.*]] = load ptr, ptr @_QMtest_0Esp_decl_tgt_ref_ptr, align 8
+  // CHECK: br i1 %{{.*}}, label %[[A:.*]], label %[[B:.*]]
+  // CHECK: [[A]]:
+  // CHECK: %[[PA:.*]] = phi ptr [ %[[L0]], %[[ENTRY:.*]] ]
+  // CHECK: br label %[[M1:.*]]
+  // CHECK: [[B]]:
+  // CHECK: %[[PB:.*]] = phi ptr [ %[[L0]], %[[ENTRY]] ]
+  // CHECK: br label %[[M1]]
+  // CHECK: [[M1]]:
+  // CHECK: %[[PHI1:.*]] = phi ptr [ %[[PB]], %[[B]] ], [ %[[PA]], %[[A]] ]
+  //
+  // CHECK: %[[L1:.*]] = load ptr, ptr @_QMtest_0Esp_decl_tgt_ref_ptr, align 8
+  // CHECK: br i1 %{{.*}}, label %[[C:.*]], label %[[D:.*]]
+  // CHECK: [[C]]:
+  // CHECK: %[[PC:.*]] = phi ptr [ %[[PHI1]], %[[M1]] ]
+  // CHECK: br label %[[M2:.*]]
+  // CHECK: [[D]]:
+  // CHECK: %[[PD:.*]] = phi ptr [ %[[L1]], %[[M1]] ]
+  // CHECK: br label %[[M2]]
+  // CHECK: [[M2]]:
+  // CHECK: %[[PHI2:.*]] = phi ptr [ %[[PD]], %[[D]] ], [ %[[PC]], %[[C]] ]
+  // CHECK: store i32 3, ptr %[[PHI2]], align 4
+  llvm.func @_QMtest_0Puse_global_nested(%cond1 : i1, %cond2 : i1) attributes {omp.declare_target = #omp.declaretarget<device_type = (any), capture_clause = (enter)>} {
+    %g = llvm.mlir.addressof @_QMtest_0Esp : !llvm.ptr<2>
+    %gc = llvm.addrspacecast %g : !llvm.ptr<2> to !llvm.ptr
+    llvm.cond_br %cond1, ^bb1(%gc : !llvm.ptr), ^bb2(%gc : !llvm.ptr)
+  ^bb1(%a : !llvm.ptr):
+    llvm.br ^merge1(%a : !llvm.ptr)
+  ^bb2(%b : !llvm.ptr):
+    llvm.br ^merge1(%b : !llvm.ptr)
+  ^merge1(%m1 : !llvm.ptr):
+    %g2 = llvm.mlir.addressof @_QMtest_0Esp : !llvm.ptr<2>
+    %g2c = llvm.addrspacecast %g2 : !llvm.ptr<2> to !llvm.ptr
+    llvm.cond_br %cond2, ^bb3(%m1 : !llvm.ptr), ^bb4(%g2c : !llvm.ptr)
+  ^bb3(%c : !llvm.ptr):
+    llvm.br ^merge2(%c : !llvm.ptr)
+  ^bb4(%d : !llvm.ptr):
+    llvm.br ^merge2(%d : !llvm.ptr)
+  ^merge2(%m2 : !llvm.ptr):
+    %v = llvm.mlir.constant(3 : i32) : i32
+    llvm.store %v, %m2 : i32, !llvm.ptr
+    llvm.return
+  }
+}
diff --git a/mlir/test/Target/LLVMIR/omptarget-declare-target-module-rewrite-to-device.mlir b/mlir/test/Target/LLVMIR/omptarget-declare-target-module-rewrite-to-device.mlir
new file mode 100644
index 0000000000000..a7e43da88756d
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/omptarget-declare-target-module-rewrite-to-device.mlir
@@ -0,0 +1,60 @@
+// RUN: mlir-translate -mlir-to-llvmir %s | FileCheck %s
+
+// This test verifies that the module-scope declare target global use rewrite
+// mechanism is NOT applied for a regular `declare target to`/`enter` variable
+// (i.e. without unified shared memory) when compiling for device. In this
+// configuration no reference pointer global is generated, so uses of the
+// original global must remain direct references to the global itself, both
+// inside of the target region and inside of a declare target function invoked
+// from the target region. No `_decl_tgt_ref_ptr` global should be created and
+// no load-from-reference-pointer should be emitted.
+
+module attributes {llvm.target_triple = "amdgcn-amd-amdhsa", omp.is_target_device = true} {
+  // CHECK-NOT: @_QMtest_0Evar_to_decl_tgt_ref_ptr
+  // CHECK-NOT: @_QMtest_0Evar_enter_decl_tgt_ref_ptr
+  // CHECK-DAG: @_QMtest_0Evar_to = global i32
+  llvm.mlir.global external @_QMtest_0Evar_to() {addr_space = 0 : i32, omp.declare_target = #omp.declaretarget<device_type = (any), capture_clause = (to)>} : i32 {
+    %0 = llvm.mlir.constant(1 : i32) : i32
+    llvm.return %0 : i32
+  }
+
+  // CHECK-DAG: @_QMtest_0Evar_enter = global i32
+  llvm.mlir.global external @_QMtest_0Evar_enter() {addr_space = 0 : i32, omp.declare_target = #omp.declaretarget<device_type = (any), capture_clause = (enter)>} : i32 {
+    %0 = llvm.mlir.constant(2 : i32) : i32
+    llvm.return %0 : i32
+  }
+
+  // CHECK-LABEL: define {{.*}} @_QMtest_0Puse_global
+  // CHECK-NOT: load ptr, ptr @_QMtest_0Evar_to_decl_tgt_ref_ptr
+  // CHECK-NOT: load ptr, ptr @_QMtest_0Evar_enter_decl_tgt_ref_ptr
+  // CHECK-DAG: store i32 100, ptr @_QMtest_0Evar_to, align 4
+  // CHECK-DAG: store i32 200, ptr @_QMtest_0Evar_enter, align 4
+  llvm.func @_QMtest_0Puse_global() attributes {omp.declare_target = #omp.declaretarget<device_type = (any), capture_clause = (enter)>} {
+    %0 = llvm.mlir.addressof @_QMtest_0Evar_to : !llvm.ptr
+    %1 = llvm.mlir.addressof @_QMtest_0Evar_enter : !llvm.ptr
+    %c100 = llvm.mlir.constant(100 : i32) : i32
+    %c200 = llvm.mlir.constant(200 : i32) : i32
+    llvm.store %c100, %0 : i32, !llvm.ptr
+    llvm.store %c200, %1 : i32, !llvm.ptr
+    llvm.return
+  }
+
+  llvm.func @test_declare_target() attributes {} {
+    %0 = llvm.mlir.addressof @_QMtest_0Evar_to : !llvm.ptr
+    %1 = llvm.mlir.addressof @_QMtest_0Evar_enter : !llvm.ptr
+    // CHECK-DAG: store i32 10, ptr @_QMtest_0Evar_to, align 4
+    // CHECK-DAG: store i32 20, ptr @_QMtest_0Evar_enter, align 4
+    // CHECK-DAG: call void @_QMtest_0Puse_global()
+    %map0 = omp.map.info var_ptr(%0 : !llvm.ptr, i32) map_clauses(tofrom) capture(ByRef) -> !llvm.ptr {name = ""}
+    %map1 = omp.map.info var_ptr(%1 : !llvm.ptr, i32) map_clauses(tofrom) capture(ByRef) -> !llvm.ptr {name = ""}
+    omp.target kernel_type(generic) map_entries(%map0 -> %arg0, %map1 -> %arg1 : !llvm.ptr, !llvm.ptr) {
+      %c10 = llvm.mlir.constant(10 : i32) : i32
+      %c20 = llvm.mlir.constant(20 : i32) : i32
+      llvm.store %c10, %arg0 : i32, !llvm.ptr
+      llvm.store %c20, %arg1 : i32, !llvm.ptr
+      llvm.call @_QMtest_0Puse_global() : () -> ()
+      omp.terminator
+    }
+    llvm.return
+  }
+}
diff --git a/mlir/test/Target/LLVMIR/omptarget-declare-target-module-rewrite-usm-device.mlir b/mlir/test/Target/LLVMIR/omptarget-declare-target-module-rewrite-usm-device.mlir
new file mode 100644
index 0000000000000..c2acb294dec94
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/omptarget-declare-target-module-rewrite-usm-device.mlir
@@ -0,0 +1,62 @@
+// RUN: mlir-translate -mlir-to-llvmir %s | FileCheck %s
+
+// This test verifies the module-scope declare target global use rewrite
+// mechanism for `declare target to` and `declare target enter` variables when
+// unified shared memory is required and compiling for device. In this
+// configuration a reference pointer is generated for the to/enter variables
+// (as with link), so uses of the original global must be rewritten to load
+// from the reference pointer at module scope.
+//
+// As with the link test, the globals are used both directly inside of a target
+// region and indirectly inside of a declare target function invoked from that
+// region, and both use-sites must be rewritten.
+
+module attributes {llvm.target_triple = "amdgcn-amd-amdhsa", omp.is_target_device = true, omp.requires = #omp<clause_requires unified_shared_memory>} {
+  // CHECK-DAG: @_QMtest_0Evar_to_usm_decl_tgt_ref_ptr = weak global ptr null, align 8
+  llvm.mlir.global external @_QMtest_0Evar_to_usm() {addr_space = 0 : i32, omp.declare_target = #omp.declaretarget<device_type = (any), capture_clause = (to)>} : i32 {
+    %0 = llvm.mlir.constant(1 : i32) : i32
+    llvm.return %0 : i32
+  }
+
+  // CHECK-DAG: @_QMtest_0Evar_enter_usm_decl_tgt_ref_ptr = weak global ptr null, align 8
+  llvm.mlir.global external @_QMtest_0Evar_enter_usm() {addr_space = 0 : i32, omp.declare_target = #omp.declaretarget<device_type = (any), capture_clause = (enter)>} : i32 {
+    %0 = llvm.mlir.constant(2 : i32) : i32
+    llvm.return %0 : i32
+  }
+
+  // CHECK-LABEL: define {{.*}} @_QMtest_0Puse_global
+  // CHECK-DAG: %[[TO_REF:.*]] = load ptr, ptr @_QMtest_0Evar_to_usm_decl_tgt_ref_ptr, align 8
+  // CHECK-DAG: store i32 100, ptr %[[TO_REF]], align 4
+  // CHECK-DAG: %[[ENTER_REF:.*]] = load ptr, ptr @_QMtest_0Evar_enter_usm_decl_tgt_ref_ptr, align 8
+  // CHECK-DAG: store i32 200, ptr %[[ENTER_REF]], align 4
+  llvm.func @_QMtest_0Puse_global() attributes {omp.declare_target = #omp.declaretarget<device_type = (any), capture_clause = (enter)>} {
+    %0 = llvm.mlir.addressof @_QMtest_0Evar_to_usm : !llvm.ptr
+    %1 = llvm.mlir.addressof @_QMtest_0Evar_enter_usm : !llvm.ptr
+    %c100 = llvm.mlir.constant(100 : i32) : i32
+    %c200 = llvm.mlir.constant(200 : i32) : i32
+    llvm.store %c100, %0 : i32, !llvm.ptr
+    llvm.store %c200, %1 : i32, !llvm.ptr
+    llvm.return
+  }
+
+  llvm.func @test_usm_declare_target() attributes {} {
+    %0 = llvm.mlir.addressof @_QMtest_0Evar_to_usm : !llvm.ptr
+    %1 = llvm.mlir.addressof @_QMtest_0Evar_enter_usm : !llvm.ptr
+    // CHECK-DAG: %[[TO_VAR:.*]] = load ptr, ptr @_QMtest_0Evar_to_usm_decl_tgt_ref_ptr, align 8
+    // CHECK-DAG: store i32 10, ptr %[[TO_VAR]], align 4
+    // CHECK-DAG: %[[ENTER_VAR:.*]] = load ptr, ptr @_QMtest_0Evar_enter_usm_decl_tgt_ref_ptr, align 8
+    // CHECK-DAG: store i32 20, ptr %[[ENTER_VAR]], align 4
+    // CHECK-DAG: call void @_QMtest_0Puse_global()
+    %map0 = omp.map.info var_ptr(%0 : !llvm.ptr, i32) map_clauses(tofrom) capture(ByRef) -> !llvm.ptr {name = ""}
+    %map1 = omp.map.info var_ptr(%1 : !llvm.ptr, i32) map_clauses(tofrom) capture(ByRef) -> !llvm.ptr {name = ""}
+    omp.target kernel_type(generic) map_entries(%map0 -> %arg0, %map1 -> %arg1 : !llvm.ptr, !llvm.ptr) {
+      %c10 = llvm.mlir.constant(10 : i32) : i32
+      %c20 = llvm.mlir.constant(20 : i32) : i32
+      llvm.store %c10, %arg0 : i32, !llvm.ptr
+      llvm.store %c20, %arg1 : i32, !llvm.ptr
+      llvm.call @_QMtest_0Puse_global() : () -> ()
+      omp.terminator
+    }
+    llvm.return
+  }
+}



More information about the llvm-commits mailing list