WIP: [FieldSensitivePL] Fix vectorization.

FieldSensitivePrunedLiveness is used as a vectorization of
PrunedLiveness.  An instance of FSPL with N elements needs to be able to
represent the same states as N instances of PL.

Previously, it failed to do that in two significant ways:

(1) It attempted to save space for which elements were live by using
    a range.  This failed to account for instructions which are users of
    non-contiguous fields of an aggregate.

    apply(
      @owned (struct_element_addr %s, #S.f1),
      @owned (struct_element_addr %s, #S.f3)
    )

(2) It used a single bit to represent whether the instruction was
    consuming.  This failed to account for instructions which consumed
    some fields and borrowed others.

    apply(
      @owned (struct_element_addr %s, #S.f1),
      @guaranteed (struct_element_addr %s, #S.f2)
    )

The fix for (1) is to use a bit vector to represent which elements
are used by the instruction.  The fix for (2) is to use a second bit
vector to represent which elements are _consumed_ by the instruction.

rdar://110676577

Author: nate-chandler

include/swift/SIL/FieldSensitivePrunedLiveness.h

@@ -384,6 +384,12 @@ struct TypeTreeLeafTypeRange {
            endEltOffset >= range.endEltOffset;
   }
 
+  void setBits(SmallBitVector &bits) {
+    for (auto element : getRange()) {
+      bits.set(element);
+    }
+  }
+
   IntRange<unsigned> getRange() const {
     return range(startEltOffset, endEltOffset);
   }
@@ -692,17 +698,66 @@ 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) {}
+    // FIXME: FIXME_NOW: Where is this called from?
+    InterestingUser() {};
 
-    InterestingUser &operator&=(bool otherValue) {
-      isConsuming &= otherValue;
-      return *this;
+    InterestingUser(unsigned bitCount, TypeTreeLeafTypeRange range, bool lifetimeEnding) : liveBits(bitCount), consumingBits(bitCount) {
+      for (auto field : range.getRange()) {
+        liveBits.set(field);
+        if (lifetimeEnding) {
+          consumingBits.set(field);
+        }
+      }
+    }
+
+    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});
+    }
+
+    // Optional<IsInterestingUser> isInterestingUserAtBits(SmallBitVector bits) const {
+    //   if (bits.count() == 0)
+    //     return NonUser;
+
+    //   Optional<IsInterestingUser> retval = llvm::None;
+    //   for (auto bit : bits.set_bits()) {
+    //     auto interesting = isInterestingUser(bit);
+    //     if (!retval) {
+    //       retval = interesting;
+    //       continue;
+    //     }
+    //     if (*retval != interesting)
+    //       return llvm::None;
+    //   }
+
+    //   return retval;
+    // }
+
+    IsInterestingUser isInterestingUser(unsigned element) const {
+      if (!liveBits.test(element))
+        return NonUser;
+      return consumingBits.test(element) ? LifetimeEndingUse : NonLifetimeEndingUse;
     }
   };
 
@@ -788,37 +843,37 @@ class FieldSensitivePrunedLiveness {
       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);
-  }
+  // 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);
-  }
+  // 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 *(
@@ -874,19 +929,41 @@ class FieldSensitivePrunedLiveness {
                         SmallBitVector &liveOutBits,
                         SmallBitVector &deadBits) const;
 
-  enum IsInterestingUser { NonUser, NonLifetimeEndingUse, LifetimeEndingUse };
+  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 {
+  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};
+      return NonUser;
+    auto &record = useIter->second;
+    return record.isInterestingUser(element);
+  }
+
+  bool isInterestingUserOfKindThroughoutRange(SILInstruction *user, IsInterestingUser kind, TypeTreeLeafTypeRange range) const {
+    assert(kind != NonUser);
+    bool wantsConsuming = (kind == LifetimeEndingUse);
+
+    auto iter = users.find(user);
+    if (iter == users.end())
+      return false;
+    auto &record = iter->second;
+    for (auto element : range.getRange()) {
+      if (!record.liveBits.test(element))
+        return false;
+      auto consuming = record.consumingBits.test(element);
+      if (consuming != wantsConsuming)
+        return false;
+    }
+    return true;
   }
 
   unsigned getNumSubElements() const { return liveBlocks.getNumBitsToTrack(); }
@@ -912,11 +989,15 @@ 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.subEltSpan.extendThrough(range);
-      iterAndSuccess.first->second &= lifetimeEnding;
+    auto iter = users.find(user);
+    if (iter == users.end()) {
+      iter = users.insert({user, InterestingUser(getNumSubElements(), range, lifetimeEnding)}).first;
+    }
+    for (auto element : range.getRange()) {
+      iter->second.liveBits.set(element);
+      if (lifetimeEnding) {
+        iter->second.consumingBits.set(element);
+      }
     }
   }
 };

lib/SIL/Utils/FieldSensitivePrunedLiveness.cpp

@@ -699,9 +699,9 @@ 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);
+          auto interestingUser = isInterestingUser(&blockInst, bit);
           bool isInteresting =
-              interestingUser.first && interestingUser.second->contains(bit);
+              interestingUser;
           PRUNED_LIVENESS_LOG(llvm::dbgs()
                      << "        Inst was dead... Is InterestingUser: "
                      << (isInteresting ? "true" : "false") << '\n');
@@ -832,8 +832,8 @@ 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)) {
+    auto interestingUser = liveness.isInterestingUser(&inst, bitNo);
+    if (interestingUser) {
       PRUNED_LIVENESS_LOG(llvm::dbgs() << "    Is interesting user for this bit!\n");
       boundary.getLastUserBits(&inst).set(bitNo);
       return;
@@ -863,8 +863,8 @@ 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)) {
+    auto interestingUser = liveness.isInterestingUser(&inst, bitNo);
+    if (interestingUser) {
       PRUNED_LIVENESS_LOG(llvm::dbgs() << "    Found interesting user: " << inst);
       boundary.getLastUserBits(&inst).set(bitNo);
       return;
@@ -999,8 +999,8 @@ 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)) {
+      auto interestingUser = isInterestingUser(&inst, bitNo);
+      if (interestingUser) {
         PRUNED_LIVENESS_LOG(
             llvm::dbgs()
             << "        Was interesting user! Moving from dead -> live!\n");