Minor edits.

Author: dan-zheng

lib/SILOptimizer/Mandatory/Differentiation.cpp

@@ -1816,7 +1816,6 @@ emitAssociatedFunctionReference(
           DifferentiationInvoker::Kind::IndirectDifferentiation)
         contextualRequirements = std::get<2>(
             invoker.getIndirectDifferentiation())->getRequirements();
-      // Parent function containing instruction
       auto *newAttr = context.getOrCreateDifferentiableAttr(
           originalFn, desiredIndices, contextualRequirements);
       bool error = context.processDifferentiableAttribute(originalFn, newAttr, invoker);
@@ -2084,8 +2083,23 @@ buildThunkSignature(SILFunction *fn,
 
 }
 
+/// The thunk kinds used in the differentiation transform.
 enum class DifferentiationThunkKind {
+  /// A reabstraction thunk.
+  ///
+  /// Reabstraction thunks transform a function-typed value to another one with
+  /// different parameter/result abstraction patterns. This is identical to the
+  /// thunks generated by SILGen.
   Reabstraction,
+
+  /// An index subset thunk.
+  ///
+  /// An index subset thunk is used transform JVP/VJPs into a version that is
+  /// "wrt" fewer differentiation parameters.
+  /// - Differentials of thunked JVPs use zero for non-requested differentiation
+  //    parameters.
+  /// - Pullbacks of thunked VJPs discard results for non-requested
+  ///   differentiation parameters.
   IndexSubset
 };
 
@@ -5822,10 +5836,6 @@ ADContext::getOrCreateAssociatedFunctionIndexSubsetThunk(
   assocSubstMap = assocSubstMap.subst(thunk->getForwardingSubstitutionMap());
   assocFnType = assocRef->getType().castTo<SILFunctionType>();
 
-  /*
-  auto arguments = map<SmallVector<SILValue, 4>>(
-      thunk->getArguments(), [](SILValue v) { return v; });
-  */
   SmallVector<SILValue, 4> arguments;
   arguments.append(thunk->getArguments().begin(), thunk->getArguments().end());
   auto *apply = builder.createApply(
@@ -5898,17 +5908,19 @@ SILValue ADContext::getCanonicalizedAutoDiffFunctionInst(
         // Construct new curry think.
         SILOptFunctionBuilder fb(transform);
         auto *newThunk = fb.getOrCreateFunction(
-            loc, newThunkName, getSpecializedLinkage(thunk, thunk->getLinkage()),
-            thunkType, thunk->isBare(), thunk->isTransparent(),
-            thunk->isSerialized(), thunk->isDynamicallyReplaceable(),
-            ProfileCounter(), thunk->isThunk());
+            loc, newThunkName,
+            getSpecializedLinkage(thunk, thunk->getLinkage()), thunkType,
+            thunk->isBare(), thunk->isTransparent(), thunk->isSerialized(),
+            thunk->isDynamicallyReplaceable(), ProfileCounter(),
+            thunk->isThunk());
         if (newThunk->empty()) {
           newThunk->setOwnershipEliminated();
           SILFunctionCloner cloner(newThunk);
           cloner.cloneFunction(thunk);
         }
 
-        auto *retInst = cast<ReturnInst>(newThunk->findReturnBB()->getTerminator());
+        auto *retInst =
+            cast<ReturnInst>(newThunk->findReturnBB()->getTerminator());
         AutoDiffFunctionInst *adfi;
         {
           SILBuilder builder(retInst);