[llvm] [AMDGPU] Try to reuse in v_cndmask register with constant from compare. (PR #131146)

[Daniil Fukalov via llvm-commits]

================
@@ -1411,15 +1411,85 @@ bool SIFoldOperandsImpl::tryFoldCndMask(MachineInstr &MI) const {
       Opc != AMDGPU::V_CNDMASK_B64_PSEUDO)
     return false;
 
+  // Try to find optimized Y == Const ? Const : Z. If Const can't be directly
+  // encoded in the cndmask, try to reuse a register already holding the Const
+  // value from the comparison instruction.
+  auto tryFoldCndMaskCmp =
+      [&](MachineOperand *SrcOp, std::optional<int64_t> SrcImm,
+          ArrayRef<unsigned> CmpOpcodes, AMDGPU::OpName CmpValName) -> bool {
+    // We'll try to process only register operands with known values.
+    if (!SrcImm || !SrcOp->isReg())
+      return false;
+
+    // Find the predicate of the cndmask instruction.
+    MachineOperand *PredOp = TII->getNamedOperand(MI, AMDGPU::OpName::src2);
+    if (!PredOp || !PredOp->isReg())
+      return false;
+
+    MachineInstr *PredI = MRI->getVRegDef(PredOp->getReg());
+    if (!PredI || !PredI->isCompare())
+      return false;
+
+    unsigned CmpOpc = PredI->getOpcode();
+
+    // Check if the comparison instruction is one of the expected ones.
+    const auto *CmpOpcI = find_if(
+        CmpOpcodes, [CmpOpc](unsigned Opcode) { return Opcode == CmpOpc; });
+    if (CmpOpcI == CmpOpcodes.end())
----------------
dfukalov wrote:

Thanks!

https://github.com/llvm/llvm-project/pull/131146