[AutoDiff] Add mangling for @differentiable function types. (#26595)

Add mangling support for `@differentiable` and `@differentiable(linear)`
function types.

This fixes debug info builds. Previously, IRGenDebugInfo crashed because
`@differentiable` function types did not have full mangling support.
`IRGenDebugInfoImpl::getOrCreateType` computes `llvm::DIType` type debug
info by demangling type names, leading to a type size inconsistency for
`@differentiable` function types.

Change `tensorflow_osx_base` build preset to use `--release-debuginfo`.
macOS toolchains (with TensorFlow) now have debug info again.

Resolves TF-597.

Author: dan-zheng

docs/ABI/Mangling.rst

@@ -583,12 +583,16 @@ mangled in to disambiguate.
   impl-function-type ::= type* 'I' FUNC-ATTRIBUTES '_'
   impl-function-type ::= type* generic-signature 'I' PSEUDO-GENERIC? FUNC-ATTRIBUTES '_'
 
-  FUNC-ATTRIBUTES ::= CALLEE-ESCAPE? CALLEE-CONVENTION FUNC-REPRESENTATION? PARAM-CONVENTION* RESULT-CONVENTION* ('z' RESULT-CONVENTION)
+  FUNC-ATTRIBUTES ::= CALLEE-ESCAPE? DIFFERENTIABILITY-KIND? CALLEE-CONVENTION FUNC-REPRESENTATION? PARAM-CONVENTION* RESULT-CONVENTION* ('z' RESULT-CONVENTION)
 
   PSEUDO-GENERIC ::= 'P'
 
   CALLEE-ESCAPE ::= 'e'                      // @escaping (inverse of SIL @noescape)
 
+  DIFFERENTIABILITY-KIND ::= DIFFERENTIABLE | LINEAR
+  DIFFERENTIABLE ::= 'd'                     // @differentiable
+  LINEAR ::= 'l'                             // @differentiable(linear)
+
   CALLEE-CONVENTION ::= 'y'                  // @callee_unowned
   CALLEE-CONVENTION ::= 'g'                  // @callee_guaranteed
   CALLEE-CONVENTION ::= 'x'                  // @callee_owned

include/swift/Demangling/DemangleNodes.def

@@ -115,6 +115,10 @@ NODE(Identifier)
 NODE(Index)
 CONTEXT_NODE(IVarInitializer)
 CONTEXT_NODE(IVarDestroyer)
+// SWIFT_ENABLE_TENSORFLOW
+NODE(ImplDifferentiable)
+NODE(ImplLinear)
+// SWIFT_ENABLE_TENSORFLOW END
 NODE(ImplEscaping)
 NODE(ImplConvention)
 NODE(ImplFunctionAttribute)

lib/AST/ASTMangler.cpp

@@ -1445,6 +1445,18 @@ void ASTMangler::appendImplFunctionType(SILFunctionType *fn) {
   if (!fn->isNoEscape())
     OpArgs.push_back('e');
 
+  // SWIFT_ENABLE_TENSORFLOW
+  switch (fn->getExtInfo().getDifferentiabilityKind()) {
+  case DifferentiabilityKind::NonDifferentiable:
+    break;
+  case DifferentiabilityKind::Normal:
+    OpArgs.push_back('d');
+    break;
+  case DifferentiabilityKind::Linear:
+    OpArgs.push_back('l');
+    break;
+  }
+
   // <impl-callee-convention>
   if (fn->getExtInfo().hasContext()) {
     OpArgs.push_back(getParamConvention(fn->getCalleeConvention()));

lib/Demangling/Demangler.cpp

@@ -1729,6 +1729,12 @@ NodePointer Demangler::demangleImplFunctionType() {
   if (nextIf('e'))
     type->addChild(createNode(Node::Kind::ImplEscaping), *this);
 
+  // SWIFT_ENABLE_TENSORFLOW
+  if (nextIf('d'))
+    type->addChild(createNode(Node::Kind::ImplDifferentiable), *this);
+  if (nextIf('l'))
+    type->addChild(createNode(Node::Kind::ImplLinear), *this);
+
   const char *CAttr = nullptr;
   switch (nextChar()) {
     case 'y': CAttr = "@callee_unowned"; break;

lib/Demangling/NodePrinter.cpp

@@ -393,6 +393,10 @@ class NodePrinter {
     case Node::Kind::Index:
     case Node::Kind::IVarInitializer:
     case Node::Kind::IVarDestroyer:
+    // SWIFT_ENABLE_TENSORFLOW
+    case Node::Kind::ImplDifferentiable:
+    case Node::Kind::ImplLinear:
+    // SWIFT_ENABLE_TENSORFLOW END
     case Node::Kind::ImplEscaping:
     case Node::Kind::ImplConvention:
     case Node::Kind::ImplFunctionAttribute:
@@ -2019,6 +2023,14 @@ NodePointer NodePrinter::print(NodePointer Node, bool asPrefixContext) {
     return nullptr;
   case Node::Kind::LabelList:
     return nullptr;
+  // SWIFT_ENABLE_TENSORFLOW
+  case Node::Kind::ImplDifferentiable:
+    Printer << "@differentiable";
+    return nullptr;
+  case Node::Kind::ImplLinear:
+    Printer << "@differentiable(linear)";
+    return nullptr;
+  // SWIFT_ENABLE_TENSORFLOW END
   case Node::Kind::ImplEscaping:
     Printer << "@escaping";
     return nullptr;

lib/Demangling/OldRemangler.cpp

@@ -1267,6 +1267,18 @@ void Remangler::mangleImplResult(Node *node) {
   mangleChildNodes(node); // impl convention, type
 }
 
+// SWIFT_ENABLE_TENSORFLOW
+void Remangler::mangleImplDifferentiable(Node *node) {
+  // The old mangler does not encode `@differentiable` function types.
+  Buffer << 'd';
+}
+
+// SWIFT_ENABLE_TENSORFLOW
+void Remangler::mangleImplLinear(Node *node) {
+  // The old mangler does not encode `@differentiable(linear)` function types.
+  Buffer << 'l';
+}
+
 void Remangler::mangleImplEscaping(Node *node) {
   // The old mangler does not encode escaping.
 }

lib/Demangling/Remangler.cpp

@@ -1379,6 +1379,16 @@ void Remangler::mangleIVarDestroyer(Node *node) {
   Buffer << "fE";
 }
 
+// SWIFT_ENABLE_TENSORFLOW
+void Remangler::mangleImplDifferentiable(Node *node) {
+  Buffer << 'd';
+}
+
+// SWIFT_ENABLE_TENSORFLOW
+void Remangler::mangleImplLinear(Node *node) {
+  Buffer << 'l';
+}
+
 void Remangler::mangleImplEscaping(Node *node) {
   Buffer << 'e';
 }
@@ -1423,6 +1433,14 @@ void Remangler::mangleImplFunctionType(Node *node) {
   Buffer << 'I' << PseudoGeneric;
   for (NodePointer Child : *node) {
     switch (Child->getKind()) {
+      // SWIFT_ENABLE_TENSORFLOW
+      case Node::Kind::ImplDifferentiable:
+        Buffer << 'd';
+        break;
+      case Node::Kind::ImplLinear:
+        Buffer << 'l';
+        break;
+      // SWIFT_ENABLE_TENSORFLOW END
       case Node::Kind::ImplEscaping:
         Buffer << 'e';
         break;

lib/IRGen/GenDiffFunc.cpp

@@ -32,6 +32,7 @@
 using namespace swift;
 using namespace irgen;
 
+/// A pair of `@differentiable` function extractee and differentiation order.
 using DiffFuncIndex =
     std::pair<AutoDiffFunctionExtractInst::Extractee, unsigned>;
 

test/AutoDiff/differentiable_func_debuginfo.swift

@@ -0,0 +1,37 @@
+// RUN: %target-swift-frontend -c %s -g -O -parse-stdlib -parse-as-library -module-name Swift
+
+// TF-597: Exact minimal reproducer for IRGenDebugInfo crash.
+// Crash occurred only with `-g` and `-O`.
+//
+// ```
+// Assertion failed: (OffsetInBits + SizeInBits <= getSizeInBits(Var) && "pars > totum"),
+// function emitVariableDeclaration, file swift/lib/IRGen/IRGenDebugInfo.cpp, line 2216.
+// Stack dump:
+// 1.	Swift version 5.1-dev (LLVM 200186e28b, Swift c09c14dec5)
+// 2.	While emitting IR SIL function "@$ss8pullback2at2inyx_q_xXFts14DifferentiableRzsADR_r0_lF".
+//  for 'pullback(at:in:)' (at swift/test/AutoDiff/differentiable_func_debuginfo.swift:21:8)
+// ```
+//
+// The crash was because `IRGenDebugInfoImpl::getOrCreateType` computes
+// `llvm::DIType` type debug info by demangling type names.
+//
+// Since `@differentiable` and `@differentiable(linear)` function types did
+// not have mangling support, `getOrCreateType` computed a regular `(A) -> B`
+// function type instead of a `@differentiable (A) -> B` function type, leading
+// to a type size inconsistency.
+//
+// Conclusion: mangling coverage is important.
+
+// Minimal dummy compiler-known `Differentiable` protocol.
+public protocol Differentiable {
+  associatedtype TangentVector
+}
+
+// This declaration is necessary to reproduce the crash for some reason.
+public func blackHole<T, U>(_: (T) -> U) {}
+
+public func pullback<T, R>(
+  at x: T, in tf597ProblematicVarDecl: @differentiable (T) -> R
+) {
+  let _ = Builtin.autodiffApply_vjp(tf597ProblematicVarDecl, x)
+}

test/Demangle/Inputs/manglings.txt

@@ -351,3 +351,6 @@ $s18keypaths_inlinable13KeypathStructV8computedSSvpACTKq  ---> key path getter f
 $s18resilient_protocol24ResilientDerivedProtocolPxAA0c4BaseE0Tn --> associated conformance descriptor for resilient_protocol.ResilientDerivedProtocol.A: resilient_protocol.ResilientBaseProtocol
 $s3red4testyAA3ResOyxSayq_GAEs5ErrorAAq_sAFHD1__HCg_GADyxq_GsAFR_r0_lF --> red.test<A, B where B: Swift.Error>(red.Res<A, B>) -> red.Res<A, [B]>
 $s3red4testyAA7OurTypeOy4them05TheirD0Vy5AssocQzGAjE0F8ProtocolAAxAA0c7DerivedH0HD1_AA0c4BaseH0HI1_AieKHA2__HCg_GxmAaLRzlF ---> red.test<A where A: red.OurDerivedProtocol>(A.Type) -> red.OurType<them.TheirType<A.Assoc>>
+// 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)

test/SILGen/mangling.swift

@@ -184,3 +184,11 @@ func varargsVsArray(arr: [Int], n: String) { }
 
 // CHECK-LABEL: sil hidden [ossa] @$s8mangling14varargsVsArray3arr1nySaySiGd_SStF : $@convention(thin) (@guaranteed Array<Array<Int>>, @guaranteed String) -> ()
 func varargsVsArray(arr: [Int]..., n: String) { }
+
+// SWIFT_ENABLE_TENSORFLOW
+// 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) {}

utils/build-presets.ini

@@ -2143,10 +2143,7 @@ libcxx
 indexstore-db
 sourcekit-lsp
 
-# SWIFT_ENABLE_TENSORFLOW
-# FIXME(TF-597): IRGenDebugInfo crash for `pullback(at:in:)`.
-# release-debuginfo
-release
+release-debuginfo
 
 lldb-no-debugserver
 lldb-use-system-debugserver