CUDAMemory.py

# Decorator to monitor memory of CUDA kernels
# Good read:
# https://github.com/numba/nvidia-cuda-tutorial/#session-5-memory-management

import gc

import numpy as np
from numba import cuda


def fmtsize(num, maxunit="TiB"):
    units = ["B", "KiB", "MiB", "GiB", "TiB"]
    units = units[: units.index(maxunit)]
    sign = np.sign(num)
    num *= sign

    for unit in units:
        if num < 1024.0:
            return f"{sign*num:.1f} {unit}"
        num /= 1024.0
    return f"{sign*num:.1f} {maxunit}"


class CUDAMemory:
    def __init__(self, device_id=None, force_deallocation=False):
        self.membefore = None
        self.memafter = None
        self.original_device = cuda.get_current_device()
        self.new_device = self.original_device
        self.force_deallocation = force_deallocation
        if device_id is not None and device_id != self.original_device.id:
            self.new_device = cuda.select_device(device_id)

    def __enter__(self):
        self.membefore = cuda.current_context().get_memory_info()
        return self

    def __exit__(self, type, value, traceback):
        if self.force_deallocation:  # No idea if these work
            gc.collect()
            cuda.current_context().deallocations.clear()

        self.memafter = cuda.current_context().get_memory_info()
        if self.original_device.id != self.new_device.id:
            cuda.select_device(self.original_device.id)

    def __str__(self):
        if self.membefore is None:
            msgbefore = f"free_before=None, total_before=None"
        else:
            msgbefore = (
                f"free_before={fmtsize(self.membefore.free)}, "
                f"total_before={fmtsize(self.membefore.total)}"
            )
        if self.memafter is None:
            msgafter = f"free_after=None, total_after=None"
        else:
            msgafter = (
                f"free_after={fmtsize(self.memafter.free)}, "
                f"total_after={fmtsize(self.memafter.total)}"
            )
        if self.membefore is not None and self.memafter is not None:
            msgdiff = (
                f"free_diff={fmtsize(self.memafter.free - self.membefore.free)}, "
                f"total_diff={fmtsize(self.memafter.total - self.membefore.total)}"
            )
            return (
                f"CUDAMemory({msgbefore};\n           {msgafter};\n           {msgdiff})"
            )
        else:
            return f"CUDAMemory({msgbefore};\n           {msgafter})"

    def __repr__(self):
        return self.__str__()


if __name__ == "__main__":
    n = 100_000_000
    dtype = np.float32
    itemsize = np.dtype(dtype).itemsize
    with CUDAMemory() as mem:
        out = cuda.device_array((n,), dtype)
    fb = fmtsize(mem.membefore.free, maxunit="MiB")
    fa = fmtsize(mem.memafter.free, maxunit="MiB")
    sz = fmtsize(n * itemsize, maxunit="MiB")
    print(f"Free before {fb} | Free after {fa} | Array size {sz}")
    print(mem)