[jit] Tree based Iterator infrastructure: for in range/list/tensor/zip/enumerate #21801

wanchaol · 2019-06-14T19:46:04Z

Stack from ghstack:

[jit] add for in dict support #22006 [jit] add for in dict support
[jit] add for in string support #21990 [jit] add for in string support
[jit] support iteration tuple unpacking #21985 [jit] support iteration tuple unpacking
[jit] refactor all for in range/tensor tests to be together with other for loop tests #21950 [jit] refactor all for in range/tensor tests to be together with other for loop tests
[jit] Tree based Iterator infrastructure: for in range/list/tensor/zip/enumerate #21801 [jit] Tree based Iterator infrastructure: for in range/list/tensor/zip/enumerate
[jit] more complete tuple assignments #21949 [jit] more complete tuple assignments

Summary:

This makes the looping infrastructure on different types of iterators be tree based. SimpleValue(List, Dict, Tensor, etc) will be leaves of the tree, they know how to calculate the loop information using len() and get_elem(), while builtins like zip or enumerate become non-leaf of the tree.

When a users writes:

for (i, (j, k), l)  in zip(a, enumerate(b), range(0, 100)):

We generate:

IterableTree {
    root = ((), ((), ()), ())
    bases = {a, range(0, math.inf, 1), b, range(0, 100)};
}

for _i in range(0, prim::min(len(a), len(b), len(range(0, 100)))):
    b_0_i = a[_i]    
    b_1_i = range(0, math.inf, 1)[_i]    
    b_2_i = b[_i]    
    b_3_i = range(0, 100)[_i]
    i, (j, k), l = (b_0_i, (b_1_i, b_2_i), b_3_i)

Follow ups:

Iteration Tuple unpacking
for in Dict
for in String
reverse builtin

Differential Revision: D15948545

for in enumerate and zip

…tensor/zip/enumerate" [jit] Tree based Iterator infrastructure: for in range, list and tensor, for in enumerate and zip gh-metadata: pytorch pytorch 21801 gh/wanchaol/18/head

zdevito

This is a good start. Next steps:

let's try to get everything into the len/getitem abstraction, include the IterableTree.
the destructuring assignment to assign to the lists iteration variables should use the same code as normal assignments, so it handles tuple restructuring correctly.

zdevito · 2019-06-17T03:00:15Z

test/test_jit.py

            return torch.ones(x), x
        self.checkScript(stuff3, ([3, 2],))

-    def test_for_in_tensors(self):


This move makes it really hard to see what tests were added for this PR. Can you tell me the test names that were added/modified or split this movement into a separate PR?

test_for_in_tensors and test_for_in_range are moved and grouped together without changes. Other tests are added, I will split this into a separate PR.

zdevito · 2019-06-17T03:00:51Z

test/test_jit.py

+        def fn(x):
+            # type: (List[int]) -> int
+            sum = 0
+            for (i, v) in enumerate(x):


for i, v in enumerate(x) is valid syntax, and will need to work. Changes for it can be in a separate PR.

Yes that is in one of my follow up list

zdevito · 2019-06-17T03:03:42Z

torch/csrc/jit/script/sugared_value.h

+// (a, (range(0, math.inf, 1), b), range(0, 100))
+// We use those base iterables to fill in the loop information like max_trip_count
+// and set the value table for loop targets
+struct TORCH_API IterableTree : SugaredValue {


Why does this not define get_elem and length?

zdevito · 2019-06-17T03:05:05Z

torch/csrc/jit/script/sugared_value.h

+    throw ErrorReport(loc) << "cannot get the length of value " << kind();
+  }
+  // expression for ith elemement for iterable value
+  virtual Value* get_elem(const SourceRange&loc, Function& m, Value* i) {


this should be call getelem to match Python's __getelem__

zdevito · 2019-06-17T03:06:55Z

torch/csrc/jit/script/sugared_value.cpp

    at::ArrayRef<NamedValue> inputs,
    at::ArrayRef<NamedValue> attributes,
    size_t n_binders) {
+


Range does not make sense as a BuiltinFunction -- it completely ignores the entire ::call infrastructure BuiltinFunction uses. It should be its own thing.

zdevito · 2019-06-17T03:14:16Z

torch/csrc/jit/script/compiler.cpp

      // to set a current element
-      if (current_element_assigner) {
-        current_element_assigner(trip_count, environment_stack);
+      if (iter_val != nullptr && targets) {


The comment above is now nonsense, and I can't follow what iter_val or targets are in this context.

zdevito · 2019-06-17T03:15:45Z

torch/csrc/jit/script/compiler.cpp

-      } else {
-        step_val = end_val->owningGraph()->insertConstant(1);
-        step_val->node()->setSourceRange(range);
+  // recursively assign the targets base on the iterable tree structure


I do not think it is a good idea to special case the iterable_tree here. There are many places in the compiler that already do tuple assignment, e.g.
a, (b, c) = foo()`

We should be reusing that code for assign the loop lhs to whatever the result of getitem is. The IterableTree can simply return a tree from that.

does that mean we should make the getitem to return SugaredValue instead of Value*?

zdevito · 2019-06-17T03:17:06Z

torch/csrc/jit/script/compiler.cpp

-    // for-in lists
-    if (siv && siv->getValue()->type()->kind() == TypeKind::ListType) {
-      emitForInListLoop(stmt, siv);
+    if ((siv && (siv->getValue()->type()->kind() == TypeKind::ListType


These branches for different iteration kinds should be avoided. We need to fit all iteraiton into the len/getitem abstraction.

I think we will still need this condition for the tuple/module list looping(or unpacking) to separate from List/Tensor SimpleValues, but I will merge the iterableTree condition to fit into one abstraction.

zdevito · 2019-06-17T03:19:33Z

torch/csrc/jit/script/compiler.cpp

+      const std::shared_ptr<IterableValue>& iterable) {
+      std::shared_ptr<IterableTree> iterable_tree = nullptr;
+      size_t input_size = inputs.size();
+      if (iterable->symbol_ == prim::enumerate) {


Could this be done on construction of the sugared value, to avoid having hold a prim::enumerate symbol to direct what to do later?

What I observed from the emitApplyExpr is that, it will first do emitSugaredExpr(apply.callee(), 1) without passing inputs in, and the builtins will emit corresponding sugared values, I could not do it there since it's missing the inputs, that's where I need the singleton IterableValue sugared value to record the symbol and do the later tree construction. I could instead do a whole special casing on the IterableTree when emitApplyExpr, but that makes it separate from the builtin registration table mechanism that we already using..

zdevito · 2019-06-17T03:23:15Z

torch/csrc/jit/script/sugared_value.h

+    Value* start_;
+    Value* end_;
+    Value* step_;
+    bool is_simple_;


What does this mean?

is_simple_ means we the input args of range only contains end, and this serves as a flag in len and getelem to not insert unnecessary builtin nodes.

maybe use has_start_and_step_ then?

softagram-bot · 2019-06-17T07:49:52Z

This impact report was requested by @villelaitila - get yours at https://softagram.com/pull-request-bot

Softagram Impact Report for pull/21801 (head commit: `d1e0d41`)

⭐ Change Overview

(Open in Softagram Desktop for full details)

⭐ Details of Dependency Changes

(Open in Softagram Desktop for full details)

📄 Full report

Permalink: Full report for pull/21801

Give feedback on this report to support@softagram.com

suo · 2019-06-17T16:57:06Z

wat

…t/tensor/zip/enumerate" [jit] Tree based Iterator infrastructure: for in range, list and tensor, for in enumerate and zip gh-metadata: pytorch pytorch 21801 gh/wanchaol/18/head

…ure: for in range/list/tensor/zip/enumerate" [jit] Tree based Iterator infrastructure: for in range, list and tensor, for in enumerate and zip gh-metadata: pytorch pytorch 21801 gh/wanchaol/18/head

…e/list/tensor/zip/enumerate" [jit] Tree based Iterator infrastructure: for in range, list and tensor, for in enumerate and zip gh-metadata: pytorch pytorch 21801 gh/wanchaol/18/head

…t/tensor/zip/enumerate" [jit] Tree based Iterator infrastructure: for in range, list and tensor, for in enumerate and zip gh-metadata: pytorch pytorch 21801 gh/wanchaol/18/head

…in range/list/tensor/zip/enumerate" [jit] Tree based Iterator infrastructure: for in range, list and tensor, for in enumerate and zip gh-metadata: pytorch pytorch 21801 gh/wanchaol/18/head

zdevito

Looks good, just minor issues

zdevito · 2019-06-21T04:13:25Z

test/test_jit.py


    def test_for_in_tensors_fail_scalar(self):
-        with self.assertRaisesRegex(RuntimeError, "cannot be used as a tuple"):
+        with self.assertRaisesRegex(RuntimeError, "cannot get length of the value type float"):


These error messages are not very good. It should say something about not being iterable. What is the python error message in this case? Can we match it?

zdevito · 2019-06-21T04:14:38Z

torch/csrc/jit/script/parser.cpp

+        // TK_FOR targets should only parse exprs prec greater than 4, which only
+        // includes subset of Exprs that suppose to be on the LHS according to the
+        // python grammer https://docs.python.org/3/reference/grammar.html
+        auto target = parseExp(4);


I think this is meant to call parseLhsExp, otherwise it is unused.

zdevito · 2019-06-21T04:14:46Z

torch/csrc/jit/script/parser.cpp


+  // parse LHS acceptable exprs, which only includes subset of Exprs that prec is
+  // greater than 4 according to the python grammer
+  Expr parseLhsExp() {


nit: parseLHSExp

zdevito · 2019-06-21T04:17:04Z

torch/csrc/jit/script/sugared_value.cpp

+      end_ = inputs[1];
+      if (inputs.size() == 3) {
+        step_ = inputs[2];
+        // error handling when step_val = 0 during runtime


Can we make __range_length just do this error handling? It is much more inline with how we already do error handling everywhere else.

zdevito · 2019-06-21T04:20:03Z

torch/csrc/jit/script/sugared_value.cpp

+  for (const SugaredValuePtr& base_iter: base_iters) {
+    lengths.emplace_back(base_iter->len(loc, m));
+  }
+  Node* list_node = g.insertNode(g.create(prim::ListConstruct, 1)->setSourceRange(loc));


use node->createList please

zdevito · 2019-06-21T04:21:46Z

torch/csrc/jit/script/sugared_value.h

  }

+  // return length of this thing, if not then it can't be iterated.
+  virtual Value* len(const SourceRange& loc, Function& m) {


Can we make these error message match what python says in these circumstances? Since these are only used in iteration, it should be talking about how something is/isn't iterable.

zdevito · 2019-06-21T04:23:53Z

torch/csrc/jit/script/sugared_value.h

+    Value* start_;
+    Value* end_;
+    Value* step_;
+    bool is_simple_;


maybe use has_start_and_step_ then?

…in range/list/tensor/zip/enumerate" [jit] Tree based Iterator infrastructure: for in range, list and tensor, for in enumerate and zip gh-metadata: pytorch pytorch 21801 gh/wanchaol/18/head

Krovatkin

Krovatkin · 2019-06-20T19:02:35Z

torch/csrc/jit/register_prim_ops.cpp

+          int64_t min_element = std::numeric_limits<int64_t>::max(); 
+
+          for(int64_t ele: int_list) {
+            if(ele < min_element) {


nitpick: we could probably use std::min_element?

Unfortunately c10::List does not allow me to use it :(

Krovatkin · 2019-06-20T19:42:50Z

torch/csrc/jit/script/sugared_value.cpp

+    return end_;
+  } else{
+    Graph& g = *m.graph();
+    return g.insert(aten::__range_length, {start_, end_, step_}, {}, loc);


the purpose of is_simple_ to expose upper bound to enable some optimizations? otherwise, __range_length should be able to handle simple_ cases as well? If that's the case, I wonder if we should explicitly state it in a comment?

i.e. like stating why we don't want to insert a length calculation and index computation

// a flag to determine if it's a simple range() call with only end_ from arguments // If true, we will not insert length calculation and index derivation

Yes we don't want to insert the node when it's not necessary to do so. I will update with more complete comments

Krovatkin · 2019-06-21T22:06:21Z

torch/csrc/jit/script/compiler.cpp

+        std::vector<Value*> input_vals = getValues(inputs, /*maybe_unpack=*/true);
+        return std::make_shared<RangeValue>(loc, method, input_vals);
+      } else if (iterable->symbol_ == prim::enumerate) {
+        // enumerate(x) can be rewrite as subtrees: (range(0, math.inf), x)


maybe we could expand your comment into IterableTree(RangeValue(0, math.inf), SimpleValue(x)

…range/list/tensor/zip/enumerate" [jit] Tree based Iterator infrastructure: for in range, list and tensor, for in enumerate and zip gh-metadata: pytorch pytorch 21801 gh/wanchaol/18/head

…t/tensor/zip/enumerate" [jit] Tree based Iterator infrastructure: for in range, list and tensor, for in enumerate and zip gh-metadata: pytorch pytorch 21801 gh/wanchaol/18/head

facebook-github-bot · 2019-06-22T08:32:55Z

@wanchaol merged this pull request in e0f5ab2.

…erate (#21801) Summary: Pull Request resolved: pytorch/pytorch#21801 ghimport-source-id: b019d3e9a6f9bf152991a01b40e424dff176ffaa Test Plan: Imported from OSS Differential Revision: D15948545 Pulled By: wanchaol fbshipit-source-id: 6110a0f3ab08cbbb398441e8330f56083ecd2d99

[jit] Tree based Iterator infrastructure: for in range, list and tensor,

2a67807

for in enumerate and zip

pytorchbot added oncall: jit Add this issue/PR to JIT oncall triage queue module: internals Related to internal abstractions in c10 and ATen labels Jun 14, 2019

wanchaol changed the title ~~[jit] Tree based Iterator infrastructure: for in range, list and tensor,~~ [jit] Tree based Iterator infrastructure: for in range/list/tensor/zip/enumerate Jun 14, 2019

wanchaol requested review from Krovatkin, eellison, suo and zdevito June 14, 2019 19:47

This was referenced Jun 14, 2019

[jit] Add for in enumerate to script #21666

Closed

[jit] New Iterable interface, refactor for loops (List, Tensor, Range) #21403

Closed

lint on "[jit] Tree based Iterator infrastructure: for in range/list/…

d1e0d41

…tensor/zip/enumerate" [jit] Tree based Iterator infrastructure: for in range, list and tensor, for in enumerate and zip gh-metadata: pytorch pytorch 21801 gh/wanchaol/18/head

zdevito reviewed Jun 17, 2019

View reviewed changes

suo removed their request for review June 17, 2019 20:16

Update on "[jit] Tree based Iterator infrastructure: for in range/lis…

ad9cfeb

…t/tensor/zip/enumerate" [jit] Tree based Iterator infrastructure: for in range, list and tensor, for in enumerate and zip gh-metadata: pytorch pytorch 21801 gh/wanchaol/18/head

This was referenced Jun 19, 2019

[jit] more complete tuple assignments #21946

Closed

[jit] refactor all for in range/tensor tests to be together with other for loop tests #21947

Closed

pytorchbot added the module: pybind Related to our Python bindings / interactions with other Python libraries label Jun 19, 2019

Update on "[jit] Tree based Iterator infrastructure: for in range/lis…

712f920

…t/tensor/zip/enumerate" [jit] Tree based Iterator infrastructure: for in range, list and tensor, for in enumerate and zip gh-metadata: pytorch pytorch 21801 gh/wanchaol/18/head

This was referenced Jun 19, 2019

[jit] more complete tuple assignments #21949

Closed

[jit] refactor all for in range/tensor tests to be together with other for loop tests #21950

Closed

Update on "[jit] Tree based Iterator infrastructure: for in range/lis…

2999f5d

…t/tensor/zip/enumerate" [jit] Tree based Iterator infrastructure: for in range, list and tensor, for in enumerate and zip gh-metadata: pytorch pytorch 21801 gh/wanchaol/18/head

wanchaol requested a review from zdevito June 19, 2019 04:43

wanchaol added 3 commits June 18, 2019 22:13

add for target tuple assign on "[jit] Tree based Iterator infrastruct…

c1e2f36

…ure: for in range/list/tensor/zip/enumerate" [jit] Tree based Iterator infrastructure: for in range, list and tensor, for in enumerate and zip gh-metadata: pytorch pytorch 21801 gh/wanchaol/18/head

use builtin on "[jit] Tree based Iterator infrastructure: for in rang…

1afa0a8

…e/list/tensor/zip/enumerate" [jit] Tree based Iterator infrastructure: for in range, list and tensor, for in enumerate and zip gh-metadata: pytorch pytorch 21801 gh/wanchaol/18/head

Update on "[jit] Tree based Iterator infrastructure: for in range/lis…

e1cf53f

…t/tensor/zip/enumerate" [jit] Tree based Iterator infrastructure: for in range, list and tensor, for in enumerate and zip gh-metadata: pytorch pytorch 21801 gh/wanchaol/18/head

This was referenced Jun 19, 2019

[jit] support iteration tuple unpacking #21985

Closed

[jit] add for in string support #21990

Closed

Update on "[jit] Tree based Iterator infrastructure: for in range/lis…

31004d8

…t/tensor/zip/enumerate" [jit] Tree based Iterator infrastructure: for in range, list and tensor, for in enumerate and zip gh-metadata: pytorch pytorch 21801 gh/wanchaol/18/head

wanchaol mentioned this pull request Jun 20, 2019

[jit] add for in dict support #22006

Closed

rebase with master on "[jit] Tree based Iterator infrastructure: for …

71fb8bf

…in range/list/tensor/zip/enumerate" [jit] Tree based Iterator infrastructure: for in range, list and tensor, for in enumerate and zip gh-metadata: pytorch pytorch 21801 gh/wanchaol/18/head

zdevito approved these changes Jun 21, 2019

View reviewed changes

wanchaol added 2 commits June 21, 2019 12:09

rebase with master on "[jit] Tree based Iterator infrastructure: for …

cdfc483

…in range/list/tensor/zip/enumerate" [jit] Tree based Iterator infrastructure: for in range, list and tensor, for in enumerate and zip gh-metadata: pytorch pytorch 21801 gh/wanchaol/18/head

rebase with master on "[jit] Tree based Iterator infrastructure: for …

30ca7bf

…in range/list/tensor/zip/enumerate" [jit] Tree based Iterator infrastructure: for in range, list and tensor, for in enumerate and zip gh-metadata: pytorch pytorch 21801 gh/wanchaol/18/head

Krovatkin approved these changes Jun 21, 2019

View reviewed changes

Krovatkin mentioned this pull request Jun 21, 2019

[WIP] for x,y in zip(a,b) #21039

Closed

wanchaol added 2 commits June 21, 2019 21:41

update comments on "[jit] Tree based Iterator infrastructure: for in …

db261d5

…range/list/tensor/zip/enumerate" [jit] Tree based Iterator infrastructure: for in range, list and tensor, for in enumerate and zip gh-metadata: pytorch pytorch 21801 gh/wanchaol/18/head

Update on "[jit] Tree based Iterator infrastructure: for in range/lis…

1ebef11

…t/tensor/zip/enumerate" [jit] Tree based Iterator infrastructure: for in range, list and tensor, for in enumerate and zip gh-metadata: pytorch pytorch 21801 gh/wanchaol/18/head

pytorchbot added the module: third_party label Jun 22, 2019

Update on "[jit] Tree based Iterator infrastructure: for in range/lis…

ab53487

…t/tensor/zip/enumerate" [jit] Tree based Iterator infrastructure: for in range, list and tensor, for in enumerate and zip gh-metadata: pytorch pytorch 21801 gh/wanchaol/18/head

wanchaol removed the module: third_party label Jun 22, 2019

facebook-github-bot closed this in e0f5ab2 Jun 22, 2019

zou3519 deleted the gh/wanchaol/18/head branch June 22, 2019 08:03

facebook-github-bot added the merged label Jun 22, 2019

mruberry added the Merged label Oct 28, 2020

[jit] Tree based Iterator infrastructure: for in range/list/tensor/zip/enumerate #21801

[jit] Tree based Iterator infrastructure: for in range/list/tensor/zip/enumerate #21801

Uh oh!

Conversation

wanchaol commented Jun 14, 2019 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

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

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!

softagram-bot commented Jun 17, 2019

Softagram Impact Report for pull/21801 (head commit: d1e0d41)

⭐ Change Overview

⭐ Details of Dependency Changes

📄 Full report

Uh oh!

suo commented Jun 17, 2019

Uh oh!

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

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

wanchaol commented Jun 14, 2019 •

edited

Loading

Softagram Impact Report for pull/21801 (head commit: `d1e0d41`)