Add check to Latch's terminator in UnrollRuntimeLoopRemainder

[rcorcs]

Summary

In this patch, I'm adding an extra check to the Latch's terminator in llvm::UnrollRuntimeLoopRemainder,
similar to how it is already done in the llvm::UnrollLoop.

The compiler would crash if this function is called with a malformed loop.

[dmgreen]

Hello. Please upload with full context to make patched easier to read. See https://llvm.org/docs/Contributing.html

lib/Transforms/Utils/LoopUnrollRuntime.cpp

559 ↗

(On Diff #163314)

For example, does the rest of this comment cover you concerns here?

[rcorcs]

Sorry, I'll make sure to upload with a full context in the next time.

The comment partially covers my concern.

The comment and the assertion that you've pointed out only covers the case where the latch block has a conditional branch that does not exit the loop.
However, if the latch block has an unconditional branch, then it will break in the command LatchBR->getSuccessor(ExitIndex) with an out of range assertion.

Since the latch has an unconditional branch, it means that:

• LatchBR->getSuccessor(0) == Header, and then ExitIndex will be 1.
• but LatchBR has no successor at position 1.

The way I see it, the two problems are related.

For this reason, I think that it would be better if we catch both cases of malformed loops.

But since the comment that you mention says that "This needs to be guaranteed by the callers of UnrollRuntimeLoopRemainder",
I guess it would be better to handle it with an assertion instead of return false (as I suggested in this patch).

If you agree with my point, I can update the patch with the assertion (instead of return false).

Thanks for the feedback.

[dmgreen]

We do seem to return false in every other case, I'm not very sure why this would be different and considered a precondition we would assert on.

Perhaps it's best to update the assert to a return false too.

Mind changing it that way instead?

[rcorcs]

I completely agree with your assessment.

I've done the suggested change and tested. All the LLVM tests pass as expected.

[anna]

Please upload patch with complete context. Could you also add a test case show casing your problem (which you've described earlier): Just a loop with unconditional latch terminator and run through runtime-unroll.

I'd added this as assert because upstream passes guarantees that by the time we reach unroll remainder, we will have a latch with conditional terminator.

[rcorcs]

I think that this patch is mainly important because these functions are externally visible in the library's API.
In fact, I experienced this problem when using the function llvm::UnrollRuntimeLoopRemainder in one of my own passes and then running it on SPEC2006.
This patch avoids unnecessarily making an implicit contract with the API's user.

I've added a test case. Please make sure that it follows the testing standards.
However, as mentioned by Anna, LLVM's loop unroll would not hit this problem as it guarantees externally that the loop latch has an exit branch.

[dmgreen]

I'm happy with this as is, but if you want the API to be tested and ensure it keeps working, you could add a unittest. Something like unittests/Transforms/Utils/BasicBlockUtilsTest.cpp or the ones is CloningTest.cpp.

test/Transforms/LoopUnroll/runtime-loop-non-exiting-latch.ll

1 ↗

(On Diff #163966)

I see that this is testing against a crash, but it's best to add a FileCheck, with some CHECK lines that the loop was not unrolled. I think the classic way to do this would be that it has one store.

[dmgreen]

Because I wanted to make sure it worked, I wrote this. Feel free to clean it up and use it as a unittest if you like:

//===- BasicBlockUtils.cpp - Unit tests for BasicBlockUtils ---------------===//
//
//                     The LLVM Compiler Infrastructure
//
// This file is distributed under the University of Illinois Open Source
// License. See LICENSE.TXT for details.
//
//===----------------------------------------------------------------------===//

#include "llvm/Transforms/Utils/UnrollLoop.h"
#include "llvm/Analysis/AssumptionCache.h"
#include "llvm/Analysis/LoopInfo.h"
#include "llvm/Analysis/ScalarEvolution.h"
#include "llvm/Analysis/TargetLibraryInfo.h"
#include "llvm/AsmParser/Parser.h"
#include "llvm/IR/BasicBlock.h"
#include "llvm/IR/Dominators.h"
#include "llvm/IR/LLVMContext.h"
#include "llvm/Support/SourceMgr.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("BasicBlockUtilsTests", errs());
  return Mod;
}

TEST(LoopUnrollRuntime, Latch) {
  LLVMContext C;

  std::unique_ptr<Module> M = parseIR(
    C,
    R"(define i32 @test(i32* %a, i32* %b, i32* %c, i64 %n) {
entry:
  br label %while.cond

while.cond:                                       ; preds = %while.body, %entry
  %i.0 = phi i64 [ 0, %entry ], [ %inc, %while.body ]
  %cmp = icmp slt i64 %i.0, %n
  br i1 %cmp, label %while.body, label %while.end

while.body:                                       ; preds = %while.cond
  %arrayidx = getelementptr inbounds i32, i32* %b, i64 %i.0
  %0 = load i32, i32* %arrayidx
  %arrayidx1 = getelementptr inbounds i32, i32* %c, i64 %i.0
  %1 = load i32, i32* %arrayidx1
  %mul = mul nsw i32 %0, %1
  %arrayidx2 = getelementptr inbounds i32, i32* %a, i64 %i.0
  store i32 %mul, i32* %arrayidx2
  %inc = add nsw i64 %i.0, 1
  br label %while.cond

while.end:                                        ; preds = %while.cond
  ret i32 0
})"
    );

  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);

  bool ret = UnrollRuntimeLoopRemainder(*LI.begin(), 4, true, true, false, &LI, &SE, &DT, &AC, true);
  EXPECT_FALSE(ret);
}

[rcorcs]

Thanks for the suggestion and the unit-test code. I'll add it to the patch.

About the check in the test file, I don't want to prevent the pass to unroll this loop.

test/Transforms/LoopUnroll/runtime-loop-non-exiting-latch.ll

1 ↗

(On Diff #163966)

I can add this check.
But it is not a problem if this loop is actually unrolled. For example, it could be runtime unrolled with a remainder.
The idea is more of having a test like the full-unroll-crashers.ll that only provides some challenging examples for the pass to deal with.

[rcorcs]

I've added the unit-test.

[anna]

test/Transforms/LoopUnroll/runtime-loop-non-exiting-latch.ll

1 ↗

(On Diff #163966)

Didn't we get an assertion failure in LLVM without the patch? Then you can add a ; REQUIRES assert at the top of the file so that it shows we're testing for an assertion failure. Some similar examples exist in other test files.

[rcorcs]

I've added the test requirement, as suggested.

[dmgreen]

LGTM. Thanks for the patch

[rcorcs]

Hi,

I'd like to confirm that someone can commit this patch for me because I believe I don't have commit access.

Thank you all for the review.

[dmgreen]

Sorry for the delay. Sure, I can do that.

Can you rebase this on trunk first? It looks like the unit test files were renamed.

[rcorcs]

I've rebased to the patch and changed the name of the unit-test UnrollLoop.cpp to UnrollLoopTest.cpp, similar to the other files in the same folder.

[dmgreen]

Thanks!

llvm/trunk/lib/Transforms/Utils/LoopUnrollRuntime.cpp

if (!L->isLoopSimplifyForm()) {
     return false;
   }
 
   // Guaranteed by LoopSimplifyForm.
   BasicBlock *Latch = L->getLoopLatch();
   BasicBlock *Header = L->getHeader();
 
   BranchInst *LatchBR = cast<BranchInst>(Latch->getTerminator());
+
+  if (!LatchBR || LatchBR->isUnconditional()) {
+    // The loop-rotate pass can be helpful to avoid this in many cases.
+    LLVM_DEBUG(
+        dbgs()
+        << "Loop latch not terminated by a conditional branch.\n");
+    return false;
+  }
+
   unsigned ExitIndex = LatchBR->getSuccessor(0) == Header ? 1 : 0;
   BasicBlock *LatchExit = LatchBR->getSuccessor(ExitIndex);
+
+  if (L->contains(LatchExit)) {
     // Cloning the loop basic blocks (`CloneLoopBlocks`) requires that one of the
-  // targets of the Latch be an exit block out of the loop. This needs
-  // to be guaranteed by the callers of UnrollRuntimeLoopRemainder.
-  assert(!L->contains(LatchExit) &&
-         "one of the loop latch successors should be the exit block!");
+    // targets of the Latch be an exit block out of the loop.
+    LLVM_DEBUG(
+        dbgs()
+        << "One of the loop latch successors must be the exit block.\n");
+    return false;
+  }
+
   // These are exit blocks other than the target of the latch exiting block.
   SmallVector<BasicBlock *, 4> OtherExits;
   bool isMultiExitUnrollingEnabled =
       canSafelyUnrollMultiExitLoop(L, OtherExits, LatchExit, PreserveLCSSA,
                                    UseEpilogRemainder) &&
       canProfitablyUnrollMultiExitLoop(L, OtherExits, LatchExit, PreserveLCSSA,
                                        UseEpilogRemainder);
   // Support only single exit and exiting block unless multi-exit loop unrolling is enabled.

llvm/trunk/test/Transforms/LoopUnroll/runtime-loop-non-exiting-latch.ll

+; REQUIRES: asserts
+; RUN: opt < %s -S -loop-unroll -unroll-runtime=true -unroll-allow-remainder=true -unroll-count=4
+
+; Make sure that the runtime unroll does not break with a non-exiting latch.
+define i32 @test(i32* %a, i32* %b, i32* %c, i64 %n) {
+entry:
+  br label %while.cond
+
+while.cond:                                       ; preds = %while.body, %entry
+  %i.0 = phi i64 [ 0, %entry ], [ %inc, %while.body ]
+  %cmp = icmp slt i64 %i.0, %n
+  br i1 %cmp, label %while.body, label %while.end
+
+while.body:                                       ; preds = %while.cond
+  %arrayidx = getelementptr inbounds i32, i32* %b, i64 %i.0
+  %0 = load i32, i32* %arrayidx
+  %arrayidx1 = getelementptr inbounds i32, i32* %c, i64 %i.0
+  %1 = load i32, i32* %arrayidx1
+  %mul = mul nsw i32 %0, %1
+  %arrayidx2 = getelementptr inbounds i32, i32* %a, i64 %i.0
+  store i32 %mul, i32* %arrayidx2
+  %inc = add nsw i64 %i.0, 1
+  br label %while.cond
+
+while.end:                                        ; preds = %while.cond
+  ret i32 0
+}

llvm/trunk/unittests/Transforms/Utils/CMakeLists.txt

add_llvm_unittest(UtilsTests
   ASanStackFrameLayoutTest.cpp
   BasicBlockUtilsTest.cpp
   CloningTest.cpp
   CodeExtractorTest.cpp
   FunctionComparatorTest.cpp
   IntegerDivisionTest.cpp
   LocalTest.cpp
   SSAUpdaterBulkTest.cpp
+  UnrollLoopTest.cpp
   ValueMapperTest.cpp
   )

llvm/trunk/unittests/Transforms/Utils/UnrollLoopTest.cpp

+//===- UnrollLoopTest.cpp - Unit tests for UnrollLoop ---------------------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Transforms/Utils/UnrollLoop.h"
+#include "llvm/Analysis/AssumptionCache.h"
+#include "llvm/Analysis/LoopInfo.h"
+#include "llvm/Analysis/ScalarEvolution.h"
+#include "llvm/Analysis/TargetLibraryInfo.h"
+#include "llvm/AsmParser/Parser.h"
+#include "llvm/IR/BasicBlock.h"
+#include "llvm/IR/Dominators.h"
+#include "llvm/IR/LLVMContext.h"
+#include "llvm/Support/SourceMgr.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("UnrollLoopTests", errs());
+  return Mod;
+}
+
+TEST(LoopUnrollRuntime, Latch) {
+  LLVMContext C;
+
+  std::unique_ptr<Module> M = parseIR(
+    C,
+    R"(define i32 @test(i32* %a, i32* %b, i32* %c, i64 %n) {
+entry:
+  br label %while.cond
+
+while.cond:                                       ; preds = %while.body, %entry
+  %i.0 = phi i64 [ 0, %entry ], [ %inc, %while.body ]
+  %cmp = icmp slt i64 %i.0, %n
+  br i1 %cmp, label %while.body, label %while.end
+
+while.body:                                       ; preds = %while.cond
+  %arrayidx = getelementptr inbounds i32, i32* %b, i64 %i.0
+  %0 = load i32, i32* %arrayidx
+  %arrayidx1 = getelementptr inbounds i32, i32* %c, i64 %i.0
+  %1 = load i32, i32* %arrayidx1
+  %mul = mul nsw i32 %0, %1
+  %arrayidx2 = getelementptr inbounds i32, i32* %a, i64 %i.0
+  store i32 %mul, i32* %arrayidx2
+  %inc = add nsw i64 %i.0, 1
+  br label %while.cond
+
+while.end:                                        ; preds = %while.cond
+  ret i32 0
+})"
+    );
+
+  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();
+
+  bool PreserveLCSSA = L->isRecursivelyLCSSAForm(DT,LI);
+
+  bool ret = UnrollRuntimeLoopRemainder(L, 4, true, false, false, &LI, &SE, &DT, &AC, PreserveLCSSA);
+  EXPECT_FALSE(ret);
+}