Const correctness changes and AArch64 option

Author: davemgreen

llvm/include/llvm/Analysis/TargetTransformInfo.h

@@ -938,7 +938,7 @@ class TargetTransformInfo {
   /// select, potentially converting it to a conditional branch. This can
   /// include select-like instructions like or(zext(c), x) that can be converted
   /// to selects.
-  bool shouldTreatInstructionLikeSelect(Instruction *I) const;
+  bool shouldTreatInstructionLikeSelect(const Instruction *I) const;
 
   /// Enable matching of interleaved access groups.
   bool enableInterleavedAccessVectorization() const;
@@ -1884,7 +1884,7 @@ class TargetTransformInfo::Concept {
   virtual MemCmpExpansionOptions
   enableMemCmpExpansion(bool OptSize, bool IsZeroCmp) const = 0;
   virtual bool enableSelectOptimize() = 0;
-  virtual bool shouldTreatInstructionLikeSelect(Instruction *I) = 0;
+  virtual bool shouldTreatInstructionLikeSelect(const Instruction *I) = 0;
   virtual bool enableInterleavedAccessVectorization() = 0;
   virtual bool enableMaskedInterleavedAccessVectorization() = 0;
   virtual bool isFPVectorizationPotentiallyUnsafe() = 0;
@@ -2422,7 +2422,7 @@ class TargetTransformInfo::Model final : public TargetTransformInfo::Concept {
   bool enableSelectOptimize() override {
     return Impl.enableSelectOptimize();
   }
-  bool shouldTreatInstructionLikeSelect(Instruction *I) override {
+  bool shouldTreatInstructionLikeSelect(const Instruction *I) override {
     return Impl.shouldTreatInstructionLikeSelect(I);
   }
   bool enableInterleavedAccessVectorization() override {

llvm/include/llvm/Analysis/TargetTransformInfoImpl.h

@@ -378,7 +378,7 @@ class TargetTransformInfoImplBase {
 
   bool enableSelectOptimize() const { return true; }
 
-  bool shouldTreatInstructionLikeSelect(Instruction *I) {
+  bool shouldTreatInstructionLikeSelect(const Instruction *I) {
     // If the select is a logical-and/logical-or then it is better treated as a
     // and/or by the backend.
     using namespace llvm::PatternMatch;

llvm/lib/Analysis/TargetTransformInfo.cpp

@@ -605,7 +605,7 @@ bool TargetTransformInfo::enableSelectOptimize() const {
 }
 
 bool TargetTransformInfo::shouldTreatInstructionLikeSelect(
-    Instruction *I) const {
+    const Instruction *I) const {
   return TTIImpl->shouldTreatInstructionLikeSelect(I);
 }
 

llvm/lib/CodeGen/SelectOptimize.cpp

@@ -318,7 +318,8 @@ class SelectOptimizeImpl {
                         CostInfo *LoopCost);
 
   // Returns a set of all the select instructions in the given select groups.
-  SmallPtrSet<const Instruction *, 2> getSIset(const SelectGroups &SIGroups);
+  SmallDenseMap<const Instruction *, SelectLike, 2>
+  getSImap(const SelectGroups &SIGroups);
 
   // Returns the latency cost of a given instruction.
   std::optional<uint64_t> computeInstCost(const Instruction *I);
@@ -1112,15 +1113,15 @@ bool SelectOptimizeImpl::computeLoopCosts(
     DenseMap<const Instruction *, CostInfo> &InstCostMap, CostInfo *LoopCost) {
   LLVM_DEBUG(dbgs() << "Calculating Latency / IPredCost / INonPredCost of loop "
                     << L->getHeader()->getName() << "\n");
-  const auto &SIset = getSIset(SIGroups);
+  const auto &SImap = getSImap(SIGroups);
   // Compute instruction and loop-critical-path costs across two iterations for
   // both predicated and non-predicated version.
   const unsigned Iterations = 2;
   for (unsigned Iter = 0; Iter < Iterations; ++Iter) {
     // Cost of the loop's critical path.
     CostInfo &MaxCost = LoopCost[Iter];
     for (BasicBlock *BB : L->getBlocks()) {
-      for (Instruction &I : *BB) {
+      for (const Instruction &I : *BB) {
         if (I.isDebugOrPseudoInst())
           continue;
         // Compute the predicated and non-predicated cost of the instruction.
@@ -1157,10 +1158,8 @@ bool SelectOptimizeImpl::computeLoopCosts(
         // BranchCost = PredictedPathCost + MispredictCost
         // PredictedPathCost = TrueOpCost * TrueProb + FalseOpCost * FalseProb
         // MispredictCost = max(MispredictPenalty, CondCost) * MispredictRate
-        if (SIset.contains(&I)) {
-          auto SI = SelectLike::match(&I);
-          assert(SI && "Expected to match an existing SelectLike");
-
+        if (SImap.contains(&I)) {
+          auto SI = SImap.at(&I);
           Scaled64 TrueOpCost = SI.getTrueOpCost(InstCostMap, TTI);
           Scaled64 FalseOpCost = SI.getFalseOpCost(InstCostMap, TTI);
           Scaled64 PredictedPathCost =
@@ -1189,13 +1188,13 @@ bool SelectOptimizeImpl::computeLoopCosts(
   return true;
 }
 
-SmallPtrSet<const Instruction *, 2>
-SelectOptimizeImpl::getSIset(const SelectGroups &SIGroups) {
-  SmallPtrSet<const Instruction *, 2> SIset;
+SmallDenseMap<const Instruction *, SelectOptimizeImpl::SelectLike, 2>
+SelectOptimizeImpl::getSImap(const SelectGroups &SIGroups) {
+  SmallDenseMap<const Instruction *, SelectLike, 2> SImap;
   for (const SelectGroup &ASI : SIGroups)
     for (SelectLike SI : ASI)
-      SIset.insert(SI.getI());
-  return SIset;
+      SImap.try_emplace(SI.getI(), SI);
+  return SImap;
 }
 
 std::optional<uint64_t>

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp

@@ -55,6 +55,9 @@ static cl::opt<unsigned> InlineCallPenaltyChangeSM(
     "inline-call-penalty-sm-change", cl::init(10), cl::Hidden,
     cl::desc("Penalty of inlining a call that requires a change to PSTATE.SM"));
 
+static cl::opt<bool> EnableOrLikeSelectOpt("enable-aarch64-or-like-select",
+                                           cl::init(true), cl::Hidden);
+
 namespace {
 class TailFoldingOption {
   // These bitfields will only ever be set to something non-zero in operator=,
@@ -4048,3 +4051,15 @@ AArch64TTIImpl::getScalingFactorCost(Type *Ty, GlobalValue *BaseGV,
     return AM.Scale != 0 && AM.Scale != 1;
   return -1;
 }
+
+bool AArch64TTIImpl::shouldTreatInstructionLikeSelect(const Instruction *I) {
+  // For the binary operators (e.g. or) we need to be more careful than
+  // selects, here we only transform them if they are already at a natural
+  // break point in the code - the end of a block with an unconditional
+  // terminator.
+  if (EnableOrLikeSelectOpt && I->getOpcode() == Instruction::Or &&
+      isa<BranchInst>(I->getNextNode()) &&
+      cast<BranchInst>(I->getNextNode())->isUnconditional())
+    return true;
+  return BaseT::shouldTreatInstructionLikeSelect(I);
+}

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.h

@@ -412,17 +412,7 @@ class AArch64TTIImpl : public BasicTTIImplBase<AArch64TTIImpl> {
 
   bool enableSelectOptimize() { return ST->enableSelectOptimize(); }
 
-  bool shouldTreatInstructionLikeSelect(Instruction *I) {
-    // For the binary operators (e.g. or) we need to be more careful than
-    // selects, here we only transform them if they are already at a natural
-    // break point in the code - the end of a block with an unconditional
-    // terminator.
-    if (I->getOpcode() == Instruction::Or &&
-        isa<BranchInst>(I->getNextNode()) &&
-        cast<BranchInst>(I->getNextNode())->isUnconditional())
-      return true;
-    return BaseT::shouldTreatInstructionLikeSelect(I);
-  }
+  bool shouldTreatInstructionLikeSelect(const Instruction *I);
 
   unsigned getStoreMinimumVF(unsigned VF, Type *ScalarMemTy,
                              Type *ScalarValTy) const {