Relay CLI and REPL

relay.py

import argparse

import tvm
from tvm import relay
from tvm.relay import testing, create_executor

SEM_VER='v0.0.1'

def parse_args():
    parser = argparse.ArgumentParser(description='Run relay program')
    parser.add_argument('prog_file', metavar='prog', type=str, default=None, nargs='?',
                        help='relay source code file name')
    return parser.parse_args()

def relay_eval(source):
    if SEM_VER not in source:
        source = SEM_VER + source

    prog = relay.fromtext(source)

    target = 'llvm'
    ctx = tvm.context(target, 0)
    if not ctx.exist:
        return
    intrp = create_executor(ctx=ctx, target=target)
    result = intrp.evaluate(prog)
    return result

def prompt_continuation(width, line_number, is_soft_wrap):
    return '.' * width

def repl():
    try:
        from prompt_toolkit import prompt
        i = 0
        while True:
            try:
                source = prompt('In [{}]: '.format(i), multiline=True, prompt_continuation=prompt_continuation)
                val = relay_eval(source)
                print('Val: {}'.format(val))
            except Exception as e:
                print("Failed: {}".format(e))
            i += 1
    except KeyboardInterrupt:
        pass

def main():
    args = parse_args()
    if args.prog_file is None:
        return repl()

    with open(args.prog_file) as f:
        source = f.read().strip()    
    print("Val: {}".format(relay_eval(source)))

if __name__ == '__main__':
    main()

(3.6.3.py) ➜  python relay.py
In [0]: let %x = 1;
........let %y = 1;
........%x
........
Val: 1