Function test_cuda_memory_alloc_noleak

tests/python/test_load_inline.py:333–364

()

    torch is None or not torch.cuda.is_available(), reason="Requires torch and CUDA"
)
def test_cuda_memory_alloc_noleak() -> None:
    assert torch is not None
    mod = tvm_ffi.cpp.load_inline(
        name="hello",
        cuda_sources=r"""
            #include <tvm/ffi/function.h>
            #include <tvm/ffi/container/tensor.h>

            namespace ffi = tvm::ffi;

            ffi::Tensor return_tensor(tvm::ffi::TensorView x) {
                ffi::Tensor y = ffi::Tensor::FromEnvAlloc(
                    TVMFFIEnvTensorAlloc, x.shape(), x.dtype(), x.device());
                return y;
            }
        """,
        functions=["return_tensor"],
    )

    def run_check() -> None:
        """Must run in a separate function to ensure deletion happens before mod unloads."""
        assert torch is not None
        x = torch.arange(1024 * 1024, dtype=torch.float32, device="cuda")
        current_allocated = torch.cuda.memory_allocated()
        repeat = 8
        for i in range(repeat):
            mod.return_tensor(x)
            diff = torch.cuda.memory_allocated() - current_allocated
            # memory should not grow as we loop over
            assert diff <= 1024**2 * 8

    run_check()