wehub-resource-sync
152 lines
4.8 KiB
Python
152 lines
4.8 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.
|
|
import numpy as np
|
|
import pytest
|
|
import tvm_ffi
|
|
|
|
import tvm
|
|
import tvm.testing
|
|
from tvm import relax
|
|
from tvm.script import relax as R
|
|
from tvm.script import tirx as T
|
|
|
|
|
|
def test_pipeline_compile():
|
|
target = tvm.target.Target("llvm", host="llvm")
|
|
pipeline = relax.pipeline.get_default_pipeline(target)
|
|
|
|
@tvm.script.ir_module
|
|
class Mod:
|
|
@R.function
|
|
def main(x: R.Tensor((3, 4), "float32"), y: R.Tensor((3, 4), "float32")):
|
|
lv0 = R.add(x, y)
|
|
return lv0
|
|
|
|
mod = Mod
|
|
mod = pipeline(mod)
|
|
|
|
ex = tvm.compile(mod, target)
|
|
x_np = np.random.rand(3, 4).astype(np.float32)
|
|
y_np = np.random.rand(3, 4).astype(np.float32)
|
|
x = tvm.runtime.tensor(x_np)
|
|
y = tvm.runtime.tensor(y_np)
|
|
|
|
vm = relax.VirtualMachine(ex, tvm.cpu())
|
|
z = vm["main"](x, y)
|
|
tvm.testing.assert_allclose(z.numpy(), x_np + y_np, rtol=1e-7, atol=1e-7)
|
|
|
|
|
|
def test_pipeline_with_kv_cache():
|
|
"""A dummy pipline that simulates KV update."""
|
|
target = tvm.target.Target("llvm", host="llvm")
|
|
pipeline = relax.pipeline.get_default_pipeline(target)
|
|
|
|
@tvm.script.ir_module
|
|
class Mod:
|
|
@R.function
|
|
def create_kv_cache(reserve_slots: R.Shape(["m"])):
|
|
# just allocate minimum slot since it is only used to signal dtype
|
|
m = T.int64()
|
|
init_data = R.ones((1, 4), "float32")
|
|
kv_cache = R.call_pure_packed(
|
|
"vm.builtin.attention_kv_cache_create",
|
|
init_data,
|
|
R.shape([m, 4]),
|
|
0,
|
|
ty_args=[R.Any()],
|
|
)
|
|
return kv_cache
|
|
|
|
@R.function(pure=False)
|
|
def main(
|
|
x: R.Tensor((1, 4), "float32"),
|
|
y: R.Tensor((1, 4), "float32"),
|
|
shape: R.Shape(["L", 4]),
|
|
kv_cache: R.Any,
|
|
):
|
|
L = T.int64()
|
|
# computation of the current value
|
|
curr_value = R.add(x, y)
|
|
# update cache
|
|
kv_cache = R.call_packed(
|
|
"vm.builtin.attention_kv_cache_append", kv_cache, curr_value, ty_args=[R.Any]
|
|
)
|
|
# return the updated cache view
|
|
kv = R.call_packed(
|
|
"vm.builtin.attention_kv_cache_view",
|
|
kv_cache,
|
|
shape,
|
|
ty_args=[R.Tensor((L, 4), "float32")],
|
|
)
|
|
return (kv, kv_cache)
|
|
|
|
mod = Mod
|
|
mod = pipeline(mod)
|
|
|
|
ex = tvm.compile(mod, target)
|
|
|
|
num_steps = 8
|
|
cache_np = np.empty((num_steps, 4), dtype="float32")
|
|
vm = relax.VirtualMachine(ex, tvm.cpu())
|
|
|
|
kv_cache = vm["create_kv_cache"](tvm_ffi.Shape([1]))
|
|
|
|
for i in range(num_steps):
|
|
x_np = np.random.rand(1, 4).astype(np.float32)
|
|
y_np = np.random.rand(1, 4).astype(np.float32)
|
|
x = tvm.runtime.tensor(x_np)
|
|
y = tvm.runtime.tensor(y_np)
|
|
np_shape = (i + 1, 4)
|
|
kv, kv_cache = vm["main"](x, y, tvm_ffi.Shape(np_shape), kv_cache)
|
|
|
|
cache_np[i, :] = x_np + y_np
|
|
tvm.testing.assert_allclose(kv.numpy(), cache_np[: np_shape[0], :], rtol=1e-7, atol=1e-7)
|
|
|
|
|
|
@pytest.mark.parametrize("target_name", ["vulkan", "webgpu"])
|
|
@pytest.mark.parametrize(
|
|
"pipeline_func",
|
|
[
|
|
relax.pipeline.library_dispatch_passes,
|
|
relax.pipeline.legalize_passes,
|
|
relax.pipeline.dataflow_lower_passes,
|
|
relax.pipeline.finalize_passes,
|
|
relax.pipeline.get_default_pipeline,
|
|
],
|
|
)
|
|
def test_gpu_generic_fallback(target_name, pipeline_func):
|
|
target = tvm.target.Target(target_name)
|
|
result = pipeline_func(target)
|
|
assert result is not None
|
|
|
|
|
|
@pytest.mark.parametrize("target_name", ["hexagon", "c"])
|
|
@pytest.mark.parametrize(
|
|
"pipeline_func",
|
|
[
|
|
relax.pipeline.library_dispatch_passes,
|
|
relax.pipeline.legalize_passes,
|
|
relax.pipeline.dataflow_lower_passes,
|
|
relax.pipeline.finalize_passes,
|
|
relax.pipeline.get_default_pipeline,
|
|
],
|
|
)
|
|
def test_non_gpu_target_raises_error(target_name, pipeline_func):
|
|
target = tvm.target.Target(target_name)
|
|
with pytest.raises(ValueError, match="not yet supported"):
|
|
pipeline_func(target)
|