Merge pull request #66690 from nate-chandler/rdar110676577

[FieldSensitivePL] Fix vectorization.

Author: nate-chandler

include/swift/Basic/FrozenMultiMap.h

@@ -71,7 +71,7 @@ class FrozenMultiMap {
     // Since our array is sorted, we need to first find the first pair with our
     // inst as the first element.
     auto start = std::lower_bound(
-        storage.begin(), storage.end(), std::make_pair(key, Value()),
+        storage.begin(), storage.end(), std::make_pair(key, llvm::None),
         [&](const std::pair<Key, Optional<Value>> &p1,
             const std::pair<Key, Optional<Value>> &p2) {
           return p1.first < p2.first;

include/swift/SIL/FieldSensitivePrunedLiveness.h

@@ -358,6 +358,13 @@ struct TypeTreeLeafTypeRange {
            endEltOffset >= range.endEltOffset;
   }
 
+  /// Sets each bit in \p bits corresponding to an element of this range.
+  void setBits(SmallBitVector &bits) {
+    for (auto element : getRange()) {
+      bits.set(element);
+    }
+  }
+
   IntRange<unsigned> getRange() const {
     return range(startEltOffset, endEltOffset);
   }
@@ -666,17 +673,60 @@ class FieldSensitivePrunedLiveness {
   FieldSensitivePrunedLiveBlocks liveBlocks;
 
 public:
+  enum IsInterestingUser { NonUser, NonLifetimeEndingUse, LifetimeEndingUse };
+
   struct InterestingUser {
-    TypeTreeLeafTypeRange subEltSpan;
-    bool isConsuming;
+    SmallBitVector liveBits;
+    SmallBitVector consumingBits;
 
-    InterestingUser() : subEltSpan(), isConsuming(false) {}
-    InterestingUser(TypeTreeLeafTypeRange subEltSpan, bool isConsuming)
-        : subEltSpan(subEltSpan), isConsuming(isConsuming) {}
+    InterestingUser(unsigned bitCount)
+        : liveBits(bitCount), consumingBits(bitCount) {}
 
-    InterestingUser &operator&=(bool otherValue) {
-      isConsuming &= otherValue;
-      return *this;
+    InterestingUser(unsigned bitCount, TypeTreeLeafTypeRange range,
+                    bool lifetimeEnding)
+        : liveBits(bitCount), consumingBits(bitCount) {
+      addUses(range, lifetimeEnding);
+    }
+
+    /// Record that the instruction uses the bits of the value in \p range.
+    void addUses(TypeTreeLeafTypeRange range, bool lifetimeEnding) {
+      range.setBits(liveBits);
+      if (lifetimeEnding) {
+        range.setBits(consumingBits);
+      }
+    }
+
+    /// Populates the provided vector with contiguous ranges of bits which are
+    /// users of the same sort.
+    void getContiguousRanges(
+        SmallVectorImpl<std::pair<TypeTreeLeafTypeRange, IsInterestingUser>>
+            &ranges) const {
+      if (liveBits.size() == 0)
+        return;
+
+      assert(ranges.empty());
+      Optional<std::pair<unsigned, IsInterestingUser>> current = llvm::None;
+      for (unsigned bit = 0, size = liveBits.size(); bit < size; ++bit) {
+        auto interesting = isInterestingUser(bit);
+        if (!current) {
+          current = {bit, interesting};
+          continue;
+        }
+        if (current->second != interesting) {
+          ranges.push_back(
+              {TypeTreeLeafTypeRange(current->first, bit), current->second});
+          current = {bit, interesting};
+        }
+      }
+      ranges.push_back({TypeTreeLeafTypeRange(current->first, liveBits.size()),
+                        current->second});
+    }
+
+    IsInterestingUser isInterestingUser(unsigned element) const {
+      if (!liveBits.test(element))
+        return NonUser;
+      return consumingBits.test(element) ? LifetimeEndingUse
+                                         : NonLifetimeEndingUse;
     }
   };
 
@@ -758,42 +808,6 @@ class FieldSensitivePrunedLiveness {
     return llvm::make_range(users.begin(), users.end());
   }
 
-  using LifetimeEndingUserRange = OptionalTransformRange<
-      UserRange,
-      function_ref<Optional<std::pair<SILInstruction *, TypeTreeLeafTypeRange>>(
-          const std::pair<SILInstruction *, InterestingUser> &)>>;
-  LifetimeEndingUserRange getAllLifetimeEndingUses() const {
-    assert(isInitialized());
-    function_ref<Optional<std::pair<SILInstruction *, TypeTreeLeafTypeRange>>(
-        const std::pair<SILInstruction *, InterestingUser> &)>
-        op;
-    op = [](const std::pair<SILInstruction *, InterestingUser> &pair)
-        -> Optional<std::pair<SILInstruction *, TypeTreeLeafTypeRange>> {
-      if (pair.second.isConsuming)
-        return {{pair.first, pair.second.subEltSpan}};
-      return None;
-    };
-    return LifetimeEndingUserRange(getAllUsers(), op);
-  }
-
-  using NonLifetimeEndingUserRange = OptionalTransformRange<
-      UserRange,
-      function_ref<Optional<std::pair<SILInstruction *, TypeTreeLeafTypeRange>>(
-          const std::pair<SILInstruction *, InterestingUser> &)>>;
-  NonLifetimeEndingUserRange getAllNonLifetimeEndingUses() const {
-    assert(isInitialized());
-    function_ref<Optional<std::pair<SILInstruction *, TypeTreeLeafTypeRange>>(
-        const std::pair<SILInstruction *, InterestingUser> &)>
-        op;
-    op = [](const std::pair<SILInstruction *, InterestingUser> &pair)
-        -> Optional<std::pair<SILInstruction *, TypeTreeLeafTypeRange>> {
-      if (!pair.second.isConsuming)
-        return {{pair.first, pair.second.subEltSpan}};
-      return None;
-    };
-    return NonLifetimeEndingUserRange(getAllUsers(), op);
-  }
-
   using UserBlockRange = TransformRange<
       UserRange, function_ref<SILBasicBlock *(
                      const std::pair<SILInstruction *, InterestingUser> &)>>;
@@ -848,19 +862,37 @@ class FieldSensitivePrunedLiveness {
                         SmallBitVector &liveOutBits,
                         SmallBitVector &deadBits) const;
 
-  enum IsInterestingUser { NonUser, NonLifetimeEndingUse, LifetimeEndingUse };
+  /// If \p user has had uses recored, return a pointer to the InterestingUser
+  /// where they've been recorded.
+  InterestingUser const *getInterestingUser(SILInstruction *user) const {
+    auto iter = users.find(user);
+    if (iter == users.end())
+      return nullptr;
+    return &iter->second;
+  }
 
-  /// Return a result indicating whether the given user was identified as an
-  /// interesting use of the current def and whether it ends the lifetime.
-  std::pair<IsInterestingUser, Optional<TypeTreeLeafTypeRange>>
-  isInterestingUser(SILInstruction *user) const {
+  /// How \p user uses the field at \p element.
+  IsInterestingUser isInterestingUser(SILInstruction *user,
+                                      unsigned element) const {
     assert(isInitialized());
-    auto useIter = users.find(user);
-    if (useIter == users.end())
-      return {NonUser, None};
-    auto isInteresting =
-        useIter->second.isConsuming ? LifetimeEndingUse : NonLifetimeEndingUse;
-    return {isInteresting, useIter->second.subEltSpan};
+    auto *record = getInterestingUser(user);
+    if (!record)
+      return NonUser;
+    return record->isInterestingUser(element);
+  }
+
+  /// Whether \p user uses the fields in \p range as indicated by \p kind.
+  bool isInterestingUserOfKind(SILInstruction *user, IsInterestingUser kind,
+                               TypeTreeLeafTypeRange range) const {
+    auto *record = getInterestingUser(user);
+    if (!record)
+      return kind == IsInterestingUser::NonUser;
+
+    for (auto element : range.getRange()) {
+      if (isInterestingUser(user, element) != kind)
+        return false;
+    }
+    return true;
   }
 
   unsigned getNumSubElements() const { return liveBlocks.getNumBitsToTrack(); }
@@ -886,10 +918,11 @@ class FieldSensitivePrunedLiveness {
   /// argument must be copied.
   void addInterestingUser(SILInstruction *user, TypeTreeLeafTypeRange range,
                           bool lifetimeEnding) {
-    auto iterAndSuccess =
-        users.insert({user, InterestingUser(range, lifetimeEnding)});
-    if (!iterAndSuccess.second)
-      iterAndSuccess.first->second &= lifetimeEnding;
+    auto iter = users.find(user);
+    if (iter == users.end()) {
+      iter = users.insert({user, InterestingUser(getNumSubElements())}).first;
+    }
+    iter->second.addUses(range, lifetimeEnding);
   }
 };
 

lib/SIL/Utils/FieldSensitivePrunedLiveness.cpp

@@ -563,6 +563,38 @@ void FieldSensitivePrunedLiveBlocks::print(llvm::raw_ostream &OS) const {
 
 void FieldSensitivePrunedLiveBlocks::dump() const { print(llvm::dbgs()); }
 
+//===----------------------------------------------------------------------===//
+//                   FieldSensitivePrunedLivenessBoundary
+//===----------------------------------------------------------------------===//
+
+void FieldSensitivePrunedLivenessBoundary::print(llvm::raw_ostream &OS) const {
+  for (auto pair : lastUsers) {
+    auto *user = pair.first;
+    auto bits = pair.second;
+    OS << "last user: " << *user 
+       << "\tat " << bits << "\n";
+  }
+  for (auto pair : boundaryEdges) {
+    auto *block = pair.first;
+    auto bits = pair.second;
+    OS << "boundary edge: ";
+    block->printAsOperand(OS);
+    OS << "\n" << "\tat " << bits << "\n";
+  }
+  if (!deadDefs.empty()) {
+    for (auto pair : deadDefs) {
+      auto *deadDef = pair.first;
+      auto bits = pair.second;
+      OS << "dead def: " << *deadDef 
+         << "\tat " << bits << "\n";
+    }
+  }
+}
+
+void FieldSensitivePrunedLivenessBoundary::dump() const {
+  print(llvm::dbgs());
+}
+
 //===----------------------------------------------------------------------===//
 //                        MARK: FieldSensitiveLiveness
 //===----------------------------------------------------------------------===//
@@ -673,9 +705,7 @@ bool FieldSensitivePrunedLiveRange<LivenessWithDefs>::isWithinBoundary(
         // If we are not live and have an interesting user that maps to our bit,
         // mark this bit as being live again.
         if (!isLive) {
-          auto interestingUser = isInterestingUser(&blockInst);
-          bool isInteresting =
-              interestingUser.first && interestingUser.second->contains(bit);
+          bool isInteresting = isInterestingUser(&blockInst, bit);
           PRUNED_LIVENESS_LOG(llvm::dbgs()
                      << "        Inst was dead... Is InterestingUser: "
                      << (isInteresting ? "true" : "false") << '\n');
@@ -806,8 +836,7 @@ void findBoundaryInNonDefBlock(SILBasicBlock *block, unsigned bitNo,
   PRUNED_LIVENESS_LOG(llvm::dbgs() << "Looking for boundary in non-def block\n");
   for (SILInstruction &inst : llvm::reverse(*block)) {
     PRUNED_LIVENESS_LOG(llvm::dbgs() << "Visiting: " << inst);
-    auto interestingUser = liveness.isInterestingUser(&inst);
-    if (interestingUser.first && interestingUser.second->contains(bitNo)) {
+    if (liveness.isInterestingUser(&inst, bitNo)) {
       PRUNED_LIVENESS_LOG(llvm::dbgs() << "    Is interesting user for this bit!\n");
       boundary.getLastUserBits(&inst).set(bitNo);
       return;
@@ -837,8 +866,7 @@ void findBoundaryInSSADefBlock(SILNode *ssaDef, unsigned bitNo,
       boundary.getDeadDefsBits(cast<SILNode>(&inst)).set(bitNo);
       return;
     }
-    auto interestingUser = liveness.isInterestingUser(&inst);
-    if (interestingUser.first && interestingUser.second->contains(bitNo)) {
+    if (liveness.isInterestingUser(&inst, bitNo)) {
       PRUNED_LIVENESS_LOG(llvm::dbgs() << "    Found interesting user: " << inst);
       boundary.getLastUserBits(&inst).set(bitNo);
       return;
@@ -973,8 +1001,7 @@ void FieldSensitiveMultiDefPrunedLiveRange::findBoundariesInBlock(
     PRUNED_LIVENESS_LOG(llvm::dbgs()
                << "    Checking if this inst is also a last user...\n");
     if (!isLive) {
-      auto interestingUser = isInterestingUser(&inst);
-      if (interestingUser.first && interestingUser.second->contains(bitNo)) {
+      if (isInterestingUser(&inst, bitNo)) {
         PRUNED_LIVENESS_LOG(
             llvm::dbgs()
             << "        Was interesting user! Moving from dead -> live!\n");