Skip to content

Commit e2ee35f

Browse files
committed
Build RMM's Python bindings with PTDS
Defines `CUDA_API_PER_THREAD_DEFAULT_STREAM` for Setuptools `Extensions` built by `setup.py`. This should make sure the per-thread default stream is used by these Python extensions. https://docs.nvidia.com/cuda/cuda-runtime-api/stream-sync-behavior.html
1 parent ba8fad5 commit e2ee35f

File tree

2 files changed

+7
-0
lines changed

2 files changed

+7
-0
lines changed

CHANGELOG.md

+1
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
## Improvements
66

77
- PR #477 Just use `None` for `strides` in `DeviceBuffer`
8+
- PR #480 Build RMM's Python bindings with PTDS
89

910
## Bug Fixes
1011

python/setup.py

+6
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,10 @@ def get_cuda_version_from_header(cuda_include_dir):
7777
except Exception:
7878
nthreads = 0
7979

80+
define_macros = [
81+
("CUDA_API_PER_THREAD_DEFAULT_STREAM", None),
82+
]
83+
8084
include_dirs = [
8185
"../include/rmm",
8286
"../include",
@@ -125,6 +129,7 @@ def get_cuda_version_from_header(cuda_include_dir):
125129
],
126130
libraries=["cuda", "rmm"],
127131
language="c++",
132+
define_macros=define_macros,
128133
extra_compile_args=["-std=c++14"],
129134
)
130135
],
@@ -148,6 +153,7 @@ def get_cuda_version_from_header(cuda_include_dir):
148153
],
149154
libraries=["cuda", "rmm"],
150155
language="c++",
156+
define_macros=define_macros,
151157
extra_compile_args=["-std=c++14"],
152158
)
153159
],

0 commit comments

Comments
 (0)