[LoopUnroll] fix cloning callbr

[nickdesaulniers]

Summary

There is currently a correctness issue when unrolling loops containing
callbr's where their indirect targets are being updated correctly to the
newly created labels, but their operands are not. This manifests in
unrolled loops where the second and subsequent copies of callbr
instructions have blockaddresses of the label from the first instance of
the unrolled loop, which would result in nonsensical runtime control
flow.

When cloning a callbr, update its blockaddress operands if they were cloned, too.

Link: https://bugs.llvm.org/show_bug.cgi?id=42489
Link: https://groups.google.com/forum/#!topic/clang-built-linux/z-hRWP9KqPI

[nickdesaulniers]

Here's how I might modify the above test if for when we can loop unroll correctly: https://gist.github.com/nickdesaulniers/7216f6e5a17c7064285190440cb88f1d

[nickdesaulniers]

I have a more complete fix for this whole issue outright. Let me get the test case working, then I'll upload a v2.

[nickdesaulniers]

• fix the loop unroller

[nickdesaulniers]

llvm/lib/Transforms/Utils/LoopUnroll.cpp

103

Notes to reviewers: let me know if the comments are excessive and I'll remove them.

Also note that I manually peeled this loop out from the above in this function. I could place the isa<CallBrInst> and rest of this block in there. Just let me know WDYT?

[srhines]

llvm/lib/Transforms/Utils/LoopUnroll.cpp

105

Fold this line into the if statement for LLVM style.

[nickdesaulniers]

• style nit

[nickdesaulniers]

llvm/lib/Transforms/Utils/LoopUnroll.cpp

111

Sorry, I think this should be a dyn_cast (we're "downcasting" from a Value* to a BasicBlock*, which may fail)

[nickdesaulniers]

• dyn_cast

[nickdesaulniers]

• reroll loops, prefer conciseness

[fhahn]

llvm/lib/Transforms/Utils/LoopUnroll.cpp

100

IIUC, for block addresses, we now set the operand first by the code above and then also by this block here. It might be slightly better to first check for the special case and continue in case we hit it.

Also, a brief comment why we need special handling here would be helpful IMO.

103

Is it possible for It->second to not be a BB here? I think if the original operand was a BB, the new one also should be a BB. So this could be a regular cast?

llvm/test/Transforms/LoopUnroll/callbr.ll

2

The test only checks the emitted IR and no debug output. So the 2>&1 should not be needed, right?

24

nit: could just be a void function.

[nickdesaulniers]

• check early, continue, fix test nits

llvm/lib/Transforms/Utils/LoopUnroll.cpp

103

Is it possible for It->second to not be a BB here?

I'm pretty sure It->second can ONLY be a BasicBlock here, since the values stored in VMap were created by llvm::CloneBasicBlock (see VMap there).

That said, I'm not super confident in my understanding of dyn_cast vs static_cast.

The VMap is an instance of ValueToValueMapTy, which simply maps Value* to Value*. As Value is the base class of a lot of stuff in LLVM, we end up storing all kinds of pointers to derived classes in VMap.

My understanding of when to use dyn_cast (or even dynamic_cast) vs static_cast is that it's always ok to static_cast from derived class to base class, but not vice versa. That's why dyn_cast or dynamic_cast might result in nullptr at runtime, and thus need to be checked.

Since It->second should ALWAYS (IIUC) be a BasicBlock* in this case, I would have thought we still need a checked call to dyn_cast.

Maybe I'm misunderstanding and that it's ok to static_cast from base to derived if you're certain that base is always derived?

[nickdesaulniers]

llvm/lib/Transforms/Utils/LoopUnroll.cpp

103

Ah, a quick reread of http://llvm.org/docs/ProgrammersManual.html#the-isa-cast-and-dyn-cast-templates makes it seem like cast<BasicBlock> should match with what I expect. Brb, enabling assertions in my build...

[fhahn]

llvm/lib/Transforms/Utils/LoopUnroll.cpp

103

Yep, cast<BasicBlock> will assert that it actually is a basic block with assertions enabled. That should be sufficient cover.

[nickdesaulniers]

• prefer cast<> to dyn_cast<>

[fhahn]

LGTM, thanks!

llvm/lib/Transforms/Utils/LoopUnroll.cpp

608

nit: unrelated change?

[nickdesaulniers]

Great! Thanks for the code review.

[nickdesaulniers]

llvm/lib/Transforms/Utils/LoopUnroll.cpp

608

Intentional; helps readability to declare as late as possible.

[efriedma]

Is this transform safe? The inline asm could stash the address of a destination in a variable in one loop iteration, and use it in a later loop iteration. Or is that not legal?

[nickdesaulniers]

In D64101#1569218, @efriedma wrote:

Is this transform safe? The inline asm could stash the address of a destination in a variable in one loop iteration, and use it in a later loop iteration. Or is that not legal?

Isn't that an issue in inline assembly regardless of whether or not it's asm goto?
GCC seems to unroll the loop, even when capturing the induction variable:
https://godbolt.org/z/HVUTC4
Should the presence of inline assembly be such an optimization barrier?

[efriedma]

If there's some rule that distinguishes blockaddresses used in callbr from general blockaddresses, we should state that explicitly somewhere in LangRef.

[fhahn]

Maybe we could make the semantics of callbr a bit more explicit in LangRef?

Currently it only mentions that the only use is to implement asm goto's like in GCC. IIUC, asm goto cannot have outputs and the inputs are explicit, those can be updated while unrolling. According to https://gcc.gnu.org/onlinedocs/gcc/Extended-Asm.html 6.47.2.7 Goto Labels, if the assembly in the asm goto modifies anything, it needs a "memory" clobber. I think this should be handled like regular inline assembly already in LLVM?

[nickdesaulniers]

it needs a "memory" clobber. I think this should be handled like regular inline assembly already in LLVM?

Neither GCC nor Clang prevent loop unrolling loops containing inline asm w/ "memory" clobber. You can observe this with my godbolt link above and manually add "memory" to the 4 position in the asm statement.

If there's some rule that distinguishes blockaddresses used in callbr from general blockaddresses, we should state that explicitly somewhere in LangRef.

I don't really understand what you're looking for @eli.friedman . Should I add a statment along the lines of "blockaddresses may be rewritten during optimization passes to refer to the address of newly created blocks"?

[efriedma]

The question is whether something like the following is legal:

void bar(void) {
    long temp = 0;
    #pragma GCC unroll 3
    for (int j = 0; j < 3; ++j) {
        asm goto("lea %l2(%%rip), %%rax\n"
                 "cmp $0, %0\n"
                 "jne 1f\n"
                 "mov %%rax, %0\n"
                 "1:\n"
                 "jmp *%0\n" :: "m"(temp), "r"(j) :"memory","rax": baz);
        baz:;
    }
}
int main() { bar(); }

I guess the answer is that no, it isn't legal, given gcc's behavior. But there's nothing in LangRef that forbids it... we need some rule that ties the value passed to the "X" constraint to the actual successor.

[efriedma]

Or actually, it might make more sense to change the way we generate/lower callbr, to make the label parameters implicit; instead of modeling them with blockaddress, we never write them in IR at all, and automatically generate them later based on the successor list.

[hfinkel]

In D64101#1569642, @efriedma wrote:

Or actually, it might make more sense to change the way we generate/lower callbr, to make the label parameters implicit; instead of modeling them with blockaddress, we never write them in IR at all, and automatically generate them later based on the successor list.

That's an appealing idea.

We do need to be careful, AFAIK, because the goto targets can be passed into the asm as parameters (where the block addresses come from the labels-as-values extension). However, in theory we already shouldn't unroll loops with blocks with their address taken, so I imagine that case is irrelevant here (*).

(*) LoopUnroll.cpp actually only seems to check Header->hasAddressTaken(), not all of the blocks, so maybe that's wrong?

[nickdesaulniers]

In D64101#1569682, @hfinkel wrote:

In D64101#1569642, @efriedma wrote:

Or actually, it might make more sense to change the way we generate/lower callbr, to make the label parameters implicit; instead of modeling them with blockaddress, we never write them in IR at all, and automatically generate them later based on the successor list.

That's an appealing idea.

I kind of agree; looking at the arguments to callbr, the label list (the final part in []) feels like duplicate information of the earlier blockaddress operands. The blockaddress operands in the common case SHOULD match or be the address of blocks in the label list, except when you're explicitly passing the address of labels around (see below link). In that first case, why specify explicitly the blockaddress at all? Surely other transforms may update the label list but not the blockaddresses, just like this patch w/ LoopUnroll. I feel like the blockaddress could be implicit based on the label list, and generated as late as possible and only when needed.

That said, I'd really like to land this patch as it addresses an observable and bad bug in the Linux kernel, and I'd like to do so to make the clang-9 release train, rather than try to rearchitect callbr here.

We do need to be careful, AFAIK, because the goto targets can be passed into the asm as parameters (where the block addresses come from the labels-as-values extension). However, in theory we already shouldn't unroll loops with blocks with their address taken, so I imagine that case is irrelevant here (*).

(*) LoopUnroll.cpp actually only seems to check Header->hasAddressTaken(), not all of the blocks, so maybe that's wrong?

Great point. Looking at GCC:
https://godbolt.org/z/lKa_HD

Notice: it does unroll/duplicate the label within the loop when the label is passed in via the final label list in the asm goto. It does not duplicate the address of label passed in, even when the induction variable is captured.

So my patch as it stands in the current revision is incorrect or will differ from GCC; I need to differentiate between is this BlockAddress operand explicitly passed in as an input, or is it from the label list. I don't think it's hard for me to change what I have to at least match GCC. To @eli.friedman 's point; then yes blockaddress is indeed handled differently in this case. Let me fix up the patch and add a test based on the godbolt link.

And thanks to everyone so far for the feedback and discussion.

[nickdesaulniers]

Here's another interesting case for me to add a test for: https://godbolt.org/z/WwKXaI

The same address being using in the label list AND explicitly passed in as an address of label. Note the difference in behavior there.

[nickdesaulniers]

In D64101#1574332, @nickdesaulniers wrote:

Here's another interesting case for me to add a test for: https://godbolt.org/z/WwKXaI

The same address being using in the label list AND explicitly passed in as an address of label. Note the difference in behavior there.

Sorry, that test case would rely on https://reviews.llvm.org/D64167 first. ;)

[fhahn]

Marking as requiring changes while the comments are being addressed.

That said, I'd really like to land this patch as it addresses an observable and bad bug in the Linux kernel, and I'd like to do so to make the clang-9 release train, rather than try to rearchitect callbr here.

IIUC the patch as is fixes an issue with unrolling asm goto, in cases unrolling is legal. But we need additional checks to prevent unrolling in the cases Eli mentioned. Maybe a safer approach would be to start with a patch to restrict the unrolling of loops with asm goto (that should also fix the bug in the linux kernel, right?) and then allow unrolling for the safe subset of cases where it is legal.

Sorry, that test case would rely on https://reviews.llvm.org/D64167 first. ;)

It's fine to have this patch depend on D64167

[nickdesaulniers]

In D64101#1574429, @fhahn wrote:

Marking as requiring changes while the comments are being addressed.

That said, I'd really like to land this patch as it addresses an observable and bad bug in the Linux kernel, and I'd like to do so to make the clang-9 release train, rather than try to rearchitect callbr here.

IIUC the patch as is fixes an issue with unrolling asm goto, in cases unrolling is legal. But we need additional checks to prevent unrolling in the cases Eli mentioned. Maybe a safer approach would be to start with a patch to restrict the unrolling of loops with asm goto (that should also fix the bug in the linux kernel, right?) and then allow unrolling for the safe subset of cases where it is legal.

Sure thing, forked off v1 of this patch into: https://reviews.llvm.org/D64368

[nickdesaulniers]

Note to self: this needs to be rebased on top of https://reviews.llvm.org/rL366130 and a test case for LoopUnswitch needs to be added.

[nickdesaulniers]

Superseded by: https://reviews.llvm.org/D129993

llvm/lib/Transforms/Utils/LoopUnroll.cpp

 #else
     cl::init(false)
 #endif
                     );
 
 /// Convert the instruction operands from referencing the current values into
 /// those specified by VMap.
 void llvm::remapInstruction(Instruction *I, ValueToValueMapTy &VMap) {
-  for (unsigned op = 0, E = I->getNumOperands(); op != E; ++op) {
-    Value *Op = I->getOperand(op);
+  for (unsigned OpNo = 0, E = I->getNumOperands(); OpNo != E; ++OpNo) {
+    Value *Op = I->getOperand(OpNo);
+
+    // If we have a BlockAddress operand where the BasicBlock of the
+    // BlockAddress was remapped and the BlockAddress is the operand of a
+    // CallBrInst, then we need to update the CallBrInst's operand to be a
+    // BlockAddress of the remapped BasicBlock, not the original BasicBlock.
+    if (auto *BA = dyn_cast<BlockAddress>(Op)) {
+      ValueToValueMapTy::iterator It = VMap.find(BA->getBasicBlock());
+      if (It != VMap.end() && isa<CallBrInst>(*I)) {
+        I->setOperand(OpNo, BlockAddress::get(cast<BasicBlock>(It->second)));
+        continue;
+      }
+    }
 
     // Unwrap arguments of dbg.value intrinsics.
     bool Wrapped = false;
     if (auto *V = dyn_cast<MetadataAsValue>(Op))
       if (auto *Unwrapped = dyn_cast<ValueAsMetadata>(V->getMetadata())) {
         Op = Unwrapped->getValue();
         Wrapped = true;
       }
 
     auto wrap = [&](Value *V) {
       auto &C = I->getContext();
       return Wrapped ? MetadataAsValue::get(C, ValueAsMetadata::get(V)) : V;
     };
 
     ValueToValueMapTy::iterator It = VMap.find(Op);
     if (It != VMap.end())
-      I->setOperand(op, wrap(It->second));
+      I->setOperand(OpNo, wrap(It->second));
   }
 
   if (PHINode *PN = dyn_cast<PHINode>(I)) {
     for (unsigned i = 0, e = PN->getNumIncomingValues(); i != e; ++i) {
       ValueToValueMapTy::iterator It = VMap.find(PN->getIncomingBlock(i));
       if (It != VMap.end())
         PN->setIncomingBlock(i, cast<BasicBlock>(It->second));
     }
   }
 }
 
 /// Check if unrolling created a situation where we need to insert phi nodes to
 /// preserve LCSSA form.
 /// \param Blocks is a vector of basic blocks representing unrolled loop.
 /// \param L is the outer loop.
 /// It's possible that some of the blocks are in L, and some are not. In this
 /// case, if there is a use is outside L, and definition is inside L, we need to
 /// insert a phi-node, otherwise LCSSA will be broken.
 /// The function is just a helper function for llvm::UnrollLoop that returns
 /// true if this situation occurs, indicating that LCSSA needs to be fixed.
 static bool needToInsertPhisForLCSSA(Loop *L, std::vector<BasicBlock *> Blocks,
LoopUnrollResult llvm::UnrollLoop(Loop *L, UnrollLoopOptions ULO, LoopInfo *LI,
   } else {
     NumUnrolledWithHeader++;
     ContinueOnTrue = L->contains(HeaderBI->getSuccessor(0));
     LoopExit = HeaderBI->getSuccessor(ContinueOnTrue);
   }
 
   // For the first iteration of the loop, we should use the precloned values for
   // PHI nodes.  Insert associations now.
-  ValueToValueMapTy LastValueMap;
   std::vector<PHINode*> OrigPHINode;
   for (BasicBlock::iterator I = Header->begin(); isa<PHINode>(I); ++I) {
     OrigPHINode.push_back(cast<PHINode>(I));
   }
 
   std::vector<BasicBlock *> Headers;
   std::vector<BasicBlock *> HeaderSucc;
   std::vector<BasicBlock *> Latches;
for (BasicBlock *BB : L->getBlocks())
             if (NewDIL)
               I.setDebugLoc(NewDIL.getValue());
             else
               LLVM_DEBUG(dbgs()
                          << "Failed to create new discriminator: "
                          << DIL->getFilename() << " Line: " << DIL->getLine());
           }
 
+  ValueToValueMapTy LastValueMap;
   for (unsigned It = 1; It != ULO.Count; ++It) {
     std::vector<BasicBlock*> NewBlocks;
     SmallDenseMap<const Loop *, Loop *, 4> NewLoops;
     NewLoops[L] = L;
 
     for (LoopBlocksDFS::RPOIterator BB = BlockBegin; BB != BlockEnd; ++BB) {
       ValueToValueMapTy VMap;
       BasicBlock *New = CloneBasicBlock(*BB, VMap, "." + Twine(It));

llvm/test/Transforms/LoopUnroll/callbr.ll

+; RUN: opt -loop-unroll -S -o - %s | FileCheck %s
+
+; CHECK-LABEL: if.then:
+; CHECK-NEXT: callbr void asm sideeffect "1: nop\0A\09.quad b, ${0:l}, $$5\0A\09", "X,~{dirflag},~{fpsr},~{flags}"(i8* blockaddress(@d, %l_yes))
+; CHECK-NEXT: to label %asm.fallthrough [label %l_yes]
+; CHECK-LABEL: l_yes:
+
+; CHECK-LABEL: if.then.1:
+; CHECK-NEXT: callbr void asm sideeffect "1: nop\0A\09.quad b, ${0:l}, $$5\0A\09", "X,~{dirflag},~{fpsr},~{flags}"(i8* blockaddress(@d, %l_yes.1))
+; CHECK-NEXT: to label %asm.fallthrough.1 [label %l_yes.1]
+; CHECK-LABLE: l_yes.1:
+
+; CHECK-LABEL: if.then.2:
+; CHECK-NEXT: callbr void asm sideeffect "1: nop\0A\09.quad b, ${0:l}, $$5\0A\09", "X,~{dirflag},~{fpsr},~{flags}"(i8* blockaddress(@d, %l_yes.2))
+; CHECK-NEXT: to label %asm.fallthrough.2 [label %l_yes.2]
+; CHECK-LABEL: l_yes.2:
+
+define dso_local void @d() {
+entry:
+  br label %for.body
+
+for.cond.cleanup:                                 ; preds = %for.inc
+  ret void
+
+for.body:                                         ; preds = %for.inc, %entry
+  %e.04 = phi i32 [ 0, %entry ], [ %inc, %for.inc ]
+  %tobool = icmp eq i32 %e.04, 0
+  br i1 %tobool, label %for.inc, label %if.then
+
+if.then:                                          ; preds = %for.body
+  callbr void asm sideeffect "1: nop\0A\09.quad b, ${0:l}, $$5\0A\09", "X,~{dirflag},~{fpsr},~{flags}"(i8* blockaddress(@d, %l_yes))
+          to label %asm.fallthrough [label %l_yes]
+
+asm.fallthrough:                                  ; preds = %if.then
+  br label %l_yes
+
+l_yes:                                            ; preds = %asm.fallthrough, %if.then
+  %call = tail call i32 (...) @g()
+  br label %for.inc
+
+for.inc:                                          ; preds = %for.body, %l_yes
+  %inc = add nuw nsw i32 %e.04, 1
+  %exitcond = icmp eq i32 %inc, 3
+  br i1 %exitcond, label %for.cond.cleanup, label %for.body
+}
+
+declare dso_local i32 @g(...) local_unnamed_addr