|
10 | 10 | import pytest
|
11 | 11 | from conftest import skipif_need_cuda_headers
|
12 | 12 |
|
| 13 | +from cuda.bindings import driver |
13 | 14 | from cuda.core.experimental import (
|
14 | 15 | Device,
|
15 | 16 | DeviceMemoryResource,
|
|
19 | 20 | ProgramOptions,
|
20 | 21 | launch,
|
21 | 22 | )
|
| 23 | +from cuda.core.experimental._memory import _SynchronousMemoryResource |
| 24 | +from cuda.core.experimental._utils.cuda_utils import handle_return |
22 | 25 |
|
23 | 26 |
|
24 | 27 | def test_launch_config_init(init_cuda):
|
@@ -211,7 +214,7 @@ def test_cooperative_launch():
|
211 | 214 | @pytest.mark.parametrize(
|
212 | 215 | "memory_resource_class",
|
213 | 216 | [
|
214 |
| - DeviceMemoryResource, |
| 217 | + "device_memory_resource", # kludgy, but can go away after #726 is resolved |
215 | 218 | pytest.param(
|
216 | 219 | LegacyPinnedMemoryResource,
|
217 | 220 | marks=pytest.mark.skipif(
|
@@ -249,9 +252,18 @@ def test_launch_with_buffers_allocated_by_memory_resource(init_cuda, memory_reso
|
249 | 252 | kernel = mod.get_kernel("memory_ops")
|
250 | 253 |
|
251 | 254 | # Create memory resource
|
252 |
| - if memory_resource_class == DeviceMemoryResource: |
253 |
| - mr = memory_resource_class(dev.device_id) |
254 |
| - else: # LegacyPinnedMemoryResource |
| 255 | + if memory_resource_class == "device_memory_resource": |
| 256 | + if ( |
| 257 | + handle_return( |
| 258 | + driver.cuDeviceGetAttribute( |
| 259 | + driver.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_MEMORY_POOLS_SUPPORTED, dev.device_id |
| 260 | + ) |
| 261 | + ) |
| 262 | + ) == 1: |
| 263 | + mr = DeviceMemoryResource(dev.device_id) |
| 264 | + else: |
| 265 | + mr = _SynchronousMemoryResource(dev.device_id) |
| 266 | + else: |
255 | 267 | mr = memory_resource_class()
|
256 | 268 |
|
257 | 269 | # Allocate memory
|
|
0 commit comments