Cython refactoring

build.sh

@@ -40,7 +40,7 @@ VERBOSE_FLAG=""
 BUILD_TYPE=Release
 INSTALL_TARGET=install
 RAN_CMAKE=0
-PYTHON_ARGS_FOR_INSTALL="--no-build-isolation --no-deps --config-settings rapidsai.disable-cuda=true"
+PYTHON_ARGS_FOR_INSTALL="-v --no-build-isolation --no-deps --config-settings rapidsai.disable-cuda=true"
 
 
 # Set defaults for vars that may not have been defined externally

python/kvikio/kvikio/__init__.py

@@ -1,18 +1,21 @@
 # Copyright (c) 2021-2024, NVIDIA CORPORATION. All rights reserved.
 # See file LICENSE for terms.
 
-from ._lib import libkvikio  # type: ignore
-from ._version import __git_commit__, __version__  # noqa: F401
-from .cufile import CuFile  # noqa: F401
+from kvikio._lib import buffer, driver_properties  # type: ignore
+from kvikio._version import __git_commit__, __version__
+from kvikio.cufile import CuFile
 
 
 def memory_register(buf) -> None:
-    return libkvikio.memory_register(buf)
+    return buffer.memory_register(buf)
 
 
 def memory_deregister(buf) -> None:
-    libkvikio.memory_deregister(buf)
+    buffer.memory_deregister(buf)
 
 
 # TODO: Wrap nicely, maybe as a dataclass?
-DriverProperties = libkvikio.DriverProperties
+DriverProperties = driver_properties.DriverProperties
+
+
+__all__ = ["__git_commit__", "__version__", "CuFile"]

python/kvikio/kvikio/_lib/CMakeLists.txt

@@ -1,5 +1,5 @@
 # =============================================================================
-# Copyright (c) 2022-2023, NVIDIA CORPORATION.
+# Copyright (c) 2022-2024, NVIDIA CORPORATION.
 #
 # Licensed 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
@@ -12,8 +12,10 @@
 # the License.
 # =============================================================================
 
-# Set the list of Cython files to build
-set(cython_modules arr.pyx libnvcomp.pyx libnvcomp_ll.pyx libkvikio.pyx)
+# Set the list of Cython files to build, one .so per file
+set(cython_modules arr.pyx buffer.pyx defaults.pyx driver_properties.pyx file_handle.pyx future.pyx
+                   libnvcomp.pyx libnvcomp_ll.pyx
+)
 
 rapids_cython_create_modules(
   CXX

python/kvikio/kvikio/_lib/arr.pxd

@@ -1,11 +1,12 @@
-# Copyright (c) 2020-2021, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2020-2024, NVIDIA CORPORATION. All rights reserved.
 # See file LICENSE for terms.
 
 # distutils: language = c++
 # cython: language_level=3
 
 
 from libc.stdint cimport uintptr_t
+from libcpp.utility cimport pair
 
 
 cdef class Array:
@@ -25,3 +26,8 @@ cdef class Array:
     cpdef bint _f_contiguous(self)
     cpdef bint _contiguous(self)
     cpdef Py_ssize_t _nbytes(self)
+
+
+cdef pair[uintptr_t, size_t] parse_buffer_argument(
+    buf, size, bint accept_host_buffer
+) except *

python/kvikio/kvikio/_lib/arr.pyx

@@ -1,4 +1,4 @@
-# Copyright (c) 2020-2021, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2020-2024, NVIDIA CORPORATION. All rights reserved.
 # See file LICENSE for terms.
 
 # cython: language_level=3
@@ -294,3 +294,24 @@ cpdef asarray(obj):
         return obj
     else:
         return Array(obj)
+
+
+cdef pair[uintptr_t, size_t] parse_buffer_argument(
+    buf, size, bint accept_host_buffer
+) except *:
+    """Parse `buf` and `size` argument and return a pointer and nbytes"""
+    if not isinstance(buf, Array):
+        buf = Array(buf)
+    cdef Array arr = buf
+    if not arr._contiguous():
+        raise ValueError("Array must be contiguous")
+    if not accept_host_buffer and not arr.cuda:
+        raise ValueError("Non-CUDA buffers not supported")
+    cdef size_t nbytes
+    if size is None:
+        nbytes = arr.nbytes
+    elif size > arr.nbytes:
+        raise ValueError("Size is greater than the size of the buffer")
+    else:
+        nbytes = size
+    return pair[uintptr_t, size_t](arr.ptr, nbytes)

python/kvikio/kvikio/_lib/buffer.pyx

@@ -0,0 +1,27 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+# See file LICENSE for terms.
+
+# distutils: language = c++
+# cython: language_level=3
+
+
+from kvikio._lib.arr cimport Array
+
+
+cdef extern from "<kvikio/buffer.hpp>" namespace "kvikio" nogil:
+    void cpp_memory_register "kvikio::memory_register"(const void* devPtr) except +
+    void cpp_memory_deregister "kvikio::memory_deregister"(const void* devPtr) except +
+
+
+def memory_register(buf) -> None:
+    if not isinstance(buf, Array):
+        buf = Array(buf)
+    cdef Array arr = buf
+    cpp_memory_register(<void*>arr.ptr)
+
+
+def memory_deregister(buf) -> None:
+    if not isinstance(buf, Array):
+        buf = Array(buf)
+    cdef Array arr = buf
+    cpp_memory_deregister(<void*>arr.ptr)

python/kvikio/kvikio/_lib/defaults.pyx

@@ -0,0 +1,54 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+# See file LICENSE for terms.
+
+# distutils: language = c++
+# cython: language_level=3
+
+from libcpp cimport bool
+
+
+cdef extern from "<kvikio/defaults.hpp>" nogil:
+    bool cpp_compat_mode "kvikio::defaults::compat_mode"() except +
+    void cpp_compat_mode_reset \
+        "kvikio::defaults::compat_mode_reset"(bool enable) except +
+    unsigned int cpp_thread_pool_nthreads \
+        "kvikio::defaults::thread_pool_nthreads"() except +
+    void cpp_thread_pool_nthreads_reset \
+        "kvikio::defaults::thread_pool_nthreads_reset" (unsigned int nthreads) except +
+    size_t cpp_task_size "kvikio::defaults::task_size"() except +
+    void cpp_task_size_reset "kvikio::defaults::task_size_reset"(size_t nbytes) except +
+    size_t cpp_gds_threshold "kvikio::defaults::gds_threshold"() except +
+    void cpp_gds_threshold_reset \
+        "kvikio::defaults::gds_threshold_reset"(size_t nbytes) except +
+
+
+def compat_mode() -> bool:
+    return cpp_compat_mode()
+
+
+def compat_mode_reset(enable: bool) -> None:
+    cpp_compat_mode_reset(enable)
+
+
+def thread_pool_nthreads() -> int:
+    return cpp_thread_pool_nthreads()
+
+
+def thread_pool_nthreads_reset(nthreads: int) -> None:
+    cpp_thread_pool_nthreads_reset(nthreads)
+
+
+def task_size() -> int:
+    return cpp_task_size()
+
+
+def task_size_reset(nbytes: int) -> None:
+    cpp_task_size_reset(nbytes)
+
+
+def gds_threshold() -> int:
+    return cpp_gds_threshold()
+
+
+def gds_threshold_reset(nbytes: int) -> None:
+    cpp_gds_threshold_reset(nbytes)

python/kvikio/kvikio/_lib/driver_properties.pyx

@@ -0,0 +1,85 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+# See file LICENSE for terms.
+
+# distutils: language = c++
+# cython: language_level=3
+
+
+from libcpp cimport bool
+
+
+cdef extern from "<kvikio/driver.hpp>" nogil:
+    cdef cppclass cpp_DriverProperties "kvikio::DriverProperties":
+        cpp_DriverProperties() except +
+        bool is_gds_available() except +
+        unsigned int get_nvfs_major_version() except +
+        unsigned int get_nvfs_minor_version() except +
+        bool get_nvfs_allow_compat_mode() except +
+        bool get_nvfs_poll_mode() except +
+        size_t get_nvfs_poll_thresh_size() except +
+        void set_nvfs_poll_mode(bool enable) except +
+        void set_nvfs_poll_thresh_size(size_t size_in_kb) except +
+        size_t get_max_device_cache_size() except +
+        void set_max_device_cache_size(size_t size_in_kb) except +
+        size_t get_per_buffer_cache_size() except +
+        size_t get_max_pinned_memory_size() except +
+        void set_max_pinned_memory_size(size_t size_in_kb) except +
+
+
+cdef class DriverProperties:
+    cdef cpp_DriverProperties _handle
+
+    @property
+    def is_gds_available(self) -> bool:
+        try:
+            return self._handle.is_gds_available()
+        except RuntimeError:
+            return False
+
+    @property
+    def major_version(self) -> bool:
+        return self._handle.get_nvfs_major_version()
+
+    @property
+    def minor_version(self) -> bool:
+        return self._handle.get_nvfs_minor_version()
+
+    @property
+    def allow_compat_mode(self) -> bool:
+        return self._handle.get_nvfs_allow_compat_mode()
+
+    @property
+    def poll_mode(self) -> bool:
+        return self._handle.get_nvfs_poll_mode()
+
+    @poll_mode.setter
+    def poll_mode(self, enable: bool) -> None:
+        self._handle.set_nvfs_poll_mode(enable)
+
+    @property
+    def poll_thresh_size(self) -> int:
+        return self._handle.get_nvfs_poll_thresh_size()
+
+    @poll_thresh_size.setter
+    def poll_thresh_size(self, size_in_kb: int) -> None:
+        self._handle.set_nvfs_poll_thresh_size(size_in_kb)
+
+    @property
+    def max_device_cache_size(self) -> int:
+        return self._handle.get_max_device_cache_size()
+
+    @max_device_cache_size.setter
+    def max_device_cache_size(self, size_in_kb: int) -> None:
+        self._handle.set_max_device_cache_size(size_in_kb)
+
+    @property
+    def per_buffer_cache_size(self) -> int:
+        return self._handle.get_per_buffer_cache_size()
+
+    @property
+    def max_pinned_memory_size(self) -> int:
+        return self._handle.get_max_pinned_memory_size()
+
+    @max_pinned_memory_size.setter
+    def max_pinned_memory_size(self, size_in_kb: int) -> None:
+        self._handle.set_max_pinned_memory_size(size_in_kb)