[llvm] [NVPTX] Fix incorrect kill flags in ProxyRegErasure (PR #198143)

[via llvm-commits]

https://github.com/andreww0210 updated https://github.com/llvm/llvm-project/pull/198143

>From a8a4ee2133041559c7778a1ea83fd5feb1c8ca6f Mon Sep 17 00:00:00 2001
From: yinghao <itsyinghao at gmail.com>
Date: Sun, 17 May 2026 11:12:43 +0800
Subject: [PATCH 1/2] [NVPTX] Clear unsafe kill flags in NVPTXProxyRegErasure

After replacing regsiter, stale `killed` flags my leave on non-terminal
uses, breaking liveness and triggering MachineVerifier errors.

Fix this by conservatively dropping the `killed` flag on `ToReg`.

Fixes #198142.
---
 .../lib/Target/NVPTX/NVPTXProxyRegErasure.cpp | 17 +++++--------
 .../NVPTX/proxy-reg-erasure-kill-flag.mir     | 25 +++++++++++++++++++
 llvm/test/CodeGen/NVPTX/proxy-reg-erasure.mir |  4 +--
 3 files changed, 33 insertions(+), 13 deletions(-)
 create mode 100644 llvm/test/CodeGen/NVPTX/proxy-reg-erasure-kill-flag.mir

diff --git a/llvm/lib/Target/NVPTX/NVPTXProxyRegErasure.cpp b/llvm/lib/Target/NVPTX/NVPTXProxyRegErasure.cpp
index bac036f3a93ae..7daea6a840e37 100644
--- a/llvm/lib/Target/NVPTX/NVPTXProxyRegErasure.cpp
+++ b/llvm/lib/Target/NVPTX/NVPTXProxyRegErasure.cpp
@@ -20,6 +20,7 @@
 
 #include "NVPTX.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
+#include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/TargetRegisterInfo.h"
 
 using namespace llvm;
@@ -87,17 +88,11 @@ bool NVPTXProxyRegErasure::runOnMachineFunction(MachineFunction &MF) {
     MI->eraseFromParent();
   }
 
-  // Now go replace the registers.
-  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);
-      }
-    }
+  // Now go replace the registers and remove kill flags conservatively.
+  MachineRegisterInfo &MRI = MF.getRegInfo();
+  for (auto [From, To] : RAUWBatch) {
+    MRI.replaceRegWith(From, To);
+    MRI.clearKillFlags(To);
   }
 
   return true;
diff --git a/llvm/test/CodeGen/NVPTX/proxy-reg-erasure-kill-flag.mir b/llvm/test/CodeGen/NVPTX/proxy-reg-erasure-kill-flag.mir
new file mode 100644
index 0000000000000..e56b3192ac704
--- /dev/null
+++ b/llvm/test/CodeGen/NVPTX/proxy-reg-erasure-kill-flag.mir
@@ -0,0 +1,25 @@
+# RUN: llc %s --run-pass=nvptx-proxyreg-erasure -verify-machineinstrs -mtriple=nvptx64 -o - | FileCheck %s
+
+--- |
+  define void @test_proxy_reg() {
+  entry:
+    ret void
+  }
+
+...
+---
+name: test_proxy_reg
+tracksRegLiveness: true
+body: |
+  bb.0:
+    %0:b32 = IMPLICIT_DEF
+    %1:b16 = IMPLICIT_DEF
+
+    ; CHECK: %2:b16 = CVT_u16_u32 %0, 0
+    ; CHECK-NEXT: %4:b16 = ADD16rr %1, %2
+    ; CHECK-NEXT: %5:b16 = ADD16rr %1, %2
+    %2:b16 = CVT_u16_u32 %0:b32, 0
+    %3:b16 = ProxyRegB16 %2:b16
+    %4:b16 = ADD16rr %1:b16, killed %3:b16
+    %5:b16 = ADD16rr %1:b16, %2:b16
+...
diff --git a/llvm/test/CodeGen/NVPTX/proxy-reg-erasure.mir b/llvm/test/CodeGen/NVPTX/proxy-reg-erasure.mir
index a84b7fcd33836..59ca24acc504c 100644
--- a/llvm/test/CodeGen/NVPTX/proxy-reg-erasure.mir
+++ b/llvm/test/CodeGen/NVPTX/proxy-reg-erasure.mir
@@ -83,7 +83,7 @@ body:             |
     %5:b32 = ProxyRegB32 killed %1
     %6:b32 = ProxyRegB32 killed %2
     %7:b32 = ProxyRegB32 killed %3
-    ; CHECK: STV_i32_v4 killed %0, killed %1, killed %2, killed %3
+    ; CHECK: STV_i32_v4 %0, %1, %2, %3
     STV_i32_v4 killed %4, killed %5, killed %6, killed %7, 0, 0, 101, 32, &func_retval0, 0 :: (store (s128), addrspace 101)
 
     %8:b32 = LD_i32 0, 0, 101, 3, 32, -1, &retval0, 0 :: (load (s32), addrspace 101)
@@ -91,7 +91,7 @@ body:             |
     %9:b32 = ProxyRegB32 killed %8
     %10:b32 = ProxyRegB32 killed %9
     %11:b32 = ProxyRegB32 killed %10
-    ; CHECK: ST_i32 killed %8
+    ; CHECK: ST_i32 %8
     ST_i32 killed %11, 0, 0, 101, 32, &func_retval0, 0 :: (store (s32), addrspace 101)
     Return
 

>From 24f57488e28c14a173ded89c6c87e69abeb28b34 Mon Sep 17 00:00:00 2001
From: yinghao <73093416+andreww0210 at users.noreply.github.com>
Date: Wed, 10 Jun 2026 21:55:24 +0800
Subject: [PATCH 2/2] Update
 llvm/test/CodeGen/NVPTX/proxy-reg-erasure-kill-flag.mir

Co-authored-by: Matt Arsenault <arsenm2 at gmail.com>
---
 llvm/test/CodeGen/NVPTX/proxy-reg-erasure-kill-flag.mir | 9 +--------
 1 file changed, 1 insertion(+), 8 deletions(-)

diff --git a/llvm/test/CodeGen/NVPTX/proxy-reg-erasure-kill-flag.mir b/llvm/test/CodeGen/NVPTX/proxy-reg-erasure-kill-flag.mir
index e56b3192ac704..73eabe0f50890 100644
--- a/llvm/test/CodeGen/NVPTX/proxy-reg-erasure-kill-flag.mir
+++ b/llvm/test/CodeGen/NVPTX/proxy-reg-erasure-kill-flag.mir
@@ -1,12 +1,5 @@
-# RUN: llc %s --run-pass=nvptx-proxyreg-erasure -verify-machineinstrs -mtriple=nvptx64 -o - | FileCheck %s
+# RUN: llc %s --run-pass=nvptx-proxyreg-erasure -mtriple=nvptx64 -o - | FileCheck %s
 
---- |
-  define void @test_proxy_reg() {
-  entry:
-    ret void
-  }
-
-...
 ---
 name: test_proxy_reg
 tracksRegLiveness: true