wehub-resource-sync
104 lines
3.4 KiB
Python
104 lines
3.4 KiB
Python
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you under the Apache License, Version 2.0 (the
|
|
# "License"); you may not use this file except in compliance
|
|
# with the License. You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
# ruff: noqa: RUF005
|
|
import numpy as np
|
|
|
|
import tvm
|
|
import tvm.testing
|
|
from tvm import relax
|
|
from tvm.relax.testing import nn
|
|
from tvm.relax.testing.lib_comparator import LibCompareVMInstrument
|
|
|
|
|
|
def get_exec(data_shape):
|
|
builder = relax.BlockBuilder()
|
|
weight1_np = np.random.randn(64, 64).astype("float32")
|
|
weight2_np = np.random.randn(64, 64).astype("float32")
|
|
|
|
with builder.function("main"):
|
|
model = nn.Sequential(
|
|
nn.Linear(data_shape[1], weight1_np.shape[0], bias=False),
|
|
nn.ReLU(),
|
|
nn.Linear(weight2_np.shape[0], weight2_np.shape[1], bias=False),
|
|
nn.ReLU(),
|
|
)
|
|
data = nn.Placeholder(data_shape, name="data")
|
|
output = model(data)
|
|
params = [data] + model.parameters()
|
|
builder.emit_func_output(output, params=params)
|
|
|
|
mod = builder.get()
|
|
|
|
params = {"linear_weight": weight1_np, "linear_weight1": weight2_np}
|
|
mod = relax.transform.BindParams("main", params)(mod)
|
|
|
|
target = "llvm"
|
|
return tvm.compile(mod, target)
|
|
|
|
|
|
def get_exec_int32(data_shape):
|
|
builder = relax.BlockBuilder()
|
|
|
|
with builder.function("main"):
|
|
model = nn.ReLU()
|
|
data = nn.Placeholder(data_shape, dtype="int32", name="data")
|
|
output = model(data)
|
|
params = [data] + model.parameters()
|
|
builder.emit_func_output(output, params=params)
|
|
|
|
mod = builder.get()
|
|
target = "llvm"
|
|
return tvm.compile(mod, target)
|
|
|
|
|
|
def test_conv2d_cpu():
|
|
data_np = np.random.randn(1, 64).astype("float32")
|
|
ex = get_exec(data_np.shape)
|
|
vm = relax.VirtualMachine(ex, tvm.cpu())
|
|
hit_count = {}
|
|
|
|
def instrument(func, name, before_run, ret_val, *args):
|
|
if (name, before_run) not in hit_count:
|
|
hit_count[(name, before_run)] = 0
|
|
hit_count[(name, before_run)] += 1
|
|
assert callable(func)
|
|
if before_run:
|
|
assert ret_val is None
|
|
if name == "matmul":
|
|
return relax.VMInstrumentReturnKind.SKIP_RUN
|
|
|
|
vm.set_instrument(instrument)
|
|
vm["main"](tvm.runtime.tensor(data_np))
|
|
assert hit_count[("matmul", True)] == 2
|
|
assert ("matmul", False) not in hit_count
|
|
assert hit_count[("relu", True)] == 2
|
|
assert hit_count[("relu", False)] == 2
|
|
|
|
|
|
def test_lib_comparator():
|
|
data_np = np.random.randn(1, 64).astype("int32")
|
|
ex = get_exec_int32(data_np.shape)
|
|
vm = relax.VirtualMachine(ex, tvm.cpu())
|
|
# compare against library module
|
|
cmp = LibCompareVMInstrument(vm.module.imports[0], tvm.cpu(), verbose=False)
|
|
vm.set_instrument(cmp)
|
|
vm["main"](tvm.runtime.tensor(data_np))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tvm.testing.main()
|