Force deabstraction of functions that take or return Tensors if referred in code. #19052

bgogul · 2018-08-29T18:23:50Z

This is a somewhat hacky patch to deal with SR-8589 related to the Jupyter tutoral. In the REPL mode (both in the in-build swift and lldb -r), the SIL module for a REPL expr is thrown away after lowering to LLVM IR. Consequently, when processing subsequent lines, TFDeabstraction does not see the body of some functions being invoked, which causes some functions not to be partitioned.

This PR performs a second round of deabstraction (and subsequent partitioning) of functions taking and returning TensorValues if they are still referred in the code after the first round.

BTW, this is not only relevant to Jupyter tutorial. We will encounter this if we allow linking to other swift-tensorflow binaries. We will need to think about how to deal with such cases in general.

@mhong, when I simply removed the condition that checks if a function operates on tensors, some tests started failing as attributes were no longer const-evaluable. That is why I am sending this instead of what we discussed.

…re referred to in the code after a first round of deabstraction.

bgogul · 2018-08-29T18:24:13Z

@swift-ci please test tensorflow linux

bgogul · 2018-08-29T18:24:33Z

@swift-ci please test tensorflow macos

lattner

Agreed that this is not optimal, but I think it should work in the near term. Please file a bug to track improving this. A sketch:

Make our existing pass over functions, deabstracting as we go.
Do a second pass where we deabstract all functions, regardless of whether they have tensor args or results. Many would be dead, but the indirectly used ones would now work.

rxwei

The fix makes sense, unblocking lots of cross-module cases.

There is a more general problem though: whenever there’s a call to an opaque function which calls a protocol method (e.g. Equatable.==) passing tensor arguments whose implementation of this protocol is inlinable, execution will fail. This is because Tensor’s protocol implementation is never getting partitioned.

To resolve the bigger problem later down the road, partitioning generic functions should be supported and will require a major architectural change.

rxwei · 2018-08-29T23:07:24Z

lib/SILOptimizer/Mandatory/TFUtilities.cpp

+      // Return true if it is referenced somewhere.
+      // TODO: There might be a better check other than fn->getRefCount() > 0.
+      // See the clean up code in TFDeabstraction::inlineCalls() function.
+      return (fn->getRefCount() > 0);


Unnecessary parentheses.

mhong

Having this as a short-term patch sounds good, but please add some unit test.

To avoid breaking unit tests, one option is to flag-control the value of forceTFFunctions, and only turn it on in REPL mode. This will buy us time to rewrite the relevant unit tests -- if we don't want a function to be partitioned (e.g. if the func takes a string/int/float param, and that param is used in a tfop attr), we can mark it " inline always". Once we finish this on all tests, we can then retire that flag.

mhong · 2018-08-30T00:50:01Z

lib/SILOptimizer/Mandatory/TFUtilities.h

@@ -277,7 +277,11 @@ class TensorFunctionClassifier {
  /// example) for inlined functions that take and return tensors, since we
  /// know that they are either unreachable or will be inlined into any
  /// clients that use them.
-  bool shouldBePartitioned(SILFunction *fn);
+  ///
+  /// If the flag forceTFFunctions is true, it forces partitioning of functions


the comment does not quite reflect the impl, because there are conditions that cause the function to return false.
e.g. if (isAvailableExternally(fn->getLinkage())). Also when a function is marked inline always.

mhong · 2018-08-30T00:51:16Z

lib/SILOptimizer/Mandatory/TFUtilities.h

+  /// If the flag forceTFFunctions is true, it forces partitioning of functions
+  /// that operate on Tensor values irrespective of whether they are inlinable,
+  /// private, etc.
+  bool shouldBePartitioned(SILFunction *fn, bool forceTFFunctions = false);


since this func does not have many call-sites, it could be less error prone if we don't use a default value for forceTFFunctions.

mhong · 2018-08-30T01:07:14Z

lib/SILOptimizer/Mandatory/TFDeabstraction.cpp

+  // are functions that are *still* referred to in the code and operate on
+  // tensor values, but have not been partitioned. It can happen in the
+  // following case, for instance, where `foo` is an external function that
+  // takes a for which we do not have the body:


what is "takes a for"?

typo. Fixed the comment.

mhong · 2018-08-30T01:07:28Z

lib/SILOptimizer/Mandatory/TFDeabstraction.cpp

+  // following case, for instance, where `foo` is an external function that
+  // takes a for which we do not have the body:
+  //   main() {
+  //     foo() { $0 -= 0.1 * $1 }


can we turn something like this into a unit test?

Added a unit test in deabstraction_finished.swift

mhong · 2018-08-30T01:09:25Z

lib/SILOptimizer/Mandatory/TFDeabstraction.cpp

+  //
+  // (Note the body of a function may be missing when we are linking against a
+  // library or in the REPL context where the function was defined on a
+  // different line.)


nit: "different line" might not be the right description.

@marcrasi showed me earlier that you can copy/paste multi-line text into REPL, and that'll be executed as one "cell".

Calling it REPL line now. Not sure if that is the right terminology.

mhong · 2018-08-30T01:14:18Z

lib/SILOptimizer/Mandatory/TFDeabstraction.cpp

+  for (auto &fn : *module) {
+    // Skip if it is already partitioned, or if it was ignored only because it
+    // operated on tensor values.
+    if (partitionedFunctions.count(&fn) > 0 ||


can you add a comment block motivating this two-round / two-pass design?

IIUC, in your example above, if we just use tfc.shouldBePartitioned(&fn, /*forceTFFunctions=*/true) in the first round, and make it the only round, we'll be able to partition the private closure, so it achieves the same correctness.

the two-round design is an optimization to minimize the set of functions to partition (would be nice to have a unit test that shows its effects -- have a function with no caller left after the first round finishes, and confirm it's not partitioned in the second round).

Added a few lines motivating the two round design.

mhong · 2018-08-30T01:16:32Z

lib/SILOptimizer/Mandatory/TFUtilities.cpp

@@ -631,7 +631,8 @@ tf::createTensorToInt1Inst(SILValue value, SILBuilder &builder,
 /// example) for inlined functions that take and return tensors, since we
 /// know that they are either unreachable or will be inlined into any
 /// clients that use them.


consider removing this comment block as it's now out of date (the new param is not documented).

(i'd also be fine with keeping it in sync with the corresponding comment block in the header file.)

Removed the block comment altogether.

mhong · 2018-08-30T01:49:27Z

@rxwei To help me understand this larger problem, can you provide a concrete example?

(It might be useful to create a bug to track that.)

rxwei · 2018-08-30T02:07:02Z

A common, concrete example is to try to use expectEqual() on two Tensors. It is guaranteed to fail, because expectEqual is opaque across module, i.e. non-inlinable, and Tensor.== is not getting partitioned because there's no inlining.

This is not a small bug, but a large hole in the architecture in that there's no way to even reject it from the compiler. These issues are not important in the initial GPE because the library ecosystem using Tensors and generics are not there yet.

I'll definitely file a bug soon. I was also planning to write a detailed document describing the issue, but it's not the best time until we hit the tutorial goal.

bgogul · 2018-08-30T06:11:46Z

@lattner , this PR already implements it in two passes as you outlined. Did I misunderstand your comment?

Agreed that this is not optimal, but I think it should work in the near term. Please file a bug to track improving this. A sketch:

Make our existing pass over functions, deabstracting as we go.

Do a second pass where we deabstract all functions, regardless of whether they have tensor args or results. Many would be dead, but the indirectly used ones would now work.

mhong · 2018-08-30T10:13:32Z

@bgogul, I believe the difference is: If we change the example from

main() {
  foo() { $0 -= 0.1 * $1 }
}

to

main() {
  foo() { #tfop("SomeOp", ...) }
}

This patch probably won't work. That is, if the private closure does not take or return tensors.

Another issue to take care (both in Chris' sketch and this PR) is that if the private closure takes a string/int/float param to be used in a tfop attr, as in:

main() {
  foo() { (s: String, i: Int) -> () in #tfop("SomeOp", intAttr: i, sttrAttr: s) }
}

The private closure cannot be partitioned due to const tfop attr requirement (as you observed). In that case it'd be useful to issue a good diagnostic message, asking user to convert the closure to a function, so that it can be annotated with something like "inline always", and as such compiler won't partition it.

rxwei · 2018-08-30T10:25:13Z

In that case it'd be useful to issue a good diagnostic message, asking user to convert the closure to a function, so that it can be annotated with something like "inline always", and as such compiler won't partition it.

In cases like this, it's also possible to make the closure function serialized (inlinable) by simply adding a [serialized] modifier to the SIL function. It'd be much simpler than asking the user to change the code and make it serialized.

…_8589

bgogul · 2018-08-30T18:26:30Z

Added unit tests and addressed comments. PTAL.

mhong

Nice!

Force partitioning of functions that take or return Tensors if they a…

add69ca

…re referred to in the code after a first round of deabstraction.

bgogul requested review from mhong, rxwei and lattner and removed request for rxwei August 29, 2018 18:32

lattner approved these changes Aug 29, 2018

View reviewed changes

rxwei approved these changes Aug 29, 2018

View reviewed changes

rxwei added the tensorflow This is for "tensorflow" branch PRs. label Aug 30, 2018

mhong approved these changes Aug 30, 2018

View reviewed changes

bgogul added 2 commits August 30, 2018 11:16

Addressed review comments and added a unit test.

fd58938

Merge branch 'tensorflow' of github.com:apple/swift into param_update…

30be6fd

…_8589

mhong approved these changes Aug 30, 2018

View reviewed changes

bgogul merged commit 7675f65 into apple:tensorflow Aug 30, 2018

bgogul deleted the param_update_8589 branch August 30, 2018 19:01

swift-ci mentioned this pull request Aug 29, 2018

[SR-8589] [blocker] model.updateParameters in the iris tutorial crashes the compiler #51107

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Force deabstraction of functions that take or return Tensors if referred in code. #19052

Force deabstraction of functions that take or return Tensors if referred in code. #19052

bgogul commented Aug 29, 2018 •

edited

bgogul commented Aug 29, 2018

bgogul commented Aug 29, 2018

lattner left a comment

rxwei left a comment

rxwei Aug 29, 2018

bgogul Aug 30, 2018

mhong left a comment

mhong Aug 30, 2018

bgogul Aug 30, 2018

mhong Aug 30, 2018

bgogul Aug 30, 2018

mhong Aug 30, 2018

bgogul Aug 30, 2018

mhong Aug 30, 2018

bgogul Aug 30, 2018

mhong Aug 30, 2018

bgogul Aug 30, 2018

mhong Aug 30, 2018

bgogul Aug 30, 2018

mhong Aug 30, 2018 •

edited

bgogul Aug 30, 2018

mhong commented Aug 30, 2018

rxwei commented Aug 30, 2018 •

edited

bgogul commented Aug 30, 2018

mhong commented Aug 30, 2018

rxwei commented Aug 30, 2018 •

edited

bgogul commented Aug 30, 2018

mhong left a comment

Force deabstraction of functions that take or return Tensors if referred in code. #19052

Force deabstraction of functions that take or return Tensors if referred in code. #19052

Conversation

bgogul commented Aug 29, 2018 • edited

bgogul commented Aug 29, 2018

bgogul commented Aug 29, 2018

lattner left a comment

Choose a reason for hiding this comment

rxwei left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

mhong left a 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

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

mhong Aug 30, 2018 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

mhong commented Aug 30, 2018

rxwei commented Aug 30, 2018 • edited

bgogul commented Aug 30, 2018

mhong commented Aug 30, 2018

rxwei commented Aug 30, 2018 • edited

bgogul commented Aug 30, 2018

mhong left a comment

Choose a reason for hiding this comment

bgogul commented Aug 29, 2018 •

edited

mhong Aug 30, 2018 •

edited

rxwei commented Aug 30, 2018 •

edited

rxwei commented Aug 30, 2018 •

edited