Fix use after free when pullback is used multiple times. (#64647)

Linear maps are captured in vjp routine via callee-guaranteed partial apply and are passed as @owned references to the enclosing pullback that finally consumes them. Necessary retains are inserted by a partial apply forwarder.

However, this is not the case when the function being differentiated contains loops as heap-allocated context is used and bare pointer is captured by the pullback partial apply. As a result, partial apply forwarder does not retain the linear maps that are owned by a heap-allocated context, however, they are still treated as @owned references and therefore are released in the pullback after the first call. As a result, subsequent pullback calls release linear maps and we'd end with possible use-after-free.

Ensure we retain values when we load values from the context.

Reproducible only when:

 * Loops (so, heap-allocated context)
 * Pullbacks of thick functions (so context is non-zero)
 * Multiple pullback calls
 * Some cleanup while there

Fixes #64257

Author: asl

include/swift/SILOptimizer/Differentiation/LinearMapInfo.h

@@ -86,15 +86,16 @@ class LinearMapInfo {
   llvm::DenseMap<std::pair<SILBasicBlock *, SILBasicBlock *>, EnumElementDecl *>
       branchingTraceEnumCases;
 
-  /// Blocks in a loop.
-  llvm::SmallSetVector<SILBasicBlock *, 4> blocksInLoop;
-
   /// A synthesized file unit.
   SynthesizedFileUnit &synthesizedFile;
 
   /// A type converter, used to compute struct/enum SIL types.
   Lowering::TypeConverter &typeConverter;
 
+  /// True, if a heap-allocated context is required. For example, when there are
+  /// any loops
+  bool heapAllocatedContext = false;
+
 private:
   /// Remaps the given type into the derivative function's context.
   SILType remapTypeInDerivative(SILType ty);
@@ -193,12 +194,8 @@ class LinearMapInfo {
     return getLinearMapTupleType(ai->getParentBlock())->getElement(idx).getType();
   }
 
-  bool hasLoops() const {
-    return !blocksInLoop.empty();
-  }
-
-  ArrayRef<SILBasicBlock *> getBlocksInLoop() const {
-    return blocksInLoop.getArrayRef();
+  bool hasHeapAllocatedContext() const {
+    return heapAllocatedContext;
   }
 };
 

lib/SILOptimizer/Differentiation/LinearMapInfo.cpp

@@ -138,11 +138,9 @@ void LinearMapInfo::populateBranchingTraceDecl(SILBasicBlock *originalBB,
     // indirectly referenced in memory owned by the context object. The payload
     // is just a raw pointer.
     if (loopInfo->getLoopFor(predBB)) {
-      blocksInLoop.insert(predBB);
+      heapAllocatedContext = true;
       decl->setInterfaceType(astCtx.TheRawPointerType);
-    }
-    // Otherwise the payload is the linear map tuple.
-    else {
+    } else { // Otherwise the payload is the linear map tuple.
       auto linearMapStructTy = getLinearMapTupleType(predBB)->getCanonicalType();
       decl->setInterfaceType(
           linearMapStructTy->hasArchetype()

lib/SILOptimizer/Differentiation/PullbackCloner.cpp

@@ -1928,7 +1928,7 @@ bool PullbackCloner::Implementation::run() {
       builder.setInsertionPoint(pullbackBB);
       // Obtain the context object, if any, and the top-level subcontext, i.e.
       // the main pullback struct.
-      if (getPullbackInfo().hasLoops()) {
+      if (getPullbackInfo().hasHeapAllocatedContext()) {
         // The last argument is the context object (`Builtin.NativeObject`).
         contextValue = pullbackBB->getArguments().back();
         assert(contextValue->getType() ==
@@ -1939,7 +1939,7 @@ bool PullbackCloner::Implementation::run() {
         SILValue mainPullbackTuple = builder.createLoad(
             pbLoc, subcontextAddr,
             pbTupleLoweredType.isTrivial(getPullback()) ?
-                LoadOwnershipQualifier::Trivial : LoadOwnershipQualifier::Take);
+                LoadOwnershipQualifier::Trivial : LoadOwnershipQualifier::Copy);
         auto *dsi = builder.createDestructureTuple(pbLoc, mainPullbackTuple);
         initializePullbackTupleElements(origBB, dsi->getAllResults());
       } else {
@@ -2023,7 +2023,7 @@ bool PullbackCloner::Implementation::run() {
   auto *pullbackEntry = pullback.getEntryBlock();
   auto pbTupleLoweredType =
     remapType(getPullbackInfo().getLinearMapTupleLoweredType(originalExitBlock));
-  unsigned numVals = (getPullbackInfo().hasLoops() ?
+  unsigned numVals = (getPullbackInfo().hasHeapAllocatedContext() ?
                       1 : pbTupleLoweredType.getAs<TupleType>()->getNumElements());
   (void)numVals;
 
@@ -2449,7 +2449,7 @@ SILBasicBlock *PullbackCloner::Implementation::buildPullbackSuccessor(
     auto predPbStructVal = pullbackTrampolineBBBuilder.createLoad(
         loc, predPbTupleAddr,
         pbTupleType.isTrivial(getPullback()) ?
-            LoadOwnershipQualifier::Trivial : LoadOwnershipQualifier::Take);
+            LoadOwnershipQualifier::Trivial : LoadOwnershipQualifier::Copy);
     trampolineArguments.push_back(predPbStructVal);
   } else {
     trampolineArguments.push_back(pullbackTrampolineBBArg);

lib/SILOptimizer/Differentiation/VJPCloner.cpp

@@ -110,7 +110,7 @@ class VJPCloner::Implementation final
 
   /// Initializes a context object if needed.
   void emitLinearMapContextInitializationIfNeeded() {
-    if (!pullbackInfo.hasLoops())
+    if (!pullbackInfo.hasHeapAllocatedContext())
       return;
 
     // Get linear map struct size.
@@ -937,7 +937,7 @@ SILFunction *VJPCloner::Implementation::createEmptyPullback() {
     pbParams.push_back(inoutParamTanParam);
   }
 
-  if (pullbackInfo.hasLoops()) {
+  if (pullbackInfo.hasHeapAllocatedContext()) {
     // Accept a `AutoDiffLinarMapContext` heap object if there are loops.
     pbParams.push_back({
       getASTContext().TheNativeObjectType,

test/AutoDiff/validation-test/issue-64257-use-after-free.swift

@@ -0,0 +1,50 @@
+// RUN: %target-run-simple-swift
+// REQUIRES: executable_test
+
+import _Differentiation
+public func simWithPullback(params: MinParams) -> (value: Output, pullback: (Output.TangentVector) -> (MinParams.TangentVector)){
+    let simulationValueAndPullback = valueWithPullback(at: params, of: run)
+    return (value: simulationValueAndPullback.value, pullback: simulationValueAndPullback.pullback)
+}
+
+@differentiable(reverse)
+public func run(params: MinParams) -> Output {
+    for t in 0 ... 1 {
+    }
+
+    let res = MiniLoop(other: params._twoDArray).results
+    return Output(results: res)
+}
+
+struct MiniLoop: Differentiable {
+    var results: Float
+    var twoDArray: Float
+  
+    @differentiable(reverse)
+    init(results: Float = 146, other: Float = 123) {self.results = results; self.twoDArray = other}
+}
+
+public struct Output: Differentiable {
+    public var results: Float
+    @differentiable(reverse)
+    public init(results: Float) {self.results = results}
+}
+
+public struct MinParams: Differentiable {
+    public var _twoDArray: Float
+    public init(foo: Float = 42) { _twoDArray = foo }
+}
+
+func test() {
+    let valueAndPullback = simWithPullback(params: MinParams())
+    let output = valueAndPullback.value
+    let resultOnes = Float(1.0)
+    var grad = valueAndPullback.pullback(Output.TangentVector(results: resultOnes))
+    print(grad)
+    grad = valueAndPullback.pullback(Output.TangentVector(results: resultOnes))
+    print(grad)
+    grad = valueAndPullback.pullback(Output.TangentVector(results: resultOnes))
+    print(grad)
+}
+
+test()