[JIT] Auto-batching IR transformation for control flow by ChunliF · Pull Request #9392 · pytorch/pytorch

ChunliF · 2018-07-12T17:34:13Z

Implement IR transformation for control flow

prim::Constant: clone to new graph directly
prim::NumToTensor: create a BatchTensor from output tensor with batch_size = 1
prim::TensorToNum: clone to new graph
prim::ListConstruct: clone to new graph
prim::If: execute both if_block and else_block and combine results from them using cond
prim::Loop:
- for loop
- while loop: change while cond to cond_any, use cond to update outputs

test case: hand-written LSTM, greedy search, beam search

facebook-github-bot

@ChunliF has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

jamesr66a

Looking good! A couple overall comments:

A lot of places in the code you have a literal 3 around to represent the size of the expanded Values. It might be better to factor that out as a constant so if the number of values used to represent a batch changes we can just change that in 1 place. In general I think it's safe to assume this is going to be a 1 -> N value transform for all values, so factoring it out as a constant might help
It would probably also be good to add some calls to self.assertExpected and pass in the batched graphs in the test cases. This way, we can manually inspect the outputs of the graph in review. That also has the nice property of failing the test if we pass it "on accident" (right answer with wrong method)

torch/csrc/jit/passes/to_batch.cpp

+  }
+}
+
+// clone prim::Constant to new graph


torch/csrc/jit/passes/to_batch.cpp

-void ToBatch::toBatch(Block* block, Block* res_block) {
-  // change inputs of a graph - expand tensor to {data, mask, dims}
+// eg: "a.1" -> {"a", "1"}; "a" -> {"a"}
+std::vector<std::string> ToBatch::get_name(std::string name) {


torch/csrc/jit/passes/to_batch.cpp

+// replace aten operator node with BatchTensor operator graph
+void ToBatch::visitAten(Node* n, Block* block, Block* res_block, std::unordered_map<std::string, Value*>& var_map){
+  if(n->outputs().size() > 1){
+    throw std::runtime_error("Cannot process multiple assignment");


torch/csrc/jit/passes/to_batch.cpp

+  // do update on assignment
+  auto name_base = get_name(n->output()->uniqueName())[0];
+  if(var_map.find(name_base) != var_map.end()){
+    std::vector<Value*> inputs(batch_map.at(var_map.at(name_base)));


torch/csrc/jit/passes/to_batch.cpp

+// elif is not supported
+// assume every variable assigned in an if statement is already defined before
+void ToBatch::visitIf(Node* n, Block* block, Block* res_block, std::unordered_map<std::string, Value*>& var_map){
+  auto res_graph = res_block->owningGraph();


torch/csrc/jit/passes/to_batch.cpp

+
+// elif is not supported
+// assume every variable assigned in an if statement is already defined before
+void ToBatch::visitIf(Node* n, Block* block, Block* res_block, std::unordered_map<std::string, Value*>& var_map){


torch/csrc/jit/passes/to_batch.cpp

+  auto res_graph = res_block->owningGraph();
+
+  // create prim::If node for res_block
+  auto add_if_node = [&](Block* block, std::shared_ptr<Graph> cond_graph, std::vector<Value*> cond, std::vector<Value*> unchanged_outputs){


…ifications

facebook-github-bot

@ChunliF has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

zdevito

Nice progress! I see some issues involving update statements that I put in the comments. Let me know if you have any questions about it.

torch/csrc/jit/passes/to_batch.cpp

+  }
+}
+
+void ToBatch::toBatch(Block* block, Block* res_block, std::unordered_map<std::string, Value*>& upper_var_map) {


torch/csrc/jit/passes/to_batch.cpp

+      visitAten(n, block, res_block, var_map);
+    }
+    else if(n->kind().is_prim()){
+      if(n->kind() == prim::Constant){


torch/csrc/jit/passes/to_batch.cpp

-    res_block->registerOutput(r_output[1]);
-    res_block->registerOutput(r_output[2]);
+  // change outputs of block - expand tensor to batchtensor(data, mask, dims)
+  // for block in prim::Loop, register outputs separately


torch/csrc/jit/passes/to_batch.cpp

+  }
+  auto outputs = script::inlineCallTo(*res_block->owningGraph(), *batch_graph, new_inputs);
+
+  // do update on assignment


test/expect/TestBatched.test_numToTensor.expect

+graph(%a.1_data : Dynamic
+      %a.1_mask : Dynamic
+      %a.1_dims : Dynamic) {
+  %3 : int = prim::Constant[value={1}]()


test/expect/TestBatched.test_if_else.expect

+  %8 : Dynamic = aten::__or__(%a.1_dims, %b_dims)
+  %9 : Dynamic = aten::mul(%6, %7)
+  %10 : Dynamic = aten::sum(%9)
+  %11 : Dynamic = aten::gt[other={0}](%10)


test/test_jit.py

+
+        @torch.jit.script
+        def batch_for(x, y):
+            for _i in range(10):


torch/jit/batchop.py

+
+@torch.jit.script
+def batch_argmax(data, mask, dims, dim, keepdim):
+    # if dim == 0:


old

facebook-github-bot

@ChunliF has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

jamesr66a

Looking good! @zdevito can you please have a look?

test/test_jit.py

+        graph = torch.to_batch_graph(batch_for.graph)
+        self.assertExpected(str(graph))
+
+    def test_lstm(self):


test/test_jit.py

+        def greedy(x, h, c, embed, w_xi, w_xf, w_xo, w_xc, w_hi, w_hf, w_ho, w_hc,
+                   b_i, b_f, b_o, b_c, w_hs, b_s, iter_num):
+            iter_count = torch.zeros_like(iter_num)
+            while(iter_count < iter_num):


torch/csrc/jit/passes/to_batch.cpp

-  auto size = block->inputs().size();
-  for(size_t i = 0; i < size; i++){
-    auto input = block->inputs()[i];
+std::shared_ptr<Graph> ToBatch::getBatchOperator(std::string name, int64_t input_num){


test/test_jit.py

        def batch_if(a, b):
            if a > b:
-                a += b
+                a = a + b


test/test_jit.py

+                s_t = s_t.view([1, -1])
+                p_t = torch.softmax(s_t, 1)
+                # print(p_t)
+                prob_t, i_t = torch.topk(p_t, k, 1)


facebook-github-bot

@ChunliF has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

ChunliF · 2018-07-27T02:30:57Z

@pytorchbot retest this please

facebook-github-bot

ChunliF has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

ChunliF · 2018-07-31T00:38:57Z

@pytorchbot retest this please

ChunliF · 2018-07-31T17:48:12Z

@pytorchbot retest this please

zdevito

I didn't do a detailed review. But I think we should merge this, assuming its tests pass so that it doesn't get broken by future jit changes.

facebook-github-bot

ChunliF has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

ChunliF · 2018-08-01T19:02:09Z

@pytorchbot retest this please

facebook-github-bot

ChunliF has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

Summary: Implement IR transformation for control flow - `prim::Constant`: clone to new graph directly - `prim::NumToTensor`: create a `BatchTensor` from output tensor with `batch_size = 1` - `prim::TensorToNum`: clone to new graph - `prim::ListConstruct`: clone to new graph - `prim::If`: execute both `if_block` and `else_block` and combine results from them using `cond` - `prim::Loop`: - for loop - while loop: change while `cond` to `cond_any`, use `cond` to update outputs test case: hand-written LSTM, greedy search, beam search Pull Request resolved: pytorch#9392 Differential Revision: D8822369 Pulled By: ChunliF fbshipit-source-id: 8f03c95757d32e8c4580eeab3974fd1bc429a1e5

batchtensor ir transformation for control flow

ce22a0d

ChunliF requested review from apaszke, colesbury, ezyang, gchanan, soumith and zdevito as code owners July 12, 2018 17:34

facebook-github-bot reviewed Jul 12, 2018

View reviewed changes

jamesr66a previously requested changes Jul 12, 2018

View reviewed changes

add assertExpected, change 3 to constant, add more comments, some mod…

6abf8af

…ifications

facebook-github-bot reviewed Jul 12, 2018

View reviewed changes

add for loop support

d840a83

zdevito reviewed Jul 13, 2018

View reviewed changes

ChunliF added 3 commits July 16, 2018 15:14

add lstm user case, support attributes for nodes

576bcb9

fix

c53133b

add topk, softmax, argmax, fix for loop

1c6a6dd

jamesr66a reviewed Jul 20, 2018

View reviewed changes

test/expect/TestBatched.test_numToTensor.expect Outdated

graph(%a.1_data : Dynamic

%a.1_mask : Dynamic

%a.1_dims : Dynamic) {

%3 : int = prim::Constant[value={1}]()

This comment was marked as off-topic.

Sign in to view

jamesr66a reviewed Jul 20, 2018

View reviewed changes

test/expect/TestBatched.test_if_else.expect Outdated

%8 : Dynamic = aten::__or__(%a.1_dims, %b_dims)

%9 : Dynamic = aten::mul(%6, %7)

%10 : Dynamic = aten::sum(%9)

%11 : Dynamic = aten::gt[other={0}](%10)

This comment was marked as off-topic.

Sign in to view

jamesr66a reviewed Jul 20, 2018

View reviewed changes

test/test_jit.py Outdated

@torch.jit.script

def batch_for(x, y):

for _i in range(10):

This comment was marked as off-topic.

Sign in to view

jamesr66a reviewed Jul 20, 2018

View reviewed changes

torch/jit/batchop.py

@torch.jit.script

def batch_argmax(data, mask, dims, dim, keepdim):

# if dim == 0:

This comment was marked as off-topic.

Sign in to view

zou3519 added the oncall: jit Add this issue/PR to JIT oncall triage queue label Jul 23, 2018

Merge https://github.com/pytorch/pytorch into control-flow

2af6e62

weiyangfb added the ready for review (this tag is deprecated) All PRs are ready for review unless they are draft, WIP, or have undismissed requested changes label Jul 24, 2018

ChunliF added 2 commits July 24, 2018 17:34

Merge https://github.com/pytorch/pytorch into control-flow

0e5b3e1

rebase based on lvalue pr, add beam search test case, add some operators

be41bd2

facebook-github-bot reviewed Jul 25, 2018

View reviewed changes

fix stray onnx-tensorrt update

293895b

jamesr66a reviewed Jul 25, 2018

View reviewed changes

fix

4fdec27

facebook-github-bot reviewed Jul 26, 2018

View reviewed changes

ChunliF and others added 3 commits July 27, 2018 14:27

rebased based on removing attr pr

8f188b1

Merge branch 'master' into control-flow

bf54541

fix merge

d02d238

facebook-github-bot reviewed Jul 30, 2018

View reviewed changes

ChunliF added 3 commits July 30, 2018 11:43

Merge https://github.com/pytorch/pytorch into control-flow

111a4a5

Merge https://github.com/pytorch/pytorch into control-flow

c1dfb0b

attempt to fix expect file fail

c7f73dd

ChunliF added 2 commits July 31, 2018 11:07

Merge https://github.com/pytorch/pytorch into control-flow

a467789

rebase based on add tensorlist pr

0b438e7

zdevito approved these changes Jul 31, 2018

View reviewed changes

facebook-github-bot reviewed Jul 31, 2018

View reviewed changes

fix merge conflict

154e4eb

facebook-github-bot reviewed Aug 2, 2018

View reviewed changes

facebook-github-bot closed this in 57061d6 Aug 2, 2018

ezyang added the merged label Jun 26, 2019

Conversation

ChunliF commented Jul 12, 2018 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

facebook-github-bot left a comment

Choose a reason for hiding this comment

Uh oh!

jamesr66a left a comment

Choose a reason for hiding this comment

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

facebook-github-bot left a comment

Choose a reason for hiding this comment

Uh oh!

zdevito left a comment

Choose a reason for hiding this comment

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

facebook-github-bot left a comment

Choose a reason for hiding this comment

Uh oh!

jamesr66a left a comment

Choose a reason for hiding this comment

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

facebook-github-bot left a comment

Choose a reason for hiding this comment

Uh oh!

ChunliF commented Jul 27, 2018

Uh oh!

facebook-github-bot left a comment

Choose a reason for hiding this comment

Uh oh!

ChunliF commented Jul 31, 2018

Uh oh!

ChunliF commented Jul 12, 2018 •

edited

Loading