Suppress ThinLTO indexing for linkstatic library in tests

This extends the scalability fix in https://github.com/bazelbuild/bazel/commit/8c5e290dfab3cab378a9ca107ecdd6267403cd4b to apply to statically linked
libraries linked into dynamically linked tests. Otherwise we are still getting
the O(N*M) LTO backend job behavior for those libraries when invoking ThinLTO
on many targets.

RELNOTES: None
PiperOrigin-RevId: 189200150

2 files changed, 72 insertions, 16 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/rules/cpp/CppLinkActionBuilder.java b/src/main/java/com/google/devtools/build/lib/rules/cpp/CppLinkActionBuilder.java
index 3facaa9a20..317e808da1 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/cpp/CppLinkActionBuilder.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/cpp/CppLinkActionBuilder.java
@@ -486,7 +486,7 @@ public class CppLinkActionBuilder {
    * used as input to the LTO indexing step.
    */
   private static NestedSet<LibraryToLink> computeLtoIndexingUniqueLibraries(
-      NestedSet<LibraryToLink> originalUniqueLibraries) {
+      NestedSet<LibraryToLink> originalUniqueLibraries, boolean includeLinkStaticInLtoIndexing) {
     NestedSetBuilder<LibraryToLink> uniqueLibrariesBuilder = NestedSetBuilder.linkOrder();
     for (LibraryToLink lib : originalUniqueLibraries) {
       if (!lib.containsObjectFiles()) {
@@ -495,6 +495,22 @@ public class CppLinkActionBuilder {
       }
       ImmutableSet.Builder<Artifact> newObjectFilesBuilder = ImmutableSet.builder();
       for (Artifact a : lib.getObjectFiles()) {
+        // If this link includes object files from another library, that library must be
+        // statically linked.
+        if (!includeLinkStaticInLtoIndexing) {
+          Preconditions.checkNotNull(lib.getSharedNonLtoBackends());
+          LtoBackendArtifacts ltoArtifacts = lib.getSharedNonLtoBackends().getOrDefault(a, null);
+          // Either we have a shared LTO artifact, or this wasn't bitcode to start with.
+          Preconditions.checkState(
+              ltoArtifacts != null || !lib.getLtoBitcodeFiles().containsKey(a));
+          if (ltoArtifacts != null) {
+            // Include the native object produced by the shared LTO backend in the LTO indexing
+            // step instead of the bitcode file. The LTO indexing step invokes the linker which
+            // must see all objects used to produce the final link output.
+            newObjectFilesBuilder.add(ltoArtifacts.getObjectFile());
+            continue;
+          }
+        }
         newObjectFilesBuilder.add(lib.getLtoBitcodeFiles().getOrDefault(a, a));
       }
       uniqueLibrariesBuilder.add(
@@ -615,6 +631,7 @@ public class CppLinkActionBuilder {
       PathFragment ltoOutputRootPrefix,
       NestedSet<LibraryToLink> uniqueLibraries,
       boolean allowLtoIndexing,
+      boolean includeLinkStaticInLtoIndexing,
       ImmutableSet<String> features) {
     Set<Artifact> compiled = new LinkedHashSet<>();
     for (LibraryToLink lib : uniqueLibraries) {
@@ -624,13 +641,18 @@ public class CppLinkActionBuilder {
     // This flattens the set of object files, so for M binaries and N .o files,
     // this is O(M*N). If we had a nested set of .o files, we could have O(M + N) instead.
     Map<PathFragment, Artifact> allBitcode = new HashMap<>();
-    for (LibraryToLink lib : uniqueLibraries) {
-      if (!lib.containsObjectFiles()) {
-        continue;
-      }
-      for (Artifact objectFile : lib.getObjectFiles()) {
-        if (compiled.contains(objectFile)) {
-          allBitcode.put(objectFile.getExecPath(), objectFile);
+    // Since this link includes object files from another library, we know that library must be
+    // statically linked, so we need to look at includeLinkStaticInLtoIndexing to decide whether
+    // to include its objects in the LTO indexing for this target.
+    if (includeLinkStaticInLtoIndexing) {
+      for (LibraryToLink lib : uniqueLibraries) {
+        if (!lib.containsObjectFiles()) {
+          continue;
+        }
+        for (Artifact objectFile : lib.getObjectFiles()) {
+          if (compiled.contains(objectFile)) {
+            allBitcode.put(objectFile.getExecPath(), objectFile);
+          }
         }
       }
     }
@@ -650,7 +672,7 @@ public class CppLinkActionBuilder {
       // each target linking this library needs a unique set of LTO backends.
       for (Artifact objectFile : lib.getObjectFiles()) {
         if (compiled.contains(objectFile)) {
-          if (allowLtoIndexing) {
+          if (includeLinkStaticInLtoIndexing) {
             List<String> backendArgv = new ArrayList<>(argv);
             backendArgv.addAll(collectPerFileLtoBackendOpts(objectFile));
             LtoBackendArtifacts ltoArtifacts =
@@ -849,11 +871,15 @@ public class CppLinkActionBuilder {
     // the targest share the dynamic libraries which were produced via smaller subsets of
     // LTO indexing/backends. ThinLTO on the tests will be different than the ThinLTO
     // optimizations applied to the associated main binaries anyway.
-    boolean allowLtoIndexing =
-        linkStaticness == LinkStaticness.DYNAMIC
-            || !(ruleContext.isTestTarget() || ruleContext.isTestOnlyTarget())
+    // Even for dynamically linked tests, disallow linkstatic libraries from participating
+    // in the test's LTO indexing step for similar reasons.
+    boolean includeLinkStaticInLtoIndexing =
+        !(ruleContext.isTestTarget() || ruleContext.isTestOnlyTarget())
             || !featureConfiguration.isEnabled(
                 CppRuleClasses.THIN_LTO_LINKSTATIC_TESTS_USE_SHARED_NONLTO_BACKENDS);
+    boolean allowLtoIndexing =
+        includeLinkStaticInLtoIndexing
+            || (linkStaticness == LinkStaticness.DYNAMIC && !ltoBitcodeFiles.isEmpty());
 
     // ruleContext can only be null during testing. This is kind of ugly.
     final ImmutableSet<String> features =
@@ -874,7 +900,11 @@ public class CppLinkActionBuilder {
       // the LTO indexing step).
       allLtoArtifacts =
           createLtoArtifacts(
-              ltoOutputRootPrefix, originalUniqueLibraries, allowLtoIndexing, features);
+              ltoOutputRootPrefix,
+              originalUniqueLibraries,
+              allowLtoIndexing,
+              includeLinkStaticInLtoIndexing,
+              features);
 
       if (!allowLtoIndexing) {
         return null;
@@ -889,7 +919,9 @@ public class CppLinkActionBuilder {
     ImmutableSet<LinkerInput> linkstampObjectFileInputs;
     if (isLtoIndexing) {
       objectFileInputs = computeLtoIndexingObjectFileInputs();
-      uniqueLibraries = computeLtoIndexingUniqueLibraries(originalUniqueLibraries);
+      uniqueLibraries =
+          computeLtoIndexingUniqueLibraries(
+              originalUniqueLibraries, includeLinkStaticInLtoIndexing);
       linkstampObjectFileInputs = ImmutableSet.of();
     } else {
       objectFileInputs = ImmutableSet.copyOf(objectFiles);
@@ -2148,6 +2180,22 @@ public class CppLinkActionBuilder {
     }
 
     /**
+     * Returns true if this artifact is produced from a bitcode file that will be input to the LTO
+     * indexing step, in which case that step will add it to the generated thinltoParamFile for
+     * inclusion in the final link step if the linker decides to include it.
+     *
+     * @param a is an artifact produced by an LTO backend.
+     */
+    private boolean handledByLtoIndexing(Artifact a) {
+      // If no LTO indexing is allowed for this link, then none are handled by LTO indexing.
+      // Otherwise, this may be from a linkstatic library that we decided not to include in
+      // LTO indexing because we are linking a test, to improve scalability when linking many tests.
+      return allowLtoIndexing
+          && !a.getRootRelativePath()
+              .startsWith(PathFragment.create(SHARED_NONLTO_BACKEND_ROOT_PREFIX));
+    }
+
+    /**
      * Adds command-line options for a static library or non-library input into options.
      *
      * @param librariesToLink - a collection that will be exposed as a build variable.
@@ -2179,7 +2227,7 @@ public class CppLinkActionBuilder {
             if (ltoMap != null && (a = ltoMap.remove(member)) != null) {
               // When ltoMap is non-null the backend artifact may be missing due to libraries that
               // list .o files explicitly, or generate .o files from assembler.
-              if (allowLtoIndexing) {
+              if (handledByLtoIndexing(a)) {
                 // The LTO artifacts that should be included in the final link
                 // are listed in the thinltoParamFile, generated by the LTO indexing.
                 continue;
@@ -2211,7 +2259,7 @@ public class CppLinkActionBuilder {
         Artifact inputArtifact = input.getArtifact();
         Artifact a;
         if (ltoMap != null && (a = ltoMap.remove(inputArtifact)) != null) {
-          if (allowLtoIndexing) {
+          if (handledByLtoIndexing(a)) {
             // The LTO artifacts that should be included in the final link
             // are listed in the thinltoParamFile, generated by the LTO indexing.
             return;
diff --git a/src/main/java/com/google/devtools/build/lib/rules/cpp/LtoBackendArtifacts.java b/src/main/java/com/google/devtools/build/lib/rules/cpp/LtoBackendArtifacts.java
index 8c4fc757ce..5aaf1c1d30 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/cpp/LtoBackendArtifacts.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/cpp/LtoBackendArtifacts.java
@@ -54,12 +54,14 @@ import java.util.Map;
 @AutoCodec
 public final class LtoBackendArtifacts {
   // A file containing mapping of symbol => bitcode file containing the symbol.
+  // It will be null when this is a shared non-lto backend.
   private final Artifact index;
 
   // The bitcode file which is the input of the compile.
   private final Artifact bitcodeFile;
 
   // A file containing a list of bitcode files necessary to run the backend step.
+  // It will be null when this is a shared non-lto backend.
   private final Artifact imports;
 
   // The result of executing the above command line, an ELF object file.
@@ -164,6 +166,12 @@ public final class LtoBackendArtifacts {
   }
 
   public void addIndexingOutputs(ImmutableSet.Builder<Artifact> builder) {
+    // For objects from linkstatic libraries, we may not be including them in the LTO indexing
+    // step when linked into a test, but rather will use shared non-LTO backends for better
+    // scalability when running large numbers of tests.
+    if (index == null) {
+      return;
+    }
     builder.add(imports);
     builder.add(index);
   }