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