[llvm] 576060f - [ReplaceConstant] Extract code for expanding users of constant (NFC)

[Nikita Popov via llvm-commits]

Author: Nikita Popov
Date: 2023-03-03T16:09:06+01:00
New Revision: 576060fb41c2de536a24d5d421b1cd8942f20b64

URL: https://github.com/llvm/llvm-project/commit/576060fb41c2de536a24d5d421b1cd8942f20b64
DIFF: https://github.com/llvm/llvm-project/commit/576060fb41c2de536a24d5d421b1cd8942f20b64.diff

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

Added: 
    

Modified: 
    llvm/include/llvm/IR/ReplaceConstant.h
    llvm/lib/IR/ReplaceConstant.cpp
    llvm/lib/Target/AMDGPU/AMDGPULowerModuleLDSPass.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/IR/ReplaceConstant.h b/llvm/include/llvm/IR/ReplaceConstant.h
index 1d6b10d9a78bc..e074ca6549aae 100644
--- 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 069da26e63b12..549883094802b 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 34dc86ecf3884..0804fda79d607 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: