Skip to content

Commit 2b7adbc

Browse files
xedinxedin
authored
Merge pull request #84800 from xedin/remove-csapply-operator-devirt
[CSApply] Don't attempt operator devirtualization
2 parents c23de59 + 2943d63 commit 2b7adbc

16 files changed

+965
-343
lines changed

‎lib/SILOptimizer/Mandatory/Differentiation.cpp‎

Lines changed: 240 additions & 117 deletions
Large diffs are not rendered by default.

‎lib/Sema/CSApply.cpp‎

Lines changed: 0 additions & 119 deletions
Original file line numberDiff line numberDiff line change
@@ -588,66 +588,6 @@ namespace {
588588
}
589589
}
590590

591-
// Returns None if the AST does not contain enough information to recover
592-
// substitutions; this is different from an Optional(SubstitutionMap()),
593-
// indicating a valid call to a non-generic operator.
594-
std::optional<SubstitutionMap> getOperatorSubstitutions(ValueDecl *witness,
595-
Type refType) {
596-
// We have to recover substitutions in this hacky way because
597-
// the AST does not retain enough information to devirtualize
598-
// calls like this.
599-
auto witnessType = witness->getInterfaceType();
600-
601-
// Compute the substitutions.
602-
auto *gft = witnessType->getAs<GenericFunctionType>();
603-
if (gft == nullptr) {
604-
if (refType->isEqual(witnessType))
605-
return SubstitutionMap();
606-
return std::nullopt;
607-
}
608-
609-
auto sig = gft->getGenericSignature();
610-
auto *env = sig.getGenericEnvironment();
611-
612-
witnessType = FunctionType::get(gft->getParams(),
613-
gft->getResult(),
614-
gft->getExtInfo());
615-
witnessType = env->mapTypeIntoContext(witnessType);
616-
617-
TypeSubstitutionMap subs;
618-
auto substType = witnessType->substituteBindingsTo(
619-
refType,
620-
[&](ArchetypeType *origType, CanType substType) -> CanType {
621-
if (auto gpType = dyn_cast<GenericTypeParamType>(
622-
origType->getInterfaceType()->getCanonicalType()))
623-
subs[gpType] = substType;
624-
625-
return substType;
626-
});
627-
628-
// If substitution failed, it means that the protocol requirement type
629-
// and the witness type did not match up. The only time that this
630-
// should happen is when the witness is defined in a base class and
631-
// the actual call uses a derived class. For example,
632-
//
633-
// protocol P { func +(lhs: Self, rhs: Self) }
634-
// class Base : P { func +(lhs: Base, rhs: Base) {} }
635-
// class Derived : Base {}
636-
//
637-
// If we enter this code path with two operands of type Derived,
638-
// we know we're calling the protocol requirement P.+, with a
639-
// substituted type of (Derived, Derived) -> (). But the type of
640-
// the witness is (Base, Base) -> (). Just bail out and make a
641-
// witness method call in this rare case; SIL mandatory optimizations
642-
// will likely devirtualize it anyway.
643-
if (!substType)
644-
return std::nullopt;
645-
646-
return SubstitutionMap::get(sig,
647-
QueryTypeSubstitutionMap{subs},
648-
LookUpConformanceInModule());
649-
}
650-
651591
/// Determine whether the given reference is to a method on
652592
/// a remote distributed actor in the given context.
653593
bool isDistributedThunk(ConcreteDeclRef ref, Expr *context);
@@ -674,65 +614,6 @@ namespace {
674614

675615
auto baseTy = getBaseType(adjustedFullType->castTo<FunctionType>());
676616

677-
// Handle operator requirements found in protocols.
678-
if (auto proto = dyn_cast<ProtocolDecl>(decl->getDeclContext())) {
679-
bool isCurried = shouldBuildCurryThunk(choice, /*baseIsInstance=*/false);
680-
681-
// If we have a concrete conformance, build a call to the witness.
682-
//
683-
// FIXME: This is awful. We should be able to handle this as a call to
684-
// the protocol requirement with Self == the concrete type, and SILGen
685-
// (or later) can devirtualize as appropriate.
686-
auto conformance = checkConformance(baseTy, proto);
687-
if (conformance.isConcrete()) {
688-
if (auto witness = conformance.getConcrete()->getWitnessDecl(decl)) {
689-
bool isMemberOperator = witness->getDeclContext()->isTypeContext();
690-
691-
if (!isMemberOperator || !isCurried) {
692-
// The fullType was computed by substituting the protocol
693-
// requirement so it always has a (Self) -> ... curried
694-
// application. Strip it off if the witness was a top-level
695-
// function.
696-
Type refType;
697-
if (isMemberOperator)
698-
refType = adjustedFullType;
699-
else
700-
refType = adjustedFullType->castTo<AnyFunctionType>()->getResult();
701-
702-
// Build the AST for the call to the witness.
703-
auto subMap = getOperatorSubstitutions(witness, refType);
704-
if (subMap) {
705-
ConcreteDeclRef witnessRef(witness, *subMap);
706-
auto declRefExpr = new (ctx) DeclRefExpr(witnessRef, loc,
707-
/*Implicit=*/false);
708-
declRefExpr->setFunctionRefInfo(choice.getFunctionRefInfo());
709-
cs.setType(declRefExpr, refType);
710-
711-
Expr *refExpr;
712-
if (isMemberOperator) {
713-
// If the operator is a type member, add the implicit
714-
// (Self) -> ... call.
715-
Expr *base =
716-
TypeExpr::createImplicitHack(loc.getBaseNameLoc(), baseTy,
717-
ctx);
718-
cs.setType(base, MetatypeType::get(baseTy));
719-
720-
refExpr =
721-
DotSyntaxCallExpr::create(ctx, declRefExpr, SourceLoc(),
722-
Argument::unlabeled(base));
723-
auto refType = adjustedFullType->castTo<FunctionType>()->getResult();
724-
cs.setType(refExpr, refType);
725-
} else {
726-
refExpr = declRefExpr;
727-
}
728-
729-
return forceUnwrapIfExpected(refExpr, locator);
730-
}
731-
}
732-
}
733-
}
734-
}
735-
736617
// Build a reference to the member.
737618
Expr *base =
738619
TypeExpr::createImplicitHack(loc.getBaseNameLoc(), baseTy, ctx);

‎test/AutoDiff/stdlib/simd.swift‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,9 @@
22
// REQUIRES: executable_test
33

44
// Would fail due to unavailability of swift_autoDiffCreateLinearMapContext.
5+
/* Temporary disabled until https://github.com/swiftlang/swift/issues/84840 is fixed
6+
We cannot use `SIMD` :( */
7+
// XFAIL: *
58

69
import _Differentiation
710
import StdlibUnittest

‎test/AutoDiff/validation-test/class_differentiation.swift‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,9 @@
22
// NOTE: Verify whether forward-mode differentiation crashes. It currently does.
33
// RUN: not --crash %target-swift-frontend -enable-experimental-forward-mode-differentiation -emit-sil %s
44
// REQUIRES: executable_test
5+
/* Temporary disabled until https://github.com/swiftlang/swift/issues/84840 is fixed
6+
We cannot use `Tracked<T>` :( */
7+
// XFAIL: *
58

69
import StdlibUnittest
710
import DifferentiationUnittest

‎test/AutoDiff/validation-test/differentiable_property.swift‎

Lines changed: 139 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,144 @@ import DifferentiationUnittest
99

1010
var E2EDifferentiablePropertyTests = TestSuite("E2EDifferentiableProperty")
1111

12+
struct TangentSpace : AdditiveArithmetic {
13+
let x, y: Float
14+
}
15+
16+
extension TangentSpace : Differentiable {
17+
typealias TangentVector = TangentSpace
18+
}
19+
20+
struct Space {
21+
/// `x` is a computed property with a custom vjp.
22+
var x: Float {
23+
@differentiable(reverse)
24+
get { storedX }
25+
set { storedX = newValue }
26+
}
27+
28+
@derivative(of: x)
29+
func vjpX() -> (value: Float, pullback: (Float) -> TangentSpace) {
30+
return (x, { v in TangentSpace(x: v, y: 0) } )
31+
}
32+
33+
private var storedX: Float
34+
35+
@differentiable(reverse)
36+
var y: Float
37+
38+
init(x: Float, y: Float) {
39+
self.storedX = x
40+
self.y = y
41+
}
42+
}
43+
44+
extension Space : Differentiable {
45+
typealias TangentVector = TangentSpace
46+
mutating func move(by offset: TangentSpace) {
47+
x.move(by: offset.x)
48+
y.move(by: offset.y)
49+
}
50+
}
51+
52+
E2EDifferentiablePropertyTests.test("computed property") {
53+
let actualGrad = gradient(at: Space(x: 0, y: 0)) { (point: Space) -> Float in
54+
return 2 * point.x
55+
}
56+
let expectedGrad = TangentSpace(x: 2, y: 0)
57+
expectEqual(expectedGrad, actualGrad)
58+
}
59+
60+
E2EDifferentiablePropertyTests.test("stored property") {
61+
let actualGrad = gradient(at: Space(x: 0, y: 0)) { (point: Space) -> Float in
62+
return 3 * point.y
63+
}
64+
let expectedGrad = TangentSpace(x: 0, y: 3)
65+
expectEqual(expectedGrad, actualGrad)
66+
}
67+
68+
struct GenericMemberWrapper<T : Differentiable> : Differentiable {
69+
// Stored property.
70+
@differentiable(reverse)
71+
var x: T
72+
73+
func vjpX() -> (T, (T.TangentVector) -> GenericMemberWrapper.TangentVector) {
74+
return (x, { TangentVector(x: $0) })
75+
}
76+
}
77+
78+
E2EDifferentiablePropertyTests.test("generic stored property") {
79+
let actualGrad = gradient(at: GenericMemberWrapper<Float>(x: 1)) { point in
80+
return 2 * point.x
81+
}
82+
let expectedGrad = GenericMemberWrapper<Float>.TangentVector(x: 2)
83+
expectEqual(expectedGrad, actualGrad)
84+
}
85+
86+
struct ProductSpaceSelfTangent : AdditiveArithmetic {
87+
let x, y: Float
88+
}
89+
90+
extension ProductSpaceSelfTangent : Differentiable {
91+
typealias TangentVector = ProductSpaceSelfTangent
92+
}
93+
94+
E2EDifferentiablePropertyTests.test("fieldwise product space, self tangent") {
95+
let actualGrad = gradient(at: ProductSpaceSelfTangent(x: 0, y: 0)) { (point: ProductSpaceSelfTangent) -> Float in
96+
return 5 * point.y
97+
}
98+
let expectedGrad = ProductSpaceSelfTangent(x: 0, y: 5)
99+
expectEqual(expectedGrad, actualGrad)
100+
}
101+
102+
struct ProductSpaceOtherTangentTangentSpace : AdditiveArithmetic {
103+
let x, y: Float
104+
}
105+
106+
extension ProductSpaceOtherTangentTangentSpace : Differentiable {
107+
typealias TangentVector = ProductSpaceOtherTangentTangentSpace
108+
}
109+
110+
struct ProductSpaceOtherTangent {
111+
var x, y: Float
112+
}
113+
114+
extension ProductSpaceOtherTangent : Differentiable {
115+
typealias TangentVector = ProductSpaceOtherTangentTangentSpace
116+
mutating func move(by offset: ProductSpaceOtherTangentTangentSpace) {
117+
x.move(by: offset.x)
118+
y.move(by: offset.y)
119+
}
120+
}
121+
122+
E2EDifferentiablePropertyTests.test("fieldwise product space, other tangent") {
123+
let actualGrad = gradient(
124+
at: ProductSpaceOtherTangent(x: 0, y: 0)
125+
) { (point: ProductSpaceOtherTangent) -> Float in
126+
return 7 * point.y
127+
}
128+
let expectedGrad = ProductSpaceOtherTangentTangentSpace(x: 0, y: 7)
129+
expectEqual(expectedGrad, actualGrad)
130+
}
131+
132+
E2EDifferentiablePropertyTests.test("computed property") {
133+
struct TF_544 : Differentiable {
134+
var value: Float
135+
@differentiable(reverse)
136+
var computed: Float {
137+
get { value }
138+
set { value = newValue }
139+
}
140+
}
141+
let actualGrad = gradient(at: TF_544(value: 2.4)) { x in
142+
return x.computed * x.computed
143+
}
144+
let expectedGrad = TF_544.TangentVector(value: 4.8)
145+
expectEqual(expectedGrad, actualGrad)
146+
}
147+
148+
/* Temporary disabled until https://github.com/swiftlang/swift/issues/84840 is fixed
149+
We cannot use `Tracked<T>` :(
12150
struct TangentSpace : AdditiveArithmetic {
13151
let x, y: Tracked<Float>
14152
}
@@ -144,5 +282,6 @@ E2EDifferentiablePropertyTests.testWithLeakChecking("computed property") {
144282
let expectedGrad = TF_544.TangentVector(value: 4.8)
145283
expectEqual(expectedGrad, actualGrad)
146284
}
285+
*/
147286

148287
runAllTests()

‎test/AutoDiff/validation-test/existential.swift‎

Lines changed: 24 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -8,19 +8,39 @@ var ExistentialTests = TestSuite("Existential")
88

99
protocol A {
1010
@differentiable(reverse, wrt: x)
11-
func a(_ x: Tracked<Float>) -> Tracked<Float>
11+
func a(_ x: Float) -> Float
1212
}
13-
func b(g: A) -> Tracked<Float> {
13+
func b(g: A) -> Float {
1414
return gradient(at: 3) { x in g.a(x) }
1515
}
1616

1717
struct B : A {
1818
@differentiable(reverse, wrt: x)
19-
func a(_ x: Tracked<Float>) -> Tracked<Float> { return x * 5 }
19+
func a(_ x: Float) -> Float { return x * 5 }
2020
}
2121

22-
ExistentialTests.testWithLeakChecking("Existential method VJP") {
22+
ExistentialTests.test("Existential method VJP-Tracked") {
2323
expectEqual(5.0, b(g: B()))
2424
}
2525

26+
/* Temporary disabled until https://github.com/swiftlang/swift/issues/84840 is fixed
27+
We cannot use `Tracked<T>` :(
28+
protocol ATracked {
29+
@differentiable(reverse, wrt: x)
30+
func a(_ x: Tracked<Float>) -> Tracked<Float>
31+
}
32+
func b(g: ATracked) -> Tracked<Float> {
33+
return gradient(at: 3) { x in g.a(x) }
34+
}
35+
36+
struct BTracked : ATracked {
37+
@differentiable(reverse, wrt: x)
38+
func a(_ x: Tracked<Float>) -> Tracked<Float> { return x * 5 }
39+
}
40+
41+
ExistentialTests.testWithLeakChecking("Existential method VJP-Tracked") {
42+
expectEqual(5.0, b(g: BTracked()))
43+
}
44+
*/
45+
2646
runAllTests()

‎test/AutoDiff/validation-test/forward_mode_simd.swift‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,8 @@
11
// RUN: %target-run-simple-swift(-Xfrontend -enable-experimental-forward-mode-differentiation)
22
// REQUIRES: executable_test
3+
/* Temporary disabled until https://github.com/swiftlang/swift/issues/84840 is fixed
4+
We cannot use `SIMD` :( */
5+
// XFAIL: *
36

47
import StdlibUnittest
58
import DifferentiationUnittest

‎test/AutoDiff/validation-test/forward_mode_simple.swift‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,8 @@
11
// RUN: %target-run-simple-swift(-Xfrontend -enable-experimental-forward-mode-differentiation)
22
// REQUIRES: executable_test
3+
/* Temporary disabled until https://github.com/swiftlang/swift/issues/84840 is fixed
4+
We cannot use `Tracked<T>` :( */
5+
// XFAIL: *
36

47
import StdlibUnittest
58
import DifferentiationUnittest

0 commit comments

Comments
 (0)