* Add cuda.tile.tune.exhaustive_search API

- Add `exhaustive_search()` and `TunedResult` to public API.
- Also add `compiler_timeout()` and `kernel.replace_hints()` to public API.
- Mark ct_experimental.autotune as deprecated.
- Add a c++ benchmark util using cudagraph

Co-authored-by: Greg Bonik <gbonik@nvidia.com>
Signed-off-by: Jay Gu <jagu@nvidia.com>

Author: haijieg

cext/cuda_helper.cpp

@@ -96,13 +96,66 @@ PyObject* get_driver_version(PyObject *self, PyObject *Py_UNUSED(ignored)) {
     return Py_BuildValue("(ii)", major, minor);
 }
 
+// ========== Context helpers ==========
+
+PyObject* synchronize_context(PyObject* self, PyObject* Py_UNUSED(ignored)) {
+    Result<const DriverApi*> driver_result = get_driver_api();
+    if (!driver_result.is_ok()) return NULL;
+    const DriverApi* d = *driver_result;
+
+    CUresult res = d->cuCtxSynchronize();
+    if (res != CUDA_SUCCESS) {
+        return PyErr_Format(PyExc_RuntimeError,
+                            "cuCtxSynchronize: %s", get_cuda_error(d, res));
+    }
+    Py_RETURN_NONE;
+}
+
+// ========== Stream helpers ==========
+
+PyObject* create_stream(PyObject* self, PyObject* Py_UNUSED(ignored)) {
+    Result<const DriverApi*> driver_result = get_driver_api();
+    if (!driver_result.is_ok()) return NULL;
+    const DriverApi* d = *driver_result;
+
+    CUstream stream;
+    CUresult res = d->cuStreamCreate(&stream, CU_STREAM_NON_BLOCKING);
+    if (res != CUDA_SUCCESS) {
+        return PyErr_Format(PyExc_RuntimeError,
+                            "cuStreamCreate: %s", get_cuda_error(d, res));
+    }
+    return PyLong_FromVoidPtr(stream);
+}
+
+PyObject* destroy_stream(PyObject* self, PyObject* arg) {
+    CUstream stream = static_cast<CUstream>(PyLong_AsVoidPtr(arg));
+    if (PyErr_Occurred()) return NULL;
+
+    Result<const DriverApi*> driver_result = get_driver_api();
+    if (!driver_result.is_ok()) return NULL;
+    const DriverApi* d = *driver_result;
+
+    CUresult res = d->cuStreamDestroy(stream);
+    if (res != CUDA_SUCCESS) {
+        return PyErr_Format(PyExc_RuntimeError,
+                            "cuStreamDestroy: %s", get_cuda_error(d, res));
+    }
+    Py_RETURN_NONE;
+}
+
 static PyMethodDef functions[] = {
     {"get_compute_capability", get_compute_capability, METH_NOARGS,
         "Get compute capability of the default CUDA device"},
     {"get_driver_version", get_driver_version, METH_NOARGS,
         "Get the cuda driver version"},
     {"_get_max_grid_size", get_max_grid_size, METH_VARARGS,
         "Get max grid size of a CUDA device, given device id"},
+    {"_synchronize_context", synchronize_context, METH_NOARGS,
+        "Synchronize the current CUDA context (drain all streams)."},
+    {"_create_stream", create_stream, METH_NOARGS,
+        "Create a non-blocking CUDA stream. Returns int handle."},
+    {"_destroy_stream", destroy_stream, METH_O,
+        "Destroy a CUDA stream given its int handle."},
     NULL
 };
 

cext/cuda_loader.h

@@ -17,6 +17,7 @@
     X(cuGetErrorString, 6000) \
     X(cuLaunchKernel, 7000) \
     X(cuPointerGetAttribute, 4000) \
+    X(cuCtxSynchronize, 2000) \
     X(cuCtxPushCurrent, 4000) \
     X(cuCtxPopCurrent, 4000) \
     X(cuCtxGetCurrent, 4000) \
@@ -37,11 +38,23 @@
     X(cuMemFreeHost, 2000) \
     X(cuMemcpyHtoDAsync, 3020) \
     X(cuStreamCreate, 2000) \
+    X(cuStreamDestroy, 4000) \
     X(cuStreamGetCtx, 9020) \
     X(cuStreamGetId, 12000) \
     X(cuStreamIsCapturing, 10000) \
     X(cuStreamSynchronize, 7000) \
-    X(cuStreamWaitEvent, 7000)
+    X(cuStreamWaitEvent, 7000) \
+    X(cuEventElapsedTime, 12080) \
+    X(cuGraphCreate, 10000) \
+    X(cuGraphDestroy, 10000) \
+    X(cuGraphAddEventRecordNode, 11010) \
+    X(cuGraphAddKernelNode, 12000) \
+    X(cuGraphAddMemsetNode, 10000) \
+    X(cuGraphAddMemAllocNode, 11040) \
+    X(cuGraphAddMemFreeNode, 11040) \
+    X(cuGraphInstantiateWithFlags, 11040) \
+    X(cuGraphExecDestroy, 10000) \
+    X(cuGraphLaunch, 10000)
 
 
 #define DECLARE_CUDA_FUNC_EXTERN(name, _cuda_version) \
@@ -52,3 +65,73 @@ struct DriverApi {
 };
 
 Result<const DriverApi*> get_driver_api();
+
+
+class CudaGraph {
+    const DriverApi* d;
+    CUgraph graph;
+public:
+    CudaGraph(const CudaGraph&) = delete;
+    void operator=(const CudaGraph&) = delete;
+
+    explicit CudaGraph(const DriverApi* d) : d(d), graph(nullptr) {}
+
+    CUresult create() {
+        CHECK(!graph);
+        return d->cuGraphCreate(&graph, 0);
+    }
+
+    CUgraph get() const {
+        return graph;
+    }
+
+    ~CudaGraph() {
+        if (graph) d->cuGraphDestroy(graph);
+    }
+};
+
+class CudaGraphExec {
+    const DriverApi* d;
+    CUgraphExec exec;
+public:
+    CudaGraphExec(const CudaGraphExec&) = delete;
+    void operator=(const CudaGraphExec&) = delete;
+
+    explicit CudaGraphExec(const DriverApi* d) : d(d), exec(nullptr) {}
+
+    CUresult instantiate(const CudaGraph& graph) {
+        CHECK(!exec);
+        return d->cuGraphInstantiateWithFlags(&exec, graph.get(), 0);
+    }
+
+    CUgraphExec get() const {
+        return exec;
+    }
+
+    ~CudaGraphExec() {
+        if (exec) d->cuGraphExecDestroy(exec);
+    }
+};
+
+class CudaEvent {
+    const DriverApi* d;
+    CUevent event;
+public:
+    CudaEvent(const CudaEvent&) = delete;
+    void operator=(const CudaEvent&) = delete;
+
+    explicit CudaEvent(const DriverApi* d) : d(d), event(nullptr) {}
+
+    CUresult create() {
+        CHECK(!event);
+        return d->cuEventCreate(&event, CU_EVENT_DEFAULT);
+    }
+
+    CUevent get() const {
+        return event;
+    }
+
+    ~CudaEvent() {
+        if (event) d->cuEventDestroy(event);
+    }
+};

cext/stream_buffer.cpp

@@ -51,6 +51,10 @@ StreamBufferPool* stream_buffer_pool_new() {
     return new StreamBufferPool();
 }
 
+void stream_buffer_pool_delete(StreamBufferPool* pool) {
+    delete pool;
+}
+
 static void delete_chunk(const DriverApi* driver, Chunk* chunk) {
     dual_ptr_free(driver, chunk->ptr);
     driver->cuEventDestroy(chunk->event);

cext/stream_buffer.h

@@ -74,6 +74,7 @@ class StreamBufferTransaction {
 };
 
 StreamBufferPool* stream_buffer_pool_new();
+void stream_buffer_pool_delete(StreamBufferPool* pool);
 
 StreamBufferTransaction stream_buffer_transaction_open(const DriverApi*,
                                                        StreamBufferPool* pool,