后向折叠负缩放的测试用例

%cd ..
import set_env

/media/pc/data/lxw/ai/tvm-book/tests/book/doc/tests

import numpy as np

import tvm
from tvm import relay
from tvm.relay import transform
# from tvm.relay.testing import create_workload
# from tvm.relay.build_module import bind_params_by_name


def initializer(_, param):
    param = np.zeros(param.shape)


def _get_positive_scale(size):
    return np.random.uniform(0.5, 1, size=size).astype("float32")


def run_opt_pass(expr, opt_pass):
    assert isinstance(opt_pass, tvm.transform.Pass)
    mod = tvm.IRModule.from_expr(expr)
    mod = opt_pass(mod)
    entry = mod["main"]
    return entry if isinstance(expr, relay.Function) else entry.body

def before(x, conv_weight, out_scale, channels, blocking):
    args = [x, conv_weight]
    y = relay.nn.conv2d(
        x,
        conv_weight,
        channels=channels,
        kernel_size=(3, 3),
        padding=(1, 1),
        data_layout="NCHW{}c".format(blocking[0]) if blocking else "NCHW",
        kernel_layout="OIHW1i{}o".format(blocking[1]) if blocking else "OIHW",
    )
    y = relay.multiply(y, out_scale)
    return relay.Function(args, y)

def expected(x, conv_weight, out_scale, channels, blocking):
    # use a fixed order of args so alpha equal check can pass
    args = [x, conv_weight]
    if blocking:
        squeezed_scale = relay.squeeze(out_scale, axis=[0, 2, 3])
        conv_weight = relay.multiply(
            conv_weight,
            relay.reshape(squeezed_scale, (channels // blocking[1], 1, 1, 1, 1, blocking[1])),
        )
    else:
        squeezed_scale = relay.squeeze(out_scale, axis=[1, 2])
        conv_weight = relay.multiply(
            conv_weight, relay.expand_dims(squeezed_scale, axis=1, num_newaxis=3)
        )
    y = relay.nn.conv2d(
        x,
        conv_weight,
        channels=channels,
        kernel_size=(3, 3),
        padding=(1, 1),
        data_layout="NCHW{}c".format(blocking[0]) if blocking else "NCHW",
        kernel_layout="OIHW1i{}o".format(blocking[1]) if blocking else "OIHW",
    )
    return relay.Function(args, y)

def check(shape, channels, blocking):
    x = relay.var("x", shape=shape)
    weight = relay.var("weight")
    if blocking:
        out_scale = relay.const(
            -_get_positive_scale((1, channels // blocking[1], 1, 1, blocking[1]))
        )
    else:
        out_scale = relay.const(-_get_positive_scale((channels, 1, 1)))
    y1 = before(x, weight, out_scale, channels, blocking)
    y1 = run_opt_pass(y1, transform.InferType())
    print("FoldScaleAxis 前：")
    tvm.IRModule.from_expr(y1).show()

    type_dict = {x.name_hint: x.checked_type for x in y1.params}
    weight = relay.var("weight", type_dict["weight"])
    y1_folded = run_opt_pass(y1, transform.BackwardFoldScaleAxis())
    print("FoldScaleAxis 后：")
    tvm.IRModule.from_expr(y1_folded).show()
    
    y1_expected = expected(x, weight, out_scale, channels, blocking)
    y1_expected = run_opt_pass(y1_expected, transform.InferType())
    tvm.ir.assert_structural_equal(y1_folded, y1_expected)

check((2, 4, 10, 10), 8, None)
check((2, 2, 10, 10, 2), 8, (2, 2))

FoldScaleAxis 前：
FoldScaleAxis 后：
FoldScaleAxis 前：
FoldScaleAxis 后：

def @main(%x: Tensor[(2, 4, 10, 10), float32] /* ty=Tensor[(2, 4, 10, 10), float32] */, %weight: Tensor[(8, 4, 3, 3), float32] /* ty=Tensor[(8, 4, 3, 3), float32] */) -> Tensor[(2, 8, 10, 10), float32] {
  %0 = nn.conv2d(%x, %weight, padding=[1, 1, 1, 1], channels=8, kernel_size=[3, 3]) /* ty=Tensor[(2, 8, 10, 10), float32] */;
  multiply(%0, meta[relay.Constant][0] /* ty=Tensor[(8, 1, 1), float32] */) /* ty=Tensor[(2, 8, 10, 10), float32] */
}

def @main(%x: Tensor[(2, 4, 10, 10), float32] /* ty=Tensor[(2, 4, 10, 10), float32] */, %weight: Tensor[(8, 4, 3, 3), float32] /* ty=Tensor[(8, 4, 3, 3), float32] */) -> Tensor[(2, 8, 10, 10), float32] {
  %0 = squeeze(meta[relay.Constant][0] /* ty=Tensor[(8, 1, 1), float32] */, axis=[1, 2]) /* ty=Tensor[(8), float32] */;
  %1 = expand_dims(%0, axis=1, num_newaxis=3) /* ty=Tensor[(8, 1, 1, 1), float32] */;
  %2 = multiply(%weight, %1) /* ty=Tensor[(8, 4, 3, 3), float32] */;
  nn.conv2d(%x, %2, padding=[1, 1, 1, 1], channels=8, kernel_size=[3, 3]) /* ty=Tensor[(2, 8, 10, 10), float32] */
}

def @main(%x: Tensor[(2, 2, 10, 10, 2), float32] /* ty=Tensor[(2, 2, 10, 10, 2), float32] */, %weight: Tensor[(4, 4, 3, 3, 1, 2), float32] /* ty=Tensor[(4, 4, 3, 3, 1, 2), float32] */) -> Tensor[(2, 4, 10, 10, 2), float32] {
  %0 = nn.conv2d(%x, %weight, padding=[1, 1, 1, 1], channels=8, kernel_size=[3, 3], data_layout="NCHW2c", kernel_layout="OIHW1i2o") /* ty=Tensor[(2, 4, 10, 10, 2), float32] */;
  multiply(%0, meta[relay.Constant][0] /* ty=Tensor[(1, 4, 1, 1, 2), float32] */) /* ty=Tensor[(2, 4, 10, 10, 2), float32] */
}

def @main(%x: Tensor[(2, 2, 10, 10, 2), float32] /* ty=Tensor[(2, 2, 10, 10, 2), float32] */, %weight: Tensor[(4, 4, 3, 3, 1, 2), float32] /* ty=Tensor[(4, 4, 3, 3, 1, 2), float32] */) -> Tensor[(2, 4, 10, 10, 2), float32] {
  %0 = squeeze(meta[relay.Constant][0] /* ty=Tensor[(1, 4, 1, 1, 2), float32] */, axis=[0, 2, 3]) /* ty=Tensor[(4, 2), float32] */;
  %1 = reshape(%0, newshape=[4, 1, 1, 1, 1, 2]) /* ty=Tensor[(4, 1, 1, 1, 1, 2), float32] */;
  %2 = multiply(%weight, %1) /* ty=Tensor[(4, 4, 3, 3, 1, 2), float32] */;
  nn.conv2d(%x, %2, padding=[1, 1, 1, 1], channels=8, kernel_size=[3, 3], data_layout="NCHW2c", kernel_layout="OIHW1i2o") /* ty=Tensor[(2, 4, 10, 10, 2), float32] */
}