[llvm] [LoopVectorize] Widen selected-base loads with masked loads (PR #213227)

[Florian Hahn via llvm-commits]

================
@@ -5439,6 +5440,118 @@ void VPlanTransforms::makeMemOpWideningDecisions(VPlan &Plan, VFRange &Range,
         return false;
       });
 
+  // Recognize loads through a pointer selected per iteration between two
+  // loop-invariant bases. Although the combined address is not consecutive,
+  // each selected address chain is. When profitable, use two complementary
+  // masked consecutive loads and blend their results instead of gathering a
+  // vector of selected pointers.
+  VPlanTransforms::runPass(
+      "widenSelectedBaseLoads", ProcessSubset, Plan, [&](VPInstruction *VPI) {
+        SelectedBaseLoadPattern Pattern;
+        if (!match(VPI, m_SelectedBaseLoad(Pattern)))
+          return false;
+
+        if (!vputils::isUniformAcrossVFsAndUFs(Pattern.TrueBase) ||
+            !vputils::isUniformAcrossVFsAndUFs(Pattern.FalseBase) ||
+            vputils::isUniformAcrossVFsAndUFs(Pattern.Cond))
+          return false;
+
+        auto *LI = cast<LoadInst>(VPI->getUnderlyingInstr());
+        Type *ScalarTy = LI->getType();
+        // The ingredient's alignment only applies to the address selected by
+        // each scalar iteration. It does not imply that either candidate base
+        // is aligned when that candidate is masked off. Use a conservative
+        // alignment for the newly-speculated base pointers.
+        const Align SplitAlignment(1);
+        if (!CostCtx.Config.isLegalMaskedLoadOrStore(
+                /*IsLoad=*/true, ScalarTy, SplitAlignment,
+                LI->getPointerAddressSpace()))
+          return false;
+
+        VPBuilder Builder(VPI);
+        SmallVector<VPValue *> TrueOps = Pattern.GEPOperands;
+        SmallVector<VPValue *> FalseOps = Pattern.GEPOperands;
+        // replace the baseOp
+        TrueOps[0] = Pattern.TrueBase;
+        FalseOps[0] = Pattern.FalseBase;
+        auto *TrueAddr =
+            Builder.insert(Pattern.Addr->cloneWithOperands(TrueOps));
+        auto *FalseAddr =
+            Builder.insert(Pattern.Addr->cloneWithOperands(FalseOps));
+
+        // Splitting the selected base must expose two unit-stride accesses.
+        // Remove the speculative address recipes again if either proof fails.
+        if (getConstantStride(TrueAddr, ScalarTy, CostCtx.PSE, CostCtx.L) !=
+                1 ||
+            getConstantStride(FalseAddr, ScalarTy, CostCtx.PSE, CostCtx.L) !=
+                1) {
+          TrueAddr->eraseFromParent();
+          FalseAddr->eraseFromParent();
+          return false;
+        }
+
+        bool HasActiveMask = RecipeBuilder.isPredicatedInst(LI);
+        auto IsProfitable = [&](ElementCount VF) {
+          // Compare against the legacy decision for each VF so a scalarization
+          // choice is not replaced just because the gather is expensive.
+          if (VF.isScalable() || CostCtx.willBeScalarized(LI, VF) ||
+              !CostCtx.Config.isLegalGatherOrScatter(LI, VF))
+            return false;
+          Type *VecTy = VectorType::get(ScalarTy, VF);
+          InstructionCost BlendedCost = slpvectorizer::getBlendedLoadCost(
----------------
fhahn wrote:

I don't think this should depend on the SLP vectorizer cost model. Please take a look at how other code in VPlan does it; various recipes expose their cost


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