numba-cuda 0.22.0__cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl

numba_cuda/numba/cuda/tests/nocuda/test_library_lookup.py

@@ -0,0 +1,252 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import sys
+import os
+import multiprocessing as mp
+import warnings
+
+
+from numba.cuda.core.config import IS_WIN32
+from numba.cuda.core.errors import NumbaWarning
+from numba.cuda.cudadrv import nvvm
+from numba.cuda.testing import (
+    unittest,
+    skip_on_cudasim,
+    skip_unless_conda_cudatoolkit,
+)
+from numba.cuda.cuda_paths import (
+    _get_libdevice_path_decision,
+    _get_nvvm_path_decision,
+    _get_cudalib_dir_path_decision,
+    get_system_ctk,
+    get_system_ctk_libdir,
+)
+
+
+has_cuda = nvvm.is_available()
+has_mp_get_context = hasattr(mp, "get_context")
+
+
+class LibraryLookupBase(unittest.TestCase):
+    def setUp(self):
+        ctx = mp.get_context("spawn")
+
+        qrecv = ctx.Queue()
+        qsend = ctx.Queue()
+        self.qsend = qsend
+        self.qrecv = qrecv
+        self.child_process = ctx.Process(
+            target=check_lib_lookup,
+            args=(qrecv, qsend),
+            daemon=True,
+        )
+        self.child_process.start()
+
+    def tearDown(self):
+        self.qsend.put(self.do_terminate)
+        self.child_process.join(3)
+        # Ensure the process is terminated
+        self.assertIsNotNone(self.child_process)
+
+    def remote_do(self, action):
+        self.qsend.put(action)
+        out = self.qrecv.get()
+        self.assertNotIsInstance(out, BaseException)
+        return out
+
+    @staticmethod
+    def do_terminate():
+        return False, None
+
+
+def remove_env(name):
+    try:
+        del os.environ[name]
+    except KeyError:
+        return False
+    else:
+        return True
+
+
+def check_lib_lookup(qout, qin):
+    status = True
+    while status:
+        try:
+            action = qin.get()
+        except Exception as e:
+            qout.put(e)
+            status = False
+        else:
+            try:
+                with warnings.catch_warnings(record=True) as w:
+                    warnings.simplefilter("always", NumbaWarning)
+                    status, result = action()
+                qout.put(result + (w,))
+            except Exception as e:
+                qout.put(e)
+                status = False
+
+
+@skip_on_cudasim("Library detection unsupported in the simulator")
+@unittest.skipUnless(has_mp_get_context, "mp.get_context not available")
+@skip_unless_conda_cudatoolkit("test assumes conda installed cudatoolkit")
+class TestLibDeviceLookUp(LibraryLookupBase):
+    def test_libdevice_path_decision(self):
+        # Check that the default is using conda environment
+        by, info, warns = self.remote_do(self.do_clear_envs)
+        if has_cuda:
+            self.assertEqual(by, "Conda environment")
+        else:
+            self.assertEqual(by, "<unknown>")
+            self.assertIsNone(info)
+        self.assertFalse(warns)
+        # Check that CUDA_HOME works by removing conda-env
+        by, info, warns = self.remote_do(self.do_set_cuda_home)
+        self.assertEqual(by, "CUDA_HOME")
+        self.assertTrue(
+            info.startswith(os.path.join("mycudahome", "nvvm", "libdevice"))
+        )
+        self.assertFalse(warns)
+
+        if get_system_ctk("nvvm", "libdevice") is None:
+            # Fake remove conda environment so no cudatoolkit is available
+            by, info, warns = self.remote_do(self.do_clear_envs)
+            self.assertEqual(by, "<unknown>")
+            self.assertIsNone(info)
+            self.assertFalse(warns)
+        else:
+            # Use system available cudatoolkit
+            by, info, warns = self.remote_do(self.do_clear_envs)
+            self.assertEqual(by, "System")
+            self.assertFalse(warns)
+
+    @staticmethod
+    def do_clear_envs():
+        remove_env("CUDA_HOME")
+        remove_env("CUDA_PATH")
+        return True, _get_libdevice_path_decision()
+
+    @staticmethod
+    def do_set_cuda_home():
+        os.environ["CUDA_HOME"] = os.path.join("mycudahome")
+        _fake_non_conda_env()
+        return True, _get_libdevice_path_decision()
+
+
+@skip_on_cudasim("Library detection unsupported in the simulator")
+@unittest.skipUnless(has_mp_get_context, "mp.get_context not available")
+@skip_unless_conda_cudatoolkit("test assumes conda installed cudatoolkit")
+class TestNvvmLookUp(LibraryLookupBase):
+    def test_nvvm_path_decision(self):
+        # Check that the default is using conda environment
+        by, info, warns = self.remote_do(self.do_clear_envs)
+        if has_cuda:
+            self.assertEqual(by, "Conda environment")
+        else:
+            self.assertEqual(by, "<unknown>")
+            self.assertIsNone(info)
+        self.assertFalse(warns)
+        # Check that CUDA_HOME works by removing conda-env
+        by, info, warns = self.remote_do(self.do_set_cuda_home)
+        self.assertEqual(by, "CUDA_HOME")
+        self.assertFalse(warns)
+        if IS_WIN32:
+            self.assertEqual(
+                os.path.dirname(info), os.path.join("mycudahome", "nvvm", "bin")
+            )
+        else:
+            self.assertEqual(
+                os.path.dirname(info),
+                os.path.join("mycudahome", "nvvm", "lib64"),
+            )
+
+        if get_system_ctk("nvvm") is None:
+            # Fake remove conda environment so no cudatoolkit is available
+            by, info, warns = self.remote_do(self.do_clear_envs)
+            self.assertEqual(by, "<unknown>")
+            self.assertIsNone(info)
+            self.assertFalse(warns)
+        else:
+            # Use system available cudatoolkit
+            by, info, warns = self.remote_do(self.do_clear_envs)
+            self.assertEqual(by, "System")
+            self.assertFalse(warns)
+
+    @staticmethod
+    def do_clear_envs():
+        remove_env("CUDA_HOME")
+        remove_env("CUDA_PATH")
+        return True, _get_nvvm_path_decision()
+
+    @staticmethod
+    def do_set_cuda_home():
+        os.environ["CUDA_HOME"] = os.path.join("mycudahome")
+        _fake_non_conda_env()
+        return True, _get_nvvm_path_decision()
+
+
+@skip_on_cudasim("Library detection unsupported in the simulator")
+@unittest.skipUnless(has_mp_get_context, "mp.get_context not available")
+@skip_unless_conda_cudatoolkit("test assumes conda installed cudatoolkit")
+class TestCudaLibLookUp(LibraryLookupBase):
+    def test_cudalib_path_decision(self):
+        # Check that the default is using conda environment
+        by, info, warns = self.remote_do(self.do_clear_envs)
+        if has_cuda:
+            self.assertEqual(by, "Conda environment")
+        else:
+            self.assertEqual(by, "<unknown>")
+            self.assertIsNone(info)
+        self.assertFalse(warns)
+
+        # Check that CUDA_HOME works by removing conda-env
+        self.remote_do(self.do_clear_envs)
+        by, info, warns = self.remote_do(self.do_set_cuda_home)
+        self.assertEqual(by, "CUDA_HOME")
+        self.assertFalse(warns)
+        if IS_WIN32:
+            # I think only wheels don't have the "Library" directory?
+            self.assertTrue(
+                info
+                in (
+                    os.path.join("mycudahome", "bin"),
+                    os.path.join("mycudahome", "Library", "bin"),
+                )
+            )
+        else:
+            self.assertEqual(info, os.path.join("mycudahome", "lib64"))
+        if get_system_ctk_libdir() is None:
+            # Fake remove conda environment so no cudatoolkit is available
+            by, info, warns = self.remote_do(self.do_clear_envs)
+            self.assertEqual(by, "<unknown>")
+            self.assertIsNone(info)
+            self.assertFalse(warns)
+        else:
+            # Use system available cudatoolkit
+            by, info, warns = self.remote_do(self.do_clear_envs)
+            self.assertEqual(by, "System")
+            self.assertFalse(warns)
+
+    @staticmethod
+    def do_clear_envs():
+        remove_env("CUDA_HOME")
+        remove_env("CUDA_PATH")
+        return True, _get_cudalib_dir_path_decision()
+
+    @staticmethod
+    def do_set_cuda_home():
+        os.environ["CUDA_HOME"] = os.path.join("mycudahome")
+        _fake_non_conda_env()
+        return True, _get_cudalib_dir_path_decision()
+
+
+def _fake_non_conda_env():
+    """
+    Monkeypatch sys.prefix to hide the fact we are in a conda-env
+    """
+    sys.prefix = ""
+
+
+if __name__ == "__main__":
+    unittest.main()

numba_cuda/numba/cuda/tests/nocuda/test_nvvm.py

@@ -0,0 +1,59 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+from numba.cuda.cudadrv import nvvm
+from numba.cuda.testing import skip_on_cudasim
+from numba.cuda import utils
+
+from llvmlite import ir
+from llvmlite import binding as llvm
+
+import unittest
+
+
+original = (
+    "call void @llvm.memset.p0i8.i64("
+    "i8* align 4 %arg.x.41, i8 0, i64 %0, i1 false)"
+)
+
+missing_align = (
+    "call void @llvm.memset.p0i8.i64(i8* %arg.x.41, i8 0, i64 %0, i1 false)"
+)
+
+
+@skip_on_cudasim("libNVVM not supported in simulator")
+@unittest.skipIf(utils.MACHINE_BITS == 32, "CUDA not support for 32-bit")
+@unittest.skipIf(not nvvm.is_available(), "No libNVVM")
+class TestNvvmWithoutCuda(unittest.TestCase):
+    def test_nvvm_accepts_encoding(self):
+        # Test that NVVM will accept a constant containing all possible 8-bit
+        # characters. Taken from the test case added in llvmlite PR #53:
+        #
+        #     https://github.com/numba/llvmlite/pull/53
+        #
+        # This test case is included in Numba to ensure that the encoding used
+        # by llvmlite (e.g. utf-8, latin1, etc.) does not result in an input to
+        # NVVM that it cannot parse correctly
+
+        # Create a module with a constant containing all 8-bit characters
+        c = ir.Constant(ir.ArrayType(ir.IntType(8), 256), bytearray(range(256)))
+        m = ir.Module()
+        m.triple = "nvptx64-nvidia-cuda"
+        nvvm.add_ir_version(m)
+        gv = ir.GlobalVariable(m, c.type, "myconstant")
+        gv.global_constant = True
+        gv.initializer = c
+        m.data_layout = nvvm.NVVM().data_layout
+
+        # Parse with LLVM then dump the parsed module into NVVM
+        parsed = llvm.parse_assembly(str(m))
+        ptx = nvvm.compile_ir(str(parsed))
+
+        # Ensure all characters appear in the generated constant array.
+        elements = ", ".join([str(i) for i in range(256)])
+        myconstant = f"myconstant[256] = {{{elements}}}".encode("utf-8")
+        self.assertIn(myconstant, ptx)
+
+
+if __name__ == "__main__":
+    unittest.main()

numba_cuda/numba/cuda/tests/nrt/__init__.py

@@ -0,0 +1,9 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+from numba.cuda.tests import load_testsuite
+import os
+
+
+def load_tests(loader, tests, pattern):
+    return load_testsuite(loader, os.path.dirname(__file__))

numba_cuda/numba/cuda/tests/nrt/test_nrt.py

@@ -0,0 +1,387 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import re
+import os
+
+import numpy as np
+import unittest
+from numba.cuda.testing import CUDATestCase, skip_on_cudasim
+from numba.cuda.tests.support import run_in_subprocess, override_config
+from numba.cuda import get_current_device
+from numba.cuda.cudadrv.nvrtc import compile
+from numba.cuda import types
+from numba.cuda.typing import signature
+from numba import cuda
+from numba.cuda import config
+from numba.cuda.typing.templates import AbstractTemplate
+from numba.cuda.cudadrv.linkable_code import (
+    CUSource,
+    PTXSource,
+    Fatbin,
+    Cubin,
+    Archive,
+    Object,
+)
+
+TEST_BIN_DIR = os.getenv("NUMBA_CUDA_TEST_BIN_DIR")
+
+if not config.ENABLE_CUDASIM:
+    from numba.cuda.memory_management.nrt import rtsys, get_include
+    from numba.cuda.cudadecl import registry as cuda_decl_registry
+    from numba.cuda.cudaimpl import lower as cuda_lower
+
+    def allocate_deallocate_handle():
+        """
+        Handle to call NRT_Allocate and NRT_Free
+        """
+        pass
+
+    @cuda_decl_registry.register_global(allocate_deallocate_handle)
+    class AllocateShimImpl(AbstractTemplate):
+        def generic(self, args, kws):
+            return signature(types.void)
+
+    device_fun_shim = cuda.declare_device(
+        "device_allocate_deallocate", types.int32()
+    )
+
+    # wrapper to turn the above into a python callable
+    def call_device_fun_shim():
+        return device_fun_shim()
+
+    @cuda_lower(allocate_deallocate_handle)
+    def allocate_deallocate_impl(context, builder, sig, args):
+        sig_ = types.int32()
+        # call the external function, passing the pointer
+        result = context.compile_internal(
+            builder,
+            call_device_fun_shim,
+            sig_,
+            (),
+        )
+
+        return result
+
+    if TEST_BIN_DIR:
+
+        def make_linkable_code(name, kind, mode):
+            path = os.path.join(TEST_BIN_DIR, name)
+            with open(path, mode) as f:
+                contents = f.read()
+            return kind(contents, nrt=True)
+
+        nrt_extern_a = make_linkable_code("nrt_extern.a", Archive, "rb")
+        nrt_extern_cubin = make_linkable_code("nrt_extern.cubin", Cubin, "rb")
+        nrt_extern_cu = make_linkable_code(
+            "nrt_extern.cu",
+            CUSource,
+            "rb",
+        )
+        nrt_extern_fatbin = make_linkable_code(
+            "nrt_extern.fatbin", Fatbin, "rb"
+        )
+        nrt_extern_fatbin_multi = make_linkable_code(
+            "nrt_extern_multi.fatbin", Fatbin, "rb"
+        )
+        nrt_extern_o = make_linkable_code("nrt_extern.o", Object, "rb")
+        nrt_extern_ptx = make_linkable_code("nrt_extern.ptx", PTXSource, "rb")
+
+
+class TestNrtBasic(CUDATestCase):
+    def run(self, result=None):
+        with override_config("CUDA_ENABLE_NRT", True):
+            super(TestNrtBasic, self).run(result)
+
+    def test_nrt_launches(self):
+        @cuda.jit
+        def f(x):
+            return x[:5]
+
+        @cuda.jit
+        def g():
+            x = np.empty(10, np.int64)
+            f(x)
+
+        g[1, 1]()
+        cuda.synchronize()
+
+    @skip_on_cudasim("CUDA Simulator does not produce PTX")
+    def test_nrt_ptx_contains_refcount(self):
+        @cuda.jit
+        def f(x):
+            return x[:5]
+
+        @cuda.jit
+        def g():
+            x = np.empty(10, np.int64)
+            f(x)
+
+        g[1, 1]()
+
+        ptx = next(iter(g.inspect_asm().values()))
+
+        # The following checks that a `call` PTX instruction is
+        # emitted for NRT_MemInfo_alloc_aligned, NRT_incref and
+        # NRT_decref
+        p1 = r"call\.uni(.|\n)*NRT_MemInfo_alloc_aligned"
+        match = re.search(p1, ptx)
+        assert match is not None
+
+        p2 = r"call\.uni.*\n?.*NRT_incref"
+        match = re.search(p2, ptx)
+        assert match is not None
+
+        p3 = r"call\.uni.*\n?.*NRT_decref"
+        match = re.search(p3, ptx)
+        assert match is not None
+
+    def test_nrt_returns_correct(self):
+        @cuda.jit
+        def f(x):
+            return x[5:]
+
+        @cuda.jit
+        def g(out_ary):
+            x = np.empty(10, np.int64)
+            x[5] = 1
+            y = f(x)
+            out_ary[0] = y[0]
+
+        out_ary = np.zeros(1, dtype=np.int64)
+
+        g[1, 1](out_ary)
+
+        self.assertEqual(out_ary[0], 1)
+
+
+class TestNrtLinking(CUDATestCase):
+    def run(self, result=None):
+        with override_config("CUDA_ENABLE_NRT", True):
+            super(TestNrtLinking, self).run(result)
+
+    @skip_on_cudasim("CUDA Simulator does not link PTX")
+    def test_nrt_detect_linked_ptx_file(self):
+        src = f"#include <{get_include()}/nrt.cuh>"
+        src += """
+                 extern "C" __device__ int device_allocate_deallocate(int* nb_retval){
+                     auto ptr = NRT_Allocate(1);
+                     NRT_Free(ptr);
+                     return 0;
+                 }
+        """
+        cc = get_current_device().compute_capability
+        ptx, _ = compile(src, "external_nrt.cu", cc)
+
+        @cuda.jit(link=[PTXSource(ptx.code, nrt=True)])
+        def kernel():
+            allocate_deallocate_handle()
+
+        kernel[1, 1]()
+
+    @unittest.skipIf(not TEST_BIN_DIR, "necessary binaries not generated.")
+    @skip_on_cudasim("CUDA Simulator does not link code")
+    def test_nrt_detect_linkable_code(self):
+        codes = (
+            nrt_extern_a,
+            nrt_extern_cubin,
+            nrt_extern_cu,
+            nrt_extern_fatbin,
+            nrt_extern_fatbin_multi,
+            nrt_extern_o,
+            nrt_extern_ptx,
+        )
+        for code in codes:
+            with self.subTest(code=code):
+
+                @cuda.jit(link=[code])
+                def kernel():
+                    allocate_deallocate_handle()
+
+                kernel[1, 1]()
+
+
+@skip_on_cudasim("CUDASIM does not have NRT statistics")
+class TestNrtStatistics(CUDATestCase):
+    def setUp(self):
+        self._stream = cuda.default_stream()
+        # Store the current stats state
+        self.__stats_state = rtsys.memsys_stats_enabled(self._stream)
+
+    def tearDown(self):
+        # Set stats state back to whatever it was before the test ran
+        if self.__stats_state:
+            rtsys.memsys_enable_stats(self._stream)
+        else:
+            rtsys.memsys_disable_stats(self._stream)
+
+    def test_stats_env_var_explicit_on(self):
+        # Checks that explicitly turning the stats on via the env var works.
+        src = """if 1:
+        from numba import cuda
+        from numba.cuda.memory_management import rtsys
+        import numpy as np
+
+        @cuda.jit
+        def foo():
+            x = np.arange(10)[0]
+
+        # initialize the NRT before use
+        rtsys.initialize()
+        assert rtsys.memsys_stats_enabled(), "Stats not enabled"
+        orig_stats = rtsys.get_allocation_stats()
+        foo[1, 1]()
+        new_stats = rtsys.get_allocation_stats()
+        total_alloc = new_stats.alloc - orig_stats.alloc
+        total_free = new_stats.free - orig_stats.free
+        total_mi_alloc = new_stats.mi_alloc - orig_stats.mi_alloc
+        total_mi_free = new_stats.mi_free - orig_stats.mi_free
+
+        expected = 1
+        assert total_alloc == expected, \\
+            f"total_alloc != expected, {total_alloc} != {expected}"
+        assert total_free == expected, \\
+            f"total_free != expected, {total_free} != {expected}"
+        assert total_mi_alloc == expected, \\
+            f"total_mi_alloc != expected, {total_mi_alloc} != {expected}"
+        assert total_mi_free == expected, \\
+            f"total_mi_free != expected, {total_mi_free} != {expected}"
+        """
+
+        # Check env var explicitly being set works
+        env = os.environ.copy()
+        env["NUMBA_CUDA_NRT_STATS"] = "1"
+        env["NUMBA_CUDA_ENABLE_NRT"] = "1"
+        run_in_subprocess(src, env=env)
+
+    def check_env_var_off(self, env):
+        src = """if 1:
+        from numba import cuda
+        import numpy as np
+        from numba.cuda.memory_management import rtsys
+
+        @cuda.jit
+        def foo():
+            arr = np.arange(10)[0]
+
+        assert rtsys.memsys_stats_enabled() == False
+        try:
+            rtsys.get_allocation_stats()
+        except RuntimeError as e:
+            assert "NRT stats are disabled." in str(e)
+        """
+        run_in_subprocess(src, env=env)
+
+    def test_stats_env_var_explicit_off(self):
+        # Checks that explicitly turning the stats off via the env var works.
+        env = os.environ.copy()
+        env["NUMBA_CUDA_NRT_STATS"] = "0"
+        self.check_env_var_off(env)
+
+    def test_stats_env_var_default_off(self):
+        # Checks that the env var not being set is the same as "off", i.e.
+        # default for Numba is off.
+        env = os.environ.copy()
+        env.pop("NUMBA_CUDA_NRT_STATS", None)
+        self.check_env_var_off(env)
+
+    def test_stats_status_toggle(self):
+        @cuda.jit
+        def foo():
+            tmp = np.ones(3)
+            arr = np.arange(5 * tmp[0])  # noqa: F841
+            return None
+
+        with (
+            override_config("CUDA_ENABLE_NRT", True),
+            override_config("CUDA_NRT_STATS", True),
+        ):
+            # Switch on stats
+            rtsys.memsys_enable_stats()
+            # check the stats are on
+            self.assertTrue(rtsys.memsys_stats_enabled())
+
+            for i in range(2):
+                # capture the stats state
+                stats_1 = rtsys.get_allocation_stats()
+                # Switch off stats
+                rtsys.memsys_disable_stats()
+                # check the stats are off
+                self.assertFalse(rtsys.memsys_stats_enabled())
+                # run something that would move the counters were they enabled
+                foo[1, 1]()
+                # Switch on stats
+                rtsys.memsys_enable_stats()
+                # check the stats are on
+                self.assertTrue(rtsys.memsys_stats_enabled())
+                # capture the stats state (should not have changed)
+                stats_2 = rtsys.get_allocation_stats()
+                # run something that will move the counters
+                foo[1, 1]()
+                # capture the stats state (should have changed)
+                stats_3 = rtsys.get_allocation_stats()
+                # check stats_1 == stats_2
+                self.assertEqual(stats_1, stats_2)
+                # check stats_2 < stats_3
+                self.assertLess(stats_2, stats_3)
+
+    def test_rtsys_stats_query_raises_exception_when_disabled(self):
+        # Checks that the standard rtsys.get_allocation_stats() query raises
+        # when stats counters are turned off.
+
+        rtsys.memsys_disable_stats()
+        self.assertFalse(rtsys.memsys_stats_enabled())
+
+        with self.assertRaises(RuntimeError) as raises:
+            rtsys.get_allocation_stats()
+
+        self.assertIn("NRT stats are disabled.", str(raises.exception))
+
+    def test_nrt_explicit_stats_query_raises_exception_when_disabled(self):
+        # Checks the various memsys_get_stats functions raise if queried when
+        # the stats counters are disabled.
+        method_variations = ("alloc", "free", "mi_alloc", "mi_free")
+        for meth in method_variations:
+            stats_func = getattr(rtsys, f"memsys_get_stats_{meth}")
+            with self.subTest(stats_func=stats_func):
+                # Turn stats off
+                rtsys.memsys_disable_stats()
+                self.assertFalse(rtsys.memsys_stats_enabled())
+                with self.assertRaises(RuntimeError) as raises:
+                    stats_func()
+                self.assertIn("NRT stats are disabled.", str(raises.exception))
+
+    def test_read_one_stat(self):
+        @cuda.jit
+        def foo():
+            tmp = np.ones(3)
+            arr = np.arange(5 * tmp[0])  # noqa: F841
+            return None
+
+        with (
+            override_config("CUDA_ENABLE_NRT", True),
+            override_config("CUDA_NRT_STATS", True),
+        ):
+            # Switch on stats
+            rtsys.memsys_enable_stats()
+
+            # Launch the kernel a couple of times to increase stats
+            foo[1, 1]()
+            foo[1, 1]()
+
+            # Get stats struct and individual stats
+            stats = rtsys.get_allocation_stats()
+            stats_alloc = rtsys.memsys_get_stats_alloc()
+            stats_mi_alloc = rtsys.memsys_get_stats_mi_alloc()
+            stats_free = rtsys.memsys_get_stats_free()
+            stats_mi_free = rtsys.memsys_get_stats_mi_free()
+
+            # Check individual stats match stats struct
+            self.assertEqual(stats.alloc, stats_alloc)
+            self.assertEqual(stats.mi_alloc, stats_mi_alloc)
+            self.assertEqual(stats.free, stats_free)
+            self.assertEqual(stats.mi_free, stats_mi_free)
+
+
+if __name__ == "__main__":
+    unittest.main()