[JIT] add CondValue to unify refinements and code emission by zdevito · Pull Request #26145 · pytorch/pytorch

zdevito · 2019-09-13T00:21:12Z

Stack from ghstack:

[jit] Allow 'Any' to appear as a type argument. #26572 [jit] Allow 'Any' to appear as a type argument.
[jit] refactor tryMatchSchema #26499 [jit] refactor tryMatchSchema
[JIT] add type refinements for isinstance checks #26271 [JIT] add type refinements for isinstance checks
[jit] refactor extra sugared values #26270 [jit] refactor extra sugared values
[jit] add dynamic isinstance #26269 [jit] add dynamic isinstance
Allow types as node attributes #26268 Allow types as node attributes
[JIT] add CondValue to unify refinements and code emission #26145 [JIT] add CondValue to unify refinements and code emission

This is step towards isinstance type refinement.
It primarily does yak shaving in compiler.cpp to unify the handling
of special case behavior that occurs in conditional expressions:

Handling type refinement as part of emission.
Handling is None static-if specialization.

It introduces a CondValue object that is a Value that also has
additional type refinements that are true when that Value is true,
and potentialy a static-true/false value that, if set, will cause if
statements to be handled statically, omitting typechecking of the other side.

This ends up expanding some behavior, for instance is None specialization
used to happen only for single expressions, but now works through
boolean logic.

Differential Revision: D17359500

This is step towards isinstance type refinement. It primarily does yak shaving in compiler.cpp to unify the handling of special case behavior that occurs in conditional expressions: * Handling type refinement as part of emission. * Handling `is None` static-if specialization. It introduces a CondValue object that is a Value that also has additional type refinements that are true when that Value is true, and potentialy a static-true/false value that, if set, will cause if statements to be handled statically, omitting typechecking of the other side. This ends up expanding some behavior, for instance `is None` specialization used to happen only for single expressions, but now works through boolean logic.

[JIT] add CondValue to unify refinements and code emission This is step towards isinstance type refinement. It primarily does yak shaving in compiler.cpp to unify the handling of special case behavior that occurs in conditional expressions: * Handling type refinement as part of emission. * Handling `is None` static-if specialization. It introduces a CondValue object that is a Value that also has additional type refinements that are true when that Value is true, and potentialy a static-true/false value that, if set, will cause if statements to be handled statically, omitting typechecking of the other side. This ends up expanding some behavior, for instance `is None` specialization used to happen only for single expressions, but now works through boolean logic. gh-metadata: pytorch pytorch 26145 gh/zdevito/107/head

eellison

Cool! CondValue will lead to less surprising user behavior around selective compilation of if branches.

I have a couple questions / comments on some decisions made. Also i don't think we should regress the refinement calculation for intersection and union given that we are opening the door for more complicated refinements with isinstance checks moving forward.

eellison · 2019-09-16T17:03:52Z

torch/csrc/jit/script/compiler.cpp


-using TypeAndRange = std::pair<TypePtr, const SourceRange*>;
+enum RefinementKind {
+  OPTIONAL_PRESENT, // the value, which has type optional[T] is present, and


What's the reason for RefinementKind instead of just containing a typeptr ?

eellison · 2019-09-16T17:42:11Z

torch/csrc/jit/script/compiler.cpp

+  bool operator==(const Refinement& rhs) const {
+    bool type_same =
+        (type_arg_ && rhs.type_arg_) ? *type_arg_ == *rhs.type_arg_ : false;
+    return type_same && identifier_ == rhs.identifier_ && kind_ == rhs.kind_;


Not sure exactly how the == operator is being used but if got rid of RefinementKind and just stored a reference to the refined type, semantically equivalent refinements would work here:

x: Optional[int] = 1 isisstance(x, int) and (x is not None)

eellison · 2019-09-16T17:53:28Z

torch/csrc/jit/script/compiler.cpp

+  // with its true value and its false value. If a boolean that has refinements
+  // associated with it is used in a conditional of an if statememt, the true
+  // and false refinements are inserted into the corresponding blocks
+  using Refinements = std::vector<Refinement>;


It's a little confusing that the struct is named RefinementSet, but is backed by a vector, and in union we add duplicates of a refinement.

fixed the union as we discussed in person, it now behaves more like a set.

eellison · 2019-09-16T18:29:57Z

torch/csrc/jit/script/compiler.cpp

+      }
+      case TK_NOT: {
+        CondValue v = emitCondExpr(Expr(expr.tree()->trees()[0]));
+        Value* input = emitToBool(v.value());


Is there any case when emitCondExpr should be called but emitToBool should not follow ? could we just move emitToBool within emitCondExpr

eellison · 2019-09-16T20:00:06Z

torch/csrc/jit/script/compiler.cpp

+    // so we take their intersection.
+    return RefinementSet(
+        unionList(true_refinements_, rhs.true_refinements_),
+        intersectList(false_refinements_, rhs.false_refinements_));


intersectList but the struct is named refinementSet

eellison · 2019-09-16T20:02:44Z

torch/csrc/jit/script/compiler.cpp

-        ret.setRefinement(name_mapping.first, name_mapping.second);
+  static Refinements intersectList(const Refinements& a, const Refinements& b) {
+    Refinements result;
+    for (const Refinement& r : a) {


What is the motivation is for removing the existing intersection and union implementations? The new implementation probably suffices for just optional refinement, but with the possibility of more complicated refinement using isinstance and an Any type i think we should use the more correct existing implementation.

Something like isinstance(x, float) or isinstance(x, int) would just accept the second type instead of unifying them.

Note: the existing intersection and union produces the correct refinements, just the insertRefinements function incorrectly inserts a refinement when the refinement is not None.

discussed in person: there may be unexpected consequences of unifyTypes refinement. For instance if we derive a type Scalar, it is another instance of injecting a Scalar in addition to the item call. This has led to surprising bugs in the past, so we should have a good example case before we make more liberal use of refinement. The behavior currently implemented is a subset of the valid refinements that we would do in the future, so it is safe to expand.

[JIT] add CondValue to unify refinements and code emission This is step towards isinstance type refinement. It primarily does yak shaving in compiler.cpp to unify the handling of special case behavior that occurs in conditional expressions: * Handling type refinement as part of emission. * Handling `is None` static-if specialization. It introduces a CondValue object that is a Value that also has additional type refinements that are true when that Value is true, and potentialy a static-true/false value that, if set, will cause if statements to be handled statically, omitting typechecking of the other side. This ends up expanding some behavior, for instance `is None` specialization used to happen only for single expressions, but now works through boolean logic. gh-metadata: pytorch pytorch 26145 gh/zdevito/107/head

eellison

Looks good!

eellison · 2019-09-17T21:00:01Z

torch/csrc/jit/script/compiler.cpp

+    // Refinement none(name, RefinementKind::OPTIONAL_NONE);
+
+    Refinement present(
+        name, lhs_value->type()->expect<OptionalType>()->getElementType());


Probably don't want to expect<OptionalType> here because it will error for something like:

a: int = 1 if a is None:

which can happen with a constant attribute that is initialized non-null

eellison · 2019-09-17T21:00:20Z

torch/csrc/jit/script/compiler.cpp

+    return a.identifier() == b.identifier();
+  }
+  static Refinements unionList(const Refinements& a, const Refinements& b) {
+    Refinements result = a;


Could consider calling this unionSet now that we are enforcing uniqueness of each element

eellison · 2019-09-17T21:00:31Z

torch/csrc/jit/script/compiler.cpp

-        ret.setRefinement(name_mapping.first, name_mapping.second);
+    return result;
+  }
+  static Refinements intersectList(const Refinements& a, const Refinements& b) {


Same comment

eellison · 2019-09-17T21:02:16Z

torch/csrc/jit/script/compiler.cpp

+    // submodule is absent, an is None check can be used to ensure the
+    // accesses to the None check, which would error, are not compiled.
+    if (cond_value.staticIf()) {
+      if (*cond_value.staticIf()) {


Not compiling one branch of an if statement with multiple staticIf checks is a new feature so maybe add tests for that

[JIT] add CondValue to unify refinements and code emission This is step towards isinstance type refinement. It primarily does yak shaving in compiler.cpp to unify the handling of special case behavior that occurs in conditional expressions: * Handling type refinement as part of emission. * Handling `is None` static-if specialization. It introduces a CondValue object that is a Value that also has additional type refinements that are true when that Value is true, and potentialy a static-true/false value that, if set, will cause if statements to be handled statically, omitting typechecking of the other side. This ends up expanding some behavior, for instance `is None` specialization used to happen only for single expressions, but now works through boolean logic. gh-metadata: pytorch pytorch 26145 gh/zdevito/107/head

facebook-github-bot · 2019-09-24T02:36:22Z

@zdevito merged this pull request in fcd1354.

zdevito requested a review from apaszke as a code owner September 13, 2019 00:21

pytorchbot added the oncall: jit Add this issue/PR to JIT oncall triage queue label Sep 13, 2019

This was referenced Sep 13, 2019

Add ivalue::type(), part 1 #25439

Closed

Use static type information to restore type tags #25447

Closed

[jit] Refactor emitIsInstance #26061

Closed

zdevito requested review from eellison and suo September 13, 2019 00:22

This was referenced Sep 16, 2019

Allow types as node attributes #26268

Closed

[jit] add dynamic isinstance #26269

Closed

[jit] refactor extra sugared values #26270

Closed

[JIT] add type refinements for isinstance checks #26271

Closed

eellison reviewed Sep 16, 2019

View reviewed changes

zdevito added 2 commits September 16, 2019 17:20

zdevito requested a review from eellison September 17, 2019 00:37

eellison approved these changes Sep 17, 2019

View reviewed changes

zdevito added 3 commits September 18, 2019 13:16

zdevito mentioned this pull request Sep 19, 2019

[jit] refactor tryMatchSchema #26499

Closed

zdevito added 2 commits September 20, 2019 10:25

zdevito mentioned this pull request Sep 20, 2019

[jit] Allow 'Any' to appear as a type argument. #26572

Closed

facebook-github-bot closed this in fcd1354 Sep 23, 2019

facebook-github-bot added the merged label Sep 24, 2019

This was referenced Sep 24, 2019

[jit] Fix circular deps in loading #26758

Closed

[jit] Remove attempToRecoverType #26767

Closed

[jit] remove dead LEGACY_PythonPrint #26786

Closed

[jit] Make PythonPrint a class #26787

Closed

[jit] clean up exported source format #26788

Merged

facebook-github-bot deleted the gh/zdevito/107/head branch October 28, 2019 22:23

mruberry added the Merged label Oct 28, 2020

Conversation

zdevito commented Sep 13, 2019 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

eellison 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!

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!

eellison 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!

facebook-github-bot commented Sep 24, 2019

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

zdevito commented Sep 13, 2019 •

edited

Loading