[AutoDiff] Parameter indices data structure overhaul.

Author: rxwei

include/swift/AST/AutoDiff.h

@@ -223,22 +223,31 @@ class AutoDiffParameterIndicesBuilder {
 };
 
 class AutoDiffIndexSubset : public llvm::FoldingSetNode {
-private:
-  using BitWord = uint64_t;
+public:
+  typedef uint64_t BitWord;
+
+  static constexpr unsigned bitWordSize = sizeof(BitWord);
+  static constexpr unsigned numBitsPerBitWord = bitWordSize * 8;
+
+  static std::pair<unsigned, unsigned>
+  getBitWordIndexAndOffset(unsigned index) {
+    auto bitWordIndex = index / numBitsPerBitWord;
+    auto bitWordOffset = index % numBitsPerBitWord;
+    return {bitWordIndex, bitWordOffset};
+  }
 
+  static unsigned getNumBitWordsNeededForCapacity(unsigned capacity) {
+    if (capacity == 0) return 0;
+    return capacity / numBitsPerBitWord + 1;
+  }
+  
+private:
   /// The total capacity of the index subset, which is `1` less than the largest
   /// index.
   unsigned capacity;
   /// The number of bit words in the index subset. in the index subset.
   unsigned numBitWords;
 
-  static std::pair<unsigned, unsigned> getBitWordIndexAndOffset(unsigned index);
-  static unsigned getNumBitWordsNeededForCapacity(unsigned capacity);
-
-  unsigned getNumBitWords() const {
-    return numBitWords;
-  }
-
   BitWord *getBitWordsData() {
     return reinterpret_cast<BitWord *>(this + 1);
   }
@@ -263,62 +272,159 @@ class AutoDiffIndexSubset : public llvm::FoldingSetNode {
     return {const_cast<BitWord *>(getBitWordsData()), getNumBitWords()};
   }
 
-  explicit AutoDiffIndexSubset(unsigned capacity, unsigned numBitWords,
-                               ArrayRef<unsigned> indices);
+  explicit AutoDiffIndexSubset(unsigned capacity, ArrayRef<unsigned> indices)
+      : capacity(capacity),
+        numBitWords(getNumBitWordsNeededForCapacity(capacity)) {
+    std::uninitialized_fill_n(getBitWordsData(), numBitWords, 0);
+    for (auto i : indices) {
+      unsigned bitWordIndex, offset;
+      std::tie(bitWordIndex, offset) = getBitWordIndexAndOffset(i);
+      getBitWord(bitWordIndex) |= (1 << offset);
+    }
+  }
 
 public:
   AutoDiffIndexSubset() = delete;
   AutoDiffIndexSubset(const AutoDiffIndexSubset &) = delete;
   AutoDiffIndexSubset &operator=(const AutoDiffIndexSubset &) = delete;
 
-  static AutoDiffIndexSubset *get(ASTContext &ctx, unsigned capacity,
-                                  bool includeAll = false);
-  static AutoDiffIndexSubset *get(ASTContext &ctx, unsigned capacity,
-                                  IntRange<> range);
-  static AutoDiffIndexSubset *get(ASTContext &ctx, unsigned capacity,
+  // Defined in ASTContext.h.
+  static AutoDiffIndexSubset *get(ASTContext &ctx,
+                                  unsigned capacity,
                                   ArrayRef<unsigned> indices);
 
+  static AutoDiffIndexSubset *getDefault(ASTContext &ctx,
+                                         unsigned capacity,
+                                         bool includeAll = false) {
+    if (includeAll)
+      return getFromRange(ctx, capacity, IntRange<>(capacity));
+    return get(ctx, capacity, {});
+  }
+
+  static AutoDiffIndexSubset *getFromRange(ASTContext &ctx,
+                                           unsigned capacity,
+                                           IntRange<> range) {
+    return get(ctx, capacity,
+               SmallVector<unsigned, 8>(range.begin(), range.end()));
+  }
+
+  unsigned getNumBitWords() const {
+    return numBitWords;
+  }
+
   unsigned getCapacity() const {
     return capacity;
   }
 
   class iterator;
 
-  iterator begin() const;
-  iterator end() const;
-  iterator_range<iterator> getIndices() const;
-  unsigned getNumIndices() const;
+  iterator begin() const {
+    return iterator(this);
+  }
+  
+  iterator end() const {
+    return iterator(this, (int)capacity);
+  }
+  
+  iterator_range<iterator> getIndices() const {
+    return make_range(begin(), end());
+  }
+
+  unsigned getNumIndices() const {
+    return (unsigned)std::distance(begin(), end());
+  }
 
   bool contains(unsigned index) const {
     unsigned bitWordIndex, offset;
     std::tie(bitWordIndex, offset) = getBitWordIndexAndOffset(index);
     return getBitWord(bitWordIndex) & (1 << offset);
   }
 
-  bool isEmpty() const;
-  bool equals(const AutoDiffIndexSubset *other) const;
-  bool isSubsetOf(const AutoDiffIndexSubset *other) const;
-  bool isSupersetOf(const AutoDiffIndexSubset *other) const;
+  bool isEmpty() const {
+    return llvm::all_of(getBitWords(), [](BitWord bw) { return !(bool)bw; });
+  }
+  
+  bool equals(const AutoDiffIndexSubset *other) const {
+    return capacity == other->getCapacity() &&
+        getBitWords().equals(other->getBitWords());
+  }
+
+  bool isSubsetOf(const AutoDiffIndexSubset *other) const {
+    assert(capacity == other->capacity);
+    for (auto index : range(numBitWords))
+      if (getBitWord(index) & ~other->getBitWord(index))
+        return false;
+    return true;
+  }
+
+  bool isSupersetOf(const AutoDiffIndexSubset *other) const {
+    assert(capacity == other->capacity);
+    for (auto index : range(numBitWords))
+      if (~getBitWord(index) & other->getBitWord(index))
+        return false;
+    return true;
+  }
 
-  AutoDiffIndexSubset *adding(unsigned index, ASTContext &ctx) const;
-  AutoDiffIndexSubset *extendingCapacity(ASTContext &ctx,
-                                         unsigned newCapacity) const;
+  AutoDiffIndexSubset *adding(
+      unsigned index, ASTContext &ctx) const {
+    assert(index < getCapacity());
+    SmallVector<unsigned, 8> newIndices;
+    newIndices.reserve(capacity + 1);
+    bool inserted = false;
+    for (auto curIndex : getIndices()) {
+      if (inserted && curIndex > index) {
+        newIndices.push_back(index);
+        inserted = false;
+      }
+      newIndices.push_back(curIndex);
+    }
+    return get(ctx, capacity, newIndices);
+  }
 
-  void Profile(llvm::FoldingSetNodeID &id) const;
+  AutoDiffIndexSubset *extendingCapacity(
+      ASTContext &ctx, unsigned newCapacity) const {
+    assert(newCapacity >= capacity);
+    if (newCapacity == capacity)
+      return const_cast<AutoDiffIndexSubset *>(this);
+    SmallVector<unsigned, 8> indices;
+    for (auto index : getIndices())
+      indices.push_back(index);
+    return AutoDiffIndexSubset::get(ctx, newCapacity, indices);
+  }
+
+  void Profile(llvm::FoldingSetNodeID &id) const {
+    id.AddInteger(capacity);
+    for (auto index : getIndices())
+      id.AddInteger(index);
+  }
+
+  void print(llvm::raw_ostream &s = llvm::outs()) const {
+    s << '{';
+    interleave(range(capacity), [this, &s](unsigned i) { s << contains(i); },
+               [&s] { s << ", "; });
+    s << '}';
+  }
+
+  void dump(llvm::raw_ostream &s = llvm::errs()) const {
+    s << "(autodiff_index_subset capacity=" << capacity << " indices=(";
+    interleave(getIndices(), [&s](unsigned i) { s << i; },
+               [&s] { s << ", "; });
+    s << "))";
+  }
 
-private:
   int findNext(int startIndex) const;
   int findFirst() const { return findNext(-1); }
   int findPrevious(int endIndex) const;
   int findLast() const { return findPrevious(capacity); }
 
-public:
   class iterator {
-  typedef unsigned value_type;
-  typedef int difference_type;
-  typedef unsigned * pointer;
-  typedef unsigned & reference;
-  typedef std::forward_iterator_tag iterator_category;
+  public:
+    typedef unsigned value_type;
+    typedef unsigned difference_type;
+    typedef unsigned * pointer;
+    typedef unsigned & reference;
+    typedef std::forward_iterator_tag iterator_category;
+
   private:
     const AutoDiffIndexSubset *parent;
     int current = 0;
@@ -349,42 +455,19 @@ class AutoDiffIndexSubset : public llvm::FoldingSetNode {
     unsigned operator*() const { return current; }
 
     bool operator==(const iterator &other) const {
-      assert(&parent == &other.parent &&
+      assert(parent == other.parent &&
              "Comparing iterators from different AutoDiffIndexSubsets");
       return current == other.current;
     }
 
     bool operator!=(const iterator &other) const {
-      assert(&parent == &other.parent &&
+      assert(parent == other.parent &&
              "Comparing iterators from different AutoDiffIndexSubsets");
       return current != other.current;
     }
   };
 };
 
-class AutoDiffFunctionParameterSubset {
-private:
-  AutoDiffIndexSubset *indexSubset;
-  bool curried;
-
-public:
-  explicit AutoDiffFunctionParameterSubset(
-      AutoDiffIndexSubset *indexSubset, bool isCurried)
-      : indexSubset(indexSubset), curried(isCurried) {}
-
-  explicit AutoDiffFunctionParameterSubset(
-      ASTContext &ctx, AutoDiffIndexSubset *parameterSubset,
-      Optional<bool> isSelfIncluded);
-
-  AutoDiffIndexSubset *getIndexSubset() const {
-    return indexSubset;
-  }
-
-  bool isCurried() const {
-    return curried;
-  }
-};
-
 /// SIL-level automatic differentiation indices. Consists of a source index,
 /// i.e. index of the dependent result to differentiate from, and parameter
 /// indices, i.e. index of independent parameters to differentiate with

include/swift/Serialization/ModuleFormat.h

@@ -52,7 +52,7 @@ const uint16_t SWIFTMODULE_VERSION_MAJOR = 0;
 /// describe what change you made. The content of this comment isn't important;
 /// it just ensures a conflict if two people change the module format.
 /// Don't worry about adhering to the 80-column limit for this line.
-const uint16_t SWIFTMODULE_VERSION_MINOR = 489; // Last change: `@differentiating` wrt
+const uint16_t SWIFTMODULE_VERSION_MINOR = 490; // Last change: `@differentiable` parameter indices layout.
 
 using DeclIDField = BCFixed<31>;
 

lib/AST/ASTContext.cpp

@@ -4536,9 +4536,9 @@ AutoDiffParameterIndices::get(llvm::SmallBitVector indices, ASTContext &C) {
   return newNode;
 }
 
-AutoDiffIndexSubset *AutoDiffIndexSubset::get(ASTContext &ctx,
-                                              unsigned capacity,
-                                              ArrayRef<unsigned> indices) {
+AutoDiffIndexSubset *
+AutoDiffIndexSubset::get(ASTContext &ctx, unsigned capacity,
+                         ArrayRef<unsigned> indices) {
   auto &foldingSet = ctx.getImpl().AutoDiffIndexSubsets;
   llvm::FoldingSetNodeID id;
   id.AddInteger(capacity);
@@ -4548,19 +4548,19 @@ AutoDiffIndexSubset *AutoDiffIndexSubset::get(ASTContext &ctx,
   for (unsigned index : indices) {
 #ifndef NDEBUG
     assert((int)index > last && "Indices must be ascending");
-    last = index;
+    last = (int)index;
 #endif
     id.AddInteger(index);
   }
   void *insertPos = nullptr;
   auto *existing = foldingSet.FindNodeOrInsertPos(id, insertPos);
   if (existing)
     return existing;
-  auto numBitWords = sizeof(AutoDiffIndexSubset) +
-  getNumBitWordsNeededForCapacity(capacity);
+  auto sizeToAlloc = sizeof(AutoDiffIndexSubset) +
+      getNumBitWordsNeededForCapacity(capacity);
   auto *buf = reinterpret_cast<AutoDiffIndexSubset *>(
-      ctx.Allocate(numBitWords, alignof(AutoDiffIndexSubset)));
-  auto *newNode = new (buf) AutoDiffIndexSubset(capacity, numBitWords, indices);
+      ctx.Allocate(sizeToAlloc, alignof(AutoDiffIndexSubset)));
+  auto *newNode = new (buf) AutoDiffIndexSubset(capacity, indices);
   foldingSet.InsertNode(newNode, insertPos);
   return newNode;
 }