Commit Graph

80 Commits

Author SHA1 Message Date
swift-ci
d33c79edaf Merge pull request #35847 from rxwei/72666310-mangle-ad-thunks 2021-02-13 05:07:32 -08:00
Richard Wei
f9ddecf459 [AutoDiff] Mangle linear map self-reordering thunks and subset parameters thunks.
Add the following new mangling rules.
```
global ::= from-type to-type 'TJO' AUTODIFF-FUNCTION-KIND // autodiff self-reordering reabstraction thunk
global ::= from-type 'TJS' AUTODIFF-FUNCTION-KIND INDEX-SUBSET 'p' INDEX-SUBSET 'r' INDEX-SUBSET 'P' // autodiff linear map subset parameters thunk
global ::= global to-type 'TJS' AUTODIFF-FUNCTION-KIND INDEX-SUBSET 'p' INDEX-SUBSET 'r' INDEX-SUBSET 'P' // autodiff derivative function subset parameters thunk
```

Example:
```console
$s13TangentVector16_Differentiation14DifferentiablePQzAaDQy_SdAFIegnnnr_TJSdSSSpSrSUSP ---> autodiff subset parameters thunk for differential from @escaping @callee_guaranteed (@in_guaranteed A._Differentiation.Differentiable.TangentVector, @in_guaranteed B._Differentiation.Differentiable.TangentVector, @in_guaranteed Swift.Double) -> (@out B._Differentiation.Differentiable.TangentVector) with respect to parameters {0, 1, 2} and results {0} to parameters {0, 2}
$sS2f8mangling3FooV13TangentVectorVIegydd_SfAESfIegydd_TJOp ---> autodiff self-reordering reabstraction thunk for pullback from @escaping @callee_guaranteed (@unowned Swift.Float) -> (@unowned Swift.Float, @unowned mangling.Foo.TangentVector) to @escaping @callee_guaranteed (@unowned Swift.Float) -> (@unowned mangling.Foo.TangentVector, @unowned Swift.Float)
```

Resolves rdar://72666310 / SR-13508.

Also fix a bug in `AutoDiffFunction` mangling where the original may be a global that contains more than 1 node (rdar://74151229 / SR-14106).
2021-02-13 01:38:35 -08:00
Erik Eckstein
542a378436 SIL: add FunctionRefInst::getReferencedFunction()
If we know that we have a FunctionRefInst (and not another variant of FunctionRefBaseInst), we know that getting the referenced function will not be null (in contrast to FunctionRefBaseInst::getReferencedFunctionOrNull).

NFC
2021-02-09 19:56:43 +01:00
Richard Wei
af8942d940 [AutoDiff] Rename '@differentiable' to '@differentiable(reverse)'.
Compiler:
- Add `Forward` and `Reverse` to `DifferentiabilityKind`.
- Expand `DifferentiabilityMask` in `ExtInfo` to 3 bits so that it now holds all 4 cases of `DifferentiabilityKind`.
- Parse `@differentiable(reverse)` and `@differentiable(_forward)` declaration attributes and type attributes.
- Emit a warning for `@differentiable` without `reverse`.
- Emit an error for `@differentiable(_forward)`.
- Rename `@differentiable(linear)` to `@differentiable(_linear)`.
- Make `@differentiable(reverse)` type lowering go through today's `@differentiable` code path. We will specialize it to reverse-mode in a follow-up patch.

ABI:
- Add `Forward` and `Reverse` to `FunctionMetadataDifferentiabilityKind`.
- Extend `TargetFunctionTypeFlags` by 1 bit to store the highest bit of differentiability kind (linear). Note that there is a 2-bit gap in `DifferentiabilityMask` which is reserved for `AsyncMask` and `ConcurrentMask`; `AsyncMask` is ABI-stable so we cannot change that.

_Differentiation module:
- Replace all occurrences of `@differentiable` with `@differentiable(reverse)`.
- Delete `_transpose(of:)`.

Resolves rdar://69980056.
2021-02-07 14:09:46 -08:00
Erik Eckstein
011358edd6 SIL: let SingleValueInstruction only inherit from a single SILNode.
This removes the ambiguity when casting from a SingleValueInstruction to SILNode, which makes the code simpler. E.g. the "isRepresentativeSILNode" logic is not needed anymore.
Also, it reduces the size of the most used instruction class - SingleValueInstruction - by one pointer.

Conceptually, SILInstruction is still a SILNode. But implementation-wise SILNode is not a base class of SILInstruction anymore.
Only the two sub-classes of SILInstruction - SingleValueInstruction and NonSingleValueInstruction - inherit from SILNode. SingleValueInstruction's SILNode is embedded into a ValueBase and its relative offset in the class is the same as in NonSingleValueInstruction (see SILNodeOffsetChecker).
This makes it possible to cast from a SILInstruction to a SILNode without knowing which SILInstruction sub-class it is.
Casting to SILNode cannot be done implicitly, but only with an LLVM `cast` or with SILInstruction::asSILNode(). But this is a rare case anyway.
2021-01-27 16:40:15 +01:00
Eric Miotto
8e7f9c9cbd Revert "SIL: let SingleValueInstruction only inherit from a single SILNode." 2021-01-26 10:02:24 -08:00
Erik Eckstein
ff1991740a SIL: let SingleValueInstruction only inherit from a single SILNode.
This removes the ambiguity when casting from a SingleValueInstruction to SILNode, which makes the code simpler. E.g. the "isRepresentativeSILNode" logic is not needed anymore.
Also, it reduces the size of the most used instruction class - SingleValueInstruction - by one pointer.

Conceptually, SILInstruction is still a SILNode. But implementation-wise SILNode is not a base class of SILInstruction anymore.
Only the two sub-classes of SILInstruction - SingleValueInstruction and NonSingleValueInstruction - inherit from SILNode. SingleValueInstruction's SILNode is embedded into a ValueBase and its relative offset in the class is the same as in NonSingleValueInstruction (see SILNodeOffsetChecker).
This makes it possible to cast from a SILInstruction to a SILNode without knowing which SILInstruction sub-class it is.
Casting to SILNode cannot be done implicitly, but only with an LLVM `cast` or with SILInstruction::asSILNode(). But this is a rare case anyway.
2021-01-25 09:30:04 +01:00
Erik Eckstein
b7351780f7 SIL: move all the block-list modifying APIs to SILFunction.
... and remove SILFunction::getBlocks().

It's just a cleanup, NFC.
2021-01-14 17:35:31 +01:00
Richard Wei
ffe6064101 Mangle derivative functions and linear maps.
- `Mangle::ASTMangler::mangleAutoDiffDerivativeFunction()` and `Mangle::ASTMangler::mangleAutoDiffLinearMap()` accept original function declarations and return a mangled name for a derivative function or linear map. This is called during SILGen and TBDGen.
- `Mangle::DifferentiationMangler` handles differentiation function mangling in the differentiation transform. This part is necessary because we need to perform demangling on the original function and remangle it as part of a differentiation function mangling tree in order to get the correct substitutions in the mangled derivative generic signature.

A mangled differentiation function name includes:
- The original function.
- The differentiation function kind.
- The parameter indices for differentiation.
- The result indices for differentiation.
- The derivative generic signature.
2021-01-07 02:21:10 -08:00
Richard Wei
8d8614058b [AudoDiff] NFC: Replace 'SILAutoDiffIndices' with 'AutoDiffConfig'. (#35079)
Resolve rdar://71678394 / SR-13889.
2020-12-14 14:32:40 -08:00
Michael Gottesman
c026e95cce [ownership] Extract out SILOwnershipKind from ValueOwnershipKind into its own type and rename Invalid -> Any.
This makes it easier to understand conceptually why a ValueOwnershipKind with
Any ownership is invalid and also allowed me to explicitly document the lattice
that relates ownership constraints/value ownership kinds.
2020-11-10 14:29:11 -08:00
Michael Gottesman
8c5737d1d5 [ownership] Change thin_to_thick function to always produce a none value.
It can already only accept values with none ownership and the merging of
ownership around ownership phis ensure that if we phi this with a partial_apply
or the like, we get the appropriate ownership on any such ownership phi values.

We are now out of SILGen emitting fewer destroy_value unnecessarily on
thin_to_thick functions. This changed some codegen and also forced me to update
some tests/fix AutoDiff.

I also deleted the DebugInfo test mandatoryinlining-wrongdebugscope.swift since:

1. It was depending on these destroys being there.

2. Given the need to improve the test @aprantl suggested I just eliminate it
solving the test failure for me.
2020-10-23 17:49:12 -07:00
Nate Chandler
94b5f76654 Revert "[SIL] Add SILFunctionType flag for async."
This reverts commit 9b8828848d.
2020-08-25 13:37:26 -07:00
Nate Chandler
9b8828848d [SIL] Add SILFunctionType flag for async. 2020-08-19 11:29:58 -07:00
Alex Efremov
40c448514a [AutoDiff] Support differentiation of non-active try_apply. (#33483)
Add differentiation support for non-active `try_apply` SIL instructions.

Notable pullback generation changes:
* Original basic blocks are now visited in a different order:
* starting from the original basic block, all its predecessors
* are visited in a breadth-first search order. This ensures that
* all successors of any block are visited before the block itself.

Resolves TF-433.
2020-08-15 10:22:28 -07:00
Dan Zheng
576d0b8d83 [AutoDiff] NFC: fix typos. (#33461) 2020-08-13 15:27:09 -07:00
Varun Gandhi
f219e58ada [NFC] Refactor ExtInfo to use a builder-pattern based API.
Since the two ExtInfos share a common ClangTypeInfo, and C++ doesn't let us
forward declare nested classes, we need to hoist out AnyFunctionType::ExtInfo
and SILFunctionType::ExtInfo to the top-level.

We also add some convenience APIs on (AST|SIL)ExtInfo for frequently used
withXYZ methods. Note that all non-default construction still goes through the
builder's build() method.

We do not add any checks for invariants here; those will be added later.
2020-07-31 13:55:55 -07:00
Alex Efremov
1a81573eed [AutoDiff] Start linear_function canonicalization skeleton (#33057)
Start `linear_function` canonicalization skeleton copying from
`differentiable_function` canonicalization. For now, transpose function
operands are filled in with `undef`.
2020-07-23 14:10:50 -07:00
Dan Zheng
c28d9cd7ea [AutoDiff] NFC: fix typos. (#33002) 2020-07-20 11:41:39 -07:00
Dan Zheng
ec5da7427b [AutoDiff] NFC: garden differentiation transform. (#32770)
Rename "emitters" to "cloners", for consistency:
- `JVPEmitter` -> `JVPCloner`
- `VJPEmitter` -> `VJPCloner`
- `PullbackEmitter` -> `PullbackCloner`

Improve `PullbackCloner` documentation.
- Document previously undocumented methods.
- Update outdated documentation.
  - For adjoint value accumulation helpers: rename "buffer access" occurrences to
    "address". Pullback generation no logner uses buffer accesses (`begin_apply`).
2020-07-08 19:26:36 -07:00
Dan Zheng
46c2073784 [AutoDiff] Start supporting loadable types with address-only tangents. (#32540)
Previously, PullbackEmitter assumed that original values' value category
matches their `TangentVector` types' value category. This was problematic
for loadable types with address-only `TangentVector` types.

Now, PullbackEmitter starts to support differentiation of loadable types with
address-only `TangentVector` types. This patch focuses on supporting and testing
class types, more support can be added incrementally.

Resolves TF-1149.
2020-06-25 08:27:03 -07:00
Dan Zheng
d3b6b89de6 [AutoDiff] Support multiple differentiability result indices in SIL. (#32206)
`DifferentiableFunctionInst` now stores result indices.
`SILAutoDiffIndices` now stores result indices instead of a source index.

`@differentiable` SIL function types may now have multiple differentiability
result indices and `@noDerivative` resutls.

`@differentiable` AST function types do not have `@noDerivative` results (yet),
so this functionality is not exposed to users.

Resolves TF-689 and TF-1256.

Infrastructural support for TF-983: supporting differentiation of `apply`
instructions with multiple active semantic results.
2020-06-05 16:25:17 -07:00
Varun Gandhi
c14e934563 [NFC] Remove redundant includes for llvm/ADT/SmallSet.h. 2020-05-31 13:07:45 -07:00
Dan Zheng
d5d076db6a [AutoDiff] Support differentiation of branching cast instructions.
Support differentiation of `is` and `as?` operators.

These operators lower to branching cast SIL instructions, requiring control
flow differentiation support.

Resolves SR-12898.
2020-05-28 12:54:38 -07:00
Dan Zheng
a5bdbd9774 [AutoDiff] Fix derivative function configurations for accessors. (#31669)
For accessors: make `AbstractFunctionDecl::getDerivativeFunctionConfigurations`
resolve configurations from parent storage declaration `@differentiable`
attributes.

Fixes "no `@differentiable` attribute" non-differentiability error for accessors
whose parent storage declaration `@differentiable` attributes have not been
type-checked (e.g. because the storage declarations are in another file).

Add protocol requirement and class member storage declaration tests.

Resolves TF-1234.
2020-05-09 07:29:15 -07:00
Dan Zheng
43715e725e [AutoDiff] NFC: reorganize differentiation SILOptimizer files. (#30969)
Move differentiation-related SILOptimizer files to
{include/swift,lib}/SILOptimizer/Differentiation/.

This reduces directory nesting and gathers files together.
2020-05-04 06:19:39 -07:00
Dan Zheng
2eb460de4d [AutoDiff upstream] Add forward-mode differentiation. (#30878)
JVP functions are forward-mode derivative functions. They take original
arguments and return original results and a differential function. Differential
functions take derivatives wrt arguments and return derivatives wrt results.

`JVPEmitter` is a cloner that emits JVP and differential functions at the same
time. In JVP functions, function applications are replaced with JVP function
applications. In differential functions, function applications are replaced
with differential function applications.

In JVP functions, each basic block takes a differential struct containing callee
differentials. These structs are consumed by differential functions.
2020-04-08 11:29:21 -07:00
Dan Zheng
1775e8ae16 [AutoDiff upstream] Add VJPEmitter.
`VJPEmitter` is a cloner that emits VJP functions. It implements reverse-mode
automatic differentiation, along with `PullbackEmitter`.

`VJPEmitter` clones an original function, replacing function applications with
VJP function applications. In VJP functions, each basic block takes a pullback
struct (containing callee pullbacks) and produces a predecessor enum: these data
structures are consumed by pullback functions.
2020-04-05 20:35:35 -07:00
Dan Zheng
146c11ec80 [AutoDiff upstream] Add differentiable_function canonicalization. (#30818)
Canonicalizes `differentiable_function` instructions by filling in missing
derivative function operands.

Derivative function emission rules, based on the original function value:

- `function_ref`: look up differentiability witness with the exact or a minimal
  superset derivative configuration. Emit a `differentiability_witness_function`
  for the derivative function.
- `witness_method`: emit a `witness_method` with the minimal superset derivative
  configuration for the derivative function.
- `class_method`: emit a `class_method` with the minimal superset derivative
  configuration for the derivative function.

If an *actual* emitted derivative function has a superset derivative
configuration versus the *desired* derivative configuration, create a "subset
parameters thunk" to thunk the actual derivative to the desired type.

For `differentiable_function` instructions formed from curry thunk applications:
clone the curry thunk (with type `(Self) -> (T, ...) -> U`) and create a new
version with type `(Self) -> @differentiable (T, ...) -> U`.

Progress towards TF-1211.
2020-04-05 20:19:10 -07:00
Dan Zheng
aa66cce808 [AutoDiff upstream] Add differentiation transform.
The differentiation transform does the following:
- Canonicalizes differentiability witnesses by filling in missing derivative
  function entries.
- Canonicalizes `differentiable_function` instructions by filling in missing
  derivative function operands.
- If necessary, performs automatic differentiation: generating derivative
  functions for original functions.
  - When encountering non-differentiability code, produces a diagnostic and
    errors out.

Partially resolves TF-1211: add the main canonicalization loop.

To incrementally stage changes, derivative functions are currently created
with empty bodies that fatal error with a nice message.

Derivative emitters will be upstreamed separately.
2020-04-02 15:43:57 -07:00