[llvm] [SandboxVec][LoadStoreVec] Vectorize load seed chains (PR #215944)

[via llvm-commits]

================
@@ -128,91 +156,111 @@ bool LoadStoreVec::vectorizeStores(ArrayRef<Instruction *> Bndl, Region &Rgn,
   // Vectorizing mixed floats and integers with external uses may not be
   // profitable on some targets, so save state here.
   saveIR(Rgn);
-
   Value *VecOp = nullptr;
   if (AllLoads) {
     // TODO: Try to avoid the extra copy to an instruction vector.
     SmallVector<Instruction *, 8> Loads;
     Loads.reserve(Operands.size());
     for (Value *Op : Operands)
       Loads.push_back(cast<Instruction>(Op));
-
-    bool Consecutive = VecUtils::areConsecutive<LoadInst, Instruction>(
-        Loads, A.getScalarEvolution(), *DL);
-    if (!Consecutive) {
-      Ctx.accept();
+    VecOp = createVectorLoad(Loads);
+    if (VecOp == nullptr) {
+      Ctx->accept();
       return false;
     }
-    if (!canVectorize(Loads, Sched)) {
-      Ctx.accept();
-      return false;
-    }
-
-    // Generate vector load.
-    Type *Ty = VecUtils::getCombinedVectorTypeFor(Bndl, *DL);
-    Value *LdPtr = cast<LoadInst>(Loads[0])->getPointerOperand();
-    // TODO: Compute alignment.
-    Align LdAlign(1);
-    auto LdWhereIt = std::next(VecUtils::getLowest(Loads)->getIterator());
-    VecOp = LoadInst::create(Ty, LdPtr, LdAlign, LdWhereIt, Ctx, "VecIinitL");
   } else if (AllConstants) {
-    SmallVector<Constant *, 8> Constants;
-    Constants.reserve(Operands.size());
-    for (Value *Op : Operands) {
-      auto *COp = cast<Constant>(Op);
-      if (auto *AggrCOp = dyn_cast<ConstantAggregate>(COp)) {
-        // If the operand is a constant aggregate, then append all its elements.
-        for (Value *Elm : AggrCOp->operands())
-          Constants.push_back(cast<Constant>(Elm));
-      } else if (auto *SeqCOp = dyn_cast<ConstantDataSequential>(COp)) {
-        for (auto ElmIdx : seq<unsigned>(SeqCOp->getNumElements()))
-          Constants.push_back(SeqCOp->getElementAsConstant(ElmIdx));
-      } else if (auto *Zero = dyn_cast<ConstantAggregateZero>(COp)) {
-        auto *ZeroElm = Zero->getSequentialElement();
-        for ([[maybe_unused]] auto Cnt :
-             seq<unsigned>(Zero->getElementCount().getFixedValue()))
-          Constants.push_back(ZeroElm);
-      } else if (isa<ConstantInt>(COp) && isa<VectorType>(COp->getType())) {
-        auto *Elm = ConstantInt::get(Ctx, cast<ConstantInt>(COp)->getValue());
-        for ([[maybe_unused]] auto Cnt :
-             seq<unsigned>(cast<VectorType>(COp->getType())
-                               ->getElementCount()
-                               .getFixedValue()))
-          Constants.push_back(Elm);
-      } else if (isa<ConstantFP>(COp) && isa<VectorType>(COp->getType())) {
-        auto *Elm = ConstantFP::get(cast<ConstantFP>(COp)->getValue(), Ctx);
-        for ([[maybe_unused]] auto Cnt :
-             seq<unsigned>(cast<VectorType>(COp->getType())
-                               ->getElementCount()
-                               .getFixedValue()))
-          Constants.push_back(Elm);
-      } else {
-        Constants.push_back(COp);
-      }
-    }
-    VecOp = ConstantVector::get(Constants);
+    VecOp = createConstantVector(Operands);
   }
 
   // Generate vector store.
-  Value *StPtr = cast<StoreInst>(Bndl[0])->getPointerOperand();
+  Value *StPtr = cast<StoreInst>(Stores[0])->getPointerOperand();
   // TODO: Compute alignment.
   Align StAlign(1);
-  auto StWhereIt = std::next(VecUtils::getLowest(Bndl)->getIterator());
-  StoreInst::create(VecOp, StPtr, StAlign, StWhereIt, Ctx);
+  auto StWhereIt = std::next(VecUtils::getLowest(Stores)->getIterator());
+  StoreInst::create(VecOp, StPtr, StAlign, StWhereIt, *Ctx);
 
-  tryEraseDeadInstrs(Bndl, Operands);
+  DeadInstrMorgue.collectPotentiallyDeadInstrs(Stores);
+  if (AllLoads)
+    DeadInstrMorgue.collectPotentiallyDeadInstrs<Value>(Operands);
+  DeadInstrMorgue.tryEraseDeadInstrs();
 
   return acceptOrRevert();
 }
 
-bool LoadStoreVec::runOnRegion(Region &Rgn, const Analyses &A) {
+LoadInst *LoadStoreVec::vectorizeLoads(ArrayRef<Instruction *> Loads,
+                                       Region &Rgn) {
+  if (!VecUtils::areConsecutive<LoadInst, Instruction>(
+          Loads, A->getScalarEvolution(), *DL))
+    return nullptr;
+  auto VecTy = canVectorize(Loads);
+  if (!VecTy)
+    return nullptr;
+
+  // TODO: Support mixed-type top-level load chains.
+  Type *VecElemTy = cast<FixedVectorType>(*VecTy)->getElementType();
+  if (!all_of(Loads, [VecElemTy](Instruction *I) {
+        return VecUtils::getElementType(I->getType()) == VecElemTy;
+      }))
+    return nullptr;
+
+  saveIR(Rgn);
+
+  auto *VecLoad = createVectorLoad(Loads);
+  if (VecLoad == nullptr) {
+    Ctx->accept();
+    return nullptr;
+  }
+
+  BasicBlock::iterator WhereIt = std::next(VecLoad->getIterator());
+  for (auto [Lane, OrigV] : VecUtils::enumerateLanes(Loads)) {
+    auto *OrigLoad = cast<LoadInst>(OrigV);
+    if (OrigLoad->hasNUses(0))
+      continue;
+    Value *Unpacked =
+        VecUtils::unpack(VecLoad, OrigLoad->getType(), Lane, WhereIt);
+    OrigLoad->replaceAllUsesWith(Unpacked);
+  }
+
+  DeadInstrMorgue.collectPotentiallyDeadInstrs(Loads);
+  DeadInstrMorgue.tryEraseDeadInstrs();
+
+  if (!acceptOrRevert())
+    return nullptr;
+  return VecLoad;
+}
+
+bool LoadStoreVec::runOnRegion(Region &Rgn, const Analyses &RegionAnalyses) {
   SmallVector<Instruction *, 8> Bndl(Rgn.getAux().begin(), Rgn.getAux().end());
   if (Bndl.size() < 2)
     return false;
   Function &F = *Bndl[0]->getParent()->getParent();
   DL = &F.getParent()->getDataLayout();
-  Scheduler Sched(A.getAA(), F.getContext(), SchedDirection::BottomUp);
-  return vectorizeStores(Bndl, Rgn, Sched, A);
+  Ctx = &F.getContext();
+  A = &RegionAnalyses;
+  Sched =
+      std::make_unique<Scheduler>(A->getAA(), *Ctx, SchedDirection::BottomUp);
+
+  // SeedCollection only ever gives us a homogeneous seed slice: stores and
+  // loads are collected in separate passes over the BB, never mixed into one
----------------
vporpo wrote:

Too much detail in this comment. We don't need to mention that they get collected in separate passes over the BB.

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