Reapply "Implement several peephole optimizations to unblock further optimizations of autodiff code" with correctness fix (#62012)

* Implement several peephole optimizations to unblock further optimizations of autodiff code

1. Simplify differentiable_function_extract of differentiable_function.
Before:
%x = differentiable_function(%orig, %jvp, %vjp)
%y = differentiable_function_extract [original] %x
After:
%y = %orig

2. Push conversion instructions inside of differentiable_function.
This unblocks inlining and specialization.
Before:
%x = differentiable_function(%orig, %jvp, %vjp)
%y = convert_escape_to_noescape %x
After:
%orig' = convert_escape_to_noescape %orig
%jvp' = convert_escape_to_noescape %jvp
%vjp' = convert_escape_to_noescape %vjp
%y = differentiable_function(%orig', %jvp', %vjp')

3. Another peephole is needed for reordering function conversion instructions to enable full inlining:
(convert_escape_to_noescape (convert_function (thin_to_thick_function x)))
=>
(convert_escape_to_noescape (thin_to_thick_function (convert_function x)))

Co-authored-by: Dan Zheng <danielzheng@google.com>

Author: asl

include/swift/AST/Types.h

@@ -4699,6 +4699,17 @@ class SILFunctionType final
   /// differentiability from all parameters.
   CanSILFunctionType getWithoutDifferentiability();
 
+  /// Given that `this` is a `@differentiable` function type, returns the type
+  /// of the given `@differentiable` function type component.
+  CanSILFunctionType getDifferentiableComponentType(
+      NormalDifferentiableFunctionTypeComponent component, SILModule &module);
+
+  /// Given that `this` is a `@differentiable(linear)` function type, returns
+  /// the type of the given `@differentiable(linear)` function type component.
+  CanSILFunctionType
+  getLinearComponentType(LinearDifferentiableFunctionTypeComponent component,
+                         SILModule &module);
+
   /// Returns the type of the derivative function for the given parameter
   /// indices, result indices, derivative function kind, derivative function
   /// generic signature (optional), and other auxiliary parameters.

include/swift/SIL/SILInstruction.h

@@ -9628,6 +9628,34 @@ class DifferentiableFunctionInst final
     }
     llvm_unreachable("invalid derivative kind");
   }
+
+  
+  /// Returns true iff the operand corresponding to the given extractee kind
+  /// exists.
+  bool hasExtractee(NormalDifferentiableFunctionTypeComponent extractee) const {
+    switch (extractee) {
+    case NormalDifferentiableFunctionTypeComponent::Original:
+      return true;
+    case NormalDifferentiableFunctionTypeComponent::JVP:
+    case NormalDifferentiableFunctionTypeComponent::VJP:
+      return hasDerivativeFunctions();
+    }
+    llvm_unreachable("invalid extractee kind");
+  }
+
+  /// Returns the operand corresponding to the given extractee kind.
+  SILValue
+  getExtractee(NormalDifferentiableFunctionTypeComponent extractee) const {
+    switch (extractee) {
+    case NormalDifferentiableFunctionTypeComponent::Original:
+      return getOriginalFunction();
+    case NormalDifferentiableFunctionTypeComponent::JVP:
+      return getJVPFunction();
+    case NormalDifferentiableFunctionTypeComponent::VJP:
+      return getVJPFunction();
+    }
+    llvm_unreachable("invalid extractee kind");
+  }
 };
 
 /// LinearFunctionInst - given a function, its derivative and transpose functions,
@@ -9668,6 +9696,31 @@ class LinearFunctionInst final
     assert(HasTransposeFunction);
     return getOperand(1);
   }
+
+  
+  /// Returns true iff the operand corresponding to the given extractee kind
+  /// exists.
+  bool hasExtractee(LinearDifferentiableFunctionTypeComponent extractee) const {
+    switch (extractee) {
+    case LinearDifferentiableFunctionTypeComponent::Original:
+      return true;
+    case LinearDifferentiableFunctionTypeComponent::Transpose:
+      return hasTransposeFunction();
+    }
+    llvm_unreachable("invalid extractee kind");
+  }
+
+  /// Returns the operand corresponding to the given extractee kind.
+  SILValue
+  getExtractee(LinearDifferentiableFunctionTypeComponent extractee) const {
+    switch (extractee) {
+    case LinearDifferentiableFunctionTypeComponent::Original:
+      return getOriginalFunction();
+    case LinearDifferentiableFunctionTypeComponent::Transpose:
+      return getTransposeFunction();
+    }
+    llvm_unreachable("invalid extractee kind");
+  }
 };
 
 /// DifferentiableFunctionExtractInst - extracts either the original or

lib/SIL/IR/SILFunctionType.cpp

@@ -263,6 +263,35 @@ IndexSubset *SILFunctionType::getDifferentiabilityResultIndices() {
   return IndexSubset::get(getASTContext(), numSemanticResults, resultIndices);
 }
 
+CanSILFunctionType SILFunctionType::getDifferentiableComponentType(
+    NormalDifferentiableFunctionTypeComponent component, SILModule &module) {
+  assert(getDifferentiabilityKind() == DifferentiabilityKind::Reverse &&
+         "Must be a `@differentiable(reverse)` function");
+  auto originalFnTy = getWithoutDifferentiability();
+  if (auto derivativeKind = component.getAsDerivativeFunctionKind()) {
+    return originalFnTy->getAutoDiffDerivativeFunctionType(
+        getDifferentiabilityParameterIndices(),
+        getDifferentiabilityResultIndices(), *derivativeKind, module.Types,
+        LookUpConformanceInModule(module.getSwiftModule()));
+  }
+  return originalFnTy;
+}
+
+CanSILFunctionType SILFunctionType::getLinearComponentType(
+    LinearDifferentiableFunctionTypeComponent component, SILModule &module) {
+  assert(getDifferentiabilityKind() == DifferentiabilityKind::Linear &&
+         "Must be a `@differentiable(linear)` function");
+  auto originalFnTy = getWithoutDifferentiability();
+  switch (component) {
+  case LinearDifferentiableFunctionTypeComponent::Original:
+    return originalFnTy;
+  case LinearDifferentiableFunctionTypeComponent::Transpose:
+    return originalFnTy->getAutoDiffTransposeFunctionType(
+        getDifferentiabilityParameterIndices(), module.Types,
+        LookUpConformanceInModule(module.getSwiftModule()));
+  }
+}
+
 CanSILFunctionType
 SILFunctionType::getWithDifferentiability(DifferentiabilityKind kind,
                                           IndexSubset *parameterIndices,

lib/SILOptimizer/Mandatory/Differentiation.cpp

@@ -62,12 +62,6 @@ using llvm::SmallSet;
 static llvm::cl::opt<bool> EnableExperimentalLinearMapTransposition(
     "enable-experimental-linear-map-transposition", llvm::cl::init(false));
 
-/// This flag is used to disable `differentiable_function_extract` instruction
-/// folding for SIL testing purposes.
-static llvm::cl::opt<bool> SkipFoldingDifferentiableFunctionExtraction(
-    "differentiation-skip-folding-differentiable-function-extraction",
-    llvm::cl::init(true));
-
 //===----------------------------------------------------------------------===//
 // Helpers
 //===----------------------------------------------------------------------===//
@@ -127,17 +121,6 @@ class DifferentiationTransformer {
   /// Process the given `linear_function` instruction, filling in the missing
   /// transpose function if necessary.
   bool processLinearFunctionInst(LinearFunctionInst *lfi);
-
-  /// Fold `differentiable_function_extract` users of the given
-  /// `differentiable_function` instruction, directly replacing them with
-  /// `differentiable_function` instruction operands. If the
-  /// `differentiable_function` instruction has no remaining uses, delete the
-  /// instruction itself after folding.
-  ///
-  /// Folding can be disabled by the
-  /// `SkipFoldingDifferentiableFunctionExtraction` flag for SIL testing
-  /// purposes.
-  void foldDifferentiableFunctionExtraction(DifferentiableFunctionInst *source);
 };
 
 } // end anonymous namespace
@@ -1239,51 +1222,6 @@ SILValue DifferentiationTransformer::promoteToLinearFunction(
   return newLinearFn;
 }
 
-/// Fold `differentiable_function_extract` users of the given
-/// `differentiable_function` instruction, directly replacing them with
-/// `differentiable_function` instruction operands. If the
-/// `differentiable_function` instruction has no remaining uses, delete the
-/// instruction itself after folding.
-///
-/// Folding can be disabled by the `SkipFoldingDifferentiableFunctionExtraction`
-/// flag for SIL testing purposes.
-// FIXME: This function is not correctly detecting the foldable pattern and
-// needs to be rewritten.
-void DifferentiationTransformer::foldDifferentiableFunctionExtraction(
-    DifferentiableFunctionInst *source) {
-  // Iterate through all `differentiable_function` instruction uses.
-  for (auto use : source->getUses()) {
-    auto *dfei = dyn_cast<DifferentiableFunctionExtractInst>(use->getUser());
-    // If user is not an `differentiable_function_extract` instruction, set flag
-    // to false.
-    if (!dfei)
-      continue;
-    // Fold original function extractors.
-    if (dfei->getExtractee() ==
-        NormalDifferentiableFunctionTypeComponent::Original) {
-      auto originalFnValue = source->getOriginalFunction();
-      dfei->replaceAllUsesWith(originalFnValue);
-      dfei->eraseFromParent();
-      continue;
-    }
-    // Fold derivative function extractors.
-    auto derivativeFnValue =
-        source->getDerivativeFunction(dfei->getDerivativeFunctionKind());
-    dfei->replaceAllUsesWith(derivativeFnValue);
-    dfei->eraseFromParent();
-  }
-  // If the `differentiable_function` instruction has no remaining uses, erase
-  // it.
-  if (isInstructionTriviallyDead(source)) {
-    SILBuilder builder(source);
-    builder.emitDestroyAddrAndFold(source->getLoc(), source->getJVPFunction());
-    builder.emitDestroyAddrAndFold(source->getLoc(), source->getVJPFunction());
-    source->eraseFromParent();
-  }
-  // Mark `source` as processed so that it won't be reprocessed after deletion.
-  context.markDifferentiableFunctionInstAsProcessed(source);
-}
-
 bool DifferentiationTransformer::processDifferentiableFunctionInst(
     DifferentiableFunctionInst *dfi) {
   PrettyStackTraceSILNode dfiTrace("canonicalizing `differentiable_function`",
@@ -1312,14 +1250,6 @@ bool DifferentiationTransformer::processDifferentiableFunctionInst(
   // Destroy the original operand.
   builder.emitDestroyValueOperation(loc, dfi->getOriginalFunction());
   dfi->eraseFromParent();
-  // If the promoted `@differentiable` function-typed value is an
-  // `differentiable_function` instruction, fold
-  // `differentiable_function_extract` instructions. If
-  // `differentiable_function_extract` folding is disabled, return.
-  if (!SkipFoldingDifferentiableFunctionExtraction)
-    if (auto *newDFI =
-            dyn_cast<DifferentiableFunctionInst>(differentiableFnValue))
-      foldDifferentiableFunctionExtraction(newDFI);
   transform.invalidateAnalysis(parent,
                                SILAnalysis::InvalidationKind::FunctionBody);
   return false;

lib/SILOptimizer/SILCombiner/SILCombiner.h

@@ -294,7 +294,9 @@ class SILCombiner :
   SILInstruction *visitConvertFunctionInst(ConvertFunctionInst *CFI);
   SILInstruction *
   visitConvertEscapeToNoEscapeInst(ConvertEscapeToNoEscapeInst *Cvt);
-
+  SILInstruction *
+  visitDifferentiableFunctionExtractInst(DifferentiableFunctionExtractInst *DFEI);
+  
   SILInstruction *legacyVisitGlobalValueInst(GlobalValueInst *globalValue);
 
 #define PASS(ID, TAG, DESCRIPTION)

lib/SILOptimizer/SILCombiner/SILCombinerCastVisitors.cpp

@@ -1065,16 +1065,94 @@ visitCheckedCastAddrBranchInst(CheckedCastAddrBranchInst *CCABI) {
 
 SILInstruction *SILCombiner::visitConvertEscapeToNoEscapeInst(
     ConvertEscapeToNoEscapeInst *Cvt) {
-  auto *OrigThinToThick =
-      dyn_cast<ThinToThickFunctionInst>(Cvt->getConverted());
-  if (!OrigThinToThick)
-    return nullptr;
-  auto origFunType = OrigThinToThick->getType().getAs<SILFunctionType>();
-  auto NewTy = origFunType->getWithExtInfo(origFunType->getExtInfo().withNoEscape(true));
+  // Rewrite conversion of `convert_function` of `thin_to_thick_function` as
+  // conversion of `thin_to_thick_function` of `convert_function`.
+  //
+  // (convert_escape_to_noescape (convert_function (thin_to_thick_function x)))
+  // =>
+  // (convert_escape_to_noescape (thin_to_thick_function (convert_function x)))
+  //
+  // This unblocks the `thin_to_thick_function` peephole optimization below.
+  if (auto *CFI = dyn_cast<ConvertFunctionInst>(Cvt->getConverted())) {
+    if (CFI->getSingleUse()) {
+      if (auto *TTTFI = dyn_cast<ThinToThickFunctionInst>(CFI->getConverted())) {
+        if (TTTFI->getSingleUse()) {
+          auto convertedThickType = CFI->getType().castTo<SILFunctionType>();
+          auto convertedThinType = convertedThickType->getWithRepresentation(
+            SILFunctionTypeRepresentation::Thin);
+          auto *newCFI = Builder.createConvertFunction(
+            CFI->getLoc(), TTTFI->getConverted(),
+            SILType::getPrimitiveObjectType(convertedThinType),
+            CFI->withoutActuallyEscaping());
+          auto *newTTTFI = Builder.createThinToThickFunction(
+            TTTFI->getLoc(), newCFI, CFI->getType());
+          replaceInstUsesWith(*CFI, newTTTFI);
+        }
+      }
+    }
+  }
+
+  // Rewrite conversion of `thin_to_thick_function` as `thin_to_thick_function`
+  // with a noescape function type.
+  //
+  // (convert_escape_to_noescape (thin_to_thick_function x)) =>
+  // (thin_to_thick_function [noescape] x)
+  if (auto *OrigThinToThick = dyn_cast<ThinToThickFunctionInst>(Cvt->getConverted())) {
+    auto origFunType = OrigThinToThick->getType().getAs<SILFunctionType>();
+    auto NewTy = origFunType->getWithExtInfo(origFunType->getExtInfo().withNoEscape(true));
 
-  return Builder.createThinToThickFunction(
+    return Builder.createThinToThickFunction(
       OrigThinToThick->getLoc(), OrigThinToThick->getOperand(),
       SILType::getPrimitiveObjectType(NewTy));
+  }
+
+  // Push conversion instructions inside `differentiable_function`. This
+  // unblocks more optimizations.
+  //
+  // Before:
+  // %x = differentiable_function(%orig, %jvp, %vjp)
+  // %y = convert_escape_to_noescape %x
+  //
+  // After:
+  // %orig' = convert_escape_to_noescape %orig
+  // %jvp' = convert_escape_to_noescape %jvp
+  // %vjp' = convert_escape_to_noescape %vjp
+  // %y = differentiable_function(%orig', %jvp', %vjp')
+  if (auto *DFI = dyn_cast<DifferentiableFunctionInst>(Cvt->getConverted())) {
+    auto createConvertEscapeToNoEscape = [&](NormalDifferentiableFunctionTypeComponent extractee) {
+      if (!DFI->hasExtractee(extractee))
+        return SILValue();
+        
+      auto operand = DFI->getExtractee(extractee);
+      auto fnType = operand->getType().castTo<SILFunctionType>();
+      auto noEscapeFnType =
+        fnType->getWithExtInfo(fnType->getExtInfo().withNoEscape());
+      auto noEscapeType = SILType::getPrimitiveObjectType(noEscapeFnType);
+      return Builder.createConvertEscapeToNoEscape(
+        operand.getLoc(), operand, noEscapeType, Cvt->isLifetimeGuaranteed())->getResult(0);
+    };
+    
+    SILValue originalNoEscape =
+      createConvertEscapeToNoEscape(NormalDifferentiableFunctionTypeComponent::Original);
+    SILValue convertedJVP = createConvertEscapeToNoEscape(
+      NormalDifferentiableFunctionTypeComponent::JVP);
+    SILValue convertedVJP = createConvertEscapeToNoEscape(
+      NormalDifferentiableFunctionTypeComponent::VJP);
+
+    Optional<std::pair<SILValue, SILValue>> derivativeFunctions;
+    if (convertedJVP && convertedVJP)
+      derivativeFunctions = std::make_pair(convertedJVP, convertedVJP);
+
+    auto *newDFI = Builder.createDifferentiableFunction(
+      DFI->getLoc(), DFI->getParameterIndices(), DFI->getResultIndices(),
+      originalNoEscape, derivativeFunctions);
+    assert(newDFI->getType() == Cvt->getType() &&
+           "New `@differentiable` function instruction should have same type "
+           "as the old `convert_escape_to_no_escape` instruction");
+    return newDFI;
+  }  
+  
+  return nullptr;
 }
 
 SILInstruction *
@@ -1207,6 +1285,54 @@ SILCombiner::visitConvertFunctionInst(ConvertFunctionInst *cfi) {
       return std::move(folder).optimizeWithSetValue(subCFI->getConverted());
   }
 
+  // Push conversion instructions inside `differentiable_function`. This
+  // unblocks more optimizations.
+  //
+  // Before:
+  // %x = differentiable_function(%orig, %jvp, %vjp)
+  // %y = convert_function %x
+  //
+  // After:
+  // %orig' = convert_function %orig
+  // %jvp' = convert_function %jvp
+  // %vjp' = convert_function %vjp
+  // %y = differentiable_function(%orig', %jvp', %vjp')
+  if (auto *DFI = dyn_cast<DifferentiableFunctionInst>(cfi->getConverted())) {
+    auto createConvertFunctionOfComponent =
+      [&](NormalDifferentiableFunctionTypeComponent extractee) {
+        if (!DFI->hasExtractee(extractee))
+          return SILValue();
+        
+        auto operand = DFI->getExtractee(extractee);
+        auto convertInstType =
+            cfi->getType().castTo<SILFunctionType>();
+        auto convertedComponentFnType =
+            convertInstType->getDifferentiableComponentType(
+                extractee, Builder.getModule());
+        auto convertedComponentType =
+        SILType::getPrimitiveObjectType(convertedComponentFnType);
+        return Builder.createConvertFunction(
+            operand.getLoc(), operand, convertedComponentType,
+            cfi->withoutActuallyEscaping())->getResult(0);
+      };
+      SILValue convertedOriginal = createConvertFunctionOfComponent(
+          NormalDifferentiableFunctionTypeComponent::Original);
+      SILValue convertedJVP = createConvertFunctionOfComponent(
+          NormalDifferentiableFunctionTypeComponent::JVP);
+      SILValue convertedVJP = createConvertFunctionOfComponent(
+          NormalDifferentiableFunctionTypeComponent::VJP);
+      Optional<std::pair<SILValue, SILValue>> derivativeFunctions;
+      if (convertedJVP && convertedVJP)
+        derivativeFunctions = std::make_pair(convertedJVP, convertedVJP);
+      auto *newDFI = Builder.createDifferentiableFunction(
+          DFI->getLoc(), DFI->getParameterIndices(), DFI->getResultIndices(),
+          convertedOriginal, derivativeFunctions);
+      assert(newDFI->getType() == cfi->getType() &&
+             "New `@differentiable` function instruction should have same type "
+             "as the old `convert_function` instruction");
+      return newDFI;
+  }
+  
   // Replace a convert_function that only has refcounting uses with its
   // operand.
   tryEliminateOnlyOwnershipUsedForwardingInst(cfi, getInstModCallbacks());