[CodeExtractor] Allow extracting allocas within simple stack{save,restore} pairs

[vsk]

Summary

This enables outlining of a pervasive kind of cold code (error logging)
on Darwin. Fixes llvm.org/PR39441.

[efriedma]

The approach seems fine.

It seems weird that this is relevant; VLAs with a fixed size should generally be rare... unless the error logging code is intentionally using a VLA to trick the compiler into dynamically allocating the space?

llvm/lib/Transforms/Utils/CodeExtractor.cpp

215

There's no reason to expect a stacksave has exactly one use in general. (For example, it could have two uses if jump threading cloned a stackrestore, or it could be dead if the use was optimized out.)

223

You need to check that the parameter of the stackrestore is the corresponding stacksave.

[vsk]

• Fix the issues Eli pointed out.

The source is essentially:

uint8_t alignas(16) Buf[__builtin_os_log_format_buffer_size(fmt, ##__VA_ARGS__)];

In general I don't think this has to have a fixed size? Maybe some other optimization was able to simplify the alloca.

[efriedma]

__builtin_os_log_format_buffer_size is a compiler builtin which always returns a constant integer.

[vsk]

In D53779#1277925, @efriedma wrote:

__builtin_os_log_format_buffer_size is a compiler builtin which always returns a constant integer.

How were you able to tell? I couldn't parse this out from the BUILTIN declaration, but, I also don't have my head wrapped around the custom type-checking code for it.

At any rate, in clang's EmitAutoVarAlloca, arrays with a size given by this builtin have !Ty->isConstantSizeType(), hence the stack{save,restore}.

I'm not sure whether I need to keep digging into this. If this patch is only generally useful to handle VLAs, I think it'd still be worthwhile.

[efriedma]

How were you able to tell?

Looked at the expansion in CGBuiltin.cpp; it returns a ConstantInt.

[vsk]

Eli points out that D53514 will improve the codegen for the use case I'm most interested in, but in a way that makes this patch a bit less relevant. While it should still be useful to deal with actual VLAs, I'll need some other solution to allow outlining calls to os_log.

[hiraditya]

Is this patch still useful now that D53514 has landed?

llvm/include/llvm/Transforms/Utils/CodeExtractor.h

class Value;
   public:
     /// Create a code extractor for a sequence of blocks.
     ///
     /// Given a sequence of basic blocks where the first block in the sequence
     /// dominates the rest, prepare a code extractor object for pulling this
     /// sequence out into its new function. When a DominatorTree is also given,
     /// extra checking and transformations are enabled. If AllowVarArgs is true,
     /// vararg functions can be extracted. This is safe, if all vararg handling
-    /// code is extracted, including vastart. If AllowAlloca is true, then
-    /// extraction of blocks containing alloca instructions would be possible,
-    /// however code extractor won't validate whether extraction is legal.
+    /// code is extracted, including vastart. If AllowAlloca is true, extraction
+    /// of blocks containing alloca instructions is always permitted, however
+    /// code extractor won't validate whether extraction is legal. Allocas
+    /// occurring between a stack{save,restore} pair within a block are always
+    /// safely extracted.
     CodeExtractor(ArrayRef<BasicBlock *> BBs, DominatorTree *DT = nullptr,
                   bool AggregateArgs = false, BlockFrequencyInfo *BFI = nullptr,
                   BranchProbabilityInfo *BPI = nullptr,
                   bool AllowVarArgs = false, bool AllowAlloca = false,
                   std::string Suffix = "");
 
     /// Create a code extractor for a loop body.
     ///

llvm/lib/Transforms/Utils/CodeExtractor.cpp

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);
     }
   }
 
-  // If explicitly requested, allow vastart and alloca. For invoke instructions
-  // verify that extraction is valid.
+  // Always allow allocas wrapped in between a stack{save,restore} pair. If
+  // explicitly requested, allow vastart and other allocas. Only verify that
+  // extraction is valid for invoke instructions.
+  SmallPtrSet<Value const *, 1> ActiveStackSaves;
   for (BasicBlock::const_iterator I = BB.begin(), E = BB.end(); I != E; ++I) {
     if (isa<AllocaInst>(I)) {
-       if (!AllowAlloca)
+       if (ActiveStackSaves.empty() && !AllowAlloca)
          return false;
        continue;
     }
 
     if (const auto *II = dyn_cast<InvokeInst>(I)) {
       // Unwind destination (either a landingpad, catchswitch, or cleanuppad)
       // must be a part of the subgraph which is being extracted.
       if (auto *UBB = II->getUnwindDest())
for (BasicBlock::const_iterator I = BB.begin(), E = BB.end(); I != E; ++I) {
     }
     if (const auto *CRI = dyn_cast<CleanupReturnInst>(I)) {
       if (auto *UBB = CRI->getUnwindDest())
         if (!Result.count(UBB))
           return false;
       continue;
     }
 
-    if (const CallInst *CI = dyn_cast<CallInst>(I))
-      if (const Function *F = CI->getCalledFunction())
-        if (F->getIntrinsicID() == Intrinsic::vastart) {
-          if (AllowVarArgs)
-            continue;
-          else
-            return false;
-        }
+    if (const CallInst *CI = dyn_cast<CallInst>(I)) {
+      const Function *F = CI->getCalledFunction();
+      if (!F)
+        continue;
+
+      auto IID = F->getIntrinsicID();
+      if (IID == Intrinsic::vastart) {
+
+      // Keep track of stack{save,restore} pairs within this block. As long
+      // as at least one of the stackrestores has not occurred, it's safe to
+      // outline any allocas seen in the block.
+      if (IID == Intrinsic::stacksave) {
+        for (auto *Restore : CI->users()) {
+          auto *RestoreCI = cast<CallInst>(Restore);
+          assert(RestoreCI->getCalledFunction()->getIntrinsicID() ==
+                     Intrinsic::stackrestore &&
+                 "Use of stacksave not a stackrestore?");
+          if (RestoreCI->getParent() == &BB) {
+            ActiveStackSaves.insert(CI);
+            break;
+          }
+        }
+      } else if (IID == Intrinsic::stackrestore) {
+        ActiveStackSaves.erase(CI->getOperand(0));
+      }
+    }
   }
 
   return true;
 }
 
 /// Build a set of blocks to extract if the input blocks are viable.
 static SetVector<BasicBlock *>
 buildExtractionBlockSet(ArrayRef<BasicBlock *> BBs, DominatorTree *DT,
                         bool AllowVarArgs, bool AllowAlloca) {
   assert(!BBs.empty() && "The set of blocks to extract must be non-empty");
   SetVector<BasicBlock *> Result;
 
   // Loop over the blocks, adding them to our set-vector, and aborting with an
   // empty set if we encounter invalid blocks.
   for (BasicBlock *BB : BBs) {
     // If this block is dead, don't process it.
     if (DT && !DT->isReachableFromEntry(BB))
       continue;
 
     if (!Result.insert(BB))
       llvm_unreachable("Repeated basic blocks in extraction input");
   }
 
   for (auto *BB : Result) {
     if (!isBlockValidForExtraction(*BB, Result, AllowVarArgs, AllowAlloca))
       return {};
 
     // Make sure that the first block is not a landing pad.
     if (BB == Result.front()) {

llvm/test/Transforms/HotColdSplit/outline-safe-alloca.ll

+; RUN: opt -S -hotcoldsplit < %s | FileCheck %s
+
+target datalayout = "e-m:o-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-apple-macosx10.14.0"
+
+; Example taken from test/Transforms/HotColdSplit/outline-if-then-else.ll.
+
+; CHECK-LABEL: define {{.*}}@no_active_stack_save(
+; CHECK-NOT: @no_active_stack_save.cold
+define void @no_active_stack_save(i32 %cond) {
+entry:
+  %cond.addr = alloca i32
+  store i32 %cond, i32* %cond.addr
+  %0 = load i32, i32* %cond.addr
+  %tobool = icmp ne i32 %0, 0
+  br i1 %tobool, label %if.then, label %if.end2
+
+if.then:                                          ; preds = %entry
+  %1 = load i32, i32* %cond.addr
+  %cmp = icmp sgt i32 %1, 10
+  br i1 %cmp, label %if.then1, label %if.else
+
+if.then1:                                         ; preds = %if.then
+  call void @sideeffect(i32 0)
+  br label %if.end
+
+if.else:                                          ; preds = %if.then
+  call void @sideeffect(i32 1)
+  br label %if.end
+
+if.end:                                           ; preds = %if.else, %if.then1
+  %save = call i8* @llvm.stacksave()
+  %vla = alloca [100 x i8]
+  call void (...) @sink()
+  ret void
+
+if.end2:                                          ; preds = %entry
+  call void @sideeffect(i32 2)
+  ret void
+}
+
+; CHECK-LABEL: define {{.*}}@no_active_stack_save2(
+; CHECK-NOT: @no_active_stack_save2.cold
+define void @no_active_stack_save2(i32 %cond) {
+entry:
+  %cond.addr = alloca i32
+  store i32 %cond, i32* %cond.addr
+  %0 = load i32, i32* %cond.addr
+  %tobool = icmp ne i32 %0, 0
+  %save2 = call i8* @llvm.stacksave()
+  br i1 %tobool, label %if.then, label %if.end2
+
+if.then:                                          ; preds = %entry
+  %1 = load i32, i32* %cond.addr
+  %cmp = icmp sgt i32 %1, 10
+  br i1 %cmp, label %if.then1, label %if.else
+
+if.then1:                                         ; preds = %if.then
+  call void @sideeffect(i32 0)
+  br label %if.end
+
+if.else:                                          ; preds = %if.then
+  call void @sideeffect(i32 1)
+  br label %if.end
+
+if.end:                                           ; preds = %if.else, %if.then1
+  %save = call i8* @llvm.stacksave()
+  %vla = alloca [100 x i8]
+  call void (...) @sink()
+  call void @llvm.stackrestore(i8* %save2)
+  ret void
+
+if.end2:                                          ; preds = %entry
+  call void @sideeffect(i32 2)
+  ret void
+}
+
+; CHECK-LABEL: define {{.*}}@has_active_stack_save(
+; CHECK: call void @has_active_stack_save.cold.1
+define void @has_active_stack_save(i32 %cond) {
+entry:
+  %cond.addr = alloca i32
+  store i32 %cond, i32* %cond.addr
+  %0 = load i32, i32* %cond.addr
+  %tobool = icmp ne i32 %0, 0
+  br i1 %tobool, label %if.then, label %if.end2
+
+if.then:                                          ; preds = %entry
+  %1 = load i32, i32* %cond.addr
+  %cmp = icmp sgt i32 %1, 10
+  br i1 %cmp, label %if.then1, label %if.else
+
+if.then1:                                         ; preds = %if.then
+  call void @sideeffect(i32 0)
+  br label %if.end
+
+if.else:                                          ; preds = %if.then
+  call void @sideeffect(i32 1)
+  br label %if.end
+
+if.end:                                           ; preds = %if.else, %if.then1
+  %save = call i8* @llvm.stacksave()
+  %save2 = call i8* @llvm.stacksave() ; Useless, but shouldn't impede outlining.
+  %vla = alloca [100 x i8]
+  call void (...) @sink()
+  call void @llvm.stackrestore(i8* %save)
+  ret void
+
+if.end2:                                          ; preds = %entry
+  call void @sideeffect(i32 2)
+  ret void
+}
+
+
+; TODO: Hoist the alloca and eliminate the stack{save,restore} pair.
+; CHECK-LABEL: define {{.*}}@has_active_stack_save.cold.1
+; CHECK: call {{.*}}@sideeffect
+; CHECK: call {{.*}}@sideeffect
+; CHECK: call {{.*}}@llvm.stacksave
+; CHECK-NEXT: call {{.*}}@llvm.stacksave
+; CHECK-NEXT: alloca [100 x i8]
+; CHECK-NEXT: call {{.*}}@sink
+; CHECK-NEXT: call {{.*}}@llvm.stackrestore
+
+declare void @sideeffect(i32)
+
+declare void @sink(...) cold
+
+declare i8* @llvm.stacksave()
+
+declare void @llvm.stackrestore(i8*)