@@ -40,11 +40,46 @@ if TYPE_CHECKING:
4040 from cuda.core.typing import DevicePointerType
4141
4242
43+ # Pool attribute reads and writes and cuMemPoolTrimTo are potentially unsafe
44+ # calls during stream capture: the driver refuses them, and invalidates the
45+ # capture, while a global-mode capture is active on any thread or a
46+ # thread-local one on the calling thread. Capture mode is a per-thread
47+ # property, so the calling thread is switched to relaxed mode around such a
48+ # call and its previous mode is restored afterwards. This has no observable
49+ # effect when nothing is capturing. The call executes immediately rather than
50+ # being recorded into a graph, which is the intent for a process-wide pool.
51+ cdef inline int _enter_relaxed_capture_mode(cydriver.CUstreamCaptureMode* mode) except ?- 1 nogil:
52+ # On return, mode holds the thread's previous capture mode.
53+ mode[0 ] = cydriver.CU_STREAM_CAPTURE_MODE_RELAXED
54+ return HANDLE_RETURN(cydriver.cuThreadExchangeStreamCaptureMode(mode))
55+
56+
57+ cdef int _exit_relaxed_capture_mode(
58+ cydriver.CUresult err, cydriver.CUstreamCaptureMode* mode) except ?- 1 nogil:
59+ """ Restore the thread's capture mode, then raise if the call or the restore failed."""
60+ cdef cydriver.CUresult restore_err = cydriver.cuThreadExchangeStreamCaptureMode(mode)
61+ if err != cydriver.CUDA_SUCCESS:
62+ with gil:
63+ try :
64+ HANDLE_RETURN(err)
65+ except :
66+ if restore_err != cydriver.CUDA_SUCCESS:
67+ # The call's error propagates with the restore failure
68+ # attached as a note (error handling policy).
69+ attach_rollback_failure(
70+ b" cuThreadExchangeStreamCaptureMode" , restore_err,
71+ b" failed to restore the thread's stream capture mode" )
72+ raise
73+ return HANDLE_RETURN(restore_err)
74+
75+
4376cdef class _MemPoolAttributes:
4477 """ Provides access to memory pool attributes.
4578
4679 Setting an attribute changes the memory pool itself, which other memory
4780 resources that share the pool, such as a device's default pool, also see.
81+ Reading or setting an attribute does not affect a stream capture in
82+ progress on any thread.
4883 """
4984 cdef:
5085 MemoryPoolHandle _h_pool
@@ -65,13 +100,21 @@ cdef class _MemPoolAttributes:
65100 )
66101
67102 cdef int _getattribute(self , cydriver.CUmemPool_attribute attr_enum , void* value ) except?-1:
103+ cdef cydriver.CUstreamCaptureMode mode
104+ cdef cydriver.CUresult err
68105 with nogil:
69- HANDLE_RETURN(cydriver.cuMemPoolGetAttribute(as_cu(self._h_pool ), attr_enum , value ))
106+ _enter_relaxed_capture_mode(&mode )
107+ err = cydriver.cuMemPoolGetAttribute(as_cu(self ._h_pool), attr_enum, value)
108+ _exit_relaxed_capture_mode(err , &mode )
70109 return 0
71110
72111 cdef int _setattribute(self , cydriver.CUmemPool_attribute attr_enum , void* value ) except?-1:
112+ cdef cydriver.CUstreamCaptureMode mode
113+ cdef cydriver.CUresult err
73114 with nogil:
74- HANDLE_RETURN(cydriver.cuMemPoolSetAttribute(as_cu(self._h_pool ), attr_enum , value ))
115+ _enter_relaxed_capture_mode(&mode )
116+ err = cydriver.cuMemPoolSetAttribute(as_cu(self ._h_pool), attr_enum, value)
117+ _exit_relaxed_capture_mode(err , &mode )
75118 return 0
76119
77120 cdef int _setflag(self , cydriver.CUmemPool_attribute attr_enum , bint flag ) except?-1:
@@ -256,7 +299,8 @@ cdef class _MemPool(MemoryResource):
256299
257300 Memory that backs live allocations is not released , nor is memory
258301 freed on a stream until the host has observed the free through a
259- synchronization.
302+ synchronization. Trimming does not affect a stream capture in progress
303+ on any thread.
260304
261305 Parameters
262306 ----------
@@ -265,8 +309,12 @@ cdef class _MemPool(MemoryResource):
265309 which releases all unused memory.
266310 """
267311 MP_check_open(self )
312+ cdef cydriver.CUstreamCaptureMode mode
313+ cdef cydriver.CUresult err
268314 with nogil:
269- HANDLE_RETURN(cydriver.cuMemPoolTrimTo(as_cu(self._h_pool ), min_bytes_to_keep ))
315+ _enter_relaxed_capture_mode(&mode )
316+ err = cydriver.cuMemPoolTrimTo(as_cu(self ._h_pool), min_bytes_to_keep)
317+ _exit_relaxed_capture_mode(err , &mode )
270318
271319 @property
272320 @cython.critical_section
@@ -388,24 +436,16 @@ cdef int MP_raise_release_threshold(_MemPool self) except? -1:
388436 the OS as soon as there are no active suballocations. Setting it to
389437 ULLONG_MAX avoids repeated OS round-trips.
390438
391- Pool attribute reads and writes are potentially unsafe calls while the
392- calling thread is inside a global or thread-local stream capture: the
393- driver refuses them and invalidates the capture. Capture mode is a
394- per-thread property , so the thread is switched to relaxed mode around the
395- two calls and its previous mode is restored afterwards. This has no
396- observable effect when the thread is not capturing. The attribute write
397- executes immediately rather than being recorded into a graph , which is the
398- intent for a process-wide pool setting.
439+ The read and the write run in relaxed capture mode , like every pool
440+ attribute access (see ``_enter_relaxed_capture_mode``).
399441 """
400442 MP_check_open(self )
401443 cdef cydriver.cuuint64_t current_threshold
402444 cdef cydriver.cuuint64_t max_threshold = ULLONG_MAX
403- cdef cydriver.CUstreamCaptureMode mode = cydriver.CU_STREAM_CAPTURE_MODE_RELAXED
445+ cdef cydriver.CUstreamCaptureMode mode
404446 cdef cydriver.CUresult err
405- cdef cydriver.CUresult restore_err
406447 with nogil:
407- # On return , mode holds the thread's previous capture mode.
408- HANDLE_RETURN(cydriver.cuThreadExchangeStreamCaptureMode(&mode ))
448+ _enter_relaxed_capture_mode(&mode )
409449 err = cydriver.cuMemPoolGetAttribute(
410450 as_cu(self ._h_pool),
411451 cydriver.CUmemPool_attribute.CU_MEMPOOL_ATTR_RELEASE_THRESHOLD,
@@ -417,18 +457,7 @@ cdef int MP_raise_release_threshold(_MemPool self) except? -1:
417457 cydriver.CUmemPool_attribute.CU_MEMPOOL_ATTR_RELEASE_THRESHOLD,
418458 & max_threshold
419459 )
420- restore_err = cydriver.cuThreadExchangeStreamCaptureMode(& mode)
421- try:
422- HANDLE_RETURN(err )
423- except:
424- if restore_err != cydriver.CUDA_SUCCESS:
425- # The attribute error propagates with the restore failure attached
426- # as a note (error handling policy ).
427- attach_rollback_failure(
428- b"cuThreadExchangeStreamCaptureMode", restore_err ,
429- b"failed to restore the thread's stream capture mode")
430- raise
431- HANDLE_RETURN(restore_err )
460+ _exit_relaxed_capture_mode(err , &mode )
432461 return 0
433462
434463
0 commit comments