[llvm] [AArch64][CostModel] Consider the cost of const vector (PR #117539)

[Sushant Gokhale via llvm-commits]

================
@@ -11019,8 +11036,54 @@ BoUpSLP::getEntryCost(const TreeEntry *E, ArrayRef<Value *> VectorizedVals,
 
   bool NeedToShuffleReuses = !E->ReuseShuffleIndices.empty();
   if (E->isGather()) {
-    if (allConstant(VL))
+    if (allConstant(VL)) {
+      auto IsAllowedScalarTy = [&](const Type *T) {
+        return T->isFloatTy() || T->isDoubleTy() || T->isIntegerTy();
+      };
+      if (IsAllowedScalarTy(E->Scalars.front()->getType())) {
+        InstructionCost ScalarCost, VectorCost;
+
+        auto IsDuplicateEntry = [&](const TreeEntry *E) {
+          auto *TE = getTreeEntry(E->Scalars[0]);
+          if (TE != E) {
+            auto It = MultiNodeScalars.find(E->Scalars[0]);
+            if (It != MultiNodeScalars.end()) {
+              auto *TEIt = find_if(It->getSecond(), [&](TreeEntry *ME) {
+                return ME->isSame(VL);
+              });
+              if (TEIt != It->getSecond().end())
+                return true;
+            }
+          }
+          return false;
+        };
+
+        // FIXME: If there is more than 1 SLP tree realizing the same const
+        // vector, codegen will realize it only once. Hence, no need to consider
+        // the cost of const vector twice. But, currently we can't check if the
+        // tree entry is present in other SLP tree.
+        if (!isSplat(E->Scalars) && !all_of(E->Scalars, IsaPred<UndefValue>) &&
+            !IsDuplicateEntry(E)) {
+          // Get unique scalars
+          SmallDenseSet<Value *> UniqScalars;
+          for (auto *V : E->Scalars)
+            UniqScalars.insert(V);
+
+          // Constant is realized by having a mov/fmov into GPR. So,
+          // ScalarCost = #UniqScalars
+          ScalarCost = (UniqScalars.size());
+
+          // FIXME: Ideally, getGatherCost API should be used for this but
+          // currently, this API returns zero cost with all constants.
+          VectorCost =
+              TTI->getConstVectCost(Instruction::Load, FinalVecTy, Align(), 0,
+                                    CostKind, TTI::OperandValueInfo(), nullptr,
+                                    /*AllConstants=*/true, ScalarCost);
+        }
+        return VectorCost - ScalarCost;
+      }
----------------
sushgokh wrote:

shifted the logic in TTI

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