apache--tvm/tests/python/relax/test_vm_cuda_graph.py

# 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: E501

import numpy as np
import pytest

import tvm
import tvm.testing
from tvm import relax
from tvm.script import ir as I
from tvm.script import relax as R
from tvm.script import tirx as T
from tvm.testing import env

# fmt: off


@I.ir_module(s_tir=True)
class Module:
    @R.function(pure=False)
    def main(x: R.Tensor((16, 16), dtype="float32")) -> R.Tensor((16, 16), dtype="float32"):
        cls = Module
        R.func_attr({"global_symbol": "main"})
        gv: R.Tuple(R.Any, R.Any) = R.call_builtin_with_ctx("vm.builtin.cuda_graph.get_cached_alloc", (cls.cuda_graph_alloc, R.prim_value(0)), ty_args=(R.Tuple(R.Any, R.Any),))
        storage: R.Any = gv[0]
        alloc = R.vm.alloc_tensor(storage, R.prim_value(0), R.shape((16, 16)), R.dtype("float32"))
        _: R.Tuple = cls.add(x, alloc)
        storage1: R.Any = gv[1]
        gv1: R.Tuple(R.Tensor(dtype="float32"), R.Any, R.Any) = (alloc, storage1, storage)
        gv2: R.Tuple(R.Tensor((16, 16), dtype="float32")) = R.call_builtin_with_ctx("vm.builtin.cuda_graph.run_or_capture", (cls.cuda_graph_capture, gv1, R.prim_value(0)), ty_args=(R.Tuple(R.Tensor((16, 16), dtype="float32")),))
        storage2: R.Any = R.vm.alloc_storage(R.shape((1024,)), R.prim_value(0), R.dtype("uint8"))
        alloc3 = R.vm.alloc_tensor(storage2, R.prim_value(0), R.shape((16, 16)), R.dtype("float32"))
        lv4: R.Tensor((16, 16), dtype="float32") = gv2[0]
        _3: R.Tuple = cls.add(lv4, alloc3)
        lv5: R.Tensor(dtype="float32") = alloc3
        return lv5

    @T.prim_func(s_tir=True)
    def add(A: T.Buffer((16, 16), "float32"), B: T.Buffer((16, 16), "float32")):
        T.func_attr({"global_symbol": "add"})
        with T.sblock("root"):
            for i in T.thread_binding(16, thread="threadIdx.x"):
                for j in range(16):
                    with T.sblock("update"):
                        vi, vj = T.axis.remap("SS", [i, j])
                        B[vi, vj] = A[vi, vj] + T.float32(1)

    @R.function
    def cuda_graph_alloc() -> R.Tuple(R.Any, R.Any):
        R.func_attr({"global_symbol": "cuda_graph_alloc"})
        storage: R.Any = R.vm.alloc_storage(R.shape((1024,)), R.prim_value(0), R.dtype("uint8"))
        storage1: R.Any = R.vm.alloc_storage(R.shape((1024,)), R.prim_value(0), R.dtype("uint8"))
        gv: R.Tuple(R.Any, R.Any) = (storage, storage1)
        return gv

    @R.function(pure=False)
    def cuda_graph_capture(alloc: R.Tensor((16, 16), dtype="float32"), storage1: R.Any, storage: R.Any) -> R.Tuple(R.Tensor((16, 16), dtype="float32")):
        cls = Module
        R.func_attr({"global_symbol": "cuda_graph_capture"})
        lv0: R.Tensor((16, 16), dtype="float32") = alloc
        alloc1 = R.vm.alloc_tensor(storage1, R.prim_value(0), R.shape((16, 16)), R.dtype("float32"))
        _1: R.Tuple = cls.add(lv0, alloc1)
        lv1: R.Tensor(dtype="float32") = alloc1
        lv2: R.Tuple(R.Tensor(dtype="float32")) = (lv1,)
        lv3: R.Tensor(dtype="float32") = lv2[0]
        alloc2 = R.vm.alloc_tensor(storage, R.prim_value(0), R.shape((16, 16)), R.dtype("float32"))
        _2: R.Tuple = cls.add(lv3, alloc2)
        lv4: R.Tensor(dtype="float32") = alloc2
        gv: R.Tuple(R.Tensor(dtype="float32")) = (lv4,)
        return gv


# fmt: on


def codegen(mod, target, exec_mode="bytecode"):
    builder = relax.ExecBuilder()
    leftover_mod = relax.vm_build._vmcodegen(builder, mod, exec_mode=exec_mode)
    tir_mod = relax.vm_build._filter_tir(leftover_mod)
    return relax.vm_build._vmlink(builder, target, tir_mod)


@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
def test_vm_run():
    mod = Module
    target = tvm.target.Target("cuda", host="llvm")
    ex = codegen(mod, target)
    x_np = np.random.uniform(size=(16, 16)).astype("float32")
    y_np = x_np + 1.0 + 1.0 + 1.0 + 1.0

    def run_and_check():
        dev = tvm.cuda(0)
        vm = relax.VirtualMachine(ex, dev)
        x = tvm.runtime.tensor(x_np, dev)
        y = vm["main"](x)
        tvm.testing.assert_allclose(y.numpy(), y_np, rtol=1e-5, atol=1e-5)

    tvm.testing.run_with_gpu_lock(run_and_check)


@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cudagraph(), reason="need cudagraph")
def test_capture_error_is_recoverable():
    """Function calls while capturing cudagraph may throw exceptions

    Calls to PackedFuncs may occur within a captured cudaGraph.  If a
    call to that PackedFunc raises an exception while capturing the
    cudaGraph, throwing exception should cleanly unwind the stack, and
    the exception may be caught in the calling scope.

    This is a regression test.  In previous implementations, an
    exception thrown while capturing a cudaGraph would skip the call
    to `cudaStreamEndCapture`, causing additional exceptions to be
    thrown while freeing memory in TVM destructors.  Since C++ does
    not support stack unwinding from multiple simultaneous exceptions,
    this would result in immediate `std::terminate`, making it
    difficult to debug the original error.

    """

    target = tvm.target.Target("cuda")

    @I.ir_module(s_tir=True)
    class Module:
        @R.function
        def main(A: R.Tensor([16], "float16")):
            B = R.add(A, A)
            C = R.call_pure_packed(
                "test_vm_cuda_graph.invalid_impl_for_cudagraph",
                B,
                ty_args=R.Tensor([16], "float16"),
            )
            D = R.add(C, C)
            return D

    with target, tvm.ir.transform.PassContext(config={"relax.backend.use_cuda_graph": True}):
        Module = tvm.ir.transform.Sequential(
            [
                tvm.relax.transform.LegalizeOps(),
                tvm.s_tir.transform.DefaultGPUSchedule(),
                tvm.relax.transform.RemovePurityChecking(),
                tvm.relax.transform.CallTIRRewrite(),
                tvm.relax.transform.StaticPlanBlockMemory(),
                tvm.relax.transform.RewriteCUDAGraph(),
            ]
        )(Module)

    assert "cuda_graph_alloc" in Module, (
        "Validity of unit test requires the call to `invalid_impl_for_cudagraph` "
        "to have been captured by RewriteCUDAGraph."
    )

    built = tvm.compile(Module, target=target)

    def run_and_check():
        dev = tvm.cuda()

        @tvm.register_global_func("test_vm_cuda_graph.invalid_impl_for_cudagraph", override=True)
        def invalid_impl_for_cudagraph(arg_tensor):
            # Memory allocation/deallocation may not be performed while
            # capturing a cudaGraph.  This passes the warm-up run
            # performed by "vm.builtin.cuda_graph.run_or_capture", but
            # throws an exception when the cudaGraph is being captured.
            _dummy_workspace = tvm.runtime.empty([16], "float16", dev)
            return arg_tensor

        vm = tvm.relax.VirtualMachine(built, dev)
        arg = tvm.runtime.tensor(np.arange(16).astype("float16"), dev)

        with pytest.raises(RuntimeError):
            vm["main"](arg)

    tvm.testing.run_with_gpu_lock(run_and_check)


if __name__ == "__main__":
    tvm.testing.main()