Prevent outlining of basicblock that uses BlockAddress

[serge-sans-paille]

Summary

CodeExtracting a block that references a blockadress breaks blockaddress usage by attaching them to other functions.

[davide]

testcase?

[serge-sans-paille]

@davide: while writing the test case, I came up with a better solution, I'll upload this as soon as it's finished.

[serge-sans-paille]

@davide Test case added, providing a fix instead of a guard.

[serge-sans-paille]

@davide Does the new change look correct to you?

[davide]

test/Transforms/CodeExtractor/BlockAddressreference.ll

1–2 ↗

(On Diff #101356)

This test has no check lines, so I guess it will pass also without your change [unless it crashed before] ? Is there a better way of testing?

[davidxl]

It is probably not legal to outline a BB with BB's address captured, so your previous version is better.

lib/Transforms/Utils/CodeExtractor.cpp

75 ↗

(On Diff #101229)

You can use BasicBlock::hasAddressTaken method

[serge-sans-paille]

@davidxl

It is probably not legal to outline a BB with BB's address captured,

I don't see anything in the lang ref that goes in that way. You could save the blockaddress in an array at some point in the program, and use it later on.

test/Transforms/CodeExtractor/BlockAddressreference.ll

1–2 ↗

(On Diff #101356)

@davide: yeah, it was crashing before, illegaly referencing blockaddress(@asterix, %for.cond) that no longer existed because of the outlining. This patch remaps it to @blockaddress(@outlined_function, %for.cond).

[davidxl]

This is assuming the referencing the saved addresses is only within the extracted function, but reference to block addresses from outside the parent function is illegal. How do you prevent the original function from branch into the middle of the outlined function?

[serge-sans-paille]

@davidxl

>  reference to block addresses from outside the parent function is illegal

If that's true, then my first patch was indeed better than this one. Can you point me to the relevant spec section?

[davidxl]

The spec probably needs update, but I remember inliner or function cloning does something to avoid this. +cc Chandler for comments.

[sanjoy]

In D33839#777673, @davidxl wrote:

The spec probably needs update, but I remember inliner or function cloning does something to avoid this. +cc Chandler for comments.

Are you looking for this (emphasis mine): "The ‘indirectbr‘ instruction implements an indirect branch to a label within the current function, whose address is specified by “address”. Address must be derived from a blockaddress constant."

[davidxl]

That looks like it - it is in the indirectbr clause, so it might be a good idea to state something in the section describing 'blockaddress' ?

[serge-sans-paille]

@davidxl , should be okay like this.

[serge-sans-paille]

@davidxl up!

[davidxl]

../lib/Transforms/Utils/CodeExtractor.cpp

75 ↗

(On Diff #102992)

Remove this comment.

80 ↗

(On Diff #102992)

Can you combine insert and count call, i.e, check the second member of the returned pair of 'insert'?

81 ↗

(On Diff #102992)

Remove braces.

82 ↗

(On Diff #102992)

return false?

84 ↗

(On Diff #102992)

You don't need to check the operands of instructions not in the BB set to be extracted. Guard this with

if (Curr->getParent() == &BB)

../test/Transforms/CodeExtractor/BlockAddressreference.ll

1 ↗

(On Diff #102992)

Need a test case covering use of block address in extracted region.

[serge-sans-paille]

@davidxl, all comments taken into account, thanks *a lot* for your review.

[davidxl]

lgtm

[efriedma]

lib/Transforms/Utils/CodeExtractor.cpp

79 ↗

(On Diff #104184)

I'm not sure I follow this check.

You can't extract a basic block whose terminator is an indirectbr; in general, you can't efficiently break the relevant CFG edges. Similarly, you can't extract a block with an indirectbr predecessor. And extracting a block whose address is taken is likely to lead to unexpected results in other cases for code which abuses block addresses for other uses (like the Linux kernel).

But I can't see why you would want to block extracting a basic block just because it refers to a blockaddress, or a global variable which contains a blockaddress.

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

 AggregateArgsOpt("aggregate-extracted-args", cl::Hidden,
                  cl::desc("Aggregate arguments to code-extracted functions"));
 
 /// \brief Test whether a block is valid for extraction.
 bool CodeExtractor::isBlockValidForExtraction(const BasicBlock &BB) {
   // Landing pads must be in the function where they were inserted for cleanup.
   if (BB.isEHPad())
     return false;
+  // taking the address of a basic block moved to another function is illegal
+  if (BB.hasAddressTaken())
+    return false;
+
+  // don't hoist code that uses another basicblock address, as it's likely to
+  // lead to unexpected behavior, like cross-function jumps
+  SmallPtrSet<User const *, 16> Visited;
+  SmallVector<User const *, 16> ToVisit;
+
+  for (Instruction const &Inst : BB)
+    ToVisit.push_back(&Inst);
+
+  while (!ToVisit.empty()) {
+    User const *Curr = ToVisit.pop_back_val();
+    if (!Visited.insert(Curr).second)
+      continue;
+    if (isa<BlockAddress const>(Curr))
+      return false; // even a reference to self is likely to be not compatible
+
+    if (isa<Instruction>(Curr) && cast<Instruction>(Curr)->getParent() != &BB)
+      continue;
+
+    for (auto const &U : Curr->operands()) {
+      if (auto *UU = dyn_cast<User>(U))
+        ToVisit.push_back(UU);
+    }
+  }
 
   // Don't hoist code containing allocas, invokes, or vastarts.
   for (BasicBlock::const_iterator I = BB.begin(), E = BB.end(); I != E; ++I) {
     if (isa<AllocaInst>(I) || isa<InvokeInst>(I))
       return false;
     if (const CallInst *CI = dyn_cast<CallInst>(I))
       if (const Function *F = CI->getCalledFunction())
         if (F->getIntrinsicID() == Intrinsic::vastart)

llvm/trunk/test/Transforms/CodeExtractor/BlockAddressReference.ll

+; RUN: opt < %s -loop-extract -S | FileCheck %s
+
+@label = common local_unnamed_addr global i8* null
+
+; CHECK: define
+; no outlined function
+; CHECK-NOT: define
+define i32 @sterix(i32 %n) {
+entry:
+  %tobool = icmp ne i32 %n, 0
+  ; this blockaddress references a basic block that goes in the extracted loop
+  %cond = select i1 %tobool, i8* blockaddress(@sterix, %for.cond), i8* blockaddress(@sterix, %exit)
+  store i8* %cond, i8** @label
+  %cmp5 = icmp sgt i32 %n, 0
+  br i1 %cmp5, label %for.body, label %exit
+
+for.cond:
+  %mul = shl nsw i32 %s.06, 1
+  %exitcond = icmp eq i32 %inc, %n
+  br i1 %exitcond, label %exit.loopexit, label %for.body
+
+for.body:
+  %i.07 = phi i32 [ %inc, %for.cond ], [ 0, %entry ]
+  %s.06 = phi i32 [ %mul, %for.cond ], [ 1, %entry ]
+  %inc = add nuw nsw i32 %i.07, 1
+  br label %for.cond
+
+exit.loopexit:
+  %phitmp = icmp ne i32 %s.06, 2
+  %phitmp8 = zext i1 %phitmp to i32
+  br label %exit
+
+exit:
+  %s.1 = phi i32 [ 1, %entry ], [ %phitmp8, %exit.loopexit ]
+  ret i32 %s.1
+}

llvm/trunk/test/Transforms/CodeExtractor/BlockAddressSelfReference.ll

+; RUN: opt < %s -loop-extract -S | FileCheck %s
+
[e-mail] = internal unnamed_addr constant [3 x i8*] [i8* blockaddress(@choum, %12), i8* blockaddress(@choum, %16), i8* blockaddress(@choum, %20)]
+
+; CHECK: define
+; no outlined function
+; CHECK-NOT: define
+
+define void @choum(i32, i32* nocapture, i32) {
+  %4 = icmp sgt i32 %0, 0
+  br i1 %4, label %5, label %26
+
+  %6 = sext i32 %2 to i64
+  %7 = getelementptr inbounds [3 x i8*], [3 x i8*]* @choum.addr, i64 0, i64 %6
+  %8 = load i8*, i8** %7
+  %9 = zext i32 %0 to i64
+  br label %10
+
+  %11 = phi i64 [ 0, %5 ], [ %24, %20 ]
+  indirectbr i8* %8, [label %12, label %16, label %20]
+
+  %13 = getelementptr inbounds i32, i32* %1, i64 %11
+  %14 = load i32, i32* %13
+  %15 = add nsw i32 %14, 1
+  store i32 %15, i32* %13
+  br label %16
+
+  %17 = getelementptr inbounds i32, i32* %1, i64 %11
+  %18 = load i32, i32* %17
+  %19 = shl nsw i32 %18, 1
+  store i32 %19, i32* %17
+  br label %20
+
+  %21 = getelementptr inbounds i32, i32* %1, i64 %11
+  %22 = load i32, i32* %21
+  %23 = add nsw i32 %22, -3
+  store i32 %23, i32* %21
+  %24 = add nuw nsw i64 %11, 1
+  %25 = icmp eq i64 %24, %9
+  br i1 %25, label %26, label %10
+
+  ret void
+}