Re-enabled mm+add tree batching in the JIT #13228
Conversation
torch/csrc/jit/passes/batch_mm.cpp
Outdated
| batch_mm->insertBefore(root.node); | ||
| root.node->output()->replaceAllUsesWith(batch_mm->output()); | ||
| auto matmuls = root.removeTransposesAndGatherMatmuls(); | ||
| WithInsertPoint insert_gurad {root.node}; |
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
|
|
||
| RegisterOperators mm_tree_reduction_reg({ | ||
| Operator( | ||
| Symbol::prim("MMTreeReduce"), |
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
There was a problem hiding this comment.
@zou3519 has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.
| } | ||
|
|
||
| bool shape_is_fast(const at::Tensor& lhs, const at::Tensor& rhs) { | ||
| size_t l = lhs.size(0); |
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
There was a problem hiding this comment.
@zou3519 has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.
I've had to generously increase the range of the CreateADSubgraphs pass, because even though it collapses the RNN loop to a single differentiable subgraphs and a few other nodes, the range uses the distances in the original graph...
cc @zdevito @zou3519