This PR aims to support inplace op it TE partially. This can enable TE fusion for patterns like "at::conv, at::relu".
"at::add, at::relu, at::add". We can get better performance by this.
Options
Option 1: Replace the in-place operator with the out-place operator
Pros: Easy to implement. Can re-use "RemoveTensorMutation" for replacement and the lowing function of outplace version. Will not have a performance penalty under "fusion" cases.
Cons: The fallback path is also outplaced version, which will have a performance penalty here. Sometimes we cannot safely replace the inplace op with its' outplace version.
Option 2: Lower the body in terms of in-place directly
Pros: More straightforward. Can support more scenarios.
Cons: Hard to develop. More engineer effort.
We choose option 1 in this PR, and we can consider option 2 if we observed that in many real-world scenarios, options1 failed.
Implement Details
We will pull inplace ops into TE fusion groups if the ops satisfy 2 conditionals:
This inplace ops can be safely replaced with its outplace version at the global graph. (Re-use the ability of RemoveTensorMutation)
The outplace version is supported by TE.
To achieve this, we extend the behavior of Operator supported check and TryMerge.
In Operator supported check.
For example, isSupported. If an inplace op can be safely replaced with its outplace version. We will create an outplace node and send this node to pass the TE checks. After the checks are done, destroy the outplace node.
We choose to make a outplace node to pass the TE checks because we do not want directly added inplace op into the supported op list. Te does not directly support these inplace ops.
We will destroy the outplace node after Operator supported check is done because Operator supported check is not all checks for TE fusion, there are many following checks later and we do not know whether all following check will pass at this stage.
In TryMerge.
Here we can know all checks are passed, we will replace an inplace op with its outplace version before the moment we merge the node into fusion groups.
Some details for safely replacement
Whether an inplace op can be replaced safely depends on the behavior of RemoveTensorMutions.
2 cases below will not be replaced.
def fn(a, b):
return a.relu_() + b.relu()
def fn(a, b):
c = a + b
return c.sigmoid().add(c.relu_())
Ops like __and__/__or__. Their inplace version are __iand__/__ior__. These ops cannot be replaced by RemoveTensorMutaion. So we do not support them yet.
Unit test details
Add 3 CPP tests, 1 should be replaced (and fused into TE fusion groups), and 2 should not be replaced (because they cannot be safely replaced.
Extend test_unary_ops, test_binary_ops, and test_ternary_ops in test_jit_fuser_te.py to include the inplace version, all should be fused by replacing inplace op with outplace version.
Summary
This PR aims to support inplace op it TE partially. This can enable TE fusion for patterns like "at::conv, at::relu". "at::add, at::relu, at::add". We can get better performance by this.
Options
Option 1: Replace the in-place operator with the out-place operator
Option 2: Lower the body in terms of in-place directly
We choose option 1 in this PR, and we can consider option 2 if we observed that in many real-world scenarios, options1 failed.
Implement Details
We will pull inplace ops into TE fusion groups if the ops satisfy 2 conditionals:
To achieve this, we extend the behavior of Operator supported check and TryMerge.
In Operator supported check.
For example, isSupported. If an inplace op can be safely replaced with its outplace version. We will create an outplace node and send this node to pass the TE checks. After the checks are done, destroy the outplace node.
In TryMerge.
Here we can know all checks are passed, we will replace an inplace op with its outplace version before the moment we merge the node into fusion groups.
Some details for safely replacement
Whether an inplace op can be replaced safely depends on the behavior of RemoveTensorMutions. 2 cases below will not be replaced.
Ops like __and__/__or__. Their inplace version are __iand__/__ior__. These ops cannot be replaced by RemoveTensorMutaion. So we do not support them yet.
Unit test details
Since only element-wise ops will own outplace version, we show the Op coverage status based on this TE support list https://github.com/pytorch/pytorch/blob/master/torch/csrc/jit/runtime/symbolic_shape_registry_util.cpp#L11.
UT covered ops
Element-wise ops with outplace support and without inplace ops support: