tvm.ir_pass.IRTransform doesn't recurse when pre_order makes a change


#1

I’m trying to use tvm.ir_pass.IRTransform to do a pre-order (top-down) transformation. The problem I’m encountering is that IRTransform doesn’t recurse anymore when pre-order function performs a transformation (see https://github.com/dmlc/tvm/blob/master/src/pass/ir_mutator.cc#L30-L46) . Is this design on purpose?

One alternative design is to change it to something like this:

    if (f_preorder_ != nullptr) {
      T pre = f_preorder_(node);
      if (pre.defined()) {
         node = pre;
      }
    }
    node = IRMutator::Mutate(node);

#2

This is on purpose. Use post order if you want recursive effect


#3

Thanks @tqchen for quick response.
If I use post-order I’ll see (and need to process) the nodes in bottom up fashion while what I want is to process the nodes top-down (i.e., preorder). In other words, for my transformation algorithm the order of traversal matters. Can you think of a way to provide this functionality with minor modification or should I write a traversal from scratch?