Use cimport where possible

Author: mdboom

cuda_core/cuda/core/system/__init__.py

@@ -25,6 +25,7 @@
     __all__.extend(
         [
             "initialize",
+            "get_nvml_version",
             "Device",
             "DeviceArchitecture",
             "UninitializedError",

cuda_core/cuda/core/system/_nvml_context.pxd

@@ -0,0 +1,7 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+cpdef initialize()
+cpdef bint is_initialized()
+cpdef validate()

cuda_core/cuda/core/system/_nvml_context.pyx

@@ -28,7 +28,7 @@ _NVML_OWNER_PID = 0
 _lock = threading.Lock()
 
 
-def initialize() -> None:
+cpdef initialize():
     """Idempotent (per-process) initialization of Nvidia Management Library (NVML).
 
     Notes
@@ -62,7 +62,7 @@ def initialize() -> None:
             raise RuntimeError(f"Unhandled initialisation state ({_NVML_STATE=}, {_NVML_OWNER_PID=})")
 
 
-def is_initialized() -> bool:
+cpdef bint is_initialized():
     """
     Check whether the NVML context is initialized on this process.
 
@@ -74,7 +74,7 @@ def is_initialized() -> bool:
     return _NVML_STATE == _NVMLState.INITIALIZED and os.getpid() == _NVML_OWNER_PID
 
 
-def validate() -> None:
+cpdef validate():
     """
     Validate NVML state.
 

cuda_core/cuda/core/system/device.pyx

@@ -10,8 +10,8 @@ from typing import Iterable
 
 from cuda.bindings import _nvml as nvml
 
-from ._nvml_context import validate
-from .utils import unpack_bitmask
+from ._nvml_context cimport validate
+from .utils cimport unpack_bitmask
 
 
 class DeviceArchitecture:

cuda_core/cuda/core/system/system.pyx

@@ -15,7 +15,7 @@ HAS_WORKING_NVML = _BINDINGS_VERSION >= (13, 1, 2) or (_BINDINGS_VERSION[0] == 1
 
 if HAS_WORKING_NVML:
     from cuda.bindings import _nvml as nvml
-    from ._nvml_context import validate
+    from ._nvml_context cimport validate
 else:
     from cuda.core._utils.cuda_utils import driver, handle_return, runtime
 

cuda_core/cuda/core/system/utils.pxd

@@ -0,0 +1,9 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+
+from cpython cimport array
+
+
+cpdef list[int] unpack_bitmask(x: list[int] | array.array)

cuda_core/cuda/core/system/utils.pyx

@@ -6,20 +6,6 @@ from cpython cimport array
 from libc.stdint cimport uint64_t
 
 
-cpdef str format_bytes(uint64_t x):
-    """Return formatted string in B, KiB, MiB, GiB or TiB"""
-    if x < 1024:
-        return f"{x} B"
-    elif x < 1024**2:
-        return f"{x / 1024:.2f} KiB"
-    elif x < 1024**3:
-        return f"{x / 1024**2:.2f} MiB"
-    elif x < 1024**4:
-        return f"{x / 1024**3:.2f} GiB"
-    else:
-        return f"{x / 1024**4:.2f} TiB"
-
-
 cpdef list[int] unpack_bitmask(x: list[int] | array.array):
     """
     Unpack a list of integers containing bitmasks.

cuda_core/tests/system/test_system_utils.py

@@ -3,19 +3,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import pytest
-from cuda.core.system.utils import format_bytes, unpack_bitmask
-
-
-def test_format_bytes():
-    assert format_bytes(0) == "0 B"
-    assert format_bytes(1) == "1 B"
-    assert format_bytes(1023) == "1023 B"
-    assert format_bytes(1024) == "1.00 KiB"
-    assert format_bytes(1024**2) == "1.00 MiB"
-    assert format_bytes(1024**3) == "1.00 GiB"
-    assert format_bytes(1024**4) == "1.00 TiB"
-    assert format_bytes(1024**5) == "1024.00 TiB"
-    assert format_bytes(1024**6) == "1048576.00 TiB"
+from cuda.core.system.utils import unpack_bitmask
 
 
 @pytest.mark.parametrize(