基础算子

基础算子#

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()

Hide 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()

Hide 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()

Hide 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()

Hide 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
目录