[llvm-extract] Expose the group extraction feature of the BlockExtractor

[qcolombet]

Summary

This patch extends the -bb option to be able to use the group
extraction feature from the BlockExtractor.
In particular, -bb=func:bb is modified to support a list of basic
blocks per function: -bb=func:bb1[;bb2...] that will be extracted
together if at all possible (region must be single entry.)

[volkan]

Thanks Quentin, LGTM.

llvm/trunk/test/tools/llvm-extract/extract-blocks-with-groups.ll

+; RUN: llvm-extract -bb 'foo:if;then;else' -bb 'bar:bb14;bb20' -S %s  | FileCheck %s
+; Extract two groups of basic blocks in two different functions.
+
+
+; The first extracted function is the region composed by the
+; blocks if, then, and else from foo.
+; CHECK: define dso_local void @foo.if.split(i32 %arg1, i32 %arg, i32* %tmp.0.ce.out) {
+; CHECK: newFuncRoot:
+; CHECK:   br label %if.split
+;
+; CHECK: end.exitStub:                                     ; preds = %end.split
+; CHECK:   ret void
+;
+; CHECK: then:                                             ; preds = %if.split
+; CHECK:   %tmp12 = shl i32 %arg1, 2
+; CHECK:   %tmp13 = add nsw i32 %tmp12, %arg
+; CHECK:   br label %end.split
+;
+; CHECK: else:                                             ; preds = %if.split
+; CHECK:   %tmp22 = mul nsw i32 %arg, 3
+; CHECK:   %tmp24 = sdiv i32 %arg1, 6
+; CHECK:   %tmp25 = add nsw i32 %tmp24, %tmp22
+; CHECK:   br label %end.split
+;
+; CHECK: if.split:                                         ; preds = %newFuncRoot
+; CHECK:   %tmp5 = icmp sgt i32 %arg, 0
+; CHECK:   %tmp8 = icmp sgt i32 %arg1, 0
+; CHECK:   %or.cond = and i1 %tmp5, %tmp8
+; CHECK:   br i1 %or.cond, label %then, label %else
+;
+; CHECK: end.split:                                        ; preds = %then, %else
+; CHECK:   %tmp.0.ce = phi i32 [ %tmp13, %then ], [ %tmp25, %else ]
+; CHECK:   store i32 %tmp.0.ce, i32* %tmp.0.ce.out
+; CHECK:   br label %end.exitStub
+; CHECK: }
+
+; The second extracted function is the region composed by the blocks
+; bb14 and bb20 from bar.
+; CHECK: define dso_local i1 @bar.bb14(i32 %arg1, i32 %arg, i32* %tmp25.out) {
+; CHECK: newFuncRoot:
+; CHECK:   br label %bb14
+;
+; CHECK: bb26.exitStub:                                    ; preds = %bb14
+; CHECK:   ret i1 true
+;
+; CHECK: bb30.exitStub:                                    ; preds = %bb20
+; CHECK:   ret i1 false
+;
+; CHECK: bb14:                                             ; preds = %newFuncRoot
+; CHECK:   %tmp0 = and i32 %arg1, %arg
+; CHECK:   %tmp1 = icmp slt i32 %tmp0, 0
+; CHECK:   br i1 %tmp1, label %bb20, label %bb26.exitStub
+;
+; CHECK: bb20:                                             ; preds = %bb14
+; CHECK:   %tmp22 = mul nsw i32 %arg, 3
+; CHECK:   %tmp24 = sdiv i32 %arg1, 6
+; CHECK:   %tmp25 = add nsw i32 %tmp24, %tmp22
+; CHECK:   store i32 %tmp25, i32* %tmp25.out
+; CHECK:   br label %bb30.exitStub
+; CHECK: }
+
+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
+}
+
+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
+  %tmp0 = and i32 %arg1, %arg
+  %tmp1 = icmp slt i32 %tmp0, 0
+  br i1 %tmp1, 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
+}
+

llvm/trunk/tools/llvm-extract/llvm-extract.cpp

for (size_t i = 0, e = ExtractRegExpFuncs.size(); i != e; ++i) {
     if (!match) {
       errs() << argv[0] << ": program doesn't contain global named '"
              << ExtractRegExpFuncs[i] << "'!\n";
       return 1;
     }
   }
 
   // Figure out which BasicBlocks we should extract.
-  SmallVector<BasicBlock *, 4> BBs;
+  SmallVector<SmallVector<BasicBlock *, 16>, 4> GroupOfBBs;
   for (StringRef StrPair : ExtractBlocks) {
     auto BBInfo = StrPair.split(':');
     // Get the function.
     Function *F = M->getFunction(BBInfo.first);
     if (!F) {
       errs() << argv[0] << ": program doesn't contain a function named '"
              << BBInfo.first << "'!\n";
       return 1;
     }
     // Do not materialize this function.
     GVs.insert(F);
-    // Get the basic block.
+    // Get the basic blocks.
+    SmallVector<BasicBlock *, 16> BBs;
+    SmallVector<StringRef, 16> BBNames;
+    BBInfo.second.split(BBNames, ';', /*MaxSplit=*/-1,
+                        /*KeepEmpty=*/false);
+    for (StringRef BBName : BBNames) {
       auto Res = llvm::find_if(*F, [&](const BasicBlock &BB) {
-      return BB.getName().equals(BBInfo.second);
+        return BB.getName().equals(BBName);
       });
       if (Res == F->end()) {
         errs() << argv[0] << ": function " << F->getName()
                << " doesn't contain a basic block named '" << BBInfo.second
                << "'!\n";
         return 1;
       }
       BBs.push_back(&*Res);
     }
+    GroupOfBBs.push_back(BBs);
+  }
 
   // Use *argv instead of argv[0] to work around a wrong GCC warning.
   ExitOnError ExitOnErr(std::string(*argv) + ": error reading input: ");
 
   if (Recursive) {
     std::vector<llvm::Function *> Workqueue;
     for (GlobalValue *GV : GVs) {
       if (auto *F = dyn_cast<Function>(GV)) {
if (!DeleteFn) {
     // FIXME: should the GVExtractionPass handle this?
     ExitOnErr(M->materializeAll());
   }
 
   // Extract the specified basic blocks from the module and erase the existing
   // functions.
   if (!ExtractBlocks.empty()) {
     legacy::PassManager PM;
-    PM.add(createBlockExtractorPass(BBs, true));
+    PM.add(createBlockExtractorPass(GroupOfBBs, true));
     PM.run(*M);
   }
 
   // In addition to deleting all other functions, we also want to spiff it
   // up a little bit.  Do this now.
   legacy::PassManager Passes;
 
   if (!DeleteFn)