[Assignment Tracking][NFC] Use BitVectors as masks for SmallVectors rather than using DenseMaps

[Orlando]

Summary

Rather than tracking 3 maps of {VariableID: SomeInfo} per block, use a BitVector indexed by VariableID to mask 3 vectors of SomeInfo.

BlockInfos now need to be initialised with a call to init which sets the BitVector width to the number of partially promoted variables in the function and fills the vectors with Top values.

Prior to this patch, in joinBlockInfo, it was necessary to insert Top values into the Join result for variables in A XOR B after joining the variables in A AND B. Now, because the vectors are pre-filled with Top values we need only join the variables A AND B and set the BitVector of tracked variables to A OR B.

The patch achieves an average of 0.25% reduction in instructions retired and a 1.1% max-rss for the CTMark suite in LTO-O3-g builds.

[scott.linder]

It might be worth including the rationale in the description/commit message, like you've done in other similar changes. I assume it is a performance improvement?

llvm/lib/CodeGen/AssignmentTrackingAnalysis.cpp

1001–1008

It's a matter of taste, but I find the algorithms approach easier to read. Alternatively, I think this loop can at least use the set_bits() range rather than find_first+find_next

1012–1016

1033–1036

1236–1244

I think this starts to get a little confusing, as there isn't a direct connection between the two arguments (Mask and M). Maybe they should be adjacent in the parameter list?

The only other suggestion I have is to hide all the data members in BlockInfo. There is already isVariableTracked and getDebugValue, I would just extend this to cover ever operation you need. All of the mask handling will then be adjacent and abstracted away, so the rest of the code reads easier and there is no danger of missing a mask check or update.

I'm leaving a bunch of edit suggestions, but feel free to ignore them! I was just going through the exercise of writing one possible implementation to prove to myself it is an actual improvement, YMMV.

1275–1277

1286

1290–1291

1303–1304

1323–1336

1520–1521

[Orlando]

Thank you @scott.linder for taking a look at this.

It might be worth including the rationale in the description/commit message, like you've done in other similar changes. I assume it is a performance improvement?

Yep there's an average of around 0.25% reduction in instructions retired and a 1.1% max-rss for the CTMark suite in LTO-O3-g builds (I'll update the description).

I'm leaving a bunch of edit suggestions, but feel free to ignore them! I was just going through the exercise of writing one possible implementation to prove to myself it is an actual improvement, YMMV

I appreciate the in-depth review and the suggestions all look good to me!

The diff has become quite large now so I'll could look at trying to factor out some of the suggestions into NFC patches?

[Orlando]

(this time with context)

[scott.linder]

In D145558#4199730, @Orlando wrote:

Thank you @scott.linder for taking a look at this.

It might be worth including the rationale in the description/commit message, like you've done in other similar changes. I assume it is a performance improvement?

Yep there's an average of around 0.25% reduction in instructions retired and a 1.1% max-rss for the CTMark suite in LTO-O3-g builds (I'll update the description).

I'm leaving a bunch of edit suggestions, but feel free to ignore them! I was just going through the exercise of writing one possible implementation to prove to myself it is an actual improvement, YMMV

I appreciate the in-depth review and the suggestions all look good to me!

The diff has become quite large now so I'll could look at trying to factor out some of the suggestions into NFC patches?

That is OK with me, however you prefer to land it! I am also OK with landing the previous version as-is, as I said I don't see any correctness issues

llvm/lib/CodeGen/AssignmentTrackingAnalysis.cpp

 #include "llvm/CodeGen/AssignmentTrackingAnalysis.h"
+#include "llvm/ADT/BitVector.h"
 #include "llvm/ADT/DenseMapInfo.h"
 #include "llvm/ADT/IntervalMap.h"
 #include "llvm/ADT/PostOrderIterator.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/ADT/UniqueVector.h"
 #include "llvm/Analysis/Interval.h"
class FunctionVarLocsBuilder {
   // Use an unordered_map so we don't invalidate iterators after
   // insert/modifications.
   std::unordered_map<const Instruction *, SmallVector<VarLocInfo>>
       VarLocsBeforeInst;
 
   SmallVector<VarLocInfo> SingleLocVars;
 
 public:
+  unsigned getNumVariables() const { return Variables.size(); }
+
   /// Find or insert \p V and return the ID.
   VariableID insertVariable(DebugVariable V) {
     return static_cast<VariableID>(Variables.insert(V));
   }
 
   /// Get a variable from its \p ID.
   const DebugVariable &getVariable(VariableID ID) const {
     return Variables[static_cast<unsigned>(ID)];
struct Assignment {
     } // Can we delete this?
     Assignment(S Status, DIAssignID *ID, DbgAssignIntrinsic *Source)
         : Status(Status), ID(ID), Source(Source) {
       // If the Status is Known then we expect there to be an assignment ID.
       assert(Status == NoneOrPhi || ID);
     }
   };
 
-  using AssignmentMap = DenseMap<VariableID, Assignment>;
-  using LocMap = DenseMap<VariableID, LocKind>;
-  using OverlapMap = DenseMap<VariableID, SmallVector<VariableID, 4>>;
+  using AssignmentMap = SmallVector<Assignment>;
+  using LocMap = SmallVector<LocKind>;
+  using OverlapMap = DenseMap<VariableID, SmallVector<VariableID>>;
   using UntaggedStoreAssignmentMap =
       DenseMap<const Instruction *,
                SmallVector<std::pair<VariableID, at::AssignmentInfo>>>;
 
 private:
+  /// The highest numbered VariableID for partially promoted variables plus 1,
+  /// the values for which start at 1.
+  unsigned TrackedVariablesVectorSize = 0;
   /// Map a variable to the set of variables that it fully contains.
   OverlapMap VarContains;
   /// Map untagged stores to the variable fragments they assign to. Used by
   /// processUntaggedInstruction.
   UntaggedStoreAssignmentMap UntaggedStoreVars;
 
   // Machinery to defer inserting dbg.values.
   using InsertMap = MapVector<Instruction *, SmallVector<VarLocInfo>>;
   InsertMap InsertBeforeMap;
   /// Clear the location definitions currently cached for insertion after /p
   /// After.
   void resetInsertionPoint(Instruction &After);
   void emitDbgValue(LocKind Kind, const DbgVariableIntrinsic *Source,
                     Instruction *After);
 
-  static bool mapsAreEqual(const AssignmentMap &A, const AssignmentMap &B) {
-    if (A.size() != B.size())
-      return false;
-    for (const auto &Pair : A) {
-      VariableID Var = Pair.first;
-      const Assignment &AV = Pair.second;
-      auto R = B.find(Var);
-      // Check if this entry exists in B, otherwise ret false.
-      if (R == B.end())
-        return false;
-      // Check that the assignment value is the same.
-      if (!AV.isSameSourceAssignment(R->second))
-        return false;
-    }
-    return true;
+  static bool mapsAreEqual(const BitVector &Mask, const AssignmentMap &A,
+                           const AssignmentMap &B) {
+    return llvm::all_of(Mask.set_bits(), [&](unsigned VarID) {
+      return A[VarID].isSameSourceAssignment(B[VarID]);
+    });
   }
 
   /// Represents the stack and debug assignments in a block. Used to describe
   /// the live-in and live-out values for blocks, as well as the "current"
   /// value as we process each instruction in a block.
   struct BlockInfo {
-    /// Dominating assignment to memory for each variable.
+    /// The set of variables (VariableID) being tracked in this block.
+    BitVector VariableIDsInBlock;
+    /// Dominating assignment to memory for each variable, indexed by
+    /// VariableID.
     AssignmentMap StackHomeValue;
-    /// Dominating assignemnt to each variable.
+    /// Dominating assignemnt to each variable, indexed by VariableID.
     AssignmentMap DebugValue;
     /// Location kind for each variable. LiveLoc indicates whether the
     /// dominating assignment in StackHomeValue (LocKind::Mem), DebugValue
     /// (LocKind::Val), or neither (LocKind::None) is valid, in that order of
     /// preference. This cannot be derived by inspecting DebugValue and
     /// StackHomeValue due to the fact that there's no distinction in
     /// Assignment (the class) between whether an assignment is unknown or a
     /// merge of multiple assignments (both are Status::NoneOrPhi). In other
     /// words, the memory location may well be valid while both DebugValue and
     /// StackHomeValue contain Assignments that have a Status of NoneOrPhi.
+    /// Indexed by VariableID.
     LocMap LiveLoc;
 
+  public:
+    enum AssignmentKind { Stack, Debug };
+    const AssignmentMap &getAssignmentMap(AssignmentKind Kind) const {
+      switch (Kind) {
+      case Stack:
+        return StackHomeValue;
+      case Debug:
+        return DebugValue;
+      }
+      llvm_unreachable("Unknown AssignmentKind");
+    }
+    AssignmentMap &getAssignmentMap(AssignmentKind Kind) {
+      return const_cast<AssignmentMap &>(
+          const_cast<const BlockInfo *>(this)->getAssignmentMap(Kind));
+    }
+
+    bool isVariableTracked(VariableID Var) const {
+      return VariableIDsInBlock[static_cast<unsigned>(Var)];
+    }
+
+    const Assignment &getAssignment(AssignmentKind Kind, VariableID Var) const {
+      assert(isVariableTracked(Var) && "Var not tracked in block");
+      return getAssignmentMap(Kind)[static_cast<unsigned>(Var)];
+    }
+
+    LocKind getLocKind(VariableID Var) const {
+      assert(isVariableTracked(Var) && "Var not tracked in block");
+      return LiveLoc[static_cast<unsigned>(Var)];
+    }
+
+    /// Set LocKind for \p Var only: does not set LocKind for VariableIDs of
+    /// fragments contained win \p Var.
+    void setLocKind(VariableID Var, LocKind K) {
+      VariableIDsInBlock.set(static_cast<unsigned>(Var));
+      LiveLoc[static_cast<unsigned>(Var)] = K;
+    }
+
+    /// Set the assignment in the \p Kind assignment map for \p Var only: does
+    /// not set the assignment for VariableIDs of fragments contained win \p
+    /// Var.
+    void setAssignment(AssignmentKind Kind, VariableID Var,
+                       const Assignment &AV) {
+      VariableIDsInBlock.set(static_cast<unsigned>(Var));
+      getAssignmentMap(Kind)[static_cast<unsigned>(Var)] = AV;
+    }
+
+    /// Return true if there is an assignment matching \p AV in the \p Kind
+    /// assignment map. Does consider assignments for VariableIDs of fragments
+    /// contained win \p Var.
+    bool hasAssignment(AssignmentKind Kind, VariableID Var,
+                       const Assignment &AV) const {
+      if (!isVariableTracked(Var))
+        return false;
+      return AV.isSameSourceAssignment(getAssignment(Kind, Var));
+    }
+
     /// Compare every element in each map to determine structural equality
     /// (slow).
     bool operator==(const BlockInfo &Other) const {
-      return LiveLoc == Other.LiveLoc &&
-             mapsAreEqual(StackHomeValue, Other.StackHomeValue) &&
-             mapsAreEqual(DebugValue, Other.DebugValue);
+      return VariableIDsInBlock == Other.VariableIDsInBlock &&
+             LiveLoc == Other.LiveLoc &&
+             mapsAreEqual(VariableIDsInBlock, StackHomeValue,
+                          Other.StackHomeValue) &&
+             mapsAreEqual(VariableIDsInBlock, DebugValue, Other.DebugValue);
     }
     bool operator!=(const BlockInfo &Other) const { return !(*this == Other); }
     bool isValid() {
       return LiveLoc.size() == DebugValue.size() &&
              LiveLoc.size() == StackHomeValue.size();
     }
+
+    /// Clear everything and initialise with ⊤-values for all variables.
+    void init(int NumVars) {
+      StackHomeValue.clear();
+      DebugValue.clear();
+      LiveLoc.clear();
+      VariableIDsInBlock = BitVector(NumVars);
+      StackHomeValue.insert(StackHomeValue.begin(), NumVars,
+                            Assignment::makeNoneOrPhi());
+      DebugValue.insert(DebugValue.begin(), NumVars,
+                        Assignment::makeNoneOrPhi());
+      LiveLoc.insert(LiveLoc.begin(), NumVars, LocKind::None);
+    }
+
+    /// Helper for join.
+    template <typename ElmtType, typename FnInputType>
+    static void joinElmt(int Index, SmallVector<ElmtType> &Target,
+                         const SmallVector<ElmtType> &A,
+                         const SmallVector<ElmtType> &B,
+                         ElmtType (*Fn)(FnInputType, FnInputType)) {
+      Target[Index] = Fn(A[Index], B[Index]);
+    }
+
+    /// See comment for AssignmentTrackingLowering::joinBlockInfo.
+    static BlockInfo join(const BlockInfo &A, const BlockInfo &B, int NumVars) {
+      // Join A and B.
+      //
+      // Intersect = join(a, b) for a in A, b in B where Var(a) == Var(b)
+      // Difference = join(x, ⊤) for x where Var(x) is in A xor B
+      // Join = Intersect ∪ Difference
+      //
+      // This is achieved by performing a join on elements from A and B with
+      // variables common to both A and B (join elements indexed by var
+      // intersect), then adding ⊤-value elements for vars in A xor B. The
+      // latter part is equivalent to performing join on elements with variables
+      // in A xor B with the ⊤-value for the map element since join(x, ⊤) = ⊤.
+      // BlockInfo::init initializes all variable entries to the ⊤ value so we
+      // don't need to explicitly perform that step as Join.VariableIDsInBlock
+      // is set to the union of the variables in A and B at the end of this
+      // function.
+      BlockInfo Join;
+      Join.init(NumVars);
+
+      BitVector Intersect = A.VariableIDsInBlock;
+      Intersect &= B.VariableIDsInBlock;
+
+      for (auto VarID : Intersect.set_bits()) {
+        joinElmt(VarID, Join.LiveLoc, A.LiveLoc, B.LiveLoc, joinKind);
+        joinElmt(VarID, Join.DebugValue, A.DebugValue, B.DebugValue,
+                 joinAssignment);
+        joinElmt(VarID, Join.StackHomeValue, A.StackHomeValue, B.StackHomeValue,
+                 joinAssignment);
+      }
+
+      Join.VariableIDsInBlock = A.VariableIDsInBlock;
+      Join.VariableIDsInBlock |= B.VariableIDsInBlock;
+      assert(Join.isValid());
+      return Join;
+    }
   };
 
   Function &Fn;
   const DataLayout &Layout;
   const DenseSet<DebugAggregate> *VarsWithStackSlot;
   FunctionVarLocsBuilder *FnVarLocs;
   DenseMap<const BasicBlock *, BlockInfo> LiveIn;
   DenseMap<const BasicBlock *, BlockInfo> LiveOut;
private:
   /// For the map-type functions, all unmapped keys in an empty map are
   /// associated with a bottom value (⊥). This represents their values being
   /// unknown. Unmapped keys in non-empty maps (joining two maps with a key
   /// only present in one) represents either a variable going out of scope or
   /// dropped debug info. It is assumed the key is associated with a top value
   /// (⊤) in this case (unknown location / assignment).
   ///@{
   static LocKind joinKind(LocKind A, LocKind B);
-  static LocMap joinLocMap(const LocMap &A, const LocMap &B);
   static Assignment joinAssignment(const Assignment &A, const Assignment &B);
-  static AssignmentMap joinAssignmentMap(const AssignmentMap &A,
-                                         const AssignmentMap &B);
-  static BlockInfo joinBlockInfo(const BlockInfo &A, const BlockInfo &B);
+  BlockInfo joinBlockInfo(const BlockInfo &A, const BlockInfo &B);
   ///@}
 
   /// Process the instructions in \p BB updating \p LiveSet along the way. \p
   /// LiveSet must be initialized with the current live-in locations before
   /// calling this.
   void process(BasicBlock &BB, BlockInfo *LiveSet);
   ///@name processMethods
   /// Methods to process instructions in order to update the LiveSet (current
private:
   ///@}
 
   /// Set the LocKind for \p Var.
   void setLocKind(BlockInfo *LiveSet, VariableID Var, LocKind K);
   /// Get the live LocKind for a \p Var. Requires addMemDef or addDbgDef to
   /// have been called for \p Var first.
   LocKind getLocKind(BlockInfo *LiveSet, VariableID Var);
   /// Return true if \p Var has an assignment in \p M matching \p AV.
-  bool hasVarWithAssignment(VariableID Var, const Assignment &AV,
-                            const AssignmentMap &M);
+  bool hasVarWithAssignment(BlockInfo *LiveSet, BlockInfo::AssignmentKind Kind,
+                            VariableID Var, const Assignment &AV);
+  /// Return the set of VariableIDs corresponding the fragments contained fully
+  /// within the variable/fragment \p Var.
+  ArrayRef<VariableID> getContainedFragments(VariableID Var) const;
+
+  /// Mark \p Var as having been touched this frame. Note, this applies only
+  /// to the exact fragment \p Var and not to any fragments contained within.
+  void touchFragment(VariableID Var);
 
   /// Emit info for variables that are fully promoted.
   bool emitPromotedVarLocs(FunctionVarLocsBuilder *FnVarLocs);
 
 public:
   AssignmentTrackingLowering(Function &Fn, const DataLayout &Layout,
                              const DenseSet<DebugAggregate> *VarsWithStackSlot)
       : Fn(Fn), Layout(Layout), VarsWithStackSlot(VarsWithStackSlot) {}
   /// Run the analysis, adding variable location info to \p FnVarLocs. Returns
   /// true if any variable locations have been added to FnVarLocs.
   bool run(FunctionVarLocsBuilder *FnVarLocs);
 };
 } // namespace
 
+ArrayRef<VariableID>
+AssignmentTrackingLowering::getContainedFragments(VariableID Var) const {
+  auto R = VarContains.find(Var);
+  if (R == VarContains.end())
+    return std::nullopt;
+  return R->second;
+}
+
+void AssignmentTrackingLowering::touchFragment(VariableID Var) {
+  VarsTouchedThisFrame.insert(Var);
+}
+
 void AssignmentTrackingLowering::setLocKind(BlockInfo *LiveSet, VariableID Var,
                                             LocKind K) {
   auto SetKind = [this](BlockInfo *LiveSet, VariableID Var, LocKind K) {
-    VarsTouchedThisFrame.insert(Var);
-    LiveSet->LiveLoc[Var] = K;
+    LiveSet->setLocKind(Var, K);
+    touchFragment(Var);
   };
   SetKind(LiveSet, Var, K);
 
   // Update the LocKind for all fragments contained within Var.
-  for (VariableID Frag : VarContains[Var])
+  for (VariableID Frag : getContainedFragments(Var))
     SetKind(LiveSet, Frag, K);
 }
 
 AssignmentTrackingLowering::LocKind
 AssignmentTrackingLowering::getLocKind(BlockInfo *LiveSet, VariableID Var) {
-  auto Pair = LiveSet->LiveLoc.find(Var);
-  assert(Pair != LiveSet->LiveLoc.end());
-  return Pair->second;
+  return LiveSet->getLocKind(Var);
 }
 
 void AssignmentTrackingLowering::addMemDef(BlockInfo *LiveSet, VariableID Var,
                                            const Assignment &AV) {
-  auto AddDef = [](BlockInfo *LiveSet, VariableID Var, Assignment AV) {
-    LiveSet->StackHomeValue[Var] = AV;
-    // Add default (Var -> ⊤) to DebugValue if Var isn't in DebugValue yet.
-    LiveSet->DebugValue.insert({Var, Assignment::makeNoneOrPhi()});
-    // Add default (Var -> ⊤) to LiveLocs if Var isn't in LiveLocs yet. Callers
-    // of addMemDef will call setLocKind to override.
-    LiveSet->LiveLoc.insert({Var, LocKind::None});
-  };
-  AddDef(LiveSet, Var, AV);
+  LiveSet->setAssignment(BlockInfo::Stack, Var, AV);
 
   // Use this assigment for all fragments contained within Var, but do not
   // provide a Source because we cannot convert Var's value to a value for the
   // fragment.
   Assignment FragAV = AV;
   FragAV.Source = nullptr;
-  for (VariableID Frag : VarContains[Var])
-    AddDef(LiveSet, Frag, FragAV);
+  for (VariableID Frag : getContainedFragments(Var))
+    LiveSet->setAssignment(BlockInfo::Stack, Frag, FragAV);
 }
 
 void AssignmentTrackingLowering::addDbgDef(BlockInfo *LiveSet, VariableID Var,
                                            const Assignment &AV) {
-  auto AddDef = [](BlockInfo *LiveSet, VariableID Var, Assignment AV) {
-    LiveSet->DebugValue[Var] = AV;
-    // Add default (Var -> ⊤) to StackHome if Var isn't in StackHome yet.
-    LiveSet->StackHomeValue.insert({Var, Assignment::makeNoneOrPhi()});
-    // Add default (Var -> ⊤) to LiveLocs if Var isn't in LiveLocs yet. Callers
-    // of addDbgDef will call setLocKind to override.
-    LiveSet->LiveLoc.insert({Var, LocKind::None});
-  };
-  AddDef(LiveSet, Var, AV);
+  LiveSet->setAssignment(BlockInfo::Debug, Var, AV);
 
   // Use this assigment for all fragments contained within Var, but do not
   // provide a Source because we cannot convert Var's value to a value for the
   // fragment.
   Assignment FragAV = AV;
   FragAV.Source = nullptr;
-  for (VariableID Frag : VarContains[Var])
-    AddDef(LiveSet, Frag, FragAV);
+  for (VariableID Frag : getContainedFragments(Var))
+    LiveSet->setAssignment(BlockInfo::Debug, Frag, FragAV);
 }
 
 static DIAssignID *getIDFromInst(const Instruction &I) {
   return cast<DIAssignID>(I.getMetadata(LLVMContext::MD_DIAssignID));
 }
 
 static DIAssignID *getIDFromMarker(const DbgAssignIntrinsic &DAI) {
   return cast<DIAssignID>(DAI.getAssignID());
 }
 
 /// Return true if \p Var has an assignment in \p M matching \p AV.
-bool AssignmentTrackingLowering::hasVarWithAssignment(VariableID Var,
-                                                      const Assignment &AV,
-                                                      const AssignmentMap &M) {
-  auto AssignmentIsMapped = [](VariableID Var, const Assignment &AV,
-                               const AssignmentMap &M) {
-    auto R = M.find(Var);
-    if (R == M.end())
-      return false;
-    return AV.isSameSourceAssignment(R->second);
-  };
-
-  if (!AssignmentIsMapped(Var, AV, M))
+bool AssignmentTrackingLowering::hasVarWithAssignment(
+    BlockInfo *LiveSet, BlockInfo::AssignmentKind Kind, VariableID Var,
+    const Assignment &AV) {
+  if (!LiveSet->hasAssignment(Kind, Var, AV))
     return false;
 
   // Check all the frags contained within Var as these will have all been
   // mapped to AV at the last store to Var.
-  for (VariableID Frag : VarContains[Var])
-    if (!AssignmentIsMapped(Frag, AV, M))
+  for (VariableID Frag : getContainedFragments(Var))
+    if (!LiveSet->hasAssignment(Kind, Frag, AV))
       return false;
   return true;
 }
 
 #ifndef NDEBUG
 const char *locStr(AssignmentTrackingLowering::LocKind Loc) {
   using LocKind = AssignmentTrackingLowering::LocKind;
   switch (Loc) {
   case LocKind::Val:
     return "Val";
   case LocKind::Mem:
for (DbgAssignIntrinsic *DAI : Linked) {
     addMemDef(LiveSet, Var, AV);
 
     LLVM_DEBUG(dbgs() << "   linked to " << *DAI << "\n");
     LLVM_DEBUG(dbgs() << "   LiveLoc " << locStr(getLocKind(LiveSet, Var))
                       << " -> ");
 
     // The last assignment to the stack is now AV. Check if the last debug
     // assignment has a matching Assignment.
-    if (hasVarWithAssignment(Var, AV, LiveSet->DebugValue)) {
+    if (hasVarWithAssignment(LiveSet, BlockInfo::Debug, Var, AV)) {
       // The StackHomeValue and DebugValue for this variable match so we can
       // emit a stack home location here.
       LLVM_DEBUG(dbgs() << "Mem, Stack matches Debug program\n";);
       LLVM_DEBUG(dbgs() << "   Stack val: "; AV.dump(dbgs()); dbgs() << "\n");
       LLVM_DEBUG(dbgs() << "   Debug val: ";
-                 LiveSet->DebugValue[Var].dump(dbgs()); dbgs() << "\n");
+                 LiveSet->DebugValue[static_cast<unsigned>(Var)].dump(dbgs());
+                 dbgs() << "\n");
       setLocKind(LiveSet, Var, LocKind::Mem);
       emitDbgValue(LocKind::Mem, DAI, &I);
       continue;
     }
 
     // The StackHomeValue and DebugValue for this variable do not match. I.e.
     // The value currently stored in the stack is not what we'd expect to
     // see, so we cannot use emit a stack home location here. Now we will
     // look at the live LocKind for the variable and determine an appropriate
     // dbg.value to emit.
     LocKind PrevLoc = getLocKind(LiveSet, Var);
     switch (PrevLoc) {
     case LocKind::Val: {
       // The value in memory in memory has changed but we're not currently
       // using the memory location. Do nothing.
       LLVM_DEBUG(dbgs() << "Val, (unchanged)\n";);
       setLocKind(LiveSet, Var, LocKind::Val);
     } break;
     case LocKind::Mem: {
       // There's been an assignment to memory that we were using as a
       // location for this variable, and the Assignment doesn't match what
       // we'd expect to see in memory.
-      if (LiveSet->DebugValue[Var].Status == Assignment::NoneOrPhi) {
+      Assignment DbgAV = LiveSet->getAssignment(BlockInfo::Debug, Var);
+      if (DbgAV.Status == Assignment::NoneOrPhi) {
         // We need to terminate any previously open location now.
         LLVM_DEBUG(dbgs() << "None, No Debug value available\n";);
         setLocKind(LiveSet, Var, LocKind::None);
         emitDbgValue(LocKind::None, DAI, &I);
       } else {
         // The previous DebugValue Value can be used here.
         LLVM_DEBUG(dbgs() << "Val, Debug value is Known\n";);
         setLocKind(LiveSet, Var, LocKind::Val);
-        Assignment PrevAV = LiveSet->DebugValue.lookup(Var);
-        if (PrevAV.Source) {
-          emitDbgValue(LocKind::Val, PrevAV.Source, &I);
+        if (DbgAV.Source) {
+          emitDbgValue(LocKind::Val, DbgAV.Source, &I);
         } else {
           // PrevAV.Source is nullptr so we must emit undef here.
           emitDbgValue(LocKind::None, DAI, &I);
         }
       }
     } break;
     case LocKind::None: {
       // There's been an assignment to memory and we currently are
void AssignmentTrackingLowering::processDbgAssign(DbgAssignIntrinsic &DAI,
   addDbgDef(LiveSet, Var, AV);
 
   LLVM_DEBUG(dbgs() << "processDbgAssign on " << DAI << "\n";);
   LLVM_DEBUG(dbgs() << "   LiveLoc " << locStr(getLocKind(LiveSet, Var))
                     << " -> ");
 
   // Check if the DebugValue and StackHomeValue both hold the same
   // Assignment.
-  if (hasVarWithAssignment(Var, AV, LiveSet->StackHomeValue)) {
+  if (hasVarWithAssignment(LiveSet, BlockInfo::Stack, Var, AV)) {
     // They match. We can use the stack home because the debug intrinsics state
     // that an assignment happened here, and we know that specific assignment
     // was the last one to take place in memory for this variable.
     LocKind Kind;
     if (DAI.isKillAddress()) {
       LLVM_DEBUG(
           dbgs()
               << "Val, Stack matches Debug program but address is killed\n";);

 
 AssignmentTrackingLowering::LocKind
 AssignmentTrackingLowering::joinKind(LocKind A, LocKind B) {
   // Partial order:
   // None > Mem, Val
   return A == B ? A : LocKind::None;
 }
 
-AssignmentTrackingLowering::LocMap
-AssignmentTrackingLowering::joinLocMap(const LocMap &A, const LocMap &B) {
-  // Join A and B.
-  //
-  // U = join(a, b) for a in A, b in B where Var(a) == Var(b)
-  // D = join(x, ⊤) for x where Var(x) is in A xor B
-  // Join = U ∪ D
-  //
-  // This is achieved by performing a join on elements from A and B with
-  // variables common to both A and B (join elements indexed by var intersect),
-  // then adding LocKind::None elements for vars in A xor B. The latter part is
-  // equivalent to performing join on elements with variables in A xor B with
-  // LocKind::None (⊤) since join(x, ⊤) = ⊤.
-  LocMap Join(std::max(A.size(), B.size()));
-  SmallVector<VariableID, 16> SymmetricDifference;
-  // Insert the join of the elements with common vars into Join. Add the
-  // remaining elements to into SymmetricDifference.
-  for (const auto &[Var, Loc] : A) {
-    // If this Var doesn't exist in B then add it to the symmetric difference
-    // set.
-    auto R = B.find(Var);
-    if (R == B.end()) {
-      SymmetricDifference.push_back(Var);
-      continue;
-    }
-    // There is an entry for Var in both, join it.
-    Join[Var] = joinKind(Loc, R->second);
-  }
-  unsigned IntersectSize = Join.size();
-  (void)IntersectSize;
-
-  // Check if A and B contain the same variables.
-  if (SymmetricDifference.empty() && A.size() == B.size())
-    return Join;
-
-  // Add the elements in B with variables that are not in A into
-  // SymmetricDifference.
-  for (const auto &Pair : B) {
-    VariableID Var = Pair.first;
-    if (A.count(Var) == 0)
-      SymmetricDifference.push_back(Var);
-  }
-
-  // Add SymmetricDifference elements to Join and return the result.
-  for (const auto &Var : SymmetricDifference)
-    Join.insert({Var, LocKind::None});
-
-  assert(Join.size() == (IntersectSize + SymmetricDifference.size()));
-  assert(Join.size() >= A.size() && Join.size() >= B.size());
-  return Join;
-}
-
 AssignmentTrackingLowering::Assignment
 AssignmentTrackingLowering::joinAssignment(const Assignment &A,
                                            const Assignment &B) {
   // Partial order:
   // NoneOrPhi(null, null) > Known(v, ?s)
 
   // If either are NoneOrPhi the join is NoneOrPhi.
   // If either value is different then the result is
auto JoinSource = [&]() -> DbgAssignIntrinsic * {
     return nullptr;
   };
   DbgAssignIntrinsic *Source = JoinSource();
   assert(A.Status == B.Status && A.Status == Assignment::Known);
   assert(A.ID == B.ID);
   return Assignment::make(A.ID, Source);
 }
 
-AssignmentTrackingLowering::AssignmentMap
-AssignmentTrackingLowering::joinAssignmentMap(const AssignmentMap &A,
-                                              const AssignmentMap &B) {
-  // Join A and B.
-  //
-  // U = join(a, b) for a in A, b in B where Var(a) == Var(b)
-  // D = join(x, ⊤) for x where Var(x) is in A xor B
-  // Join = U ∪ D
-  //
-  // This is achieved by performing a join on elements from A and B with
-  // variables common to both A and B (join elements indexed by var intersect),
-  // then adding LocKind::None elements for vars in A xor B. The latter part is
-  // equivalent to performing join on elements with variables in A xor B with
-  // Status::NoneOrPhi (⊤) since join(x, ⊤) = ⊤.
-  AssignmentMap Join(std::max(A.size(), B.size()));
-  SmallVector<VariableID, 16> SymmetricDifference;
-  // Insert the join of the elements with common vars into Join. Add the
-  // remaining elements to into SymmetricDifference.
-  for (const auto &[Var, AV] : A) {
-    // If this Var doesn't exist in B then add it to the symmetric difference
-    // set.
-    auto R = B.find(Var);
-    if (R == B.end()) {
-      SymmetricDifference.push_back(Var);
-      continue;
-    }
-    // There is an entry for Var in both, join it.
-    Join[Var] = joinAssignment(AV, R->second);
-  }
-  unsigned IntersectSize = Join.size();
-  (void)IntersectSize;
-
-  // Check if A and B contain the same variables.
-  if (SymmetricDifference.empty() && A.size() == B.size())
-    return Join;
-
-  // Add the elements in B with variables that are not in A into
-  // SymmetricDifference.
-  for (const auto &Pair : B) {
-    VariableID Var = Pair.first;
-    if (A.count(Var) == 0)
-      SymmetricDifference.push_back(Var);
-  }
-
-  // Add SymmetricDifference elements to Join and return the result.
-  for (auto Var : SymmetricDifference)
-    Join.insert({Var, Assignment::makeNoneOrPhi()});
-
-  assert(Join.size() == (IntersectSize + SymmetricDifference.size()));
-  assert(Join.size() >= A.size() && Join.size() >= B.size());
-  return Join;
-}
-
 AssignmentTrackingLowering::BlockInfo
 AssignmentTrackingLowering::joinBlockInfo(const BlockInfo &A,
                                           const BlockInfo &B) {
-  BlockInfo Join;
-  Join.LiveLoc = joinLocMap(A.LiveLoc, B.LiveLoc);
-  Join.StackHomeValue = joinAssignmentMap(A.StackHomeValue, B.StackHomeValue);
-  Join.DebugValue = joinAssignmentMap(A.DebugValue, B.DebugValue);
-  assert(Join.isValid());
-  return Join;
+  return BlockInfo::join(A, B, TrackedVariablesVectorSize);
 }
 
 bool AssignmentTrackingLowering::join(
     const BasicBlock &BB, const SmallPtrSet<BasicBlock *, 16> &Visited) {
   BlockInfo BBLiveIn;
   bool FirstJoin = true;
   // LiveIn locs for BB is the join of the already-processed preds' LiveOut
   // locs.
for (auto I = pred_begin(&BB), E = pred_end(&BB); I != E; I++) {
     // Perform the join of BBLiveIn (current live-in info) and PrevLiveOut.
     if (FirstJoin)
       BBLiveIn = PredLiveOut->second;
     else
       BBLiveIn = joinBlockInfo(std::move(BBLiveIn), PredLiveOut->second);
     FirstJoin = false;
   }
 
+  if (FirstJoin)
+    BBLiveIn.init(TrackedVariablesVectorSize);
+
   auto CurrentLiveInEntry = LiveIn.find(&BB);
   // Check if there isn't an entry, or there is but the LiveIn set has changed
   // (expensive check).
   if (CurrentLiveInEntry == LiveIn.end() ||
       BBLiveIn != CurrentLiveInEntry->second) {
     LiveIn[&BB] = std::move(BBLiveIn);
     // A change has occured.
     return true;
getUntaggedStoreAssignmentInfo(const Instruction &I, const DataLayout &Layout) {
   return std::nullopt;
 }
 
 /// Build a map of {Variable x: Variables y} where all variable fragments
 /// contained within the variable fragment x are in set y. This means that
 /// y does not contain all overlaps because partial overlaps are excluded.
 ///
 /// While we're iterating over the function, add single location defs for
-/// dbg.declares to \p FnVarLocs
+/// dbg.declares to \p FnVarLocs.
+///
+/// Variables that are interesting to this pass in are added to
+/// FnVarLocs->Variables first. TrackedVariablesVectorSize is set to the ID of
+/// the last interesting variable plus 1, meaning variables with ID 1
+/// (inclusive) to TrackedVariablesVectorSize (exclusive) are interesting. The
+/// subsequent variables are either stack homed or fully promoted.
 ///
 /// Finally, populate UntaggedStoreVars with a mapping of untagged stores to
 /// the stored-to variable fragments.
 ///
 /// These tasks are bundled together to reduce the number of times we need
 /// to iterate over the function as they can be achieved together in one pass.
 static AssignmentTrackingLowering::OverlapMap buildOverlapMapAndRecordDeclares(
     Function &Fn, FunctionVarLocsBuilder *FnVarLocs,
-    AssignmentTrackingLowering::UntaggedStoreAssignmentMap &UntaggedStoreVars) {
+    AssignmentTrackingLowering::UntaggedStoreAssignmentMap &UntaggedStoreVars,
+    unsigned &TrackedVariablesVectorSize) {
   DenseSet<DebugVariable> Seen;
   // Map of Variable: [Fragments].
   DenseMap<DebugAggregate, SmallVector<DebugVariable, 8>> FragmentMap;
   // Iterate over all instructions:
   // - dbg.declare    -> add single location variable record
   // - dbg.*          -> Add fragments to FragmentMap
   // - untagged store -> Add fragments to FragmentMap and update
   //                     UntaggedStoreVars.
   // We need to add fragments for untagged stores too so that we can correctly
   // clobber overlapped fragment locations later.
+  SmallVector<DbgDeclareInst *> Declares;
   for (auto &BB : Fn) {
     for (auto &I : BB) {
       if (auto *DDI = dyn_cast<DbgDeclareInst>(&I)) {
-        FnVarLocs->addSingleLocVar(DebugVariable(DDI), DDI->getExpression(),
-                                   DDI->getDebugLoc(),
-                                   DDI->getWrappedLocation());
+        Declares.push_back(DDI);
       } else if (auto *DII = dyn_cast<DbgVariableIntrinsic>(&I)) {
         DebugVariable DV = DebugVariable(DII);
         DebugAggregate DA = {DV.getVariable(), DV.getInlinedAt()};
         if (Seen.insert(DV).second)
           FragmentMap[DA].push_back(DV);
       } else if (auto Info = getUntaggedStoreAssignmentInfo(
                      I, Fn.getParent()->getDataLayout())) {
         // Find markers linked to this alloca.
for (auto It = Frags.begin(), IEnd = Frags.end(); It != IEnd; ++It) {
         DIExpression::FragmentInfo OtherFrag = OtherIt->getFragmentOrDefault();
         VariableID OtherVar = FnVarLocs->insertVariable(*OtherIt);
         if (fullyContains(OtherFrag, Frag))
           Map[OtherVar].push_back(ThisVar);
       }
     }
   }
 
+  // VariableIDs are 1-based so the variable-tracking bitvector needs
+  // NumVariables plus 1 bits.
+  TrackedVariablesVectorSize = FnVarLocs->getNumVariables() + 1;
+
+  // Finally, insert the declares afterwards, so the first IDs are all
+  // partially stack homed vars.
+  for (auto *DDI : Declares)
+    FnVarLocs->addSingleLocVar(DebugVariable(DDI), DDI->getExpression(),
+                               DDI->getDebugLoc(), DDI->getWrappedLocation());
   return Map;
 }
 
 bool AssignmentTrackingLowering::run(FunctionVarLocsBuilder *FnVarLocsBuilder) {
   if (Fn.size() > MaxNumBlocks) {
     LLVM_DEBUG(dbgs() << "[AT] Dropping var locs in: " << Fn.getName()
                       << ": too many blocks (" << Fn.size() << ")\n");
     at::deleteAll(&Fn);
     return false;
   }
 
   FnVarLocs = FnVarLocsBuilder;
 
   // The general structure here is inspired by VarLocBasedImpl.cpp
   // (LiveDebugValues).
 
   // Build the variable fragment overlap map.
   // Note that this pass doesn't handle partial overlaps correctly (FWIW
   // neither does LiveDebugVariables) because that is difficult to do and
   // appears to be rare occurance.
-  VarContains =
-      buildOverlapMapAndRecordDeclares(Fn, FnVarLocs, UntaggedStoreVars);
+  VarContains = buildOverlapMapAndRecordDeclares(
+      Fn, FnVarLocs, UntaggedStoreVars, TrackedVariablesVectorSize);
 
   // Prepare for traversal.
   ReversePostOrderTraversal<Function *> RPOT(&Fn);
   std::priority_queue<unsigned int, std::vector<unsigned int>,
                       std::greater<unsigned int>>
       Worklist;
   std::priority_queue<unsigned int, std::vector<unsigned int>,
                       std::greater<unsigned int>>