[llvm] [SPIR-V] Erase all instructions after OpKill, not just the next one (PR #219022)

[Arseniy Obolenskiy via llvm-commits]

https://github.com/aobolensk updated https://github.com/llvm/llvm-project/pull/219022

>From 39de94770b6fb39e66f0a3b43be569e86354090d Mon Sep 17 00:00:00 2001
From: Arseniy Obolenskiy <arseniy.obolenskiy at amd.com>
Date: Wed, 26 Aug 2026 21:51:18 +0200
Subject: [PATCH 1/2] [SPIR-V] Erase all instructions after OpKill, not just
 the next one

---
 .../Target/SPIRV/SPIRVInstructionSelector.cpp |  8 +++---
 .../CodeGen/SPIRV/hlsl-intrinsics/discard.ll  | 26 +++++++++++++++++++
 2 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/llvm/lib/Target/SPIRV/SPIRVInstructionSelector.cpp b/llvm/lib/Target/SPIRV/SPIRVInstructionSelector.cpp
index f3f9984190de5..1607583611e82 100644
--- a/llvm/lib/Target/SPIRV/SPIRVInstructionSelector.cpp
+++ b/llvm/lib/Target/SPIRV/SPIRVInstructionSelector.cpp
@@ -4483,9 +4483,11 @@ bool SPIRVInstructionSelector::selectDiscard(Register ResVReg,
   } else {
     Opcode = SPIRV::OpKill;
     // OpKill must be the last operation of any basic block.
-    if (MachineInstr *NextI = I.getNextNode()) {
-      GR.invalidateMachineInstr(NextI);
-      NextI->eraseFromParent();
+    for (MachineInstr *NextI = I.getNextNode(); NextI;) {
+      MachineInstr *ToErase = NextI;
+      NextI = NextI->getNextNode();
+      GR.invalidateMachineInstr(ToErase);
+      ToErase->eraseFromParent();
     }
   }
 
diff --git a/llvm/test/CodeGen/SPIRV/hlsl-intrinsics/discard.ll b/llvm/test/CodeGen/SPIRV/hlsl-intrinsics/discard.ll
index 39be59bbc7332..fb13464010115 100644
--- a/llvm/test/CodeGen/SPIRV/hlsl-intrinsics/discard.ll
+++ b/llvm/test/CodeGen/SPIRV/hlsl-intrinsics/discard.ll
@@ -47,6 +47,32 @@ end:                                              ; preds = %lt0, %entry
 }
 declare void @llvm.spv.discard()
 
+define void @test_scalar_multi_instr_after_discard(float noundef %Buf, ptr %Out) {
+entry:
+; CHECK-LABEL: ; -- Begin function test_scalar_multi_instr_after_discard
+; CHECK:       OpBranchConditional %[[#]] %[[#truel:]] %[[#endl:]]
+; CHECK:       %[[#truel]] = OpLabel
+; SPIRV15:     OpKill
+; SPIRV15-NOT: OpStore
+; SPIRV16:     OpDemoteToHelperInvocation
+; SPIRV16:     OpBranch %[[#endl]]
+; CHECK:       %[[#endl]] = OpLabel
+  %Buf.addr = alloca float, align 4
+  store float %Buf, ptr %Buf.addr, align 4
+  %1 = load float, ptr %Buf.addr, align 4
+  %2 = fcmp olt float %1, 0.000000e+00
+  br i1 %2, label %lt0, label %end
+
+lt0:                                              ; preds = %entry
+  call void @llvm.spv.discard()
+  store volatile float 1.0, ptr %Out, align 4
+  store volatile float 2.0, ptr %Out, align 4
+  br label %end
+
+end:                                              ; preds = %lt0, %entry
+  ret void
+}
+
 define void @test_vector(<4 x float> noundef %Buf) {
 entry:
 ; CHECK-LABEL: ; -- Begin function test_vector

>From f5c872aa83b1a0722048cb44d573c71735fb5eca Mon Sep 17 00:00:00 2001
From: Arseniy Obolenskiy <arseniy.obolenskiy at amd.com>
Date: Thu, 27 Aug 2026 17:06:28 +0200
Subject: [PATCH 2/2] address comments

---
 llvm/lib/Target/SPIRV/SPIRVInstructionSelector.cpp | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/llvm/lib/Target/SPIRV/SPIRVInstructionSelector.cpp b/llvm/lib/Target/SPIRV/SPIRVInstructionSelector.cpp
index 1607583611e82..d87e16117cf13 100644
--- a/llvm/lib/Target/SPIRV/SPIRVInstructionSelector.cpp
+++ b/llvm/lib/Target/SPIRV/SPIRVInstructionSelector.cpp
@@ -22,6 +22,7 @@
 #include "SPIRVTypeInst.h"
 #include "SPIRVUtils.h"
 #include "llvm/ADT/APFloat.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/STLFunctionalExtras.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/StringExtras.h"
@@ -4475,6 +4476,7 @@ bool SPIRVInstructionSelector::selectDiscard(Register ResVReg,
                                              MachineInstr &I) const {
 
   unsigned Opcode;
+  MachineBasicBlock &BB = *I.getParent();
 
   if (STI.canUseExtension(
           SPIRV::Extension::SPV_EXT_demote_to_helper_invocation) ||
@@ -4483,15 +4485,13 @@ bool SPIRVInstructionSelector::selectDiscard(Register ResVReg,
   } else {
     Opcode = SPIRV::OpKill;
     // OpKill must be the last operation of any basic block.
-    for (MachineInstr *NextI = I.getNextNode(); NextI;) {
-      MachineInstr *ToErase = NextI;
-      NextI = NextI->getNextNode();
-      GR.invalidateMachineInstr(ToErase);
-      ToErase->eraseFromParent();
+    for (MachineInstr &ToErase : make_early_inc_range(
+             make_range(std::next(I.getIterator()), BB.instr_end()))) {
+      GR.invalidateMachineInstr(&ToErase);
+      ToErase.eraseFromParent();
     }
   }
 
-  MachineBasicBlock &BB = *I.getParent();
   BuildMI(BB, I, I.getDebugLoc(), TII.get(Opcode))
       .constrainAllUses(TII, TRI, RBI);
   return true;