Adding @asl's fix for subset parameters thunks involving functions with multiple results, and an activity analysis test representing code that had exposed that issue.

Author: BradLarson

lib/SILOptimizer/Differentiation/Common.cpp

@@ -232,7 +232,8 @@ void collectMinimalIndicesForFunctionCall(
     auto &param = paramAndIdx.value();
     if (!param.isIndirectMutating())
       continue;
-    unsigned idx = paramAndIdx.index();
+    unsigned idx =
+        paramAndIdx.index() + calleeFnTy->getNumIndirectFormalResults();
     auto inoutArg = ai->getArgument(idx);
     results.push_back(inoutArg);
     resultIndices.push_back(inoutParamResultIndex++);

lib/SILOptimizer/Differentiation/Thunk.cpp

@@ -472,6 +472,12 @@ getOrCreateSubsetParametersThunkForLinearMap(
     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:
@@ -480,9 +486,29 @@ getOrCreateSubsetParametersThunkForLinearMap(
   //     indices).
   //   - Zeros (when parameter is not in desired indices).
   case AutoDiffDerivativeFunctionKind::JVP: {
-    // Forward all indirect results.
-    arguments.append(thunk->getIndirectResults().begin(),
-                     thunk->getIndirectResults().end());
+    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.
@@ -507,10 +533,6 @@ getOrCreateSubsetParametersThunkForLinearMap(
   //   - Zeros (when parameter is not in desired indices).
   // - All actual arguments.
   case AutoDiffDerivativeFunctionKind::VJP: {
-    auto toIndirectResultsIter = thunk->getIndirectResults().begin();
-    auto useNextIndirectResult = [&]() {
-      arguments.push_back(*toIndirectResultsIter++);
-    };
     // Collect pullback arguments.
     unsigned pullbackResultIndex = 0;
     for (unsigned i : actualConfig.parameterIndices->getIndices()) {
@@ -539,8 +561,18 @@ getOrCreateSubsetParametersThunkForLinearMap(
       arguments.push_back(indirectResult);
     }
     // Forward all actual non-indirect-result arguments.
-    arguments.append(thunk->getArgumentsWithoutIndirectResults().begin(),
-                     thunk->getArgumentsWithoutIndirectResults().end() - 1);
+    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;
   }
   }
@@ -550,10 +582,33 @@ getOrCreateSubsetParametersThunkForLinearMap(
   auto *ai = builder.createApply(loc, linearMap, SubstitutionMap(), arguments);
 
   // If differential thunk, deallocate local allocations and directly return
-  // `apply` result.
+  // `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();
-    builder.createReturn(loc, ai);
+    auto result = joinElements(results, builder, loc);
+    builder.createReturn(loc, result);
     return {thunk, interfaceSubs};
   }
 

test/AutoDiff/SILOptimizer/activity_analysis.swift

@@ -533,6 +533,44 @@ func activeInoutArgNonactiveInitialResult(_ x: Float) -> Float {
 // CHECK: [ACTIVE]   %13 = begin_access [read] [static] %2 : $*Float
 // CHECK: [ACTIVE]   %14 = load [trivial] %13 : $*Float
 
+public struct ArrayWrapper: Differentiable {
+  var values: [Float]
+
+  @differentiable(reverse)
+  mutating func get(index: Int) -> Float {
+    self.values[index]
+  }
+
+  // Check `inout` with result.
+
+  // CHECK-LABEL: [AD] Activity info for ${{.*}}get{{.*}} at parameter indices (1) and result indices (0, 1)
+  // CHECK: bb0:
+  // CHECK: [USEFUL] %0 = argument of bb0 : $Int
+  // CHECK: [ACTIVE] %1 = argument of bb0 : $*ArrayWrapper
+  // CHECK: [ACTIVE]   %4 = begin_access [read] [static] %1 : $*ArrayWrapper
+  // CHECK: [ACTIVE]   %5 = struct_element_addr %4 : $*ArrayWrapper, #ArrayWrapper.values
+  // CHECK: [ACTIVE]   %6 = load_borrow %5 : $*Array<Float>
+  // CHECK: [ACTIVE]   %7 = alloc_stack $Float
+  // CHECK: [NONE]   // function_ref Array.subscript.getter
+  // CHECK:   %8 = function_ref @$sSayxSicig : $@convention(method) <τ_0_0> (Int, @guaranteed Array<τ_0_0>) -> @out τ_0_0
+  // CHECK: [NONE]   %9 = apply %8<Float>(%7, %0, %6) : $@convention(method) <τ_0_0> (Int, @guaranteed Array<τ_0_0>) -> @out τ_0_0
+  // CHECK: [ACTIVE]   %10 = load [trivial] %7 : $*Float
+}
+
+@differentiable(reverse)
+func testInoutAndResult(x: Int, y: inout ArrayWrapper) {
+  let _ = y.get(index: x)
+}
+
+// CHECK-LABEL: [AD] Activity info for ${{.*}}testInoutAndResult{{.*}} at parameter indices (1) and result indices (0)
+// CHECK: bb0:
+// CHECK: [USEFUL] %0 = argument of bb0 : $Int
+// CHECK: [ACTIVE] %1 = argument of bb0 : $*ArrayWrapper
+// CHECK: [ACTIVE]   %4 = begin_access [modify] [static] %1 : $*ArrayWrapper
+// CHECK: [NONE]   // function_ref ArrayWrapper.get(index:)
+// CHECK:   %5 = function_ref @$s17activity_analysis12ArrayWrapperV3get5indexSfSi_tF : $@convention(method) (Int, @inout ArrayWrapper) -> Float
+// CHECK: [VARIED]   %6 = apply %5(%0, %4) : $@convention(method) (Int, @inout ArrayWrapper) -> Float
+
 //===----------------------------------------------------------------------===//
 // Throwing function differentiation (`try_apply`)
 //===----------------------------------------------------------------------===//