feat(python): Implement CUDA build in Python bindings (#547)

This PR updates the build system such that nanoarrow's CUDA capability
can be exposed in Python (it doesn't quite expose any of the
functionality yet). I'm not sure this is the best long-term strategy...I
think we might be able to get CUDA plugged in at runtime because we
don't use much of the API...but I think this would at least enable
distributing a cuda-enabled build on conda-forge.

To build, one has to do something like:

```shell
export NANOARROW_PYTHON_CUDA_HOME=/usr/local/cuda
pip install .
```

On Windows, this would be something like:

```powershell
$Env:NANOARROW_PYTHON_CUDA_HOME = 'C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v12.5'
python -m pip install .
```

Then, you should be able to resolve a CUDA device:

```python
from nanoarrow import device

device.resolve(device.DeviceType.CUDA.value, 0)
#> <nanoarrow.device.Device>
#> - device_type: CUDA <2>
#> - device_id: 0
```

Author: paleolimbot

.github/workflows/build-and-test-device.yaml

@@ -132,6 +132,15 @@ jobs:
           cd build
           ctest -T memcheck .
 
+      - name: Check Python bindings with CUDA
+        if: matrix.config.label == 'with-cuda'
+        env:
+          NANOARROW_PYTHON_CUDA_HOME: "/usr/local/cuda"
+        run: |
+          cd python
+          python3 -m pip install ".[test]" -vv
+          python3 -m pytest -vv
+
       - name: Upload memcheck results
         if: failure() && matrix.config.label == 'default-build'
         uses: actions/upload-artifact@main

CMakeLists.txt

@@ -284,12 +284,12 @@ if(NANOARROW_DEVICE)
 
   # If the bundler didn't already assign the source files for the library, do so here
   if(NOT NANOARROW_BUNDLE)
-    set(NANOARROW_DEVICE_BUILD_SOURCES src/nanoarrow/device/device.c)
+    set(NANOARROW_DEVICE_BUILD_SOURCES src/nanoarrow/device/device.c
+                                       ${NANOARROW_DEVICE_SOURCES_CUDA})
   endif()
 
-  add_library(nanoarrow_device
-              ${NANOARROW_DEVICE_BUILD_SOURCES} ${NANOARROW_DEVICE_SOURCES_METAL}
-              ${NANOARROW_DEVICE_SOURCES_CUDA})
+  add_library(nanoarrow_device ${NANOARROW_DEVICE_BUILD_SOURCES}
+                               ${NANOARROW_DEVICE_SOURCES_METAL})
 
   target_include_directories(nanoarrow_device
                              PUBLIC $<BUILD_INTERFACE:${NANOARROW_BUILD_INCLUDE_DIR}>

ci/scripts/bundle.py

@@ -171,10 +171,11 @@ def bundle_nanoarrow_device(
         yield f"{output_include_dir}/{filename}", content
 
     # Generate sources
-    for filename in ["device.c"]:
-        content = read_content(src_dir / "device" / filename)
-        content = namespace_nanoarrow_includes(content, header_namespace)
-        yield f"{output_source_dir}/nanoarrow_{filename}", content
+    content = concatenate_content(
+        [src_dir / "device" / "device.c", src_dir / "device" / "cuda.c"]
+    )
+    content = namespace_nanoarrow_includes(content, header_namespace)
+    yield f"{output_source_dir}/nanoarrow_device.c", content
 
 
 def bundle_nanoarrow_ipc(

python/bootstrap.py

@@ -22,14 +22,14 @@
 
 
 # Generate the nanoarrow_c.pxd file used by the Cython extensions
-class NanoarrowPxdGenerator:
+class PxdGenerator:
     def __init__(self):
         self._define_regexes()
 
-    def generate_nanoarrow_pxd(self, file_in, file_out):
+    def generate_pxd(self, file_in, file_out):
         file_in_name = pathlib.Path(file_in).name
 
-        # Read the nanoarrow.h header
+        # Read the header
         content = None
         with open(file_in, "r") as input:
             content = input.read()
@@ -38,7 +38,7 @@ def generate_nanoarrow_pxd(self, file_in, file_out):
         content = self.re_comment.sub("", content)
 
         # Replace NANOARROW_MAX_FIXED_BUFFERS with its value
-        content = self.re_max_buffers.sub("3", content)
+        content = self._preprocess_content(content)
 
         # Find typedefs, types, and function definitions
         typedefs = self._find_typedefs(content)
@@ -62,13 +62,7 @@ def generate_nanoarrow_pxd(self, file_in, file_out):
             )
 
             # A few things we add in manually
-            output.write(b"\n")
-            output.write(b"    cdef int NANOARROW_OK\n")
-            output.write(b"    cdef int NANOARROW_MAX_FIXED_BUFFERS\n")
-            output.write(b"    cdef int ARROW_FLAG_DICTIONARY_ORDERED\n")
-            output.write(b"    cdef int ARROW_FLAG_NULLABLE\n")
-            output.write(b"    cdef int ARROW_FLAG_MAP_KEYS_SORTED\n")
-            output.write(b"\n")
+            self._write_defs(output)
 
             for type in types_cython:
                 output.write(type.encode("UTF-8"))
@@ -77,15 +71,21 @@ def generate_nanoarrow_pxd(self, file_in, file_out):
             for typedef in typedefs_cython:
                 output.write(typedef.encode("UTF-8"))
                 output.write(b"\n")
+
             output.write(b"\n")
 
             for func_def in func_defs_cython:
                 output.write(func_def.encode("UTF-8"))
                 output.write(b"\n")
 
+    def _preprocess_content(self, content):
+        return content
+
+    def _write_defs(self, output):
+        pass
+
     def _define_regexes(self):
         self.re_comment = re.compile(r"\s*//[^\n]*")
-        self.re_max_buffers = re.compile(r"NANOARROW_MAX_FIXED_BUFFERS")
         self.re_typedef = re.compile(r"typedef(?P<typedef>[^;]+)")
         self.re_type = re.compile(
             r"(?P<type>struct|union|enum) (?P<name>Arrow[^ ]+) {(?P<body>[^}]*)}"
@@ -167,11 +167,58 @@ def _pxd_header(self):
 
         # cython: language_level = 3
 
-        from libc.stdint cimport int8_t, uint8_t, int16_t, uint16_t
-        from libc.stdint cimport int32_t, uint32_t, int64_t, uint64_t
         """
 
 
+class NanoarrowPxdGenerator(PxdGenerator):
+    def _preprocess_content(self, content):
+        return re.sub(r"NANOARROW_MAX_FIXED_BUFFERS", "3", content)
+
+    def _pxd_header(self):
+        return (
+            super()._pxd_header()
+            + """
+    from libc.stdint cimport int8_t, uint8_t, int16_t, uint16_t
+    from libc.stdint cimport int32_t, uint32_t, int64_t, uint64_t
+    """
+        )
+
+    def _write_defs(self, output):
+        output.write(b"\n")
+        output.write(b"    cdef int NANOARROW_OK\n")
+        output.write(b"    cdef int NANOARROW_MAX_FIXED_BUFFERS\n")
+        output.write(b"    cdef int ARROW_FLAG_DICTIONARY_ORDERED\n")
+        output.write(b"    cdef int ARROW_FLAG_NULLABLE\n")
+        output.write(b"    cdef int ARROW_FLAG_MAP_KEYS_SORTED\n")
+        output.write(b"\n")
+
+
+class NanoarrowDevicePxdGenerator(PxdGenerator):
+    def _preprocess_content(self, content):
+        self.device_names = re.findall("#define (ARROW_DEVICE_[A-Z0-9_]+)", content)
+        return super()._preprocess_content(content)
+
+    def _find_typedefs(self, content):
+        return []
+
+    def _pxd_header(self):
+        return (
+            super()._pxd_header()
+            + """
+    from libc.stdint cimport int32_t, int64_t
+    from nanoarrow_c cimport *
+    """
+        )
+
+    def _write_defs(self, output):
+        output.write(b"\n")
+        output.write(b"    ctypedef int32_t ArrowDeviceType\n")
+        output.write(b"\n")
+        for name in self.device_names:
+            output.write(f"    cdef ArrowDeviceType {name}\n".encode())
+        output.write(b"\n")
+
+
 # Runs cmake -DNANOARROW_BUNDLE=ON if cmake exists or copies nanoarrow.c/h
 # from ../dist if it does not. Running cmake is safer because it will sync
 # any changes from nanoarrow C library sources in the checkout but is not
@@ -228,16 +275,18 @@ def copy_or_generate_nanoarrow_c():
 
 
 # Runs the pxd generator with some information about the file name
-def generate_nanoarrow_pxd():
+def generate_nanoarrow_pxds():
     this_dir = pathlib.Path(__file__).parent.resolve()
-    maybe_nanoarrow_h = this_dir / "vendor/nanoarrow.h"
-    maybe_nanoarrow_pxd = this_dir / "vendor/nanoarrow_c.pxd"
 
-    NanoarrowPxdGenerator().generate_nanoarrow_pxd(
-        maybe_nanoarrow_h, maybe_nanoarrow_pxd
+    NanoarrowPxdGenerator().generate_pxd(
+        this_dir / "vendor" / "nanoarrow.h", this_dir / "vendor" / "nanoarrow_c.pxd"
+    )
+    NanoarrowDevicePxdGenerator().generate_pxd(
+        this_dir / "vendor" / "nanoarrow_device.h",
+        this_dir / "vendor" / "nanoarrow_device_c.pxd",
     )
 
 
 if __name__ == "__main__":
     copy_or_generate_nanoarrow_c()
-    generate_nanoarrow_pxd()
+    generate_nanoarrow_pxds()

python/setup.py

@@ -20,6 +20,7 @@
 import os
 import subprocess
 import sys
+from pathlib import Path
 
 from setuptools import Extension, setup
 
@@ -52,18 +53,45 @@ def get_version(pkg_path):
 
 
 # Set some extra flags for compiling with coverage support
+extra_include_dirs = []
+extra_compile_args = []
+extra_link_args = []
+extra_define_macros = []
+library_dirs = []
+libraries = []
+
 if os.getenv("NANOARROW_PYTHON_COVERAGE") == "1":
-    extra_compile_args = ["--coverage"]
-    extra_link_args = ["--coverage"]
-    extra_define_macros = [("CYTHON_TRACE", 1)]
-elif os.getenv("NANOARROW_DEBUG_EXTENSION") == "1":
-    extra_compile_args = ["-g", "-O0"]
-    extra_link_args = []
-    extra_define_macros = []
-else:
-    extra_compile_args = []
-    extra_link_args = []
-    extra_define_macros = []
+    extra_compile_args.append("--coverage")
+    extra_link_args.append("--coverage")
+    extra_define_macros.append(("CYTHON_TRACE", 1))
+
+if os.getenv("NANOARROW_DEBUG_EXTENSION") == "1":
+    extra_compile_args.extend(["-g", "-O0"])
+
+cuda_toolkit_root = os.getenv("NANOARROW_PYTHON_CUDA_HOME")
+if cuda_toolkit_root:
+    cuda_lib = "cuda.lib" if os.name == "nt" else "libcuda.so"
+    include_dir = Path(cuda_toolkit_root) / "include"
+    possible_libs = [
+        Path(cuda_toolkit_root) / "lib" / cuda_lib,
+        Path(cuda_toolkit_root) / "lib64" / cuda_lib,
+        Path(cuda_toolkit_root) / "lib" / "x64" / cuda_lib,
+        Path("/usr/lib/wsl/lib") / cuda_lib,
+    ]
+
+    if not include_dir.is_dir():
+        raise ValueError(f"CUDA include directory does not exist: '{include_dir}'")
+
+    lib_dirs = [d for d in possible_libs if d.exists()]
+    if not lib_dirs:
+        lib_dirs_err = ", ".join(f"'{d}" for d in possible_libs)
+        raise ValueError(f"Can't find CUDA library directory. Checked {lib_dirs_err}")
+
+    extra_include_dirs.append(str(include_dir))
+    library_dirs.append(str(lib_dirs[0].parent))
+    libraries.append("cuda")
+    extra_define_macros.append(("NANOARROW_DEVICE_WITH_CUDA", 1))
+
 
 setup(
     ext_modules=[
@@ -80,7 +108,7 @@ def get_version(pkg_path):
         ),
         Extension(
             name="nanoarrow._utils",
-            include_dirs=["src/nanoarrow", "vendor"],
+            include_dirs=extra_include_dirs + ["src/nanoarrow", "vendor"],
             language="c",
             sources=[
                 "src/nanoarrow/_utils.pyx",
@@ -92,7 +120,7 @@ def get_version(pkg_path):
         ),
         Extension(
             name="nanoarrow._lib",
-            include_dirs=["src/nanoarrow", "vendor"],
+            include_dirs=extra_include_dirs + ["src/nanoarrow", "vendor"],
             language="c",
             sources=[
                 "src/nanoarrow/_lib.pyx",
@@ -102,10 +130,12 @@ def get_version(pkg_path):
             extra_compile_args=extra_compile_args,
             extra_link_args=extra_link_args,
             define_macros=extra_define_macros,
+            library_dirs=library_dirs,
+            libraries=libraries,
         ),
         Extension(
             name="nanoarrow._ipc_lib",
-            include_dirs=["src/nanoarrow", "vendor"],
+            include_dirs=extra_include_dirs + ["src/nanoarrow", "vendor"],
             language="c",
             sources=[
                 "src/nanoarrow/_ipc_lib.pyx",

python/src/nanoarrow/_lib.pyx

@@ -248,9 +248,13 @@ cdef class Device:
     def resolve(device_type, int64_t device_id):
         if int(device_type) == ARROW_DEVICE_CPU:
             return DEVICE_CPU
-        else:
+
+        cdef ArrowDevice* c_device = ArrowDeviceResolve(device_type, device_id)
+        if c_device == NULL:
             raise ValueError(f"Device not found for type {device_type}/{device_id}")
 
+        return Device(None, <uintptr_t>c_device)
+
 
 # Cache the CPU device
 # The CPU device is statically allocated (so base is None)