comaniac opened a new issue #5662:
URL: https://github.com/apache/incubator-tvm/issues/5662
In #5656, we found that `pattern.partition` will not lift the bind constant
nodes to the partitioned function arguments. This results in argument mismatch
and could be a potential problem when applying to op fusion.
Here is an illustration example:
```python
import tvm
from tvm import relay
from tvm.relay.dataflow_pattern import *
from tvm.relay.build_module import bind_params_by_name
import numpy as np
x = relay.var('x', shape=(1, 3, 224, 224))
w = relay.var('w', shape=(3, 3, 3, 3))
b = relay.var('b', shape=(3,))
conv2d = relay.op.nn.conv2d(x, w)
out = relay.op.nn.bias_add(conv2d, b)
func = relay.Function([x, w, b], out)
mod = tvm.IRModule.from_expr(func)
mod["main"] = bind_params_by_name(mod["main"],
{'w': tvm.nd.array(np.ones(shape=(3, 3, 3,
3)))})
print('=== Fuse ====')
print(relay.transform.FuseOps()(mod))
conv2d = is_op('nn.conv2d')(wildcard(), wildcard())
pattern = is_op('nn.bias_add')(conv2d, wildcard())
print('=== Partition ===')
print(pattern.partition(mod['main'].body, {'Composite': 'aa'}))
```
Output:
```
=== Fuse ====
def @main(%x: Tensor[(1, 3, 224, 224), float32], %b: Tensor[(3), float32])
-> Tensor[(1, 3, 222, 222), float32] {
%1 = fn (%p0: Tensor[(1, 3, 224, 224), float32], %p1: Tensor[(3, 3, 3, 3),
float64], %p2: Tensor[(3), float32], Primitive=1) -> Tensor[(1, 3, 222, 222),
float32] {
%0 = nn.conv2d(%p0, %p1, padding=[0, 0, 0, 0]) /* ty=Tensor[(1, 3, 222,
222), float32] */;
nn.bias_add(%0, %p2) /* ty=Tensor[(1, 3, 222, 222), float32] */
};
%1(%x, meta[relay.Constant][0] /* ty=Tensor[(3, 3, 3, 3), float64] */ /*
ty=Tensor[(3, 3, 3, 3), float64] */, %b) /* ty=Tensor[(1, 3, 222, 222),
float32] */
}
// meta data omitted. you can use show_meta_data=True to include meta data
=== Partition ===
free_var %x: Tensor[(1, 3, 224, 224), float32]
free_var %b: Tensor[(3), float32]
%1 = fn (%FunctionVar_0_0, %FunctionVar_0_1, Composite="aa",
PartitionedFromPattern="nn.conv2d_nn.bias_add_") {
%0 = nn.conv2d(%FunctionVar_0_0, meta[relay.Constant][0] /* ty=Tensor[(3,
3, 3, 3), float64] */ /* ty=Tensor[(3, 3, 3, 3), float64] */, padding=[0, 0, 0,
0]);
nn.bias_add(%0, %FunctionVar_0_1)
};
%1(%x, %b)
// meta data omitted. you can use show_meta_data=True to include meta data
```
We can see that the function generated by the op fusion keeps the original
arguments and refers to the constant node in the function call. However, the
partitioned function directly accesses the constant node from inside of the
function body. Ideally, the partitioned should be same as the fused function.
cc @mbrookhart @masahi @zhiics
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
