Merge pull request #35521 from eeckstein/sil-bitfields

SIL: add a utility which can manage per-block bitfields and flags efficiently.

Author: eeckstein

include/swift/SIL/SILBasicBlock.h

@@ -34,6 +34,7 @@ public llvm::ilist_node<SILBasicBlock>, public SILAllocated<SILBasicBlock> {
   friend class SILFunction;
   friend class SILGlobalVariable;
   template <typename Data, typename Vector> friend class BasicBlockData;
+  friend class BasicBlockBitfield;
 
 public:
   using InstListType = llvm::iplist<SILInstruction>;
@@ -57,6 +58,25 @@ public llvm::ilist_node<SILBasicBlock>, public SILAllocated<SILBasicBlock> {
   /// A value of -1 means that the index is not initialized yet.
   int index = -1;
 
+  /// Custom bits managed by BasicBlockBitfield.
+  uint32_t customBits = 0;
+  
+  /// The BasicBlockBitfield ID of the last initialized bitfield in customBits.
+  /// Example:
+  ///
+  ///                   Last initialized field:
+  ///           lastInitializedBitfieldID == C.bitfieldID
+  ///                              |
+  ///                              V
+  /// customBits:  <unused> EE DDD C BB AAA
+  ///              31         ...         0
+  ///
+  /// -> AAA, BB and C are initialized,
+  ///    DD and EEE are uninitialized
+  ///
+  /// See also: BasicBlockBitfield::bitfieldID, SILFunction::currentBitfieldID.
+  uint64_t lastInitializedBitfieldID = 0;
+
   friend struct llvm::ilist_traits<SILBasicBlock>;
   SILBasicBlock() : Parent(nullptr) {}
   void operator=(const SILBasicBlock &) = delete;

include/swift/SIL/SILBitfield.h

@@ -0,0 +1,169 @@
+//===--- SILBitField.h - Defines the bitfield utilities ---------*- C++ -*-===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2021 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+//
+// This file defines the BasicBlockBitfield and BasicBlockFlag utilities.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef SWIFT_SIL_SILBITFIELD_H
+#define SWIFT_SIL_SILBITFIELD_H
+
+#include "swift/SIL/SILFunction.h"
+
+namespace swift {
+
+/// Utility to add a custom bitfield to a function's basic blocks.
+///
+/// This can be used by transforms to store temporary flags or tiny values per
+/// basic block.
+/// It is very efficient: no memory allocation is needed, no hash set or map is
+/// needed for lookup and there is no initialization cost (in contrast to
+/// BasicBlockData which needs to iterate over all blocks at initialization).
+///
+/// Restrictions:
+/// * BasicBlockBitfield instances must be allocated and deallocated
+///   following a strict stack discipline. This means, it's fine to use them as
+///   (or in) local variables in transformations. But it's e.g. not possible to
+///   store a BasicBlockBitfield in an Analysis.
+/// * The total number of bits which are alive at the same time must not exceed
+///   32.
+class BasicBlockBitfield {
+  /// The bitfield is "added" to the blocks of this function.
+  SILFunction *function;
+
+  /// A single linked list of currently alive BasicBlockBitfields (oldest is
+  /// last, newest is first).
+  /// The head of the list is function->lastAllocatedBitfield.
+  BasicBlockBitfield *parent;
+  
+  /// Initialized with the monotonically increasing currentBitfieldID of the
+  /// function.
+  /// Used to check if the bitfield in a block is initialized.
+  /// If a block's lastInitializedBitfieldID is less than this ID, it means
+  /// that the bits of that block are not initialized yet.
+  /// See also: SILBasicBlock::lastInitializedBitfieldID,
+  ///           SILFunction::currentBitfieldID
+  unsigned bitfieldID;
+
+  short startBit;
+  short endBit;
+  uint32_t mask;
+
+public:
+  BasicBlockBitfield(SILFunction *function, int size) :
+      function(function),
+      parent(function->newestAliveBitfield),
+      bitfieldID(function->currentBitfieldID),
+      startBit(parent ? parent->endBit : 0),
+      endBit(startBit + size),
+      mask(0xffffffffu >> (32 - size) << startBit) {
+    assert(size > 0 && "bit field size must be > 0");
+    assert(endBit <= 32 && "too many/large bit fields allocated in function");
+    assert((!parent || bitfieldID > parent->bitfieldID) &&
+           "BasicBlockBitfield indices are not in order");
+    function->newestAliveBitfield = this;
+    ++function->currentBitfieldID;
+    assert(function->currentBitfieldID != 0 && "currentBitfieldID overflow");
+  }
+
+  ~BasicBlockBitfield() {
+    assert(function->newestAliveBitfield == this &&
+           "BasicBlockBitfield destructed too early");
+    function->newestAliveBitfield = parent;
+  }
+
+  BasicBlockBitfield(const BasicBlockBitfield &) = delete;
+  BasicBlockBitfield(BasicBlockBitfield &&) = delete;
+  BasicBlockBitfield &operator=(const BasicBlockBitfield &) = delete;
+  BasicBlockBitfield &operator=(BasicBlockBitfield &&) = delete;
+
+  SILFunction *getFunction() const { return function; }
+
+  unsigned get(SILBasicBlock *block) const {
+    if (bitfieldID > block->lastInitializedBitfieldID) {
+      // The bitfield is not initialized yet in this block.
+      return 0;
+    }
+    return (block->customBits & mask) >> startBit;
+  }
+
+  void set(SILBasicBlock *block, unsigned value) {
+    assert(((value << startBit) & ~mask) == 0 &&
+           "value too large for BasicBlockBitfield");
+    unsigned clearMask = mask;
+    if (bitfieldID > block->lastInitializedBitfieldID) {
+
+      // The bitfield is not initialized yet in this block.
+      // Initialize the bitfield, and also initialize all parent bitfields,
+      // which are not initialized, yet. Example:
+      //
+      //  This field   Last initialized field
+      //    |            |
+      //    V            V
+      //   EE DDD C BB AAA
+      //
+      // block->lastInitializedBitfieldID == AAA.bitfieldID
+      // -> we have to initialize the fields: BB, C, DDD and EE
+      //
+      BasicBlockBitfield *bf = parent;
+      while (bf && bf->bitfieldID > block->lastInitializedBitfieldID) {
+        clearMask |= bf->mask;
+        bf = bf->parent;
+      }
+      block->lastInitializedBitfieldID = bitfieldID;
+    }
+    block->customBits = (block->customBits & ~clearMask) | (value << startBit);
+  }
+};
+
+/// A BasicBlockBitfield containing a single bit - a flag.
+class BasicBlockFlag {
+  BasicBlockBitfield bit;
+
+public:
+  BasicBlockFlag(SILFunction *function) : bit(function, 1) {}
+
+  SILFunction *getFunction() const { return bit.getFunction(); }
+  
+  bool get(SILBasicBlock *block) const { return (bool)bit.get(block); }
+
+  void set(SILBasicBlock *block) { bit.set(block, 1); }
+  void reset(SILBasicBlock *block) { bit.set(block, 0); }
+
+  /// Sets the flag and returns the old value.
+  bool testAndSet(SILBasicBlock *block) {
+    bool oldValue = get(block);
+    set(block);
+    return oldValue;
+  }
+};
+
+/// A BasicBlockFlag with a set-like API.
+class BasicBlockSet {
+  BasicBlockFlag flag;
+
+public:
+  BasicBlockSet(SILFunction *function) : flag(function) {}
+
+  SILFunction *getFunction() const { return flag.getFunction(); }
+
+  bool contains(SILBasicBlock *block) const { return flag.get(block); }
+
+  /// Returns true if \p block was not contained in the set before inserting.
+  bool insert(SILBasicBlock *block) { return !flag.testAndSet(block); }
+
+  void remove(SILBasicBlock *block) { flag.reset(block); }
+};
+
+} // namespace swift
+
+#endif

include/swift/SIL/SILFunction.h

@@ -37,6 +37,7 @@ class SILInstruction;
 class SILModule;
 class SILFunctionBuilder;
 class SILProfiler;
+class BasicBlockBitfield;
 
 namespace Lowering {
 class TypeLowering;
@@ -148,6 +149,7 @@ class SILFunction
   friend class SILModule;
   friend class SILFunctionBuilder;
   template <typename Data, typename Vector> friend class BasicBlockData;
+  friend class BasicBlockBitfield;
 
   /// Module - The SIL module that the function belongs to.
   SILModule &Module;
@@ -192,6 +194,16 @@ class SILFunction
 
   Identifier ObjCReplacementFor;
 
+  /// The head of a single-linked list of currently alive BasicBlockBitfield.
+  BasicBlockBitfield *newestAliveBitfield = nullptr;
+
+  /// A monotonically increasing ID which is incremented whenever a
+  /// BasicBlockBitfield is constructed.
+  /// Usually this stays below 1000, so a 32-bit unsigned is more than
+  /// sufficient.
+  /// For details see BasicBlockBitfield::bitfieldID;
+  unsigned currentBitfieldID = 1;
+
   /// The function's set of semantics attributes.
   ///
   /// TODO: Why is this using a std::string? Why don't we use uniqued
@@ -1012,9 +1024,6 @@ class SILFunction
   /// Transfer all blocks of \p F into this function, at the begin of the block
   /// list.
   void moveAllBlocksFromOtherFunction(SILFunction *F) {
-    for (SILBasicBlock &block : *F) {
-      block.index = -1;
-    }
     BlockList.splice(begin(), F->BlockList);
   }
 
@@ -1026,7 +1035,6 @@ class SILFunction
     assert(otherFunc != this);
     BlockList.splice(insertPointInThisFunction, otherFunc->BlockList,
                      blockInOtherFunction);
-    blockInOtherFunction->index = -1;
   }
 
   /// Move block \p BB to immediately before the iterator \p IP.

include/swift/SILOptimizer/Utils/BasicBlockOptUtils.h

@@ -23,6 +23,7 @@
 #define SWIFT_SILOPTIMIZER_UTILS_BASICBLOCKOPTUTILS_H
 
 #include "swift/SIL/SILBasicBlock.h"
+#include "swift/SIL/SILBitfield.h"
 #include "swift/SIL/SILCloner.h"
 #include "swift/SIL/SILInstruction.h"
 #include "swift/SILOptimizer/Utils/InstOptUtils.h"
@@ -35,20 +36,22 @@ class SILLoopInfo;
 
 /// Compute the set of reachable blocks.
 class ReachableBlocks {
-  SmallPtrSet<SILBasicBlock *, 32> visited;
+  BasicBlockSet visited;
 
 public:
+  ReachableBlocks(SILFunction *function) : visited(function) {}
+
   /// Invoke \p visitor for each reachable block in \p f in worklist order (at
   /// least one predecessor has been visited--defs are always visited before
   /// uses except for phi-type block args). The \p visitor takes a block
   /// argument, which is already marked visited, and must return true to
   /// continue visiting blocks.
   ///
   /// Returns true if all reachable blocks were visited.
-  bool visit(SILFunction *f, function_ref<bool(SILBasicBlock *)> visitor);
+  bool visit(function_ref<bool(SILBasicBlock *)> visitor);
 
   /// Return true if \p bb has been visited.
-  bool isVisited(SILBasicBlock *bb) const { return visited.count(bb); }
+  bool isVisited(SILBasicBlock *bb) const { return visited.contains(bb); }
 };
 
 /// Remove all instructions in the body of \p bb in safe manner by using

include/swift/SILOptimizer/Utils/ValueLifetime.h

@@ -20,6 +20,7 @@
 #include "swift/Basic/STLExtras.h"
 #include "swift/SIL/SILBuilder.h"
 #include "swift/SIL/SILInstruction.h"
+#include "swift/SIL/SILBitfield.h"
 #include "swift/SILOptimizer/Utils/InstOptUtils.h"
 
 namespace swift {
@@ -34,7 +35,10 @@ class ValueLifetimeAnalysis {
   PointerUnion<SILInstruction *, SILArgument *> defValue;
 
   /// The set of blocks where the value is live.
-  llvm::SmallSetVector<SILBasicBlock *, 16> liveBlocks;
+  llvm::SmallVector<SILBasicBlock *, 16> liveBlocks;
+  
+  /// True for blocks which are in liveBlocks.
+  BasicBlockFlag inLiveBlocks;
 
   /// The set of instructions where the value is used, or the users-list
   /// provided with the constructor.
@@ -66,7 +70,7 @@ class ValueLifetimeAnalysis {
   /// iterators.
   template <typename RangeTy>
   ValueLifetimeAnalysis(SILArgument *def, const RangeTy &useRange)
-      : defValue(def), userSet() {
+      : defValue(def), inLiveBlocks(def->getFunction()), userSet() {
     for (SILInstruction *use : useRange)
       userSet.insert(use);
     propagateLiveness();
@@ -75,7 +79,7 @@ class ValueLifetimeAnalysis {
   ValueLifetimeAnalysis(
       SILArgument *def,
       llvm::iterator_range<ValueBaseUseIterator> useRange)
-      : defValue(def), userSet() {
+      : defValue(def), inLiveBlocks(def->getFunction()), userSet() {
     for (Operand *use : useRange)
       userSet.insert(use->getUser());
     propagateLiveness();
@@ -84,7 +88,7 @@ class ValueLifetimeAnalysis {
   template <typename RangeTy>
   ValueLifetimeAnalysis(
       SILInstruction *def, const RangeTy &useRange)
-      : defValue(def), userSet() {
+      : defValue(def), inLiveBlocks(def->getFunction()), userSet() {
     for (SILInstruction *use : useRange)
       userSet.insert(use);
     propagateLiveness();
@@ -93,7 +97,7 @@ class ValueLifetimeAnalysis {
   ValueLifetimeAnalysis(
       SILInstruction *def,
       llvm::iterator_range<ValueBaseUseIterator> useRange)
-      : defValue(def), userSet() {
+      : defValue(def), inLiveBlocks(def->getFunction()), userSet() {
     for (Operand *use : useRange)
       userSet.insert(use->getUser());
     propagateLiveness();
@@ -144,7 +148,7 @@ class ValueLifetimeAnalysis {
 
   /// Returns true if the value is alive at the begin of block \p bb.
   bool isAliveAtBeginOfBlock(SILBasicBlock *bb) {
-    return liveBlocks.count(bb) &&
+    return inLiveBlocks.get(bb) &&
            (hasUsersBeforeDef || bb != getDefValueParentBlock());
   }
 

lib/SIL/IR/SILBasicBlock.cpp

@@ -309,6 +309,8 @@ transferNodesFromList(llvm::ilist_traits<SILBasicBlock> &SrcTraits,
   // If splicing blocks not in the same function, update the parent pointers.
   for (; First != Last; ++First) {
     First->Parent = Parent;
+    First->index = -1;
+    First->lastInitializedBitfieldID = 0;
     for (auto &II : *First)
       II.setDebugScope(ScopeCloner.getOrCreateClonedScope(II.getDebugScope()));
   }

lib/SIL/IR/SILFunction.cpp

@@ -210,6 +210,8 @@ SILFunction::~SILFunction() {
 
   assert(RefCount == 0 &&
          "Function cannot be deleted while function_ref's still exist");
+  assert(!newestAliveBitfield &&
+         "Not all BasicBlockBitfields deleted at function destruction");
 }
 
 void SILFunction::createProfiler(ASTNode Root, SILDeclRef forDecl,

lib/SILOptimizer/Analysis/EscapeAnalysis.cpp

@@ -1619,8 +1619,8 @@ void EscapeAnalysis::ConnectionGraph::verify() const {
   // Verify that all pointer nodes are still mapped, otherwise the process of
   // merging nodes may have lost information. Only visit reachable blocks,
   // because the graph builder only mapped values from reachable blocks.
-  ReachableBlocks reachable;
-  reachable.visit(F, [this](SILBasicBlock *bb) {
+  ReachableBlocks reachable(F);
+  reachable.visit([this](SILBasicBlock *bb) {
     for (auto &i : *bb) {
       if (isNonWritableMemoryAddress(&i))
         continue;
@@ -1748,8 +1748,8 @@ void EscapeAnalysis::buildConnectionGraph(FunctionInfo *FInfo,
   assert(ConGraph->isEmpty());
 
   // Visit the blocks in dominance order.
-  ReachableBlocks reachable;
-  reachable.visit(ConGraph->F, [&](SILBasicBlock *bb) {
+  ReachableBlocks reachable(ConGraph->F);
+  reachable.visit([&](SILBasicBlock *bb) {
     // Create edges for the instructions.
     for (auto &i : *bb) {
       analyzeInstruction(&i, FInfo, BottomUpOrder, RecursionDepth);