Expand loop peeling phi computation to handle binary ops and casts

[jamieschmeiser]

Summary

Expand the capabilities of the code for computing how many peels are
needed to make phis determined. A cast gets the peel count for the
value being casted while a binary op gets the maximum of the operands.

[syzaara]

llvm/lib/Transforms/Utils/LoopPeel.cpp

236

I think these asserts are not needed since BinaryOp and CmpInst will always have 2 operands. Same for the assert below with the cast instruction.

[jamieschmeiser]

Respond to review comments: remove unnecessary asserts

[mkazantsev]

Looks reasonable.

llvm/lib/Transforms/Utils/LoopPeel.cpp

if (const PHINode *Phi = dyn_cast<PHINode>(&V)) {
     }
     // We need to analyze the input from the back edge and add 1.
     Value *Input = Phi->getIncomingValueForBlock(L.getLoopLatch());
     PeelCounter Iterations = calculate(*Input);
     assert(IterationsToInvariance[Input] == Iterations &&
            "unexpected value saved");
     return (IterationsToInvariance[Phi] = addOne(Iterations));
   }
-  // TODO: handle expressions
+  if (const Instruction *I = dyn_cast<Instruction>(&V)) {
+    if (isa<CmpInst>(I) || I->isBinaryOp()) {
+      // Binary instructions get the max of the operands.
+      PeelCounter LHS = calculate(*I->getOperand(0));
+      if (LHS == Unknown)
+        return Unknown;
+      PeelCounter RHS = calculate(*I->getOperand(1));
+      if (RHS == Unknown)
+        return Unknown;
+      return (IterationsToInvariance[I] = {std::max(*LHS, *RHS)});
+    }
+    if (I->isCast())
+      // Cast instructions get the value of the operand.
+      return (IterationsToInvariance[I] = calculate(*I->getOperand(0)));
+  }
+  // TODO: handle more expressions
 
   // Everything else is Unknown.
   assert(IterationsToInvariance[&V] == Unknown && "unexpected value saved");
   return Unknown;
 }
 
 Optional<unsigned> PhiAnalyzer::calculateIterationsToPeel() {
   unsigned Iterations = 0;

llvm/test/Transforms/LoopUnroll/peel-loop-phi-analysis.ll

 ; are removed since x is based on a cast of y, which is based on a, which is
 ; set on the backedge.
 ; Consider the calls to g and f.
 ; First iteration: g(0), x=0, f(0.0), y=0.0, a=5.0
 ; Second iteration: g(0), x=0, f(0.0), y=5.0, a=5.0
 ; Third iteration: g(0), x=5 (requires cast), f(5.0), a=5.0
 ; Fourth iteration (and subsequent): g(5), x=5, f(5.0), a=5.0
 ; Therefore, peeling 3 times removes the phi nodes, so check for 3 peels.
-; CURRENT LIMITATION: only peels twice because cannot handle cast
 ;
 ; void castTest() {
 ;   int x = 0;
 ;   float y = 0.0;
 ;   float a = 0.0;
 ;   for(int i = 0; i <100000; ++i) {
 ;     g(x);
 ;     x = y;

 ; CHECK:       for.body.peel2:
 ; CHECK-NEXT:    tail call void @_Z1gi(i32 noundef signext [[CONV_PEEL]])
 ; CHECK-NEXT:    [[CONV_PEEL3:%.*]] = fptosi float 0.000000e+00 to i32
 ; CHECK-NEXT:    tail call void @_Z1ff(float noundef 0.000000e+00)
 ; CHECK-NEXT:    [[INC_PEEL4:%.*]] = add nuw nsw i32 [[INC_PEEL]], 1
 ; CHECK-NEXT:    [[EXITCOND_PEEL5:%.*]] = icmp ne i32 [[INC_PEEL4]], 100000
 ; CHECK-NEXT:    br i1 [[EXITCOND_PEEL5]], label [[FOR_BODY_PEEL_NEXT1:%.*]], label [[FOR_COND_CLEANUP]]
 ; CHECK:       for.body.peel.next1:
-; CHECK-NEXT:    br label [[FOR_BODY_PEEL_NEXT6:%.*]]
+; CHECK-NEXT:    br label [[FOR_BODY_PEEL7:%.*]]
+; CHECK:       for.body.peel7:
+; CHECK-NEXT:    tail call void @_Z1gi(i32 noundef signext [[CONV_PEEL3]])
+; CHECK-NEXT:    [[CONV_PEEL8:%.*]] = fptosi float 5.000000e+00 to i32
+; CHECK-NEXT:    tail call void @_Z1ff(float noundef 5.000000e+00)
+; CHECK-NEXT:    [[INC_PEEL9:%.*]] = add nuw nsw i32 [[INC_PEEL4]], 1
+; CHECK-NEXT:    [[EXITCOND_PEEL10:%.*]] = icmp ne i32 [[INC_PEEL9]], 100000
+; CHECK-NEXT:    br i1 [[EXITCOND_PEEL10]], label [[FOR_BODY_PEEL_NEXT6:%.*]], label [[FOR_COND_CLEANUP]]
 ; CHECK:       for.body.peel.next6:
+; CHECK-NEXT:    br label [[FOR_BODY_PEEL_NEXT11:%.*]]
+; CHECK:       for.body.peel.next11:
 ; CHECK-NEXT:    br label [[ENTRY_PEEL_NEWPH:%.*]]
 ; CHECK:       entry.peel.newph:
 ; CHECK-NEXT:    br label [[FOR_BODY:%.*]]
 ; CHECK:       for.cond.cleanup.loopexit:
 ; CHECK-NEXT:    br label [[FOR_COND_CLEANUP]]
 ; CHECK:       for.cond.cleanup:
 ; CHECK-NEXT:    ret void
 ; CHECK:       for.body:
-; CHECK-NEXT:    [[I:%.*]] = phi i32 [ [[INC_PEEL4]], [[ENTRY_PEEL_NEWPH]] ], [ [[INC:%.*]], [[FOR_BODY]] ]
-; CHECK-NEXT:    [[X:%.*]] = phi i32 [ [[CONV_PEEL3]], [[ENTRY_PEEL_NEWPH]] ], [ 5, [[FOR_BODY]] ]
+; CHECK-NEXT:    [[I:%.*]] = phi i32 [ [[INC_PEEL9]], [[ENTRY_PEEL_NEWPH]] ], [ [[INC:%.*]], [[FOR_BODY]] ]
+; CHECK-NEXT:    [[X:%.*]] = phi i32 [ [[CONV_PEEL8]], [[ENTRY_PEEL_NEWPH]] ], [ 5, [[FOR_BODY]] ]
 ; CHECK-NEXT:    tail call void @_Z1gi(i32 noundef signext [[X]])
 ; CHECK-NEXT:    tail call void @_Z1ff(float noundef 5.000000e+00)
 ; CHECK-NEXT:    [[INC]] = add nuw nsw i32 [[I]], 1
 ; CHECK-NEXT:    [[EXITCOND:%.*]] = icmp ne i32 [[INC]], 100000
 ; CHECK-NEXT:    br i1 [[EXITCOND]], label [[FOR_BODY]], label [[FOR_COND_CLEANUP_LOOPEXIT:%.*]], !llvm.loop [[LOOP0:![0-9]+]]
 ;
 entry:
   br label %for.body

 ; are removed since x is based on y, which is based on a, which is based
 ; on a binary add of a phi and a constant.
 ; Consider the calls to g:
 ; First iteration: g(0), x=0, g(0), y=1, a=5
 ; Second iteration: g(0), x=1, g(5), y=6(binary operator), a=5
 ; Third iteration: g(1), x=6, g(5), y=6, a=5
 ; Fourth iteration (and subsequent): g(6), x=6, g(5), y=6, a=5
 ; Therefore, peeling 3 times removes the phi nodes.
-; CURRENT_LIMITATION: only peels once because cannot handle binary operator
 ;
 ; void g(int);
 ; void binary() {
 ;   int x = 0;
 ;   int y = 0;
 ;   int a = 0;
 ;   for(int i = 0; i <100000; ++i) {
 ;     g(x);

 ; CHECK:       for.body.peel:
 ; CHECK-NEXT:    tail call void @_Z1gi(i32 signext 0)
 ; CHECK-NEXT:    tail call void @_Z1gi(i32 signext 0)
 ; CHECK-NEXT:    [[ADD_PEEL:%.*]] = add nuw nsw i32 0, 1
 ; CHECK-NEXT:    [[INC_PEEL:%.*]] = add nuw nsw i32 0, 1
 ; CHECK-NEXT:    [[EXITCOND_PEEL:%.*]] = icmp eq i32 [[INC_PEEL]], 100000
 ; CHECK-NEXT:    br i1 [[EXITCOND_PEEL]], label [[FOR_COND_CLEANUP:%.*]], label [[FOR_BODY_PEEL_NEXT:%.*]]
 ; CHECK:       for.body.peel.next:
-; CHECK-NEXT:    br label [[FOR_BODY_PEEL_NEXT1:%.*]]
+; CHECK-NEXT:    br label [[FOR_BODY_PEEL2:%.*]]
+; CHECK:       for.body.peel2:
+; CHECK-NEXT:    tail call void @_Z1gi(i32 signext 0)
+; CHECK-NEXT:    tail call void @_Z1gi(i32 signext 5)
+; CHECK-NEXT:    [[INC_PEEL4:%.*]] = add nuw nsw i32 [[INC_PEEL]], 1
+; CHECK-NEXT:    [[EXITCOND_PEEL5:%.*]] = icmp eq i32 [[INC_PEEL4]], 100000
+; CHECK-NEXT:    br i1 [[EXITCOND_PEEL5]], label [[FOR_COND_CLEANUP]], label [[FOR_BODY_PEEL_NEXT1:%.*]]
 ; CHECK:       for.body.peel.next1:
+; CHECK-NEXT:    br label [[FOR_BODY_PEEL7:%.*]]
+; CHECK:       for.body.peel7:
+; CHECK-NEXT:    tail call void @_Z1gi(i32 signext [[ADD_PEEL]])
+; CHECK-NEXT:    tail call void @_Z1gi(i32 signext 5)
+; CHECK-NEXT:    [[INC_PEEL9:%.*]] = add nuw nsw i32 [[INC_PEEL4]], 1
+; CHECK-NEXT:    [[EXITCOND_PEEL10:%.*]] = icmp eq i32 [[INC_PEEL9]], 100000
+; CHECK-NEXT:    br i1 [[EXITCOND_PEEL10]], label [[FOR_COND_CLEANUP]], label [[FOR_BODY_PEEL_NEXT6:%.*]]
+; CHECK:       for.body.peel.next6:
+; CHECK-NEXT:    br label [[FOR_BODY_PEEL_NEXT11:%.*]]
+; CHECK:       for.body.peel.next11:
 ; CHECK-NEXT:    br label [[ENTRY_PEEL_NEWPH:%.*]]
 ; CHECK:       entry.peel.newph:
 ; CHECK-NEXT:    br label [[FOR_BODY:%.*]]
 ; CHECK:       for.cond.cleanup.loopexit:
 ; CHECK-NEXT:    br label [[FOR_COND_CLEANUP]]
 ; CHECK:       for.cond.cleanup:
 ; CHECK-NEXT:    ret void
 ; CHECK:       for.body:
-; CHECK-NEXT:    [[I:%.*]] = phi i32 [ [[INC_PEEL]], [[ENTRY_PEEL_NEWPH]] ], [ [[INC:%.*]], [[FOR_BODY]] ]
-; CHECK-NEXT:    [[Y:%.*]] = phi i32 [ [[ADD_PEEL]], [[ENTRY_PEEL_NEWPH]] ], [ 6, [[FOR_BODY]] ]
-; CHECK-NEXT:    [[X:%.*]] = phi i32 [ 0, [[ENTRY_PEEL_NEWPH]] ], [ [[Y]], [[FOR_BODY]] ]
-; CHECK-NEXT:    tail call void @_Z1gi(i32 signext [[X]])
+; CHECK-NEXT:    [[I:%.*]] = phi i32 [ [[INC_PEEL9]], [[ENTRY_PEEL_NEWPH]] ], [ [[INC:%.*]], [[FOR_BODY]] ]
+; CHECK-NEXT:    tail call void @_Z1gi(i32 signext 6)
 ; CHECK-NEXT:    tail call void @_Z1gi(i32 signext 5)
 ; CHECK-NEXT:    [[INC]] = add nuw nsw i32 [[I]], 1
 ; CHECK-NEXT:    [[EXITCOND:%.*]] = icmp eq i32 [[INC]], 100000
 ; CHECK-NEXT:    br i1 [[EXITCOND]], label [[FOR_COND_CLEANUP_LOOPEXIT:%.*]], label [[FOR_BODY]], !llvm.loop [[LOOP2:![0-9]+]]
 ;
 entry:
   br label %for.body