[AutoDiff] Bump-pointer allocate pullback structs in loops.

[rxwei]

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.

The following entry points are added to the runtime:

/// Creates a linear map context with a tail-allocated top-level subcontext.
SWIFT_EXPORT_FROM(swift_Differentiation) SWIFT_CC(swift)
AutoDiffLinearMapContext *swift_autoDiffCreateLinearMapContext(
    size_t topLevelSubcontextSize);

/// Returns the address of the tail-allocated top-level subcontext.
SWIFT_EXPORT_FROM(swift_Differentiation) SWIFT_CC(swift)
void *swift_autoDiffProjectTopLevelSubcontext(AutoDiffLinearMapContext *);

/// Allocates memory for a new subcontext.
SWIFT_EXPORT_FROM(swift_Differentiation) SWIFT_CC(swift)
void *swift_autoDiffAllocateSubcontext(AutoDiffLinearMapContext *, size_t size);

This is paving the road for a series of optimizations on linear map closure context allocations. For example, a pass can be run on all user-registered derivatives to allocate their closure contexts as a subcontext in a swift::AutoDiffLinearMapContext.

As a result, differentiating loops over 1 million iterations no longer segfaults, and derivatives with loops have a consistent small performance increase. More work to be done later:

• Replace llvm::BumpPtrAllcoator with something like the task allocator in [concurrency] Implement the Task allocator as bump-pointer allocator. #34880 so that we can deallocate things in a stack discipline and tail-allocate the initial slab.
• Apply bump-pointer allocation to all linear map contexts.

[rxwei]

@swift-ci please test

[swift-ci]

Build failed
Swift Test OS X Platform
Git Sha - 4a19baaa4e66b1f34275d16f11acb225658cec4d

[dan-zheng]

What a nice holiday gift! Incredible.

[dan-zheng]

Can we start autodiff benchmarks to evaluate performance-impacting compiler changes?

Plugging into Swift compiler benchmark suite

@marcrasi added autodiff tests to the benchmark suite, I think reviving and landing that is a good start: #31108.

Dedicated Swift differentiation benchmark library

Personally, I like google/swift-benchmark as a benchmarking library. I found it gives more information than XCTest utilities (XCTestCase.measure) and the output is quite human-interpretable.

I used google/swift-benchmark for various autodiff benchmarking experiments but didn't know a good home for the code (an AutoDiffBenchmarks SwiftPM package). Maybe it could be a GitHub repository called swift-autodiff-benchmarking.

It's a bit harder to test compiler changes with a SwiftPM benchmark suite, but it should be possible with a bit of work.

[rxwei]

I'm interested in using the Swift compiler benchmark suite. Thanks for pointing to #31108 — I'll revive this.

From my local benchmarks, I've seen a consistent 2x-10x speedup on loops, but it's small compared to the rest of the issues to be fixed later. A more significant outcome is that loops over 1 million iterations no longer segfault.

[rxwei]

@swift-ci please test