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

It's used to implement `InstructionSet` and `ValueSet`: sets of SILValues and SILInstructions.
Just like `BasicBlockSet` for basic blocks, the set is implemented by setting bits directly in SILNode.
This is super efficient because insertion and deletion to/from the set are basic bit operations.

The cost is an additional word in SILNode. But this is basically negligible: it just adds ~0.7% of memory used for SILInstructions.
In my experiments, I didn't see any relevant changes in memory consumption or compile time.

Author: eeckstein

SwiftCompilerSources/Sources/Optimizer/DataStructures/BasicBlockSet.swift

@@ -17,118 +17,46 @@
 #ifndef SWIFT_SIL_BASICBLOCKBITS_H
 #define SWIFT_SIL_BASICBLOCKBITS_H
 
-#include "swift/SIL/SILFunction.h"
-#include "llvm/ADT/SmallVector.h"
+#include "swift/SIL/SILBitfield.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.
-/// The memory managed is a 32 bit field within each basic block (\see
-/// BasicBlock::customBits) and thus is very efficient: no memory allocation is
-/// needed, no hash set or map is needed for lookup and there is no
+/// The bits are stored in a 32 bit field within each basic block (\see
+/// SILBasicBlock::customBits) which 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).
 ///
 /// Invariants:
 /// * BasicBlockBitfield instances must be allocated and deallocated
 ///   following a strict stack discipline, because bit-positions in
-///   BasicBlock::customBits are "allocated" and "freed" with a stack-allocation
+///   SILBasicBlock::customBits are "allocated" and "freed" with a stack-allocation
 ///   algorithm. This means, it's fine to use a BasicBlockBitfield as (or in)
 ///   local variables, e.g. in transformations. But it's 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 (the size of BasicBlock::customBits).
-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;
+///   32 (the size of SILBasicBlock::customBits).
+class BasicBlockBitfield : public SILBitfield<BasicBlockBitfield, SILBasicBlock> {
+  template <class, class> friend class SILBitfield;
+
+  BasicBlockBitfield *insertInto(SILFunction *function) {
+    BasicBlockBitfield *oldParent = function->newestAliveBlockBitfield;
+    function->newestAliveBlockBitfield = this;
+    return oldParent;
+  }
 
 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");
-  }
+    SILBitfield(function, size, insertInto(function)) {}
 
   ~BasicBlockBitfield() {
-    assert(function->newestAliveBitfield == this &&
+    assert(function->newestAliveBlockBitfield == 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 {
-    assert(block->getParent() == function);
-    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(block->getParent() == function);
-    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);
+    function->newestAliveBlockBitfield = parent;
   }
 };
 

include/swift/SIL/BasicBlockBits.h

@@ -0,0 +1,130 @@
+//===--- NodeBits.h ---------------------------------------------*- C++ -*-===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2022 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 utilities for SILNode bit fields and sets.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef SWIFT_SIL_NODEBITS_H
+#define SWIFT_SIL_NODEBITS_H
+
+#include "swift/SIL/SILBitfield.h"
+
+namespace swift {
+
+/// Utility to add a custom bitfield to a SILInstruction or SILValue.
+///
+/// This can be used by transforms to store temporary flags or tiny values per
+/// instruction or value.
+/// The bits are stored in a small bit field (8 bits) within each node (\see
+/// SILNode::Bits::customBits) which is very efficient: no memory allocation
+/// is needed, no hash set or map is needed for lookup and there is no
+/// initialization cost.
+///
+/// Invariants:
+/// * NodeBitfield instances must be allocated and deallocated
+///   following a strict stack discipline, because bit-positions in the node's
+///   bit field are "allocated" and "freed" with a stack-allocation
+///   algorithm. This means, it's fine to use a NodeBitfield as (or in)
+///   local variables, e.g. in transformations. But it's not possible to store
+///   a NodeBitfield in an Analysis.
+/// * The total number of bits which are alive at the same time must not exceed
+///   8 (the size of SILNode::Bits::customBits).
+class NodeBitfield : public SILBitfield<NodeBitfield, SILNode> {
+  template <class, class> friend class SILBitfield;
+
+  NodeBitfield *insertInto(SILFunction *function) {
+    NodeBitfield *oldParent = function->newestAliveNodeBitfield;
+    function->newestAliveNodeBitfield = this;
+    return oldParent;
+  }
+
+public:
+  NodeBitfield(SILFunction *function, int size) :
+    SILBitfield(function, size, insertInto(function)) {}
+
+  ~NodeBitfield() {
+    assert(function->newestAliveNodeBitfield == this &&
+           "BasicBlockBitfield destructed too early");
+    function->newestAliveNodeBitfield = parent;
+  }
+};
+
+/// A set of SILNodes.
+///
+/// For details see NodeBitfield.
+class NodeSet {
+  NodeBitfield bit;
+  
+public:
+  NodeSet(SILFunction *function) : bit(function, 1) {}
+
+  SILFunction *getFunction() const { return bit.getFunction(); }
+
+  bool contains(SILNode *node) const { return (bool)bit.get(node); }
+
+  /// Returns true if \p node was not contained in the set before inserting.
+  bool insert(SILNode *node) {
+    bool wasContained = contains(node);
+    if (!wasContained) {
+      bit.set(node, 1);
+    }
+    return !wasContained;
+  }
+
+  void erase(SILNode *node) {
+    bit.set(node, 0);
+  }
+};
+
+
+/// A set of SILInstructions.
+///
+/// For details see NodeBitfield.
+class InstructionSet {
+  NodeSet nodeSet;
+  
+public:
+  InstructionSet(SILFunction *function) : nodeSet(function) {}
+
+  SILFunction *getFunction() const { return nodeSet.getFunction(); }
+
+  bool contains(SILInstruction *inst) const { return nodeSet.contains(inst->asSILNode()); }
+
+  /// Returns true if \p inst was not contained in the set before inserting.
+  bool insert(SILInstruction *inst) { return nodeSet.insert(inst->asSILNode()); }
+
+  void erase(SILInstruction *inst) { nodeSet.erase(inst->asSILNode()); }
+};
+
+/// A set of SILValues.
+///
+/// For details see NodeBitfield.
+class ValueSet {
+  NodeSet nodeSet;
+  
+public:
+  ValueSet(SILFunction *function) : nodeSet(function) {}
+
+  SILFunction *getFunction() const { return nodeSet.getFunction(); }
+
+  bool contains(SILValue value) const { return nodeSet.contains(value); }
+
+  /// Returns true if \p value was not contained in the set before inserting.
+  bool insert(SILValue value) { return nodeSet.insert(value); }
+
+  void erase(SILValue value) { nodeSet.erase(value); }
+};
+
+} // namespace swift
+
+#endif

include/swift/SIL/NodeBits.h

@@ -38,10 +38,12 @@ public SwiftObjectHeader {
   friend class SILFunction;
   friend class SILGlobalVariable;
   template <typename, unsigned> friend class BasicBlockData;
-  friend class BasicBlockBitfield;
+  template <class, class> friend class SILBitfield;
 
   static SwiftMetatype registeredMetatype;
-    
+  
+  using CustomBitsType = uint32_t;
+  
 public:
   using InstListType = llvm::iplist<SILInstruction>;
 private:
@@ -67,7 +69,7 @@ public SwiftObjectHeader {
   int index = -1;
 
   /// Custom bits managed by BasicBlockBitfield.
-  uint32_t customBits = 0;
+  CustomBitsType customBits = 0;
 
   /// The BasicBlockBitfield ID of the last initialized bitfield in customBits.
   /// Example:
@@ -82,9 +84,17 @@ public SwiftObjectHeader {
   /// -> AAA, BB and C are initialized,
   ///    DD and EEE are uninitialized
   ///
-  /// See also: BasicBlockBitfield::bitfieldID, SILFunction::currentBitfieldID.
+  /// See also: SILBitfield::bitfieldID, SILFunction::currentBitfieldID.
   uint64_t lastInitializedBitfieldID = 0;
 
+  // Used by `BasicBlockBitfield`.
+  unsigned getCustomBits() const { return customBits; }
+  // Used by `BasicBlockBitfield`.
+  void setCustomBits(unsigned value) { customBits = value; }
+
+  // Used by `BasicBlockBitfield`.
+  enum { numCustomBits = std::numeric_limits<CustomBitsType>::digits };
+
   friend struct llvm::ilist_traits<SILBasicBlock>;
 
   SILBasicBlock();
@@ -111,6 +121,7 @@ public SwiftObjectHeader {
   Optional<llvm::StringRef> getDebugName() const;
 
   SILFunction *getParent() { return Parent; }
+  SILFunction *getFunction() { return getParent(); }
   const SILFunction *getParent() const { return Parent; }
 
   SILModule &getModule() const;