IRGen: -enable-split-cold-code frontend feature

Author: kavon

include/swift/AST/IRGenOptions.h

@@ -475,6 +475,9 @@ class IRGenOptions {
 
   unsigned UseFragileResilientProtocolWitnesses : 1;
 
+  // Whether to run the HotColdSplitting pass when optimizing.
+  unsigned EnableHotColdSplit : 1;
+
   /// The number of threads for multi-threaded code generation.
   unsigned NumThreads = 0;
 
@@ -564,6 +567,7 @@ class IRGenOptions {
         DisableReadonlyStaticObjects(false), CollocatedMetadataFunctions(false),
         ColocateTypeDescriptors(true), UseRelativeProtocolWitnessTables(false),
         UseFragileResilientProtocolWitnesses(false),
+        EnableHotColdSplit(false),
         CmdArgs(), SanitizeCoverage(llvm::SanitizerCoverageOptions()),
         TypeInfoFilter(TypeInfoDumpFilter::All),
         PlatformCCallingConvention(llvm::CallingConv::C), UseCASBackend(false),

include/swift/Option/FrontendOptions.td

@@ -1310,6 +1310,13 @@ def disable_fragile_resilient_protocol_witnesses :
   Flag<["-"], "disable-fragile-relative-protocol-tables">,
   HelpText<"Disable relative protocol witness tables">;
 
+def enable_split_cold_code :
+  Flag<["-"], "enable-split-cold-code">,
+  HelpText<"Enable splitting of cold code when optimizing">;
+def disable_split_cold_code :
+  Flag<["-"], "disable-split-cold-code">,
+  HelpText<"Disable splitting of cold code when optimizing">;
+
 def enable_new_llvm_pass_manager :
   Flag<["-"], "enable-new-llvm-pass-manager">,
   HelpText<"Enable the new llvm pass manager">;

lib/Frontend/CompilerInvocation.cpp

@@ -3362,6 +3362,10 @@ static bool ParseIRGenArgs(IRGenOptions &Opts, ArgList &Args,
     Args.hasFlag(OPT_enable_fragile_resilient_protocol_witnesses,
                  OPT_disable_fragile_resilient_protocol_witnesses,
                  Opts.UseFragileResilientProtocolWitnesses);
+  Opts.EnableHotColdSplit =
+      Args.hasFlag(OPT_enable_split_cold_code,
+                   OPT_disable_split_cold_code,
+                   Opts.EnableHotColdSplit);
   Opts.EnableLargeLoadableTypesReg2Mem =
       Args.hasFlag(OPT_enable_large_loadable_types_reg2mem,
                    OPT_disable_large_loadable_types_reg2mem,

lib/IRGen/IRGen.cpp

@@ -211,6 +211,7 @@ void swift::performLLVMOptimizations(const IRGenOptions &Opts,
   bool RunSwiftSpecificLLVMOptzns =
       !Opts.DisableSwiftSpecificLLVMOptzns && !Opts.DisableLLVMOptzns;
 
+  bool DoHotColdSplit = false;
   PTO.CallGraphProfile = false;
 
   llvm::OptimizationLevel level = llvm::OptimizationLevel::O0;
@@ -221,6 +222,7 @@ void swift::performLLVMOptimizations(const IRGenOptions &Opts,
     PTO.LoopVectorization = true;
     PTO.SLPVectorization = true;
     PTO.MergeFunctions = true;
+    DoHotColdSplit = Opts.EnableHotColdSplit;
     level = llvm::OptimizationLevel::Os;
   } else {
     level = llvm::OptimizationLevel::O0;
@@ -259,6 +261,8 @@ void swift::performLLVMOptimizations(const IRGenOptions &Opts,
   PB.crossRegisterProxies(LAM, FAM, CGAM, MAM);
   ModulePassManager MPM;
 
+  PB.setEnableHotColdSplitting(DoHotColdSplit);
+
   if (RunSwiftSpecificLLVMOptzns) {
     PB.registerScalarOptimizerLateEPCallback(
         [](FunctionPassManager &FPM, OptimizationLevel Level) {

test/IRGen/cold_split.swift

@@ -0,0 +1,37 @@
+// RUN: %target-swift-frontend %s -module-name=test -emit-assembly \
+// RUN:   -enable-throws-prediction -O -enable-split-cold-code \
+// RUN:       | %FileCheck --check-prefix CHECK-ENABLED %s
+
+//// Test disabling just the pass doesn't yield a split.
+// RUN: %target-swift-frontend %s -module-name=test -emit-assembly \
+// RUN:   -enable-throws-prediction -O -disable-split-cold-code \
+// RUN:       | %FileCheck --check-prefix CHECK-DISABLED %s
+
+//// Test disabling optimization entirely doesn't yield a split.
+// RUN: %target-swift-frontend %s -module-name=test -emit-assembly \
+// RUN:   -enable-throws-prediction -enable-split-cold-code \
+// RUN:       | %FileCheck --check-prefix CHECK-DISABLED %s
+
+
+// CHECK-ENABLED: cold
+
+// CHECK-DISABLED-NOT: cold
+
+enum MyError: Error { case err }
+
+func getRandom(_ b: Bool) throws -> Int {
+    if b {
+        return Int.random(in: 0..<1024)
+    } else {
+        throw MyError.err
+    }
+}
+
+public func numberWithLogging(_ b: Bool) -> Int {
+    do {
+        return try getRandom(b)
+    } catch {
+        print("Log: random number generator failed with b=\(b)")
+        return 1337
+    }
+}