How to traverse HalideIR by tvm C++ API

Application
#1

The topic about Where can I find any c++ sample? will be helpful to someone need to traverse IR node in TVM, so I think it’s necessary to create a new topic.
Thanks a lot for @wweic

#2
#3

Continuing the discussion from How to traverse HalideIR by tvm C++ API:

@merrymercy Thanks for your reply
I have used IRVisitor to traverse all IR node, I don’t know what’s the best position to pass a parameter to IRVisitor, so I have to modify build_module.cc, the following is some code snippet about how do I traverse all IR node.


Stmt BuildStmt(Schedule sch,
    const Array<Tensor>& args,
    const std::unordered_map<Tensor, Buffer>& binds,
    bool loop_partition,
    Array<NodeRef> *out_arg_list,
    const BuildConfig& config) {
    Map<Tensor, Buffer> out_binds;
    GetBinds(args, binds, &out_binds, out_arg_list, config);
    
    sch = sch.normalize();
    
    // Phase 0
    auto bounds = schedule::InferBound(sch);
    auto stmt = schedule::ScheduleOps(sch, bounds, false);
        
    /* -- Traverse code begin -- */
    // HalideIRCollector inherit from IRVisitor
    HalideIRCollector hIRC(stmt);  // stmt is the traverse entry
    hIRC.RunVisit(filepath);       // filepath is used to store information of ir node
    /* -- Traverse code end -- */
    
    stmt = ir::InjectPrefetch(stmt);
    
    // Phase 1
    stmt = ir::StorageFlatten(stmt, out_binds, 64);
    stmt = ir::CanonicalSimplify(stmt);

Is this the correct way to traverse all of IR nodes?
How can I traverse all IR nodes in my source code instead of modifying some code of TVM?

#4

In python api, you can add your own pass. See tutorial here
https://docs.tvm.ai/tutorials/dev/low_level_custom_pass.html#sphx-glr-tutorials-dev-low-level-custom-pass-py

In c++ api, you can add your pass here (the doc string is wrong)

#5

Thanks, we don’t use python API, is there any c++ api sample?