def build_and_eval(lib, s, bufs, target, dev_id, rpc_info: RpcInfo = None, number=1): if rpc_info is not None: target_host = rpc_info.target_host fcompile = rpc_info.fcompile use_rpc = rpc_info.use_rpc else: target_host, fcompile, use_rpc = None, None, None # mod = tvm.lower(s, bufs, simple_mode=True) # print("Building...") func = tvm.build(s, bufs, target=target, target_host=target_host) tvm_arys = [] try: func.export_library(lib, fcompile) # print("Connecting...") remote = rpc_info.get_remote() # print("Allocating...") ctx = (remote if remote else tvm).device(target, dev_id) for buf in bufs: shape = to_tuple(buf.shape) tmp = np.random.uniform(0, 1, size=shape).astype(buf.dtype) tmp = tvm.nd.array(tmp, ctx) tvm_arys.append(tmp) if use_rpc: # print("Uploading...") remote.upload(lib) func = remote.load_module(os.path.split(lib)[-1]) else: func = tvm.runtime.module.load_module(lib) # print("Evaluating...") evaluator = func.time_evaluator(func.entry_name, ctx, number=number) time_cost = evaluator(*tvm_arys).mean * 1e3 finally: while len(tvm_arys) > 0: del tvm_arys[-1] return time_cost
if __name__ == "__main__": N, C, H, W, K, _, k, _, _, st, pad, dilation, group = yolo_shapes_b1[5] use_bias = False vlen = 8 target = "llvm" dev_id = 0 trials = 100 timeout = 10 parallel = 20 method = "searching" force_inline = True use_model = False logfile = open("tmp.log", "w") rpc_info = RpcInfo("0.0.0.0", 9090, target_host="llvm") args = (N, C, H, W, K, k, use_bias, st, pad, dilation, group) task = Task("conv2d_nchwc", "yolo_conv6", conv2d_nchwc_compute_avx2, args, target, dev_id=dev_id) register_task(task, override=False) beg = time.time() s, bufs, configs = schedule(task.key, op_trial=trials, timeout=timeout, op_stop=30,
parser.add_argument("--timeout", help="timeout", type=float, default=4.0) parser.add_argument("--parallel", help="parallel", type=int, default=1) parser.add_argument("--use_model", help="use performance model", action="store_true") parser.add_argument("--method", help="how to schedule", type=str, default="searching") parser.add_argument("--host", type=str, default="0.0.0.0") parser.add_argument("--target_host", type=str, default="llvm") parser.add_argument("--port", type=int, default=9090) parser.add_argument("--force_inline", action="store_true") # parser.add_argument("--op_hint", type=str, default="split_fuse") args = parser.parse_args() rpc_info = RpcInfo(args.host, args.port, target_host=args.target_host) shapes = PixelCNN_shape if args.to < 0: end = len(shapes) else: end = args.to if args.log != "": with open(args.log, "a") as flog: ret = optimize( args.shapes, args.from_, shapes[args.from_:end], target=args.target, dev_id=args.device, timeout=args.timeout,
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 if __name__ == "__main__": parser = argparse.ArgumentParser() init_arg_parser(parser) args = parser.parse_args() rpc_info = RpcInfo(args.host, args.port, args.target_host, args.device_key, args.use_rpc, args.fcompile, args.timeout) shapes = conv1d_shapes if args.to < 0: end = len(shapes) else: end = args.to if args.test != "": with open(args.test, "r") as fin: for line in fin: name, string = line.split(":", 1) obj = json.loads(string) configs = Config(obj[0], obj[1]) test(name, configs, args.device, rpc_info, args.check)