循环依赖

循环依赖#

import numpy as np

import tvm
from tvm import relax
from tvm.relax.backend.cuda.cublas import partition_for_cublas
from tvm.relax.backend.cuda.cutlass import partition_for_cutlass
from tvm.relax.dpl.pattern import (
    is_op,
    is_tuple_get_item,
    make_fused_bias_activation_pattern,
    wildcard,
)
from tvm.relax.transform import PatternCheckContext
from tvm.script import ir as I
from tvm.script import relax as R
from tvm.script import tir as T

@tvm.script.ir_module
class Branch:
    @R.function
    def main(
        data: R.Tensor((1, 64, 56, 56), "float32"),
        weight: R.Tensor((64, 64, 3, 3), "float32"),
    ):
        with R.dataflow():
            conv1 = R.nn.conv2d(data, weight)
            relu1 = R.nn.relu(conv1)
            gelu1 = R.nn.gelu(conv1)

            out = relax.op.add(relu1, gelu1)
            R.output(out)

        return out

Branch.show()

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

@I.ir_module
class Module:
    @R.function
    def main(data: R.Tensor((1, 64, 56, 56), dtype="float32"), weight: R.Tensor((64, 64, 3, 3), dtype="float32")) -> R.Tensor((1, 64, 54, 54), dtype="float32"):
        with R.dataflow():
            conv1: R.Tensor((1, 64, 54, 54), dtype="float32") = R.nn.conv2d(data, weight, strides=[1, 1], padding=[0, 0, 0, 0], dilation=[1, 1], groups=1, data_layout="NCHW", kernel_layout="OIHW", out_layout="NCHW", out_dtype="void")
            relu1: R.Tensor((1, 64, 54, 54), dtype="float32") = R.nn.relu(conv1)
            gelu1: R.Tensor((1, 64, 54, 54), dtype="float32") = R.nn.gelu(conv1)
            out: R.Tensor((1, 64, 54, 54), dtype="float32") = R.add(relu1, gelu1)
            R.output(out)
        return out

import pytest
conv_pat = make_fused_bias_activation_pattern("relax.nn.conv2d")
relu_pat = is_op("relax.nn.relu")(conv_pat)
add_pat = is_op("relax.add")(relu_pat, wildcard())

with pytest.raises(tvm.error.TVMError) as err:
    relax.transform.FuseOpsByPattern(
        [("compiler_A.conv2d_relu_add", add_pat)], bind_constants=True
    )(Branch)

assert "A cyclic dependency detected" in str(err.value)