Daniil Kovalev a0907a0f78 [SILOptimizer] Enable iterative specialization of nested closure chains (#88920) ...

Enhance both ClosureSpecialization and AutoDiffClosureSpecialization
passes to specialize chains of `partial_apply`s that capture each other,
instead of stopping after the outermost closure.

Whether a closure gets specialized depends on it being fully applied in
the callee. In a chain of nested closures, only the least-nested one is
fully applied at first; specializing it exposes the next. This pattern
is common in AutoDiff pullbacks (see `pb_bar` -> `pb_foo` -> `pb_sin` in
test/AutoDiff/SILOptimizer/closure_specialization/single_bb_nested.sil)
and in plain nested capturing closures (see `wrapper2` -> `wrapper1` ->
`wrapper0` in test/SILOptimizer/closure_specialization_nested.sil).

This patch introduces the following logic.

1. Fold `partial_apply`+`apply` after `partial_apply` movement to the
specialized callee. Previously, it was expected that SILCombiner does
this job. Do the fold early to allow further specialization rounds since
specialization is only happening when full apply in callee is detected.

2. Run specialization logic recursively for the callee just after it is
specialized. This handles nested closures chains.

3. Run up to 5 rounds of ClosureSpecialization pass while changes keep
occurring (mirroring the AutoDiffClosureSpecialization pass)

**Note 1.** One CHECK in `closure_specialize_loop.swift` is now a FIXME.
The enhanced ClosureSpecialization pass handles that test case without
the ClosureSpecialization/CapturePropagation ping-pong it was guarding.

**Note 2.** This patch is a prerequisite for a subsequent patch allowing
for more AutoDiff VJP inlining decisions.

2026-06-14 07:43:13 +03:00

3.9 KiB

Raw Permalink Blame History

View Raw