[LoopPeel] Check for non-LCSSA form loops

[anna]

Summary

Currently, the way loop peeling is done, it requires LCSSA form.
We need to add this as a check in canPeel (similar to the requirement
for LoopSimplifyForm in canPeel). Without the check if we peel
non-LCSSA loops, we end up breaking SSA form on uses outside the loop.
This cannot be seen upstream because all uses of loop peeling is done on
LCSSA form loops.

Added a unit test for this.

[mkazantsev]

This patch does not account for the fact that DT and LI can both be nullptr. I think it may be reasonable to lift this restriction, but so far the pass supports this case. If you can't go w/o DT and LI, then we first need to rework the pass in a way that it always expects DT and LI to be provided.

llvm/lib/Transforms/Scalar/LoopFuse.cpp

195

Loop peeling pass is not guaranteed to be provided with DT and LI. They can both be (and sometimes are) nullptr, at least in old PM.

llvm/lib/Transforms/Utils/LoopPeel.cpp

748

This should also work for DT == nullptr, see checks below (line 764).

[reames]

Basic idea looks reasonable, though we appear to have bit of cleanup to do before this can reasonably land.

llvm/lib/Transforms/Scalar/LoopFuse.cpp

195

Max, your comment may be correct for LoopPeeling (haven't checked yet), but the LoopFuze code does appear to require DT and LI today. As such, they can't be null along this particular codepath.

llvm/lib/Transforms/Utils/LoopPeel.cpp

748

Chasing through the code, I think the handling for nullptr DT is simply dead. It definitely used to be supported, but it looks like the callers (both fuse and unroll) now require DT.

I also see some code which looks like it would crash if a null DT was actually passed in.

Anna, could you do a cleanup which folded the DT null checks and asserted it was never null? I think we need that before we can land this one.

llvm/unittests/Transforms/Utils/LoopPeelTest.cpp

1

Any reason this can't be a lit test? Strongly preferred.

[anna]

llvm/lib/Transforms/Utils/LoopPeel.cpp

748

Yes, I'll do the cleanup. Note that the only use case of concern maybe downstream users with peeling without DT? But we definitely need DT to make this change.

llvm/unittests/Transforms/Utils/LoopPeelTest.cpp

1

Currently, Loop peeling is not a pass in itself and the passes which use it in upstream requires LCSSA form or constructs it (LoopUnroll).

Now, that I mention it, I don't see LCSSA being a requirement in loop Fusion. So, maybe I can construct a testcase with fusion.

[mkazantsev]

llvm/lib/Transforms/Scalar/LoopFuse.cpp

195

Then the first step would be to replace DominatorTree * with DominatorTree & to be adamant. Now this fact is at least not obvious.

llvm/lib/Transforms/Utils/LoopPeel.cpp

748

I vaguely remember that we've seen some real cases of null DT while working on some peeling patch. Maybe it only happens in unit tests, though. I don't mind if we rework this code so that DT is always non-null, but let's clearly state it (and replace pointer with reference).

[mkazantsev]

llvm/unittests/Transforms/Utils/LoopPeelTest.cpp

1

I think LoopUnroll can be forced to only do peeling by setting relevant options.

[reames]

llvm/lib/Transforms/Scalar/LoopFuse.cpp

195

The first step is to add asserts, run tests, and check them in if nothing fails. The *second* step is to convert the assert into a refactoring change to use references where possible.

llvm/lib/Transforms/Utils/LoopPeel.cpp

748

I vaguely remember this too. That's why I want asserts to find our counter example if one exists. :)

[anna]

llvm/lib/Transforms/Utils/LoopPeel.cpp

748

DT is definitely non-null and we have an (unintentional) assert for this at line 915 already :). This makes the code cleanup here simpler.

[fhahn]

llvm/lib/Transforms/Utils/LoopPeel.cpp

748

Thanks for pushing on this cleanup!

llvm/include/llvm/Transforms/Utils/LoopPeel.h

 
 #ifndef LLVM_TRANSFORMS_UTILS_LOOPPEEL_H
 #define LLVM_TRANSFORMS_UTILS_LOOPPEEL_H
 
 #include "llvm/Analysis/TargetTransformInfo.h"
 
 namespace llvm {
 
-bool canPeel(Loop *L);
+bool canPeel(Loop *L, DominatorTree &DT, LoopInfo &LI);
 
 bool peelLoop(Loop *L, unsigned PeelCount, LoopInfo *LI, ScalarEvolution *SE,
               DominatorTree *DT, AssumptionCache *AC, bool PreserveLCSSA);
 
 TargetTransformInfo::PeelingPreferences
 gatherPeelingPreferences(Loop *L, ScalarEvolution &SE,
                          const TargetTransformInfo &TTI,
                          Optional<bool> UserAllowPeeling,
                          Optional<bool> UserAllowProfileBasedPeeling,
                          bool UnrollingSpecficValues = false);
 
 void computePeelCount(Loop *L, unsigned LoopSize,
                       TargetTransformInfo::PeelingPreferences &PP,
-                      unsigned TripCount, DominatorTree &DT,
+                      unsigned TripCount, DominatorTree &DT, LoopInfo &LI,
                       ScalarEvolution &SE, unsigned Threshold = UINT_MAX);
 
 } // end namespace llvm
 
 #endif // LLVM_TRANSFORMS_UTILS_LOOPPEEL_H

llvm/lib/Transforms/Scalar/LoopFuse.cpp

struct FusionCandidate {
   bool Peeled;
 
   /// Dominator and PostDominator trees are needed for the
   /// FusionCandidateCompare function, required by FusionCandidateSet to
   /// determine where the FusionCandidate should be inserted into the set. These
   /// are used to establish ordering of the FusionCandidates based on dominance.
   const DominatorTree *DT;
   const PostDominatorTree *PDT;
+  /// Needed to identify if peeling is valid.
+  LoopInfo *LI;
 
   OptimizationRemarkEmitter &ORE;
 
   FusionCandidate(Loop *L, const DominatorTree *DT,
-                  const PostDominatorTree *PDT, OptimizationRemarkEmitter &ORE,
-                  TTI::PeelingPreferences PP)
+                  const PostDominatorTree *PDT, LoopInfo *LI,
+                  OptimizationRemarkEmitter &ORE, TTI::PeelingPreferences PP)
       : Preheader(L->getLoopPreheader()), Header(L->getHeader()),
         ExitingBlock(L->getExitingBlock()), ExitBlock(L->getExitBlock()),
         Latch(L->getLoopLatch()), L(L), Valid(true),
-        GuardBranch(L->getLoopGuardBranch()), PP(PP), AbleToPeel(canPeel(L)),
-        Peeled(false), DT(DT), PDT(PDT), ORE(ORE) {
+        GuardBranch(L->getLoopGuardBranch()), PP(PP),
+        AbleToPeel(canPeel(L, *(const_cast<DominatorTree *>(DT)), *LI)),
+        Peeled(false), DT(DT), PDT(PDT), LI(LI), ORE(ORE) {
 
     // Walk over all blocks in the loop and check for conditions that may
     // prevent fusion. For each block, walk over all instructions and collect
     // the memory reads and writes If any instructions that prevent fusion are
     // found, invalidate this object and return.
     for (BasicBlock *BB : L->blocks()) {
       if (BB->hasAddressTaken()) {
         invalidate();
private:
 
   /// Iterate over all loops in the given loop set and identify the loops that
   /// are eligible for fusion. Place all eligible fusion candidates into Control
   /// Flow Equivalent sets, sorted by dominance.
   void collectFusionCandidates(const LoopVector &LV) {
     for (Loop *L : LV) {
       TTI::PeelingPreferences PP =
           gatherPeelingPreferences(L, SE, TTI, None, None);
-      FusionCandidate CurrCand(L, &DT, &PDT, ORE, PP);
+      FusionCandidate CurrCand(L, &DT, &PDT, &LI, ORE, PP);
       if (!CurrCand.isEligibleForFusion(SE))
         continue;
 
       // Go through each list in FusionCandidates and determine if L is control
       // flow equivalent with the first loop in that list. If it is, append LV.
       // If not, go to the next list.
       // If no suitable list is found, start another list and add it to
       // FusionCandidates.
for (auto &CandidateSet : FusionCandidates) {
           // Report fusion to the Optimization Remarks.
           // Note this needs to be done *before* performFusion because
           // performFusion will change the original loops, making it not
           // possible to identify them after fusion is complete.
           reportLoopFusion<OptimizationRemark>((Peel ? FC0Copy : *FC0), *FC1,
                                                FuseCounter);
 
           FusionCandidate FusedCand(
-              performFusion((Peel ? FC0Copy : *FC0), *FC1), &DT, &PDT, ORE,
+              performFusion((Peel ? FC0Copy : *FC0), *FC1), &DT, &PDT, &LI, ORE,
               FC0Copy.PP);
           FusedCand.verify();
           assert(FusedCand.isEligibleForFusion(SE) &&
                  "Fused candidate should be eligible for fusion!");
 
           // Notify the loop-depth-tree that these loops are not valid objects
           LDT.removeLoop(FC1->L);
 

llvm/lib/Transforms/Scalar/LoopUnrollPass.cpp

if (auto UnrollFactor = shouldFullUnroll(L, TTI, DT, SE, EphValues,
                                              MaxTripCount, UCE, UP)) {
       UP.Count = *UnrollFactor;
       UseUpperBound = true;
       return ExplicitUnroll;
     }
   }
 
   // 5th priority is loop peeling.
-  computePeelCount(L, LoopSize, PP, TripCount, DT, SE, UP.Threshold);
+  computePeelCount(L, LoopSize, PP, TripCount, DT, *LI, SE, UP.Threshold);
   if (PP.PeelCount) {
     UP.Runtime = false;
     UP.Count = 1;
     return ExplicitUnroll;
   }
 
   // Before starting partial unrolling, set up.partial to true,
   // if user explicitly asked  for unrolling

llvm/lib/Transforms/Utils/LoopPeel.cpp

 
 static cl::opt<unsigned> UnrollForcePeelCount(
     "unroll-force-peel-count", cl::init(0), cl::Hidden,
     cl::desc("Force a peel count regardless of profiling information."));
 
 static const char *PeeledCountMetaData = "llvm.loop.peeled.count";
 
 // Check whether we are capable of peeling this loop.
-bool llvm::canPeel(Loop *L) {
+bool llvm::canPeel(Loop *L, DominatorTree &DT, LoopInfo &LI) {
   // Make sure the loop is in simplified form
   if (!L->isLoopSimplifyForm())
     return false;
 
+  // We require the loop to be in LCSSA form when we peel and map values from
+  // peeled region to their uses outside the loop.
+  if (!L->isRecursivelyLCSSAForm(DT, LI))
+    return false;
+
   // Don't try to peel loops where the latch is not the exiting block.
   // This can be an indication of two different things:
   // 1) The loop is not rotated.
   // 2) The loop contains irreducible control flow that involves the latch.
   const BasicBlock *Latch = L->getLoopLatch();
   if (!L->isLoopExiting(Latch))
     return false;
 
return any_of(ExitBlocks, [](const BasicBlock *EB) {
       return !EB->getTerminatingDeoptimizeCall();
     });
 }
 
 
 // Return the number of iterations we want to peel off.
 void llvm::computePeelCount(Loop *L, unsigned LoopSize,
                             TargetTransformInfo::PeelingPreferences &PP,
-                            unsigned TripCount, DominatorTree &DT,
+                            unsigned TripCount, DominatorTree &DT, LoopInfo &LI,
                             ScalarEvolution &SE, unsigned Threshold) {
   assert(LoopSize > 0 && "Zero loop size is not allowed!");
   // Save the PP.PeelCount value set by the target in
   // TTI.getPeelingPreferences or by the flag -unroll-peel-count.
   unsigned TargetPeelCount = PP.PeelCount;
   PP.PeelCount = 0;
-  if (!canPeel(L))
+  if (!canPeel(L, DT, LI))
     return;
 
   // Only try to peel innermost loops by default.
   // The constraint can be relaxed by the target in TTI.getPeelingPreferences
   // or by the flag -unroll-allow-loop-nests-peeling.
   if (!PP.AllowLoopNestsPeeling && !L->isInnermost())
     return;
 

 /// For loops that dynamically execute \p PeelCount iterations or less
 /// this provides a benefit, since the peeled off iterations, which account
 /// for the bulk of dynamic execution, can be further simplified by scalar
 /// optimizations.
 bool llvm::peelLoop(Loop *L, unsigned PeelCount, LoopInfo *LI,
                     ScalarEvolution *SE, DominatorTree *DT, AssumptionCache *AC,
                     bool PreserveLCSSA) {
   assert(PeelCount > 0 && "Attempt to peel out zero iterations?");
-  assert(canPeel(L) && "Attempt to peel a loop which is not peelable?");
+  assert(canPeel(L, *DT, *LI) && "Attempt to peel a loop which is not peelable?");
 
   LoopBlocksDFS LoopBlocks(L);
   LoopBlocks.perform(LI);
 
   BasicBlock *Header = L->getHeader();
   BasicBlock *PreHeader = L->getLoopPreheader();
   BasicBlock *Latch = L->getLoopLatch();
   SmallVector<std::pair<BasicBlock *, BasicBlock *>, 4> ExitEdges;

llvm/unittests/Transforms/Utils/CMakeLists.txt

add_llvm_unittest(UtilsTests
   CallPromotionUtilsTest.cpp
   CloningTest.cpp
   CodeExtractorTest.cpp
   CodeMoverUtilsTest.cpp
   DebugifyTest.cpp
   FunctionComparatorTest.cpp
   IntegerDivisionTest.cpp
   LocalTest.cpp
+  LoopPeelTest.cpp
   LoopRotationUtilsTest.cpp
   LoopUtilsTest.cpp
   ModuleUtilsTest.cpp
   ScalarEvolutionExpanderTest.cpp
   SizeOptsTest.cpp
   SSAUpdaterBulkTest.cpp
   UnrollLoopTest.cpp
   ValueMapperTest.cpp
   VFABIUtils.cpp
   )

llvm/unittests/Transforms/Utils/LoopPeelTest.cpp

+//===- LoopPeelUtilsTest.cpp - Unit tests for LoopRotation utility ----===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Transforms/Utils/LoopPeel.h"
+#include "llvm/Analysis/AssumptionCache.h"
+#include "llvm/Analysis/InstructionSimplify.h"
+#include "llvm/Analysis/LoopInfo.h"
+#include "llvm/Analysis/ScalarEvolution.h"
+#include "llvm/Analysis/TargetLibraryInfo.h"
+#include "llvm/Analysis/TargetTransformInfo.h"
+#include "llvm/AsmParser/Parser.h"
+#include "llvm/IR/Dominators.h"
+#include "llvm/IR/LLVMContext.h"
+#include "llvm/Support/SourceMgr.h"
+#include "llvm/Transforms/Utils/LoopUtils.h"
+#include "gtest/gtest.h"
+
+using namespace llvm;
+
+static std::unique_ptr<Module> parseIR(LLVMContext &C, const char *IR) {
+  SMDiagnostic Err;
+  std::unique_ptr<Module> Mod = parseAssemblyString(IR, Err, C);
+  if (!Mod)
+    Err.print("LoopPeelUtilsTest", errs());
+  return Mod;
+}
+
+/// This test contains loop that is not in LCSSA form. We cannot peel such
+/// loops, unless they are in LCSSA form.
+TEST(LoopPeel, NonLCSSAForm) {
+  LLVMContext C;
+
+  std::unique_ptr<Module> M = parseIR(C,
+                                      R"(
+declare i32 @llvm.experimental.deoptimize.i32(...)
+
+define i32 @test(i32 * nonnull %a, i64 %x) {
+entry:
+  br label %for.cond1
+
+for.cond1:
+  %idx = phi i64 [ 0, %entry ], [ %idx.next, %for.tail ]
+  %sum = phi i32 [ 0, %entry ], [ %sum.next, %for.tail ]
+  %a.idx = getelementptr inbounds i32, i32 *%a, i64 %idx
+  %val.a.idx = load i32, i32* %a.idx, align 4
+  %zero.check = icmp eq i32 %val.a.idx, 0
+  br i1 %zero.check, label %deopt.exit, label %for.tail
+
+for.tail:
+  %sum.next = add i32 %sum, %val.a.idx
+  %idx.next = add nuw nsw i64 %idx, 1
+  %for.check = icmp ult i64 %idx, %x
+  br i1 %for.check, label %for.cond1, label %return
+
+return:
+  ret i32 %sum
+
+deopt.exit:
+  %deopt.val = call i32(...) @llvm.experimental.deoptimize.i32() [ "deopt"(i32 %val.a.idx) ]
+  ret i32 %deopt.val
+})");
+
+  auto *F = M->getFunction("test");
+  DominatorTree DT(*F);
+  LoopInfo LI(DT);
+  AssumptionCache AC(*F);
+  TargetLibraryInfoImpl TLII;
+  TargetLibraryInfo TLI(TLII);
+  ScalarEvolution SE(*F, TLI, AC, DT, LI);
+
+  Loop *L = *LI.begin();
+
+  EXPECT_FALSE(L->isLCSSAForm(DT));
+  bool ret = canPeel(L, DT, LI);
+  // We cannot peel loops that are not in LCSSA form.
+  EXPECT_FALSE(ret);
+  formLCSSARecursively(*L, DT, &LI, &SE);
+  bool ret2 = canPeel(L, DT, LI);
+  EXPECT_TRUE(ret2);
+  bool result =
+      peelLoop(L, 1 /*PeelCount*/, &LI, &SE, &DT, &AC, /*preserveLCSSA*/ true);
+  EXPECT_TRUE(result);
+}