[IndVarSimplify][LoopUtils] Avoid TOCTOU/ordering issues (PR45835)

[lebedev.ri]

Summary

Currently, rewriteLoopExitValues()'s logic is roughly as following:

Loop over each incoming value in each PHI node.
Query whether the SCEV for that incoming value is high-cost.
Expand the SCEV.
Perform sanity check (isValidRewrite(), D51582)
Record the info
Afterwards, see if we can drop the loop given replacements.
Maybe perform replacements.

The problem is that we interleave SCEV cost checking and expansion.
This is A Problem, because isHighCostExpansion() takes special care
to not bill for the expansions that were already expanded, and we can reuse.

While it makes sense in general - if we know that we will expand some SCEV,
all the other SCEV's costs should account for that, which might cause
some of them to become non-high-cost too, and cause chain reaction.

But that isn't what we are doing here. We expand *all* SCEV's, unconditionally.
So every next SCEV's cost will be affected by the already-performed expansions
for previous SCEV's. Even if we are not planning on keeping
some of the expansions we performed.

Worse yet, this current "bonus" depends on the exact PHI node
incoming value processing order. This is completely wrong.

As an example of an issue, see @dmajor's pr45835.ll - if we happen to have
a PHI node with two(!) identical high-cost incoming values for the same basic blocks,
we would decide first time around that it is high-cost, expand it,
and immediately decide that it is not high-cost because we have an expansion
that we could reuse (because we expanded it right before, temporarily),
and replace the second incoming value but not the first one;
thus resulting in a broken PHI.

What we instead should do for now, is not perform any expansions
until after we've queried all the costs.

Later, in particular after isValidRewrite() is an assertion (D51582)
we could improve upon that, but in a more coherent fashion.

See PR45835

[dmajor]

Is it a problem that the test case doesn't fail on trunk before the patch, since SCEVMinMaxExpr is no longer automatically high cost? I noted this in D79720 but wasn't sure how to nudge the code around it.

[lebedev.ri]

In D79787#2031598, @dmajor wrote:

Is it a problem that the test case doesn't fail on trunk before the patch, since SCEVMinMaxExpr is no longer automatically high cost? I noted this in D79720 but wasn't sure how to nudge the code around it.

Oh right, fixed.

[dmajor]

Thanks a lot for taking this over! Looks good from my side but you may wish to wait for a more expert review.

I also confirmed that this still fixes our issue on the 10.0 branch. I had to remove -scev-cheap-expansion-budget=1 since that didn't exist in 10, but fortunately it's not needed there. Since it wasn't a clean merge, I put the diff in P8220 if it makes your/Tom's life easier.

[lebedev.ri]

In D79787#2031793, @dmajor wrote:

Thanks a lot for taking this over! Looks good from my side but you may wish to wait for a more expert review.

I also confirmed that this still fixes our issue on the 10.0 branch. I had to remove -scev-cheap-expansion-budget=1 since that didn't exist in 10, but fortunately it's not needed there. Since it wasn't a clean merge, I put the diff in P8220 if it makes your/Tom's life easier.

Thank you for the review!
Waiting on @reames / @mkazantsev / ...

[lebedev.ri]

ping

[mkazantsev]

In general, I'm ok with nits regarding naming. However this FIXME is concerning. Could you please add a XFAIL test that shows the difference between current state and when this check is turned into assert?

llvm/lib/Transforms/Utils/LoopUtils.cpp

1219

Please give it either more clear name or more clear comment. The actual SCEV of what? Besides, please avoid including word Value into the name of SCEVs in a code that has both, it creates confusion. :)

1220

Nit: maybe call it insertion/expansion point?

1226

Initialization?

[lebedev.ri]

Thank you for taking a look!

In D79787#2045773, @mkazantsev wrote:

In general, I'm ok with nits regarding naming.

However this FIXME is concerning. Could you please add a XFAIL test that shows the difference between current state and when this check is turned into assert?

I don't have an actionable example at hand. I could just drop FIXME if it's too weird.

The thought is:
since the current as we have just seen in the new testcase the order of expansion affects the cost calculation,
then

1. may it be the case that we might have two exit values (A and B) that we can rewrite
2. one is high-cost (A) while other one is not high-cost (B)
3. the high-cost A is based on low-cost B (something like A = B + 42)

Therefore, if we'd expand low-cost B first, then the cost of high-cost A would be lower,
potentially becoming low-cost, and we could be okay expanding it too.
Let me know if that does/doesn't make sense.

[lebedev.ri]

Attempted to address naming nits. Is this any better?

[mkazantsev]

I don't have an actionable example at hand. I could just drop FIXME if it's too weird.

Then please add a debug printout that it happened.

Otherwise I don't have much objections. LGTM.

llvm/lib/Transforms/Utils/LoopUtils.cpp

1226

Better call it ValidRewrite because of HighCost without Is above.

[lebedev.ri]

In D79787#2048569, @mkazantsev wrote:

I don't have an actionable example at hand. I could just drop FIXME if it's too weird.

Then please add a debug printout that it happened.

I'm sorry, i'm not sure i follow. It's not a bug, it's a potential improvement.
There is no point in trying to detect such a situation right now
because it will bring no benefit but will cost us compile time.

Otherwise I don't have much objections. LGTM.

Thank you for the review!

llvm/lib/Transforms/Utils/LoopUtils.cpp

while (!WorkList.empty()) {
     }
   }
   return false;
 }
 
 // Collect information about PHI nodes which can be transformed in
 // rewriteLoopExitValues.
 struct RewritePhi {
-  PHINode *PN;
-  unsigned Ith;   // Ith incoming value.
-  Value *Val;     // Exit value after expansion.
-  bool HighCost;  // High Cost when expansion.
-
-  RewritePhi(PHINode *P, unsigned I, Value *V, bool H)
-      : PN(P), Ith(I), Val(V), HighCost(H) {}
+  PHINode *PN;               // For which PHI node is this replacement?
+  unsigned Ith;              // For which incoming value?
+  const SCEV *ExpansionSCEV; // The SCEV of the incoming value we are rewriting.
+  Instruction *ExpansionPoint; // Where we'd like to expand that SCEV?
+  bool HighCost;               // Is this expansion a high-cost?
+
+  Value *Expansion = nullptr;
+  bool IsValidRewrite = false;
+
+  RewritePhi(PHINode *P, unsigned I, const SCEV *Val, Instruction *ExpansionPt,
+             bool H)
+      : PN(P), Ith(I), ExpansionSCEV(Val), ExpansionPoint(ExpansionPt),
+        HighCost(H) {}
 };
 
 // Check whether it is possible to delete the loop after rewriting exit
 // value. If it is possible, ignore ReplaceExitValue and do rewriting
 // aggressively.
 static bool canLoopBeDeleted(Loop *L, SmallVector<RewritePhi, 8> &RewritePhiSet) {
   BasicBlock *Preheader = L->getLoopPreheader();
   // If there is no preheader, the loop will not be deleted.
static bool canLoopBeDeleted(Loop *L, SmallVector<RewritePhi, 8> &RewritePhiSet) {
   while (PHINode *P = dyn_cast<PHINode>(BI)) {
     Value *Incoming = P->getIncomingValueForBlock(ExitingBlocks[0]);
 
     // If the Incoming value of P is found in RewritePhiSet, we know it
     // could be rewritten to use a loop invariant value in transformation
     // phase later. Skip it in the loop invariant check below.
     bool found = false;
     for (const RewritePhi &Phi : RewritePhiSet) {
+      if (!Phi.IsValidRewrite)
+        continue;
       unsigned i = Phi.Ith;
       if (Phi.PN == P && (Phi.PN)->getIncomingValue(i) == Incoming) {
         found = true;
         break;
       }
     }
 
     Instruction *I;
while ((PN = dyn_cast<PHINode>(BBI++))) {
         // definitely used inside the loop in a way which can not be optimized
         // away. Avoid doing so unless we know we have a value which computes
         // the ExitValue already. TODO: This should be merged into SCEV
         // expander to leverage its knowledge of existing expressions.
         if (ReplaceExitValue != AlwaysRepl && !isa<SCEVConstant>(ExitValue) &&
             !isa<SCEVUnknown>(ExitValue) && hasHardUserWithinLoop(L, Inst))
           continue;
 
+        // Check if expansions of this SCEV would count as being high cost.
         bool HighCost = Rewriter.isHighCostExpansion(
             ExitValue, L, SCEVCheapExpansionBudget, TTI, Inst);
-        Value *ExitVal = Rewriter.expandCodeFor(ExitValue, PN->getType(), Inst);
+
+        // Note that we must not perform expansions until after
+        // we query *all* the costs, because if we perform temporary expansion
+        // inbetween, one that we might not intend to keep, said expansion
+        // *may* affect cost calculation of the the next SCEV's we'll query,
+        // and next SCEV may errneously get smaller cost.
+
+        // Collect all the candidate PHINodes to be rewritten.
+        RewritePhiSet.emplace_back(PN, i, ExitValue, Inst, HighCost);
+      }
+    }
+  }
+
+  // Now that we've done preliminary filtering and billed all the SCEV's,
+  // we can perform the last sanity check - the expansion must be valid.
+  for (RewritePhi &Phi : RewritePhiSet) {
+    Phi.Expansion = Rewriter.expandCodeFor(Phi.ExpansionSCEV, Phi.PN->getType(),
+                                           Phi.ExpansionPoint);
 
-        LLVM_DEBUG(dbgs() << "rewriteLoopExitValues: AfterLoopVal = "
-                          << *ExitVal << '\n' << "  LoopVal = " << *Inst
-                          << "\n");
+                      << *(Phi.Expansion) << '\n'
+                      << "  LoopVal = " << *(Phi.ExpansionPoint) << "\n");
 
-        if (!isValidRewrite(SE, Inst, ExitVal)) {
-          DeadInsts.push_back(ExitVal);
-          continue;
-        }
+    // FIXME: isValidRewrite() is a hack. it should be an assert, eventually.
+    Phi.IsValidRewrite = isValidRewrite(SE, Phi.ExpansionPoint, Phi.Expansion);
+    if (!Phi.IsValidRewrite) {
+      DeadInsts.push_back(Phi.Expansion);
 
 #ifndef NDEBUG
-        // If we reuse an instruction from a loop which is neither L nor one of
-        // its containing loops, we end up breaking LCSSA form for this loop by
-        // creating a new use of its instruction.
-        if (auto *ExitInsn = dyn_cast<Instruction>(ExitVal))
-          if (auto *EVL = LI->getLoopFor(ExitInsn->getParent()))
-            if (EVL != L)
-              assert(EVL->contains(L) && "LCSSA breach detected!");
+    if (auto *ExitInsn = dyn_cast<Instruction>(Phi.Expansion))
 #endif
-
-        // Collect all the candidate PHINodes to be rewritten.
-        RewritePhiSet.emplace_back(PN, i, ExitVal, HighCost);
-      }
-    }
   }
 
+  // TODO: after isValidRewrite() is an assertion, evaluate whether
+  // it is beneficial to change how we calculate high-cost:
+  // if we have SCEV 'A' which we know we will expand, should we calculate
+  // the cost of other SCEV's after expanding SCEV 'A',
+  // thus potentially giving cost bonus to those other SCEV's?
+
   bool LoopCanBeDel = canLoopBeDeleted(L, RewritePhiSet);
   int NumReplaced = 0;
 
   // Transformation.
   for (const RewritePhi &Phi : RewritePhiSet) {
+    if (!Phi.IsValidRewrite)
+      continue;
+
     PHINode *PN = Phi.PN;
-    Value *ExitVal = Phi.Val;
+    Value *ExitVal = Phi.Expansion;
 
     // Only do the rewrite when the ExitValue can be expanded cheaply.
     // If LoopCanBeDel is true, rewrite exit value aggressively.
     if (ReplaceExitValue == OnlyCheapRepl && !LoopCanBeDel && Phi.HighCost) {
       DeadInsts.push_back(ExitVal);
       continue;
     }
 

llvm/test/Transforms/IndVarSimplify/pr45835.ll

+; RUN: opt < %s -indvars -replexitval=always -S | FileCheck %s --check-prefix=ALWAYS
+; RUN: opt < %s -indvars -replexitval=never -S | FileCheck %s --check-prefix=NEVER
+; RUN: opt < %s -indvars -replexitval=cheap -scev-cheap-expansion-budget=1 -S | FileCheck %s --check-prefix=CHEAP
+
+; rewriteLoopExitValues() must rewrite all or none of a PHI's values from a given block.
+
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+
+@a = common global i8 0, align 1
+
+define internal fastcc void @d(i8* %c) unnamed_addr #0 {
+entry:
+  %cmp = icmp ule i8* %c, getelementptr inbounds (i8, i8* @a, i64 65535)
+  %add.ptr = getelementptr inbounds i8, i8* %c, i64 -65535
+  br label %while.cond
+
+while.cond:
+  br i1 icmp ne (i8 0, i8 0), label %cont, label %while.end
+
+cont:
+  %a.mux = select i1 %cmp, i8* @a, i8* %add.ptr
+  switch i64 0, label %while.cond [
+    i64 -1, label %handler.pointer_overflow.i
+    i64 0, label %handler.pointer_overflow.i
+  ]
+
+handler.pointer_overflow.i:
+  %a.mux.lcssa4 = phi i8* [ %a.mux, %cont ], [ %a.mux, %cont ]
+; ALWAYS: [ %scevgep, %cont ], [ %scevgep, %cont ]
+; NEVER: [ %a.mux, %cont ], [ %a.mux, %cont ]
+; In cheap mode, use either one as long as it's consistent.
+; CHEAP: [ %[[VAL:.*]], %cont ], [ %[[VAL]], %cont ]
+  %x5 = ptrtoint i8* %a.mux.lcssa4 to i64
+  br label %while.end
+
+while.end:
+  ret void
+}