Verilator 代码生成测试

import set_env

from infrastructure import (
    compile_hardware,
    compiler_opts,
    run_module,
    offload,
    clear_stats,
    stats
)

import numpy as np

import tvm
import tvm.testing
from tvm import relay

def create_module_add(shape, dtype):
    """Create add module.

    Paramters
    ---------
    shape : Tuple
        The shape tuple.

    dtype : Str
        The data type.

    Returns
    -------
    mod: Module
        The relay module.
    """
    x = relay.var("x", shape=shape, dtype=dtype)
    y = relay.var("y", shape=shape, dtype=dtype)
    z = relay.add(x, y)
    f = relay.Function([x, y], z)
    mod = tvm.IRModule()
    mod["main"] = f
    return mod


def create_module_bias_add(xshape, yshape, dtype):
    """Create bias_add module.

    Paramters
    ---------
    xshape : Tuple
        The x shape tuple.

    yshape : Tuple
        The y shape tuple.

    dtype : Str
        The data type.

    Returns
    -------
    mod: Module
        The relay module.
    """
    x = relay.var("x", shape=xshape, dtype=dtype)
    y = relay.var("y", shape=yshape, dtype=dtype)
    z = relay.nn.bias_add(x, y, axis=3)
    f = relay.Function([x, y], z)
    mod = tvm.IRModule()
    mod["main"] = f
    return mod


def run_and_check(xshape, yshape, dtype, mod, opts):
    """Run and check values.

    Paramters
    ---------
    xshape : Tuple
        The x shape tuple.

    yshape : Tuple
        The y shape tuple.

    dtype : Str
        The data type.

    mod: Module
        The relay module.

    opts: Dict
        The compiler options.

    Returns
    -------
    cycles: Int
        The number of cycles.
    """
    x_data = np.random.randint(5, size=xshape, dtype=dtype)
    y_data = np.random.randint(5, size=yshape, dtype=dtype)
    ref = x_data + y_data
    inp = {"x": x_data, "y": y_data}
    clear_stats()
    out = run_module(inp, mod, params=None, opts=opts)
    values = stats()
    tvm.testing.assert_allclose(out.numpy(), ref, rtol=1e-5, atol=1e-5)
    return values["cycle_counter"]


def print_test_info(test, lanes, cycles):
    """Print counter

    Paramters
    ---------
    test : Str
        The name of the test.

    lanes : Int
        The number of vector lanes.

    cycles : Int
        The number of cycles.
    """
    print("test:{} vector-lanes:{} number of cycles:{}".format(test, lanes, cycles))

def tadd(lanes, verilator_app_path):
    """Print counter

    Paramters
    ---------
    lanes : Int
        The number of vector lanes.
    """
    dtype = "int32"
    shape = (8, 4)
    mod = create_module_add(shape, dtype)
    mod = offload(mod)
    lib = compile_hardware(lanes, verilator_app_path)
    opts = compiler_opts(lib)
    cycles = run_and_check(shape, shape, dtype, mod, opts)
    print_test_info("add", lanes, cycles)

def tbias(lanes, verilator_app_path):
    """Print counter

    Paramters
    ---------
    lanes : Int
        The number of vector lanes.
    """
    dtype = "int32"
    xshape = (1, 112, 112, 32)
    yshape = (32,)
    mod = create_module_bias_add(xshape, yshape, dtype)
    mod = offload(mod)
    lib = compile_hardware(lanes, verilator_app_path)
    opts = compiler_opts(lib)
    cycles = run_and_check(xshape, yshape, dtype, mod, opts)
    print_test_info("nn.bias_add", lanes, cycles)

def tbias(lanes, verilator_app_path):
    """Print counter

    Paramters
    ---------
    lanes : Int
        The number of vector lanes.
    """
    dtype = "int32"
    xshape = (1, 112, 112, 32)
    yshape = (32,)
    mod = create_module_bias_add(xshape, yshape, dtype)
    mod = offload(mod)
    lib = compile_hardware(lanes, verilator_app_path)
    opts = compiler_opts(lib)
    cycles = run_and_check(xshape, yshape, dtype, mod, opts)
    print_test_info("nn.bias_add", lanes, cycles)

from set_env import TVM_ROOT
# add tests
verilator_app_path = f"{TVM_ROOT}/3rdparty/vta-hw/apps/verilator/add"
tadd(1, verilator_app_path)
tadd(4, verilator_app_path)

test:add vector-lanes:1 number of cycles:32
test:add vector-lanes:4 number of cycles:8

[15:23:02] /media/pc/data/lxw/ai/tvm/src/relay/backend/vm/compiler.cc:1199: All lowered functions have been build by BYOC -- generating an empty TVM module
[15:23:05] /media/pc/data/lxw/ai/tvm/src/relay/backend/vm/compiler.cc:1199: All lowered functions have been build by BYOC -- generating an empty TVM module

# bias_add tests
tbias(1, verilator_app_path)
tbias(32, verilator_app_path)

test:nn.bias_add vector-lanes:1 number of cycles:401408
test:nn.bias_add vector-lanes:32 number of cycles:12544

[15:23:05] /media/pc/data/lxw/ai/tvm/src/relay/backend/vm/compiler.cc:1199: All lowered functions have been build by BYOC -- generating an empty TVM module
[15:23:08] /media/pc/data/lxw/ai/tvm/src/relay/backend/vm/compiler.cc:1199: All lowered functions have been build by BYOC -- generating an empty TVM module