[RFC] [Hybrid Script] A hybrid script backend

[were]

Sorry to change my mind again and again, and I will never use the term "finally". Making ideas is endless.

An IR Dumper #2288 , is more heavy-weighted than I supposed before.
I now decide to break this issue into 2 steps:

We already have Halide's IRPrinter to do a statical sanity check, and also the lowered IRs are too low-level to give good readability. Overall, what I want becomes something which is both runnable and gives full control to the IR body.

1. The current Halide IRPrinter has bad support to the double-name thing.
   This is because currently, TVM uses "unique pointer" to differentiate objects.
   For example, these 2 tensors below share the same name, but they are essentially two different tensors.
   However, if users use these 2 tensors to build a module, and see the lowered body, it will be extremely confusing. IRPrinter just dumps these tensors without checking double-name.

a = tvm.placeholder(shape)
b = tvm.placeholder(shape)

Also, he current IRDump dumps results of each pass to files whose suffixes are .cc. I guess the point of doing this is to borrow the syntax highlight of .cc. However, the contents dumped have nothing to do with C, so I believe migrating the text format to Python, and using Python suffix will be better.

2. Thus, What's more, dumping IR as Python snippet will not suffice. Dumping a complete hybrid script function will be highly desirable.
   Something which is both runnable and contains its source code is exactly what module is in TVM.
   Finally, I decided to extend a hybrid module looks like this:

#sch is a schedule of abitary op node
with tvm.target.create('hybrid'): #this with stmt is required by special lower passes
    module = tvm.build(sch, tvm_args) #a hybrid module is returned
hybrid_op = module(*tvm_args) #hybrid compilation
np_res = module(*np_args) #software emulation

[tqchen]

maybe we want to have a pass instead of a backend, if we want to use it as a default debug printer.

[were]

@tqchen A debug printer is far away from what we truly want.

As I mention above, Halide IRPrinter can already do a static sanity check. There is no point to build another static dumper, so a "runnable" one is desirable. To make it runnable, it requires much runtime information passed to build modules.

To sum up, it is built like a "module" (arguments and buffers), stores information like a "module" (hybrid source code), and runs like a "module" (pass either tensors/placeholders to drive it), why do not we just make it a module?

[were]

@tqchen I now know what you mean.
It generates infinite files named .cc.
It just borrows .cc suffix for code highlight.
We do need a more mature dumper.

[xqdan]

Normally we have two kinds dump in traditional compilers,

• one is just for reading, just like gcc, tvm's dump pass ir plays just like this right now.

• another is for reading and executing, just like llvm ir or text ir in relay, which can be as input for another phase of compilers.

is the second one what you want?
if so, we can do it like this, we can save two dump for each pass, one is cc files for reading, another is json file, json file can be loaded as input for another pass.

but we can not have single readable and runnable dump

this is our internal hacking

from contextlib import nested

   def decorate(self, func):
        """ decorate the pass function"""
        def dump(*args, **kwargs):
            """dump function"""
            start = time()
            retv = func(*args, **kwargs)
            dur = time() - start
            if not isinstance(retv, (_stmt.Stmt, container.LoweredFunc, container.Array)):
                return retv
            fname = func.func_name if hasattr(func, 'func_name') else func.__name__
            pname = str(self._pass_id) + "_" + fname + "_ir.cc"
            jpname = str(self._pass_id) + "_" + fname + "_ir.json"
            print(pname + " time   "+ str(dur))
            with nested(open(pname, "a"),  open(jpname, "a")) as (f,jf):
                out = retv.body if isinstance(retv, container.LoweredFunc) else retv
                f.write(str(out))
                st = api.save_json(out)
                jf.write(st)
                if isinstance(retv, container.Array):
                    for x in retv:
                        out = x.body if isinstance(x, container.LoweredFunc) else x
                        f.write("---------%s\n%s\n-----------\n"%(x.name, str(out)))
                self._pass_id += 1
            return retv
        return dump

and then load json as input,

import tvm                                                                                                                     

def test():
"""

"""
    scope_tb = "local.UB"
    @tvm.register_func("tvm.info.mem.%s" % scope_tb)
    def mem_info_inp_buffer():
        return tvm.make.node("MemoryInfo",
                        unit_bits= 256,
                        max_simd_bits=256,
                        max_num_bits=2031616,
                        head_address=None)
    scope_tb = "local.REG"
    @tvm.register_func("tvm.info.mem.%s" % scope_tb)
    def mem_info_inp_buffer():
        return tvm.make.node("MemoryInfo",
                        unit_bits= 16, 
                        max_simd_bits=64,
                        max_num_bits=204800,
                        head_address=None)
    f = open("28_InjectDoubleBuffer_ir.json")
    str_ = f.read()
    stmt = tvm.load_json(str_)
    print stmt
    print "=============="
    stmt = tvm.ir_pass.StorageRewrite(stmt)
    print stmt

test()

[kevinthesun]

@were Can you provide more details about how this backend is used by developer?

[were]

@kevinthesun I can elaborate it here:
You can compile a schedule of arbitrary op node to hybrid target, which returns a hybrid module.
The usage of this module is basically the same as hybrid functions and a normal module.
You can call it by tvm data structures to compile it to hybrid op.
You can software emulate it by passing numpy data structures.
You can get the source code by using .get_source method.

[tqchen]

#2477