[Dev API] Possible to remove AST nodes?

Questions
#1

Hello all,

I have found myself having to do some AST manipulation (via a custom ir_pass).
Reading the Writing a Customized Pass Tutorial, I know that I can manipulate the AST via the tvm.ir_pass.IRTransform(input_stmt, preorder_func, postorder_func,[list_of_AST_nodes]) construct.
Replacing an AST node is quite easy but I have no clue how to remove a node

define postorder_func(stmt):
    irb = tvm.ir_builder.create()
    if some_condition:
        #replace the AST node
        irb.emit(tvm.stmt.SomeStatementType(...))
        return irb.get()
    else:
        #remove the AST node
        irb.emit(...)
        return irb.get() #not returning a value doesn't help eiter

All the tvm.stmt types that I can find are “non-NOP” and therefore generate an AST node.
I tried tvm.stmt.Eval(0) but I obviously get a 0 printed in my code (which I can deal with but I’d rather have it removed).

Is there something completely obvious that I am missing?

Thanks

2 Likes
A CallNode has relay ir but the primfunc has gone
#2

Because all the objects are immutable, we cannot directly mutate a node. Instead, in order to remove something, we might want to simply return an Evaluate(0)(which is a nop) in post order transform

1 Like
#3

Thanks for the reply.

Yeah I guess code generators will need to catch tvm.stmt.Evaluate(0) and interpret them as to being NOPs or actually skippable code (I guess dependening at what level you work at, a NOP is not really a skippable instruction ex: assembly). I think this is what you meant, right?
It just feels like its an assumption that is not really documented/exposed.

  • What would the interpretation of an tvm.stmt.Evaluate(some_int_other_than_0) be?

Adding a new tvm.stmt.NOP or allowing a node of the AST to be deleted, I guess would be a clear way to signalize our intent.