[llvm] IROutliner: Use ValueMapper to remap constants in a function (PR #134850)

[Matt Arsenault via llvm-commits]

https://github.com/arsenm updated https://github.com/llvm/llvm-project/pull/134850

>From 59edaaee2ea921062bdc4745207462efc3f53878 Mon Sep 17 00:00:00 2001
From: Matt Arsenault <Matthew.Arsenault at amd.com>
Date: Tue, 8 Apr 2025 19:24:12 +0700
Subject: [PATCH 1/2] IROutliner: Use ValueMapper to remap constants in a
 function

Stop relying on the uselist of constants, which also required filtering
the uses that are in the correct function. I'm not sure why this pass is
doing its own cloning instead of building the value map while doing its
cloning.
---
 llvm/lib/Transforms/IPO/IROutliner.cpp | 17 +++++++++--------
 1 file changed, 9 insertions(+), 8 deletions(-)

diff --git a/llvm/lib/Transforms/IPO/IROutliner.cpp b/llvm/lib/Transforms/IPO/IROutliner.cpp
index fd76d3a32049c..7e9ec4a04d5e3 100644
--- a/llvm/lib/Transforms/IPO/IROutliner.cpp
+++ b/llvm/lib/Transforms/IPO/IROutliner.cpp
@@ -24,6 +24,7 @@
 #include "llvm/IR/PassManager.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Transforms/IPO.h"
+#include "llvm/Transforms/Utils/ValueMapper.h"
 #include <optional>
 #include <vector>
 
@@ -1930,10 +1931,12 @@ replaceArgumentUses(OutlinableRegion &Region,
 /// \param Region [in] - The region of extracted code to be changed.
 void replaceConstants(OutlinableRegion &Region) {
   OutlinableGroup &Group = *Region.Parent;
+  Function *OutlinedFunction = Group.OutlinedFunction;
+  ValueToValueMapTy VMap;
+
   // Iterate over the constants that need to be elevated into arguments
   for (std::pair<unsigned, Constant *> &Const : Region.AggArgToConstant) {
     unsigned AggArgIdx = Const.first;
-    Function *OutlinedFunction = Group.OutlinedFunction;
     assert(OutlinedFunction && "Overall Function is not defined?");
     Constant *CST = Const.second;
     Argument *Arg = Group.OutlinedFunction->getArg(AggArgIdx);
@@ -1943,16 +1946,14 @@ void replaceConstants(OutlinableRegion &Region) {
     // TODO: If in the future constants do not have one global value number,
     // i.e. a constant 1 could be mapped to several values, this check will
     // have to be more strict.  It cannot be using only replaceUsesWithIf.
-
+    VMap[CST] = Arg;
     LLVM_DEBUG(dbgs() << "Replacing uses of constant " << *CST
                       << " in function " << *OutlinedFunction << " with "
-                      << *Arg << "\n");
-    CST->replaceUsesWithIf(Arg, [OutlinedFunction](Use &U) {
-      if (Instruction *I = dyn_cast<Instruction>(U.getUser()))
-        return I->getFunction() == OutlinedFunction;
-      return false;
-    });
+                      << *Arg << '\n');
   }
+
+  RemapFunction(*OutlinedFunction, VMap,
+                RF_NoModuleLevelChanges | RF_IgnoreMissingLocals);
 }
 
 /// It is possible that there is a basic block that already performs the same

>From 34bf25a5277f48cfeaecf66c63643382b995b798 Mon Sep 17 00:00:00 2001
From: Matt Arsenault <Matthew.Arsenault at amd.com>
Date: Wed, 9 Apr 2025 12:51:56 +0800
Subject: [PATCH 2/2] Remove todo

---
 llvm/lib/Transforms/IPO/IROutliner.cpp | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/llvm/lib/Transforms/IPO/IROutliner.cpp b/llvm/lib/Transforms/IPO/IROutliner.cpp
index 7e9ec4a04d5e3..8c5946fb07c65 100644
--- a/llvm/lib/Transforms/IPO/IROutliner.cpp
+++ b/llvm/lib/Transforms/IPO/IROutliner.cpp
@@ -1942,10 +1942,6 @@ void replaceConstants(OutlinableRegion &Region) {
     Argument *Arg = Group.OutlinedFunction->getArg(AggArgIdx);
     // Identify the argument it will be elevated to, and replace instances of
     // that constant in the function.
-
-    // TODO: If in the future constants do not have one global value number,
-    // i.e. a constant 1 could be mapped to several values, this check will
-    // have to be more strict.  It cannot be using only replaceUsesWithIf.
     VMap[CST] = Arg;
     LLVM_DEBUG(dbgs() << "Replacing uses of constant " << *CST
                       << " in function " << *OutlinedFunction << " with "