[ReplaceConstant] Extract code for expanding users of constant (NFC)

AMDGPU implements some handy code for expanding all constexpr
users of LDS globals. Extract the core logic into ReplaceConstant,
so that it can be reused elsewhere.

2 files changed, 62 insertions, 54 deletions

diff --git a/llvm/lib/IR/ReplaceConstant.cpp b/llvm/lib/IR/ReplaceConstant.cpp
index 069da26..54988309 100644
--- a/llvm/lib/IR/ReplaceConstant.cpp
+++ b/llvm/lib/IR/ReplaceConstant.cpp
@@ -12,6 +12,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/IR/ReplaceConstant.h"
+#include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/Instructions.h"
@@ -133,4 +134,61 @@ void collectConstantExprPaths(
   }
 }
 
+bool convertUsersOfConstantsToInstructions(ArrayRef<Constant *> Consts) {
+  // Find all ConstantExpr that are direct users Consts.
+  SmallVector<ConstantExpr *> Stack;
+  for (Constant *C : Consts)
+    for (User *U : C->users())
+      if (ConstantExpr *CE = dyn_cast<ConstantExpr>(U))
+        Stack.push_back(CE);
+
+  // Expand to include constexpr users of direct users
+  SetVector<ConstantExpr *> ConstExprUsers;
+  while (!Stack.empty()) {
+    ConstantExpr *V = Stack.pop_back_val();
+    if (ConstExprUsers.contains(V))
+      continue;
+
+    ConstExprUsers.insert(V);
+
+    for (auto *Nested : V->users())
+      if (ConstantExpr *CE = dyn_cast<ConstantExpr>(Nested))
+        Stack.push_back(CE);
+  }
+
+  // Find all instructions that use any of the ConstExpr users
+  SetVector<Instruction *> InstructionWorklist;
+  for (ConstantExpr *CE : ConstExprUsers)
+    for (User *U : CE->users())
+      if (auto *I = dyn_cast<Instruction>(U))
+        InstructionWorklist.insert(I);
+
+  // Replace those ConstExpr operands with instructions
+  bool Changed = false;
+  while (!InstructionWorklist.empty()) {
+    Instruction *I = InstructionWorklist.pop_back_val();
+    for (Use &U : I->operands()) {
+      auto *BI = I;
+      if (auto *Phi = dyn_cast<PHINode>(I)) {
+        BasicBlock *BB = Phi->getIncomingBlock(U);
+        BasicBlock::iterator It = BB->getFirstInsertionPt();
+        assert(It != BB->end() && "Unexpected empty basic block");
+        BI = &*It;
+      }
+
+      if (ConstantExpr *C = dyn_cast<ConstantExpr>(U.get())) {
+        if (ConstExprUsers.contains(C)) {
+          Changed = true;
+          Instruction *NI = C->getAsInstruction(BI);
+          InstructionWorklist.insert(NI);
+          U.set(NI);
+          C->removeDeadConstantUsers();
+        }
+      }
+    }
+  }
+
+  return Changed;
+}
+
 } // namespace llvm
diff --git a/llvm/lib/Target/AMDGPU/AMDGPULowerModuleLDSPass.cpp b/llvm/lib/Target/AMDGPU/AMDGPULowerModuleLDSPass.cpp
index 34dc86e..0804fda 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPULowerModuleLDSPass.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPULowerModuleLDSPass.cpp
@@ -134,6 +134,7 @@
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicsAMDGPU.h"
 #include "llvm/IR/MDBuilder.h"
+#include "llvm/IR/ReplaceConstant.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/CommandLine.h"
@@ -248,63 +249,12 @@ class AMDGPULowerModuleLDS : public ModulePass {
     //         if (constantExprUsesLDS(Op))
     //           replaceConstantExprInFunction(I, Op);
 
-    bool Changed = false;
-
-    // Find all ConstantExpr that are direct users of an LDS global
-    SmallVector<ConstantExpr *> Stack;
+    SmallVector<Constant *> LDSGlobals;
     for (auto &GV : M.globals())
       if (AMDGPU::isLDSVariableToLower(GV))
-        for (User *U : GV.users())
-          if (ConstantExpr *C = dyn_cast<ConstantExpr>(U))
-            Stack.push_back(C);
-
-    // Expand to include constexpr users of direct users
-    SetVector<ConstantExpr *> ConstExprUsersOfLDS;
-    while (!Stack.empty()) {
-      ConstantExpr *V = Stack.pop_back_val();
-      if (ConstExprUsersOfLDS.contains(V))
-        continue;
-
-      ConstExprUsersOfLDS.insert(V);
-
-      for (auto *Nested : V->users())
-        if (ConstantExpr *CE = dyn_cast<ConstantExpr>(Nested))
-          Stack.push_back(CE);
-    }
-
-    // Find all instructions that use any of the ConstExpr users of LDS
-    SetVector<Instruction *> InstructionWorklist;
-    for (ConstantExpr *CE : ConstExprUsersOfLDS)
-      for (User *U : CE->users())
-        if (auto *I = dyn_cast<Instruction>(U))
-          InstructionWorklist.insert(I);
-
-    // Replace those ConstExpr operands with instructions
-    while (!InstructionWorklist.empty()) {
-      Instruction *I = InstructionWorklist.pop_back_val();
-      for (Use &U : I->operands()) {
-
-        auto *BI = I;
-        if (auto *Phi = dyn_cast<PHINode>(I)) {
-          BasicBlock *BB = Phi->getIncomingBlock(U);
-          BasicBlock::iterator It = BB->getFirstInsertionPt();
-          assert(It != BB->end() && "Unexpected empty basic block");
-          BI = &(*(It));
-        }
+        LDSGlobals.push_back(&GV);
 
-        if (ConstantExpr *C = dyn_cast<ConstantExpr>(U.get())) {
-          if (ConstExprUsersOfLDS.contains(C)) {
-            Changed = true;
-            Instruction *NI = C->getAsInstruction(BI);
-            InstructionWorklist.insert(NI);
-            U.set(NI);
-            C->removeDeadConstantUsers();
-          }
-        }
-      }
-    }
-
-    return Changed;
+    return convertUsersOfConstantsToInstructions(LDSGlobals);
   }
 
 public: