[llvm] [IROutliner] Correctly Replace Outlined Constants with Arguments (PR #179885)

[Chuanqi Xu via llvm-commits]

================
@@ -1915,24 +1914,30 @@ replaceArgumentUses(OutlinableRegion &Region,
 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;
     assert(OutlinedFunction && "Overall Function is not defined?");
     Constant *CST = Const.second;
     Argument *Arg = Group.OutlinedFunction->getArg(AggArgIdx);
-    // Identify the argument it will be elevated to, and replace instances of
-    // that constant in the function.
-    VMap[CST] = Arg;
-    LLVM_DEBUG(dbgs() << "Replacing uses of constant " << *CST
+    LLVM_DEBUG(dbgs() << "Replacing instruction uses of constant " << *CST
                       << " in function " << *OutlinedFunction << " with "
                       << *Arg << '\n');
-  }
 
-  RemapFunction(*OutlinedFunction, VMap,
-                RF_NoModuleLevelChanges | RF_IgnoreMissingLocals);
+    for (auto &BB : *OutlinedFunction) {
+      for (Instruction &I : BB) {
+        // We specifically iterate over the direct operands of instructions, to
+        // avoid replacing constants that appear in `ConstantExpr`s
+        for (auto &U : I.operands()) {
+          if (auto *CCandidate = dyn_cast<Constant>(U)) {
+            if (CCandidate == CST)
+              U = Arg;
+          }
+        }
+      }
+    }
+  }
----------------
ChuanqiXu9 wrote:

The '}' looks suprising : ) I feel we should change the implementation of `RemapFunction` and add new flags for it (if they don't already have)

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