Add missing conditional node support

Fixes issue #55

Author: vzhurba01

cuda/_lib/ccudart/utils.pyx.in

@@ -3410,6 +3410,15 @@ cdef cudaError_t toDriverGraphNodeParams(const cudaGraphNodeParams *rtParams, cc
     elif rtParams[0].type == cudaGraphNodeType.cudaGraphNodeTypeMemFree:
         driverParams[0].type = ccuda.CUgraphNodeType_enum.CU_GRAPH_NODE_TYPE_MEM_FREE
         driverParams[0].free.dptr = <ccuda.CUdeviceptr>rtParams[0].free.dptr
+    elif rtParams[0].type == cudaGraphNodeType.cudaGraphNodeTypeConditional:
+        driverParams[0].type = ccuda.CUgraphNodeType_enum.CU_GRAPH_NODE_TYPE_CONDITIONAL
+        # RT params mirror the driver params except the RT struct lacks the ctx at the end.
+        memcpy(&driverParams[0].conditional, &rtParams[0].conditional, sizeof(rtParams[0].conditional))
+        err = <cudaError_t>ccuda._cuCtxGetCurrent(&context)
+        if err != cudaSuccess:
+            _setLastError(err)
+            return err
+        driverParams[0].conditional.ctx = context
     else:
         return cudaErrorInvalidValue
     return cudaSuccess
@@ -3418,6 +3427,8 @@ cdef cudaError_t toDriverGraphNodeParams(const cudaGraphNodeParams *rtParams, cc
 cdef void toCudartGraphNodeOutParams(const ccuda.CUgraphNodeParams *driverParams, cudaGraphNodeParams *rtParams) nogil:
     if driverParams[0].type == ccuda.CUgraphNodeType_enum.CU_GRAPH_NODE_TYPE_MEM_ALLOC:
         rtParams[0].alloc.dptr = <void *>driverParams[0].alloc.dptr
+    elif driverParams[0].type == ccuda.CUgraphNodeType_enum.CU_GRAPH_NODE_TYPE_CONDITIONAL:
+        rtParams[0].conditional.phGraph_out = <cudaGraph_t *>driverParams[0].conditional.phGraph_out
 
 
 cdef cudaError_t toDriverKernelNodeParams(const cudaKernelNodeParams nodeParams[0], ccuda.CUDA_KERNEL_NODE_PARAMS *driverNodeParams) except ?cudaErrorCallRequiresNewerDriver nogil:

cuda/cuda.pyx.in

@@ -8441,8 +8441,7 @@ cdef class CUDA_CONDITIONAL_NODE_PARAMS:
         self._handle = CUgraphConditionalHandle(_ptr=<void_ptr>&self._ptr[0].handle)
         self._ctx = CUcontext(_ptr=<void_ptr>&self._ptr[0].ctx)
     def __dealloc__(self):
-        if self._phGraph_out is not NULL:
-            free(self._phGraph_out)
+        pass
     def getPtr(self):
         return <void_ptr>self._ptr
     def __repr__(self):
@@ -8501,25 +8500,8 @@ cdef class CUDA_CONDITIONAL_NODE_PARAMS:
         self._ptr[0].size = size
     @property
     def phGraph_out(self):
-        arrs = [<void_ptr>self._ptr[0].phGraph_out + x*sizeof(ccuda.CUgraph) for x in range(self._phGraph_out_length)]
+        arrs = [<void_ptr>self._ptr[0].phGraph_out + x*sizeof(ccuda.CUgraph) for x in range(self.size)]
         return [CUgraph(_ptr=arr) for arr in arrs]
-    @phGraph_out.setter
-    def phGraph_out(self, val):
-        if len(val) == 0:
-            free(self._phGraph_out)
-            self._phGraph_out_length = 0
-            self._ptr[0].phGraph_out = NULL
-        else:
-            if self._phGraph_out_length != <size_t>len(val):
-                free(self._phGraph_out)
-                self._phGraph_out = <ccuda.CUgraph*> calloc(len(val), sizeof(ccuda.CUgraph))
-                if self._phGraph_out is NULL:
-                    raise MemoryError('Failed to allocate length x size memory: ' + str(len(val)) + 'x' + str(sizeof(ccuda.CUgraph)))
-                self._phGraph_out_length = <size_t>len(val)
-                self._ptr[0].phGraph_out = self._phGraph_out
-            for idx in range(len(val)):
-                self._phGraph_out[idx] = (<CUgraph>val[idx])._ptr[0]
-
     @property
     def ctx(self):
         return self._ctx

cuda/cudart.pyx.in

@@ -10878,8 +10878,7 @@ cdef class cudaConditionalNodeParams:
     def __init__(self, void_ptr _ptr = 0):
         self._handle = cudaGraphConditionalHandle(_ptr=<void_ptr>&self._ptr[0].handle)
     def __dealloc__(self):
-        if self._phGraph_out is not NULL:
-            free(self._phGraph_out)
+        pass
     def getPtr(self):
         return <void_ptr>self._ptr
     def __repr__(self):
@@ -10934,25 +10933,8 @@ cdef class cudaConditionalNodeParams:
         self._ptr[0].size = size
     @property
     def phGraph_out(self):
-        arrs = [<void_ptr>self._ptr[0].phGraph_out + x*sizeof(ccudart.cudaGraph_t) for x in range(self._phGraph_out_length)]
+        arrs = [<void_ptr>self._ptr[0].phGraph_out + x*sizeof(ccudart.cudaGraph_t) for x in range(self.size)]
         return [cudaGraph_t(_ptr=arr) for arr in arrs]
-    @phGraph_out.setter
-    def phGraph_out(self, val):
-        if len(val) == 0:
-            free(self._phGraph_out)
-            self._phGraph_out_length = 0
-            self._ptr[0].phGraph_out = NULL
-        else:
-            if self._phGraph_out_length != <size_t>len(val):
-                free(self._phGraph_out)
-                self._phGraph_out = <ccudart.cudaGraph_t*> calloc(len(val), sizeof(ccudart.cudaGraph_t))
-                if self._phGraph_out is NULL:
-                    raise MemoryError('Failed to allocate length x size memory: ' + str(len(val)) + 'x' + str(sizeof(ccudart.cudaGraph_t)))
-                self._phGraph_out_length = <size_t>len(val)
-                self._ptr[0].phGraph_out = self._phGraph_out
-            for idx in range(len(val)):
-                self._phGraph_out[idx] = (<cudaGraph_t>val[idx])._ptr[0]
-
 {{endif}}
 {{if 'struct cudaChildGraphNodeParams' in found_types}}
 

cuda/tests/test_cuda.py

@@ -817,3 +817,33 @@ def test_graph_poly():
     assert(err == cuda.CUresult.CUDA_SUCCESS)
     err, = cuda.cuCtxDestroy(ctx)
     assert(err == cuda.CUresult.CUDA_SUCCESS)
+
+@pytest.mark.skipif(driverVersionLessThan(12030)
+                    or not supportsCudaAPI('cuGraphConditionalHandleCreate'), reason='Conditional graph APIs required')
+def test_conditional():
+    err, = cuda.cuInit(0)
+    assert(err == cuda.CUresult.CUDA_SUCCESS)
+    err, device = cuda.cuDeviceGet(0)
+    assert(err == cuda.CUresult.CUDA_SUCCESS)
+    err, ctx = cuda.cuCtxCreate(0, device)
+    assert(err == cuda.CUresult.CUDA_SUCCESS)
+
+    err, graph = cuda.cuGraphCreate(0)
+    assert(err == cuda.CUresult.CUDA_SUCCESS)
+    err, handle = cuda.cuGraphConditionalHandleCreate(graph, ctx, 0, 0)
+    assert(err == cuda.CUresult.CUDA_SUCCESS)
+
+    params = cuda.CUgraphNodeParams()
+    params.type = cuda.CUgraphNodeType.CU_GRAPH_NODE_TYPE_CONDITIONAL
+    params.conditional.handle = handle
+    params.conditional.type = cuda.CUgraphConditionalNodeType.CU_GRAPH_COND_TYPE_IF
+    params.conditional.size = 1
+    params.conditional.ctx = ctx
+
+    assert(len(params.conditional.phGraph_out) == 1)
+    assert(int(params.conditional.phGraph_out[0]) == 0)
+    err, node = cuda.cuGraphAddNode(graph, None, 0, params)
+    assert(err == cuda.CUresult.CUDA_SUCCESS)
+
+    assert(len(params.conditional.phGraph_out) == 1)
+    assert(int(params.conditional.phGraph_out[0]) != 0)

cuda/tests/test_cudart.py

@@ -32,6 +32,9 @@ def supportsSparseTexturesDeviceFilter():
     err, isSupported = cudart.cudaDeviceGetAttribute(cudart.cudaDeviceAttr.cudaDevAttrSparseCudaArraySupported, 0)
     return isSuccess(err) and isSupported
 
+def supportsCudaAPI(name):
+    return name in dir(cuda) or dir(cudart)
+
 def test_cudart_memcpy():
     # Allocate dev memory
     size = 1024 * np.uint8().itemsize
@@ -1275,3 +1278,26 @@ def task_callback_stream(stream, status, userData):
 def test_cudart_func_callback():
     cudart_func_stream_callback(use_host_api=False)
     cudart_func_stream_callback(use_host_api=True)
+
+
+@pytest.mark.skipif(driverVersionLessThan(12030)
+                    or not supportsCudaAPI('cudaGraphConditionalHandleCreate'), reason='Conditional graph APIs required')
+def test_cudart_conditional():
+    err, graph = cudart.cudaGraphCreate(0)
+    assertSuccess(err)
+    err, handle = cudart.cudaGraphConditionalHandleCreate(graph, 0, 0)
+    assertSuccess(err)
+
+    params = cudart.cudaGraphNodeParams()
+    params.type = cudart.cudaGraphNodeType.cudaGraphNodeTypeConditional
+    params.conditional.handle = handle
+    params.conditional.type = cudart.cudaGraphConditionalNodeType.cudaGraphCondTypeIf
+    params.conditional.size = 1
+
+    assert(len(params.conditional.phGraph_out) == 1)
+    assert(int(params.conditional.phGraph_out[0]) == 0)
+    err, node = cudart.cudaGraphAddNode(graph, None, 0, params)
+    assertSuccess(err)
+
+    assert(len(params.conditional.phGraph_out) == 1)
+    assert(int(params.conditional.phGraph_out[0]) != 0)