[DWARFLinker] Fix placement lost-update crash in parallel marking (#211009)

llvm/lib/DWARFLinker/Parallel/DWARFLinkerCompileUnit.h

@@ -220,6 +220,36 @@ public:
       return false;
     }
 
+    /// Atomically joins \p Placement into the current placement: the
+    /// least-upper-bound of the lattice NotSet < {TypeTable, PlainDwarf} <
+    /// Both, which is a plain OR because the values are bit flags. The join is
+    /// monotone and never clears a bit, so unlike setPlacement it composes
+    /// correctly when applied concurrently from several marks.
+    void joinPlacement(DieOutputPlacement Placement) {
+      auto InputData = Flags.load();
+      while (!Flags.compare_exchange_weak(InputData, (InputData | Placement))) {
+      }
+    }
+
+    /// Atomically joins \p Placement for a DW_TAG_variable, for which
+    /// PlainDwarf is absorbing because a variable cannot occupy the type table
+    /// and plain DWARF at once. Once the placement is (or concurrently becomes)
+    /// PlainDwarf it stays PlainDwarf, otherwise \p Placement is OR-joined.
+    /// Recomputing inside the compare_exchange loop keeps a racing PlainDwarf
+    /// mark from turning the variable into Both.
+    void joinVariablePlacement(DieOutputPlacement Placement) {
+      auto InputData = Flags.load();
+      uint16_t Desired;
+      do {
+        DieOutputPlacement Current = DieOutputPlacement(InputData & 0x7);
+        DieOutputPlacement Joined =
+            (Current == PlainDwarf || Current == Both)
+                ? PlainDwarf
+                : DieOutputPlacement(Current | Placement);
+        Desired = (InputData & ~0x7) | Joined;
+      } while (!Flags.compare_exchange_weak(InputData, Desired));
+    }
+
 #define SINGLE_FLAG_METHODS_SET(Name, Value)                                   \
   bool get##Name() const { return Flags & Value; }                             \
   void set##Name() {                                                           \

llvm/lib/DWARFLinker/Parallel/DependencyTracker.cpp

@@ -409,19 +409,44 @@ void DependencyTracker::markParentsAsKeepingChildren(
   }
 }
 
-// This function tries to set specified \p Placement for the \p Entry.
-// Depending on the concrete entry, the placement could be:
-//  a) changed to another.
-//  b) joined with current entry placement.
-//  c) set as requested.
-static CompileUnit::DieOutputPlacement
+namespace {
+struct FinalPlacement {
+  CompileUnit::DieOutputPlacement Placement;
+
+  /// How Placement combines with the DIE's current placement when applied.
+  enum ApplyMode {
+    /// Overwrite the current placement. Used for entries whose placement is
+    /// fully determined regardless of how they were reached, so every mark
+    /// agrees on the value (ODR-unavailable entries and static data member
+    /// declarations).
+    Overwrite,
+    /// OR-join into the current placement (the common monotone-lattice case):
+    /// a DIE reached by both a live and a type mark ends up in Both.
+    Join,
+    /// Join for a DW_TAG_variable, which cannot occupy the type table and plain
+    /// DWARF at once: PlainDwarf is absorbing so the variable never lands in
+    /// Both.
+    JoinVariable,
+  } Mode;
+};
+} // namespace
+
+// Computes the placement to apply to \p Entry for a mark requesting \p
+// Placement (PlainDwarf for a live action, TypeTable for a type action), along
+// with how it combines with the DIE's current placement. Most entries join, so
+// a DIE reached by both actions ends up in Both. Entries whose placement is
+// fully determined regardless of how they were reached instead overwrite with
+// an exact placement: ODR-unavailable entries cannot be deduplicated into the
+// type table, and a DW_TAG_variable cannot occupy the type table and plain
+// DWARF at once.
+static FinalPlacement
 getFinalPlacementForEntry(const UnitEntryPairTy &Entry,
                           CompileUnit::DieOutputPlacement Placement) {
   assert((Placement != CompileUnit::NotSet) && "Placement is not set");
   CompileUnit::DIEInfo &EntryInfo = Entry.CU->getDIEInfo(Entry.DieEntry);
 
   if (!EntryInfo.getODRAvailable())
-    return CompileUnit::PlainDwarf;
+    return {CompileUnit::PlainDwarf, FinalPlacement::Overwrite};
 
   if (Entry.DieEntry->getTag() == dwarf::DW_TAG_variable) {
     // In-class static member declarations (e.g. "static constexpr int x = 1;")
@@ -447,35 +472,19 @@ getFinalPlacementForEntry(const UnitEntryPairTy &Entry,
     if (IsDeclaration && ParentIsType) {
       // Pure declarations have no runtime address; they belong with the class
       // type. Always place in TypeTable regardless of how they were reached.
-      return CompileUnit::TypeTable;
+      return {CompileUnit::TypeTable, FinalPlacement::Overwrite};
     }
 
-    // Do not put variable into the "TypeTable" and "PlainDwarf" at the same
-    // time.
-    if (EntryInfo.getPlacement() == CompileUnit::PlainDwarf ||
-        EntryInfo.getPlacement() == CompileUnit::Both)
-      return CompileUnit::PlainDwarf;
-
+    // A live (PlainDwarf) mark pins the variable to plain DWARF.
     if (Placement == CompileUnit::PlainDwarf || Placement == CompileUnit::Both)
-      return CompileUnit::PlainDwarf;
+      return {CompileUnit::PlainDwarf, FinalPlacement::Overwrite};
+
+    // Only a type-table mark reaches here. The variable join keeps a PlainDwarf
+    // mark racing this one from turning the variable into Both.
+    return {Placement, FinalPlacement::JoinVariable};
   }
 
-  switch (EntryInfo.getPlacement()) {
-  case CompileUnit::NotSet:
-    return Placement;
-
-  case CompileUnit::TypeTable:
-    return Placement == CompileUnit::PlainDwarf ? CompileUnit::Both : Placement;
-
-  case CompileUnit::PlainDwarf:
-    return Placement == CompileUnit::TypeTable ? CompileUnit::Both : Placement;
-
-  case CompileUnit::Both:
-    return CompileUnit::Both;
-  };
-
-  llvm_unreachable("Unknown placement type.");
-  return Placement;
+  return {Placement, FinalPlacement::Join};
 }
 
 bool DependencyTracker::markDIEEntryAsKeptRec(
@@ -487,10 +496,11 @@ bool DependencyTracker::markDIEEntryAsKeptRec(
 
   CompileUnit::DIEInfo &Info = Entry.CU->getDIEInfo(Entry.DieEntry);
 
-  // Calculate final placement placement.
-  CompileUnit::DieOutputPlacement Placement = getFinalPlacementForEntry(
+  // Calculate final placement.
+  FinalPlacement Final = getFinalPlacementForEntry(
       Entry,
       isLiveAction(Action) ? CompileUnit::PlainDwarf : CompileUnit::TypeTable);
+  CompileUnit::DieOutputPlacement Placement = Final.Placement;
   assert((Info.getODRAvailable() || isLiveAction(Action) ||
           Placement == CompileUnit::PlainDwarf) &&
          "Wrong kind of placement for ODR unavailable entry");
@@ -515,7 +525,20 @@ bool DependencyTracker::markDIEEntryAsKeptRec(
   if (!RecordDepsOnly) {
     // Mark current DIE as kept.
     Info.setKeep();
-    Info.setPlacement(Placement);
+    // Marks compose monotonically so no interleaving loses an update: a general
+    // mark only raises the placement in the lattice, and a forced placement is
+    // a value every mark agrees on.
+    switch (Final.Mode) {
+    case FinalPlacement::Overwrite:
+      Info.setPlacement(Placement);
+      break;
+    case FinalPlacement::Join:
+      Info.joinPlacement(Placement);
+      break;
+    case FinalPlacement::JoinVariable:
+      Info.joinVariablePlacement(Placement);
+      break;
+    }
 
     // Set keep children property for parents.
     markParentsAsKeepingChildren(Entry);

llvm/unittests/DWARFLinkerParallel/CMakeLists.txt

@@ -6,6 +6,7 @@ set(LLVM_LINK_COMPONENTS
   )
 
 add_llvm_unittest(DWARFLinkerParallelTests
+  DIEInfoPlacementTest.cpp
   DWARFLinkerTest.cpp
   StringPoolTest.cpp
   )
@@ -14,4 +15,10 @@ target_link_libraries(DWARFLinkerParallelTests PRIVATE
   LLVMTestingSupport
   )
 
+# DIEInfoPlacementTest.cpp exercises CompileUnit::DIEInfo, whose definition
+# lives in the library's private headers.
+target_include_directories(DWARFLinkerParallelTests PRIVATE
+  ${LLVM_MAIN_SRC_DIR}/lib/DWARFLinker/Parallel
+  )
+
 add_dependencies(DWARFLinkerParallelTests intrinsics_gen)

llvm/unittests/DWARFLinkerParallel/DIEInfoPlacementTest.cpp

@@ -0,0 +1,104 @@
+//===----------------------------------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "DWARFLinkerCompileUnit.h"
+#include "gtest/gtest.h"
+#include <thread>
+
+using namespace llvm;
+using namespace llvm::dwarf_linker::parallel;
+
+namespace {
+
+using DIEInfo = CompileUnit::DIEInfo;
+
+// joinPlacement is the least-upper-bound over the lattice
+// NotSet < {TypeTable, PlainDwarf} < Both. It only ever raises the placement,
+// so no interleaving of marks can demote a DIE.
+TEST(DIEInfoPlacement, JoinIsMonotoneLatticeLUB) {
+  DIEInfo Info;
+  EXPECT_EQ(Info.getPlacement(), CompileUnit::NotSet);
+
+  Info.joinPlacement(CompileUnit::TypeTable);
+  EXPECT_EQ(Info.getPlacement(), CompileUnit::TypeTable);
+
+  // The other mark raises TypeTable to Both.
+  Info.joinPlacement(CompileUnit::PlainDwarf);
+  EXPECT_EQ(Info.getPlacement(), CompileUnit::Both);
+
+  // Re-applying either mark never demotes.
+  Info.joinPlacement(CompileUnit::TypeTable);
+  Info.joinPlacement(CompileUnit::PlainDwarf);
+  EXPECT_EQ(Info.getPlacement(), CompileUnit::Both);
+}
+
+// joinPlacement leaves the non-placement flag bits untouched.
+TEST(DIEInfoPlacement, JoinPreservesOtherFlags) {
+  DIEInfo Info;
+  Info.setKeep();
+  Info.setODRAvailable();
+  Info.joinPlacement(CompileUnit::TypeTable);
+  Info.joinPlacement(CompileUnit::PlainDwarf);
+  EXPECT_EQ(Info.getPlacement(), CompileUnit::Both);
+  EXPECT_TRUE(Info.getKeep());
+  EXPECT_TRUE(Info.getODRAvailable());
+}
+
+// For a DW_TAG_variable, PlainDwarf is absorbing: a variable cannot occupy the
+// type table and plain DWARF at once, so joining a TypeTable mark into a
+// PlainDwarf placement keeps PlainDwarf rather than producing Both.
+TEST(DIEInfoPlacement, VariableJoinKeepsPlainDwarfAbsorbing) {
+  // TypeTable-then-PlainDwarf: the plain mark forces PlainDwarf (handled by the
+  // overwrite path), and a subsequent type mark must not resurrect TypeTable.
+  {
+    DIEInfo Info;
+    Info.joinVariablePlacement(CompileUnit::TypeTable);
+    EXPECT_EQ(Info.getPlacement(), CompileUnit::TypeTable);
+    Info.setPlacement(CompileUnit::PlainDwarf);
+    Info.joinVariablePlacement(CompileUnit::TypeTable);
+    EXPECT_EQ(Info.getPlacement(), CompileUnit::PlainDwarf);
+  }
+
+  // A stale Both is collapsed back to PlainDwarf, never left as Both.
+  {
+    DIEInfo Info;
+    Info.setPlacement(CompileUnit::Both);
+    Info.joinVariablePlacement(CompileUnit::TypeTable);
+    EXPECT_EQ(Info.getPlacement(), CompileUnit::PlainDwarf);
+  }
+
+  // Only a TypeTable mark ever reaches an unset variable: it lands in
+  // TypeTable.
+  {
+    DIEInfo Info;
+    Info.joinVariablePlacement(CompileUnit::TypeTable);
+    EXPECT_EQ(Info.getPlacement(), CompileUnit::TypeTable);
+  }
+}
+
+// Concurrently applying a PlainDwarf overwrite and a TypeTable variable-join to
+// one DIEInfo must never leave it in Both, regardless of interleaving. A plain
+// joinPlacement (OR) would produce Both here. joinVariablePlacement recomputes
+// on each attempt so PlainDwarf stays absorbing under the race.
+TEST(DIEInfoPlacement, VariableJoinRaceNeverProducesBoth) {
+  for (unsigned Trial = 0; Trial < 2000; ++Trial) {
+    DIEInfo Info;
+    std::thread Plain([&] { Info.setPlacement(CompileUnit::PlainDwarf); });
+    std::thread Type(
+        [&] { Info.joinVariablePlacement(CompileUnit::TypeTable); });
+    Plain.join();
+    Type.join();
+    // The final placement is one of the single values, never Both.
+    CompileUnit::DieOutputPlacement P = Info.getPlacement();
+    EXPECT_TRUE(P == CompileUnit::PlainDwarf || P == CompileUnit::TypeTable)
+        << "placement was " << static_cast<unsigned>(P) << " on trial "
+        << Trial;
+  }
+}
+
+} // namespace