[AutoDiff] Bump-pointer allocate pullback structs in loops. (#34886)

In derivatives of loops, no longer allocate boxes for indirect case payloads. Instead, use a custom pullback context in the runtime which contains a bump-pointer allocator.

When a function contains a differentiated loop, the closure context is a `Builtin.NativeObject`, which contains a `swift::AutoDiffLinearMapContext` and a tail-allocated top-level linear map struct (which represents the linear map struct that was previously directly partial-applied into the pullback). In branching trace enums, the payloads of previously indirect cases will be allocated by `swift::AutoDiffLinearMapContext::allocate` and stored as a `Builtin.RawPointer`.

Author: rxwei

include/swift/AST/ASTContext.h

@@ -713,6 +713,9 @@ class ASTContext final {
   /// Get the runtime availability of support for concurrency.
   AvailabilityContext getConcurrencyAvailability();
 
+  /// Get the runtime availability of support for differentiation.
+  AvailabilityContext getDifferentiationAvailability();
+
   /// Get the runtime availability of features introduced in the Swift 5.2
   /// compiler for the target platform.
   AvailabilityContext getSwift52Availability();

include/swift/AST/Builtins.def

@@ -752,6 +752,15 @@ BUILTIN_MISC_OPERATION_WITH_SILGEN(CreateAsyncTaskFuture,
 /// is a pure value and therefore we can consider it as readnone).
 BUILTIN_MISC_OPERATION_WITH_SILGEN(GlobalStringTablePointer, "globalStringTablePointer", "n", Special)
 
+// autoDiffCreateLinearMapContext: (Builtin.Word) -> Builtin.NativeObject
+BUILTIN_MISC_OPERATION_WITH_SILGEN(AutoDiffCreateLinearMapContext, "autoDiffCreateLinearMapContext", "n", Special)
+
+// autoDiffProjectTopLevelSubcontext: (Builtin.NativeObject) -> Builtin.RawPointer
+BUILTIN_MISC_OPERATION_WITH_SILGEN(AutoDiffProjectTopLevelSubcontext, "autoDiffProjectTopLevelSubcontext", "n", Special)
+
+// autoDiffAllocateSubcontext: (Builtin.NativeObject, Builtin.Word) -> Builtin.RawPointer
+BUILTIN_MISC_OPERATION_WITH_SILGEN(AutoDiffAllocateSubcontext, "autoDiffAllocateSubcontext", "", Special)
+
 #undef BUILTIN_MISC_OPERATION_WITH_SILGEN
 
 #undef BUILTIN_MISC_OPERATION

include/swift/Runtime/RuntimeFunctions.def

@@ -1518,6 +1518,30 @@ FUNCTION(TaskCreateFutureFunc,
               TaskContinuationFunctionPtrTy, SizeTy),
          ATTRS(NoUnwind, ArgMemOnly))
 
+// AutoDiffLinearMapContext *swift_autoDiffCreateLinearMapContext(size_t);
+FUNCTION(AutoDiffCreateLinearMapContext,
+         swift_autoDiffCreateLinearMapContext, SwiftCC,
+         DifferentiationAvailability,
+         RETURNS(RefCountedPtrTy),
+         ARGS(SizeTy),
+         ATTRS(NoUnwind, ArgMemOnly))
+
+// void *swift_autoDiffProjectTopLevelSubcontext(AutoDiffLinearMapContext *);
+FUNCTION(AutoDiffProjectTopLevelSubcontext,
+         swift_autoDiffProjectTopLevelSubcontext, SwiftCC,
+         DifferentiationAvailability,
+         RETURNS(Int8PtrTy),
+         ARGS(RefCountedPtrTy),
+         ATTRS(NoUnwind, ArgMemOnly))
+
+// void *swift_autoDiffAllocateSubcontext(AutoDiffLinearMapContext *, size_t);
+FUNCTION(AutoDiffAllocateSubcontext,
+         swift_autoDiffAllocateSubcontext, SwiftCC,
+         DifferentiationAvailability,
+         RETURNS(Int8PtrTy),
+         ARGS(RefCountedPtrTy, SizeTy),
+         ATTRS(NoUnwind, ArgMemOnly))
+
 #undef RETURNS
 #undef ARGS
 #undef ATTRS

include/swift/SILOptimizer/Differentiation/Common.h

@@ -192,6 +192,16 @@ void extractAllElements(SILValue value, SILBuilder &builder,
 void emitZeroIntoBuffer(SILBuilder &builder, CanType type,
                         SILValue bufferAccess, SILLocation loc);
 
+/// Emit a `Builtin.Word` value that represents the given type's memory layout
+/// size.
+SILValue emitMemoryLayoutSize(
+    SILBuilder &builder, SILLocation loc, CanType type);
+
+/// Emit a projection of the top-level subcontext from the context object.
+SILValue emitProjectTopLevelSubcontext(
+    SILBuilder &builder, SILLocation loc, SILValue context,
+    SILType subcontextType);
+
 //===----------------------------------------------------------------------===//
 // Utilities for looking up derivatives of functions
 //===----------------------------------------------------------------------===//

include/swift/SILOptimizer/Differentiation/LinearMapInfo.h

@@ -63,6 +63,9 @@ class LinearMapInfo {
   /// Activity info of the original function.
   const DifferentiableActivityInfo &activityInfo;
 
+  /// The original function's loop info.
+  SILLoopInfo *loopInfo;
+
   /// Differentiation indices of the function.
   const SILAutoDiffIndices indices;
 
@@ -86,6 +89,9 @@ class LinearMapInfo {
   /// Mapping from linear map structs to their branching trace enum fields.
   llvm::DenseMap<StructDecl *, VarDecl *> linearMapStructEnumFields;
 
+  /// Blocks in a loop.
+  llvm::SmallSetVector<SILBasicBlock *, 4> blocksInLoop;
+
   /// A synthesized file unit.
   SynthesizedFileUnit &synthesizedFile;
 
@@ -144,7 +150,8 @@ class LinearMapInfo {
   explicit LinearMapInfo(ADContext &context, AutoDiffLinearMapKind kind,
                          SILFunction *original, SILFunction *derivative,
                          SILAutoDiffIndices indices,
-                         const DifferentiableActivityInfo &activityInfo);
+                         const DifferentiableActivityInfo &activityInfo,
+                         SILLoopInfo *loopInfo);
 
   /// Returns the linear map struct associated with the given original block.
   StructDecl *getLinearMapStruct(SILBasicBlock *origBB) const {
@@ -200,20 +207,28 @@ class LinearMapInfo {
 
   /// Returns the branching trace enum field for the linear map struct of the
   /// given original block.
-  VarDecl *lookUpLinearMapStructEnumField(SILBasicBlock *origBB) {
+  VarDecl *lookUpLinearMapStructEnumField(SILBasicBlock *origBB) const {
     auto *linearMapStruct = getLinearMapStruct(origBB);
     return linearMapStructEnumFields.lookup(linearMapStruct);
   }
 
   /// Finds the linear map declaration in the pullback struct for the given
   /// `apply` instruction in the original function.
-  VarDecl *lookUpLinearMapDecl(ApplyInst *ai) {
+  VarDecl *lookUpLinearMapDecl(ApplyInst *ai) const {
     assert(ai->getFunction() == original);
     auto lookup = linearMapFieldMap.find(ai);
     assert(lookup != linearMapFieldMap.end() &&
            "No linear map field corresponding to the given `apply`");
     return lookup->getSecond();
   }
+
+  bool hasLoops() const {
+    return !blocksInLoop.empty();
+  }
+
+  ArrayRef<SILBasicBlock *> getBlocksInLoop() const {
+    return blocksInLoop.getArrayRef();
+  }
 };
 
 } // end namespace autodiff

include/swift/SILOptimizer/Differentiation/VJPCloner.h

@@ -21,6 +21,7 @@
 #include "swift/SILOptimizer/Analysis/DifferentiableActivityAnalysis.h"
 #include "swift/SILOptimizer/Differentiation/DifferentiationInvoker.h"
 #include "swift/SILOptimizer/Differentiation/LinearMapInfo.h"
+#include "swift/SIL/LoopInfo.h"
 
 namespace swift {
 namespace autodiff {
@@ -52,6 +53,7 @@ class VJPCloner final {
   const SILAutoDiffIndices getIndices() const;
   DifferentiationInvoker getInvoker() const;
   LinearMapInfo &getPullbackInfo() const;
+  SILLoopInfo *getLoopInfo() const;
   const DifferentiableActivityInfo &getActivityInfo() const;
 
   /// Performs VJP generation on the empty VJP function. Returns true if any

lib/AST/ASTMangler.cpp

@@ -785,8 +785,8 @@ static StringRef getPrivateDiscriminatorIfNecessary(const ValueDecl *decl) {
 
   // Mangle non-local private declarations with a textual discriminator
   // based on their enclosing file.
-  auto topLevelContext = decl->getDeclContext()->getModuleScopeContext();
-  auto fileUnit = cast<FileUnit>(topLevelContext);
+  auto topLevelSubcontext = decl->getDeclContext()->getModuleScopeContext();
+  auto fileUnit = cast<FileUnit>(topLevelSubcontext);
 
   Identifier discriminator =
       fileUnit->getDiscriminatorForPrivateValue(decl);
@@ -2900,17 +2900,17 @@ void ASTMangler::appendEntity(const ValueDecl *decl) {
 void
 ASTMangler::appendProtocolConformance(const ProtocolConformance *conformance) {
   GenericSignature contextSig;
-  auto topLevelContext =
+  auto topLevelSubcontext =
       conformance->getDeclContext()->getModuleScopeContext();
-  Mod = topLevelContext->getParentModule();
+  Mod = topLevelSubcontext->getParentModule();
 
   auto conformingType = conformance->getType();
   appendType(conformingType->getCanonicalType());
 
   appendProtocolName(conformance->getProtocol());
 
   bool needsModule = true;
-  if (auto *file = dyn_cast<FileUnit>(topLevelContext)) {
+  if (auto *file = dyn_cast<FileUnit>(topLevelSubcontext)) {
     if (file->getKind() == FileUnitKind::ClangModule ||
         file->getKind() == FileUnitKind::DWARFModule) {
       if (conformance->getProtocol()->hasClangNode())

lib/AST/ASTVerifier.cpp

@@ -229,7 +229,7 @@ class Verifier : public ASTWalker {
   typedef llvm::PointerIntPair<DeclContext *, 1, bool> ClosureDiscriminatorKey;
   llvm::DenseMap<ClosureDiscriminatorKey, SmallBitVector>
       ClosureDiscriminators;
-  DeclContext *CanonicalTopLevelContext = nullptr;
+  DeclContext *CanonicalTopLevelSubcontext = nullptr;
 
   Verifier(PointerUnion<ModuleDecl *, SourceFile *> M, DeclContext *DC)
       : M(M),
@@ -898,9 +898,9 @@ class Verifier : public ASTWalker {
     DeclContext *getCanonicalDeclContext(DeclContext *DC) {
       // All we really need to do is use a single TopLevelCodeDecl.
       if (auto topLevel = dyn_cast<TopLevelCodeDecl>(DC)) {
-        if (!CanonicalTopLevelContext)
-          CanonicalTopLevelContext = topLevel;
-        return CanonicalTopLevelContext;
+        if (!CanonicalTopLevelSubcontext)
+          CanonicalTopLevelSubcontext = topLevel;
+        return CanonicalTopLevelSubcontext;
       }
 
       // TODO: check for uniqueness of initializer contexts?

lib/AST/Availability.cpp

@@ -327,6 +327,10 @@ AvailabilityContext ASTContext::getConcurrencyAvailability() {
   return getSwiftFutureAvailability();
 }
 
+AvailabilityContext ASTContext::getDifferentiationAvailability() {
+  return getSwiftFutureAvailability();
+}
+
 AvailabilityContext ASTContext::getSwift52Availability() {
   auto target = LangOpts.Target;
 

lib/AST/Builtins.cpp

@@ -1383,6 +1383,25 @@ static ValueDecl *getCreateAsyncTaskFuture(ASTContext &ctx, Identifier id) {
   return builder.build(id);
 }
 
+static ValueDecl *getAutoDiffCreateLinearMapContext(ASTContext &ctx,
+                                                    Identifier id) {
+  return getBuiltinFunction(
+      id, {BuiltinIntegerType::getWordType(ctx)}, ctx.TheNativeObjectType);
+}
+
+static ValueDecl *getAutoDiffProjectTopLevelSubcontext(ASTContext &ctx,
+                                                       Identifier id) {
+  return getBuiltinFunction(
+      id, {ctx.TheNativeObjectType}, ctx.TheRawPointerType);
+}
+
+static ValueDecl *getAutoDiffAllocateSubcontext(ASTContext &ctx,
+                                                Identifier id) {
+  return getBuiltinFunction(
+      id, {ctx.TheNativeObjectType, BuiltinIntegerType::getWordType(ctx)},
+      ctx.TheRawPointerType);
+}
+
 static ValueDecl *getPoundAssert(ASTContext &Context, Identifier Id) {
   auto int1Type = BuiltinIntegerType::get(1, Context);
   auto optionalRawPointerType = BoundGenericEnumType::get(
@@ -2549,6 +2568,15 @@ ValueDecl *swift::getBuiltinValueDecl(ASTContext &Context, Identifier Id) {
 
   case BuiltinValueKind::TriggerFallbackDiagnostic:
     return getTriggerFallbackDiagnosticOperation(Context, Id);
+
+  case BuiltinValueKind::AutoDiffCreateLinearMapContext:
+    return getAutoDiffCreateLinearMapContext(Context, Id);
+
+  case BuiltinValueKind::AutoDiffProjectTopLevelSubcontext:
+    return getAutoDiffProjectTopLevelSubcontext(Context, Id);
+
+  case BuiltinValueKind::AutoDiffAllocateSubcontext:
+    return getAutoDiffAllocateSubcontext(Context, Id);
   }
 
   llvm_unreachable("bad builtin value!");