wehub-resource-sync
92 lines
3.1 KiB
Python
92 lines
3.1 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
|
|
import tvm.script
|
|
import tvm.testing
|
|
from tvm import relax
|
|
from tvm.script import ir as I
|
|
from tvm.script import relax as R
|
|
|
|
|
|
def test_multinomial_from_uniform():
|
|
@I.ir_module
|
|
class CallSample:
|
|
@R.function
|
|
def foo(x: R.Tensor((3, 5), "float32"), y: R.Tensor((3, 1), "float32")):
|
|
z = R.call_pure_packed(
|
|
"vm.builtin.multinomial_from_uniform",
|
|
x,
|
|
y,
|
|
ty_args=(R.Tensor((3, 1), dtype="int64")),
|
|
)
|
|
return z
|
|
|
|
mod = CallSample
|
|
target = tvm.target.Target("llvm", host="llvm")
|
|
ex = tvm.compile(mod, target)
|
|
np_rand = np.random.rand(3, 5).astype(np.float32)
|
|
# normalize it to get the random prob
|
|
np_prob = np_rand / np_rand.sum(axis=1, keepdims=True)
|
|
nd_prob = tvm.runtime.tensor(np_prob)
|
|
# special sample to get deterministic results
|
|
nd_sample = tvm.runtime.tensor(np.array([[1.0], [0], [1]]).astype(np.float32))
|
|
|
|
vm = relax.VirtualMachine(ex, tvm.cpu())
|
|
res = vm["foo"](nd_prob, nd_sample)
|
|
tvm.testing.assert_allclose(res.numpy(), np.array([[4], [0], [4]]).astype(np.int64))
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not tvm.testing.device_enabled("cuda"), reason="cuda not enabled")
|
|
def test_alloc_tensor_raises_out_of_memory():
|
|
"""Out-of-memory exceptions may be raised from VM
|
|
|
|
This is a regression test. In previous implementations, the Relax
|
|
VM would segfault if the built-in function
|
|
"vm.builtin.alloc_storage" was unable to allocate the requested
|
|
buffer.
|
|
"""
|
|
target = "cuda"
|
|
|
|
@I.ir_module
|
|
class Module:
|
|
@R.function
|
|
def main():
|
|
# Allocate a 1-petabyte tensor to trigger OOM. If the CI
|
|
# ever runs on a device with more than 1 petabyte of GPU
|
|
# memory, this test will need to be updated.
|
|
output = R.builtin.alloc_tensor(R.shape([1024, 1024, 1024, 1024, 1024]), "uint8", 0)
|
|
return output
|
|
|
|
built = tvm.compile(Module, target=target)
|
|
|
|
def run_and_check():
|
|
dev = tvm.device(target)
|
|
vm = relax.VirtualMachine(built, dev)
|
|
with pytest.raises(Exception, match="CUDA.*out of memory"):
|
|
vm["main"]()
|
|
|
|
tvm.testing.run_with_gpu_lock(run_and_check)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tvm.testing.main()
|