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

[Sairudra More]

================
@@ -1051,6 +1054,73 @@ 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;
+
+      Builder.SetInsertPoint(Insn);
----------------
Saieiei wrote:

Could you handle `PHINode` users separately here? `convertUsersOfConstantsToInstructions()` does not expand a PHI whose incoming value is directly `OldGV`, so this inserts the load immediately before the PHI and produces invalid LLVM IR. An MLIR block argument fed by llvm.mlir.addressof can create this shape. Please materialize the load on each matching incoming edge and add regression coverage.

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