[AudoDiff] NFC: Replace 'SILAutoDiffIndices' with 'AutoDiffConfig'. (swiftlang#35079)

Resolve rdar://71678394 / SR-13889.

Author: rxwei

include/swift/AST/AutoDiff.h

@@ -173,80 +173,59 @@ enum class AutoDiffGeneratedDeclarationKind : uint8_t {
   BranchingTraceEnum
 };
 
-/// SIL-level automatic differentiation indices. Consists of:
-/// - The differentiability parameter indices.
-/// - The differentiability result indices.
-// TODO(TF-913): Remove `SILAutoDiffIndices` in favor of `AutoDiffConfig`.
-// `AutoDiffConfig` additionally stores a derivative generic signature.
-struct SILAutoDiffIndices {
-  /// The indices of independent parameters to differentiate with respect to.
-  IndexSubset *parameters;
-  /// The indices of dependent results to differentiate from.
-  IndexSubset *results;
-
-  /*implicit*/ SILAutoDiffIndices(IndexSubset *parameters, IndexSubset *results)
-      : parameters(parameters), results(results) {
-    assert(parameters && "Parameter indices must be non-null");
-    assert(results && "Result indices must be non-null");
-  }
-
-  bool operator==(const SILAutoDiffIndices &other) const;
-
-  bool operator!=(const SILAutoDiffIndices &other) const {
-    return !(*this == other);
-  };
+/// Identifies an autodiff derivative function configuration:
+/// - Parameter indices.
+/// - Result indices.
+/// - Derivative generic signature (optional).
+struct AutoDiffConfig {
+  IndexSubset *parameterIndices;
+  IndexSubset *resultIndices;
+  GenericSignature derivativeGenericSignature;
+
+  /*implicit*/ AutoDiffConfig(
+      IndexSubset *parameterIndices, IndexSubset *resultIndices,
+      GenericSignature derivativeGenericSignature = GenericSignature())
+      : parameterIndices(parameterIndices), resultIndices(resultIndices),
+        derivativeGenericSignature(derivativeGenericSignature) {}
 
   /// Returns true if `parameterIndex` is a differentiability parameter index.
   bool isWrtParameter(unsigned parameterIndex) const {
-    return parameterIndex < parameters->getCapacity() &&
-           parameters->contains(parameterIndex);
+    return parameterIndex < parameterIndices->getCapacity() &&
+           parameterIndices->contains(parameterIndex);
   }
 
-  void print(llvm::raw_ostream &s = llvm::outs()) const;
-  SWIFT_DEBUG_DUMP;
+  /// Returns true if `resultIndex` is a differentiability result index.
+  bool isWrtResult(unsigned resultIndex) const {
+    return resultIndex < resultIndices->getCapacity() &&
+           resultIndices->contains(resultIndex);
+  }
 
+  AutoDiffConfig withGenericSignature(GenericSignature signature) const {
+    return AutoDiffConfig(parameterIndices, resultIndices, signature);
+  }
+
+  // TODO(SR-13506): Use principled mangling for AD-generated symbols.
   std::string mangle() const {
     std::string result = "src_";
     interleave(
-        results->getIndices(),
+        resultIndices->getIndices(),
         [&](unsigned idx) { result += llvm::utostr(idx); },
         [&] { result += '_'; });
     result += "_wrt_";
     llvm::interleave(
-        parameters->getIndices(),
+        parameterIndices->getIndices(),
         [&](unsigned idx) { result += llvm::utostr(idx); },
         [&] { result += '_'; });
     return result;
   }
-};
-
-/// Identifies an autodiff derivative function configuration:
-/// - Parameter indices.
-/// - Result indices.
-/// - Derivative generic signature (optional).
-struct AutoDiffConfig {
-  IndexSubset *parameterIndices;
-  IndexSubset *resultIndices;
-  GenericSignature derivativeGenericSignature;
-
-  /*implicit*/ AutoDiffConfig(IndexSubset *parameterIndices,
-                              IndexSubset *resultIndices,
-                              GenericSignature derivativeGenericSignature)
-      : parameterIndices(parameterIndices), resultIndices(resultIndices),
-        derivativeGenericSignature(derivativeGenericSignature) {}
-
-  /// Returns the `SILAutoDiffIndices` corresponding to this config's indices.
-  // TODO(TF-913): This is a temporary shim for incremental removal of
-  // `SILAutoDiffIndices`. Eventually remove this.
-  SILAutoDiffIndices getSILAutoDiffIndices() const;
 
   void print(llvm::raw_ostream &s = llvm::outs()) const;
   SWIFT_DEBUG_DUMP;
 };
 
 inline llvm::raw_ostream &operator<<(llvm::raw_ostream &s,
-                                     const SILAutoDiffIndices &indices) {
-  indices.print(s);
+                                     const AutoDiffConfig &config) {
+  config.print(s);
   return s;
 }
 

include/swift/AST/IndexSubset.h

@@ -134,7 +134,7 @@ class IndexSubset : public llvm::FoldingSetNode {
   /// Returns the number of bit words used to store the index subset.
   // Note: Use `getCapacity()` to get the total index subset capacity.
   // This is public only for unit testing
-  // (in unittests/AST/SILAutoDiffIndices.cpp).
+  // (in unittests/AST/IndexSubsetTests.cpp).
   unsigned getNumBitWords() const {
     return numBitWords;
   }

include/swift/SIL/SILDifferentiabilityWitness.h

@@ -133,11 +133,6 @@ class SILDifferentiabilityWitness
   bool isSerialized() const { return IsSerialized; }
   const DeclAttribute *getAttribute() const { return Attribute; }
 
-  /// Returns the `SILAutoDiffIndices` corresponding to this config's indices.
-  // TODO(TF-893): This is a temporary shim for incremental removal of
-  // `SILAutoDiffIndices`. Eventually remove this.
-  SILAutoDiffIndices getSILAutoDiffIndices() const;
-
   /// Verify that the differentiability witness is well-formed.
   void verify(const SILModule &module) const;
 

include/swift/SILOptimizer/Analysis/DifferentiableActivityAnalysis.h

@@ -200,21 +200,47 @@ class DifferentiableActivityInfo {
   /// (dependent variable) indices.
   bool isUseful(SILValue value, IndexSubset *dependentVariableIndices) const;
 
-  /// Returns true if the given value is active for the given
-  /// `SILAutoDiffIndices` (parameter indices and result index).
-  bool isActive(SILValue value, const SILAutoDiffIndices &indices) const;
+  /// Returns true if the given value is active for the given parameter indices
+  /// and result indices.
+  bool isActive(SILValue value,
+                IndexSubset *parameterIndices,
+                IndexSubset *resultIndices) const;
+
+  /// Returns true if the given value is active for the given config.
+  bool isActive(SILValue value, AutoDiffConfig config) const {
+    return isActive(value, config.parameterIndices, config.resultIndices);
+  }
+
+  /// Returns the activity of the given value for the given config.
+  Activity getActivity(SILValue value,
+                       IndexSubset *parameterIndices,
+                       IndexSubset *resultIndices) const;
 
-  /// Returns the activity of the given value for the given `SILAutoDiffIndices`
-  /// (parameter indices and result index).
-  Activity getActivity(SILValue value, const SILAutoDiffIndices &indices) const;
+  /// Returns the activity of the given value for the given config.
+  Activity getActivity(SILValue value, AutoDiffConfig config) const {
+    return getActivity(value, config.parameterIndices, config.resultIndices);
+  }
 
-  /// Prints activity information for the `indices` of the given `value`.
-  void dump(SILValue value, const SILAutoDiffIndices &indices,
+  /// Prints activity information for the config of the given value.
+  void dump(SILValue value,
+            IndexSubset *parameterIndices, IndexSubset *resultIndices,
             llvm::raw_ostream &s = llvm::dbgs()) const;
 
-  /// Prints activity information for the given `indices`.
-  void dump(SILAutoDiffIndices indices,
+  /// Prints activity information for the config of the given value.
+  void dump(SILValue value, AutoDiffConfig config,
+            llvm::raw_ostream &s = llvm::dbgs()) const {
+    return dump(value, config.parameterIndices, config.resultIndices, s);
+  }
+
+  /// Prints all activity information for the given parameter indices and result
+  /// indices.
+  void dump(IndexSubset *parameterIndices, IndexSubset *resultIndices,
             llvm::raw_ostream &s = llvm::dbgs()) const;
+
+  /// Prints all activity information for the given config.
+  void dump(AutoDiffConfig config, llvm::raw_ostream &s = llvm::dbgs()) const {
+    return dump(config.parameterIndices, config.resultIndices, s);
+  }
 };
 
 class DifferentiableActivityCollection {

include/swift/SILOptimizer/Differentiation/ADContext.h

@@ -45,9 +45,9 @@ namespace autodiff {
 
 /// Stores `apply` instruction information calculated by VJP generation.
 struct NestedApplyInfo {
-  /// The differentiation indices that are used to differentiate this `apply`
+  /// The differentiation config that is used to differentiate this `apply`
   /// instruction.
-  SILAutoDiffIndices indices;
+  AutoDiffConfig config;
   /// The original pullback type before reabstraction. `None` if the pullback
   /// type is not reabstracted.
   Optional<CanSILFunctionType> originalPullbackType;
@@ -120,6 +120,9 @@ class ADContext {
   /// Construct an ADContext for the given module.
   explicit ADContext(SILModuleTransform &transform);
 
+  // No copying.
+  ADContext(const ADContext &) = delete;
+
   //--------------------------------------------------------------------------//
   // General utilities
   //--------------------------------------------------------------------------//

include/swift/SILOptimizer/Differentiation/Common.h

@@ -119,7 +119,7 @@ void collectAllActualResultsInTypeOrder(
 /// - The set of minimal parameter and result indices for differentiating the
 ///   `apply` instruction.
 void collectMinimalIndicesForFunctionCall(
-    ApplyInst *ai, SILAutoDiffIndices parentIndices,
+    ApplyInst *ai, AutoDiffConfig parentConfig,
     const DifferentiableActivityInfo &activityInfo,
     SmallVectorImpl<SILValue> &results, SmallVectorImpl<unsigned> &paramIndices,
     SmallVectorImpl<unsigned> &resultIndices);

include/swift/SILOptimizer/Differentiation/LinearMapInfo.h

@@ -67,7 +67,7 @@ class LinearMapInfo {
   SILLoopInfo *loopInfo;
 
   /// Differentiation indices of the function.
-  const SILAutoDiffIndices indices;
+  const AutoDiffConfig config;
 
   /// Mapping from original basic blocks to linear map structs.
   llvm::DenseMap<SILBasicBlock *, StructDecl *> linearMapStructs;
@@ -149,7 +149,7 @@ class LinearMapInfo {
 
   explicit LinearMapInfo(ADContext &context, AutoDiffLinearMapKind kind,
                          SILFunction *original, SILFunction *derivative,
-                         SILAutoDiffIndices indices,
+                         AutoDiffConfig config,
                          const DifferentiableActivityInfo &activityInfo,
                          SILLoopInfo *loopInfo);
 

include/swift/SILOptimizer/Differentiation/Thunk.h

@@ -106,8 +106,8 @@ SILValue reabstractFunction(
 std::pair<SILFunction *, SubstitutionMap>
 getOrCreateSubsetParametersThunkForDerivativeFunction(
     SILOptFunctionBuilder &fb, SILValue origFnOperand, SILValue derivativeFn,
-    AutoDiffDerivativeFunctionKind kind, SILAutoDiffIndices desiredIndices,
-    SILAutoDiffIndices actualIndices);
+    AutoDiffDerivativeFunctionKind kind, AutoDiffConfig desiredConfig,
+    AutoDiffConfig actualConfig);
 
 /// Get or create a derivative function parameter index subset thunk from
 /// `actualIndices` to `desiredIndices` for the given associated function
@@ -119,7 +119,7 @@ getOrCreateSubsetParametersThunkForLinearMap(
     SILOptFunctionBuilder &fb, SILFunction *assocFn,
     CanSILFunctionType origFnType, CanSILFunctionType linearMapType,
     CanSILFunctionType targetType, AutoDiffDerivativeFunctionKind kind,
-    SILAutoDiffIndices desiredIndices, SILAutoDiffIndices actualIndices);
+    AutoDiffConfig desiredConfig, AutoDiffConfig actualConfig);
 
 } // end namespace autodiff
 

include/swift/SILOptimizer/Differentiation/VJPCloner.h

@@ -50,7 +50,7 @@ class VJPCloner final {
   SILFunction &getVJP() const;
   SILFunction &getPullback() const;
   SILDifferentiabilityWitness *getWitness() const;
-  const SILAutoDiffIndices getIndices() const;
+  AutoDiffConfig getConfig() const;
   DifferentiationInvoker getInvoker() const;
   LinearMapInfo &getPullbackInfo() const;
   SILLoopInfo *getLoopInfo() const;

lib/AST/ASTMangler.cpp

@@ -420,7 +420,7 @@ std::string ASTMangler::mangleAutoDiffDerivativeFunctionHelper(
     Buffer << "_vjp_";
     break;
   }
-  Buffer << config.getSILAutoDiffIndices().mangle();
+  Buffer << config.mangle();
   if (config.derivativeGenericSignature) {
     Buffer << '_';
     appendGenericSignature(config.derivativeGenericSignature);
@@ -445,7 +445,7 @@ std::string ASTMangler::mangleAutoDiffLinearMapHelper(
     Buffer << "_pullback_";
     break;
   }
-  Buffer << config.getSILAutoDiffIndices().mangle();
+  Buffer << config.mangle();
   if (config.derivativeGenericSignature) {
     Buffer << '_';
     appendGenericSignature(config.derivativeGenericSignature);
@@ -484,7 +484,7 @@ std::string ASTMangler::mangleAutoDiffGeneratedDeclaration(
     }
     break;
   }
-  Buffer << config.getSILAutoDiffIndices().mangle();
+  Buffer << config.mangle();
   if (config.derivativeGenericSignature) {
     Buffer << '_';
     appendGenericSignature(config.derivativeGenericSignature);