1  # src/examples/python-cuda/BUCK
 2  #
 3  # CUDA kernels exposed to Python via pybind11
 4  #
 5  # NOTE: Requires nvidia-sdk in the toolchain.
 6  # These targets are disabled in aleph-0xff (minimal repo).
 7  # Enable in downstream flakes with nv toolchain configured.
 8  #
 9  # Demonstrates:
10  #   - nvidia_library for CUDA kernels
11  #   - pybind11_extension with nvidia_deps
12  #   - numpy array integration with CUDA
13  #   - Device memory management
14  #
15  # Build: buck2 build //src/examples/python-cuda:gpu_module
16  # Test:  buck2 run //src/examples/python-cuda:test_gpu
17  
18  load("@toolchains//:nv.bzl", "nv_library")
19  load("@toolchains//:python.bzl", "pybind11_extension", "python_script")
20  
21  # # CUDA kernels compiled with clang
22  nv_library(
23      name = "kernels",
24      srcs = ["kernels.cu"],
25      exported_headers = ["kernels.cuh"],
26      visibility = ["PUBLIC"],
27  )
28  
29  # # pybind11 extension linking CUDA kernels
30  pybind11_extension(
31      name = "gpu_module",
32      srcs = ["bindings.cpp"],
33      nv_deps = [":kernels"],
34      visibility = ["PUBLIC"],
35  )
36  
37  # # Python test script
38  python_script(
39      name = "test_gpu",
40      main = "test_gpu.py",
41      deps = [":gpu_module"],
42      visibility = ["PUBLIC"],
43  )