Closure specialization optimization specific to Swift Autodiff

Development Compiler
, ,
1

Hi,
I am working on improving the performance of Swift Autodiff (AD). Through some profiling we determined that the derivatives we generate are spending quite a bit of time allocating memory, so I'm working on optimizations that could help us reduce memory allocations.

The biggest chunk of our memory allocations are coming from the usage of closures and we have identified that a closure specialization (closure-spec) optimization could help us out here. We will however need a closure-spec optimization pass specific to Swift AD due to our specific code patterns which the general closure-spec optimization either does not or is not able to handle.

I wanted to understand the reasoning behind some of the limitations of the general closure-spec optimization.

  1. Callsites with multiple closure arguments are not handled [link].

  2. Partial applies of generic functions are not supported [link].

  3. Closures with non-object SILType arguments that are passed by address and
    are not @inout are not supported [link].

AFAIK, (1) can happen frequently in Swift AD. (2) and (3) may or may not be possible, but I still need to get a definitive answer. Mainly the motivation behind my questions is to understand whether or not I should be handling these optimizations in Swift AD specific closure-spec optimization.

3 Likes
2

@Erik_Eckstein would you happen to know the answer to these? Or know someone who might know the answer to these?

3

Correction:
Only case (1) can happen in AD. (2) and (3) do not happen in AD.

4

The closure specialization pass has not been touched in a long time just because no one has had the express need to do so.

Some thoughts:

  1. I think this is because we were trying to be conservative and were targeting things like map (which take a single closure).

  2. The generic function thing was just a time limit thing IIRC.

  3. Even in the link it says this is temporary.

I would say it is better to change closure specialization than make a specific pass.

2 Likes
5

Thank you for the information Michael!

I would say it is better to change closure specialization than make a specific pass.

That's good to hear and I'll surely keep this as a preference. I'm drawing out the design/specifics of the things we'd need. I can ping you once that's ready, to take a quick look and point out early if something might just not fit in the existing optimization pass.

1 Like
6

@Michael_Gottesman another real quick request.

Could you very briefly elaborate what makes it harder to handle cases (2) and (3), such that they weren't handled in the compiler at the time the optimization was written?

1 Like
7

It was just very early in the time period of developing the compiler and needed to do other things. I don't think there was anything specific about it.

1 Like