[SPIR-V] Ensure that cleaning of temporary constants doesn't purge tracked constants (#95303)

This PR fixes a problem in logics of cleaning unused constants, ensuring
that cleaning of temporary constants doesn't purge tracked constants. On
a rare occasion when this happens SPIR-V Backend emits a code that
refers to a non-existent register, earlier related with a constant.
Attached to the PR test case is a minimal reproducer where names of
variables and instructions lead to such a rare coincidence.

Author: VyacheslavLevytskyy

llvm/lib/Target/SPIRV/SPIRVPreLegalizer.cpp

@@ -41,7 +41,8 @@ class SPIRVPreLegalizer : public MachineFunctionPass {
 static void
 addConstantsToTrack(MachineFunction &MF, SPIRVGlobalRegistry *GR,
                     const SPIRVSubtarget &STI,
-                    DenseMap<MachineInstr *, Type *> &TargetExtConstTypes) {
+                    DenseMap<MachineInstr *, Type *> &TargetExtConstTypes,
+                    SmallSet<Register, 4> &TrackedConstRegs) {
   MachineRegisterInfo &MRI = MF.getRegInfo();
   DenseMap<MachineInstr *, Register> RegsAlreadyAddedToDT;
   SmallVector<MachineInstr *, 10> ToErase, ToEraseComposites;
@@ -80,6 +81,7 @@ addConstantsToTrack(MachineFunction &MF, SPIRVGlobalRegistry *GR,
             }
           }
           GR->add(Const, &MF, SrcReg);
+          TrackedConstRegs.insert(SrcReg);
           if (Const->getType()->isTargetExtTy()) {
             // remember association so that we can restore it when assign types
             MachineInstr *SrcMI = MRI.getVRegDef(SrcReg);
@@ -121,7 +123,9 @@ addConstantsToTrack(MachineFunction &MF, SPIRVGlobalRegistry *GR,
     MI->eraseFromParent();
 }
 
-static void foldConstantsIntoIntrinsics(MachineFunction &MF) {
+static void
+foldConstantsIntoIntrinsics(MachineFunction &MF,
+                            const SmallSet<Register, 4> &TrackedConstRegs) {
   SmallVector<MachineInstr *, 10> ToErase;
   MachineRegisterInfo &MRI = MF.getRegInfo();
   const unsigned AssignNameOperandShift = 2;
@@ -137,7 +141,8 @@ static void foldConstantsIntoIntrinsics(MachineFunction &MF) {
         MI.removeOperand(NumOp);
         MI.addOperand(MachineOperand::CreateImm(
             ConstMI->getOperand(1).getCImm()->getZExtValue()));
-        if (MRI.use_empty(ConstMI->getOperand(0).getReg()))
+        Register DefReg = ConstMI->getOperand(0).getReg();
+        if (MRI.use_empty(DefReg) && !TrackedConstRegs.contains(DefReg))
           ToErase.push_back(ConstMI);
       }
     }
@@ -836,8 +841,10 @@ bool SPIRVPreLegalizer::runOnMachineFunction(MachineFunction &MF) {
   MachineIRBuilder MIB(MF);
   // a registry of target extension constants
   DenseMap<MachineInstr *, Type *> TargetExtConstTypes;
-  addConstantsToTrack(MF, GR, ST, TargetExtConstTypes);
-  foldConstantsIntoIntrinsics(MF);
+  // to keep record of tracked constants
+  SmallSet<Register, 4> TrackedConstRegs;
+  addConstantsToTrack(MF, GR, ST, TargetExtConstTypes, TrackedConstRegs);
+  foldConstantsIntoIntrinsics(MF, TrackedConstRegs);
   insertBitcasts(MF, GR, MIB);
   generateAssignInstrs(MF, GR, MIB, TargetExtConstTypes);
   processSwitches(MF, GR, MIB);

llvm/test/CodeGen/SPIRV/keep-tracked-const.ll

@@ -0,0 +1,25 @@
+; This test case ensures that cleaning of temporary constants doesn't purge tracked ones.
+
+; RUN: llc -O0 -mtriple=spirv64-unknown-unknown %s -o - | FileCheck %s --check-prefix=CHECK-SPIRV
+; RUN: %if spirv-tools %{ llc -O0 -mtriple=spirv64-unknown-unknown %s -o - -filetype=obj | spirv-val %}
+
+; CHECK-SPIRV: %[[#Int:]] = OpTypeInt 32 0
+; CHECK-SPIRV: %[[#C0:]] = OpConstant %[[#Int]] 0
+; CHECK-SPIRV: %[[#C1:]] = OpConstant %[[#Int]] 1
+; CHECK-SPIRV: OpSelect %[[#Int]] %[[#]] %[[#C1]] %[[#C0]]
+
+
+define spir_kernel void @foo() {
+entry:
+  %addr = alloca i32
+  %r1 = call i8 @_Z20__spirv_SpecConstantia(i32 0, i8 1)
+  ; The name '%conv17.i' is important for the test case,
+  ; because it includes i32 0 when encoded for SPIR-V usage.
+  %conv17.i = sext i8 %r1 to i64
+  %tobool = trunc i8 %r1 to i1
+  %r2 = zext i1 %tobool to i32
+  store i32 %r2, ptr %addr
+  ret void
+}
+
+declare i8 @_Z20__spirv_SpecConstantia(i32, i8)