Hi All,
Thanks @masahi, @comaniac and @matt-arm. It worked now. That being said, I'd
like to confirm 1) if I understand the functionality of PartitionGraph()
function in the relay, 2) I'd like to understand if ParitioGraph() can be used
for my specific use case.
Here is my understanding of how PartitionGraph() function works:
* List item PartitionGraph() function partitions the IRModule into functions
based on the annotations.
* List item Annotations are achieved by adding attributes of the same kind
operators (e.g., add, subtract) by function AnnotateTarget() function.
Here is my main question:
- Annotations are done per operator kind such as "add", not operator instances.
For example, if I have two "add" operators in my true and false branch, and I'd
like to partition the true and false branches separately, can PartitionGraph()
can help me? I know I can overwrite visit_if() function in ExprMutator class
to achieve what I just described, but I am looking for a more high-level
solutions for more complex problems.
To me, it looks like ParitioGraph() seems limited because it partitions based
on annotations that are attached to per operator kind.
Ideally, I'd like to have a solution, that does the followings:
- Partition relay IRModule based on user-provided annotations on expressions
into separate Relay IR functions (or IRModules)
After mod = tvm.IRModule({"main": f})
```
print(mod)
def @main(%c, %x: Tensor[(10, 1), float32], %y: Tensor[(10, 1), float32], %x1:
Tensor[(10, 1), float32], %y1: Tensor[(10, 1), float32]) {
if (%c) {
add(%x, %y)
} else {
multiply(%x1, %y1)
}
}
```
After annotation: mod = relay.transform.AnnotateTarget(["special"])(mod)
```
print(mod)
def @main(%c: bool, %x: Tensor[(10, 1), float32], %y: Tensor[(10, 1), float32],
%x1: Tensor[(10, 1), float32], %y1: Tensor[(10, 1), float32]) -> Tensor[(10,
1), float32] {
%0 = annotation.compiler_begin(%c, meta[relay.attrs.CompilerAttrs][0]) /*
ty=bool */;
%9 = if (%0) {
%1 = annotation.compiler_begin(%x, meta[relay.attrs.CompilerAttrs][1]) /*
ty=Tensor[(10, 1), float32] */;
%2 = annotation.compiler_begin(%y, meta[relay.attrs.CompilerAttrs][2]) /*
ty=Tensor[(10, 1), float32] */;
%3 = add(%1, %2) /* ty=Tensor[(10, 1), float32] */;
%4 = annotation.compiler_end(%3, meta[relay.attrs.CompilerAttrs][3]) /*
ty=Tensor[(10, 1), float32] */;
annotation.compiler_begin(%4, meta[relay.attrs.CompilerAttrs][4]) /*
ty=Tensor[(10, 1), float32] */
} else {
%5 = annotation.compiler_begin(%x1, meta[relay.attrs.CompilerAttrs][5]) /*
ty=Tensor[(10, 1), float32] */;
%6 = annotation.compiler_begin(%y1, meta[relay.attrs.CompilerAttrs][6]) /*
ty=Tensor[(10, 1), float32] */;
%7 = multiply(%5, %6) /* ty=Tensor[(10, 1), float32] */;
%8 = annotation.compiler_end(%7, meta[relay.attrs.CompilerAttrs][7]) /*
ty=Tensor[(10, 1), float32] */;
annotation.compiler_begin(%8, meta[relay.attrs.CompilerAttrs][8]) /*
ty=Tensor[(10, 1), float32] */
};
annotation.compiler_end(%9, meta[relay.attrs.CompilerAttrs][9]) /*
ty=Tensor[(10, 1), float32] */
}
```
After mod = relay.transform.PartitionGraph()(mod)
```
def @special_0(%special_0_i0: Tensor[(10, 1), float32], %special_0_i1:
Tensor[(10, 1), float32], global_symbol="special_0", Primitive=1,
Compiler="special", Inline=1) -> Tensor[(10, 1), float32] {
add(%special_0_i0, %special_0_i1) /* ty=Tensor[(10, 1), float32] */
}
def @main(%c: bool, %x: Tensor[(10, 1), float32], %y: Tensor[(10, 1), float32],
%x1: Tensor[(10, 1), float32], %y1: Tensor[(10, 1), float32]) -> Tensor[(10,
1), float32] {
if (%c) {
@special_0(%x, %y) /* ty=Tensor[(10, 1), float32] */
} else {
multiply(%x1, %y1) /* ty=Tensor[(10, 1), float32] */
}
}
```
---
[Visit
Topic](https://discuss.tvm.apache.org/t/understanding-tvm-relays-partitiongraph-mod-function/8290/6)
to respond.
You are receiving this because you enabled mailing list mode.
To unsubscribe from these emails, [click
here](https://discuss.tvm.apache.org/email/unsubscribe/7ce3e2cf8831eaba26741a1c2c88a49f727aabfad7968c49444fdd7248dab810).
