[FRONTEND] Fix Python error handling in launch (#2334)

This was regressed by #2185 because we didn't realise CUDA_CHECK macro
could do Python calls (similar to what led to #2225). I think the
PyErr_Occurred got removed in that PR because there was missing error
handling before the call to _launch, so it looked like it was just in
the wrong place.

It looks like there are also potentially a couple places in cuda.c that
can return with error set, e.g. getDeviceProperties, memAlloc,
memcpyHtoD, memFree, tensorMapEncodeTiled etc, but those are all
pre-existing and not affected by recent changes.
This commit is contained in:
Shantanu
2023-09-19 00:12:00 -07:00
committed by GitHub
parent 3a848e2729
commit 8e75e392ae

View File

@@ -255,6 +255,9 @@ static PyObject* launch(PyObject* self, PyObject* args) {{
Py_BEGIN_ALLOW_THREADS;
_launch(gridX, gridY, gridZ, num_warps, num_ctas, clusterDimX, clusterDimY, clusterDimZ, shared_memory, (CUstream)_stream, (CUfunction)_function{', ' + ', '.join(f"ptr_info{i}.dev_ptr" if ty[0]=="*" else f"_arg{i}"for i, ty in signature.items()) if len(signature) > 0 else ''});
Py_END_ALLOW_THREADS;
if (PyErr_Occurred()) {{
return NULL;
}}
if (launch_exit_hook != Py_None && !PyObject_CallObject(launch_exit_hook, args)) {{
return NULL;