mul 重写

import tvm
from tvm import relax
from tvm.relax.testing import nn
from tvm_book.nn import modules
from tvm.script import ir as I, relax as R, tir as T

测试当乘法算子的右操作数不是常量时的量化处理逻辑。

bb = relax.BlockBuilder()

with bb.function("main"):
    model = modules.Conv2D(16, 16, kernel_size=(3, 3), padding=(1, 1))
    x = nn.Placeholder((1, 16, 64, 64), dtype="float32", name="x")
    data = nn.Placeholder((1, 16, 1, 1), dtype="float32", name="const")
    multiplier = R.sigmoid(data)
    conv = model(x)
    act = R.nn.relu(data=conv)
    output = act * multiplier
    params = [x, data] + model.parameters()
    bb.emit_func_output(output, params)
mod = bb.get()

mod.show()

# from tvm.script import ir as I
# from tvm.script import relax as R

@I.ir_module
class Module:
    @R.function
    def main(x: R.Tensor((1, 16, 64, 64), dtype="float32"), const: R.Tensor((1, 16, 1, 1), dtype="float32"), bias: R.Tensor((16,)), weight: R.Tensor((16, 16, 3, 3))) -> R.Tensor((1, 16, 64, 64)):
        gv: R.Tensor((1, 16, 64, 64)) = conv2_d(x, bias, weight)
        gv1: R.Tensor((1, 16, 64, 64)) = R.nn.relu(gv)
        gv2: R.Tensor((1, 16, 1, 1), dtype="float32") = R.sigmoid(const)
        gv3: R.Tensor((1, 16, 64, 64)) = R.multiply(gv1, gv2)
        return gv3