分析提取中间表达式

import tvm
from tvm import relay

def get_conv_net():
    """This gets the net for:
          conv2d
          /  |
         /   |
    conv2d   |
        \    |
         \   |
        elemwise add
             |
             |
             |
           split
             |
             |
             |
        elemwise add
    """
    dshape = (1, 1, 5, 1)
    x = relay.var("x", shape=dshape)
    y = relay.nn.conv2d(x, relay.var("w1"), kernel_size=(3, 3), padding=(1, 1), channels=1)
    x1 = relay.nn.conv2d(y, relay.var("w2"), kernel_size=(3, 3), padding=(1, 1), channels=1)

    z = relay.add(y, x1)

    tuple_out = relay.op.split(z, indices_or_sections=1, axis=0)

    tuple_0_add = relay.add(tuple_out[0], relay.const(1, dtype="float32"))

    return tvm.IRModule.from_expr(tuple_0_add)


def get_conv2d():
    x = relay.var("x", shape=(1, 56, 56, 64))
    weight1 = relay.var("weight1", shape=(3, 3, 64, 32))
    y = relay.nn.conv2d(
        x,
        weight1,
        channels=32,
        kernel_size=(3, 3),
        padding=(1, 1),
        data_layout="NHWC",
        kernel_layout="HWIO",
    )
    return tvm.IRModule.from_expr(y)

dshape = (1, 1, 5, 1)

mod = get_conv_net()
mod.show()

def @main(%x: Tensor[(1, 1, 5, 1), float32], %w1, %w2) {
  %0 = nn.conv2d(%x, %w1, padding=[1, 1, 1, 1], channels=1, kernel_size=[3, 3]);
  %1 = nn.conv2d(%0, %w2, padding=[1, 1, 1, 1], channels=1, kernel_size=[3, 3]);
  %2 = add(%0, %1);
  %3 = split(%2, indices_or_sections=meta[runtime.BoxInt][0]);
  %4 = %3.0;
  add(%4, 1f)
}

relay.analysis.extract_intermdeiate_expr(mod, 0).show()

def @main(%x: Tensor[(1, 1, 5, 1), float32], %w1) {
  nn.conv2d(%x, %w1, padding=[1, 1, 1, 1], channels=1, kernel_size=[3, 3])
}

relay.analysis.extract_intermdeiate_expr(mod, 1).show()

def @main(%x: Tensor[(1, 1, 5, 1), float32], %w1, %w2) {
  %0 = nn.conv2d(%x, %w1, padding=[1, 1, 1, 1], channels=1, kernel_size=[3, 3]);
  nn.conv2d(%0, %w2, padding=[1, 1, 1, 1], channels=1, kernel_size=[3, 3])
}