[AutoDiff] Closure specialization: specialize branch tracing enums

kovdan01 · kovdan01 · commit cf717ce56b93 · 2025-10-02T17:45:17.000+03:00
This patch contains part of the changes intended to resolve #68944. 1. Closure info gathering logic. 2. Branch tracing enum specialization logic. 3. Specialization of branch tracing enum basic block arguments in VJP. 4. Specialization of branch tracing enum payload basic block arguments in pullback. 5. C++ helpers for things which are not yet bridged to Swift, such as demangling-related logic.
diff --git a/SwiftCompilerSources/Sources/Optimizer/FunctionPasses/ClosureSpecialization.swift b/SwiftCompilerSources/Sources/Optimizer/FunctionPasses/ClosureSpecialization.swift
diff --git a/SwiftCompilerSources/Sources/Optimizer/Utilities/FunctionTest.swift b/SwiftCompilerSources/Sources/Optimizer/Utilities/FunctionTest.swift
@@ -44,6 +44,7 @@ public func registerOptimizerTests() {
     addressOwnershipLiveRangeTest,
     argumentConventionsTest,
     getPullbackClosureInfoTest,
+    getPullbackClosureInfoMultiBBTest,
     interiorLivenessTest,
     lifetimeDependenceRootTest,
     lifetimeDependenceScopeTest,
@@ -53,6 +54,9 @@ public func registerOptimizerTests() {
     localVariableReachingAssignmentsTest,
     rangeOverlapsPathTest,
     rewrittenCallerBodyTest,
+    specializeBranchTracingEnums,
+    specializeBTEArgInVjpBB,
+    specializePayloadArgInPullbackBB,
     specializedFunctionSignatureAndBodyTest,
     variableIntroducerTest
   )
diff --git a/include/module.modulemap b/include/module.modulemap
@@ -26,3 +26,8 @@ module OptimizerBridging {
   header "swift/SILOptimizer/OptimizerBridging.h"
   export *
 }
+
+module AutoDiffClosureSpecializationBridging {
+  header "swift/SILOptimizer/AutoDiffClosureSpecializationBridging.h"
+  export *
+}
diff --git a/include/swift/SILOptimizer/AutoDiffClosureSpecializationBridging.h b/include/swift/SILOptimizer/AutoDiffClosureSpecializationBridging.h
@@ -0,0 +1,38 @@
+//===--- AutoDiffClosureSpecializationBridging.h --------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2025 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef SWIFT_SILOPTIMIZER_ADCSBRIDGING_H
+#define SWIFT_SILOPTIMIZER_ADCSBRIDGING_H
+
+#include "swift/AST/ASTBridging.h"
+#include "swift/Basic/BasicBridging.h"
+#include "swift/SIL/SILBridging.h"
+
+SWIFT_BEGIN_NULLABILITY_ANNOTATIONS
+
+SWIFT_IMPORT_UNSAFE BridgedType
+getBranchingTraceEnumLoweredType(BridgedDeclObj ed, BridgedFunction vjp);
+
+SWIFT_IMPORT_UNSAFE BridgedType
+getBranchingTraceEnumLoweredType(BridgedEnumDecl ed, BridgedFunction vjp);
+
+SWIFT_IMPORT_UNSAFE BridgedNullableGenericParamList
+cloneGenericParameters(BridgedASTContext ctx, BridgedDeclContext dc,
+                       BridgedCanGenericSignature sig);
+
+SWIFT_IMPORT_UNSAFE BridgedSourceFile autodiffGetSourceFile(BridgedFunction f);
+
+SWIFT_IMPORT_UNSAFE BridgedOwnedString getEnumDeclAsString(BridgedType bteType);
+
+SWIFT_END_NULLABILITY_ANNOTATIONS
+
+#endif
diff --git a/include/swift/SILOptimizer/Differentiation/ADContext.h b/include/swift/SILOptimizer/Differentiation/ADContext.h
@@ -418,6 +418,9 @@ ADContext::emitNondifferentiabilityError(SourceLoc loc,
   llvm_unreachable("Invalid invoker kind"); // silences MSVC C4715
 }
 
+/// Get the source file for the given `SILFunction`.
+SourceFile &getSourceFile(SILFunction *f);
+
 } // end namespace autodiff
 } // end namespace swift
 
diff --git a/include/swift/SILOptimizer/Differentiation/LinearMapInfo.h b/include/swift/SILOptimizer/Differentiation/LinearMapInfo.h
@@ -30,6 +30,14 @@ class SILLoopInfo;
 
 namespace autodiff {
 
+inline SILType getLoweredTypeImpl(CanType ty, SILFunction *derivative,
+                                  Lowering::TypeConverter &typeConverter) {
+  Lowering::AbstractionPattern pattern(
+      derivative->getLoweredFunctionType()->getSubstGenericSignature(), ty);
+  return typeConverter.getLoweredType(pattern, ty,
+                                      TypeExpansionContext::minimal());
+}
+
 class ADContext;
 
 /// Linear map struct and branching trace enum information for an original
@@ -161,13 +169,9 @@ class LinearMapInfo {
   /// the given original block.
   SILType getBranchingTraceEnumLoweredType(SILBasicBlock *origBB) const {
     auto *traceDecl = getBranchingTraceDecl(origBB);
-    auto traceDeclType =
-        traceDecl->getDeclaredInterfaceType()->getCanonicalType();
-    Lowering::AbstractionPattern pattern(
-        derivative->getLoweredFunctionType()->getSubstGenericSignature(),
-        traceDeclType);
-    return typeConverter.getLoweredType(pattern, traceDeclType,
-                                        TypeExpansionContext::minimal());
+    return getLoweredTypeImpl(
+        traceDecl->getDeclaredInterfaceType()->getCanonicalType(), derivative,
+        typeConverter);
   }
 
   /// Returns the enum element in the given successor block's branching trace
@@ -199,6 +203,11 @@ class LinearMapInfo {
   }
 };
 
+/// Clone the generic parameters of the given generic signature and return a new
+/// `GenericParamList`.
+GenericParamList *cloneGenericParameters(ASTContext &ctx, DeclContext *dc,
+                                         CanGenericSignature sig);
+
 } // end namespace autodiff
 } // end namespace swift
 
diff --git a/lib/SILOptimizer/Differentiation/ADContext.cpp b/lib/SILOptimizer/Differentiation/ADContext.cpp
@@ -60,8 +60,7 @@ ADContext::ADContext(SILModuleTransform &transform)
     : transform(transform), module(*transform.getModule()),
       passManager(*transform.getPassManager()) {}
 
-/// Get the source file for the given `SILFunction`.
-static SourceFile &getSourceFile(SILFunction *f) {
+SourceFile &getSourceFile(SILFunction *f) {
   if (f->hasLocation())
     if (auto *declContext = f->getLocation().getAsDeclContext())
       if (auto *parentSourceFile = declContext->getParentSourceFile())
diff --git a/lib/SILOptimizer/Differentiation/LinearMapInfo.cpp b/lib/SILOptimizer/Differentiation/LinearMapInfo.cpp
@@ -33,11 +33,8 @@ namespace autodiff {
 // Local helpers
 //===----------------------------------------------------------------------===//
 
-/// Clone the generic parameters of the given generic signature and return a new
-/// `GenericParamList`.
-static GenericParamList *cloneGenericParameters(ASTContext &ctx,
-                                                DeclContext *dc,
-                                                CanGenericSignature sig) {
+GenericParamList *cloneGenericParameters(ASTContext &ctx, DeclContext *dc,
+                                         CanGenericSignature sig) {
   SmallVector<GenericTypeParamDecl *, 2> clonedParams;
   for (auto paramType : sig.getGenericParams()) {
     auto *clonedParam = GenericTypeParamDecl::createImplicit(
diff --git a/lib/SILOptimizer/Utils/AutoDiffClosureSpecializationBridging.cpp b/lib/SILOptimizer/Utils/AutoDiffClosureSpecializationBridging.cpp
@@ -0,0 +1,56 @@
+//===--- AutoDiffClosureSpecializationBridging.cpp ------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2025 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+#define DEBUG_TYPE "autodiff-closure-specialization-bridging"
+
+#include "swift/SILOptimizer/AutoDiffClosureSpecializationBridging.h"
+#include "swift/SILOptimizer/Differentiation/ADContext.h"
+#include "swift/SILOptimizer/Differentiation/LinearMapInfo.h"
+
+using namespace swift;
+
+static SILType getBranchingTraceEnumLoweredTypeImpl(EnumDecl *ed,
+                                                    SILFunction &vjp) {
+  return autodiff::getLoweredTypeImpl(
+      ed->getDeclaredInterfaceType()->getCanonicalType(), &vjp,
+      vjp.getModule().Types);
+}
+
+BridgedType getBranchingTraceEnumLoweredType(BridgedEnumDecl ed,
+                                              BridgedFunction vjp) {
+  return getBranchingTraceEnumLoweredTypeImpl(ed.unbridged(),
+                                              *vjp.getFunction());
+}
+
+BridgedType getBranchingTraceEnumLoweredType(BridgedDeclObj ed,
+                                             BridgedFunction vjp) {
+  return getBranchingTraceEnumLoweredTypeImpl(ed.getAs<EnumDecl>(),
+                                              *vjp.getFunction());
+}
+
+BridgedNullableGenericParamList
+cloneGenericParameters(BridgedASTContext ctx, BridgedDeclContext dc,
+                       BridgedCanGenericSignature sig) {
+  return autodiff::cloneGenericParameters(ctx.unbridged(), dc.unbridged(),
+                                          sig.unbridged());
+}
+
+BridgedSourceFile autodiffGetSourceFile(BridgedFunction f) {
+  return {&autodiff::getSourceFile(f.getFunction())};
+}
+
+BridgedOwnedString getEnumDeclAsString(BridgedType bteType) {
+  std::string str;
+  llvm::raw_string_ostream out(str);
+  bteType.unbridged().getEnumOrBoundGenericEnum()->print(out);
+  return BridgedOwnedString(/*stringToCopy=*/StringRef(str));
+}
diff --git a/lib/SILOptimizer/Utils/CMakeLists.txt b/lib/SILOptimizer/Utils/CMakeLists.txt
@@ -1,4 +1,5 @@
 target_sources(swiftSILOptimizer PRIVATE
+  AutoDiffClosureSpecializationBridging.cpp
   BasicBlockOptUtils.cpp
   CFGOptUtils.cpp
   CanonicalizeInstruction.cpp
diff --git a/lib/SILOptimizer/Utils/OptimizerBridging.cpp b/lib/SILOptimizer/Utils/OptimizerBridging.cpp
@@ -13,6 +13,7 @@
 #include "swift/SILOptimizer/OptimizerBridging.h"
 #include "../../IRGen/IRGenModule.h"
 #include "swift/AST/SemanticAttrs.h"
+#include "swift/Demangling/ManglingMacros.h"
 #include "swift/SIL/DynamicCasts.h"
 #include "swift/SIL/OSSACompleteLifetime.h"
 #include "swift/SIL/SILCloner.h"
@@ -527,6 +528,73 @@ bool BridgedFunction::isAutodiffVJP() const {
       getFunction(), swift::AutoDiffFunctionComponent::VJP);
 }
 
+bool BridgedFunction::isAutodiffSubsetParametersThunk() const {
+  Demangle::Context Ctx;
+  if (auto *root = Ctx.demangleSymbolAsNode(getFunction()->getName())) {
+    // root node has Global kind, the AutoDiffSubsetParametersThunk node (if
+    // present) is direct child of root.
+    return root->findByKind(Demangle::Node::Kind::AutoDiffSubsetParametersThunk,
+                            /*maxDepth=*/1) != nullptr;
+  }
+  return false;
+}
+
+// See also ASTMangler::mangleAutoDiffGeneratedDeclaration.
+bool BridgedType::isAutodiffBranchTracingEnumInVJP(BridgedFunction vjp) const {
+  assert(vjp.isAutodiffVJP());
+  EnumDecl *ed = unbridged().getEnumOrBoundGenericEnum();
+  if (ed == nullptr)
+    return false;
+
+  llvm::StringRef edName = ed->getNameStr();
+  if (!edName.starts_with("_AD__"))
+    return false;
+  if (!llvm::StringRef(edName.data() + 5, edName.size() - 5)
+           .starts_with(MANGLING_PREFIX_STR))
+    return false;
+
+  // At this point, we know that the type is indeed a branch tracing enum.
+  // Now we need to ensure that it is the enum related to the given VJP.
+
+  std::size_t idx = edName.rfind("__Pred__");
+  assert(idx != std::string::npos);
+
+  // Before "__Pred__", we have "_bbX", where X is a number.
+  // The loop calculates the start position of X.
+  for (; idx != 0 && std::isdigit(edName[idx - 1]); --idx)
+    ;
+
+  assert(std::isdigit(edName[idx]));
+  assert(!std::isdigit(edName[idx - 1]));
+
+  // The branch tracing enum decl name has the following components:
+  // 1) "_AD__";
+  // 2) MANGLING_PREFIX;
+  // 3) original function name;
+  // 4) "_bb";
+  // 5) X at position idx (see above);
+  // 6) the rest of the enum decl name.
+  // Thus, "_AD__", MANGLING_PREFIX and "_bb" must have total length less than
+  // idx.
+  std::size_t manglingPrefixSize = std::strlen(MANGLING_PREFIX_STR);
+  assert(idx > 5 + manglingPrefixSize + 3);
+  assert(std::string_view(edName.data() + idx - 3, 3) == "_bb");
+  assert(std::string_view(edName.data(), 5 + manglingPrefixSize) == "_AD__$s");
+
+  llvm::StringRef enumOrigFuncName =
+      std::string_view(edName.data() + 5 + manglingPrefixSize,
+                       idx - (5 + manglingPrefixSize + 3));
+
+  Demangle::Context Ctx;
+  if (auto *root = Ctx.demangleSymbolAsNode(vjp.getFunction()->getName()))
+    if (auto *node =
+            root->findByKind(Demangle::Node::Kind::Function, /*maxDepth=*/3))
+      if (mangleNode(node).result() == enumOrigFuncName)
+        return true;
+
+  return false;
+}
+
 SwiftInt BridgedFunction::specializationLevel() const {
   return swift::getSpecializationLevel(getFunction());
 }
diff --git a/test/AutoDiff/SILOptimizer/closure_specialization/multi_bb_bte.sil b/test/AutoDiff/SILOptimizer/closure_specialization/multi_bb_bte.sil

Original file line number	Diff line number	Diff line change
`@@ -26,3 +26,8 @@ module OptimizerBridging {`
`26`	`26`	`header "swift/SILOptimizer/OptimizerBridging.h"`
`27`	`27`	`export *`
`28`	`28`	`}`
	`29`	`+`
	`30`	`+module AutoDiffClosureSpecializationBridging {`
	`31`	`+ header "swift/SILOptimizer/AutoDiffClosureSpecializationBridging.h"`
	`32`	`+ export *`
	`33`	`+}`
Original file line number	Diff line number	Diff line change
`@@ -418,6 +418,9 @@ ADContext::emitNondifferentiabilityError(SourceLoc loc,`
`418`	`418`	`llvm_unreachable("Invalid invoker kind"); // silences MSVC C4715`
`419`	`419`	`}`
`420`	`420`
	`421`	+/// Get the source file for the given `SILFunction`.
	`422`	`+SourceFile &getSourceFile(SILFunction *f);`
	`423`	`+`
`421`	`424`	`} // end namespace autodiff`
`422`	`425`	`} // end namespace swift`
`423`	`426`
Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,5 @@`
`1`	`1`	`target_sources(swiftSILOptimizer PRIVATE`
	`2`	`+ AutoDiffClosureSpecializationBridging.cpp`
`2`	`3`	`BasicBlockOptUtils.cpp`
`3`	`4`	`CFGOptUtils.cpp`
`4`	`5`	`CanonicalizeInstruction.cpp`