diff --git a/llvm/lib/Target/NVPTX/NVPTXProxyRegErasure.cpp b/llvm/lib/Target/NVPTX/NVPTXProxyRegErasure.cpp index 258ae97a20d582..f3a3362addb0ea 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 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 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()) + continue; + 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() {