[Autodiff] Separates out changes for using non-ossa instructions in the closure-spec optimization pass

The OSSA elimination pass has not yet been moved below all high level
function passes. Until that work has been completed the Autodiff
closure-spec optimization pass cannot solely support OSSA instructions.

Author: jkshtj

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/ClosureSpecialization.swift

@@ -255,21 +255,23 @@ private func rewriteApplyInstruction(using specializedCallee: Function, callSite
       let closureArgDesc = callSite.closureArgDesc(at: parentClosureArgIndex)!
       var builder = Builder(before: closureArgDesc.closure, context)
 
+      // TODO: Support only OSSA instructions once the OSSA elimination pass is moved after all function optimization 
+      // passes.
       if callSite.applySite.parentBlock != closureArgDesc.closure.parentBlock {
         // Emit the retain and release that keeps the argument live across the callee using the closure.
-        builder.createCopyValue(operand: capturedArg)
+        builder.createRetainValue(operand: capturedArg)
 
         for instr in closureArgDesc.lifetimeFrontier {
           builder = Builder(before: instr, context)
-          builder.createDestroyValue(operand: capturedArg)
+          builder.createReleaseValue(operand: capturedArg)
         }
 
         // Emit the retain that matches the captured argument by the partial_apply in the callee that is consumed by
         // the partial_apply.
         builder = Builder(before: callSite.applySite, context)
-        builder.createCopyValue(operand: capturedArg)
+        builder.createRetainValue(operand: capturedArg)
       } else {
-        builder.createCopyValue(operand: capturedArg)
+        builder.createRetainValue(operand: capturedArg)
       }
     }
   }
@@ -286,12 +288,14 @@ private func rewriteApplyInstruction(using specializedCallee: Function, callSite
                                             isOnStack: oldApply.isOnStack)
 
   builder = Builder(before: callSite.applySite.next!, context)
+  // TODO: Support only OSSA instructions once the OSSA elimination pass is moved after all function optimization 
+  // passes.
   for closureArgDesc in callSite.closureArgDescriptors {
     if closureArgDesc.isClosureConsumed,
        !closureArgDesc.isPartialApplyOnStack,
        !closureArgDesc.parameterInfo.isTrivialNoescapeClosure
     {
-      builder.createDestroyValue(operand: closureArgDesc.closure)
+      builder.createReleaseValue(operand: closureArgDesc.closure)
     }
   }
 
@@ -825,7 +829,7 @@ private extension SpecializationCloner {
             {
               builder.destroyPartialApplyOnStack(paiOnStack: pai, self.context)  
             } else {
-              builder.createDestroyValue(operand: closure)
+              builder.createReleaseValue(operand: closure)
             }
           }
         }
@@ -992,11 +996,25 @@ private extension Builder {
   func destroyPartialApplyOnStack(paiOnStack: PartialApplyInst, _ context: FunctionPassContext){
     precondition(paiOnStack.isOnStack, "Function must only be called for `partial_apply`s on stack!")
 
-    for arg in paiOnStack.arguments {
-      self.createDestroyValue(operand: arg)
-    }
+    // TODO: Support only OSSA instructions once the OSSA elimination pass is moved after all function optimization 
+    // passes.
+    //
+    // for arg in paiOnStack.arguments {
+    //   self.createDestroyValue(operand: arg)
+    // }
 
-    self.createDestroyValue(operand: paiOnStack)
+    // self.createDestroyValue(operand: paiOnStack)
+
+    if paiOnStack.parentFunction.hasOwnership {
+      // Under OSSA, the closure acts as an owned value whose lifetime is a borrow scope for the captures, so we need to
+      // end the borrow scope before ending the lifetimes of the captures themselves.
+      self.createDestroyValue(operand: paiOnStack)
+      self.destroyCapturedArgs(for: paiOnStack)
+    } else {
+      self.destroyCapturedArgs(for: paiOnStack)
+      self.createDeallocStack(paiOnStack)
+      context.notifyInvalidatedStackNesting()
+    }
   }
 }
 

SwiftCompilerSources/Sources/SIL/Builder.swift

@@ -151,6 +151,18 @@ public struct Builder {
     return notifyNew(endInit.getAs(EndInitLetRefInst.self))
   }
 
+  @discardableResult
+  public func createRetainValue(operand: Value) -> RetainValueInst {
+    let retain = bridged.createRetainValue(operand.bridged)
+    return notifyNew(retain.getAs(RetainValueInst.self))
+  }
+
+  @discardableResult
+  public func createReleaseValue(operand: Value) -> ReleaseValueInst {
+    let release = bridged.createReleaseValue(operand.bridged)
+    return notifyNew(release.getAs(ReleaseValueInst.self))
+  }
+
   @discardableResult
   public func createStrongRetain(operand: Value) -> StrongRetainInst {
     let retain = bridged.createStrongRetain(operand.bridged)

include/swift/SIL/SILBridging.h

@@ -1210,6 +1210,10 @@ struct BridgedBuilder{
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedInstruction createBeginDeallocRef(BridgedValue reference,
                                                                               BridgedValue allocation) const;
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedInstruction createEndInitLetRef(BridgedValue op) const;
+  SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedInstruction
+  createRetainValue(BridgedValue op) const;
+  SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedInstruction
+  createReleaseValue(BridgedValue op) const;
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedInstruction createStrongRetain(BridgedValue op) const;
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedInstruction createStrongRelease(BridgedValue op) const;
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedInstruction createUnownedRetain(BridgedValue op) const;

include/swift/SIL/SILBridgingImpl.h

@@ -1590,6 +1590,18 @@ BridgedInstruction BridgedBuilder::createEndInitLetRef(BridgedValue op) const {
   return {unbridged().createEndInitLetRef(regularLoc(), op.getSILValue())};
 }
 
+BridgedInstruction BridgedBuilder::createRetainValue(BridgedValue op) const {
+  auto b = unbridged();
+  return {b.createRetainValue(regularLoc(), op.getSILValue(),
+                              b.getDefaultAtomicity())};
+}
+
+BridgedInstruction BridgedBuilder::createReleaseValue(BridgedValue op) const {
+  auto b = unbridged();
+  return {b.createReleaseValue(regularLoc(), op.getSILValue(),
+                               b.getDefaultAtomicity())};
+}
+
 BridgedInstruction BridgedBuilder::createStrongRetain(BridgedValue op) const {
   auto b = unbridged();
   return {b.createStrongRetain(regularLoc(), op.getSILValue(), b.getDefaultAtomicity())};