[llvm] [llvm][NVPTX] Fix quadratic runtime in ProxyRegErasure (PR #105730)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-backend-nvptx

Author: Jeff Niu (Mogball)

<details>
<summary>Changes</summary>

This pass performs RAUW by walking the machine function for each RAUW operation. For large functions, this runtime in this pass starts to blow up. Linearize the pass by batching the RAUW ops at once.

---
Full diff: https://github.com/llvm/llvm-project/pull/105730.diff


1 Files Affected:

- (modified) llvm/lib/Target/NVPTX/NVPTXProxyRegErasure.cpp (+28-32) 


``````````diff
diff --git a/llvm/lib/Target/NVPTX/NVPTXProxyRegErasure.cpp b/llvm/lib/Target/NVPTX/NVPTXProxyRegErasure.cpp
index 258ae97a20d582..7fe6993fd33e4c 100644
--- a/llvm/lib/Target/NVPTX/NVPTXProxyRegErasure.cpp
+++ b/llvm/lib/Target/NVPTX/NVPTXProxyRegErasure.cpp
@@ -34,7 +34,6 @@ void initializeNVPTXProxyRegErasurePass(PassRegistry &);
 namespace {
 
 struct NVPTXProxyRegErasure : public MachineFunctionPass {
-public:
   static char ID;
   NVPTXProxyRegErasure() : MachineFunctionPass(ID) {
     initializeNVPTXProxyRegErasurePass(*PassRegistry::getPassRegistry());
@@ -49,23 +48,22 @@ struct NVPTXProxyRegErasure : public MachineFunctionPass {
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     MachineFunctionPass::getAnalysisUsage(AU);
   }
-
-private:
-  void replaceMachineInstructionUsage(MachineFunction &MF, MachineInstr &MI);
-
-  void replaceRegisterUsage(MachineInstr &Instr, MachineOperand &From,
-                            MachineOperand &To);
 };
 
 } // namespace
 
 char NVPTXProxyRegErasure::ID = 0;
 
-INITIALIZE_PASS(NVPTXProxyRegErasure, "nvptx-proxyreg-erasure", "NVPTX ProxyReg Erasure", false, false)
+INITIALIZE_PASS(NVPTXProxyRegErasure, "nvptx-proxyreg-erasure",
+                "NVPTX ProxyReg Erasure", false, false)
 
 bool NVPTXProxyRegErasure::runOnMachineFunction(MachineFunction &MF) {
   SmallVector<MachineInstr *, 16> RemoveList;
 
+  // ProxyReg instructions forward a register as another: `%dst = mov.iN %src`.
+  // Bulk RAUW the `%dst` registers in two passes over the machine function.
+  DenseMap<Register, Register> RAUWBatch;
+
   for (auto &BB : MF) {
     for (auto &MI : BB) {
       switch (MI.getOpcode()) {
@@ -74,44 +72,42 @@ bool NVPTXProxyRegErasure::runOnMachineFunction(MachineFunction &MF) {
       case NVPTX::ProxyRegI32:
       case NVPTX::ProxyRegI64:
       case NVPTX::ProxyRegF32:
-      case NVPTX::ProxyRegF64:
-        replaceMachineInstructionUsage(MF, MI);
+      case NVPTX::ProxyRegF64: {
+        auto &InOp = *MI.uses().begin();
+        auto &OutOp = *MI.defs().begin();
+        assert(InOp.isReg() && "ProxyReg input should be a register.");
+        assert(OutOp.isReg() && "ProxyReg output should be a register.");
         RemoveList.push_back(&MI);
+        RAUWBatch.try_emplace(OutOp.getReg(), InOp.getReg());
         break;
       }
+      }
     }
   }
 
+  // If there were no proxy instructions, exit early.
+  if (RemoveList.empty())
+    return false;
+
+  // Erase the proxy instructions first.
   for (auto *MI : RemoveList) {
     MI->eraseFromParent();
   }
 
-  return !RemoveList.empty();
-}
-
-void NVPTXProxyRegErasure::replaceMachineInstructionUsage(MachineFunction &MF,
-                                                          MachineInstr &MI) {
-  auto &InOp = *MI.uses().begin();
-  auto &OutOp = *MI.defs().begin();
-
-  assert(InOp.isReg() && "ProxyReg input operand should be a register.");
-  assert(OutOp.isReg() && "ProxyReg output operand should be a register.");
-
+  // Now go replace the registers.
   for (auto &BB : MF) {
-    for (auto &I : BB) {
-      replaceRegisterUsage(I, OutOp, InOp);
+    for (auto &MI : BB) {
+      for (auto &Op : MI.uses()) {
+        if (Op.isReg()) {
+          auto it = RAUWBatch.find(Op.getReg());
+          if (it != RAUWBatch.end())
+            Op.setReg(it->second);
+        }
+      }
     }
   }
-}
 
-void NVPTXProxyRegErasure::replaceRegisterUsage(MachineInstr &Instr,
-                                                MachineOperand &From,
-                                                MachineOperand &To) {
-  for (auto &Op : Instr.uses()) {
-    if (Op.isReg() && Op.getReg() == From.getReg()) {
-      Op.setReg(To.getReg());
-    }
-  }
+  return true;
 }
 
 MachineFunctionPass *llvm::createNVPTXProxyRegErasurePass() {

``````````

</details>


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