chore: import upstream snapshot with attribution

tests/python/relax/test_vm_builtin.py

@@ -0,0 +1,91 @@
+# 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()