[AutoDiff] Parameter indices data structure overhaul (#24761)

### Introduce `AutoDiffIndexSubset`

This PR overhauls data structures used for parameter indices in the AutoDiff infrastructure in SIL.

Previously, we used `llvm::SmallBitVector` to represent differentiation parameter indices in both AST and SIL. It was not efficient, and most importantly there's no way to put this in an instruction without causing memory leaks.

This change replaces all uses of `llvm::SmallBitVector` in SIL AutoDiff code paths with a `ASTContext`-uniqued `AutoDiffIndexSubset *` where bits are stored as trailing objects.

`AutoDiffIndexSubset` does not have "parameter indices" in its name because it is not only designed for parameter indices, but also for result indices as we move to supporting multi-result differentiation. `AutoDiffIndexSubset` has set operations like `isSubsetOf`, `isSupersetOf`, and `contains`, but it also has a special _capacity_ property. All differentiable function's parameter indices data should store the number of parameters as capacity, so that the differentiation transform won't need special logic to check whether an index is out of range.

Another minor change is the module format layout of `SILDifferentiableAttr`. It used to store parameter indices as consecutive `bool` bits, but now stores numeric parameter indices directly for efficiency.

It will be necessary to refactor or eliminate `AutoDiffParameterIndices` to make use of `AutoDiffIndexSubset`. `AutoDiffParameterIndices` is at the AST level, so it is not in the scope for this PR.

Unblocks #23482. Partially resolves [TF-67](https://bugs.swift.org/browse/TF-67).

Author: rxwei

include/swift/AST/AutoDiff.h

@@ -20,6 +20,7 @@
 
 #include "ASTContext.h"
 #include "llvm/ADT/SmallBitVector.h"
+#include "swift/Basic/Range.h"
 
 namespace swift {
 
@@ -73,6 +74,7 @@ class ParsedAutoDiffParameter {
 };
 
 class AnyFunctionType;
+class AutoDiffIndexSubset;
 class AutoDiffParameterIndicesBuilder;
 class Type;
 
@@ -173,7 +175,8 @@ class AutoDiffParameterIndices : public llvm::FoldingSetNode {
   ///   ==> returns 1110
   ///   (because the lowered SIL type is (A, B, C, D) -> R)
   ///
-  llvm::SmallBitVector getLowered(AnyFunctionType *functionType) const;
+  AutoDiffIndexSubset *getLowered(ASTContext &ctx,
+                                  AnyFunctionType *functionType) const;
 
   void Profile(llvm::FoldingSetNodeID &ID) const {
     ID.AddInteger(parameters.size());
@@ -219,6 +222,216 @@ class AutoDiffParameterIndicesBuilder {
   unsigned size() { return parameters.size(); }
 };
 
+class AutoDiffIndexSubset : public llvm::FoldingSetNode {
+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.
+  unsigned numBitWords;
+
+  BitWord *getBitWordsData() {
+    return reinterpret_cast<BitWord *>(this + 1);
+  }
+
+  const BitWord *getBitWordsData() const {
+    return reinterpret_cast<const BitWord *>(this + 1);
+  }
+
+  ArrayRef<BitWord> getBitWords() const {
+    return {getBitWordsData(), getNumBitWords()};
+  }
+
+  BitWord getBitWord(unsigned i) const {
+    return getBitWordsData()[i];
+  }
+
+  BitWord &getBitWord(unsigned i) {
+    return getBitWordsData()[i];
+  }
+
+  MutableArrayRef<BitWord> getMutableBitWords() {
+    return {const_cast<BitWord *>(getBitWordsData()), getNumBitWords()};
+  }
+
+  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;
+
+  // 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 {
+    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 {
+    return llvm::all_of(getBitWords(), [](BitWord bw) { return !(bool)bw; });
+  }
+  
+  bool equals(AutoDiffIndexSubset *other) const {
+    return capacity == other->getCapacity() &&
+        getBitWords().equals(other->getBitWords());
+  }
+
+  bool isSubsetOf(AutoDiffIndexSubset *other) const;
+  bool isSupersetOf(AutoDiffIndexSubset *other) const;
+
+  AutoDiffIndexSubset *adding(unsigned index, ASTContext &ctx) const;
+  AutoDiffIndexSubset *extendingCapacity(ASTContext &ctx,
+                                         unsigned newCapacity) const;
+
+  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 << "))";
+  }
+
+  int findNext(int startIndex) const;
+  int findFirst() const { return findNext(-1); }
+  int findPrevious(int endIndex) const;
+  int findLast() const { return findPrevious(capacity); }
+
+  class iterator {
+  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;
+
+    void advance() {
+      assert(current != -1 && "Trying to advance past end.");
+      current = parent->findNext(current);
+    }
+
+  public:
+    iterator(const AutoDiffIndexSubset *parent, int current)
+        : parent(parent), current(current) {}
+    explicit iterator(const AutoDiffIndexSubset *parent)
+        : iterator(parent, parent->findFirst()) {}
+    iterator(const iterator &) = default;
+
+    iterator operator++(int) {
+      auto prev = *this;
+      advance();
+      return prev;
+    }
+
+    iterator &operator++() {
+      advance();
+      return *this;
+    }
+
+    unsigned operator*() const { return current; }
+
+    bool operator==(const iterator &other) const {
+      assert(parent == other.parent &&
+             "Comparing iterators from different AutoDiffIndexSubsets");
+      return current == other.current;
+    }
+
+    bool operator!=(const iterator &other) const {
+      assert(parent == other.parent &&
+             "Comparing iterators from different AutoDiffIndexSubsets");
+      return current != other.current;
+    }
+  };
+};
+
 /// 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
@@ -242,38 +455,33 @@ struct SILAutoDiffIndices {
   ///   Function type: (A, B) -> (C, D) -> R
   ///   Bits: [C][D][A][B]
   ///
-  llvm::SmallBitVector parameters;
+  AutoDiffIndexSubset *parameters;
 
   /// Creates a set of AD indices from the given source index and a bit vector
   /// representing parameter indices.
   /*implicit*/ SILAutoDiffIndices(unsigned source,
-                                  llvm::SmallBitVector parameters)
+                                  AutoDiffIndexSubset *parameters)
       : source(source), parameters(parameters) {}
 
-  /// Creates a set of AD indices from the given source index and an array of
-  /// parameter indices. Elements in `parameters` must be ascending integers.
-  /*implicit*/ SILAutoDiffIndices(unsigned source,
-                                  ArrayRef<unsigned> parameters);
-
   bool operator==(const SILAutoDiffIndices &other) const;
 
   /// Queries whether the function's parameter with index `parameterIndex` is
   /// one of the parameters to differentiate with respect to.
   bool isWrtParameter(unsigned parameterIndex) const {
-    return parameterIndex < parameters.size() &&
-           parameters.test(parameterIndex);
+    return parameterIndex < parameters->getCapacity() &&
+           parameters->contains(parameterIndex);
   }
 
   void print(llvm::raw_ostream &s = llvm::outs()) const {
     s << "(source=" << source << " parameters=(";
-    interleave(parameters.set_bits(),
+    interleave(parameters->getIndices(),
                [&s](unsigned p) { s << p; }, [&s]{ s << ' '; });
     s << "))";
   }
 
   std::string mangle() const {
     std::string result = "src_" + llvm::utostr(source) + "_wrt_";
-    interleave(parameters.set_bits(),
+    interleave(parameters->getIndices(),
                [&](unsigned idx) { result += llvm::utostr(idx); },
                [&] { result += '_'; });
     return result;
@@ -449,19 +657,18 @@ template<typename T> struct DenseMapInfo;
 
 template<> struct DenseMapInfo<SILAutoDiffIndices> {
   static SILAutoDiffIndices getEmptyKey() {
-    return { DenseMapInfo<unsigned>::getEmptyKey(), SmallBitVector() };
+    return { DenseMapInfo<unsigned>::getEmptyKey(), nullptr };
   }
 
   static SILAutoDiffIndices getTombstoneKey() {
-    return { DenseMapInfo<unsigned>::getTombstoneKey(),
-             SmallBitVector(sizeof(intptr_t), true) };
+    return { DenseMapInfo<unsigned>::getTombstoneKey(), nullptr };
   }
 
   static unsigned getHashValue(const SILAutoDiffIndices &Val) {
-    auto params = Val.parameters.set_bits();
     unsigned combinedHash =
       hash_combine(~1U, DenseMapInfo<unsigned>::getHashValue(Val.source),
-                   hash_combine_range(params.begin(), params.end()));
+                   hash_combine_range(Val.parameters->begin(),
+                                      Val.parameters->end()));
     return combinedHash;
   }
 

include/swift/AST/DiagnosticsParse.def

@@ -1535,7 +1535,9 @@ ERROR(sil_inst_autodiff_operand_list_expected_rbrace,PointsToFirstBadToken,
 ERROR(sil_inst_autodiff_num_operand_list_order_mismatch,PointsToFirstBadToken,
       "the number of operand lists does not match the order", ())
 ERROR(sil_inst_autodiff_expected_associated_function_kind_attr,PointsToFirstBadToken,
-      "expects an assoiacted function kind attribute, e.g. '[jvp]'", ())
+      "expected an associated function kind attribute, e.g. '[jvp]'", ())
+ERROR(sil_inst_autodiff_expected_function_type_operand,PointsToFirstBadToken,
+      "expected an operand of a function type", ())
 
 //------------------------------------------------------------------------------
 // MARK: Generics parsing diagnostics

include/swift/AST/Types.h

@@ -4132,14 +4132,14 @@ class SILFunctionType final : public TypeBase, public llvm::FoldingSetNode,
 
   // SWIFT_ENABLE_TENSORFLOW
   CanSILFunctionType getWithDifferentiability(
-      unsigned differentiationOrder, const SmallBitVector &parameterIndices);
+      unsigned differentiationOrder, AutoDiffIndexSubset *parameterIndices);
 
   CanSILFunctionType getWithoutDifferentiability();
 
   /// Returns the type of a differentiation function that is associated with
   /// a function of this type.
   CanSILFunctionType getAutoDiffAssociatedFunctionType(
-      const SmallBitVector &parameterIndices, unsigned resultIndex,
+      AutoDiffIndexSubset *parameterIndices, unsigned resultIndex,
       unsigned differentiationOrder, AutoDiffAssociatedFunctionKind kind,
       SILModule &module, LookupConformanceFn lookupConformance,
       GenericSignature *whereClauseGenericSignature = nullptr);
@@ -4148,7 +4148,7 @@ class SILFunctionType final : public TypeBase, public llvm::FoldingSetNode,
   /// differentiate with respect to for this differentiable function type. (e.g.
   /// which parameters are not @nondiff). The function type must be
   /// differentiable.
-  SmallBitVector getDifferentiationParameterIndices() const;
+  AutoDiffIndexSubset *getDifferentiationParameterIndices();
 
   /// If this is a @convention(witness_method) function with a class
   /// constrained self parameter, return the class constraint for the

include/swift/SIL/SILBuilder.h

@@ -504,7 +504,7 @@ class SILBuilder {
 
   /// SWIFT_ENABLE_TENSORFLOW
   AutoDiffFunctionInst *createAutoDiffFunction(
-      SILLocation loc, const llvm::SmallBitVector &parameterIndices,
+      SILLocation loc, AutoDiffIndexSubset *parameterIndices,
       unsigned differentiationOrder, SILValue original,
       ArrayRef<SILValue> associatedFunctions = {}) {
     return insert(AutoDiffFunctionInst::create(getModule(),

include/swift/SIL/SILFunction.h

@@ -174,6 +174,9 @@ class SILDifferentiableAttr final {
   SILFunction *getOriginal() const { return Original; }
 
   const SILAutoDiffIndices &getIndices() const { return indices; }
+  void setIndices(const SILAutoDiffIndices &indices) {
+    this->indices = indices;
+  }
 
   TrailingWhereClause *getWhereClause() const { return WhereClause; }