[WIP] TF-1021: verify that retrodiff works.

dan-zheng · dan-zheng · commit 8c38065ac74a · 2019-12-09T12:16:06.000-08:00
Added a simple four file test:
- A: defines original function.
- B: registers derivative for original function.
- C: also registers derivative for original function.
- D: imports A, B, and C, and differentiates original function.
diff --git a/lib/Sema/TypeCheckAttr.cpp b/lib/Sema/TypeCheckAttr.cpp
@@ -3906,6 +3906,7 @@ void AttributeChecker::visitDerivativeAttr(DerivativeAttr *attr) {
     return;
   }
 
+#if 0
   // Reject different-file retroactive derivatives.
   // TODO(TF-136): Lift this restriction now that SIL differentiability witness
   // infrastructure is ready.
@@ -3914,6 +3915,7 @@ void AttributeChecker::visitDerivativeAttr(DerivativeAttr *attr) {
                           diag::derivative_attr_not_in_same_file_as_original);
     return;
   }
+#endif
 
   // Valid `@derivative` attributes are uniqued by original function and
   // parameter indices. Reject duplicate attributes.
diff --git a/test/AutoDiff/Inputs/derivative_registration_derivative_module.swift b/test/AutoDiff/Inputs/derivative_registration_derivative_module.swift
@@ -0,0 +1,6 @@
+import derivative_registration_original_module
+
+@derivative(of: id)
+func vjpId(_ x: Float) -> (value: Float, pullback: (Float) -> Float) {
+  return (id(x), { $0 })
+}
diff --git a/test/AutoDiff/Inputs/derivative_registration_original_module.swift b/test/AutoDiff/Inputs/derivative_registration_original_module.swift
@@ -0,0 +1 @@
+public func id(_ x: Float) -> Float { x }
diff --git a/test/AutoDiff/derivative_registration/main.swift b/test/AutoDiff/derivative_registration/main.swift
@@ -0,0 +1,18 @@
+// Verify that cross-file derivative registration works.
+
+// RUN: %empty-directory(%t)
+// RUN: %target-swift-frontend -emit-module -primary-file %S/../Inputs/derivative_registration_original_module.swift -emit-module-path %t/derivative_registration_original_module.swiftmodule
+// RUN: %target-swift-frontend -I %t -emit-module -primary-file %S/../Inputs/derivative_registration_derivative_module.swift -emit-module-path %t/derivative_registration_derivative_module.swiftmodule
+// RUN: %target-swift-frontend -I %t -emit-module -primary-file %S/../Inputs/derivative_registration_derivative_module2.swift -emit-module-path %t/derivative_registration_derivative_module2.swiftmodule
+// RUN: %target-swift-emit-sil -I %t -emit-module %s
+
+import derivative_registration_original_module
+import derivative_registration_derivative_module
+import derivative_registration_derivative_module2
+
+func x(_ x: Float) -> Float {
+  // TODO: Check which derivative for `id` is used:
+  // - The one from `derivative_registration_derivative_module`, or
+  // - The one from `derivative_registration_derivative_module2`.
+  return gradient(at: x, in: id)
+}

Original file line number	Diff line number	Diff line change
`@@ -3906,6 +3906,7 @@ void AttributeChecker::visitDerivativeAttr(DerivativeAttr *attr) {`
`3906`	`3906`	`return;`
`3907`	`3907`	`}`
`3908`	`3908`
	`3909`	`+#if 0`
`3909`	`3910`	`// Reject different-file retroactive derivatives.`
`3910`	`3911`	`// TODO(TF-136): Lift this restriction now that SIL differentiability witness`
`3911`	`3912`	`// infrastructure is ready.`
`@@ -3914,6 +3915,7 @@ void AttributeChecker::visitDerivativeAttr(DerivativeAttr *attr) {`
`3914`	`3915`	`diag::derivative_attr_not_in_same_file_as_original);`
`3915`	`3916`	`return;`
`3916`	`3917`	`}`
	`3918`	`+#endif`
`3917`	`3919`
`3918`	`3920`	// Valid `@derivative` attributes are uniqued by original function and
`3919`	`3921`	`// parameter indices. Reject duplicate attributes.`
Original file line number	Diff line number	Diff line change
`@@ -0,0 +1 @@`
	`1`	`+public func id(_ x: Float) -> Float { x }`