[llvm] [FuncSpec] Improve estimation of select instruction. (PR #111176)

[Alexandros Lamprineas via llvm-commits]

https://github.com/labrinea created https://github.com/llvm/llvm-project/pull/111176

When propagating a constant to a select instruction we only consider the condition operand as the use. I am extending the logic to consider the true and false values too, in case the condition had been found to be constant in a previous propagation but halted.

>From 4744e5922e8ef1e960ec085cd5a85cea9170e724 Mon Sep 17 00:00:00 2001
From: Alexandros Lamprineas <alexandros.lamprineas at arm.com>
Date: Fri, 4 Oct 2024 16:17:51 +0100
Subject: [PATCH] [FuncSpec] Improve estimation of select instruction.

When propagating a constant to a select instruction we only consider
the condition operand as the use. I am extending the logic to consider
the true and false values too, in case the condition had been found
to be constant in a previous propagation but halted.
---
 .../Transforms/IPO/FunctionSpecialization.cpp   | 17 ++++++++++-------
 1 file changed, 10 insertions(+), 7 deletions(-)

diff --git a/llvm/lib/Transforms/IPO/FunctionSpecialization.cpp b/llvm/lib/Transforms/IPO/FunctionSpecialization.cpp
index 548335d750e33d..7d109af9091479 100644
--- a/llvm/lib/Transforms/IPO/FunctionSpecialization.cpp
+++ b/llvm/lib/Transforms/IPO/FunctionSpecialization.cpp
@@ -423,13 +423,16 @@ Constant *InstCostVisitor::visitGetElementPtrInst(GetElementPtrInst &I) {
 Constant *InstCostVisitor::visitSelectInst(SelectInst &I) {
   assert(LastVisited != KnownConstants.end() && "Invalid iterator!");
 
-  if (I.getCondition() != LastVisited->first)
-    return nullptr;
-
-  Value *V = LastVisited->second->isZeroValue() ? I.getFalseValue()
-                                                : I.getTrueValue();
-  Constant *C = findConstantFor(V, KnownConstants);
-  return C;
+  if (I.getCondition() == LastVisited->first) {
+    Value *V = LastVisited->second->isZeroValue() ? I.getFalseValue()
+                                                  : I.getTrueValue();
+    return findConstantFor(V, KnownConstants);
+  }
+  if (Constant *Condition = findConstantFor(I.getCondition(), KnownConstants))
+    if (I.getTrueValue() == LastVisited->first && Condition->isOneValue() ||
+        I.getFalseValue() == LastVisited->first && Condition->isZeroValue())
+      return LastVisited->second;
+  return nullptr;
 }
 
 Constant *InstCostVisitor::visitCastInst(CastInst &I) {