[llvm][NVPTX] Fix RAUW bug in NVPTXProxyRegErasure

[Mogball]

Fix bug introduced in #105730

The bug is in how the batch RAUW is implemented. If we have

%0 = mov %src
%1 = mov %0

use %0
use %1

The use of %1 is rewritten to %0, not %src. This PR just looks for a replacement when it maps to the src register, which should transitively propagate the replacements.

[llvmbot]

@llvm/pr-subscribers-backend-nvptx

Author: Jeff Niu (Mogball)

Changes

Fix bug introduced in #105730

---

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

2 Files Affected:

• (modified) llvm/lib/Target/NVPTX/NVPTXProxyRegErasure.cpp (+5-1)
• (added) llvm/test/CodeGen/NVPTX/bug105730.ll (+17)

diff --git a/llvm/lib/Target/NVPTX/NVPTXProxyRegErasure.cpp b/llvm/lib/Target/NVPTX/NVPTXProxyRegErasure.cpp
index f3a3362addb0ea..16c2b307efabfb 100644
--- a/llvm/lib/Target/NVPTX/NVPTXProxyRegErasure.cpp
+++ b/llvm/lib/Target/NVPTX/NVPTXProxyRegErasure.cpp
@@ -78,7 +78,11 @@ bool NVPTXProxyRegErasure::runOnMachineFunction(MachineFunction &MF) {
         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());
+        Register replacement = InOp.getReg();
+        // Check if the replacement itself has been replaced.
+        if (auto it = RAUWBatch.find(replacement); it != RAUWBatch.end())
+          replacement = it->second;
+        RAUWBatch.try_emplace(OutOp.getReg(), replacement);
         break;
       }
       }
diff --git a/llvm/test/CodeGen/NVPTX/bug105730.ll b/llvm/test/CodeGen/NVPTX/bug105730.ll
new file mode 100644
index 00000000000000..718e7ca6b80fd8
--- /dev/null
+++ b/llvm/test/CodeGen/NVPTX/bug105730.ll
@@ -0,0 +1,17 @@
+; RUN: llc < %s -verify-machineinstrs
+
+; Check that llc doesn't crash.
+
+target triple = "nvptx64-nvidia-cuda"
+
+define void @__builtin_splat_i8(i32 %0) {
+.lr.ph:
+  %1 = trunc i32 %0 to i8
+  %broadcast.splatinsert = insertelement <4 x i8> poison, i8 %1, i64 0
+  %broadcast.splat = shufflevector <4 x i8> %broadcast.splatinsert, <4 x i8> poison, <4 x i32> zeroinitializer
+  br label %vector.body
+
+vector.body:
+  store <4 x i8> %broadcast.splat, ptr addrspace(1) poison, align 1
+  br label %vector.body
+}

[Mogball]

I'm getting more reports of breakages coming from downstream projects, so I'm going to go ahead and land this to fix-forward. If you have any more comments on the PR, I am happy to address them in follow ups! Thanks for the review and testing tips