[llvm] [SLP] Initial support for masked stores (PR #204893)

[Alexey Bataev via llvm-commits]

================
@@ -17923,7 +18031,31 @@ BoUpSLP::getEntryCost(const TreeEntry *E, ArrayRef<Value *> VectorizedVals,
           VecStCost +=
               TTI->getCastInstrCost(Instruction::BitCast, VecTy, StridedStoreTy,
                                     getCastContextHint(*E), CostKind);
-
+      } else if (E->State == TreeEntry::ExpandVectorize) {
+        const StridedPtrInfo &SPtrInfo = TreeEntryToStridedPtrInfoMap.at(E);
+        FixedVectorType *MaskedStoreTy = SPtrInfo.Ty;
+        assert(MaskedStoreTy && "Missing StridedPointerInfo for tree entry.");
+        Align CommonAlignment =
+            computeCommonAlignment<StoreInst>(UniqueValues.getArrayRef());
+        // Masked store: the values are expanded into the widened vector and
+        // stored with a constant mask.
+        VecStCost = TTI->getMemIntrinsicInstrCost(
+            MemIntrinsicCostAttributes(Intrinsic::masked_store, MaskedStoreTy,
+                                       CommonAlignment,
+                                       BaseSI->getPointerAddressSpace()),
+            CostKind);
+        const TreeEntry *ValOpTE = getOperandEntry(E, 0);
+        if (auto ValBWIt = MinBWs.find(ValOpTE);
+            ValBWIt != MinBWs.end() && ValOpTE->hasState() &&
+            UnaryInstruction::isCast(ValOpTE->getOpcode())) {
+          auto *SrcVecTy = getWidenedType(
+              IntegerType::get(ScalarTy->getContext(), ValBWIt->second.first),
+              VL.size());
+          if (SrcVecTy != VecTy)
+            VecStCost += TTI->getCastInstrCost(
+                ValBWIt->second.second ? Instruction::SExt : Instruction::ZExt,
+                VecTy, SrcVecTy, getCastContextHint(*ValOpTE), CostKind);
+        }
----------------
alexey-bataev wrote:

Ah, no, not a bug, need to remove it here, already handled for the value, just forgot about it

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