Revert "[LV][VPlan] Remove any-of reduction from precomputeCost. NFC"

[ElvisWang123]

Reverts #117109

Some test cases need to update.

[llvmbot]

@llvm/pr-subscribers-vectorizers

@llvm/pr-subscribers-llvm-transforms

Author: Elvis Wang (ElvisWang123)

Changes

Reverts llvm/llvm-project#117109

Some test cases need to update.

---

Full diff: https://github.com/llvm/llvm-project/pull/117280.diff

1 Files Affected:

• (modified) llvm/lib/Transforms/Vectorize/LoopVectorize.cpp (+22-2)

diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
index 5b556058cc762c..d13770a35c108f 100644
--- a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
+++ b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
@@ -7303,14 +7303,34 @@ LoopVectorizationPlanner::precomputeCosts(VPlan &Plan, ElementCount VF,
 
   // The legacy cost model has special logic to compute the cost of in-loop
   // reductions, which may be smaller than the sum of all instructions involved
-  // in the reduction.
+  // in the reduction. For AnyOf reductions, VPlan codegen may remove the select
+  // which the legacy cost model uses to assign cost. Pre-compute their costs
+  // for now.
   // TODO: Switch to costing based on VPlan once the logic has been ported.
   for (const auto &[RedPhi, RdxDesc] : Legal->getReductionVars()) {
     if (ForceTargetInstructionCost.getNumOccurrences())
       continue;
 
-    if (!CM.isInLoopReduction(RedPhi))
+    if (!CM.isInLoopReduction(RedPhi) &&
+        !RecurrenceDescriptor::isAnyOfRecurrenceKind(
+            RdxDesc.getRecurrenceKind()))
+      continue;
+
+    // AnyOf reduction codegen may remove the select. To match the legacy cost
+    // model, pre-compute the cost for AnyOf reductions here.
+    if (RecurrenceDescriptor::isAnyOfRecurrenceKind(
+            RdxDesc.getRecurrenceKind())) {
+      auto *Select = cast<SelectInst>(*find_if(
+          RedPhi->users(), [](User *U) { return isa<SelectInst>(U); }));
+      assert(!CostCtx.SkipCostComputation.contains(Select) &&
+             "reduction op visited multiple times");
+      CostCtx.SkipCostComputation.insert(Select);
+      auto ReductionCost = CostCtx.getLegacyCost(Select, VF);
+      LLVM_DEBUG(dbgs() << "Cost of " << ReductionCost << " for VF " << VF
+                        << ":\n any-of reduction " << *Select << "\n");
+      Cost += ReductionCost;
       continue;
+    }
 
     const auto &ChainOps = RdxDesc.getReductionOpChain(RedPhi, OrigLoop);
     SetVector<Instruction *> ChainOpsAndOperands(ChainOps.begin(),

[llvm-ci]

LLVM Buildbot has detected a new failure on builder arc-builder running on arc-worker while building llvm at step 6 "test-build-unified-tree-check-all".

Full details are available at: https://lab.llvm.org/buildbot/#/builders/3/builds/8060

Here is the relevant piece of the build log for the reference

Step 6 (test-build-unified-tree-check-all) failure: 1200 seconds without output running [b'ninja', b'check-all'], attempting to kill
0.450 [59/18/1] Linking CXX executable tools/clang/unittests/InstallAPI/InstallAPITests
0.869 [58/18/2] Linking CXX executable tools/clang/unittests/Basic/BasicTests
1.216 [57/18/3] Linking CXX executable tools/clang/unittests/Format/FormatTests
5.821 [56/18/4] Linking CXX executable tools/clang/unittests/Lex/LexTests
6.323 [55/18/5] Linking CXX executable tools/clang/unittests/libclang/libclangTests
6.505 [54/18/6] Linking CXX executable tools/clang/unittests/AST/ByteCode/InterpTests
6.540 [53/18/7] Linking CXX executable tools/clang/unittests/CrossTU/CrossTUTests
6.810 [52/18/8] Linking CXX executable tools/clang/unittests/Analysis/ClangAnalysisTests
6.920 [51/18/9] Linking CXX executable tools/clang/unittests/Tooling/Syntax/SyntaxTests
6.964 [50/18/10] Linking CXX executable tools/clang/unittests/libclang/CrashTests/libclangCrashTests
command timed out: 1200 seconds without output running [b'ninja', b'check-all'], attempting to kill
process killed by signal 9
program finished with exit code -1
elapsedTime=1270.064332