Thunk.cpp

//===--- Thunk.cpp - Automatic differentiation thunks ---------*- C++ -*---===//
//
// This source file is part of the Swift.org open source project
//
// Copyright (c) 2019 - 2020 Apple Inc. and the Swift project authors
// Licensed under Apache License v2.0 with Runtime Library Exception
//
// See https://swift.org/LICENSE.txt for license information
// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
//
//===----------------------------------------------------------------------===//
//
// Automatic differentiation thunk generation utilities.
//
//===----------------------------------------------------------------------===//

#define DEBUG_TYPE "differentiation"

#include "swift/SILOptimizer/Differentiation/Thunk.h"
#include "swift/SILOptimizer/Differentiation/Common.h"

#include "swift/AST/AnyFunctionRef.h"
#include "swift/AST/Requirement.h"
#include "swift/AST/SubstitutionMap.h"
#include "swift/AST/TypeCheckRequests.h"
#include "swift/SILOptimizer/Utils/SILOptFunctionBuilder.h"
#include "swift/SILOptimizer/Utils/DifferentiationMangler.h"

namespace swift {
namespace autodiff {

//===----------------------------------------------------------------------===//
// Thunk helpers
//===----------------------------------------------------------------------===//
// These helpers are copied/adapted from SILGen. They should be refactored and
// moved to a shared location.
//===----------------------------------------------------------------------===//

CanSILFunctionType buildThunkType(SILFunction *fn,
                                  CanSILFunctionType &sourceType,
                                  CanSILFunctionType &expectedType,
                                  GenericEnvironment *&genericEnv,
                                  SubstitutionMap &interfaceSubs,
                                  bool withoutActuallyEscaping,
                                  DifferentiationThunkKind thunkKind) {
  CanType inputSubstType;
  CanType outputSubstType;
  CanType dynamicSelfType;
  return buildSILFunctionThunkType(
      fn, sourceType, expectedType, inputSubstType, outputSubstType, genericEnv,
      interfaceSubs, dynamicSelfType, withoutActuallyEscaping, thunkKind);
}

/// Forward function arguments, handling ownership convention mismatches.
/// Adapted from `forwardFunctionArguments` in SILGenPoly.cpp.
///
/// Forwarded arguments are appended to `forwardedArgs`.
///
/// Local allocations are appended to `localAllocations`. They need to be
/// deallocated via `dealloc_stack`.
///
/// Local values requiring cleanup are appended to `valuesToCleanup`.
static void forwardFunctionArgumentsConvertingOwnership(
    SILBuilder &builder, SILLocation loc, CanSILFunctionType fromTy,
    CanSILFunctionType toTy, ArrayRef<SILArgument *> originalArgs,
    SmallVectorImpl<SILValue> &forwardedArgs,
    SmallVectorImpl<AllocStackInst *> &localAllocations,
    SmallVectorImpl<SILValue> &valuesToCleanup) {
  auto fromParameters = fromTy->getParameters();
  auto toParameters = toTy->getParameters();
  assert(fromParameters.size() == toParameters.size());
  assert(fromParameters.size() == originalArgs.size());
  for (auto index : indices(originalArgs)) {
    auto &arg = originalArgs[index];
    auto fromParam = fromParameters[index];
    auto toParam = toParameters[index];
    // To convert guaranteed argument to be owned, create a copy.
    if (fromParam.isConsumed() && !toParam.isConsumed()) {
      // If the argument has an object type, create a `copy_value`.
      if (arg->getType().isObject()) {
        auto argCopy = builder.emitCopyValueOperation(loc, arg);
        forwardedArgs.push_back(argCopy);
        continue;
      }
      // If the argument has an address type, create a local allocation and
      // `copy_addr` its contents to the local allocation.
      auto *alloc = builder.createAllocStack(loc, arg->getType());
      builder.createCopyAddr(loc, arg, alloc, IsNotTake, IsInitialization);
      localAllocations.push_back(alloc);
      forwardedArgs.push_back(alloc);
      continue;
    }
    // To convert owned argument to be guaranteed, borrow the argument.
    if (fromParam.isGuaranteed() && !toParam.isGuaranteed()) {
      auto bbi = builder.emitBeginBorrowOperation(loc, arg);
      forwardedArgs.push_back(bbi);
      valuesToCleanup.push_back(bbi);
      valuesToCleanup.push_back(arg);
      continue;
    }
    // Otherwise, simply forward the argument.
    forwardedArgs.push_back(arg);
  }
}

SILFunction *getOrCreateReabstractionThunk(SILOptFunctionBuilder &fb,
                                           SILModule &module, SILLocation loc,
                                           SILFunction *caller,
                                           CanSILFunctionType fromType,
                                           CanSILFunctionType toType) {
  assert(!fromType->getCombinedSubstitutions());
  assert(!toType->getCombinedSubstitutions());

  SubstitutionMap interfaceSubs;
  GenericEnvironment *genericEnv = nullptr;
  auto thunkType =
      buildThunkType(caller, fromType, toType, genericEnv, interfaceSubs,
                     /*withoutActuallyEscaping*/ false,
                     DifferentiationThunkKind::Reabstraction);
  auto thunkDeclType =
      thunkType->getWithExtInfo(thunkType->getExtInfo().withNoEscape(false));

  auto fromInterfaceType = fromType->mapTypeOutOfContext()->getCanonicalType();
  auto toInterfaceType = toType->mapTypeOutOfContext()->getCanonicalType();

  Mangle::ASTMangler mangler;
  std::string name = mangler.mangleReabstractionThunkHelper(
      thunkType, fromInterfaceType, toInterfaceType, Type(), Type(),
      module.getSwiftModule());

  auto *thunk = fb.getOrCreateSharedFunction(
      loc, name, thunkDeclType, IsBare, IsTransparent, IsSerialized,
      ProfileCounter(), IsReabstractionThunk, IsNotDynamic, IsNotDistributed,
      IsNotRuntimeAccessible);
  if (!thunk->empty())
    return thunk;

  thunk->setGenericEnvironment(genericEnv);
  auto *entry = thunk->createBasicBlock();
  SILBuilder builder(entry);
  createEntryArguments(thunk);

  SILFunctionConventions fromConv(fromType, module);
  SILFunctionConventions toConv(toType, module);
  assert(toConv.useLoweredAddresses());

  // Forward thunk arguments, handling ownership convention mismatches.
  SmallVector<SILValue, 4> forwardedArgs;
  for (auto indRes : thunk->getIndirectResults())
    forwardedArgs.push_back(indRes);
  SmallVector<AllocStackInst *, 4> localAllocations;
  SmallVector<SILValue, 4> valuesToCleanup;
  forwardFunctionArgumentsConvertingOwnership(
      builder, loc, fromType, toType,
      thunk->getArgumentsWithoutIndirectResults().drop_back(), forwardedArgs,
      localAllocations, valuesToCleanup);

  SmallVector<SILValue, 4> arguments;
  auto toArgIter = forwardedArgs.begin();
  auto useNextArgument = [&]() { arguments.push_back(*toArgIter++); };

  auto createAllocStack = [&](SILType type) {
    auto *alloc = builder.createAllocStack(loc, type);
    localAllocations.push_back(alloc);
    return alloc;
  };

  // Handle indirect results.
  assert(fromType->getNumResults() == toType->getNumResults());
  for (unsigned resIdx : range(toType->getNumResults())) {
    auto fromRes = fromConv.getResults()[resIdx];
    auto toRes = toConv.getResults()[resIdx];
    // No abstraction mismatch.
    if (fromRes.isFormalIndirect() == toRes.isFormalIndirect()) {
      // If result types are indirect, directly pass as next argument.
      if (toRes.isFormalIndirect())
        useNextArgument();
      continue;
    }
    // Convert indirect result to direct result.
    if (fromRes.isFormalIndirect()) {
      SILType resultTy =
          fromConv.getSILType(fromRes, builder.getTypeExpansionContext());
      assert(resultTy.isAddress());
      auto *indRes = createAllocStack(resultTy);
      arguments.push_back(indRes);
      continue;
    }
    // Convert direct result to indirect result.
    // Increment thunk argument iterator; reabstraction handled later.
    ++toArgIter;
  }

  // Reabstract parameters.
  assert(toType->getNumParameters() == fromType->getNumParameters());
  for (unsigned paramIdx : range(toType->getNumParameters())) {
    auto fromParam = fromConv.getParameters()[paramIdx];
    auto toParam = toConv.getParameters()[paramIdx];
    // No abstraction mismatch. Directly use next argument.
    if (fromParam.isFormalIndirect() == toParam.isFormalIndirect()) {
      useNextArgument();
      continue;
    }
    // Convert indirect parameter to direct parameter.
    if (fromParam.isFormalIndirect()) {
      auto paramTy = fromConv.getSILType(fromType->getParameters()[paramIdx],
                                         builder.getTypeExpansionContext());
      if (!paramTy.hasArchetype())
        paramTy = thunk->mapTypeIntoContext(paramTy);
      assert(paramTy.isAddress());
      auto toArg = *toArgIter++;
      auto *buf = createAllocStack(toArg->getType());
      toArg = builder.emitCopyValueOperation(loc, toArg);
      builder.emitStoreValueOperation(loc, toArg, buf,
                                      StoreOwnershipQualifier::Init);
      valuesToCleanup.push_back(buf);
      arguments.push_back(buf);
      continue;
    }
    // Convert direct parameter to indirect parameter.
    assert(toParam.isFormalIndirect());
    auto toArg = *toArgIter++;
    auto load = builder.emitLoadBorrowOperation(loc, toArg);
    if (isa<LoadBorrowInst>(load))
      valuesToCleanup.push_back(load);
    arguments.push_back(load);
  }

  auto *fnArg = thunk->getArgumentsWithoutIndirectResults().back();
  auto *apply = builder.createApply(loc, fnArg, SubstitutionMap(), arguments);

  // Get return elements.
  SmallVector<SILValue, 4> results;
  // Extract all direct results.
  SmallVector<SILValue, 4> directResults;
  extractAllElements(apply, builder, directResults);

  auto fromDirResultsIter = directResults.begin();
  auto fromIndResultsIter = apply->getIndirectSILResults().begin();
  auto toIndResultsIter = thunk->getIndirectResults().begin();
  // Reabstract results.
  for (unsigned resIdx : range(toType->getNumResults())) {
    auto fromRes = fromConv.getResults()[resIdx];
    auto toRes = toConv.getResults()[resIdx];
    // Check function-typed results.
    if (isa<SILFunctionType>(fromRes.getInterfaceType()) &&
        isa<SILFunctionType>(toRes.getInterfaceType())) {
      auto fromFnType = cast<SILFunctionType>(fromRes.getInterfaceType());
      auto toFnType = cast<SILFunctionType>(toRes.getInterfaceType());
      auto fromUnsubstFnType = fromFnType->getUnsubstitutedType(module);
      auto toUnsubstFnType = toFnType->getUnsubstitutedType(module);
      // If unsubstituted function types are not equal, perform reabstraction.
      if (fromUnsubstFnType != toUnsubstFnType) {
        auto fromFn = *fromDirResultsIter++;
        auto newFromFn = reabstractFunction(
            builder, fb, loc, fromFn, toFnType,
            [](SubstitutionMap substMap) { return substMap; });
        results.push_back(newFromFn);
        continue;
      }
    }
    // No abstraction mismatch.
    if (fromRes.isFormalIndirect() == toRes.isFormalIndirect()) {
      // If result types are direct, add call result as direct thunk result.
      if (toRes.isFormalDirect())
        results.push_back(*fromDirResultsIter++);
      // If result types are indirect, increment indirect result iterators.
      else {
        ++fromIndResultsIter;
        ++toIndResultsIter;
      }
      continue;
    }
    // Load direct results from indirect results.
    if (fromRes.isFormalIndirect()) {
      auto indRes = *fromIndResultsIter++;
      auto load = builder.emitLoadValueOperation(loc, indRes,
                                                 LoadOwnershipQualifier::Take);
      results.push_back(load);
      continue;
    }
    // Store direct results to indirect results.
    assert(toRes.isFormalIndirect());
#ifndef NDEBUG
    SILType resultTy =
        toConv.getSILType(toRes, builder.getTypeExpansionContext());
    assert(resultTy.isAddress());
#endif
    auto indRes = *toIndResultsIter++;
    auto dirRes = *fromDirResultsIter++;
    builder.emitStoreValueOperation(loc, dirRes, indRes,
                                    StoreOwnershipQualifier::Init);
  }
  auto retVal = joinElements(results, builder, loc);

  // Clean up local values.
  // Guaranteed values need an `end_borrow`.
  // Owned values need to be destroyed.
  for (auto arg : valuesToCleanup) {
    switch (arg->getOwnershipKind()) {
    case OwnershipKind::Any:
      llvm_unreachable("value with any ownership kind?!");
    case OwnershipKind::Guaranteed:
      builder.emitEndBorrowOperation(loc, arg);
      break;
    case OwnershipKind::Owned:
    case OwnershipKind::Unowned:
    case OwnershipKind::None:
      builder.emitDestroyOperation(loc, arg);
      break;
    }
  }

  // Deallocate local allocations.
  for (auto *alloc : llvm::reverse(localAllocations))
    builder.createDeallocStack(loc, alloc);

  // Create return.
  builder.createReturn(loc, retVal);

  LLVM_DEBUG(auto &s = getADDebugStream() << "Created reabstraction thunk.\n";
             s << "  From type: " << fromType << '\n';
             s << "  To type: " << toType << '\n'; s << '\n'
                                                     << *thunk);

  return thunk;
}

SILValue reabstractFunction(
    SILBuilder &builder, SILOptFunctionBuilder &fb, SILLocation loc,
    SILValue fn, CanSILFunctionType toType,
    std::function<SubstitutionMap(SubstitutionMap)> remapSubstitutions) {
  auto &module = *fn->getModule();
  auto fromType = fn->getType().getAs<SILFunctionType>();
  auto unsubstFromType = fromType->getUnsubstitutedType(module);
  auto unsubstToType = toType->getUnsubstitutedType(module);

  auto *thunk = getOrCreateReabstractionThunk(fb, module, loc,
                                              /*caller*/ fn->getFunction(),
                                              unsubstFromType, unsubstToType);
  auto *thunkRef = builder.createFunctionRef(loc, thunk);

  if (fromType != unsubstFromType)
    fn = builder.createConvertFunction(
        loc, fn, SILType::getPrimitiveObjectType(unsubstFromType),
        /*withoutActuallyEscaping*/ false);

  fn = builder.createPartialApply(
      loc, thunkRef, remapSubstitutions(thunk->getForwardingSubstitutionMap()),
      {fn}, fromType->getCalleeConvention());

  if (toType != unsubstToType)
    fn = builder.createConvertFunction(loc, fn,
                                       SILType::getPrimitiveObjectType(toType),
                                       /*withoutActuallyEscaping*/ false);

  return fn;
}

std::pair<SILFunction *, SubstitutionMap>
getOrCreateSubsetParametersThunkForLinearMap(
    SILOptFunctionBuilder &fb, SILFunction *parentThunk,
    CanSILFunctionType origFnType, CanSILFunctionType linearMapType,
    CanSILFunctionType targetType, AutoDiffDerivativeFunctionKind kind,
    const AutoDiffConfig &desiredConfig, const AutoDiffConfig &actualConfig,
    ADContext &adContext) {
  LLVM_DEBUG(getADDebugStream()
             << "Getting a subset parameters thunk for "
             << (kind == AutoDiffDerivativeFunctionKind::JVP ? "jvp" : "vjp")
             << " linear map " << linearMapType
             << " from " << actualConfig << " to " << desiredConfig << '\n');

  assert(!linearMapType->getCombinedSubstitutions());
  assert(!targetType->getCombinedSubstitutions());
  SubstitutionMap interfaceSubs;
  GenericEnvironment *genericEnv = nullptr;
  auto thunkType = buildThunkType(parentThunk, linearMapType, targetType,
                                  genericEnv, interfaceSubs,
                                  /*withoutActuallyEscaping*/ true,
                                  DifferentiationThunkKind::Reabstraction);

  Mangle::DifferentiationMangler mangler;
  auto fromInterfaceType =
      linearMapType->mapTypeOutOfContext()->getCanonicalType();

  auto thunkName = mangler.mangleLinearMapSubsetParametersThunk(
      fromInterfaceType, kind.getLinearMapKind(),
      actualConfig.parameterIndices, actualConfig.resultIndices,
      desiredConfig.parameterIndices);

  auto loc = parentThunk->getLocation();
  auto *thunk = fb.getOrCreateSharedFunction(
      loc, thunkName, thunkType, IsBare, IsTransparent, IsSerialized,
      ProfileCounter(), IsThunk, IsNotDynamic, IsNotDistributed,
      IsNotRuntimeAccessible);

  if (!thunk->empty())
    return {thunk, interfaceSubs};

  thunk->setGenericEnvironment(genericEnv);
  auto *entry = thunk->createBasicBlock();
  TangentBuilder builder(entry, adContext);
  createEntryArguments(thunk);

  // Get arguments.
  SmallVector<SILValue, 4> arguments;
  SmallVector<AllocStackInst *, 4> localAllocations;
  SmallVector<SILValue, 4> valuesToCleanup;
  auto cleanupValues = [&]() {
    for (auto value : llvm::reverse(valuesToCleanup))
      builder.emitDestroyOperation(loc, value);

    for (auto *alloc : llvm::reverse(localAllocations))
      builder.createDeallocStack(loc, alloc);
  };

  // Build a `.zero` argument for the given `Differentiable`-conforming type.
  auto buildZeroArgument = [&](SILParameterInfo zeroSILParameter) {
    auto zeroSILType = zeroSILParameter.getSILStorageInterfaceType();
    auto zeroSILObjType = zeroSILType.getObjectType();
    auto zeroType = zeroSILType.getASTType();
    auto *swiftMod = parentThunk->getModule().getSwiftModule();
    auto tangentSpace =
        zeroType->getAutoDiffTangentSpace(LookUpConformanceInModule(swiftMod));
    assert(tangentSpace && "No tangent space for this type");
    switch (tangentSpace->getKind()) {
    case TangentSpace::Kind::TangentVector: {
      auto *buf = builder.createAllocStack(loc, zeroSILObjType);
      localAllocations.push_back(buf);
      builder.emitZeroIntoBuffer(loc, buf, IsInitialization);
      if (zeroSILType.isAddress()) {
        arguments.push_back(buf);
        if (zeroSILParameter.isGuaranteed()) {
          valuesToCleanup.push_back(buf);
        }
      } else {
        auto arg = builder.emitLoadValueOperation(loc, buf,
                                                  LoadOwnershipQualifier::Take);
        arguments.push_back(arg);
        if (zeroSILParameter.isGuaranteed()) {
          valuesToCleanup.push_back(arg);
        }
      }
      break;
    }
    case TangentSpace::Kind::Tuple: {
      llvm_unreachable("Unimplemented: Handle zero initialization for tuples");
    }
    }
  };

  // The indices in `actualConfig` and `desiredConfig` 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(
      actualConfig.parameterIndices->getCapacity(), UINT_MAX);
  {
    unsigned indexInBitVec = 0;
    for (auto index : actualConfig.parameterIndices->getIndices()) {
      actualParamIndicesMap[index] = indexInBitVec;
      ++indexInBitVec;
    }
  }
  auto mapOriginalParameterIndex = [&](unsigned index) -> unsigned {
    auto mappedIndex = actualParamIndicesMap[index];
    assert(mappedIndex < actualConfig.parameterIndices->getCapacity());
    return mappedIndex;
  };

  auto toIndirectResultsIter = thunk->getIndirectResults().begin();
  auto useNextIndirectResult = [&]() {
      assert(toIndirectResultsIter != thunk->getIndirectResults().end());
      arguments.push_back(*toIndirectResultsIter++);
  };

  switch (kind) {
  // Differential arguments are:
  // - All indirect results, followed by:
  // - An interleaving of:
  //   - Thunk arguments (when parameter index is in both desired and actual
  //     indices).
  //   - Zeros (when parameter is not in desired indices).
  case AutoDiffDerivativeFunctionKind::JVP: {
    unsigned numIndirectResults = linearMapType->getNumIndirectFormalResults();
    // Forward desired indirect results
    for (unsigned idx : *actualConfig.resultIndices) {
      if (idx >= numIndirectResults)
        break;

      auto resultInfo = linearMapType->getResults()[idx];
      assert(idx < linearMapType->getNumResults());

      // Forward result argument in case we do not need to thunk it away
      if (desiredConfig.resultIndices->contains(idx)) {
        useNextIndirectResult();
        continue;
      }

      // Otherwise, allocate and use an uninitialized indirect result
      auto *indirectResult = builder.createAllocStack(
        loc, resultInfo.getSILStorageInterfaceType());
      localAllocations.push_back(indirectResult);
      arguments.push_back(indirectResult);
    }
    assert(toIndirectResultsIter == thunk->getIndirectResults().end());

    auto toArgIter = thunk->getArgumentsWithoutIndirectResults().begin();
    auto useNextArgument = [&]() { arguments.push_back(*toArgIter++); };
    // Iterate over actual indices.
    for (unsigned i : actualConfig.parameterIndices->getIndices()) {
      // If index is desired, use next argument.
      if (desiredConfig.isWrtParameter(i)) {
        useNextArgument();
      }
      // Otherwise, construct and use a zero argument.
      else {
        auto zeroSILParameter =
            linearMapType->getParameters()[mapOriginalParameterIndex(i)];
        buildZeroArgument(zeroSILParameter);
      }
    }
    break;
  }
  // Pullback arguments are:
  // - An interleaving of:
  //   - Thunk indirect results (when parameter index is in both desired and
  //     actual indices).
  //   - Zeros (when parameter is not in desired indices).
  // - All actual arguments.
  case AutoDiffDerivativeFunctionKind::VJP: {
    // Collect pullback arguments.
    unsigned pullbackResultIndex = 0;
    for (unsigned i : actualConfig.parameterIndices->getIndices()) {
      auto origParamInfo = origFnType->getParameters()[i];
      // Skip original semantic result parameters. All non-indirect-result pullback
      // arguments (including semantic result` arguments) are appended to `arguments`
      // later.
      if (origParamInfo.isAutoDiffSemanticResult())
        continue;
      auto resultInfo = linearMapType->getResults()[pullbackResultIndex];
      assert(pullbackResultIndex < linearMapType->getNumResults());
      ++pullbackResultIndex;
      // Skip pullback direct results. Only indirect results are relevant as
      // arguments.
      if (resultInfo.isFormalDirect())
        continue;
      // If index is desired, use next pullback indirect result.
      if (desiredConfig.isWrtParameter(i)) {
        useNextIndirectResult();
        continue;
      }
      // Otherwise, allocate and use an uninitialized pullback indirect result.
      auto *indirectResult = builder.createAllocStack(
          loc, resultInfo.getSILStorageInterfaceType());
      localAllocations.push_back(indirectResult);
      arguments.push_back(indirectResult);
    }
    // Forward all actual non-indirect-result arguments.
    auto thunkArgs = thunk->getArgumentsWithoutIndirectResults();
    // Slice out the function to be called
    thunkArgs = thunkArgs.slice(0, thunkArgs.size() - 1);
    unsigned thunkArg = 0;
    for (unsigned idx : *actualConfig.resultIndices) {
      // Forward result argument in case we do not need to thunk it away
      if (desiredConfig.resultIndices->contains(idx))
        arguments.push_back(thunkArgs[thunkArg++]);
      else // otherwise, zero it out
        buildZeroArgument(linearMapType->getParameters()[arguments.size()]);
    }
    break;
  }
  }

  // Get the linear map thunk argument and apply it.
  auto *linearMap = thunk->getArguments().back();
  auto *ai = builder.createApply(loc, linearMap, SubstitutionMap(), arguments);

  // If differential thunk, deallocate local allocations and directly return
  // `apply` result (if it is desired).
  if (kind == AutoDiffDerivativeFunctionKind::JVP) {
    SmallVector<SILValue, 8> differentialDirectResults;
    extractAllElements(ai, builder, differentialDirectResults);
    SmallVector<SILValue, 8> allResults;
    collectAllActualResultsInTypeOrder(ai, differentialDirectResults, allResults);
    unsigned numResults = thunk->getConventions().getNumDirectSILResults() +
     thunk->getConventions().getNumDirectSILResults();
    SmallVector<SILValue, 8> results;
    for (unsigned idx : *actualConfig.resultIndices) {
      if (idx >= numResults)
        break;

      auto result = allResults[idx];
      if (desiredConfig.isWrtResult(idx))
        results.push_back(result);
      else {
        if (result->getType().isAddress())
          builder.emitDestroyAddrAndFold(loc, result);
        else
          builder.emitDestroyValueOperation(loc, result);
      }
    }

    cleanupValues();
    auto result = joinElements(results, builder, loc);
    builder.createReturn(loc, result);
    return {thunk, interfaceSubs};
  }

  // If pullback thunk, return only the desired results and clean up the
  // undesired results.
  SmallVector<SILValue, 8> pullbackDirectResults;
  extractAllElements(ai, builder, pullbackDirectResults);
  SmallVector<SILValue, 8> allResults;
  collectAllActualResultsInTypeOrder(ai, pullbackDirectResults, allResults);
  // Collect pullback semantic result arguments in type order.
  unsigned semanticResultArgIdx = 0;
  SILFunctionConventions origConv(origFnType, thunk->getModule());
  for (auto paramIdx : actualConfig.parameterIndices->getIndices()) {
    auto paramInfo = origConv.getParameters()[paramIdx];
    if (!paramInfo.isAutoDiffSemanticResult())
      continue;
    auto semanticResultArg =
      *std::next(ai->getAutoDiffSemanticResultArguments().begin(),
                 semanticResultArgIdx++);
    unsigned mappedParamIdx = mapOriginalParameterIndex(paramIdx);
    allResults.insert(allResults.begin() + mappedParamIdx, semanticResultArg);
  }
  assert(allResults.size() == actualConfig.parameterIndices->getNumIndices() &&
         "Number of pullback results should match number of differentiability "
         "parameters");

  SmallVector<SILValue, 8> results;
  for (unsigned i : actualConfig.parameterIndices->getIndices()) {
    unsigned mappedIndex = mapOriginalParameterIndex(i);
    // If result is desired:
    // - 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[mappedIndex];
    if (desiredConfig.isWrtParameter(i)) {
      if (result->getType().isObject())
        results.push_back(result);
    }
    // Otherwise, cleanup the unused results.
    else {
      if (result->getType().isAddress())
        builder.emitDestroyAddrAndFold(loc, result);
      else
        builder.emitDestroyValueOperation(loc, result);
    }
  }
  // Deallocate local allocations and return final direct result.
  cleanupValues();
  auto result = joinElements(results, builder, loc);
  builder.createReturn(loc, result);

  return {thunk, interfaceSubs};
}

std::pair<SILFunction *, SubstitutionMap>
getOrCreateSubsetParametersThunkForDerivativeFunction(
    SILOptFunctionBuilder &fb, SILValue origFnOperand, SILValue derivativeFn,
    AutoDiffDerivativeFunctionKind kind, const AutoDiffConfig &desiredConfig,
    const AutoDiffConfig &actualConfig, ADContext &adContext) {
  LLVM_DEBUG(getADDebugStream()
             << "Getting a subset parameters thunk for derivative "
             << (kind == AutoDiffDerivativeFunctionKind::JVP ? "jvp" : "vjp")
             << " function " << derivativeFn
             << " of the original function " << origFnOperand
             << " from " << actualConfig << " to " << desiredConfig << '\n');

  auto origFnType = origFnOperand->getType().castTo<SILFunctionType>();
  auto &module = fb.getModule();
  auto lookupConformance = LookUpConformanceInModule(module.getSwiftModule());

  // Compute target type for thunking.
  auto derivativeFnType = derivativeFn->getType().castTo<SILFunctionType>();
  auto targetType = origFnType->getAutoDiffDerivativeFunctionType(
      desiredConfig.parameterIndices, desiredConfig.resultIndices, kind,
      module.Types, lookupConformance);
  auto *caller = derivativeFn->getFunction();
  if (targetType->hasArchetype()) {
    auto substTargetType =
        caller->mapTypeIntoContext(targetType->mapTypeOutOfContext())
            ->getCanonicalType();
    targetType = SILType::getPrimitiveObjectType(substTargetType)
                     .castTo<SILFunctionType>();
  }
  assert(derivativeFnType->getNumParameters() ==
         targetType->getNumParameters());
  assert(derivativeFnType->getNumResults() == targetType->getNumResults());

  // Build thunk type.
  SubstitutionMap interfaceSubs;
  GenericEnvironment *genericEnv = nullptr;
  auto thunkType = buildThunkType(derivativeFn->getFunction(), derivativeFnType,
                                  targetType, genericEnv, interfaceSubs,
                                  /*withoutActuallyEscaping*/ false,
                                  DifferentiationThunkKind::IndexSubset);

  // FIXME: The logic for resolving `assocRef` does not reapply function
  // conversions, which is problematic if `derivativeFn` is a `partial_apply`
  // instruction.
  StringRef origName;
  if (auto *origFnRef =
          peerThroughFunctionConversions<FunctionRefInst>(origFnOperand)) {
    origName = origFnRef->getReferencedFunction()->getName();
  } else if (auto *origMethodInst =
                 peerThroughFunctionConversions<MethodInst>(origFnOperand)) {
    origName = origMethodInst->getMember()
                   .getAnyFunctionRef()
                   ->getAbstractFunctionDecl()
                   ->getNameStr();
  }
  assert(!origName.empty() && "Original function name could not be resolved");
  Mangle::DifferentiationMangler mangler;
  auto thunkName = mangler.mangleDerivativeFunctionSubsetParametersThunk(
      origName, targetType->mapTypeOutOfContext()->getCanonicalType(),
      kind, actualConfig.parameterIndices, actualConfig.resultIndices,
      desiredConfig.parameterIndices);

  auto loc = origFnOperand.getLoc();
  auto *thunk = fb.getOrCreateSharedFunction(
      loc, thunkName, thunkType, IsBare, IsTransparent, caller->isSerialized(),
      ProfileCounter(), IsThunk, IsNotDynamic, IsNotDistributed,
      IsNotRuntimeAccessible);

  if (!thunk->empty())
    return {thunk, interfaceSubs};

  thunk->setGenericEnvironment(genericEnv);
  auto *entry = thunk->createBasicBlock();
  SILBuilder builder(entry);
  createEntryArguments(thunk);

  SubstitutionMap assocSubstMap;
  if (auto *partialApply = dyn_cast<PartialApplyInst>(derivativeFn))
    assocSubstMap = partialApply->getSubstitutionMap();

  // FIXME: The logic for resolving `assocRef` does not reapply function
  // conversions, which is problematic if `derivativeFn` is a `partial_apply`
  // instruction.
  SILValue assocRef;
  if (auto *derivativeFnRef =
          peerThroughFunctionConversions<FunctionRefInst>(derivativeFn)) {
    auto *assoc = derivativeFnRef->getReferencedFunction();
    assocRef = builder.createFunctionRef(loc, assoc);
  } else if (auto *assocMethodInst =
                 peerThroughFunctionConversions<WitnessMethodInst>(
                     derivativeFn)) {
    assocRef = builder.createWitnessMethod(
        loc, assocMethodInst->getLookupType(),
        assocMethodInst->getConformance(), assocMethodInst->getMember(),
        thunk->mapTypeIntoContext(assocMethodInst->getType()));
  } else if (auto *assocMethodInst =
                 peerThroughFunctionConversions<ClassMethodInst>(
                     derivativeFn)) {
    auto classOperand = thunk->getArgumentsWithoutIndirectResults().back();
#ifndef NDEBUG
    auto classOperandType = assocMethodInst->getOperand()->getType();
    assert(classOperand->getType() == classOperandType);
#endif
    assocRef = builder.createClassMethod(
        loc, classOperand, assocMethodInst->getMember(),
        thunk->mapTypeIntoContext(assocMethodInst->getType()));
  } else if (auto *diffWitFn = peerThroughFunctionConversions<
                 DifferentiabilityWitnessFunctionInst>(derivativeFn)) {
    assocRef = builder.createDifferentiabilityWitnessFunction(
        loc, diffWitFn->getWitnessKind(), diffWitFn->getWitness());
  }
  assert(assocRef && "Expected derivative function to be resolved");

  assocSubstMap = assocSubstMap.subst(thunk->getForwardingSubstitutionMap());
  derivativeFnType = assocRef->getType().castTo<SILFunctionType>();

  SmallVector<SILValue, 4> arguments;
  arguments.append(thunk->getArguments().begin(), thunk->getArguments().end());
  assert(arguments.size() ==
         derivativeFnType->getNumParameters() +
             derivativeFnType->getNumIndirectFormalResults());
  auto *apply = builder.createApply(loc, assocRef, assocSubstMap, arguments);

  // Extract all direct results.
  SmallVector<SILValue, 8> directResults;
  extractAllElements(apply, builder, directResults);
  auto linearMap = directResults.back();
  directResults.pop_back();

  auto linearMapType = linearMap->getType().castTo<SILFunctionType>();
  auto linearMapTargetType = targetType->getResults()
                                 .back()
                                 .getSILStorageInterfaceType()
                                 .castTo<SILFunctionType>();
  auto unsubstLinearMapType = linearMapType->getUnsubstitutedType(module);
  auto unsubstLinearMapTargetType =
      linearMapTargetType->getUnsubstitutedType(module);

  SILFunction *linearMapThunk;
  SubstitutionMap linearMapSubs;
  std::tie(linearMapThunk, linearMapSubs) =
      getOrCreateSubsetParametersThunkForLinearMap(
          fb, thunk, origFnType, unsubstLinearMapType,
          unsubstLinearMapTargetType, kind, desiredConfig, actualConfig,
          adContext);

  auto *linearMapThunkFRI = builder.createFunctionRef(loc, linearMapThunk);
  SILValue thunkedLinearMap = linearMap;
  if (linearMapType != unsubstLinearMapType) {
    thunkedLinearMap = builder.createConvertFunction(
        loc, thunkedLinearMap,
        SILType::getPrimitiveObjectType(unsubstLinearMapType),
        /*withoutActuallyEscaping*/ false);
  }
  thunkedLinearMap = builder.createPartialApply(
      loc, linearMapThunkFRI, linearMapSubs, {thunkedLinearMap},
      ParameterConvention::Direct_Guaranteed);
  if (linearMapTargetType != unsubstLinearMapTargetType) {
    thunkedLinearMap = builder.createConvertFunction(
        loc, thunkedLinearMap,
        SILType::getPrimitiveObjectType(linearMapTargetType),
        /*withoutActuallyEscaping*/ false);
  }
  assert(origFnType->getNumAutoDiffSemanticResults() > 0);
  if (origFnType->getNumResults() > 0 &&
      origFnType->getResults().front().isFormalDirect()) {
    directResults.push_back(thunkedLinearMap);
    auto result = joinElements(directResults, builder, loc);
    builder.createReturn(loc, result);
  } else {
    builder.createReturn(loc, thunkedLinearMap);
  }

  return {thunk, interfaceSubs};
}

} // end namespace autodiff
} // end namespace swift