nvmath-python 1.0.0__cp314-cp314t-win_amd64.whl

nvmath/internal/__init__.pxd

@@ -0,0 +1,3 @@
+# Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0

nvmath/internal/__init__.py

@@ -0,0 +1,10 @@
+# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+
+"""APIs in this module are for NVIDIA-use only.
+
+APIs in internal are not intended for third-party developers, but are shared between NVIDIA
+packages.
+"""

nvmath/internal/_bindings.pxd

@@ -0,0 +1,18 @@
+# Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from libc.stdint cimport intptr_t, int64_t, uint64_t
+
+cpdef memcpy_async(intptr_t dst_ptr, intptr_t src_ptr, int64_t size, intptr_t stream)
+cpdef memset_async(int itemsize, intptr_t dst_ptr, intptr_t src_ptr, int64_t count, intptr_t stream)
+cpdef stream_sync(intptr_t stream)
+cpdef uint64_t get_memory_pool_release_threshold(pool) except? -1
+cpdef int set_memory_pool_release_threshold(pool, uint64_t threshold) except -1
+cpdef uint64_t get_memory_pool_reserved_memory_size(pool) except? -1
+cpdef uint64_t get_memory_pool_used_memory_size(pool) except? -1
+cpdef free_memory_pool_reserved_memory(pool)
+cpdef launch_kernel(intptr_t f, intptr_t kernel_params, int gx, int gy, int gz, int bx, int by, int bz, unsigned int shared_mem_bytes, intptr_t stream_handle)
+cpdef int get_runtime_version() except? -1
+cpdef int get_driver_version() except? -1
+cpdef int set_memory_pool_access(intptr_t pool, int device_id) except -1

nvmath/internal/_device_utils.py

@@ -0,0 +1,45 @@
+# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import threading
+
+from cuda.core import Device, system
+
+thread_local = threading.local()
+
+
+def get_device(device_id: int) -> Device:
+    """
+    NOTE: the function DOES NOT change the current cuda device. If it needs
+    to call .set_current(), it immediately restores the previous device
+    before exiting. To temporarily change the current device,
+    please use utils.device_ctx().
+
+    A counterpart to cuda.core.Device(int), with extra initialization.
+    Returns initialized cuda.core.Device(device_id) object, i.e.
+    makes sure to call .set_current() on the instance once per thread.
+    Calling .set_current() is not just about a valid CUDA context in a thread,
+    but specifically cuda.core requirement to initialize the device instance.
+
+    The function keeps a thread-local cache of "initialized"
+    device instances. It mirrors cuda.core.Device(int) caching,
+    adding the set_current() call on the first access.
+    """
+    try:
+        devices = thread_local.devices
+    except AttributeError:
+        thread_local.devices = devices = [None] * system.get_num_devices()
+
+    device = devices[device_id]
+    if device is None:
+        current_device = Device()
+        current_device_id = current_device.device_id
+        try:
+            device = Device(device_id)
+            device.set_current()
+            devices[device_id] = device
+        finally:
+            if current_device_id != device_id:
+                current_device.set_current()
+    return device

nvmath/internal/_layout/__init__.pxd

@@ -0,0 +1,3 @@
+# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0

nvmath/internal/_layout/__init__.py

@@ -0,0 +1,7 @@
+# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from ._layout import SlicedLayout, StridedLayout
+
+__all__ = ["SlicedLayout", "StridedLayout"]