[BlockExtractor] Extend the file format to support the grouping of basic blocks

[qcolombet]

Summary

Prior to this patch, each basic block listed in the extrack-blocks-file
would be extracted to a different function.

This patch adds the support for comma separated list of basic blocks
to form group.

When the region formed by a group is not extractable, e.g., not single
entry, all the blocks of that group are left untouched.

Let us see this new format in action (comments are not part of the
file format):
;; funcName bbName[,bbName...]
foo bb1 ;; Extract bb1 in its own function
foo bb2,bb3 ;; Extract bb2,bb3 in their own function
bar bb1,bb4 ;; Extract bb1,bb4 in their own function
bar bb2 ;; Extract bb2 in its own function

Assuming all regions are extractable, this will create one function and
thus one call per region.

[volkan]

Hi Quentin,

This looks good to me, but seems like a bit complicated. I think we can simplify the logic by getting basic block groups as a comma separated string.
For example:

foo bb1 ;; Extract bb1 in its own function
foo bb2,bb3 ;; Extract bb2 and bb3 in their own function

What do you think?

[qcolombet]

Hi Volkan,

Good point!

I ended up tailoring the syntax to my current use case (which basically adds group ids at the end of the exiting file with sed), but your syntax is much easier to understand!

Updating!
Cheers,
-Quentin

[qcolombet]

Update:

• Use comma separated list instead of group naming.

[volkan]

LGTM, thanks Quentin. Could you please update the summary before pushing this?

[qcolombet]

Could you please update the summary before pushing this?

Done!

Thanks!!

llvm/lib/Transforms/IPO/BlockExtractor.cpp

 #include "llvm/IR/Module.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Transforms/IPO.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include "llvm/Transforms/Utils/CodeExtractor.h"
+
+#include <map>
+
 using namespace llvm;
 
 #define DEBUG_TYPE "block-extractor"
 
 STATISTIC(NumExtracted, "Number of basic blocks extracted");
 
 static cl::opt<std::string> BlockExtractorFile(
     "extract-blocks-file", cl::value_desc("filename"),
     cl::desc("A file containing list of basic blocks to extract"), cl::Hidden);
 
 cl::opt<bool> BlockExtractorEraseFuncs("extract-blocks-erase-funcs",
                                        cl::desc("Erase the existing functions"),
                                        cl::Hidden);
-
 namespace {
 class BlockExtractor : public ModulePass {
-  SmallVector<BasicBlock *, 16> Blocks;
+  SmallVector<SmallVector<BasicBlock *, 16>, 4> GroupsOfBlocks;
   bool EraseFunctions;
-  SmallVector<std::pair<std::string, std::string>, 32> BlocksByName;
+  /// Map a function name to groups of blocks.
+  /// The groups are formed according to the group ID given
+  /// in the input file.
+  SmallVector<std::pair<std::string, SmallVector<std::string, 4>>, 4>
+      BlocksByName;
+
+  /// Create a new entry in the BlocksByName and set the first element
+  /// of the created pair to \p FuncName.
+  /// \return the index in BlocksByName of the newly created entry.
+  size_t createNewBlocksByNameEntry(StringRef FuncName) {
+    size_t NextEntryIdx = BlocksByName.size();
+    BlocksByName.push_back({FuncName, SmallVector<std::string, 4>()});
+    return NextEntryIdx;
+  }
 
 public:
   static char ID;
   BlockExtractor(const SmallVectorImpl<BasicBlock *> &BlocksToExtract,
                  bool EraseFunctions)
-      : ModulePass(ID), Blocks(BlocksToExtract.begin(), BlocksToExtract.end()),
-        EraseFunctions(EraseFunctions) {
+      : ModulePass(ID), EraseFunctions(EraseFunctions) {
+    // We want one group per element of the input list.
+    for (BasicBlock *BB : BlocksToExtract) {
+      SmallVector<BasicBlock *, 16> NewGroup;
+      NewGroup.push_back(BB);
+      GroupsOfBlocks.push_back(NewGroup);
+    }
     if (!BlockExtractorFile.empty())
       loadFile();
   }
   BlockExtractor() : BlockExtractor(SmallVector<BasicBlock *, 0>(), false) {}
   bool runOnModule(Module &M) override;
 
 private:
   void loadFile();

 
 /// Gets all of the blocks specified in the input file.
 void BlockExtractor::loadFile() {
   auto ErrOrBuf = MemoryBuffer::getFile(BlockExtractorFile);
   if (ErrOrBuf.getError())
     report_fatal_error("BlockExtractor couldn't load the file.");
   // Read the file.
   auto &Buf = *ErrOrBuf;
+  std::map<std::pair<std::string, std::string>, unsigned> GroupToID;
   SmallVector<StringRef, 16> Lines;
   Buf->getBuffer().split(Lines, '\n', /*MaxSplit=*/-1,
                          /*KeepEmpty=*/false);
   for (const auto &Line : Lines) {
-    auto FBPair = Line.split(' ');
-    BlocksByName.push_back({FBPair.first, FBPair.second});
+    SmallVector<StringRef, 4> LineSplit;
+    Line.split(LineSplit, ' ', /*MaxSplit=*/-1,
+               /*KeepEmpty=*/false);
+    if (LineSplit.empty())
+      continue;
+    std::string FuncName = LineSplit[0];
+    std::string BBName = LineSplit[1];
+    unsigned GroupIdx;
+    if (LineSplit.size() == 3) {
+      std::string SpecifiedGroupID = LineSplit[2];
+      auto FuncAndID = std::make_pair<>(FuncName, SpecifiedGroupID);
+      auto GroupToIDIt = GroupToID.find(FuncAndID);
+      if (GroupToIDIt == GroupToID.end())
+        GroupToIDIt =
+            GroupToID.insert({FuncAndID, createNewBlocksByNameEntry(FuncName)})
+                .first;
+      GroupIdx = GroupToIDIt->second;
+    } else {
+      assert(LineSplit.size() == 2 && "Invalid file format");
+      GroupIdx = createNewBlocksByNameEntry(FuncName);
+    }
+    BlocksByName[GroupIdx].second.push_back(BBName);
   }
 }
 
 /// Extracts the landing pads to make sure all of them have only one
 /// predecessor.
 void BlockExtractor::splitLandingPadPreds(Function &F) {
   for (BasicBlock &BB : F) {
     for (Instruction &I : BB) {
bool BlockExtractor::runOnModule(Module &M) {
   // Get all the functions.
   SmallVector<Function *, 4> Functions;
   for (Function &F : M) {
     splitLandingPadPreds(F);
     Functions.push_back(&F);
   }
 
   // Get all the blocks specified in the input file.
+  unsigned NextGroupIdx = GroupsOfBlocks.size();
+  GroupsOfBlocks.resize(NextGroupIdx + BlocksByName.size());
   for (const auto &BInfo : BlocksByName) {
     Function *F = M.getFunction(BInfo.first);
     if (!F)
       report_fatal_error("Invalid function name specified in the input file");
+    for (const auto &BBInfo : BInfo.second) {
       auto Res = llvm::find_if(*F, [&](const BasicBlock &BB) {
-      return BB.getName().equals(BInfo.second);
+        return BB.getName().equals(BBInfo);
       });
       if (Res == F->end())
         report_fatal_error("Invalid block name specified in the input file");
-    Blocks.push_back(&*Res);
+      GroupsOfBlocks[NextGroupIdx].push_back(&*Res);
+    }
+    ++NextGroupIdx;
   }
 
-  // Extract basic blocks.
-  for (BasicBlock *BB : Blocks) {
+  // Extract each group of basic blocks.
+  for (auto &BBs : GroupsOfBlocks) {
+    SmallVector<BasicBlock *, 32> BlocksToExtractVec;
+    for (BasicBlock *BB : BBs) {
       // Check if the module contains BB.
       if (BB->getParent()->getParent() != &M)
         report_fatal_error("Invalid basic block");
       LLVM_DEBUG(dbgs() << "BlockExtractor: Extracting "
                         << BB->getParent()->getName() << ":" << BB->getName()
                         << "\n");
-    SmallVector<BasicBlock *, 2> BlocksToExtractVec;
       BlocksToExtractVec.push_back(BB);
       if (const InvokeInst *II = dyn_cast<InvokeInst>(BB->getTerminator()))
         BlocksToExtractVec.push_back(II->getUnwindDest());
-    CodeExtractor(BlocksToExtractVec).extractCodeRegion();
       ++NumExtracted;
       Changed = true;
     }
+    Function *F = CodeExtractor(BlocksToExtractVec).extractCodeRegion();
+    if (F)
+      LLVM_DEBUG(dbgs() << "Extracted group '" << (*BBs.begin())->getName()
+                        << "' in: " << F->getName() << '\n');
+    else
+      LLVM_DEBUG(dbgs() << "Failed to extract for group '"
+                        << (*BBs.begin())->getName() << "'\n");
+  }
 
   // Erase the functions.
   if (EraseFunctions || BlockExtractorEraseFuncs) {
     for (Function *F : Functions) {
       LLVM_DEBUG(dbgs() << "BlockExtractor: Trying to delete " << F->getName()
                         << "\n");
       F->deleteBody();
     }
     // Set linkage as ExternalLinkage to avoid erasing unreachable functions.
     for (Function &F : M)
       F.setLinkage(GlobalValue::ExternalLinkage);
     Changed = true;
   }
 
   return Changed;
 }

llvm/test/Transforms/BlockExtractor/extract-blocks-with-groups.ll

+; Extract the 'if', 'then', and 'else' blocks into the same function that
+; we call 'diamond'
+; RUN: echo 'foo if diamond' > %t
+; RUN: echo 'foo then diamond' >> %t
+; RUN: echo 'foo else diamond' >> %t
+; Make sure we can still extract 'end' without specifying a group name
+; RUN: echo 'foo end' >> %t
+; Make sure we reuse the group name just fine across different functions.
+; RUN: echo 'bar bb14 diamond' >> %t
+; RUN: echo 'bar bb20 diamond' >> %t
+; RUN: opt -S -extract-blocks -extract-blocks-file=%t %s | FileCheck %s
+
+; CHECK-LABEL: foo
+;
+; CHECK: if:
+; The diamond will produce an i32 value, make sure we account for that.
+; CHECK: [[RES_VAL_ADDR:%[^ ]*]] = alloca i32
+; CHECK-NEXT: br label %[[FOO_DIAMOND_LABEL:.*$]]
+;
+; The if-then-else blocks should be in just one function.
+; CHECK: [[FOO_DIAMOND_LABEL]]:
+; CHECK: call void [[FOO_DIAMOND:@[^(]*]](i32 %arg1, i32 %arg, i32* [[RES_VAL_ADDR]])
+; CHECK-NEXT: [[RES_VAL:%[^ ]*]] = load i32, i32* [[RES_VAL_ADDR]]
+; Then it should directly jump to end.
+; CHECK: br label %[[FOO_END_LABEL:.*$]]
+;
+; End should have been extracted into its own function.
+; CHECK: [[FOO_END_LABEL]]:
+; CHECK: [[CMP:%[^ ]*]] = call i1 [[FOO_END:@[^(]*]](i32 [[RES_VAL]], i32 %arg)
+; CHECK-NEXT: br i1 [[CMP]], label %ret0, label %ret1
+define i32 @foo(i32 %arg, i32 %arg1) {
+if:
+  %tmp5 = icmp sgt i32 %arg, 0
+  %tmp8 = icmp sgt i32 %arg1, 0
+  %or.cond = and i1 %tmp5, %tmp8
+  br i1 %or.cond, label %then, label %else
+
+then:
+  %tmp12 = shl i32 %arg1, 2
+  %tmp13 = add nsw i32 %tmp12, %arg
+  br label %end
+
+else:
+  %tmp22 = mul nsw i32 %arg, 3
+  %tmp24 = sdiv i32 %arg1, 6
+  %tmp25 = add nsw i32 %tmp24, %tmp22
+  br label %end
+
+end:
+  %tmp.0 = phi i32 [ %tmp13, %then ], [ %tmp25, %else ]
+  %and0 = and i32 %tmp.0, %arg
+  %cmp1 = icmp slt i32 %and0, 0
+  br i1 %cmp1, label %ret0, label %ret1
+
+ret0:
+  ret i32 0
+
+ret1:
+  ret i32 1
+}
+
+; CHECK-LABEL: bar
+;
+; Check that we extracted bb14 and bb20 in their own (shared) function.
+; CHECK: bb
+; CHECK: br i1 %or.cond, label %bb9, label %[[BAR_DIAMOND_LABEL:.*$]]
+;
+; CHECK: [[BAR_DIAMOND_LABEL]]:
+; CHECK: [[CMP:%[^ ]*]] = call i1 [[BAR_DIAMOND:@[^(]*]](i32 %arg1, i32 %arg, i32*
+; CHECK: br i1 [[CMP]], label %bb26, label %bb30
+define i32 @bar(i32 %arg, i32 %arg1) {
+bb:
+  %tmp5 = icmp sgt i32 %arg, 0
+  %tmp8 = icmp sgt i32 %arg1, 0
+  %or.cond = and i1 %tmp5, %tmp8
+  br i1 %or.cond, label %bb9, label %bb14
+
+bb9:                                              ; preds = %bb
+  %tmp12 = shl i32 %arg1, 2
+  %tmp13 = add nsw i32 %tmp12, %arg
+  br label %bb30
+
+bb14:                                             ; preds = %bb
+  %0 = and i32 %arg1, %arg
+  %1 = icmp slt i32 %0, 0
+  br i1 %1, label %bb20, label %bb26
+
+bb20:                                             ; preds = %bb14
+  %tmp22 = mul nsw i32 %arg, 3
+  %tmp24 = sdiv i32 %arg1, 6
+  %tmp25 = add nsw i32 %tmp24, %tmp22
+  br label %bb30
+
+bb26:                                             ; preds = %bb14
+  %tmp29 = sub nsw i32 %arg, %arg1
+  br label %bb30
+
+bb30:                                             ; preds = %bb26, %bb20, %bb9
+  %tmp.0 = phi i32 [ %tmp13, %bb9 ], [ %tmp25, %bb20 ], [ %tmp29, %bb26 ]
+  ret i32 %tmp.0
+}
+
+; Check that we extracted the three asked basic blocks.
+; CHECK: [[FOO_DIAMOND]]
+; CHECK: then:
+; Make sure the function doesn't end in the middle of the
+; list of the blocks we are checking.
+; CHECK-NOT: }
+; CHECK: else:
+; CHECK-NOT: }
+; The name of the if block is weird because it had to be split
+; since it was the name of the entry of the original function.
+; CHECK: if.split:
+; CHECK: }
+
+; CHECK: [[FOO_END]]
+; CHECK-NOT: }
+; CHECK: end:
+; CHECK: }
+
+; Check that we extracted the two asked basic blocks.
+; CHECK: [[BAR_DIAMOND]]
+; CHECK-NOT: }
+; CHECK: bb14:
+; CHECK-NOT: }
+; CHECK: bb20:
+; CHECK: }