[HotColdSplit] Add options for splitting cold functions in separate section

Add support for (if enabled) splitting cold functions into a separate section
in order to further boost locality of hot code.

Authored by: rjf (Ruijie Fang)
Reviewed by: hiraditya,rcorcs,vsk

Differential Revision: https://reviews.llvm.org/D85331

llvm/lib/Transforms/IPO/HotColdSplitting.cpp

@@ -69,6 +69,7 @@
 #include "llvm/Transforms/Utils/ValueMapper.h"
 #include <algorithm>
 #include <cassert>
+#include <string>
 
 #define DEBUG_TYPE "hotcoldsplit"
 
@@ -85,6 +86,17 @@ static cl::opt<int>
                        cl::desc("Base penalty for splitting cold code (as a "
                                 "multiple of TCC_Basic)"));
 
+static cl::opt<bool> EnableColdSection(
+    "enable-cold-section", cl::init(false), cl::Hidden,
+    cl::desc("Enable placement of extracted cold functions"
+             " into a separate section after hot-cold splitting."));
+
+static cl::opt<std::string>
+    ColdSectionName("hotcoldsplit-cold-section-name", cl::init("__llvm_cold"),
+                    cl::Hidden,
+                    cl::desc("Name for the section containing cold functions "
+                             "extracted by hot-cold splitting."));
+
 namespace {
 // Same as blockEndsInUnreachable in CodeGen/BranchFolding.cpp. Do not modify
 // this function unless you modify the MBB version as well.
@@ -339,8 +351,12 @@ Function *HotColdSplitting::extractColdRegion(
     }
     CI->setIsNoInline();
 
-    if (OrigF->hasSection())
-      OutF->setSection(OrigF->getSection());
+    if (EnableColdSection)
+      OutF->setSection(ColdSectionName);
+    else {
+      if (OrigF->hasSection())
+        OutF->setSection(OrigF->getSection());
+    }
 
     markFunctionCold(*OutF, BFI != nullptr);
 

llvm/test/Transforms/HotColdSplit/section-splitting-custom.ll

@@ -0,0 +1,40 @@
+; RUN: opt -hotcoldsplit -hotcoldsplit-threshold=-1 -pass-remarks=hotcoldsplit -enable-cold-section=true -hotcoldsplit-cold-section-name="__cold_custom" -S < %s 2>&1 | FileCheck %s
+; RUN: opt -passes=hotcoldsplit -hotcoldsplit-threshold=-1 -pass-remarks=hotcoldsplit -enable-cold-section=true -hotcoldsplit-cold-section-name="__cold_custom" -S < %s 2>&1 | FileCheck %s
+
+; This test case is copied over from split-cold-2.ll, modified
+; to test the `-enable-cold-section` and `-hotcoldsplit-cold-section-name`
+; parameters in opt.
+; Make sure this compiles. This test used to fail with an invalid phi node: the
+; two predecessors were outlined and the SSA representation was invalid.
+
+; CHECK: remark: <unknown>:0:0: fun split cold code into fun.cold.1
+; CHECK-LABEL: @fun
+; CHECK: codeRepl:
+; CHECK-NEXT: call void @fun.cold.1
+
+; CHECK: define {{.*}}@fun.cold.1{{.*}} [[cold_attr:#[0-9]+]] section "__cold_custom"
+; CHECK: attributes [[cold_attr]] = { {{.*}}noreturn
+
+define void @fun() {
+entry:
+  br i1 undef, label %if.then, label %if.else
+
+if.then:
+  ret void
+
+if.else:
+  br label %if.then4
+
+if.then4:
+  br i1 undef, label %if.then5, label %if.end
+
+if.then5:
+  br label %cleanup
+
+if.end:
+  br label %cleanup
+
+cleanup:
+  %cleanup.dest.slot.0 = phi i32 [ 1, %if.then5 ], [ 0, %if.end ]
+  unreachable
+}

llvm/test/Transforms/HotColdSplit/section-splitting-default.ll

@@ -0,0 +1,39 @@
+; RUN: opt -hotcoldsplit -hotcoldsplit-threshold=-1 -pass-remarks=hotcoldsplit -enable-cold-section=true -S < %s 2>&1 | FileCheck %s
+; RUN: opt -passes=hotcoldsplit -hotcoldsplit-threshold=-1 -pass-remarks=hotcoldsplit -enable-cold-section=true -S < %s 2>&1 | FileCheck %s
+
+; This test case is copied over from split-cold-2.ll, modified
+; to test the `-enable-cold-section` parameter in opt.
+; Make sure this compiles. This test used to fail with an invalid phi node: the
+; two predecessors were outlined and the SSA representation was invalid.
+
+; CHECK: remark: <unknown>:0:0: fun split cold code into fun.cold.1
+; CHECK-LABEL: @fun
+; CHECK: codeRepl:
+; CHECK-NEXT: call void @fun.cold.1
+
+; CHECK: define {{.*}}@fun.cold.1{{.*}} [[cold_attr:#[0-9]+]] section "__llvm_cold"
+; CHECK: attributes [[cold_attr]] = { {{.*}}noreturn
+
+define void @fun() {
+entry:
+  br i1 undef, label %if.then, label %if.else
+
+if.then:
+  ret void
+
+if.else:
+  br label %if.then4
+
+if.then4:
+  br i1 undef, label %if.then5, label %if.end
+
+if.then5:
+  br label %cleanup
+
+if.end:
+  br label %cleanup
+
+cleanup:
+  %cleanup.dest.slot.0 = phi i32 [ 1, %if.then5 ], [ 0, %if.end ]
+  unreachable
+}