wehub-resource-sync
191 lines
6.4 KiB
Python
191 lines
6.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.
|
|
"""Test eliminate common subexpr pass"""
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
import tvm
|
|
import tvm.testing
|
|
from tvm import relax
|
|
from tvm.ir.module import IRModule
|
|
from tvm.script.parser import ir as I
|
|
from tvm.script.parser import relax as R
|
|
from tvm.testing import env
|
|
|
|
|
|
def compile(mod: IRModule):
|
|
# compile the model
|
|
mod = relax.transform.RealizeVDevice()(mod)
|
|
mod = relax.transform.LegalizeOps()(mod)
|
|
mod = tvm.s_tir.transform.DefaultGPUSchedule()(mod)
|
|
# no need to feed target argument for mult-target compilation
|
|
return tvm.compile(mod)
|
|
|
|
|
|
def test_multi_cpu():
|
|
@I.ir_module
|
|
class Example:
|
|
I.module_attrs({"attr": 10})
|
|
I.module_global_infos(
|
|
{
|
|
"vdevice": [
|
|
I.vdevice("llvm", 0),
|
|
I.vdevice("llvm", 1),
|
|
]
|
|
}
|
|
)
|
|
|
|
@R.function
|
|
def foo(
|
|
x: R.Tensor((2, 3), "float32"),
|
|
y: R.Tensor((3, 4), "float32"),
|
|
z: R.Tensor((4, 5), "float32"),
|
|
) -> R.Tensor((2, 5), "float32"):
|
|
with R.dataflow():
|
|
lv0 = R.matmul(x, y)
|
|
lv0 = R.hint_on_device(lv0, tvm.cpu(0))
|
|
lv1: R.Tensor((2, 4), "float32", "llvm:1") = R.to_vdevice(lv0, "llvm:1")
|
|
gv = R.matmul(lv1, z)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
devices = [tvm.cpu(0), tvm.cpu(1)]
|
|
vm = relax.VirtualMachine(compile(Example), devices)
|
|
|
|
np_ipt0 = np.random.rand(2, 3).astype(np.float32)
|
|
np_ipt1 = np.random.rand(3, 4).astype(np.float32)
|
|
np_ipt2 = np.random.rand(4, 5).astype(np.float32)
|
|
np_res = np.matmul(np.matmul(np_ipt0, np_ipt1), np_ipt2)
|
|
|
|
ipt0 = tvm.runtime.tensor(np_ipt0, devices[0])
|
|
ipt1 = tvm.runtime.tensor(np_ipt1, devices[0])
|
|
ipt2 = tvm.runtime.tensor(np_ipt2, devices[1])
|
|
res = vm["foo"](ipt0, ipt1, ipt2)
|
|
tvm.testing.assert_allclose(res.numpy(), np_res)
|
|
|
|
|
|
@pytest.mark.skipif(not env.has_multi_gpu(), reason="need multiple gpus")
|
|
def test_multi_gpu():
|
|
@I.ir_module
|
|
class Example:
|
|
I.module_attrs({"attr": 10})
|
|
I.module_global_infos(
|
|
{
|
|
"vdevice": [
|
|
I.vdevice("cuda", 1),
|
|
I.vdevice("cuda", 0),
|
|
I.vdevice("cuda", 2),
|
|
]
|
|
}
|
|
)
|
|
|
|
@R.function
|
|
def foo(
|
|
a: R.Tensor((2, 3), "float32"),
|
|
b: R.Tensor((3, 4), "float32"),
|
|
c: R.Tensor((4, 5), "float32"),
|
|
d: R.Tensor((5, 6), "float32"),
|
|
) -> R.Tensor((2, 6), "float32"):
|
|
with R.dataflow():
|
|
lv0: R.Tensor((2, 4), "float32", "cuda:0") = R.matmul(a, b)
|
|
lv1: R.Tensor((2, 4), "float32", "cuda:1") = R.to_vdevice(
|
|
lv0,
|
|
"cuda:1",
|
|
)
|
|
lv2: R.Tensor((2, 5), "float32", "cuda:1") = R.matmul(lv1, c)
|
|
lv3: R.Tensor((2, 5), "float32", "cuda:2") = R.to_vdevice(
|
|
lv2,
|
|
"cuda:2",
|
|
)
|
|
gv: R.Tensor((2, 6), "float32", "cuda:2") = R.matmul(lv3, d)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
np_ipt0 = np.random.rand(2, 3).astype(np.float32)
|
|
np_ipt1 = np.random.rand(3, 4).astype(np.float32)
|
|
np_ipt2 = np.random.rand(4, 5).astype(np.float32)
|
|
np_ipt3 = np.random.rand(5, 6).astype(np.float32)
|
|
np_res = np.matmul(np.matmul(np.matmul(np_ipt0, np_ipt1), np_ipt2), np_ipt3)
|
|
|
|
ex = compile(Example)
|
|
|
|
def run_and_check():
|
|
if not tvm.cuda(2).exist:
|
|
pytest.skip("requires at least 3 visible CUDA devices")
|
|
devices = [tvm.cuda(1), tvm.cuda(0), tvm.cuda(2)]
|
|
vm = relax.VirtualMachine(ex, devices)
|
|
ipt0 = tvm.runtime.tensor(np_ipt0, devices[0])
|
|
ipt1 = tvm.runtime.tensor(np_ipt1, devices[0])
|
|
ipt2 = tvm.runtime.tensor(np_ipt2, devices[1])
|
|
ipt3 = tvm.runtime.tensor(np_ipt3, devices[2])
|
|
res = vm["foo"](ipt0, ipt1, ipt2, ipt3)
|
|
tvm.testing.assert_allclose(res.numpy(), np_res)
|
|
|
|
tvm.testing.run_with_gpu_lock(run_and_check)
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_gpu(), reason="need gpu")
|
|
def test_multi_device():
|
|
@I.ir_module
|
|
class Example:
|
|
I.module_attrs({"attr": 10})
|
|
I.module_global_infos(
|
|
{
|
|
"vdevice": [
|
|
I.vdevice("cuda", 0),
|
|
I.vdevice("llvm"),
|
|
]
|
|
}
|
|
)
|
|
|
|
@R.function
|
|
def foo(
|
|
x: R.Tensor((2, 3), "float32"),
|
|
y: R.Tensor((3, 4), "float32"),
|
|
z: R.Tensor((4, 5), "float32"),
|
|
) -> R.Tensor((2, 5), "float32"):
|
|
with R.dataflow():
|
|
lv0: R.Tensor((2, 4), "float32", "llvm") = R.matmul(x, y)
|
|
lv1: R.Tensor((2, 4), "float32", "cuda") = R.to_vdevice(lv0, "cuda")
|
|
gv: R.Tensor((2, 5), "float32", "cuda") = R.matmul(lv1, z)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
np_ipt0 = np.random.rand(2, 3).astype(np.float32)
|
|
np_ipt1 = np.random.rand(3, 4).astype(np.float32)
|
|
np_ipt2 = np.random.rand(4, 5).astype(np.float32)
|
|
np_res = np.matmul(np.matmul(np_ipt0, np_ipt1), np_ipt2)
|
|
|
|
ex = compile(Example)
|
|
|
|
def run_and_check():
|
|
devices = [tvm.cuda(0), tvm.cpu(0)]
|
|
vm = relax.VirtualMachine(ex, devices)
|
|
ipt0 = tvm.runtime.tensor(np_ipt0, devices[1])
|
|
ipt1 = tvm.runtime.tensor(np_ipt1, devices[1])
|
|
ipt2 = tvm.runtime.tensor(np_ipt2, devices[0])
|
|
res = vm["foo"](ipt0, ipt1, ipt2)
|
|
tvm.testing.assert_allclose(res.numpy(), np_res, rtol=1e-4, atol=1e-4)
|
|
|
|
tvm.testing.run_with_gpu_lock(run_and_check)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tvm.testing.main()
|