Preserve loop metadata when splitting exit blocks

[clin1]

Summary

LoopUtils.cpp contains a utility that splits an loop exit block, so that the new block contains only edges coming from the loop. In the case of nested loops, the exit path for the inner loop might also be the back-edge of the outer loop. The new block which is inserted on this path, is now a latch for the outer loop, and it needs to hold the loop metadata for the outer loop. (The test case gives a more concrete view of the situation.)

[mkazantsev]

lib/Transforms/Utils/LoopUtils.cpp

83

{ } not needed here.

85

Dot in the end of sentence.

96

Is it possible that their loops don't match? If no, please turn it into an assert. If yes, please add a test on this situation.

test/Transforms/LoopSimplify/preserve-llvm-loop-metadata2.ll

35

Can the test be reduced to make it clear that it's only supposed to test metadata and not fp computations? I think it can be just reduced to having just blocks with counters.

Maybe also makes sense to rename blocks to smth like outer.loop, inner.loop.

[dmgreen]

lib/Transforms/Utils/LoopUtils.cpp

39

These should be in order. I think running clang-format on the patch will do that for you.

test/Transforms/LoopSimplify/preserve-llvm-loop-metadata2.ll

7

Is it worth-while to make sure the LOOPMD is the correct of the two metadata nodes (still "llvm.loop.unroll_and_jam.count")? And that the inner loops still has nounroll?

64

You can probably remove the tbaa info too, to make things simpler.

[clin1]

lib/Transforms/Utils/LoopUtils.cpp

39

Added by mistake; thanks.

96

It's taken care of by the assert below, but it needs a better explanation. Comments have all been rewritten.

[clin1]

Rewrote comments, fixed formatting, simplified test and added checks for specific metadata

[clin1]

Previous patch was null, this is the right one.

[mkazantsev]

LGTM

[clin1]

Thank you for the thorough review. I have one more request to the reviewers on this list: could someone give a hand to commit the code? I don't have write access to the repo yet.

[jdoerfert]

Thanks for the patch!

In case it comes ever up and is actually still relevant (I haven't checked), there are similar problems all over the place.
The one I tried to fix at some point but never actually did can be found here: https://reviews.llvm.org/D5344

[rnk]

I think this broke the windows selfhost:
http://lab.llvm.org:8011/builders/clang-x64-windows-msvc/builds/1457/steps/stage%202%20build/logs/stdio

Assertion failed: TI->getMetadata(LLVMContext::MD_loop) == OuterLoopMD && "exit edge to other loop doesn't contain expected metadata", file C:\b\slave\clang-x64-windows-msvc\build\llvm.src\lib\Transforms\Utils\LoopUtils.cpp, line 105
#0 0x00007ff647c07f05 HandleAbort c:\b\slave\clang-x64-windows-msvc\build\llvm.src\lib\support\windows\signals.inc:409:0
#1 0x00007ff99ee4dc17 (C:\windows\SYSTEM32\ucrtbase.DLL+0x6dc17)
#2 0x00007ff99ee4eaa1 (C:\windows\SYSTEM32\ucrtbase.DLL+0x6eaa1)
#3 0x00007ff99ee50751 (C:\windows\SYSTEM32\ucrtbase.DLL+0x70751)
#4 0x00007ff99ee50a5f (C:\windows\SYSTEM32\ucrtbase.DLL+0x70a5f)
#5 0x00007ff647cf0c9e <lambda_e0b0a017e173a90c4e1e4410b36fad3c>::operator() c:\b\slave\clang-x64-windows-msvc\build\llvm.src\lib\transforms\utils\looputils.cpp:106:0
#6 0x00007ff647cf329a llvm::formDedicatedExitBlocks(class llvm::Loop *,class llvm::DominatorTree *,class llvm::LoopInfo *,bool) c:\b\slave\clang-x64-windows-msvc\build\llvm.src\lib\transforms\utils\looputils.cpp:135:0
#7 0x00007ff647c749d8 simplifyOneLoop c:\b\slave\clang-x64-windows-msvc\build\llvm.src\lib\transforms\utils\loopsimplify.cpp:535:0
#8 0x00007ff647c73d4a llvm::simplifyLoop(class llvm::Loop *,class llvm::DominatorTree *,class llvm::LoopInfo *,class llvm::ScalarEvolution *,class llvm::AssumptionCache *,bool) c:\b\slave\clang-x64-windows-msvc\build\llvm.src\lib\transforms\utils\loopsimplify.cpp:706:0
#9 0x00007ff647c72b6a `anonymous namespace'::LoopSimplify::runOnFunction c:\b\slave\clang-x64-windows-msvc\build\llvm.src\lib\transforms\utils\loopsimplify.cpp:776:0
...

I'll get in a revert.

[craig.topper]

Reopening as this was reverted in r346823

[rnk]

It went green with my change here:
http://lab.llvm.org:8011/builders/clang-x64-windows-msvc/builds/1463

Can you gather the reproducer? I'm a bit busy for that right now.

[clin1]

Sorry for the trouble. I'll figure out what happened.

[clin1]

Remove duplicate blocks from the list of exit predecessors. This happens when the predecessor inst is a switch with duplicate targets. Avoids block double-processing and maintains the metadata consistency assertion.

[efriedma]

You're calling std::unique on an unsorted array?

[clin1]

Missed the sort call -- thanks.

[efriedma]

lib/Transforms/Utils/LoopUtils.cpp

72

You're calling std::sort on a list of pointers; the result is non-deterministic, which could affect the compiler output.

[clin1]

Keep the InLoopPredecessors vector in the original order.

[efriedma]

Maybe add a separate testcase, instead of modifying an existing testcase? Otherwise looks fine, although maybe wait for @mkazantsev to look again.

[clin1]

Did not include the reverted part of the code (the important part :) ) in the new patch. Added separate testcase for switch-backedge. Thanks for the fix suggestions.

[clin1]

@mkazantsev, would you mind taking one more look at this patch? One significant fix was made since your last review, to handle duplicate edges from SwitchInsts. The patch has now been rebased to latest (which hasn't changed it much). Thanks!

[efriedma]

Is this patch still relevant?

[Meinersbur]

I think yes, together with all the other places were loop metadata is not preserved, such as D66892.

[craig.topper]

I'm trying to make sense of the history here. It looks like I commited it in November last year. Then it got reverted. So I reopened it but maybe didn't move it to changes required and instead it stayed in accepted. Then @clin1 updated it, but it never got re-reviewed. Then it was rebased in January. And I guess sat in accepted state until today when phabricator lost its mind and closed it to the original commit? And probably erased the January version of the diff and went back to November 18?

[Meinersbur]

The patch was closed by Phabricator since it discovered a "new" commit from git for this patch while assuming that the committed version is the most recent one. Eli re-opened it suspecting that we still might want to fix this.

[clin1]

Restoring latest patch.

[craig.topper]

lib/Transforms/Utils/LoopUtils.cpp

67

Can we use 'if (PushBlocks.insert(PredBB).second)'

[fhahn]

Looks like this still hasn't landed and probably requires at least a rebase.

lib/Transforms/Utils/LoopUtils.cpp

 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/PatternMatch.h"
 #include "llvm/IR/ValueHandle.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/KnownBits.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 
 using namespace llvm;
 using namespace llvm::PatternMatch;
 
 #define DEBUG_TYPE "loop-utils"
 
 bool llvm::formDedicatedExitBlocks(Loop *L, DominatorTree *DT, LoopInfo *LI,
                                    bool PreserveLCSSA) {
   bool Changed = false;
auto RewriteExit = [&](BasicBlock *BB) {
     bool IsDedicatedExit = true;
     for (auto *PredBB : predecessors(BB))
       if (L->contains(PredBB)) {
         if (isa<IndirectBrInst>(PredBB->getTerminator()))
           // We cannot rewrite exiting edges from an indirectbr.
           return false;
 
         InLoopPredecessors.push_back(PredBB);
       } else {
         IsDedicatedExit = false;
       }
 
     assert(!InLoopPredecessors.empty() && "Must have *some* loop predecessor!");
 
     // Nothing to do if this is already a dedicated exit.
     if (IsDedicatedExit)
       return false;
 
+    // With nested loops, the inner loop might exit to the header of an
+    // enclosing loop, and the in-loop-predecessor is a latch for that
+    // enclosing loop. If we insert a block between the latch and the header,
+    // that block becomes the new latch. Any loop metadata from the old latch
+    // needs to be moved to the new one.
+    MDNode *OuterLoopMD = nullptr;
+
+    // If the exit block is a header of a different loop, get that loop's
+    // metadata before we split the block.
+    if (LI->isLoopHeader(BB))
+      OuterLoopMD = LI->getLoopFor(BB)->getLoopID();
+
     auto *NewExitBB = SplitBlockPredecessors(
         BB, InLoopPredecessors, ".loopexit", DT, LI, nullptr, PreserveLCSSA);
 
+    // If OuterLoopMD is non-null, we know that the exit block BB is a
+    // loop header for a different loop, with metadata on its back edges.
+    // If NewExitBB is a member of that loop, then NewExitBB is a latch,
+    // and the loop's metadata needs to be copied to NewExitBB.
+    if (NewExitBB && OuterLoopMD &&
+        LI->getLoopFor(NewExitBB) == LI->getLoopFor(BB)) {
+      // The preds of NewExitBB are all former latches of the outer loop.
+      // Remove their metadata.
+      for (auto *PredLoopBB : InLoopPredecessors) {
+        Instruction *TI = PredLoopBB->getTerminator();
+        // All the latches should have the same metadata (ensured by
+        // getLoopID()).
+        assert(TI->getMetadata(LLVMContext::MD_loop) == OuterLoopMD &&
+               "exit edge to other loop doesn't contain expected metadata");
+        TI->setMetadata(LLVMContext::MD_loop, nullptr);
+      }
+      NewExitBB->getTerminator()->setMetadata(LLVMContext::MD_loop,
+                                              OuterLoopMD);
+    }
+
     if (!NewExitBB)
       LLVM_DEBUG(
           dbgs() << "WARNING: Can't create a dedicated exit block for loop: "
                  << *L << "\n");
     else
       LLVM_DEBUG(dbgs() << "LoopSimplify: Creating dedicated exit block "
                         << NewExitBB->getName() << "\n");
     return true;

test/Transforms/LoopSimplify/preserve-llvm-loop-metadata2.ll

+; RUN: opt -S -loop-simplify < %s | FileCheck %s
+
+; Two-loop nest with llvm.loop metadata on each loop.
+; inner.header exits to outer.header. inner.header is a latch for the outer
+; loop, and contains the outer loop's metadata.
+; After loop-simplify, a new block "outer.header.loopexit" is created between
+; inner.header and outer.header. The metadata from inner.header must be moved
+; to the new block, as the new block becomes the outer loop latch.
+; The metadata on the inner loop's latch should be untouched.
+
+; CHECK: outer.header.loopexit:
+; CHECK-NEXT: llvm.loop [[UJAMTAG:.*]]
+; CHECK-NOT: br i1 {{.*}}, label {{.*}}, label %outer.header.loopexit, !llvm.loop
+; CHECK: br label %inner.header, !llvm.loop [[UNROLLTAG:.*]]
+
+; CHECK: distinct !{[[UJAMTAG]], [[UJAM:.*]]}
+; CHECK: [[UJAM]] = !{!"llvm.loop.unroll_and_jam.count", i32 17}
+; CHECK: distinct !{[[UNROLLTAG]], [[UNROLL:.*]]}
+; CHECK: [[UNROLL]] = !{!"llvm.loop.unroll.count", i32 1}
+
+
+define dso_local void @loopnest() local_unnamed_addr #0 {
+entry:
+  br label %outer.header
+
+outer.header:                                         ; preds = %inner.header, %entry
+  %ii.0 = phi i64 [ 2, %entry ], [ %add, %inner.header ]
+  %cmp = icmp ult i64 %ii.0, 64
+  br i1 %cmp, label %inner.header, label %outer.header.cleanup
+
+outer.header.cleanup:                                 ; preds = %outer.header
+  ret void
+
+inner.header:                                        ; preds = %outer.header, %inner.body
+  %j.0 = phi i64 [ %add10, %inner.body ], [ %ii.0, %outer.header ]
+  %add = add nuw nsw i64 %ii.0, 16
+  %cmp2 = icmp ult i64 %j.0, %add
+  br i1 %cmp2, label %inner.body, label %outer.header, !llvm.loop !2
+
+inner.body:                                        ; preds = %inner.header
+  %add10 = add nuw nsw i64 %j.0, 1
+  br label %inner.header, !llvm.loop !4
+}
+
+!2 = distinct !{!2, !3}
+!3 = !{!"llvm.loop.unroll_and_jam.count", i32 17}
+!4 = distinct !{!4, !5}
+!5 = !{!"llvm.loop.unroll.count", i32 1}