Revert "Implement several peephole optimizations to unblock further optimizations of autodiff code (#60520)"

This reverts commit 2f5492f.

Author: meg-gupta

include/swift/AST/Types.h

@@ -4683,17 +4683,6 @@ 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,34 +9628,6 @@ 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,
@@ -9696,31 +9668,6 @@ 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,35 +263,6 @@ 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,6 +62,12 @@ 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
 //===----------------------------------------------------------------------===//
@@ -121,6 +127,17 @@ 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
@@ -1222,6 +1239,51 @@ 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`",
@@ -1250,6 +1312,14 @@ 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,8 +294,8 @@ class SILCombiner :
   SILInstruction *visitConvertFunctionInst(ConvertFunctionInst *CFI);
   SILInstruction *
   visitConvertEscapeToNoEscapeInst(ConvertEscapeToNoEscapeInst *Cvt);
-  SILInstruction *
-  visitDifferentiableFunctionExtractInst(DifferentiableFunctionExtractInst *DFEI);
+
+  SILInstruction *legacyVisitGlobalValueInst(GlobalValueInst *globalValue);
 
 #define PASS(ID, TAG, DESCRIPTION)
 #define SWIFT_FUNCTION_PASS(ID, TAG, DESCRIPTION)

lib/SILOptimizer/SILCombiner/SILCombinerCastVisitors.cpp

@@ -1065,92 +1065,16 @@ visitCheckedCastAddrBranchInst(CheckedCastAddrBranchInst *CCABI) {
 
 SILInstruction *SILCombiner::visitConvertEscapeToNoEscapeInst(
     ConvertEscapeToNoEscapeInst *Cvt) {
-  // 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 (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));
+  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));
 
-    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 *
@@ -1283,54 +1207,6 @@ 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());