def create_executor(kind="debug", mod=None, device=None, target="llvm", params=None):
"""Factory function to create an executor.
Example
-------
.. code-block:: python
import tvm.relay
import numpy as np
x = tvm.relay.var("x", tvm.relay.TensorType([1], dtype="float32"))
expr = tvm.relay.add(x, tvm.relay.Constant(tvm.nd.array(np.array([1], dtype="float32"))))
tvm.relay.create_executor(
kind="vm", mod=tvm.IRModule.from_expr(tvm.relay.Function([x], expr))
).evaluate()(np.array([2], dtype="float32"))
# returns `array([3.], dtype=float32)`
Parameters
----------
kind : str
The type of executor. Avaliable options are `debug` for the
interpreter, `graph` for the graph executor, and `vm` for the virtual
machine.
mod : :py:class:`~tvm.IRModule`
The Relay module containing collection of functions
device : :py:class:`Device`
The device to execute the code.
target : :py:class:`tvm.Target`
The corresponding context
params : dict of str to NDArray
Input parameters to the graph that do not change
during inference time.
Returns
-------
executor : :py:class:`~tvm.relay.backend.interpreter.Executor`
"""
if mod is None:
mod = IRModule()
if device is not None:
assert device.device_type == _nd.device(str(target), 0).device_type
else:
device = _nd.device(str(target), 0)
if params is not None:
mod = IRModule.from_expr(bind_params_by_name(mod["main"], params))
if isinstance(target, str):
target = Target(target)
if kind == "debug":
return _interpreter.Interpreter(mod, device, target)
if kind == "graph":
return GraphExecutor(mod, device, target)
if kind == "vm":
return VMExecutor(mod, device, target)
raise RuntimeError("unknown execution strategy: {0}".format(kind))
def _mod(mod: Union[PrimFunc, IRModule]) -> IRModule:
if isinstance(mod, PrimFunc):
mod = mod.with_attr("global_symbol", "main")
mod = mod.with_attr("tir.noalias", True)
mod = IRModule({"main": mod})
if not isinstance(mod, IRModule):
raise TypeError(f"Expected `mod` to be PrimFunc or IRModule, but gets: {mod}")
# in order to make sure the mod can be found in ApplyHistoryBest
# different func name can cause structural unequal
func_names = mod.get_global_vars()
(func_name,) = func_names
if len(func_names) == 1 and func_name != "main":
mod = IRModule({"main": mod[func_name]})
return mod
def from_auto_ml(model, shape=None, dtype="float32", func_name="transform"):
"""
Import scikit-learn model to Relay.
"""
try:
import sklearn # pylint: disable=unused-import
except ImportError as e:
raise ImportError(
"Unable to import scikit-learn which is required {}".format(e))
if func_name == "transform":
inexpr_float = _expr.var("input_float", shape=shape, dtype=dtype)
inexpr_string = _expr.var("input_string", shape=shape, dtype=dtype)
inexpr = [inexpr_float, inexpr_string]
if type(model.feature_transformer.steps[0]
[1]).__name__ != "ColumnTransformer":
raise NameError(
"The First Transformer must be an ColumnTransformer, but {} is given"
.format(type(model.feature_transformer.steps[0][1]).__name__))
outexpr = inexpr
for _, transformer in model.feature_transformer.steps:
outexpr = sklearn_op_to_relay(transformer, outexpr, shape, dtype,
func_name, None)
else:
inexpr = _expr.var("input", shape=shape, dtype=dtype)
transformer = model.target_transformer
outexpr = sklearn_op_to_relay(transformer, inexpr, shape, dtype,
func_name, None)
func = _function.Function(analysis.free_vars(outexpr), outexpr)
return IRModule.from_expr(func), []
def check(dim, axis, nstep):
eps = 0.01
ttype1 = rly.TensorType(tuple(10 for i in range(dim)), dtype)
ttype2 = rly.TensorType((10, ), dtype)
x = rly.var("x", ttype1)
beta = rly.var("beta", ttype2)
gamma = rly.var("gamma", ttype2)
moving_var = rly.var("moving_var", ttype2)
moving_mean = rly.var("moving_mean", ttype2)
y1, y2 = x, x
for _ in range(nstep):
y1, _, _ = rly.nn.batch_norm(y1 + rly.const(1, dtype),
gamma,
beta,
moving_mean,
moving_var,
epsilon=eps,
axis=axis)
y1 = rly.nn.dropout(y1)
y2 = simple_bn(y2 + rly.const(1, dtype),
gamma,
beta,
moving_mean,
moving_var,
epsilon=eps,
axis=axis,
shape=ttype1.shape)
mod = IRModule.from_expr(y1)
simplify = SimplifyInference()
mod = simplify(mod)
y1 = mod["main"].body
assert rly.analysis.graph_equal(y1, y2)
def _parse_mod(mod: Union[PrimFunc, IRModule]) -> IRModule:
if isinstance(mod, PrimFunc):
mod = IRModule({"main": mod})
if not isinstance(mod, IRModule):
raise TypeError(
f"Expected `mod` to be PrimFunc or IRModule, but gets: {mod}")
return mod
def from_program(self, program, shape_dict, scope):
"""Construct the TVM relay expression from PaddlePaddle program."""
self.shape_dict = shape_dict
if scope is None:
import paddle
scope = paddle.fluid.global_scope()
self.check_unsupported_ops(program)
self.extract_parameters(program, scope)
self.ops_to_relay(program)
output_names = list()
for block in program.blocks:
for op in block.ops:
if op.type == "fetch":
output_names.append(op.input("X")[0])
outputs = [self.nodes[name] for name in output_names]
outputs = outputs[0] if len(outputs) == 1 else _expr.Tuple(outputs)
free_vars = analysis.free_vars(outputs)
func = _function.Function(free_vars, outputs)
mod = IRModule.from_expr(func)
return mod, self.params
def create_executor(kind="debug", mod=None, ctx=None, target="llvm"):
"""Factory function to create an executor.
Parameters
----------
kind : str
The type of executor
mod : :py:class:`~tvm.IRModule`
The Relay module containing collection of functions
ctx : :py:class:`tvmContext`
The context to execute the code.
target : :py:class:`tvm.Target`
The corresponding context
"""
if mod is None:
mod = IRModule()
if ctx is not None:
assert ctx.device_type == _nd.context(str(target), 0).device_type
else:
ctx = _nd.context(str(target), 0)
if isinstance(target, str):
target = _target.create(target)
if kind == "debug":
return _interpreter.Interpreter(mod, ctx, target)
if kind == "graph":
return GraphExecutor(mod, ctx, target)
if kind == "vm":
return VMExecutor(mod, ctx, target)
raise RuntimeError("unknown execution strategy: {0}".format(kind))
def replace_ir_builder_module(deep_copy=False, realize=False):
new_func = tvm.script.from_source(tvm.script.asscript(elementwise))
other_func = tvm.script.from_source(tvm.script.asscript(elementwise))
mod = IRModule(functions={"main": new_func, "other": other_func})
s = tir.ScheduleState(mod, debug_mode=True)
target = tvm.tir.Block(
iter_vars=[],
reads=[],
writes=[],
name_hint="target",
body=s.mod["main"].body.block.body[1],
init=None,
alloc_buffers=None,
match_buffers=None,
annotations=None,
)
if realize:
target = tvm.tir.BlockRealize(
iter_values=[],
predicate=True,
block=target,
)
if deep_copy:
target.__setstate__(target.__getstate__())
gc.collect()
return s, target
def check_sketches(
mod: IRModule,
sketches: List[Schedule],
expected_mods: List[IRModule],
expected_decisions: List[List[Tuple[str, List[int]]]],
*,
debug_mask="all",
):
assert len(expected_mods) == len(expected_decisions)
assert len(sketches) == len(expected_mods)
expected_mods = [
IRModule({"main": m}) if not isinstance(m, IRModule) else m for m in expected_mods
]
sketches = list(sketches)
for expected_id, (expected_mod, expected_decision) in enumerate(
zip(expected_mods, expected_decisions)
):
sketch_id = _find_match_sketch_id(
mod,
sketches,
expected_mod,
expected_decision,
debug_mask=debug_mask,
)
if sketch_id is None:
raise AssertionError(
f"Expected sketch #{expected_id} doesn't exist in the generated sketches."
)
sketches.pop(sketch_id)
def from_translated_layer(self, layer, shape_dict):
"""Construct the TVM relay expression from PaddlePaddle TranslatedLayer."""
self.shape_dict = shape_dict
program = layer.program()
parameters = dict()
for param in layer.parameters():
parameters[param.name] = np.array(param.value().get_tensor())
self.check_unsupported_ops(program)
self.extract_parameters(program, parameters)
input_specs = layer._input_spec()
self.ops_to_relay(program, input_specs)
output_names = [x.name for x in layer._output_spec()]
outputs = [self.nodes[name] for name in output_names]
outputs = outputs[0] if len(outputs) == 1 else _expr.Tuple(outputs)
free_vars = analysis.free_vars(outputs)
func = _function.Function(free_vars, outputs)
mod = IRModule.from_expr(func)
# remove unused parameters
final_params = dict()
for var in free_vars:
if var.name_hint in self.params:
final_params[var.name_hint] = self.params[var.name_hint]
self.params = final_params
return mod, self.params
def __init__(
self,
mod: Union[PrimFunc, IRModule],
debug_mode: Union[bool, int] = False,
) -> None:
"""Construct a schedule state from an IRModule or a PrimFunc
Parameters
----------
mod : Union[PrimFunc, IRModule]
The IRModule or PrimFunc to be scheduled
debug_mode : Union[bool, int]
Do extra correctness checking after the class creation and each time
after calling the Replace method.
Possible choices of `debug_mode`:
1) True - Turn on all the checks
2) False - Turn off all the checks
3) An integer - Turn on checks according to the bitmasks provided in ScheduleDebugMask
"""
if isinstance(mod, PrimFunc):
mod = IRModule({"main": mod})
if isinstance(debug_mode, bool):
if debug_mode:
debug_mode = -1
else:
debug_mode = 0
if not isinstance(debug_mode, int):
raise TypeError(
f"`debug_mode` should be integer or boolean, but gets: {debug_mode}"
)
self.__init_handle_by_constructor__(
_ffi_api.ScheduleState, # type: ignore # pylint: disable=no-member
mod,
debug_mode,
)
def from_darknet(self):
"""To convert the darknet symbol to relay functions."""
for i in range(self._net.n):
layer = self._net.layers[i]
need_skip, sym = self._preproc_layer(layer, i)
if need_skip:
continue
processed, sym = self._handle_darknet_rnn_layers(i, sym)
if processed:
continue
attr = self._get_darknet_attrs(layer, i)
op_name = self._get_opname(layer)
prefix = _get_params_prefix(op_name, i)
params = self._get_darknet_params(self._net.layers[i], prefix)
sym = _darknet_convert_symbol(op_name, _as_list(sym), params, attr,
prefix)
if params:
self._tvmparams.update(params)
self._sym_array[i] = sym
self._make_outlist(sym, prefix, layer, i)
outputs = _as_list(sym) + self._outs
outputs = outputs[0] if len(outputs) == 1 else _expr.Tuple(outputs)
sym = _function.Function(analysis.free_vars(outputs), outputs)
return IRModule.from_expr(sym), self._tvmparams
def _set_params(mod, input_scale_func, weight_scale_func):
quantize_op = _op.get("relay.op.annotation.simulated_quantize")
cfg = quantize.current_qconfig()
const_params = {}
def visit_func(expr):
'''visitor function for traverse'''
if isinstance(expr, _expr.Call) and expr.op == quantize_op:
_, ndom_scale, nclip_min, nclip_max = expr.args
attrs = expr.attrs
kind = attrs.kind
nbit = cfg.get_nbit_by_kind(kind)
valid_bit = nbit - attrs.sign
# set scale
if kind == quantize.QAnnotateKind.WEIGHT:
assert isinstance(expr.args[0], _expr.Constant)
scale = weight_scale_func(expr)
else:
scale = input_scale_func(expr)
def _make_const(val):
return _expr.const(val, 'float32')
valid_range = 2**valid_bit
const_params[ndom_scale] = _make_const(scale / valid_range)
const_params[nclip_min] = _make_const(- (valid_range - 1))
const_params[nclip_max] = _make_const((valid_range - 1))
func = mod['main']
_analysis.post_order_visit(func, visit_func)
func = _expr.bind(func, const_params)
return IRModule.from_expr(func)
def _interp_wrapper(*args, **kwargs):
if expr is None:
args = self._convert_args(self.mod["main"], args, kwargs)
else:
args = self._convert_args(expr, args, kwargs)
relay_args = []
for arg in args:
relay_args.append(_arg_to_ast(self.mod, arg))
# Set the entry function for the module.
if expr is None:
pass
elif isinstance(expr, GlobalVar):
self.mod["main"] = self.mod[expr]
else:
assert isinstance(expr, Function)
func = Function([], Call(expr, relay_args))
relay_args = []
if self.mod:
self.mod["main"] = func
else:
self.mod = IRModule.from_expr(func)
mod = self.optimize()
opt_expr = Call(mod["main"], relay_args)
return _intrp(opt_expr)
def infer_value(input_val, params, mod=None):
"""A hack for getting the value of an expression by evaluating a
portion of the relay graph. This is often needed for functions that
whose output shape depends on the value of a tensor.
"""
# Check that all free variables have associated parameters.
assert all(
var.name_hint in params.keys() for var in analysis.free_vars(input_val)
), "All inputs to infer must be available in params."
try:
# TODO(kevinthesun): Use VM for all cases.
# pylint: disable=import-outside-toplevel
from tvm.contrib import graph_runtime
func = _function.Function(analysis.free_vars(input_val), input_val)
with tvm.transform.PassContext(opt_level=0):
lib = tvm.relay.build(func, target="llvm", params=params)
ctx = tvm.cpu(0)
m = graph_runtime.GraphModule(lib["default"](ctx))
m.run()
return m.get_output(0)
except Exception:
if isinstance(mod, IRModule):
mod["main"] = _function.Function(analysis.free_vars(input_val), input_val)
else:
mod = IRModule.from_expr(input_val)
exc = tvm.relay.create_executor("debug", mod=mod, ctx=tvm.cpu(), target="llvm")
inputs = []
for param in mod["main"].params:
inputs.append(params[param.name_hint])
result = exc.evaluate()(*inputs)
return result
def __init__(self, shape, dtype):
self._nodes = {}
self._params = {}
self._visited_nodes = set()
self._ops = {}
self._shape = shape
self._dtype = dtype
self._mod = IRModule({})
def infer_type(node, mod=None):
"""A method to infer the type of an intermediate node in the relay graph."""
new_mod = IRModule.from_expr(node)
if mod is not None:
new_mod.update(mod)
new_mod = _transform.InferType()(new_mod)
entry = new_mod["main"]
return entry if isinstance(node, _function.Function) else entry.body
def test_tir_schedule_creation():
# Tests:
# - Schedule.__init__ for PrimFunc and IRModule
# - Schedule.mod
# - Schedule.state
sch_1 = tir.Schedule(matmul, debug_mode=True)
sch_2 = tir.Schedule(IRModule({"main": matmul}), debug_mode=True)
assert sch_1.mod["main"].same_as(sch_2.mod["main"])
assert sch_1.state.mod["main"].same_as(sch_2.state.mod["main"])
def __init__(self, shape, dtype, net, model):
#网络和参数
self._NetLayer = self.__getNetLayer(net)
self._ModelLayer = self.__getModelLayer(model)
self._params = {}
self._nodes = {}
self._LayerList = []
self._shape = shape
self._dtype = dtype
self._mod = IRModule({})
def create_executor(kind="debug", mod=None, ctx=None, target="llvm"):
"""Factory function to create an executor.
Example
-------
.. code-block:: python
import tvm.relay
import numpy as np
x = tvm.relay.var("x", tvm.relay.TensorType([1], dtype="float32"))
expr = tvm.relay.add(x, tvm.relay.Constant(tvm.nd.array(np.array([1], dtype="float32"))))
tvm.relay.create_executor(
kind="vm", mod=tvm.IRModule.from_expr(tvm.relay.Function([x], expr))
).evaluate()(np.array([2], dtype="float32"))
# returns `array([3.], dtype=float32)`
Parameters
----------
kind : str
The type of executor. Avaliable options are `debug` for the
interpreter, `graph` for the graph runtime, and `vm` for the virtual
machine.
mod : :py:class:`~tvm.IRModule`
The Relay module containing collection of functions
ctx : :py:class:`tvmContext`
The context to execute the code.
target : :py:class:`tvm.Target`
The corresponding context
Returns
-------
executor : :py:class:`~tvm.relay.backend.interpreter.Executor`
"""
if mod is None:
mod = IRModule()
if ctx is not None:
assert ctx.device_type == _nd.context(str(target), 0).device_type
else:
ctx = _nd.context(str(target), 0)
if isinstance(target, str):
target = Target(target)
if kind == "debug":
return _interpreter.Interpreter(mod, ctx, target)
if kind == "graph":
return GraphExecutor(mod, ctx, target)
if kind == "vm":
return VMExecutor(mod, ctx, target)
raise RuntimeError("unknown execution strategy: {0}".format(kind))
def __init__(self, source_name: str) -> None:
self.source_name = source_name
self.module = IRModule({}) # type: IRModule
# Adding an empty scope allows naked lets without pain.
self.var_scopes = deque([deque()]) # type: Scopes[expr.Var]
self.global_vars = {} # type: Scope[expr.GlobalVar]
self.type_var_scopes = deque([deque()]) # type: Scopes[ty.TypeVar]
self.global_type_vars = {} # type: Scope[expr.GlobalVar]
self.graph_expr = [] # type: List[expr.Expr]
super(ParseTreeToRelayIR, self).__init__()
def optimize(mod, target=None, params=None):
"""Helper function that optimizes a Relay module.
Parameters
----------
mod : :py:class:`~tvm.IRModule`
The module to build. Using relay.Function is deprecated.
target : None, or any multi-target like object, see Target.canon_multi_target
For homogeneous compilation, the unique build target.
For heterogeneous compilation, a dictionary or list of possible build targets.
Defaults to the current target in the environment if None.
params : dict of str to NDArray
Input parameters to the graph that do not change
during inference time. Used for constant folding.
Returns
-------
mod : :py:class:`~tvm.IRModule`
The optimized relay module.
params : dict
The parameters of the final graph.
"""
if not isinstance(mod, (IRModule, _function.Function)):
raise ValueError("Type of input parameter mod must be tvm.IRModule")
if isinstance(mod, _function.Function):
if params:
mod = bind_params_by_name(mod, params)
mod = IRModule.from_expr(mod)
warnings.warn(
"Please use input parameter mod (tvm.IRModule) "
"instead of deprecated parameter func (tvm.relay.function.Function)",
DeprecationWarning,
)
raw_targets = Target.canon_multi_target_and_host(
Target.target_or_current(target))
# If current dispatch context is fallback context (the default root context),
# then load pre-tuned parameters from TopHub
if isinstance(autotvm.DispatchContext.current, autotvm.FallbackContext):
tophub_context = autotvm.tophub.context(raw_targets)
else:
tophub_context = autotvm.utils.EmptyContext()
with tophub_context:
bld_mod = BuildModule()
mod, params = bld_mod.optimize(mod, target=raw_targets, params=params)
return mod, params
def optimize(mod, target=None, params=None):
"""Helper function that optimizes a Relay module.
Parameters
----------
mod : :py:class:`~tvm.IRModule`
The module to build. Using relay.Function is deprecated.
target : str, :any:`tvm.target.Target`, or dict of str(i.e. device/context
name) to str/tvm.target.Target, optional
For heterogeneous compilation, it is a dictionary indicating context to
target mapping. For homogeneous compilation, it is a build target.
params : dict of str to NDArray
Input parameters to the graph that do not change
during inference time. Used for constant folding.
Returns
-------
mod : :py:class:`~tvm.IRModule`
The optimized relay module.
params : dict
The parameters of the final graph.
"""
if not isinstance(mod, (IRModule, _function.Function)):
raise ValueError("Type of input parameter mod must be tvm.IRModule")
if isinstance(mod, _function.Function):
if params:
mod = bind_params_by_name(mod, params)
mod = IRModule.from_expr(mod)
warnings.warn(
"Please use input parameter mod (tvm.IRModule) "
"instead of deprecated parameter func (tvm.relay.function.Function)",
DeprecationWarning,
)
target = _update_target(target)
# If current dispatch context is fallback context (the default root context),
# then load pre-tuned parameters from TopHub
if isinstance(autotvm.DispatchContext.current, autotvm.FallbackContext):
tophub_context = autotvm.tophub.context(list(target.values()))
else:
tophub_context = autotvm.utils.EmptyContext()
with tophub_context:
bld_mod = BuildModule()
mod, params = bld_mod.optimize(mod, target, params)
return mod, params
def visit_call(self, call):
# Keep track of our current depth and layer count
# so we can know whether to skip this layer or not.
current_depth = self.depth_count
current_layer = self.valid_op_count - current_depth - 1
if call.op in self.valid_ops:
self.depth_count += 1
# Visit current call operation
new_fn = self.visit(call.op)
# Visit current arguments
args = []
for arg in call.args:
args.append(self.visit(arg))
self.depth_count = current_depth
# Downcast this op if its the correct type and not skipped.
if call.op in self.valid_ops and current_layer not in self.skip_layers:
# Recast inputs to specified type.
args = [self.visit(arg) for arg in call.args]
new_args = list()
for arg in args:
new_args.append(relay.cast(arg, dtype=self.dtype))
# If out_dtype is in the attributes, we need to update it.
orig_dtype = None
if call.attrs is not None and "out_dtype" in call.attrs.keys():
new_attr_dict = {}
for attr in call.attrs.keys():
attr_value = call.attrs[attr]
if isinstance(attr_value, tvm.ir.container.Array):
attr_value = tuple(attr_value)
new_attr_dict[str(attr)] = attr_value
new_attr_dict["out_dtype"] = self.out_dtype
attr_type = str(call.attrs).split("(")[0]
new_attrs = tvm.ir.make_node(attr_type, **new_attr_dict)
if call.attrs["out_dtype"] != "":
orig_dtype = call.attrs["out_dtype"]
else:
new_attrs = call.attrs
if orig_dtype is None:
# Perform type inference to determine the original type.
new_mod = IRModule.from_expr(call)
new_mod = InferType()(new_mod)
checked_arg = new_mod["main"].body
orig_dtype = checked_arg.checked_type.dtype
# Recast the output for compatibility with other graph operations.
return relay.cast(Call(new_fn, new_args, new_attrs), orig_dtype)
# Otherwise return the unchanged call.
return Call(new_fn, args, call.attrs)
def preprocess_mod(mod, params):
if not isinstance(mod, (IRModule, _function.Function)):
raise ValueError("Type of input parameter mod must be tvm.IRModule")
if isinstance(mod, _function.Function):
if params:
mod = bind_params_by_name(mod, params)
mod = IRModule.from_expr(mod)
warnings.warn(
"Please use input parameter mod (tvm.IRModule) "
"instead of deprecated parameter mod (tvm.relay.function.Function)",
DeprecationWarning,
)
return mod
def __init__(
self,
mod: Union[PrimFunc, IRModule],
*,
debug_mode: Union[bool, int] = False,
error_render_level: ERROR_RENDER_LEVEL_CANDIDATES = "detail",
) -> None:
"""Construct a concrete TensorIR schedule from an IRModule or a PrimFunc
Parameters
----------
mod : Union[PrimFunc, IRModule]
The IRModule or PrimFunc to be scheduled
debug_mode : Union[bool, int]
Do extra correctness checking after the class creation and each time
scheduling primitive
error_render_level : str = "detail"
The level of error rendering. Choices: "detail", "fast", "none".
"detail": Render a detailed error message, with the TIR and error locations printed
"fast: Show a simple error message without rendering or string manipulation
"none": Do not show any error message.
Note
----
The checks performed includes:
1) VerifySRefTree
2) VerifyCachedFlags
"""
if isinstance(mod, PrimFunc):
mod = IRModule({"main": mod})
if isinstance(debug_mode, bool):
if debug_mode:
debug_mode = -1
else:
debug_mode = 0
if not isinstance(debug_mode, int):
raise TypeError(
f"`debug_mode` should be integer or boolean, but gets: {debug_mode}"
)
if error_render_level not in Schedule.ERROR_RENDER_LEVEL:
raise ValueError(
'error_render_level can be "detail", "fast", or "none", but got: '
+ f"{error_render_level}")
self.__init_handle_by_constructor__(
_ffi_api.ConcreteSchedule, # type: ignore # pylint: disable=no-member
mod,
debug_mode,
Schedule.ERROR_RENDER_LEVEL.get(error_render_level),
)
def test_conv2d_winograd_cuda():
mod = conv2d_winograd_cuda
mod = IRModule({"main": mod})
context = TuneContext(
mod=mod,
target=Target("nvidia/geforce-rtx-3090", host="llvm"),
task_name="Custom Search Space Task",
sch_rules=DefaultCUDA._sch_rules(), # pylint: disable=protected-access
)
for sch_rule in context.sch_rules:
sch_rule.initialize_with_tune_context(context)
post_order_apply = PostOrderApply()
post_order_apply.initialize_with_tune_context(context)
(sch,) = post_order_apply.generate_design_space(mod)
decisions = dict(
zip(
[i for i in sch.trace.insts if i.kind.name.startswith("Sample")],
[
# data_pack
[3, 3],
[64, 2],
2,
# inverse
[3, 3],
[2, 64],
2,
# bgemm
[1, 1, 1, 1, 6],
[1, 1, 1, 3, 2],
[3, 1, 1, 1, 3],
[4, 2, 1, 4, 4],
[32, 1, 4],
1,
1,
# root anno
2,
# conv2d
2,
],
)
)
trace = Trace(sch.trace.insts, decisions=decisions)
sch = Schedule(mod=mod)
trace.apply_to_schedule(sch, remove_postproc=False)
answer = sch.mod
expected = _get_mod()
tvm.ir.assert_structural_equal(answer, expected)
def from_sklearn(model, shape=None, dtype="float32", func_name="transform", columns=None):
"""
Import scikit-learn model to Relay.
"""
try:
import sklearn # pylint: disable=unused-import
except ImportError as e:
raise ImportError("Unable to import scikit-learn which is required {}".format(e))
if type(model).__name__ == "ColumnTransformer":
raise NameError("ColumnTransformer is not supported for single op compilation.")
inexpr = _expr.var("input", shape=shape, dtype=dtype)
outexpr = sklearn_op_to_relay(model, inexpr, shape, dtype, func_name, columns)
func = _function.Function(analysis.free_vars(outexpr), outexpr)
return IRModule.from_expr(func), []
def infer_type(node, mod=None):
"""A method to infer the type of an intermediate node in the relay graph."""
if isinstance(mod, IRModule):
mod["main"] = _function.Function(tvm.relay.analysis.free_vars(node), node)
mod = _transform.InferType()(mod)
entry = mod["main"]
ret = entry.body
else:
new_mod = IRModule.from_expr(node)
if mod is not None:
new_mod.update(mod)
new_mod = _transform.InferType()(new_mod)
entry = new_mod["main"]
ret = entry if isinstance(node, _function.Function) else entry.body
return ret
def all_type_vars(expr, mod=None):
"""Get all type variables from expression/type e
Parameters
----------
expr : Union[tvm.relay.Expr,tvm.relay.Type]
The input expression/type
mod : Optional[tvm.IRModule]
The global module
Returns
-------
free : List[tvm.relay.TypeVar]
The list of all type variables in post-DFS order
"""
use_mod = mod if mod is not None else IRModule()
return _analysis.all_type_vars(expr, use_mod)
