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

[Vyacheslav Levytskyy via llvm-commits]

https://github.com/VyacheslavLevytskyy created https://github.com/llvm/llvm-project/pull/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.

>From da30d725b5324dedc64a034b4eda2e6ad10fe273 Mon Sep 17 00:00:00 2001
From: "Levytskyy, Vyacheslav" <vyacheslav.levytskyy at intel.com>
Date: Wed, 12 Jun 2024 13:22:53 -0700
Subject: [PATCH] ensure that cleaning of temporary constants doesn't purge
 tracked ones

---
 llvm/lib/Target/SPIRV/SPIRVPreLegalizer.cpp   | 16 ++++++++----
 llvm/test/CodeGen/SPIRV/keep-tracked-const.ll | 25 +++++++++++++++++++
 2 files changed, 36 insertions(+), 5 deletions(-)
 create mode 100644 llvm/test/CodeGen/SPIRV/keep-tracked-const.ll

diff --git a/llvm/lib/Target/SPIRV/SPIRVPreLegalizer.cpp b/llvm/lib/Target/SPIRV/SPIRVPreLegalizer.cpp
index 53e0432192ca9..d899dde3d1393 100644
--- a/llvm/lib/Target/SPIRV/SPIRVPreLegalizer.cpp
+++ b/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,8 @@ addConstantsToTrack(MachineFunction &MF, SPIRVGlobalRegistry *GR,
     MI->eraseFromParent();
 }
 
-static void foldConstantsIntoIntrinsics(MachineFunction &MF) {
+static void foldConstantsIntoIntrinsics(MachineFunction &MF,
+                                        SmallSet<Register, 4> &TrackedConstRegs) {
   SmallVector<MachineInstr *, 10> ToErase;
   MachineRegisterInfo &MRI = MF.getRegInfo();
   const unsigned AssignNameOperandShift = 2;
@@ -137,7 +140,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 +840,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);
diff --git a/llvm/test/CodeGen/SPIRV/keep-tracked-const.ll b/llvm/test/CodeGen/SPIRV/keep-tracked-const.ll
new file mode 100644
index 0000000000000..0c25d8a67aca0
--- /dev/null
+++ b/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)