Skip to content

Commit bcc7ed9

Browse files
authored
[AutoDiff] Define derivative for concrete SIMD.init(repeating:) (#81864)
The PR #81766 introduced some concrete SIMD operations, but `init(repeating:)` was temporarily disabled due to differentiation testing break. See: stephentyrone@7a00619 This PR contains two changes: 1. Define custom derivatives for concrete `init(repeating:)` so we do not fall into non-differentiability error diagnostic. 2. Add a fix to SIL linker so differentiability witness lookup is done when the original function has both `@_alwaysEmitIntoClient` and `@_transparent`. Similar changes were introduced previously in #78908, but they only handled `@_alwaysEmitIntoClient` without `@_transparent`. <!-- If this pull request is targeting a release branch, please fill out the following form: https://github.com/swiftlang/.github/blob/main/PULL_REQUEST_TEMPLATE/release.md?plain=1 Otherwise, replace this comment with a description of your changes and rationale. Provide links to external references/discussions if appropriate. If this pull request resolves any GitHub issues, link them like so: Resolves <link to issue>, resolves <link to another issue>. For more information about linking a pull request to an issue, see: https://docs.github.com/issues/tracking-your-work-with-issues/linking-a-pull-request-to-an-issue --> <!-- Before merging this pull request, you must run the Swift continuous integration tests. For information about triggering CI builds via @swift-ci, see: https://github.com/apple/swift/blob/main/docs/ContinuousIntegration.md#swift-ci Thank you for your contribution to Swift! -->
1 parent dad81fb commit bcc7ed9

File tree

4 files changed

+38
-7
lines changed

4 files changed

+38
-7
lines changed

lib/SIL/IR/Linker.cpp

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -141,6 +141,20 @@ void SILLinkerVisitor::maybeAddFunctionToWorklist(
141141
// are also set to IsSerialized.
142142
Worklist.push_back(F);
143143
}
144+
145+
if (F->markedAsAlwaysEmitIntoClient()) {
146+
// For @_alwaysEmitIntoClient functions, we need to lookup its
147+
// differentiability witness and, if present, ask SILLoader to obtain its
148+
// definition. Otherwise, a linker error would occur due to undefined
149+
// reference to these symbols.
150+
for (SILDifferentiabilityWitness *witness :
151+
F->getModule().lookUpDifferentiabilityWitnessesForFunction(
152+
F->getName())) {
153+
F->getModule().getSILLoader()->lookupDifferentiabilityWitness(
154+
witness->getKey());
155+
}
156+
}
157+
144158
return;
145159
}
146160

stdlib/public/Differentiation/SIMDDifferentiation.swift.gyb

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -455,3 +455,27 @@ where
455455
return (Self(repeating: value), { v in Self(repeating: v) })
456456
}
457457
}
458+
459+
%for (Scalar, bits) in [('Float',32), ('Double',64)]:
460+
% for n in vectorscalarCounts:
461+
extension SIMD${n} where Scalar == ${Scalar} {
462+
@inlinable
463+
@_alwaysEmitIntoClient
464+
@derivative(of: init(repeating:))
465+
static func _vjpInit(repeating value: Scalar)
466+
-> (value: Self, pullback: (TangentVector) -> Scalar.TangentVector)
467+
{
468+
return (Self(repeating: value), { v in v.sum() })
469+
}
470+
471+
@inlinable
472+
@_alwaysEmitIntoClient
473+
@derivative(of: init(repeating:))
474+
static func _jvpInit(repeating value: Scalar)
475+
-> (value: Self, differential: (Scalar.TangentVector) -> TangentVector)
476+
{
477+
return (Self(repeating: value), { v in Self(repeating: v) })
478+
}
479+
}
480+
% end
481+
%end

stdlib/public/core/SIMDFloatConcreteOperations.swift.gyb

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -34,8 +34,6 @@ extension SIMD${n} where Scalar == ${Scalar} {
3434
_storage = ${Scalar}.SIMD${storageN}Storage(_builtin)
3535
}
3636

37-
/* Breaks differentiation testing, commented out while we figure out
38-
what to do about that.
3937
@_alwaysEmitIntoClient @_transparent
4038
public init(repeating scalar: ${Scalar}) {
4139
let asVector = Builtin.insertelement_${Builtin}_FPIEEE${bits}_Int32(
@@ -52,7 +50,6 @@ extension SIMD${n} where Scalar == ${Scalar} {
5250
))
5351
%end
5452
}
55-
*/
5653

5754
% if n >= 4:
5855
@_alwaysEmitIntoClient @_transparent

test/stdlib/SIMDFloatInitializers.swift.gyb

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -14,10 +14,6 @@
1414
// RUN: %target-swift-frontend -primary-file %t/SIMDFloatInitializers.swift -S | %FileCheck %t/SIMDFloatInitializers.swift --check-prefix=CHECK --check-prefix=CHECK-%target-cpu --check-prefix=CHECKOnone-%target-cpu
1515
// RUN: %target-swift-frontend -primary-file %t/SIMDFloatInitializers.swift -S -O | %FileCheck %t/SIMDFloatInitializers.swift --check-prefix=CHECK --check-prefix=CHECK-%target-cpu --check-prefix=CHECKO-%target-cpu
1616

17-
// Disable this test for now because aEIC/transparent functions still are not
18-
// correctly differentiable, and so these inits are suppressed in the stdlib.
19-
// REQUIRES: differentiable-aEIC-transparent
20-
2117
import Swift
2218

2319
%for bits in [16,32,64]:

0 commit comments

Comments
 (0)