[AutoDiff upstream] Add `differentiable_function` canonicalization. #30818

dan-zheng · 2020-04-05T23:38:11Z

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.

@differentiable

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.

dan-zheng · 2020-04-05T23:39:56Z

lib/SILOptimizer/Mandatory/Differentiation.cpp

+      std::tie(thunk, interfaceSubs) =
+          getOrCreateSubsetParametersThunkForDerivativeFunction(
+              fb, origFnOperand, derivativeFn, derivativeFnKind, desiredIndices,
+              actualIndices);


Note: subset parameters thunk tests are blocked by derivative function emitters, which will be upstreamed next.

dan-zheng · 2020-04-05T23:41:16Z

lib/SILOptimizer/Mandatory/Differentiation.cpp

+/// Check whether the given requirements are satisfied, with the given
+/// derivative generic signature (containing requirements), and substitution
+/// map. Returns true if error is emitted.
+static bool diagnoseUnsatisfiedRequirements(ADContext &context,


Todo: upstream differentiation_transform_diagnostics.swift, which tests diagnoseUnsatisfiedRequirements.

dan-zheng · 2020-04-05T23:42:52Z

@gottesmm: would you like to help review this patch, as a SILOptimizer code owner?

dan-zheng · 2020-04-05T23:43:01Z

@swift-ci Please test

swift-ci · 2020-04-06T01:12:11Z

Build failed
Swift Test Linux Platform
Git Sha - 5430142

dan-zheng · 2020-04-06T01:26:33Z

@swift-ci Please test Linux

dan-zheng · 2020-04-06T03:19:04Z

Merging to unblock progress. Happy to address any feedback later!

gottesmm

@dan-zheng some comments. Please fix:

Also, are these used anywhere but the differentiation pass? If not, we should just put it in the differentiation pass until we find multiple uses. A bunch of the helpers that you added need to be refactored. Please fix them. I may look at this again in a bit.

gottesmm · 2020-04-10T22:59:51Z

include/swift/SILOptimizer/Utils/Differentiation/Common.h

+/// Emit a zero value into the given buffer access by calling
+/// `AdditiveArithmetic.zero`. The given type must conform to
+/// `AdditiveArithmetic`.
+void emitZeroIntoBuffer(SILBuilder &builder, CanType type,


Please fix these APIs to not take a SILBuilder. Instead, it should take an insertion point and a SILBuilderContext.

This change seems important for correctness. It can probably help clean up this code:
https://github.com/apple/swift/blob/13d5a8addbe3605984edc4ce7c6cbf1f0649a9e1/lib/SILOptimizer/Utils/Differentiation/PullbackEmitter.cpp#L486-L497

I started looking into this but ran into some difficulties. The user above wants to insert at the end of a basic block (SILBasicBlock::iterator), but no appropriate constructor exists. Only SILBuilderWithScope(SILInstruction *I, SILBuilderContext &C) is recommended.

SILBuilderWithScope(SILBasicBlock::iterator I, SILBuilderContext &C) would work but doesn't exist. Can we add it, or are there other considerations?

gottesmm · 2020-04-10T23:02:02Z

include/swift/SILOptimizer/Utils/Differentiation/Common.h

@@ -33,6 +34,12 @@ namespace autodiff {
 /// This is being used to print short debug messages within the AD pass.
 raw_ostream &getADDebugStream();

+/// Given a function call site, gathers all of its actual results (both direct
+/// and indirect) in an order defined by its result type.
+void collectAllActualResultsInTypeOrder(


I have a number of problems with this. I am pretty sure we already have something like this.

gottesmm · 2020-04-10T23:04:28Z

lib/SILOptimizer/Mandatory/Differentiation.cpp

+  auto diffResultFnTy = resultFnTy->getWithExtInfo(
+      resultFnTy->getExtInfo().withDifferentiabilityKind(
+          DifferentiabilityKind::Normal));
+  auto newThunkResult = thunkResult.getWithInterfaceType(diffResultFnTy);


If there is a common pattern here, maybe put the construction of this function onto a helper on SILOptFunctionBuilder? The hope with that was to hide the low level SILFunction creation functions in favor of higher level routines so we don't have to have huge invocations of SILFunctionType::get.

The code here takes a $(Self) -> (T, ...) -> U function (curry thunk) and clones it into a $(Self) -> @differentiable (T, ...) -> U function.

SILOptFunctionBuilder looks very minimal, so I'm not sure adding an entry point makes sense. We should add a type calculation helper to SILFunctionType, but that seems ad-hoc because this is the only user.

Let me know if you have further suggestions!

@dan-zheng I wrote that code and the reason why it is minimal now is that I was trying to create the scaffolding for what I just mentioned. Just no one has had the time to fix it up old uses. That being said, we should not make the problem worse.

gottesmm · 2020-04-10T23:09:28Z

lib/SILOptimizer/Utils/Differentiation/Common.cpp

+  }
+}
+
+SILValue joinElements(ArrayRef<SILValue> elements, SILBuilder &builder,


Is this really needed? Could we put this onto SILBuilder or something? Seems wrong to have such helper utilities.

I would give it a name that makes it explicit that you are using a tuple. Otherwise, how do I know I am getting back a tuple?

Addressing in #30970 (comment), let's continue discussion there.

gottesmm · 2020-04-10T23:12:15Z

lib/SILOptimizer/Utils/Differentiation/Common.cpp

+  return builder.createTuple(loc, elements);
+}
+
+void extractAllElements(SILValue value, SILBuilder &builder,


This is void SILBuilder::emitDestructureValueOperation(SILLocation loc, SILValue operand,
SmallVectorImpl &result);

Addressing in #30970 (comment), let's continue discussion there.

gottesmm · 2020-04-10T23:15:32Z

lib/SILOptimizer/Utils/Differentiation/Common.cpp

+    results.push_back(builder.createTupleExtract(value.getLoc(), value, i));
+}
+
+void emitZeroIntoBuffer(SILBuilder &builder, CanType type,


This again doesn't feel like a method that has anything to do with differentiation.

Currently, this code currently emits AdditiveArithmetic.zero, but we plan to change it to emit Differentiable.zeroTangentVector sometime (TF-1005), so I'd prefer to keep it differentiation-specific.

Callers of this function could be deduped though: PullbackEmitter::emitZeroIndirect and JVPEmitter::emitZeroIndirect. I can look into that, but I'm not sure how to best use SILBuilderContext & since a SILBasicBlock::iterator is needed as an insertion point: #30818 (comment)

marcrasi · 2020-04-10T23:35:33Z

Also, are these used anywhere but the differentiation pass? If not, we should just put it in the differentiation pass until we find multiple uses.

Would moving all the contents of include/SILOptimizer/Utils/Differentiation and lib/SILOptimizer/Utils/Differentiation to a new directory named lib/SILOptimizer/Mandatory/Differentiation be an appropriate way to do this?

I would very much like to avoid moving everything into a single .cpp file because differentiation has many distinct bits of analysis and synthesis that are easier to comprehend separately. (Also having multiple files is really nice for incremental compilation time.)

dan-zheng · 2020-04-11T01:20:43Z

Thanks @gottesmm! I'll create a PR addressing your feedback, we can continue discussion there.

dan-zheng requested review from rxwei and marcrasi April 5, 2020 23:38

dan-zheng commented Apr 5, 2020

View reviewed changes

rxwei approved these changes Apr 5, 2020

View reviewed changes

dan-zheng merged commit 146c11e into swiftlang:master Apr 6, 2020

dan-zheng deleted the differentiation-transform branch April 6, 2020 03:22

gottesmm reviewed Apr 10, 2020

View reviewed changes

This was referenced Apr 11, 2020

[AutoDiff] NFC: reorganize differentiation SILOptimizer files. #30969

Merged

[AutoDiff] NFC: gardening. #30970

Closed

[AutoDiff upstream] Add differentiable_function canonicalization. #30818

[AutoDiff upstream] Add differentiable_function canonicalization. #30818

Uh oh!

Conversation

dan-zheng commented Apr 5, 2020

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dan-zheng commented Apr 5, 2020

Uh oh!

dan-zheng commented Apr 5, 2020

Uh oh!

swift-ci commented Apr 6, 2020

Uh oh!

dan-zheng commented Apr 6, 2020

Uh oh!

dan-zheng commented Apr 6, 2020

Uh oh!

gottesmm left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dan-zheng Apr 11, 2020 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dan-zheng Apr 11, 2020 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

marcrasi commented Apr 10, 2020

Uh oh!

dan-zheng commented Apr 11, 2020

Uh oh!

Uh oh!

[AutoDiff upstream] Add `differentiable_function` canonicalization. #30818

[AutoDiff upstream] Add `differentiable_function` canonicalization. #30818

dan-zheng Apr 11, 2020 •

edited

Loading

dan-zheng Apr 11, 2020 •

edited

Loading