基础算子#
import numpy as np
import tvm
from tvm.relax.frontend.nn import Module, Tensor, spec
from tvm.script import relax as R
张量二元运算符测试#
class Model(Module):
def test(self, x: Tensor, y: Tensor):
z0 = x + y
z1 = x * y
z2 = x / y
z3 = x.maximum(y)
z4 = x.minimum(y)
return (z0, z1, z2, z3, z4)
m = Model()
irmodule, _ = m.export_tvm(
spec={"test": {"x": spec.Tensor([1, 10], "float32"), "y": spec.Tensor([2, 1], "float32")}},
debug=True,
)
irmodule.show()
Show code cell output
# from tvm.script import ir as I
# from tvm.script import relax as R
@I.ir_module
class Module:
@R.function
def _initialize_effect() -> R.Tuple(R.Object):
with R.dataflow():
_io: R.Object = R.null_value()
lv: R.Tuple(R.Object) = (_io,)
gv: R.Tuple(R.Object) = lv
R.output(gv)
return gv
@R.function
def test(x: R.Tensor((1, 10), dtype="float32"), y: R.Tensor((2, 1), dtype="float32"), _io: R.Object) -> R.Tuple(R.Tuple(R.Tensor((2, 10), dtype="float32"), R.Tensor((2, 10), dtype="float32"), R.Tensor((2, 10), dtype="float32"), R.Tensor((2, 10), dtype="float32"), R.Tensor((2, 10), dtype="float32")), R.Tuple(R.Object)):
R.func_attr({"num_input": 3})
with R.dataflow():
add: R.Tensor((2, 10), dtype="float32") = R.add(x, y)
mul: R.Tensor((2, 10), dtype="float32") = R.multiply(x, y)
divide: R.Tensor((2, 10), dtype="float32") = R.divide(x, y)
maximum: R.Tensor((2, 10), dtype="float32") = R.maximum(x, y)
minimum: R.Tensor((2, 10), dtype="float32") = R.minimum(x, y)
gv1: R.Tuple(R.Tuple(R.Tensor((2, 10), dtype="float32"), R.Tensor((2, 10), dtype="float32"), R.Tensor((2, 10), dtype="float32"), R.Tensor((2, 10), dtype="float32"), R.Tensor((2, 10), dtype="float32")), R.Tuple(R.Object)) = (add, mul, divide, maximum, minimum), (_io,)
R.output(gv1)
return gv1
张量与标量的二元运算符测试#
class Model(Module):
def test(self, x: Tensor):
y = 10
z0 = x + y
z1 = y + x
z2 = x * y
z3 = x / y
z4 = x.maximum(y)
z5 = x.minimum(y)
return (z0, z1, z2, z3, z4, z5)
m = Model()
irmodule, _ = m.export_tvm(spec={"test": {"x": spec.Tensor([1, 10], "float32")}}, debug=True)
irmodule.show()
Show code cell output
# from tvm.script import ir as I
# from tvm.script import relax as R
@I.ir_module
class Module:
@R.function
def _initialize_effect() -> R.Tuple(R.Object):
with R.dataflow():
_io: R.Object = R.null_value()
lv: R.Tuple(R.Object) = (_io,)
gv: R.Tuple(R.Object) = lv
R.output(gv)
return gv
@R.function
def test(x: R.Tensor((1, 10), dtype="float32"), _io: R.Object) -> R.Tuple(R.Tuple(R.Tensor((1, 10), dtype="float32"), R.Tensor((1, 10), dtype="float32"), R.Tensor((1, 10), dtype="float32"), R.Tensor((1, 10), dtype="float32"), R.Tensor((1, 10), dtype="float32"), R.Tensor((1, 10), dtype="float32")), R.Tuple(R.Object)):
R.func_attr({"num_input": 2})
with R.dataflow():
add: R.Tensor((1, 10), dtype="float32") = R.add(x, R.const(10.0, "float32"))
add1: R.Tensor((1, 10), dtype="float32") = R.add(x, R.const(10.0, "float32"))
mul: R.Tensor((1, 10), dtype="float32") = R.multiply(x, R.const(10.0, "float32"))
divide: R.Tensor((1, 10), dtype="float32") = R.divide(x, R.const(10.0, "float32"))
maximum: R.Tensor((1, 10), dtype="float32") = R.maximum(x, R.const(10.0, "float32"))
minimum: R.Tensor((1, 10), dtype="float32") = R.minimum(x, R.const(10.0, "float32"))
gv1: R.Tuple(R.Tuple(R.Tensor((1, 10), dtype="float32"), R.Tensor((1, 10), dtype="float32"), R.Tensor((1, 10), dtype="float32"), R.Tensor((1, 10), dtype="float32"), R.Tensor((1, 10), dtype="float32"), R.Tensor((1, 10), dtype="float32")), R.Tuple(R.Object)) = (add, add1, mul, divide, maximum, minimum), (_io,)
R.output(gv1)
return gv1
类型转换#
class Model(Module):
def test(self, x: Tensor):
z0 = x.astype(dtype="float16")
return z0
m = Model()
irmodule, _ = m.export_tvm(spec={"test": {"x": spec.Tensor([1, 10], "float32")}}, debug=True)
irmodule.show()
Show code cell output
# from tvm.script import ir as I
# from tvm.script import relax as R
@I.ir_module
class Module:
@R.function
def _initialize_effect() -> R.Tuple(R.Object):
with R.dataflow():
_io: R.Object = R.null_value()
lv: R.Tuple(R.Object) = (_io,)
gv: R.Tuple(R.Object) = lv
R.output(gv)
return gv
@R.function
def test(x: R.Tensor((1, 10), dtype="float32"), _io: R.Object) -> R.Tuple(R.Tensor((1, 10), dtype="float16"), R.Tuple(R.Object)):
R.func_attr({"num_input": 2})
with R.dataflow():
astype: R.Tensor((1, 10), dtype="float16") = R.astype(x, dtype="float16")
gv1: R.Tuple(R.Tensor((1, 10), dtype="float16"), R.Tuple(R.Object)) = astype, (_io,)
R.output(gv1)
return gv1
算子运算#
class Model(Module):
def test(self, x: Tensor):
z0 = x.reshape(2, 5, 2)
z1 = x.permute_dims(2, 1, 0)
z2 = x.repeat(2, axis=1)
return (z0, z1, z2)
m = Model()
irmodule, _ = m.export_tvm(spec={"test": {"x": spec.Tensor([2, 1, 10], "float32")}}, debug=True)
irmodule.show()
Show code cell output
# from tvm.script import ir as I
# from tvm.script import relax as R
@I.ir_module
class Module:
@R.function
def _initialize_effect() -> R.Tuple(R.Object):
with R.dataflow():
_io: R.Object = R.null_value()
lv: R.Tuple(R.Object) = (_io,)
gv: R.Tuple(R.Object) = lv
R.output(gv)
return gv
@R.function
def test(x: R.Tensor((2, 1, 10), dtype="float32"), _io: R.Object) -> R.Tuple(R.Tuple(R.Tensor((2, 5, 2), dtype="float32"), R.Tensor((10, 1, 2), dtype="float32"), R.Tensor((2, 2, 10), dtype="float32")), R.Tuple(R.Object)):
R.func_attr({"num_input": 2})
with R.dataflow():
reshape: R.Tensor((2, 5, 2), dtype="float32") = R.reshape(x, R.shape([2, 5, 2]))
permute_dims: R.Tensor((10, 1, 2), dtype="float32") = R.permute_dims(x, axes=[2, 1, 0])
repeat: R.Tensor((2, 2, 10), dtype="float32") = R.repeat(x, repeats=2, axis=1)
gv1: R.Tuple(R.Tuple(R.Tensor((2, 5, 2), dtype="float32"), R.Tensor((10, 1, 2), dtype="float32"), R.Tensor((2, 2, 10), dtype="float32")), R.Tuple(R.Object)) = (reshape, permute_dims, repeat), (_io,)
R.output(gv1)
return gv1