3 files changed, 67 insertions, 23 deletions

diff --git a/llvm/lib/LTO/LTO.cpp b/llvm/lib/LTO/LTO.cpp
index 23a4a4011f7..8e0d532fa7e 100644
--- a/llvm/lib/LTO/LTO.cpp
+++ b/llvm/lib/LTO/LTO.cpp
@@ -20,6 +20,7 @@
 #include "llvm/Config/llvm-config.h"
 #include "llvm/IR/AutoUpgrade.h"
 #include "llvm/IR/DiagnosticPrinter.h"
+#include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/LegacyPassManager.h"
 #include "llvm/IR/Mangler.h"
 #include "llvm/IR/Metadata.h"
@@ -808,6 +809,45 @@ unsigned LTO::getMaxTasks() const {
   return RegularLTO.ParallelCodeGenParallelismLevel + ThinLTO.ModuleMap.size();
 }
 
+// If only some of the modules were split, we cannot correctly handle
+// code that contains type tests or type checked loads.
+Error LTO::checkPartiallySplit() {
+  if (!ThinLTO.CombinedIndex.partiallySplitLTOUnits())
+    return Error::success();
+
+  Function *TypeTestFunc = RegularLTO.CombinedModule->getFunction(
+      Intrinsic::getName(Intrinsic::type_test));
+  Function *TypeCheckedLoadFunc = RegularLTO.CombinedModule->getFunction(
+      Intrinsic::getName(Intrinsic::type_checked_load));
+
+  // First check if there are type tests / type checked loads in the
+  // merged regular LTO module IR.
+  if ((TypeTestFunc && !TypeTestFunc->use_empty()) ||
+      (TypeCheckedLoadFunc && !TypeCheckedLoadFunc->use_empty()))
+    return make_error<StringError>(
+        "inconsistent LTO Unit splitting (recompile with -fsplit-lto-unit)",
+        inconvertibleErrorCode());
+
+  // Otherwise check if there are any recorded in the combined summary from the
+  // ThinLTO modules.
+  for (auto &P : ThinLTO.CombinedIndex) {
+    for (auto &S : P.second.SummaryList) {
+      auto *FS = dyn_cast<FunctionSummary>(S.get());
+      if (!FS)
+        continue;
+      if (!FS->type_test_assume_vcalls().empty() ||
+          !FS->type_checked_load_vcalls().empty() ||
+          !FS->type_test_assume_const_vcalls().empty() ||
+          !FS->type_checked_load_const_vcalls().empty() ||
+          !FS->type_tests().empty())
+        return make_error<StringError>(
+            "inconsistent LTO Unit splitting (recompile with -fsplit-lto-unit)",
+            inconvertibleErrorCode());
+    }
+  }
+  return Error::success();
+}
+
 Error LTO::run(AddStreamFn AddStream, NativeObjectCache Cache) {
   // Compute "dead" symbols, we don't want to import/export these!
   DenseSet<GlobalValue::GUID> GUIDPreservedSymbols;
@@ -850,6 +890,17 @@ Error LTO::run(AddStreamFn AddStream, NativeObjectCache Cache) {
     StatsFile->keep();
   }
 
+  // Finalize linking of regular LTO modules containing summaries now that
+  // we have computed liveness information.
+  for (auto &M : RegularLTO.ModsWithSummaries)
+    if (Error Err = linkRegularLTO(std::move(M),
+                                   /*LivenessFromIndex=*/true))
+      return Err;
+
+  // Ensure we don't have inconsistently split LTO units with type tests.
+  if (Error Err = checkPartiallySplit())
+    return Err;
+
   Error Result = runRegularLTO(AddStream);
   if (!Result)
     Result = runThinLTO(AddStream, Cache);
@@ -861,11 +912,6 @@ Error LTO::run(AddStreamFn AddStream, NativeObjectCache Cache) {
 }
 
 Error LTO::runRegularLTO(AddStreamFn AddStream) {
-  for (auto &M : RegularLTO.ModsWithSummaries)
-    if (Error Err = linkRegularLTO(std::move(M),
-                                   /*LivenessFromIndex=*/true))
-      return Err;
-
   // Make sure commons have the right size/alignment: we kept the largest from
   // all the prevailing when adding the inputs, and we apply it here.
   const DataLayout &DL = RegularLTO.CombinedModule->getDataLayout();
diff --git a/llvm/lib/Transforms/IPO/LowerTypeTests.cpp b/llvm/lib/Transforms/IPO/LowerTypeTests.cpp
index 2f8a96be875..398005d2234 100644
--- a/llvm/lib/Transforms/IPO/LowerTypeTests.cpp
+++ b/llvm/lib/Transforms/IPO/LowerTypeTests.cpp
@@ -1692,6 +1692,14 @@ void LowerTypeTestsModule::replaceDirectCalls(Value *Old, Value *New) {
 }
 
 bool LowerTypeTestsModule::lower() {
+  // If only some of the modules were split, we cannot correctly perform
+  // this transformation. We already checked for the presense of type tests
+  // with partially split modules during the thin link, and would have emitted
+  // an error if any were found, so here we can simply return.
+  if ((ExportSummary && ExportSummary->partiallySplitLTOUnits()) ||
+      (ImportSummary && ImportSummary->partiallySplitLTOUnits()))
+    return false;
+
   Function *TypeTestFunc =
       M.getFunction(Intrinsic::getName(Intrinsic::type_test));
   Function *ICallBranchFunnelFunc =
@@ -1701,13 +1709,6 @@ bool LowerTypeTestsModule::lower() {
       !ExportSummary && !ImportSummary)
     return false;
 
-  // If only some of the modules were split, we cannot correctly handle
-  // code that contains type tests.
-  if (TypeTestFunc && !TypeTestFunc->use_empty() &&
-      ((ExportSummary && ExportSummary->partiallySplitLTOUnits()) ||
-       (ImportSummary && ImportSummary->partiallySplitLTOUnits())))
-    report_fatal_error("inconsistent LTO Unit splitting with llvm.type.test");
-
   if (ImportSummary) {
     if (TypeTestFunc) {
       for (auto UI = TypeTestFunc->use_begin(), UE = TypeTestFunc->use_end();
diff --git a/llvm/lib/Transforms/IPO/WholeProgramDevirt.cpp b/llvm/lib/Transforms/IPO/WholeProgramDevirt.cpp
index ab6f0eb660f..6b6dd6194e1 100644
--- a/llvm/lib/Transforms/IPO/WholeProgramDevirt.cpp
+++ b/llvm/lib/Transforms/IPO/WholeProgramDevirt.cpp
@@ -1563,23 +1563,20 @@ void DevirtModule::removeRedundantTypeTests() {
 }
 
 bool DevirtModule::run() {
+  // If only some of the modules were split, we cannot correctly perform
+  // this transformation. We already checked for the presense of type tests
+  // with partially split modules during the thin link, and would have emitted
+  // an error if any were found, so here we can simply return.
+  if ((ExportSummary && ExportSummary->partiallySplitLTOUnits()) ||
+      (ImportSummary && ImportSummary->partiallySplitLTOUnits()))
+    return false;
+
   Function *TypeTestFunc =
       M.getFunction(Intrinsic::getName(Intrinsic::type_test));
   Function *TypeCheckedLoadFunc =
       M.getFunction(Intrinsic::getName(Intrinsic::type_checked_load));
   Function *AssumeFunc = M.getFunction(Intrinsic::getName(Intrinsic::assume));
 
-  // If only some of the modules were split, we cannot correctly handle
-  // code that contains type tests or type checked loads.
-  if ((ExportSummary && ExportSummary->partiallySplitLTOUnits()) ||
-      (ImportSummary && ImportSummary->partiallySplitLTOUnits())) {
-    if ((TypeTestFunc && !TypeTestFunc->use_empty()) ||
-        (TypeCheckedLoadFunc && !TypeCheckedLoadFunc->use_empty()))
-      report_fatal_error("inconsistent LTO Unit splitting with llvm.type.test "
-                         "or llvm.type.checked.load");
-    return false;
-  }
-
   // Normally if there are no users of the devirtualization intrinsics in the
   // module, this pass has nothing to do. But if we are exporting, we also need
   // to handle any users that appear only in the function summaries.