Reimplement run_coroutine() in C++ to simplify tracebacks

Signed-off-by: Greg Bonik <gbonik@nvidia.com>

Author: gbonik

cext/CMakeLists.txt

@@ -45,6 +45,7 @@ set(cext_include_dirs
 # Build a static library first, so that we could reuse it for several build targets
 
 add_library(_cext_static STATIC
+    coroutine_util.cpp
     cuda_loader.cpp
     cuda_helper.cpp
     memory.cpp

cext/coroutine_util.cpp

@@ -0,0 +1,123 @@
+// SPDX-FileCopyrightText: Copyright (c) <2026> NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#include "py.h"
+#include "vec.h"
+
+#include <frameobject.h>  // For PyFrame_Check on Python 3.10
+
+
+static void best_effort_cleanup_on_internal_error(Vec<PyPtr>& stack) {
+    if (stack.empty())
+        return;
+    CHECK(PyErr_Occurred());
+    PyPtr frame = try_getattr(stack.back(), "cr_frame");
+    if (frame && PyFrame_Check(frame.get()))
+        PyTraceBack_Here(reinterpret_cast<PyFrameObject*>(frame.get()));
+
+    ErrorGuard guard;
+    while (!stack.empty()) {
+        PyPtr coro = stack.back();
+        stack.pop_back();
+        Py_XDECREF(PyObject_CallMethod(coro.get(), "close", ""));
+        if (PyErr_Occurred())
+            PyErr_Print();
+    }
+}
+
+
+// Run a coroutine using a software stack to bypass the Python's recursion limit.
+// Use resume_after() to break the call chain and push a new frame to the software stack.
+static PyObject* run_coroutine(PyObject* self, PyObject* main_coro) {
+    if (!PyCoro_CheckExact(main_coro)) {
+        raise(PyExc_TypeError, "Expected a coroutine");
+        return nullptr;
+    }
+
+    Vec<PyPtr> stack;
+    stack.push_back(newref(main_coro));
+
+    PyPtr ret = newref(Py_None);
+    SavedException exc;
+
+    while (!stack.empty()) {
+        PyObject* coro = stack.back().get();
+        if (!exc) {
+            // Happy path: use PyIter_Send() C API for efficiency
+            PyObject* res = nullptr;
+            PySendResult send_res = PyIter_Send(coro, ret.get(), &res);
+            if (send_res == PYGEN_RETURN) {
+                ret = steal(res);
+                stack.pop_back();
+            } else if (send_res == PYGEN_NEXT) {
+                PyPtr next_value = steal(res);
+                ret = newref(Py_None);
+                if (!PyCoro_CheckExact(next_value.get())) {
+                    raise(PyExc_TypeError, "Expected a continuation coroutine");
+                    best_effort_cleanup_on_internal_error(stack);
+                    return nullptr;
+                }
+                stack.push_back(std::move(next_value));
+            } else {
+                CHECK(send_res == PYGEN_ERROR);
+                exc = save_raised_exception();
+                exc.normalize();
+                ret = newref(Py_None);
+                stack.pop_back();
+            }
+            continue;
+        }
+
+        // Slow path: need to call .throw() since there is no public C API to _gen_throw()
+        PyPtr continuation = steal(PyObject_CallMethod(coro, "throw", "(O)", exc.value.get()));
+        if (continuation) {
+            ret = newref(Py_None);
+            exc = {};
+            if (!PyCoro_CheckExact(continuation.get())) {
+                raise(PyExc_TypeError, "Expected a continuation coroutine");
+                best_effort_cleanup_on_internal_error(stack);
+                return nullptr;
+            }
+            stack.push_back(std::move(continuation));
+        } else {
+            CHECK(PyErr_Occurred());
+            exc = save_raised_exception();
+            exc.normalize();
+            CHECK(exc.value);
+            if (PyErr_GivenExceptionMatches(exc.value.get(), PyExc_StopIteration)) {
+                ret = getattr(exc.value, "value");
+                if (!ret) {
+                    best_effort_cleanup_on_internal_error(stack);
+                    return nullptr;
+                }
+                exc = {};
+            } else {
+                ret = newref(Py_None);
+            }
+            stack.pop_back();
+        }
+    }
+
+    if (exc) {
+        exc.restore();
+        return nullptr;
+    }
+
+    return ret.release();
+}
+
+
+static PyMethodDef functions[] = {
+    {"run_coroutine", run_coroutine, METH_O, ""},
+    {}
+};
+
+
+Status coroutine_util_init(PyObject* m) {
+    if (PyModule_AddFunctions(m, functions) < 0)
+        return ErrorRaised;
+
+    return OK;
+}
+

cext/coroutine_util.h

@@ -0,0 +1,8 @@
+// SPDX-FileCopyrightText: Copyright (c) <2026> NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#pragma once
+
+Status coroutine_util_init(PyObject* m);
+

cext/module.cpp

@@ -6,6 +6,7 @@
 
 #include "tile_kernel.h"
 #include "cuda_helper.h"
+#include "coroutine_util.h"
 #include "xla_ffi_py.h"
 
 #ifdef _WIN32
@@ -48,6 +49,9 @@ PyMODINIT_FUNC PyInit__cext() {
     if (!cuda_helper_init(m.get()))
         return nullptr;
 
+    if (!coroutine_util_init(m.get()))
+        return nullptr;
+
     if (!xla_ffi_init(m.get()))
         return nullptr;
 

cext/py.h

@@ -202,6 +202,10 @@ ErrorRaised_t raise(PyObject* exctype, const char* fmt, Args&&... args) {
 struct SavedException {
     PyPtr type, value, traceback;
 
+    operator bool() const {
+        return bool(type);
+    }
+
     void normalize() {
         PyObject* tmp_type = type.release();
         PyObject* tmp_value = value.release();
@@ -210,6 +214,15 @@ struct SavedException {
         type = steal(tmp_type);
         value = steal(tmp_value);
         traceback = steal(tmp_traceback);
+        if (traceback)
+            PyException_SetTraceback(value.get(), traceback.get());
+    }
+
+    void restore() {
+        PyObject* tmp_type = type.release();
+        PyObject* tmp_value = value.release();
+        PyObject* tmp_traceback = traceback.release();
+        PyErr_Restore(tmp_type, tmp_value, tmp_traceback);
     }
 };
 
@@ -250,10 +263,7 @@ struct ErrorGuard {
     void operator=(const ErrorGuard&) = delete;
 
     ~ErrorGuard() {
-        PyObject* tmp_type = exc.type.release();
-        PyObject* tmp_value = exc.value.release();
-        PyObject* tmp_traceback = exc.traceback.release();
-        PyErr_Restore(tmp_type, tmp_value, tmp_traceback);
+        exc.restore();
     }
 };
 

cext/vec.h

@@ -158,6 +158,19 @@ class Vec {
         return data_[i];
     }
 
+    T& back() {
+        return data_[size_ - 1];
+    }
+
+    const T& back() const {
+        return data_[size_ - 1];
+    }
+
+    void pop_back() {
+        back().~T();
+        --size_;
+    }
+
     bool operator== (const Vec& other) const {
         size_t n = size_;
         if (n != other.size_) return false;

src/cuda/tile/_cext.pyi

@@ -105,6 +105,13 @@ def _benchmark(stream: int,
                /) -> float: ...
 
 
+def run_coroutine(coro):
+    """
+    Run a coroutine using a software stack to bypass the Python's recursion limit.
+    Use resume_after() to break the call chain and push a new frame to the software stack.
+    """
+
+
 CU_TENSOR_MAP_DATA_TYPE_UINT8: int
 CU_TENSOR_MAP_DATA_TYPE_UINT16: int
 CU_TENSOR_MAP_DATA_TYPE_UINT32: int

src/cuda/tile/_coroutine_util.py

@@ -2,43 +2,10 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-import sys
-from contextlib import ExitStack
 from dataclasses import dataclass
 from typing import Awaitable
 
-
-# Run a coroutine using a software stack to bypass the Python's recursion limit.
-# Use resume_after() to break the call chain and push a new frame to the software stack.
-def run_coroutine(awaitable: Awaitable):
-    ret = None
-    exc_info = None
-    stack = []
-    with ExitStack() as es:
-        try:
-            stack.append(awaitable.__await__())
-            while stack:
-                top = stack[-1]
-                try:
-                    continuation = top.send(ret) if exc_info is None else top.throw(exc_info[1])
-                except StopIteration as s:
-                    ret = s.value
-                    exc_info = None
-                    stack.pop()
-                except Exception:
-                    ret = None
-                    exc_info = sys.exc_info()
-                    stack.pop()
-                else:
-                    ret = exc_info = None
-                    stack.append(continuation.__await__())
-            if exc_info is None:
-                return ret
-            else:
-                raise exc_info[1]
-        finally:
-            for c in stack:
-                es.callback(c.close)
+from cuda.tile._cext import run_coroutine
 
 
 # Replace `await foo()` with `await resume_after(foo())` to bypass the recursion limit.
@@ -48,3 +15,6 @@ class resume_after:
 
     def __await__(self):
         return (yield self.awaitable)
+
+
+__all__ = ["run_coroutine", "resume_after"]

test/test_coroutine_util.py

@@ -52,6 +52,32 @@ def test_raise_then_catch():
     assert res == 100 + 2 + 3 + 4
 
 
+async def return_123():
+    return 123
+
+
+async def raise_then_catch_and_call_another(n):
+    if n == 0:
+        raise ValueError("Hello")
+
+    if n == 1:
+        try:
+            await resume_after(raise_then_catch_and_call_another(0))
+        except ValueError as e:
+            assert str(e) == "Hello"
+            x = await resume_after(return_123())
+            return x
+        assert False
+
+    r = await resume_after(raise_then_catch_and_call_another(n - 1))
+    return r + n
+
+
+def test_raise_then_catch_and_call_another():
+    res = run_coroutine(raise_then_catch_and_call_another(4))
+    assert res == 123 + 2 + 3 + 4
+
+
 async def two_calls():
     t1 = await resume_after(series(3))
     t2 = await resume_after(series(4))
@@ -75,5 +101,48 @@ def test_traceback_preserved():
     try:
         run_coroutine(call_leaf())
     except ValueError as e:
+        traceback.print_tb(e.__traceback__)
         frame_names = [f.name for f in traceback.extract_tb(e.__traceback__)]
         assert "raise_in_leaf" in frame_names
+    else:
+        assert False
+
+
+class WeirdAwaitable:
+    def __await__(self):
+        return iter([123])
+
+
+async def weird_await():
+    await WeirdAwaitable()
+
+
+async def call_weird_await():
+    await weird_await()
+
+
+def test_unexpected_awaitable():
+    try:
+        run_coroutine(call_weird_await())
+    except TypeError as e:
+        assert "Expected a continuation coroutine" in str(e)
+        traceback.print_tb(e.__traceback__)
+        frame_names = [f.name for f in traceback.extract_tb(e.__traceback__)]
+        assert "call_weird_await" in frame_names
+    else:
+        assert False
+
+
+async def resume_after_call_weird_await(flag):
+    try:
+        await resume_after(call_weird_await())
+    finally:
+        flag[0] = True
+
+
+def test_cleanup_after_internal_error():
+    flag = [False]
+    coro = resume_after_call_weird_await(flag)
+    with pytest.raises(TypeError, match="Expected a continuation coroutine"):
+        run_coroutine(coro)
+    assert flag[0] is True