Retrieve AST of a network

trusira · November 28, 2018, 10:41pm

How do I obtain the IR of a network written in NNVM for custom IR optimization?
That is, assume I have the following network

net, params = nnvm.testing.resnet.get_workload(num_layers=18, batch_size=batch_size, image_shape=image_shape)

and I would like to obtain the IR for writing custom optimizations as described in the tutorial here [1]

[1] https://docs.tvm.ai/tutorials/dev/low_level_custom_pass.html

Thanks,

vinx13 · November 29, 2018, 5:11am

You can hook nnvm.compiler.lower by calling tvm.register_func with override=True

github.com

dmlc/tvm/blob/bbc78221777a06e3d6b6ee0c07f51f9f37201928/nnvm/python/nnvm/compiler/build_module.py#L103




Returns
-------
config: BuildConfig
    The build configuration
"""
return BuildConfig(**kwargs)




@tvm.register_func("nnvm.compiler.lower")
def _lower(sch, inputs, func_name, graph):
import traceback
# pylint: disable=broad-except
try:
    f = tvm.lower(sch, inputs, name=func_name)
    logging.debug("lower function %s", func_name)
    logging.debug("%s", tvm.lower(sch, inputs, simple_mode=True))
except Exception:
    msg = traceback.format_exc()
    msg += "Error during compile graph\n"
    msg += "--------------------------\n"

The lower function is called in

github.com

dmlc/tvm/blob/bbc78221777a06e3d6b6ee0c07f51f9f37201928/nnvm/src/compiler/compile_engine.cc#L272




   std::tie(sch, all_args, graph) = GetScheduleArgs(
       graph, inputs, target, master_idx,
       &readable_name, &outputs);


   auto gf = tvm::make_node<GraphFuncNode>();
   gf->target = target;
   gf->func_name = GetUniqeName(readable_name);
   gf->inputs = inputs;
   gf->outputs = outputs;
   static const PackedFunc& flower = GetPackedFunc("nnvm.compiler.lower");
   gf->funcs = flower(sch, all_args, gf->func_name, graph);
   return GraphFunc(gf);
 }


private:
 // Get unique name
 std::string GetUniqeName(std::string name) {
   while (true) {
     auto it = name_map_.find(name);
     if (it == name_map_.end()) {

trusira · November 29, 2018, 9:37pm

Thanks, but I’m still not clear how to generate a schedule of the network to call the ‘lower’ function.

vinx13 · November 30, 2018, 2:14am

DoLower calls into GetScheduleArgs https://github.com/dmlc/tvm/blob/bbc78221777a06e3d6b6ee0c07f51f9f37201928/nnvm/src/compiler/compile_engine.cc#L141, which will call TOPI schedule of opperators (the schedules are registered in nnvm/python/nnvm/top)