Skip to content

Commit c08d080

Browse files
authored
[AutoDiff] Fix pullback subset thunk generation crash. (#33969)
Fix pullback subset thunk generation crash due to unmapped parameter index for `inout` differentiability parameters. Resolves TF-1315.
1 parent 54fe511 commit c08d080

File tree

2 files changed

+29
-1
lines changed

2 files changed

+29
-1
lines changed

lib/SILOptimizer/Differentiation/Thunk.cpp

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -768,7 +768,8 @@ getOrCreateSubsetParametersThunkForLinearMap(
768768
if (!paramInfo.isIndirectMutating())
769769
continue;
770770
auto inoutArg = *std::next(ai->getInoutArguments().begin(), inoutArgIdx++);
771-
allResults.insert(allResults.begin() + paramIdx, inoutArg);
771+
unsigned mappedParamIdx = mapOriginalParameterIndex(paramIdx);
772+
allResults.insert(allResults.begin() + mappedParamIdx, inoutArg);
772773
}
773774
assert(allResults.size() == actualIndices.parameters->getNumIndices() &&
774775
"Number of pullback results should match number of differentiability "
Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
// RUN: %target-swift-frontend -emit-sil %s
2+
// REQUIRES: asserts
3+
4+
// TF-1315: Pullback subset thunk generation crash due to unmapped parameter
5+
// index for `inout` differentiability parameters.
6+
7+
import _Differentiation
8+
9+
func foo(_ x: Int, _ y: Float, _ z: inout Float) {}
10+
11+
@derivative(of: foo, wrt: (y, z))
12+
func vjpFoo(_ x: Int, _ y: Float, _ z: inout Float) -> (
13+
value: Void, pullback: (inout Float) -> Float
14+
) {
15+
fatalError()
16+
}
17+
18+
@differentiable
19+
func TF_1315(_ x: Float) -> Float {
20+
var x = x
21+
// The call to `foo` below triggers pullback subset parameter thunk generation.
22+
// `foo` original function type: `(Int, Float, inout Float) -> ()`
23+
// Actual parameter indices: 1, 2
24+
// Desired parameter indices: 2
25+
foo(1, 2, &x)
26+
return x
27+
}

0 commit comments

Comments
 (0)