def schedule_conv2d(attrs, outs, target):
""" 2D convolution schedule.
"""
layout = attrs["layout"]
if is_packed_layout(layout):
target = tvm.target.create(target)
if target.device_name == "vta":
return schedule_packed_conv2d(outs)
if str(target).startswith("llvm"):
return tvm.create_schedule([x.op for x in outs])
raise RuntimeError("not support target %s" % target)
return _nn.schedule_conv2d(attrs, outs, target)
def schedule_conv2d(attrs, outs, target):
""" 2D convolution schedule.
"""
layout = attrs["layout"]
if is_packed_layout(layout):
target = tvm.target.create(target)
if target.device_name == "vta":
return schedule_packed_conv2d(outs)
if str(target).startswith("llvm"):
return tvm.create_schedule([x.op for x in outs])
raise RuntimeError("not support target %s" % target)
return _nn.schedule_conv2d(attrs, outs, target)
def schedule_conv2d(attrs, outs, target):
""" Schedule definition of conv2d """
layout = attrs["layout"]
groups = attrs.get_int('groups')
if is_packed_layout(layout):
target = tvm.target.create(target)
if target.device_name == "vta":
if groups == 1:
return topi.generic.schedule_conv2d_nchw(outs)
return topi.generic.schedule_group_conv2d_nchw(outs)
elif str(target).startswith("llvm"):
return tvm.create_schedule([x.op for x in outs])
else:
raise RuntimeError("not support target %s" % target)
with tvm.target.arm_cpu(tvm.target.current_target().model):
return _nn.schedule_conv2d(attrs, outs, tvm.target.current_target())