Teach SILGen to reabstract @autodiff functions #21367

marcrasi · 2018-12-17T02:38:43Z

This PR teaches SILGen to reabstract @autodiff functions by splitting them into their components, reabstracting all the components, and then putting them back together again. cc @slavapestov for your thoughts on this approach.

I pulled in some additional fixes from @rxwei:

rxwei/swift@372748e: Update autodiff-associated function type calculation to work with abstraction thunks.
rxwei/swift@fe47c71: Add a conversion rule for convert_function where the converted function has @noescape.
rxwei/swift@ed8e7d0: Retain arguments before partial_apply, fixing a crasher.

I also switched the order of autodiff_function and function_conversion_expr in CSApply, to work around the problem where we can't convert escaping @autodiff functions to noescape.

As a result of these changes, differential operators work, so I enabled the tests.

rxwei · 2018-12-17T02:44:55Z

Awesome! I've just experimented this (reversed function conversion and peer-through for convert_escape_to_noescape in the AD pass) on my local branch. It's great that you've finished this!

rxwei · 2018-12-17T02:45:22Z

include/swift/SIL/SILInstruction.h

@@ -7659,6 +7659,9 @@ class AutoDiffFunctionExtractInst :
    Extractee(AutoDiffAssociatedFunctionKind kind);
    explicit Extractee(StringRef name);
    operator innerty() const { return rawValue; }
+
+    llvm::Optional<AutoDiffAssociatedFunctionKind>


No need for llvm::, I think.

I think calling it getExtracteeAsAssociatedFunction() would provide more clarity.

I'm going to split this out into a separate PR because it's a nice refactoring but it's no longer important for this PR after I include your assoc fn ty calculation fix.

rxwei · 2018-12-17T02:45:32Z

lib/SIL/SILInstructions.cpp

@@ -688,32 +688,53 @@ AutoDiffFunctionExtractInst::Extractee::Extractee(StringRef string) {
  rawValue = *result;
 }

+llvm::Optional<AutoDiffAssociatedFunctionKind>


Suggested change

llvm::Optional<AutoDiffAssociatedFunctionKind>

Optional<AutoDiffAssociatedFunctionKind>

rxwei · 2018-12-17T02:50:02Z

lib/Sema/CSApply.cpp

+        expr = cs.cacheType(new (tc.Context)
+                            AutoDiffFunctionExpr(expr, toFunc));
+      }
+      else if (fromEI.isDifferentiable() && !toEI.isDifferentiable()) {


If we have an @escaping @autodiff function and we need to convert it to a non-escaping non-differentiable function, won't this emit a (autodiff_function_extract_original_expr (convert_function_expr f)) which would then trigger a reference counting crasher again? In this case, I think autodiff_function_extract_original should happen before function_conversion_expr.

rxwei · 2018-12-17T02:52:23Z

lib/SILOptimizer/Mandatory/TFDifferentiation.cpp

@@ -1668,6 +1670,15 @@ reapplyFunctionConversion(SILValue newFunc, SILValue oldFunc,
        loc, innerNewFunc, pai->getSubstitutionMap(), newArgs,
        ParameterConvention::Direct_Guaranteed);
  }
+  // convert_escape_to_noescape
+  if (auto *cetn = dyn_cast<ConvertEscapeToNoEscapeInst>(oldConvertedFunc)) {
+    auto innerNewFunc = reapplyFunctionConversion(newFunc, oldFunc, cetn->getOperand(), builder, loc, substituteOperand);


rxwei · 2018-12-17T04:05:34Z

lib/SIL/SILVerifier.cpp

+    // If the associated function comes from a reabstraction thunk, then it is
+    // impossible to determine the type of the associated function from the
+    // type of the original function, because we also need to know the
+    // abstraction pattern that the reabstraction made. So we currently


I see what's going on here:

The autodiff-associated function type calculation produces the following type:

$@noescape @callee_guaranteed (@in_guaranteed Float) -> (@out Float, @owned @callee_guaranteed (Float) -> Float)

But reabstraction makes the closure result's parameters and results indirect:

$@noescape @callee_guaranteed (@in_guaranteed Float) -> (@out Float, @owned @callee_guaranteed (@in_guaranteed Float) -> @out Float)

Currently SILFunctionType::getAutoDiffAssociatedFunctionType uses getFormalResultInfo and getFormalParameterInfo to determine the parameter convention and result convention of the differential/pullback. If we instead mirror their conventions in the original function type, then associated fn types can still be correctly calculated even if they are reabstracted, I think.

I updated SILFunctionType::getAutoDiffAssociatedFunctionType in 372748e and it's passing all tests. Feel free to pull that in.

rxwei · 2018-12-17T05:11:34Z

I made a bunch of additional fixes based on this PR, please feel free to pull those in.

rxwei@372748e: Update autodiff-associated function type calculation to work with abstraction thunks.
rxwei@fe47c71: Add a conversion rule for convert_function where the converted function has @noescape.
rxwei@ed8e7d0: Retain arguments before partial_apply, fixing a crasher.

With these fixes, generic functional differential operators are working end-to-end on tensors.

import Swift
import TensorFlow

func valueWithPullback<T, R>(
  at x: T, in f: @autodiff (T) -> R
) -> (R, (R.CotangentVector) -> T.CotangentVector)
  where T : Differentiable, R : Differentiable {
  return Builtin.autodiffApplyVJP(f, x)
}

let (y, pullback) = valueWithPullback(at: Tensor<Float>(1)) { x in
  return sin(x)
}
print(y) // 0.84147096
print(pullback(Tensor(1))) // 0.5403023

…traction thunks

marcrasi · 2018-12-17T21:08:07Z

I addressed your comments, pulled in all your fixes, and updated the PR description to describe the pieces you added.

marcrasi · 2018-12-17T21:08:15Z

@swift-ci please test tensorflow

marcrasi · 2018-12-17T21:08:26Z

@swift-ci please test tensorflow

slavapestov

The approach looks fine. Couple of nitpicks

slavapestov · 2018-12-17T21:16:43Z

lib/SILGen/SILGenPoly.cpp

+  // If this is differentiable, then we need to apply a thunk to all the
+  // components. We extract them, apply a thunk to them, and then combine them
+  // back into a bundle.
+  if (sourceType->isDifferentiable()) {


Maybe we could factor this out into a separate method?

slavapestov · 2018-12-17T21:17:10Z

lib/SILGen/SILGenPoly.cpp

+        SGF, loc, original, inputOrigTypeNotDiff, inputSubstTypeNotDiff,
+        outputOrigTypeNotDiff, outputSubstTypeNotDiff, expectedTLNotDiff);
+
+    // TODO: Use parameter indices specified in the funciton type.


Typo: funciton

marcrasi · 2018-12-17T21:38:54Z

@swift-ci please test tensorflow

marcrasi · 2018-12-17T21:39:11Z

@swift-ci please test tensorflow

marcrasi · 2018-12-17T22:47:17Z

@swift-ci please test tensorflow

marcrasi · 2018-12-17T22:47:21Z

@swift-ci please test tensorflow

marcrasi · 2018-12-17T22:47:42Z

@swift-ci please test tensorflow

marcrasi · 2018-12-17T22:48:13Z

@swift-ci please test tensorflow

marcrasi · 2018-12-18T00:24:03Z

@swift-ci please test tensorflow linux

rxwei · 2018-12-18T01:02:08Z

Merging this for rebasing my next PR.

marcrasi requested a review from rxwei December 17, 2018 02:38

marcrasi force-pushed the reabstraction-thunks branch from 5597098 to 2876e23 Compare December 17, 2018 02:40

rxwei reviewed Dec 17, 2018

View reviewed changes

Marc Rasi and others added 5 commits December 17, 2018 11:29

reabstraction thunks

b1fc37e

Update autodiff-associated function type calculation to work with abs…

fd0eec2

…traction thunks

Handle convert_function in reapplyFunctionConversion

e374500

Retain arguments before partial_apply, fixing a crasher.

d240fa8

fixes

5dd0ec8

marcrasi force-pushed the reabstraction-thunks branch from 2876e23 to 5dd0ec8 Compare December 17, 2018 21:05

rxwei approved these changes Dec 17, 2018

View reviewed changes

slavapestov reviewed Dec 17, 2018

View reviewed changes

address comments

728e9e4

fix build failure

a70af49

rxwei merged commit dab0804 into apple:tensorflow Dec 18, 2018

rxwei mentioned this pull request Dec 17, 2018

[SR-9458] [AD] Crash when emitting @autodiff function arguments in IRGen #51921

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Teach SILGen to reabstract @autodiff functions #21367

Teach SILGen to reabstract @autodiff functions #21367

marcrasi commented Dec 17, 2018 •

edited

rxwei commented Dec 17, 2018

rxwei Dec 17, 2018

rxwei Dec 17, 2018

marcrasi Dec 17, 2018

rxwei Dec 17, 2018

rxwei Dec 17, 2018

rxwei Dec 17, 2018

rxwei Dec 17, 2018

rxwei Dec 17, 2018

rxwei commented Dec 17, 2018 •

edited

marcrasi commented Dec 17, 2018

marcrasi commented Dec 17, 2018

marcrasi commented Dec 17, 2018

slavapestov left a comment

slavapestov Dec 17, 2018

slavapestov Dec 17, 2018

marcrasi commented Dec 17, 2018

marcrasi commented Dec 17, 2018

marcrasi commented Dec 17, 2018

marcrasi commented Dec 17, 2018

marcrasi commented Dec 17, 2018

marcrasi commented Dec 17, 2018

marcrasi commented Dec 18, 2018

rxwei commented Dec 18, 2018

	llvm::Optional<AutoDiffAssociatedFunctionKind>
	Optional<AutoDiffAssociatedFunctionKind>

Teach SILGen to reabstract @autodiff functions #21367

Teach SILGen to reabstract @autodiff functions #21367

Conversation

marcrasi commented Dec 17, 2018 • edited

rxwei commented Dec 17, 2018

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

rxwei commented Dec 17, 2018 • edited

marcrasi commented Dec 17, 2018

marcrasi commented Dec 17, 2018

marcrasi commented Dec 17, 2018

slavapestov left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

marcrasi commented Dec 17, 2018

marcrasi commented Dec 17, 2018

marcrasi commented Dec 17, 2018

marcrasi commented Dec 17, 2018

marcrasi commented Dec 17, 2018

marcrasi commented Dec 17, 2018

marcrasi commented Dec 18, 2018

rxwei commented Dec 18, 2018

marcrasi commented Dec 17, 2018 •

edited

rxwei commented Dec 17, 2018 •

edited