[NFC] [AutoDiff] Gardening: @differentiable function type mangling. (#26600)

- Add todo comments tracking TF-750.
- Optimize `SWIFT_ENABLE_TENSORFLOW` comments.
- Clarify comment in test/AutoDiff/differentiable_func_debuginfo.swift.

Addresses feedback from #26595.

Author: dan-zheng

lib/Demangling/OldRemangler.cpp

@@ -1269,15 +1269,15 @@ void Remangler::mangleImplResult(Node *node) {
 
 // SWIFT_ENABLE_TENSORFLOW
 void Remangler::mangleImplDifferentiable(Node *node) {
-  // The old mangler does not encode `@differentiable` function types.
+  // TODO(TF-750): Check if this code path actually triggers and add a test.
   Buffer << 'd';
 }
 
-// SWIFT_ENABLE_TENSORFLOW
 void Remangler::mangleImplLinear(Node *node) {
-  // The old mangler does not encode `@differentiable(linear)` function types.
+  // TODO(TF-750): Check if this code path actually triggers and add a test.
   Buffer << 'l';
 }
+// SWIFT_ENABLE_TENSORFLOW END
 
 void Remangler::mangleImplEscaping(Node *node) {
   // The old mangler does not encode escaping.

lib/Demangling/Remangler.cpp

@@ -1204,23 +1204,21 @@ void Remangler::mangleDifferentiableFunctionType(Node *node) {
   Buffer << "XF";
 }
 
-// SWIFT_ENABLE_TENSORFLOW
 void Remangler::mangleEscapingDifferentiableFunctionType(Node *node) {
   mangleFunctionSignature(node);
   Buffer << "XG";
 }
 
-// SWIFT_ENABLE_TENSORFLOW
 void Remangler::mangleLinearFunctionType(Node *node) {
   mangleFunctionSignature(node);
   Buffer << "XH";
 }
 
-// SWIFT_ENABLE_TENSORFLOW
 void Remangler::mangleEscapingLinearFunctionType(Node *node) {
   mangleFunctionSignature(node);
   Buffer << "XI";
 }
+// SWIFT_ENABLE_TENSORFLOW END
 
 void Remangler::mangleGenericProtocolWitnessTable(Node *node) {
   mangleSingleChildNode(node);
@@ -1384,10 +1382,10 @@ void Remangler::mangleImplDifferentiable(Node *node) {
   Buffer << 'd';
 }
 
-// SWIFT_ENABLE_TENSORFLOW
 void Remangler::mangleImplLinear(Node *node) {
   Buffer << 'l';
 }
+// SWIFT_ENABLE_TENSORFLOW END
 
 void Remangler::mangleImplEscaping(Node *node) {
   Buffer << 'e';

test/AutoDiff/differentiable_func_debuginfo.swift

@@ -27,7 +27,14 @@ public protocol Differentiable {
   associatedtype TangentVector
 }
 
-// This declaration is necessary to reproduce the crash for some reason.
+// This declaration is necessary to reproduce the crash.
+// `Builtin.autodiffApply_vjp` constructs a use of the `tf597ProblematicVarDecl`
+// type, which was mangled without `@differentiable` attribute. The parameter
+// for `blackHole` is of type `$@noescape @callee_guaranteed (@in_guaranteed T) -> @out U`,
+// which matched the mangled name for the type of the parameter of `blackHole`.
+// As a result, the types are uniqued when generating debug info. The type of
+// the parameter of `blackHole` is smaller than the `@differentiable` function
+// type, causing IRGenDebugInfo to crash.
 public func blackHole<T, U>(_: (T) -> U) {}
 
 public func pullback<T, R>(

test/Demangle/Inputs/manglings.txt

@@ -354,3 +354,4 @@ $s3red4testyAA7OurTypeOy4them05TheirD0Vy5AssocQzGAjE0F8ProtocolAAxAA0c7DerivedH0
 // SWIFT_ENABLE_TENSORFLOW
 $sxq_Idgnr_D ---> @differentiable @callee_guaranteed (@in_guaranteed A) -> (@out B)
 $sxq_Ilgnr_D ---> @differentiable(linear) @callee_guaranteed (@in_guaranteed A) -> (@out B)
+// SWIFT_ENABLE_TENSORFLOW END

test/SILGen/mangling.swift

@@ -189,6 +189,6 @@ func varargsVsArray(arr: [Int]..., n: String) { }
 // CHECK-LABEL: sil hidden [ossa] @$s8mangling15funcVsDiffFunc12fnyS2fXE_tF : $@convention(thin) (@noescape @callee_guaranteed (Float) -> Float) -> ()
 func funcVsDiffFunc1(fn: (Float) -> Float) {}
 
-// SWIFT_ENABLE_TENSORFLOW
 // CHECK-LABEL: sil hidden [ossa] @$s8mangling15funcVsDiffFunc22fnyS2fXF_tF : $@convention(thin) (@differentiable @noescape @callee_guaranteed (Float) -> Float) -> ()
 func funcVsDiffFunc2(fn: @differentiable (Float) -> Float) {}
+// SWIFT_ENABLE_TENSORFLOW END

test/TypeDecoder/structural_types.swift

@@ -140,7 +140,6 @@ do {
   _ = f
 }
 
-// SWIFT_ENABLE_TENSORFLOW
 do {
   let f: (@escaping @differentiable (Float) -> Float) -> () = { _ in }
   // FIXME(TF-123): `@differentiable` function type + opaque abstraction
@@ -151,7 +150,6 @@ do {
 
 // TODO: Uncomment when `@differentiable(linear)` function types are enabled.
 /*
-// SWIFT_ENABLE_TENSORFLOW
 do {
   let f: @differentiable(linear) (Float) -> Float = { $0 }
   // FIXME(TF-123): `@differentiable` function type + opaque abstraction
@@ -160,7 +158,6 @@ do {
   _ = f
 }
 
-// SWIFT_ENABLE_TENSORFLOW
 do {
   let f: (@escaping @differentiable(linear) (Float) -> Float) -> () = { _ in }
   // FIXME(TF-123): `@differentiable` function type + opaque abstraction
@@ -169,6 +166,7 @@ do {
   _ = f
 }
 */
+// SWIFT_ENABLE_TENSORFLOW END
 
 // DEMANGLE: $syycD
 // DEMANGLE: $sySSzcD
@@ -188,10 +186,13 @@ do {
 // DEMANGLE: $syyyccD
 // DEMANGLE: $sSayyyXCGD
 // DEMANGLE: $sSayyyyXL_yyXBtcGD
+
+// SWIFT_ENABLE_TENSORFLOW
 // DEMANGLE: $sS2fXFD
 // DEMANGLE: $sS2fXGD
 // DEMANGLE: $sS2fXHD
 // DEMANGLE: $sS2fXID
+// SWIFT_ENABLE_TENSORFLOW END
 
 // CHECK: () -> ()
 // CHECK: (inout String) -> ()
@@ -211,10 +212,13 @@ do {
 // CHECK: (@escaping () -> ()) -> ()
 // CHECK: Array<@convention(c) () -> ()>
 // CHECK: Array<(@escaping @convention(block) () -> (), @convention(block) () -> ()) -> ()>
+
+// SWIFT_ENABLE_TENSORFLOW
 // CHECK: @differentiable (Float) -> Float
 // CHECK: @differentiable (Float) -> Float
 // CHECK: @differentiable(linear) (Float) -> Float
 // CHECK: @differentiable(linear) (Float) -> Float
+// SWIFT_ENABLE_TENSORFLOW END
 
 // DEMANGLE: $sSimD
 // DEMANGLE: $syycmD
@@ -235,10 +239,13 @@ do {
 // DEMANGLE: $syyyccmD
 // DEMANGLE: $sSayyyXCGmD
 // DEMANGLE: $sSayyyyXL_yyXBtcGmD
+
+// SWIFT_ENABLE_TENSORFLOW
 // DEMANGLE: $sS2fXFmD
 // DEMANGLE: $sS2fXGmD
 // DEMANGLE: $sS2fXHmD
 // DEMANGLE: $sS2fXImD
+// SWIFT_ENABLE_TENSORFLOW END
 
 // CHECK: Int.Type
 // CHECK: ((inout String) -> ()).Type
@@ -258,7 +265,10 @@ do {
 // CHECK: ((@escaping () -> ()) -> ()).Type
 // CHECK: Array<@convention(c) () -> ()>.Type
 // CHECK: Array<(@escaping @convention(block) () -> (), @convention(block) () -> ()) -> ()>.Type
+
+// SWIFT_ENABLE_TENSORFLOW
 // CHECK: (@differentiable (Float) -> Float).Type
 // CHECK: (@differentiable (Float) -> Float).Type
 // CHECK: (@differentiable(linear) (Float) -> Float).Type
 // CHECK: (@differentiable(linear) (Float) -> Float).Type
+// SWIFT_ENABLE_TENSORFLOW END