[AutoDiff] Fix a subset-parameters thunk parameter indexing crasher. (#25699)

The subset parameters thunk generation logic is not handling cases where the original JVP/JVP is already producing a subset-parameters linear map. This patch fixes it by computing indices w.r.t. the linear map from indices w.r.t. the original function. 

Resolves [TF-594](https://bugs.swift.org/browse/TF-594) and unblocks tensorflow/swift-apis#147.

Author: rxwei

lib/SILOptimizer/Mandatory/Differentiation.cpp

@@ -5853,6 +5853,10 @@ ADContext::getOrCreateSubsetParametersThunkForLinearMap(
     SILFunction *parentThunk, CanSILFunctionType linearMapType,
     CanSILFunctionType targetType, AutoDiffAssociatedFunctionKind kind,
     SILAutoDiffIndices desiredIndices, SILAutoDiffIndices actualIndices) {
+  LLVM_DEBUG(getADDebugStream() << "Getting a subset parameters thunk for " <<
+             linearMapType << " from " << actualIndices << " to " <<
+             desiredIndices << '\n');
+
   SubstitutionMap interfaceSubs = parentThunk->getForwardingSubstitutionMap();
   GenericEnvironment *genericEnv = parentThunk->getGenericEnvironment();
   auto thunkType = buildThunkType(
@@ -5931,6 +5935,32 @@ ADContext::getOrCreateSubsetParametersThunkForLinearMap(
     }
   };
 
+  // `actualIndices` and `desiredIndices` are with respect to the original
+  // function. However, the differential parameters and pullback results may
+  // already be w.r.t. a subset. We create a map between the original function's
+  // actual parameter indices and the linear map's actual indices.
+  // Example:
+  //   Original: (T0, T1, T2) -> R
+  //   Actual indices: 0, 2
+  //   Original differential: (T0, T2) -> R
+  //   Original pullback: R -> (T0, T2)
+  //   Desired indices w.r.t. original: 2
+  //   Desired indices w.r.t. linear map: 1
+  SmallVector<unsigned, 4> actualParamIndicesMap(
+      actualIndices.parameters->getCapacity(), UINT_MAX);
+  {
+    unsigned indexInBitVec = 0;
+    for (auto index : actualIndices.parameters->getIndices()) {
+      actualParamIndicesMap[index] = indexInBitVec;
+      indexInBitVec++;
+    }
+  }
+  auto mapOriginalParameterIndex = [&](unsigned index) -> unsigned {
+    auto mappedIndex = actualParamIndicesMap[index];
+    assert(mappedIndex < actualIndices.parameters->getCapacity());
+    return mappedIndex;
+  };
+
   switch (kind) {
   // Differential arguments are:
   // - All indirect results, followed by:
@@ -5955,7 +5985,8 @@ ADContext::getOrCreateSubsetParametersThunkForLinearMap(
       // Otherwise, construct and use a zero argument.
       else {
         auto zeroSILType =
-            linearMapType->getParameters()[i].getSILStorageType();
+            linearMapType->getParameters()[mapOriginalParameterIndex(i)]
+                .getSILStorageType();
         buildZeroArgument(zeroSILType);
       }
     }
@@ -5974,7 +6005,8 @@ ADContext::getOrCreateSubsetParametersThunkForLinearMap(
     };
     // Iterate over actual indices.
     for (unsigned i : actualIndices.parameters->getIndices()) {
-      auto resultInfo = linearMapType->getResults()[i];
+      auto resultInfo =
+          linearMapType->getResults()[mapOriginalParameterIndex(i)];
       // Skip direct results. Only indirect results are relevant as arguments.
       if (resultInfo.isFormalDirect())
         continue;
@@ -6022,14 +6054,15 @@ ADContext::getOrCreateSubsetParametersThunkForLinearMap(
     // - Do nothing if result is indirect.
     //   (It was already forwarded to the `apply` instruction).
     // - Push it to `results` if result is direct.
+    auto result = allResults[mapOriginalParameterIndex(i)];
     if (desiredIndices.isWrtParameter(i)) {
-      if (allResults[i]->getType().isAddress())
+      if (result->getType().isAddress())
         continue;
-      results.push_back(allResults[i]);
+      results.push_back(result);
     }
     // Otherwise, cleanup the unused results.
     else {
-      emitCleanup(builder, loc, allResults[i]);
+      emitCleanup(builder, loc, result);
     }
   }
   // Deallocate local allocations and return final direct result.
@@ -6047,6 +6080,11 @@ ADContext::getOrCreateSubsetParametersThunkForAssociatedFunction(
     SILValue origFnOperand, SILValue assocFn,
     AutoDiffAssociatedFunctionKind kind, SILAutoDiffIndices desiredIndices,
     SILAutoDiffIndices actualIndices) {
+  LLVM_DEBUG(getADDebugStream() << "Getting a subset parameters thunk for "
+             "associated function " << assocFn << " of the original function "
+             << origFnOperand << " from " << actualIndices << " to " <<
+             desiredIndices << '\n');
+
   auto origFnType = origFnOperand->getType().castTo<SILFunctionType>();
   auto &module = getModule();
   auto lookupConformance = LookUpConformanceInModule(module.getSwiftModule());
@@ -6368,6 +6406,11 @@ void ADContext::foldAutoDiffFunctionExtraction(AutoDiffFunctionInst *source) {
 }
 
 bool ADContext::processAutoDiffFunctionInst(AutoDiffFunctionInst *adfi) {
+  LLVM_DEBUG({
+    auto &s = getADDebugStream() << "Processing AutoDiffFunctionInst:\n";
+    adfi->printInContext(s);
+  });
+
   if (adfi->getNumAssociatedFunctions() ==
       autodiff::getNumAutoDiffAssociatedFunctions(
           adfi->getDifferentiationOrder()))

test/AutoDiff/superset_adjoint.swift

@@ -33,6 +33,14 @@ SupersetVJPTests.test("CrossModuleClosure") {
   expectEqual(1, gradient(at: Float(1)) { x in x + 2 })
 }
 
+SupersetVJPTests.test("SubsetOfSubset") {
+  @differentiable(wrt: (x, z))
+  func foo(_ x: Float, _ y: Float, _ z: Float) -> Float {
+    withoutDerivative(at: 0)
+  }
+  expectEqual(0, gradient(at: 0, in: { x in foo(x, 0, 0) }))
+}
+
 // FIXME: The expression `(+) as @differentiable (Float, @nondiff Float) -> Float)`
 // forms a curry thunk of `Float.+` before conversion to @differentiable, and AD
 // doesn't know how to differentiate the curry thunk, so it produces a