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

This reverts commit 08e5a1d.

Author: Groverkss

llvm/lib/Target/NVPTX/NVPTXProxyRegErasure.cpp

@@ -34,6 +34,7 @@ void initializeNVPTXProxyRegErasurePass(PassRegistry &);
 namespace {
 
 struct NVPTXProxyRegErasure : public MachineFunctionPass {
+public:
   static char ID;
   NVPTXProxyRegErasure() : MachineFunctionPass(ID) {
     initializeNVPTXProxyRegErasurePass(*PassRegistry::getPassRegistry());
@@ -48,22 +49,23 @@ 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()) {
@@ -72,42 +74,44 @@ bool NVPTXProxyRegErasure::runOnMachineFunction(MachineFunction &MF) {
       case NVPTX::ProxyRegI32:
       case NVPTX::ProxyRegI64:
       case NVPTX::ProxyRegF32:
-      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.");
+      case NVPTX::ProxyRegF64:
+        replaceMachineInstructionUsage(MF, MI);
         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();
   }
 
-  // Now go replace the registers.
+  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.");
+
   for (auto &BB : MF) {
-    for (auto &MI : BB) {
-      for (auto &Op : MI.uses()) {
-        if (!Op.isReg())
-          continue;
-        auto it = RAUWBatch.find(Op.getReg());
-        if (it != RAUWBatch.end())
-          Op.setReg(it->second);
-      }
+    for (auto &I : BB) {
+      replaceRegisterUsage(I, OutOp, InOp);
     }
   }
+}
 
-  return true;
+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());
+    }
+  }
 }
 
 MachineFunctionPass *llvm::createNVPTXProxyRegErasurePass() {