def _setup(self, mod, target): tgts = {} if isinstance(target, dict): for dev, tgt in target.items(): if not isinstance(tgt, (str, _target.Target)): raise Exception("Unknown target type") tgts[dev] = _target.create(tgt) elif isinstance(target, (str, _target.Target)): tgts[_expr.IntImm("int32", 0)] = _target.create(target) self._init(mod, tgts)
def set_task(self, task): self.task = task if check_remote(task.target, self.key, self.host, self.port): logger.info("Get devices for measurement successfully!") else: raise RuntimeError( "Cannot get remote devices from the tracker. " "Please check the status of tracker by " "'python -m tvm.exec.query_rpc_tracker --port [THE PORT YOU USE]' " "and make sure you have free devices on the queue status.") if self.check_correctness: # use llvm cpu to generate a reference input/output # this option works for tuning topi, but might not work for you custom op with _target.create("llvm"): s, arg_bufs = task.instantiate(task.config_space.get(0)) self.ref_input = [ np.random.uniform(size=get_const_tuple(x.shape)).astype( x.dtype) for x in arg_bufs ] func = build(s, arg_bufs, "llvm") tvm_buf = [nd.array(x) for x in self.ref_input] func(*tvm_buf) self.ref_output = [x.asnumpy() for x in tvm_buf]
def get_flatten_name(fea): """ Get names of feature after flatten. Parameters ---------- fea: list or str return value of get_itervar_feature or a line of logfile Returns ------- feature_names: Array of str """ feature_name = { "_attr_": ["length", "nest_level", "topdown", "bottomup"] + ["ann_%d" % i for i in range(20)], "_arith_": ["add", "mul", "div"], "buf_touch": ["stride", "mod", "count", "reuse", "T_count", "T_reuse"], } if isinstance(fea, str): # pylint: disable=import-outside-toplevel from .record import decode # flatten line to feature line = fea ret = decode(line) if ret is None: raise ValueError("Unsupported AutoTVM log format") inp, _ = ret target = _target.create(inp.target) with target: s, args = inp.template.instantiate(inp.config) fea = get_itervar_feature(s, args) names = [] ct = 0 for row in fea: var_name = str(row[0][1]) for pair in row[1:]: key = pair[0] if key in feature_name: name_list = feature_name[key] else: name_list = feature_name["buf_touch"] for i in range(len((pair[1:]))): names.append(".".join( ["f%d" % ct, var_name, key, name_list[i]])) ct += 1 return names
def get_flatten_name(fea): """ Get names of feature after flatten. Parameters ---------- fea: list or str return value of get_itervar_feature or a line of logfile Returns ------- feature_names: Array of str """ feature_name = { "_attr_": ["length", "nest_level", "topdown", "bottomup"] + ["ann_%d" % i for i in range(20)], "_arith_": ["add", "mul", "div"], "buf_touch": ["stride", "mod", "count", "reuse", "T_count", "T_reuse"], } if isinstance(fea, str): from .record import decode # flatten line to feature line = fea inp, _ = decode(line) target = _target.create(inp.target) with target: s, args = inp.template.instantiate(inp.config) fea = get_itervar_feature(s, args) names = [] ct = 0 for row in fea: var_name = str(row[0][1]) for pair in row[1:]: key = pair[0] if key in feature_name: name_list = feature_name[key] else: name_list = feature_name["buf_touch"] for i in range(len((pair[1:]))): names.append(".".join(["f%d" % ct, var_name, key, name_list[i]])) ct += 1 return names
def create(task_name, args, target, target_host=None): """Create a tuning task and initialize its search space Parameters ---------- task_name : str The AutoTVM task name args : List Positional arguments target : Target The compilation target target_host: Target, optional The compilation target for host side Returns ------- tsk: Task a task object """ args = serialize_args(args) ret = Task(task_name, args) if isinstance(target, str): target = _target.create(target) # init config space ret.config_space = ConfigSpace() ctx = ApplyConfig(ret.config_space) with ctx: with target: sch, _ = ret.func(*args) ret.config_space.code_hash = getattr(sch, 'code_hash', None) ret.flop = ret.config_space.flop or compute_flop(sch) ret.target = target ret.target_host = target_host return ret
def build(inputs, args=None, target=None, target_host=None, name="default_function", binds=None): """Build a function with arguments as signature. Code will be generated for devices coupled with target information. Parameters ---------- inputs : tvm.te.Schedule, IRModule, or dict of target to IRModule The schedule to be built args : list of Buffer or Tensor or Var, optional The argument lists to the function. target : str or :any:`tvm.target.Target`, optional The target and option of the compilation. target_host : str or :any:`tvm.target.Target` optional Host compilation target, if target is device. When TVM compiles device specific program such as CUDA, we also need host(CPU) side code to interact with the driver setup the dimensions and parameters correctly. target_host is used to specify the host side codegen target. By default, llvm is used if it is enabled, otherwise a stackvm intepreter is used. name : str, optional The name of result function. binds : dict, optional Dictionary that maps the binding of symbolic buffer to Tensor. By default, a new buffer is created for each tensor in the argument. Returns ------- ret : tvm.module A module that combines both host and device code. Examples ________ There are two typical example uses of this function depending on the type of the argument `inputs`: 1. it is an IRModule. .. code-block:: python n = 2 A = te.placeholder((n,), name='A') B = te.placeholder((n,), name='B') C = te.compute(A.shape, lambda *i: A(*i) + B(*i), name='C') s = tvm.te.create_schedule(C.op) m = tvm.lower(s, [A, B, C], name="test_add") rt_mod = tvm.build(m, target="llvm") 2. it is a dict of compilation target to IRModule. .. code-block:: python n = 2 A = te.placeholder((n,), name='A') B = te.placeholder((n,), name='B') C = te.compute(A.shape, lambda *i: A(*i) + B(*i), name='C') s1 = tvm.te.create_schedule(C.op) with tvm.target.cuda() as cuda_tgt: s2 = topi.cuda.schedule_injective(cuda_tgt, [C]) m1 = tvm.lower(s1, [A, B, C], name="test_add1") m2 = tvm.lower(s2, [A, B, C], name="test_add2") rt_mod = tvm.build({"llvm": m1, "cuda": m2}, target_host="llvm") Note ---- See the note on :any:`tvm.target` on target string format. """ if isinstance(inputs, schedule.Schedule): if args is None: raise ValueError("args must be given for build from schedule") input_mod = lower(inputs, args, name=name, binds=binds) elif isinstance(inputs, (list, tuple, container.Array)): merged_mod = tvm.IRModule({}) for x in inputs: merged_mod.update(x) input_mod = merged_mod elif isinstance(inputs, tvm.IRModule): input_mod = inputs elif not isinstance(inputs, (dict, container.Map)): raise ValueError( "inputs must be Schedule, IRModule or dict of target to IRModule") if not isinstance(inputs, (dict, container.Map)): target = _target.Target.current() if target is None else target target = target if target else "llvm" target_input_mod = {target: input_mod} else: target_input_mod = inputs for tar, mod in target_input_mod.items(): if not isinstance(tar, (str, _target.Target)): raise ValueError("The key of inputs must be str or " "_target.Target when inputs is dict.") if not isinstance(mod, tvm.IRModule): raise ValueError("inputs must be Schedule, IRModule," "or dict of str to IRModule.") if not target_host: for tar, _ in target_input_mod.items(): tar = _target.create(tar) device_type = ndarray.context(tar.target_name, 0).device_type if device_type == ndarray.cpu(0).device_type: target_host = tar break if not target_host: target_host = "llvm" if tvm.runtime.enabled("llvm") else "stackvm" mod_host_all = tvm.IRModule({}) device_modules = [] for tar, input_mod in target_input_mod.items(): mod_host, mdev = _build_for_device(input_mod, tar, target_host) mod_host_all.update(mod_host) device_modules.append(mdev) # Generate a unified host module. rt_mod_host = codegen.build_module(mod_host_all, target_host) # Import all modules. for mdev in device_modules: if mdev: rt_mod_host.import_module(mdev) return rt_mod_host
def _build_for_device(input_mod, target, target_host): """Build the lowered functions for a device with the given compilation target. Parameters ---------- input_mod : IRModule The schedule to be built. target : str or :any:`tvm.target.Target` The target and option of the compilation. target_host : str or :any:`tvm.target.Target` The host compilation target. Returns ------- fhost : IRModule The host IRModule. mdev : tvm.module A module that contains device code. """ target = _target.create(target) target_host = _target.create(target_host) device_type = ndarray.context(target.target_name, 0).device_type mod_mixed = input_mod mod_mixed = tvm.tir.transform.Apply( lambda f: f.with_attr("target", target))(mod_mixed) opt_mixed = [tvm.tir.transform.VerifyMemory()] if len(mod_mixed.functions) == 1: opt_mixed += [ tvm.tir.transform.Apply( lambda f: f.with_attr("tir.is_entry_func", True)) ] if BuildConfig.current().detect_global_barrier: opt_mixed += [tvm.tir.transform.ThreadSync("global")] opt_mixed += [ tvm.tir.transform.ThreadSync("shared"), tvm.tir.transform.ThreadSync("warp"), tvm.tir.transform.InferFragment(), tvm.tir.transform.LowerThreadAllreduce(), tvm.tir.transform.MakePackedAPI(), tvm.tir.transform.SplitHostDevice() ] mod_mixed = tvm.transform.Sequential(opt_mixed)(mod_mixed) # device optimizations opt_device = tvm.transform.Sequential([ tvm.tir.transform.Filter( lambda f: "calling_conv" in f.attrs and f.attrs[ "calling_conv"].value == CallingConv.DEVICE_KERNEL_LAUNCH), tvm.tir.transform.LowerWarpMemory(), tvm.tir.transform.Simplify(), tvm.tir.transform.LowerDeviceStorageAccessInfo(), tvm.tir.transform.LowerIntrin() ]) mod_dev = opt_device(mod_mixed) # host optimizations opt_host = tvm.transform.Sequential([ tvm.tir.transform.Filter( lambda f: "calling_conv" not in f.attrs or f.attrs[ "calling_conv"].value != CallingConv.DEVICE_KERNEL_LAUNCH), tvm.tir.transform.Apply(lambda f: f.with_attr("target", target)), tvm.tir.transform.LowerTVMBuiltin(), tvm.tir.transform.LowerDeviceStorageAccessInfo(), tvm.tir.transform.LowerIntrin(), tvm.tir.transform.CombineContextCall() ]) mod_host = opt_host(mod_mixed) if device_type == ndarray.cpu(0).device_type and target_host == target: assert len(mod_dev.functions) == 0 if "gpu" in target.keys and len(mod_dev.functions) == 0: warnings.warn( "Specified target %s, but cannot find device code, did you do " "bind?" % target) rt_mod_dev = codegen.build_module( mod_dev, target) if len(mod_dev.functions) != 0 else None return mod_host, rt_mod_dev
def _build_for_device(flist, target, target_host): """Build the lowered functions for a device with the given compilation target. Parameters ---------- flist : list of LoweredFunc The schedule to be built. target : str or :any:`tvm.target.Target` The target and option of the compilation. target_host : str or :any:`tvm.target.Target` The host compilation target. Returns ------- fhost : list of LoweredFunc A list of lowered functions for the host. mdev : tvm.module A module that contains device code. """ target = _target.create(target) target_host = _target.create(target_host) device_type = ndarray.context(target.target_name, 0).device_type for func in flist: if not ir_pass.VerifyMemory(func, device_type): raise ValueError( "Direct host side access to device memory is detected in %s. " "Did you forget to bind?" % func.name) mod_mixed = tvm.testing.LoweredFuncsToIRModule(flist) opt_mixed = [tvm.tir.transform.Apply(lambda f: f.with_attr("target", target))] if BuildConfig.current().detect_global_barrier: opt_mixed += [tvm.tir.transform.ThreadSync("global")] opt_mixed += [tvm.tir.transform.ThreadSync("shared"), tvm.tir.transform.ThreadSync("warp"), tvm.tir.transform.InferFragment(), tvm.tir.transform.LowerThreadAllreduce(), tvm.tir.transform.BindDeviceType(), tvm.tir.transform.SplitHostDevice()] mod_mixed = tvm.ir.transform.Sequential(opt_mixed)(mod_mixed) # device optimizations opt_device = tvm.ir.transform.Sequential( [tvm.tir.transform.Filter( lambda f: "calling_conv" in f.attrs and f.attrs["calling_conv"].value == CallingConv.DEVICE_KERNEL_LAUNCH), tvm.tir.transform.LowerWarpMemory(), tvm.tir.transform.LowerDeviceStorageAccessInfo(), tvm.tir.transform.LowerIntrin()]) mod_dev = opt_device(mod_mixed) # host optimizations opt_host = tvm.ir.transform.Sequential( [tvm.tir.transform.Filter( lambda f: "calling_conv" not in f.attrs or f.attrs["calling_conv"].value != CallingConv.DEVICE_KERNEL_LAUNCH), tvm.tir.transform.Apply(lambda f: f.with_attr("target", target)), tvm.tir.transform.LowerTVMBuiltin(), tvm.tir.transform.LowerDeviceStorageAccessInfo(), tvm.tir.transform.LowerIntrin(), tvm.tir.transform.CombineContextCall()]) mod_host = opt_host(mod_mixed) if device_type == ndarray.cpu(0).device_type and target_host == target: assert len(mod_dev.functions) == 0 if "gpu" in target.keys and len(mod_dev.functions) == 0: warnings.warn( "Specified target %s, but cannot find device code, did you do " "bind?" % target) rt_mod_dev = codegen.build_module(mod_dev, target) if len(mod_dev.functions) != 0 else None return mod_host, rt_mod_dev
def _build_for_device(flist, target, target_host): """Build the lowered functions for a device with the given compilation target. Parameters ---------- flist : list of LoweredFunc The schedule to be built. target : str or :any:`tvm.target.Target` The target and option of the compilation. target_host : str or :any:`tvm.target.Target` The host compilation target. Returns ------- fhost : list of LoweredFunc A list of lowered functions for the host. mdev : tvm.module A module that contains device code. """ @tvm.tir.transform.prim_func_pass(opt_level=0) class BindTarget: def __init__(self, target): self.target = target # pylint: disable=unused-argument def transform_function(self, func, mod, ctx): return func.with_attr("target", self.target) target = _target.create(target) device_type = ndarray.context(target.target_name, 0).device_type fhost = [] fdevice = [] for func in flist: if not ir_pass.VerifyMemory(func, device_type): raise ValueError( "Direct host side access to device memory is detected in %s. " "Did you forget to bind?" % func.name) if func.func_type == LoweredFunc.MixedFunc: if BuildConfig.current().detect_global_barrier: func = ir_pass.ThreadSync(func, "global") func = ir_pass.ThreadSync(func, "shared") func = ir_pass.ThreadSync(func, "warp") func = ir_pass.InferFragment(func) warp_size = target.thread_warp_size func = ir_pass.LowerThreadAllreduce(func, warp_size) fsplits = list(ir_pass.SplitHostDevice(func)) fhost.append(fsplits[0]) for x in fsplits[1:]: fdevice.append(x) elif func.func_type == LoweredFunc.HostFunc: fhost.append(func) elif func.func_type == LoweredFunc.DeviceFunc: fdevice.append(func) else: raise ValueError("unknown function type %d" % func.func_type) if "gpu" in target.keys and not fdevice: warnings.warn( "Specified target %s, but cannot find device code, did you do " "bind?" % target) fhost = [ir_pass.BindDeviceType(x, device_type) for x in fhost] if device_type == ndarray.cpu(0).device_type and target_host == target: assert not fdevice target_host = _target.create(target_host) # device optimizations mod_dev = tvm.testing.LoweredFuncsToIRModule(fdevice) opt_device = tvm.ir.transform.Sequential( [BindTarget(target), tvm.tir.transform.LowerWarpMemory(), tvm.tir.transform.LowerDeviceStorageAccessInfo(), tvm.tir.transform.LowerIntrin()]) mod_dev = opt_device(mod_dev) # host optimizations mod_host = tvm.testing.LoweredFuncsToIRModule(fhost) opt_host = tvm.ir.transform.Sequential( [BindTarget(target_host), tvm.tir.transform.LowerTVMBuiltin(), tvm.tir.transform.LowerDeviceStorageAccessInfo(), tvm.tir.transform.LowerIntrin(), tvm.tir.transform.CombineContextCall()]) mod_host = opt_host(mod_host) rt_mod_dev = codegen.build_module(mod_dev, target) if fdevice else None return mod_host, rt_mod_dev
def _build_for_device(flist, target, target_host): """Build the lowered functions for a device with the given compilation target. Parameters ---------- flist : list of LoweredFunc The schedule to be built. target : str or :any:`tvm.target.Target` The target and option of the compilation. target_host : str or :any:`tvm.target.Target` The host compilation target. Returns ------- fhost : list of LoweredFunc A list of lowered functions for the host. mdev : tvm.module A module that contains device code. """ target = _target.create(target) device_type = ndarray.context(target.target_name, 0).device_type fhost = [] fdevice = [] for func in flist: if not ir_pass.VerifyMemory(func, device_type): raise ValueError( "Direct host side access to device memory is detected in %s. " "Did you forget to bind?" % func.name) if func.func_type == LoweredFunc.MixedFunc: if BuildConfig.current().detect_global_barrier: func = ir_pass.ThreadSync(func, "global") func = ir_pass.ThreadSync(func, "shared") func = ir_pass.ThreadSync(func, "warp") func = ir_pass.InferFragment(func) warp_size = target.thread_warp_size func = ir_pass.LowerThreadAllreduce(func, warp_size) fsplits = list(ir_pass.SplitHostDevice(func)) fhost.append(fsplits[0]) for x in fsplits[1:]: fdevice.append(x) elif func.func_type == LoweredFunc.HostFunc: fhost.append(func) elif func.func_type == LoweredFunc.DeviceFunc: fdevice.append(func) else: raise ValueError("unknown function type %d" % func.func_type) for i, func in enumerate(fdevice): warp_size = target.thread_warp_size fdevice[i] = ir_pass.LowerWarpMemory(func, warp_size) if "gpu" in target.keys and not fdevice: warnings.warn( "Specified target %s, but cannot find device code, did you do " "bind?" % target) fhost = [ir_pass.BindDeviceType(x, device_type) for x in fhost] fhost = [ir_pass.LowerTVMBuiltin(x) for x in fhost] if device_type == ndarray.cpu(0).device_type and target_host == target: assert not fdevice target_host = _target.create(target_host) fdevice = [ir_pass.LowerDeviceStorageAccessInfo(x) for x in fdevice] fhost = [ir_pass.LowerDeviceStorageAccessInfo(x) for x in fhost] fdevice = [ir_pass.LowerIntrin(x, target.target_name) for x in fdevice] fhost = [ir_pass.LowerIntrin(x, target_host.target_name) for x in fhost] fhost = [ir_pass.CombineContextCall(x) for x in fhost] mdev = codegen.build_module(fdevice, str(target)) if fdevice else None return fhost, mdev