Refactor Dynamic to Static

[mbrookhart]

I recently spent a lot of time fighting dynamic rank issues in a kind of crazy ONNX model. Fixing it required doing incremental dynamic-to-static before type inference. This PR basically changes the logic of dynamic to static from this:

infer type on the whole function
constant fold the whole function
then go find dynamic ops who's inputs are constant
    replace them with static ops
Rinse and repeat

to this:

go find dynamic ops
    infer type and constant fold their inputs
    If the inputs are now constant, replace them with static ops

This has the advantage that it can analyze those crazy dynamic rank and control flow graphs and simplify them, but it has the disadvantage that it's slower than the previous version because we call infertype and constant folding many more times.

Performance checking shows that this takes a BERT compile from ~15 seconds to ~60 seconds. This should be fixable when incremental type inference becomes available.

Thanks,
Matthew

cc @masahi @jroesch @tmoreau89 @jwfromm @electriclilies

[masahi]

Since you always run PrepareArgs when you find a dynamic op, I'd run PrepareArgs here

tvm/src/relay/transforms/dynamic_to_static.cc

Line 209 in 384714b

auto out = op_map_[call_node->op](call_node);

[masahi]

Can we assume that compile time regression is the worst for BERT? I don't recall infer type or fold constant being slow on other models.

[masahi]

cc @t-vi who might be interested in incremental type inference

[t-vi]

cc @t-vi who might be interested in incremental type inference

He, yeah, I had hoped you fixed it. 😉 I think @jroesch had looked into it more than I did at some point (in the context of #6274).

My impression is that part of the difficulty is that in-place graph operations are not a good fit with how things work in TVM in general and the frequent copying we do removes the type info. If memory serves me well, this was the main reason for doing the incremental type inference "manually" in the PyTorch frontend.

[mbrookhart]

Since you always run PrepareArgs when you find a dynamic op, I'd run PrepareArgs here

tvm/src/relay/transforms/dynamic_to_static.cc

Line 209 in 384714b

auto out = op_map_[call_node->op](call_node);

I tried this early on, unfortunately, PrepareArgs ends up making a copy of the IR to do type inference, and then we end up with two different versions of input variables depending on whether or not the op that uses them has a dynamic op before it or not, this breaks several unit tests.

To fix it, I would need to do infer_type/constant folding on every op during traversal, but without incremental type inference, that's impossibly slow. This is a middle ground that fixes the problem without too much of a performance hit.

[mbrookhart]

@t-vi I feel your pain, a few people in OctoML are looking at a possible rewrite of the type inferencer in the coming months to fix some of these issues.

[mbrookhart]

Can we assume that compile time regression is the worst for BERT? I don't recall infer type or fold constant being slow on other models.

The worst model I've seen with this pass is ONNX SSD-Mobilenet, which takes about 3 minutes and prompted all of the dynamic rank fixes.

[masahi]

thanks @mbrookhart