def annotated(): add = relay.add(x, y) sub = relay.subtract(add, z) func = relay.Function([x, y, z], sub) func = run_opt_pass(func, transform.RewriteAnnotatedOps(dev1.device_type)) return func
def annotated(): conv2d_1 = relay.nn.conv2d(data1, weight, channels=64, kernel_size=(3, 3), padding=(1, 1)) _conv2d_1 = relay.annotation.on_device(conv2d_1, dev2) conv2d_2 = relay.nn.conv2d(data2, weight, channels=64, kernel_size=(3, 3), padding=(1, 1)) _conv2d_2 = relay.annotation.on_device(conv2d_2, dev2) add = relay.add(_conv2d_1, _conv2d_2) _add = relay.annotation.on_device(add, dev1) conv2d_3 = relay.nn.conv2d(_add, weight, channels=64, kernel_size=(3, 3), padding=(1, 1)) _conv2d_3 = relay.annotation.on_device(conv2d_3, dev2) func = relay.Function([data1, data2, weight], _conv2d_3) func = run_opt_pass( func, transform.RewriteAnnotatedOps(tvm.device(3).device_type)) return func
def annotated(): add = relay.add(x, y) _add = relay.annotation.on_device(add, dev1) sub = relay.subtract(_add, z) _sub = relay.annotation.on_device(sub, dev2) expr = run_opt_pass(_sub, transform.RewriteAnnotatedOps(dev1.device_type)) return expr
def original(): conv2d_1 = relay.nn.conv2d( data1, weight, channels=64, kernel_size=(3, 3), padding=(1, 1)) conv2d_2 = relay.nn.conv2d( data2, weight, channels=64, kernel_size=(3, 3), padding=(1, 1)) add = relay.add(conv2d_1, conv2d_2) conv2d_3 = relay.nn.conv2d( add, weight, channels=64, kernel_size=(3, 3), padding=(1, 1)) func = relay.Function([data1, data2, weight], conv2d_3) func = run_opt_pass( func, transform.RewriteAnnotatedOps(tvm.context(3).device_type)) return func
def annotated(): add = relay.add(x, y) _add = relay.annotation.on_device(add, dev2) sub = relay.subtract(_add, z) _sub = relay.annotation.on_device(sub, dev2) func = relay.Function([x, y, z], _sub) func = run_opt_pass(func, transform.RewriteAnnotatedOps(dev1.device_type)) return func
def annotated(): add = relay.add(x, y) _add1 = relay.annotation.on_device(add, dev2) _add2 = relay.annotation.on_device(add, dev2) sub1 = relay.subtract(_add1, z) sub2 = relay.subtract(_add2, z) func = relay.Function([x, y, z], relay.Tuple([sub1, sub2])) func = run_opt_pass(func, transform.RewriteAnnotatedOps(dev1.device_type)) return func
def annotated(): x = relay.var("x", relay.ty.TensorType((3, 3, 4), "float32")) split = relay.op.split(x, 3) split = split.astuple() split = relay.annotation.on_device(split, gpu_dev) split = relay.TupleWrapper(split, 3) sub = split[0] - split[1] func = relay.Function(relay.analysis.free_vars(sub), sub) func = run_opt_pass(func, transform.RewriteAnnotatedOps(cpu_dev.device_type)) return func
def annotated(): add = relay.add(a, b) _add = relay.annotation.on_device(add, dev_dev) mul = relay.multiply(c, d) _mul = relay.annotation.on_device(mul, cpu_dev) sub = relay.subtract(_add, _mul) _sub = relay.annotation.on_device(sub, dev_dev) func = relay.Function([a, b, c, d], _sub) func = run_opt_pass(func, transform.RewriteAnnotatedOps(dev_dev.device_type)) return func
def annotated(): add = relay.add(x, y) sqrt = relay.sqrt(add) log = relay.log(add) subtract = relay.subtract(sqrt, log) exp = relay.exp(subtract) _exp = relay.annotation.on_device(exp, cpu_dev) func = relay.Function([x, y], _exp) func = run_opt_pass(func, transform.RewriteAnnotatedOps(dev_dev.device_type)) return func
def annotated(): log = relay.log(x) _log = relay.annotation.on_device(log, expected_dev_type["log"]) log2 = relay.log2(_log) _log2 = relay.annotation.on_device(log2, expected_dev_type["log2"]) log10 = relay.log10(_log) _log10 = relay.annotation.on_device(log10, expected_dev_type["log10"]) add = relay.add(_log2, _log10) _add = relay.annotation.on_device(add, expected_dev_type["add"]) tan = relay.tan(_add) _tan = relay.annotation.on_device(tan, expected_dev_type["tan"]) func = run_opt_pass(_tan, transform.RewriteAnnotatedOps(dev1.device_type)) return func
def annotated(): log = relay.log(x) _log = relay.annotation.on_device(log, expected_dev_type['log']) log2 = relay.log2(_log) _log2 = relay.annotation.on_device(log2, expected_dev_type['log2']) log10 = relay.log10(_log) _log10 = relay.annotation.on_device(log10, expected_dev_type['log10']) add = relay.add(_log2, _log10) _add = relay.annotation.on_device(add, expected_dev_type['add']) tan = relay.tan(_add) _tan = relay.annotation.on_device(tan, expected_dev_type['tan']) func = run_opt_pass(_tan, transform.RewriteAnnotatedOps(ctx1.device_type)) return func
def annotated(): add = relay.add(x, y) _add = relay.annotation.on_device(add, dev_dev) sqrt = relay.sqrt(_add) _sqrt = relay.annotation.on_device(sqrt, dev_dev) log = relay.log(_add) _log = relay.annotation.on_device(log, dev_dev) subtract = relay.subtract(_sqrt, _log) _subtract = relay.annotation.on_device(subtract, dev_dev) exp = relay.exp(_subtract) _exp = relay.annotation.on_device(exp, dev_dev) func = relay.Function([x, y], _exp) func = run_opt_pass(func, transform.RewriteAnnotatedOps(cpu_dev.device_type)) return func
def check_storage_and_device_types(): func = annotated() func = run_opt_pass(func, [transform.RewriteAnnotatedOps(3), transform.FuseOps(2)]) smap = relay.backend._backend.GraphPlanMemory(func) storage_ids = [] device_types = [] for _, storage_dev_type in smap.items(): assert len(storage_dev_type) == 2 for sid in storage_dev_type[0]: storage_ids.append(sid.value) for did in storage_dev_type[1]: device_types.append(did.value) assert len(storage_ids) == 10 assert len(set(storage_ids)) == 8 assert len(set(device_types)) == 2 assert set(device_types) == {1, 2}
def annotate_with_visitor(func): sched = ScheduleConv2d(dev2) func = sched.visit(func) func = run_opt_pass(func, transform.RewriteAnnotatedOps(dev1.device_type)) return func