def test(task_key, configs, dev_id=None):
task = TASK_TABLE[task_key]
s, bufs = schedule_with_config(task_key, configs)
dev_id = dev_id if dev_id is not None else task.dev_id
time_cost = _evaluate(s, bufs, task.target, dev_id, 10)
print(task_key, "use", time_cost, "ms")
print()
def test():
# create an empty task but has the correct key we want
task = Task("yolo1", None, (1, 3, 448, 448, 64, 7, 2, 3, 1, 1), "llvm", 0)
beg = time.time()
# s, bufs, configs = schedule(task.key)
end = time.time()
# print(tvm.lower(s, bufs, simple_mode=True))
# print("######################################")
# print("op schedules:")
# for config in configs.op_config_lst:
# print("----------------------------------")
# for name, value in config.items():
# if value:
# print(name, value)
# print("graph schedules:")
# for name, value in configs.graph_config.items():
# if value:
# print(name, value)
op_configs = [{
"spatial": [[1, 1, 1, 1], [1, 1, 1, 3], [454, 1, 1, 1], [1, 227, 2,
1]],
"unroll": [[1500, 1]]
}, {
"spatial": [[1, 1, 1, 1], [2, 4, 2, 4], [8, 1, 4, 7], [7, 1, 16, 2]],
"reduce": [[1, 3, 1], [7, 1, 1], [7, 1, 1]],
"unroll": [[1500, 1]]
}]
graph_config = {"inline": [[0, 0]]}
configs = Config(op_configs, graph_config)
s, bufs = schedule_with_config(task.key, configs)
time_cost = _evaluate(s, bufs, "llvm", 0, 10)
print("Use", time_cost, "ms")
print("Cost", end - beg, "s")
def check_result(configs, shape, target="cuda", dev_id=0):
ctx = tvm.context(target, dev_id)
name, configs = configs
batch, in_channel, H, W, out_channel, k, _, stride, padding, dilation, groups = shape
A_np = np.random.uniform(-10, 10, size=[batch, in_channel, H,
W]).astype("float32")
A_tvm = tvm.nd.array(A_np, ctx)
A_torch = torch.tensor(A_np) # .cuda("cuda:" + str(dev_id))
W_np = np.random.uniform(-10,
10,
size=[out_channel, in_channel // groups, k,
k]).astype("float32")
W_tvm = tvm.nd.array(W_np, ctx)
W_torch = torch.tensor(W_np) # .cuda("cuda:" + str(dev_id))
Output_torch = torch.nn.functional.conv2d(A_torch,
W_torch,
stride=stride,
padding=padding,
dilation=dilation,
groups=groups)
Output_np = np.zeros(Output_torch.shape).astype(np.float32)
Output_tvm = tvm.nd.array(Output_np, ctx)
s, bufs = schedule_with_config(name, configs)
func = tvm.build(s, bufs, target)
func(A_tvm, W_tvm, Output_tvm)
passed = test_allclose(Output_tvm.asnumpy(),
Output_torch.cpu().numpy(),
rtol=1e-5,
print_diff=True)
if passed == 1:
print("Passed!")
else:
print("Failed!")
def optimize(shapes,
slevel=4,
rlevel=3,
target="llvm",
dev_id=0,
timeout=4.0,
trials=100,
parallel=1,
method="searching",
use_model=False,
rpc_info=None,
logfile=sys.stdout):
ret = dict()
for i, shape in enumerate(shapes):
print("Optimize depthwise conv2d shape {}".format(shape), flush=True)
batch, in_channel, H, W, factor, k, _, stride, padding, dilation = shape
# create an empty task but has the correct key we want
task = Task(
"conv2d", "depthwise", None,
(batch, in_channel, H, W, factor, k, stride, padding, dilation),
target, dev_id)
beg = time.time()
s, bufs, configs = schedule(
task.key,
slevel=slevel,
rlevel=rlevel,
op_trial=trials,
timeout=timeout,
op_stop=trials // 2,
method=method,
use_model=use_model,
parallel=parallel,
rpc_info=rpc_info,
trials=[trials // 10, trials],
number=10,
)
end = time.time()
# print(tvm.lower(s, bufs, simple_mode=True))
print("######################################")
print("op schedules:")
for config in configs.op_config_lst:
print("----------------------------------")
for name, value in config.items():
if value:
print(name, value)
print("graph schedules:")
for name, value in configs.graph_config.items():
if value:
print(name, value)
ret[task.key] = configs
string = json.dumps(configs)
line = task.key + ":" + string
print(line, file=logfile, flush=True)
s, bufs = schedule_with_config(task.key, configs)
time_cost = evaluate(task.key, s, bufs, target, task.dev_id, 10,
rpc_info)
print("Use", time_cost, "ms")
print("Cost", end - beg, "s")
print()
return ret
def optimize(prefix,
from_,
shapes,
target="llvm",
dev_id=0,
trials=100,
timeout=4.0,
parallel=1,
method="searching",
use_model=False,
logfile=sys.stdout):
ret = dict()
for i, shape in enumerate(shapes):
print("Optimize {} conv_transpose2d layer {} shape {}".format(
prefix, i + 1 + from_, shape),
flush=True)
batch, in_channel, height, width, out_channel, _, k_h, k_w, _, stride, padding, dilation, groups = shape
rout_channel = in_channel
rin_channel = out_channel
rheight = (height + 2 * padding - dilation *
(k_h - 1) - 1) // stride + 1
rwidth = (width + 2 * padding - dilation * (k_w - 1) - 1) // stride + 1
# create an empty task but has the correct key we want
task = Task("conv_transpose2d", prefix + str(i + from_), None,
(batch, rin_channel, rheight, rwidth, rout_channel, k_h,
stride, padding, dilation, groups), target, dev_id)
beg = time.time()
s, bufs, configs = schedule(
task.key,
op_trial=trials,
timeout=timeout,
op_stop=30,
parallel=parallel,
method=method,
use_model=use_model,
trials=[trials // 10, trials // 10, trials])
end = time.time()
# print(tvm.lower(s, bufs, simple_mode=True))
print("######################################")
print("op schedules:")
for config in configs.op_config_lst:
print("----------------------------------")
for name, value in config.items():
if value:
print(name, value)
print("graph schedules:")
for name, value in configs.graph_config.items():
if value:
print(name, value)
ret[task.key] = configs
string = json.dumps(configs)
line = task.key + ":" + string
print(line, file=logfile, flush=True)
s, bufs = schedule_with_config(task.key, configs)
time_cost = _evaluate(s, bufs, target, task.dev_id, 10)
print("Use", time_cost, "ms")
print("Cost", end - beg, "s")
print()
return ret
def optimize(shapes,
slevel=4,
rlevel=3,
target="llvm",
dev_id=0,
timeout=4.0,
trials=100,
parallel=1,
method="searching",
use_model=False,
rpc_info=None,
logfile=sys.stdout):
ret = dict()
for i, shape in enumerate(shapes):
print("Optimize gemm shape %s [%.6f]" % (str(shape), time.time()),
flush=True)
N, K, M = shape
# create an empty task but has the correct key we want
task = Task("gemm", "gemm", None, (N, K, M), target, dev_id)
beg = time.time()
s, bufs, configs = schedule(task.key,
slevel=slevel,
rlevel=rlevel,
op_trial=trials,
timeout=timeout,
op_stop=trials // 2,
method=method,
use_model=use_model,
parallel=parallel,
rpc_info=rpc_info,
number=10)
end = time.time()
# print(tvm.lower(s, bufs, simple_mode=True))
print("###################################### [%.6f]" % time.time())
print("op schedules:")
for config in configs.op_config_lst:
print("----------------------------------")
for name, value in config.items():
if value:
print(name, value)
print("graph schedules:")
for name, value in configs.graph_config.items():
if value:
print(name, value)
ret[task.key] = configs
string = json.dumps(configs)
line = task.key + ":" + string
print(line, file=logfile, flush=True)
s, bufs = schedule_with_config(task.key, configs)
time_cost = evaluate(task.key,
s,
bufs,
target,
dev_id,
rpc_info=rpc_info)
print("Use", time_cost, "ms")
print("Cost", end - beg, "s")
print()
return ret
def optimize(prefix,
from_,
shapes,
target="llvm",
dev_id=0,
trials=100,
timeout=4.0,
parallel=1,
method="searching",
use_model=False,
rpc_info=None,
force_inline=False,
logfile=sys.stdout):
ret = dict()
for i, shape in enumerate(shapes):
print("Optimize {} pixelCNN layer {} shape {}".format(
prefix, i + 1 + from_, shape),
flush=True)
# create an empty task but has the correct key we want
task = Task("gatedpixelcnn", "gatedpixelcnn", None, shape, target,
dev_id)
beg = time.time()
s, bufs, configs = schedule(
task.key,
op_trial=trials,
timeout=timeout,
op_stop=30,
parallel=parallel,
method=method,
use_model=use_model,
# trials=[trials//10, trials],
force_inline=force_inline,
rpc_info=rpc_info,
)
end = time.time()
# print(tvm.lower(s, bufs, simple_mode=True))
print("######################################")
print("op schedules:")
for config in configs.op_config_lst:
print("----------------------------------")
for name, value in config.items():
if value:
print(name, value)
print("graph schedules:")
for name, value in configs.graph_config.items():
if value:
print(name, value)
ret[task.key] = configs
string = json.dumps(configs)
line = task.key + ":" + string
print(line, file=logfile, flush=True)
s, bufs = schedule_with_config(task.key, configs)
time_cost = evaluate(task.key, s, bufs, target, task.dev_id, 10,
rpc_info)
print("Use", time_cost, "ms")
print("Cost", end - beg, "s")
print()
return ret
def test(task_key, configs, dev_id=None, rpc_info=None):
task = TASK_TABLE[task_key]
s, bufs = schedule_with_config(task_key, configs)
# print(tvm.lower(s, bufs, simple_mode=True))
func = tvm.build(s, bufs, "cuda")
print(func.imported_modules[0].get_source())
dev_id = dev_id if dev_id is not None else task.dev_id
time_cost = evaluate(task_key, s, bufs, task.target, dev_id, 10, rpc_info)
print(task_key, "use", time_cost, "ms")
print()
def optimize(shapes,
slevel=4,
rlevel=3,
target="llvm",
dev_id=0,
timeout=4.0,
trials=100,
parallel=1,
method="searching",
use_model=False,
logfile=sys.stdout):
ret = dict()
for i, shape in enumerate(shapes):
print("Optimize block_circulant_matrix shape {}".format(shape),
flush=True)
ROW, COL, FFT = shape
# create an empty task but has the correct key we want
task = Task("block_circulant_matrix", "block_circulant_matrix", None,
(ROW, COL, FFT), target, dev_id)
beg = time.time()
s, bufs, configs = schedule(task.key,
slevel=slevel,
rlevel=rlevel,
op_trial=trials,
timeout=timeout,
op_stop=30,
method=method,
use_model=use_model,
parallel=parallel)
end = time.time()
# print(tvm.lower(s, bufs, simple_mode=True))
print("######################################")
print("op schedules:")
for config in configs.op_config_lst:
print("----------------------------------")
for name, value in config.items():
if value:
print(name, value)
print("graph schedules:")
for name, value in configs.graph_config.items():
if value:
print(name, value)
ret[task.key] = configs
string = json.dumps(configs)
line = task.key + ":" + string
print(line, file=logfile, flush=True)
s, bufs = schedule_with_config(task.key, configs)
time_cost = _evaluate(s, bufs, target, task.dev_id, 10)
print("Use", time_cost, "ms")
print("Cost", end - beg, "s")
print()
return ret
def test(task_key, configs, dev_id=None, rpc_info=None, check_result=False):
task = TASK_TABLE[task_key]
result_generator = None
if check_result:
def _generator(np_arys):
ops, bufs = task.func(*task.args)
s = te.create_schedule(ops)
func = tvm.build(s, bufs, target="llvm")
ctx = tvm.cpu(0)
tvm_arys = [tvm.nd.array(arr, ctx) for arr in np_arys]
func(*tvm_arys)
return tvm_arys[-1].asnumpy()
result_generator = _generator
s, bufs = schedule_with_config(task_key, configs)
dev_id = dev_id if dev_id is not None else task.dev_id
time_cost = evaluate(task_key, s, bufs, task.target, dev_id, 100, rpc_info,
result_generator)
print(task_key, "use", time_cost, "ms")
print()
use_model=use_model,
trials=[trials // 10, trials],
force_inline=force_inline,
rpc_info=rpc_info,
slevel=2,
rlevel=2)
end = time.time()
print("######################################")
print("op schedules:")
for config in configs.op_config_lst:
print("----------------------------------")
for name, value in config.items():
if value:
print(name, value)
print("graph schedules:")
for name, value in configs.graph_config.items():
if value:
print(name, value)
string = json.dumps(configs)
line = task.key + ":" + string
print(line, file=logfile, flush=True)
s, bufs = schedule_with_config(task.key, configs)
time_cost = _evaluate(s, bufs, target, dev_id, 10)
print(
"Use", time_cost, "ms", "throughput: %f GFLOPS" %
(N * C * H * W * K * k * k / st / st / group / 1e6 / time_cost))
print("Cost", end - beg, "s")
logfile.close()
