Daniil Kovalev 1f77138afe [AutoDiff] Closure specialization: specialize branch tracing enums (#85757) ...

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.

Note that mangling-related logic is implemented in C++ since at this
moment we have no Swift bridged for that.

Here is a simplified example of how branch tracing enum (BTE)
specialization looks like.

Before specialization:

```
enum $_AD__xxx {
  case bb0(((Float) -> Float))
}

func vjp(...) {
  // ...
  %foo      = function_ref $foo         : (Float, Float) -> Float
  %pa1      = partial_apply %foo(%arg1) : (Float) -> Float
  %payload1 = tuple (%pa1)              : ((Float) -> Float)
  %bte      = enum $_AD__xxx.bb0!enumelt, %payload1
  // ...
}

func pullback(%bte, ...) {
  // ...
  %payload2 = unchecked_enum_data %bte, $_AD__xxx.bb0!enumelt : ((Float) -> Float)
  %pa2      = tuple_extract %payload2, 0                      : (Float) -> Float
  %res      = apply %pa2(%arg2)                               : Float
  // ...
}
```

After specialization:

```
enum $_AD__xxx_spec_bb0_0 {
  case bb0(((Float)))
}

func vjp(...) {
  // ...
  %captured1 = tuple (%arg1)      : (Float)
  %payload1  = tuple (%captured1) : ((Float))
  %bte_spec  = enum $_AD__xxx_spec_bb0_0.bb0!enumelt, %payload1
  // ...
}

func pullback_spec(%bte_spec, ...) {
  // ...
  %payload2  = unchecked_enum_data %bte, $_AD__xxx_spec_bb0_0.bb0!enumelt : ((Float))
  %captured2 = tuple_extract %payload2, 0                                 : (Float)
  %arg1      = tuple_extract %captured2, 0                                : Float
  %foo       = function_ref $foo                                          : (Float, Float) -> Float
  %res       = apply %foo(%arg2, %arg1)                                   : Float
  // ...
}
```

2025-12-21 00:33:50 +00:00

52 KiB

Raw Permalink Blame History

View Raw