[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.

diff --git a/llvm/include/llvm/IR/ReplaceConstant.h b/llvm/include/llvm/IR/ReplaceConstant.h
index 1d6b10d9a78bc81c4ad669b3a7890d5ac9fe6999..e074ca6549aae203824c6672248fbf7a03dc2518 100644 (file)
--- a/llvm/include/llvm/IR/ReplaceConstant.h
+++ b/llvm/include/llvm/IR/ReplaceConstant.h
@@ -19,6 +19,8 @@
 
 namespace llvm {
 
+template <typename T> class ArrayRef;
+class Constant;
 class ConstantExpr;
 class Instruction;
 class Use;
@@ -54,6 +56,10 @@ void collectConstantExprPaths(
     Instruction *I, ConstantExpr *CE,
     std::map<Use *, std::vector<std::vector<ConstantExpr *>>> &CEPaths);
 
+/// Replace constant expressions users of the given constants with
+/// instructions. Return whether anything was changed.
+bool convertUsersOfConstantsToInstructions(ArrayRef<Constant *> Consts);
+
 } // end namespace llvm
 
 #endif // LLVM_IR_REPLACECONSTANT_H

diff --git a/llvm/lib/IR/ReplaceConstant.cpp b/llvm/lib/IR/ReplaceConstant.cpp
index 069da26e63b121ce959c6b6b3aebdffef130898c..549883094802b56b4d5bf7f6967a1aa1256a17b0 100644 (file)
--- 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 34dc86ecf3884f24c7819090742092a3ce4c87fe..0804fda79d6073dd6756d5d43ed0163b75dd7fb7 100644 (file)
--- 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: