Relay Transform

device_copy 是由 在 LowerTE 时处理的:

1. 一方面它被用来确定 expr 所在的 target, 以确定 topi strategy
2. 另一方面运行时 graph_executor 会根据这个标记做真正的数据拷贝.

tvm/src/relay/backend/te_compiler_cache.cc::// Set the name to `__copy`. It will be detected in graph runtime to perform

例如, graph_executor 在执行时碰到 `__copy` 时会调用设备相关的 CopyDataFromTo 等函数, 对于 opencl 来说就是 clEnqueueCopyBuffer:

void OpenCLWorkspace::CopyDataFromTo(
    DLTensor* from, DLTensor* to, TVMStreamHandle stream) {
    if (IsOpenCLDevice(from->device) && IsOpenCLDevice(to->device)) {
        const auto* from_desc =
            static_cast<const cl::BufferDescriptor*>(from->data);
        auto* to_desc = static_cast<cl::BufferDescriptor*>(to->data);
        clEnqueueCopyBuffer(
            this->GetQueue(to->device), from_desc->buffer, to_desc->buffer,
            from->byte_offset, to->byte_offset, nbytes, 0, nullptr, nullptr);
    }
    // ...
}

除了上面 __copy, 使用 opencl 等 target 在 set_input/get_output 时也会通过 CopyDataFromTo 与设备交换数据.

通过 FuseOps 显然可以减少函数调用的开销. 那么, 每次都把所有的 op 都合并到一个函数中不就可以了么?

每个 op 都有一个 schedule, 如果 schedule 不是 `相容` 的, 则它们是无法合并在一起的, 例如:

import tvm
from tvm import relay
from tvm.relay import transform
from tvm.relay.testing import run_opt_pass
import numpy as np


def test_fuse_simple():
    def before():
        x = relay.var("x", shape=(20, 20))
        y1 = relay.add(x, x)
        y2 = relay.add(y1, y1)
        y3 = relay.sum(y2, axis=-1)
        y = relay.add(y3, y3)

        return relay.Function([x], y)

    z = before()
    z = run_opt_pass(z, transform.FuseOps())
    print(z)

    with tvm.transform.PassContext(opt_level=0):
        graph, lib, params = relay.build(z, target="c", params=None)

    print(lib.get_source())


if __name__ == "__main__":
    test_fuse_simple()

fn (%x: Tensor[(20, 20), float32]) -> Tensor[(20), float32] { %2 = fn (%p01: Tensor[(20, 20), float32], Primitive=1) -> Tensor[(20), float32] { %0 = add(%p01, %p01) * ty=Tensor[(20, 20), float32] *; %1 = add(%0, %0) * ty=Tensor[(20, 20), float32] *; sum(%1, axis=[-1]) * ty=Tensor[(20), float32] * }; %3 = %2(%x) * ty=Tensor[(20), float32] *; %4 = fn (%p0: Tensor[(20), float32], Primitive=1) -> Tensor[(20), float32] { add(%p0, %p0) * ty=Tensor[(20), float32] * }; %4(%3) * ty=Tensor[(20), float32] * } // tvm target: c -keys=cpu -link-params=0 #define TVM_EXPORTS #include "tvm/runtime/c_runtime_api.h" #include "tvm/runtime/c_backend_api.h" #include <math.h> #ifdef __cplusplus extern "C" #endif TVM_DLL int32_t tvmgen_default_fused_add_add_sum(void* args, void* arg_type_ids, int32_t num_args, void* out_ret_value, void* out_ret_tcode, void* resource_handle) { void* arg0 = (((TVMValue*)args)[0].v_handle); int32_t arg0_code = ((int32_t*)arg_type_ids)[(0)]; void* arg1 = (((TVMValue*)args)[1].v_handle); int32_t arg1_code = ((int32_t*)arg_type_ids)[(1)]; void* placeholder = (((DLTensor*)arg0)[0].data); void* arg0_shape = (((DLTensor*)arg0)[0].shape); void* arg0_strides = (((DLTensor*)arg0)[0].strides); int32_t dev_id = (((DLTensor*)arg0)[0].device.device_id); void* T_add_red = (((DLTensor*)arg1)[0].data); void* arg1_shape = (((DLTensor*)arg1)[0].shape); void* arg1_strides = (((DLTensor*)arg1)[0].strides); if (!(arg0_strides = NULL)) { } if (!(arg1_strides = NULL)) { } for (int32_t ax0 = 0; ax0 < 20; ++ax0) { ((float*)T_add_red)[(ax0)] = 0.000000e+00f; for (int32_t k1 = 0; k1 < 20; ++k1) { ((float*)T_add_red)[(ax0)] = (((float*)T_add_red)[(ax0)] + ((((float*)placeholder)[(((ax0 * 20) + k1))] + ((float*)placeholder)[(((ax0 * 20) + k1))]) + (((float*)placeholder)[(((ax0 * 20) + k1))] + ((float*)placeholder)[(((ax0 * 20) + k1))]))); } } return 0; }

#ifdef __cplusplus extern "C" #endif TVM_DLL int32_t tvmgen_default_fused_add(void* args, void* arg_type_ids, int32_t num_args, void* out_ret_value, void* out_ret_tcode, void* resource_handle) { void* arg0 = (((TVMValue*)args)[0].v_handle); int32_t arg0_code = ((int32_t*)arg_type_ids)[(0)]; void* arg1 = (((TVMValue*)args)[1].v_handle); int32_t arg1_code = ((int32_t*)arg_type_ids)[(1)]; void* placeholder = (((DLTensor*)arg0)[0].data); void* arg0_shape = (((DLTensor*)arg0)[0].shape); void* arg0_strides = (((DLTensor*)arg0)[0].strides); int32_t dev_id = (((DLTensor*)arg0)[0].device.device_id); void* T_add = (((DLTensor*)arg1)[0].data); void* arg1_shape = (((DLTensor*)arg1)[0].shape); void* arg1_strides = (((DLTensor*)arg1)[0].strides); if (!(arg0_strides = NULL)) { } if (!(arg1_strides = NULL)) { } for (int32_t ax0_outer = 0; ax0_outer < 2; ++ax0_outer) { for (int32_t ax0_inner_s = 0; ax0_inner_s < 16; ++ax0_inner_s) { if (((ax0_outer * 16) + ax0_inner_s) < 20) { ((float*)T_add)[(((ax0_outer * 16) + ax0_inner_s))] = (((float*)placeholder)[(((ax0_outer * 16) + ax0_inner_s))] + ((float*)placeholder)[(((ax0_outer * 16) + ax0_inner_s))]); } } } return 0; }

y1, y2, y3 三个操作使用的 schedule 都是要处理一个 20x20 的循环, 但 y 的 schdule 处理的是 20x1 的循环, 所以它们无法合并成一个 op.

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# 2021-10-14 22:36
import numpy as np

import tvm
from tvm import te
from tvm import relay
from tvm.relay import transform

I = N = 1
O = C = 2
H = W = 3


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 test_fold():
    def get_model(x, weight, scale):
        args = [x]
        # x = relay.multiply(x, scale)
        y = relay.nn.conv2d(
            x,
            weight,
            channels=C,
            kernel_size=(3, 3),
            padding=(1, 1),
        )
        y = relay.multiply(y, scale)

        return relay.Function(args, y)

    def check(shape):
        x = relay.var("x", shape=shape)

        weight = relay.const(np.random.randn(O, I, H, W).astype("float32"))
        scale = relay.const(np.random.randn(C, 1, 1).astype("float32"))

        y = get_model(x, weight, scale)
        y = run_opt_pass(y, transform.InferType())
        print(y)

        y_folded = run_opt_pass(y, transform.BackwardFoldScaleAxis())
        print(y_folded)

        y_folded = run_opt_pass(y_folded, transform.FoldConstant())
        print(y_folded)

    check((1, N, 10, 10))


if __name__ == "__main__":
    test_fold()

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

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

fn (%x: Tensor[(1, 1, 10, 10), float32]) -> Tensor[(1, 2, 10, 10), float32] { nn.conv2d(%x, meta[relay.Constant][0] * ty=Tensor[(2, 1, 3, 3), float32] *, padding=[1, 1, 1, 1], channels=2, kernel_size=[3, 3]) * ty=Tensor[(1, 2, 10, 10), float32] * }

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# 2021-10-14 22:36
import numpy as np

import tvm
from tvm import te
from tvm import relay
from tvm.relay import transform


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


def test_combine_parallel_dense():
    """Simple testcase. One dense cannot be combined due to shape mismatch"""

    def before(x, w1, w2):
        args = [x, w1, w2]
        y1 = relay.nn.dense(x, w1)
        y2 = relay.nn.dense(x, w2)

        y = relay.Tuple((y1, y2))
        return relay.Function(args, y)

    def check(i, j, k):
        x = relay.var("x", shape=(i, k))
        w1 = relay.var("w1", shape=(j, k))
        w2 = relay.var("w2", shape=(j, k))

        y_before = before(x, w1, w2)
        print(y_before)
        y = run_opt_pass(y_before, transform.CombineParallelDense(min_num_branches=2))
        print(y)

    # 3x4 分别和两个 4x5 相乘, 得到两个 3x5
    check(3, 5, 4)


if __name__ == "__main__":
    test_combine_parallel_dense()

fn (%x: Tensor[(3, 4), float32], %w1: Tensor[(5, 4), float32], %w2: Tensor[(5, 4), float32]) { %0 = nn.dense(%x, %w1, units=None); %1 = nn.dense(%x, %w2, units=None); (%0, %1) } fn (%x: Tensor[(3, 4), float32], %w1: Tensor[(5, 4), float32], %w2: Tensor[(5, 4), float32]) -> (Tensor[(3, 5), float32], Tensor[(3, 5), float32]) { %0 = (%x, %x); %1 = (%w1, %w2); %2 = stack(%0) * ty=Tensor[(2, 3, 4), float32] *; %3 = stack(%1) * ty=Tensor[(2, 5, 4), float32] *; %4 = nn.batch_matmul(%2, %3, transpose_b=True) * ty=Tensor[(2, 3, 5), float32] *; %5 = split(%4, indices_or_sections=2) * ty=(Tensor[(1, 3, 5), float32], Tensor[(1, 3, 5), float32]) *; %6 = %5.0; %7 = %5.1; %8 = squeeze(%6, axis=[0]) * ty=Tensor[(3, 5), float32] *; %9 = squeeze(%7, axis=[0]) * ty=Tensor[(3, 5), float32] *; (%8, %9) }

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# 2021-10-22 16:46

import tvm
from tvm import relay
from tvm.relay import transform


def test_combine_parallel_conv2d():
    def before(x, w1, w2):
        args = [x, w1, w2]
        y1 = relay.nn.conv2d(x, w1)
        y2 = relay.nn.conv2d(x, w2)
        y3 = relay.nn.max_pool2d(x)
        y = relay.Tuple((y1, y2, y3))
        func = relay.Function(args, y)
        mod = tvm.IRModule.from_expr(func)
        mod = tvm.relay.transform.InferType()(mod)
        return mod

    def check(x_shape, channels1, channels2):
        x = relay.var("x", shape=x_shape)
        in_c = x_shape[1]
        w1 = relay.var("w1", shape=(channels1, in_c, 1, 1))
        w2 = relay.var("w2", shape=(channels2, in_c, 1, 1))

        mod = before(x, w1, w2)
        print("------before------")
        print(mod)
        mod = transform.CombineParallelConv2D(min_num_branches=2)(mod)
        print("------after------")
        print(mod)

    check((1, 4, 16, 16), 4, 4)


if __name__ == "__main__":
    test_combine_parallel_conv2d()

-–—before-–— def @main(%x: Tensor[(1, 4, 16, 16), float32], %w1: Tensor[(4, 4, 1, 1), float32], %w2: Tensor[(4, 4, 1, 1), float32]) -> (Tensor[(1, 4, 16, 16), float32], Tensor[(1, 4, 16, 16), float32], Tensor[(1, 4, 16, 16), float32]) { %0 = nn.conv2d(%x, %w1, padding=[0, 0, 0, 0]) * ty=Tensor[(1, 4, 16, 16), float32] *; %1 = nn.conv2d(%x, %w2, padding=[0, 0, 0, 0]) * ty=Tensor[(1, 4, 16, 16), float32] *; %2 = nn.max_pool2d(%x, pool_size=[1, 1], padding=[0, 0, 0, 0]) * ty=Tensor[(1, 4, 16, 16), float32] *; (%0, %1, %2) }

-–—after-–— def @main(%x: Tensor[(1, 4, 16, 16), float32], %w1: Tensor[(4, 4, 1, 1), float32], %w2: Tensor[(4, 4, 1, 1), float32]) -> (Tensor[(1, 4, 16, 16), float32], Tensor[(1, 4, 16, 16), float32], Tensor[(1, 4, 16, 16), float32]) { %0 = (%w1, %w2); %1 = concatenate(%0) * ty=Tensor[(8, 4, 1, 1), float32] *; %2 = nn.conv2d(%x, %1, padding=[0, 0, 0, 0], channels=8) * ty=Tensor[(1, 8, 16, 16), float32] *; %3 = strided_slice(%2, begin=[0, 0], end=[-1, 4], strides=[1, 1], slice_mode="size", axes=None) * ty=Tensor[(1, 4, 16, 16), float32] *; %4 = strided_slice(%2, begin=[0, 4], end=[-1, 4], strides=[1, 1], slice_mode="size", axes=None) * ty=Tensor[(1, 4, 16, 16), float32] *; %5 = nn.max_pool2d(%x, pool_size=[1, 1], padding=[0, 0, 0, 0]) * ty=Tensor[(1, 4, 16, 16), float32] *; (%3, %4, %5) }