[llvm] [VPlan] Wrap SCEVAddRecExprs in VPExpandSCEVRecipe in VPSCEVExpander. (PR #210013)

[Ramkumar Ramachandra via llvm-commits]

================
@@ -873,22 +876,37 @@ VPValue *VPSCEVExpander::tryToExpand(const SCEV *S) {
       llvm_unreachable("Unexpected min/max SCEV type");
     }
     // Chain operands in reverse order matching SCEVExpander's expansion of
-    // min/max expressions.
-    SmallVector<VPValue *, 2> Ops;
-    for (const SCEVUse &Op : reverse(MinMax->operands())) {
-      VPValue *OpV = tryToExpand(Op);
-      if (!OpV)
-        return nullptr;
-      Ops.push_back(OpV);
-    }
+    // min/max expressions. In SafeUDivMode freeze expansion results of operands
+    // other than the first for sequential UMins, to avoid short-circuiting
+    // divide-by-0/poison.
+    bool IsSequential = S->getSCEVType() == scSequentialUMinExpr;
     Type *ResultTy = MinMax->getType();
-    VPValue *Result = Ops.front();
-    for (VPValue *Op : drop_begin(Ops))
-      Result = Builder.createScalarIntrinsic(IntrinsicID, {Result, Op},
-                                             ResultTy, DL);
+    bool PrevSafeMode = SafeUDivMode;
+    VPValue *Result = nullptr;
+    for (const auto &[I, SCEVOp] : enumerate(reverse(MinMax->operands()))) {
+      bool MayShortCircuit = IsSequential && I != MinMax->getNumOperands() - 1;
+      SafeUDivMode = MayShortCircuit || PrevSafeMode;
+      VPValue *Op = expand(SCEVOp);
+      SafeUDivMode = PrevSafeMode;
+      if (MayShortCircuit)
+        Op = Builder.createScalarFreeze(Op, ResultTy, DL);
+      Result = Result ? Builder.createScalarIntrinsic(IntrinsicID, {Result, Op},
+                                                      ResultTy, DL)
+                      : Op;
+    }
     return Result;
   }
-  default:
-    return nullptr;
+  case scAddRecExpr: {
+    [[maybe_unused]] BasicBlock *PH =
+        cast<VPIRBasicBlock>(Builder.getPlan().getEntry())->getIRBasicBlock();
+    assert(
+        SE.DT.dominates(cast<SCEVAddRecExpr>(S)->getLoop()->getHeader(), PH) &&
+        "can only expand AddRecs for loops outside VPlan's scope");
+    // AddRecs outside VPlan's scope must be expanded via VPExpandSCEV.
----------------
artagnon wrote:

I figured out what happens: this assert is trivially true, because we're checking the underlying IR: we simply append IR BasicBlocks to the underlying IR, and wrap it in a VPIRBasicBlock in the Plan. The VPIRBasicBlock corresponding to Loop's AR would not dominate the Plan's entry. Yes, the IR BasicBlocks are manipulated under VPIRBasicBlock (which is simply a transparent handle), and the AR's loop doesn't need to correspond to any VPIRBasicBlock's underlying IR BasicBlock, and in fact, the AR's loop header may refer to a BasicBlock that's completely absent in the plan -- so yes, your patch is correct is general case (modulo the misleading assert), but my patch refines it further when the AR's loop does correspond to the VPIRBasicBlock's underlying IR BasicBlock (and it does happen in a few cases).

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