def after():
data = relay.var("data", shape=(1, 32))
eq1 = relay.var("e1", shape=[], dtype="float32")
eq2 = relay.var("e2", shape=[], dtype="float32")
cb_1 = relay.annotation.compiler_begin(eq1, target)
cb_2 = relay.annotation.compiler_begin(eq2, target)
equality_condition = relay.equal(cb_1, cb_2)
ce_1 = relay.annotation.compiler_end(equality_condition, target)
# if condition
cb_3 = relay.annotation.compiler_begin(data, target)
true_branch = relay.tanh(cb_3)
ce_2 = relay.annotation.compiler_end(true_branch, target)
# else condition
cb_4 = relay.annotation.compiler_begin(data, target)
false_branch = relay.sigmoid(cb_4)
ce_3 = relay.annotation.compiler_end(false_branch, target)
if_condition = relay.If(ce_1, ce_2, ce_3)
cb_5 = relay.annotation.compiler_begin(if_condition, target)
erf_out = relay.erf(cb_5)
ce_4 = relay.annotation.compiler_end(erf_out, target)
func = relay.Function([data, eq1, eq2], ce_4)
mod = tvm.IRModule.from_expr(func)
return mod
def _erf_legalize(attrs, inputs, arg_types):
"""Legalizes ERF op if needed.
Parameters
----------
attrs : tvm.ir.Attrs
Attributes of current convolution
inputs : list of tvm.relay.Expr
The args of the Relay expr to be legalized
types : list of types
List of input and output types
Returns
-------
result : tvm.relay.Expr
The legalized expr
"""
# Extract types and expressions.
data = inputs[0]
data_tensor = arg_types[0]
# Check if the input type is supported.
data_dtype = data_tensor.dtype
# If input is not fp32, we must cast to it.
if data_dtype != "float32":
data = relay.cast(data, "float32")
output = relay.erf(data)
return relay.cast(output, data_dtype)
# Otherwise do nothing.
return None
def get_func_with_control_flow():
data = relay.var("data", shape=(1, 3, 224, 224))
weight = relay.var("weight", shape=(32, 3, 3, 3))
eq1 = relay.var("e1", shape=[], dtype="float32")
eq2 = relay.var("e2", shape=[], dtype="float32")
eq = relay.equal(eq1, eq2)
true_branch = relay.zeros(shape=(1, 32, 222, 222), dtype="float32")
false_branch = relay.nn.conv2d(data, weight, kernel_size=(3, 3), channels=32)
ife = relay.If(eq, true_branch, false_branch)
out = relay.erf(ife)
return relay.Function([data, weight, eq1, eq2], out)
def test_erf():
x = relay.var("x", shape=(1, 16, 16, 16), dtype="float32")
y = relay.erf(x)
func = relay.Function([x], y)
mod = tvm.IRModule.from_expr(func)
fast_mod = FastMath()(mod)
assert "fast_erf" in fast_mod.astext()
# Check that FastMath option works for relay.build.
with tvm.transform.PassContext(opt_level=3, required_pass=["FastMath"]):
fast_mod = relay.optimize(mod, target="llvm", params=None)
assert "fast_erf" in fast_mod[0].astext()
def before():
data = relay.var("data", shape=(1, 32))
eq1 = relay.var("e1", shape=[], dtype="float32")
eq2 = relay.var("e2", shape=[], dtype="float32")
eq = relay.equal(eq1, eq2)
true_branch = relay.tanh(data)
false_branch = relay.sigmoid(data)
ife = relay.If(eq, true_branch, false_branch)
out = relay.erf(ife)
func = relay.Function([data, eq1, eq2], out)
mod = tvm.IRModule.from_expr(func)
return mod
def test_tensorrt_not_compatible():
if should_skip():
return
dtype = 'float32'
xshape = (1, 32, 14, 14)
x = relay.var('x', shape=(xshape), dtype=dtype)
y = relay.add(x, x)
z = relay.erf(y)
out = relay.nn.relu(z)
f = relay.Function([x], out)
mod = relay.Module()
mod['main'] = f
mod = relay.tensorrt.EnableTrt(mod)
assert not mod['main'].attrs
def test_tensorrt_not_compatible():
if skip_codegen_test():
return
dtype = "float32"
xshape = (1, 32, 14, 14)
x_data = np.random.uniform(-1, 1, xshape).astype(dtype)
x = relay.var("x", shape=(xshape), dtype=dtype)
y = relay.add(x, x)
z = relay.erf(y)
out = relay.nn.relu(z)
f = relay.Function([x], out)
mod = tvm.IRModule()
mod["main"] = f
mod, config = tensorrt.partition_for_tensorrt(mod)
for mode in ["graph", "vm"]:
with tvm.transform.PassContext(opt_level=3, config={"relay.ext.tensorrt.options": config}):
exec = relay.create_executor(mode, mod=mod, ctx=tvm.gpu(0), target="cuda")
if not skip_runtime_test():
results = exec.evaluate()(x_data)
def test_tensorrt_not_compatible():
if skip_codegen_test():
return
dtype = "float32"
xshape = (1, 32, 14, 14)
x = relay.var("x", shape=(xshape), dtype=dtype)
y = relay.add(x, x)
z = relay.erf(y)
out = relay.nn.relu(z)
f = relay.Function([x], out)
mod = tvm.IRModule()
mod["main"] = f
mod, config = tensorrt.partition_for_tensorrt(mod)
with tvm.transform.PassContext(opt_level=3, config={"relay.ext.tensorrt.options": config}):
graph, lib, params = relay.build(mod, "cuda")
if skip_runtime_test():
return
mod = graph_runtime.create(graph, lib, ctx=tvm.gpu(0))
x_data = np.random.uniform(-1, 1, xshape).astype(dtype)
mod.run(x=x_data)
results = [mod.get_output(i).asnumpy() for i in range(mod.get_num_outputs())]
