Merge pull request #1306 from undingen/generator_opt

generator optimizations

src/core/threading.cpp

@@ -83,24 +83,6 @@ private:
 public:
     void* stack_start;
 
-    struct StackInfo {
-        BoxedGenerator* next_generator;
-        void* stack_start;
-        void* stack_limit;
-
-        StackInfo(BoxedGenerator* next_generator, void* stack_start, void* stack_limit)
-            : next_generator(next_generator), stack_start(stack_start), stack_limit(stack_limit) {
-#if STACK_GROWS_DOWN
-            assert(stack_start > stack_limit);
-            assert((char*)stack_start - (char*)stack_limit < (1L << 30));
-#else
-            assert(stack_start < stack_limit);
-            assert((char*)stack_limit - (char*)stack_start < (1L << 30));
-#endif
-        }
-    };
-
-    std::vector<StackInfo> previous_stacks;
     pthread_t pthread_id;
 
     PyThreadState* public_thread_state;
@@ -137,36 +119,10 @@ public:
     }
 
     ucontext_t* getContext() { return &ucontext; }
-
-    void pushGenerator(BoxedGenerator* g, void* new_stack_start, void* old_stack_limit) {
-        previous_stacks.emplace_back(g, this->stack_start, old_stack_limit);
-        this->stack_start = new_stack_start;
-    }
-
-    void popGenerator() {
-        assert(previous_stacks.size());
-        StackInfo& stack = previous_stacks.back();
-        stack_start = stack.stack_start;
-        previous_stacks.pop_back();
-    }
-
-    void assertNoGenerators() { assert(previous_stacks.size() == 0); }
 };
 static std::unordered_map<pthread_t, ThreadStateInternal*> current_threads;
 static __thread ThreadStateInternal* current_internal_thread_state = 0;
 
-void pushGenerator(BoxedGenerator* g, void* new_stack_start, void* old_stack_limit) {
-    assert(new_stack_start);
-    assert(old_stack_limit);
-    assert(current_internal_thread_state);
-    current_internal_thread_state->pushGenerator(g, new_stack_start, old_stack_limit);
-}
-
-void popGenerator() {
-    assert(current_internal_thread_state);
-    current_internal_thread_state->popGenerator();
-}
-
 // These are guarded by threading_lock
 static int signals_waiting(0);
 
@@ -243,8 +199,6 @@ static void tstate_delete_common(PyThreadState* tstate) {
 }
 
 static void unregisterThread() {
-    current_internal_thread_state->assertNoGenerators();
-
     tstate_delete_common(current_internal_thread_state->public_thread_state);
     PyThreadState_Clear(current_internal_thread_state->public_thread_state);
 
@@ -438,7 +392,6 @@ static void wait_for_thread_shutdown(void) noexcept {
 
 void finishMainThread() {
     assert(current_internal_thread_state);
-    current_internal_thread_state->assertNoGenerators();
 
     wait_for_thread_shutdown();
 }

src/core/threading.h

@@ -47,13 +47,6 @@ void finishMainThread();
 
 bool isMainThread();
 
-// Some hooks to keep track of the list of stacks that this thread has been using.
-// Every time we switch to a new generator, we need to pass a reference to the generator
-// itself (so we can access the registers it is saving), the location of the new stack, and
-// where we stopped executing on the old stack.
-void pushGenerator(BoxedGenerator* g, void* new_stack_start, void* old_stack_limit);
-void popGenerator();
-
 #ifndef THREADING_USE_GIL
 #define THREADING_USE_GIL 1
 #define THREADING_USE_GRWL 0

src/runtime/generator.cpp

@@ -84,7 +84,7 @@ Context* getReturnContextForGeneratorFrame(void* frame_addr) {
     return generator->returnContext;
 }
 
-void generatorEntry(BoxedGenerator* g) {
+void generatorEntry(BoxedGenerator* g) noexcept {
     {
         assert(g->cls == generator_cls);
         assert(g->function->cls == function_cls);
@@ -92,8 +92,7 @@ void generatorEntry(BoxedGenerator* g) {
         assert(g->returnValue == Py_None);
         Py_CLEAR(g->returnValue);
 
-        threading::pushGenerator(g, g->stack_begin, g->returnContext);
-        try {
+        {
             RegisterHelper context_registerer(g, __builtin_frame_address(0));
 
             g->top_caller_frame_info = (FrameInfo*)cur_thread_state.frame_info;
@@ -104,18 +103,20 @@ void generatorEntry(BoxedGenerator* g) {
             // KEEP_ALIVE(func);
 
             Box** args = g->args ? &g->args->elts[0] : nullptr;
-            auto r = callCLFunc<ExceptionStyle::CXX, NOT_REWRITABLE>(func->md, nullptr, func->md->numReceivedArgs(),
-                                                                     func->closure, g, func->globals, g->arg1, g->arg2,
-                                                                     g->arg3, args);
-            Py_DECREF(r);
-        } catch (ExcInfo e) {
-            // unhandled exception: propagate the exception to the caller
-            g->exception = e;
+            auto r = callCLFunc<ExceptionStyle::CAPI, NOT_REWRITABLE>(func->md, nullptr, func->md->numReceivedArgs(),
+                                                                      func->closure, g, func->globals, g->arg1, g->arg2,
+                                                                      g->arg3, args);
+            if (r)
+                Py_DECREF(r);
+            else {
+                // unhandled exception: propagate the exception to the caller
+                PyErr_Fetch(&g->exception.type, &g->exception.value, &g->exception.traceback);
+                PyErr_Clear();
+            }
         }
 
         // we returned from the body of the generator. next/send/throw will notify the caller
         g->entryExited = true;
-        threading::popGenerator();
     }
     assert(g->top_caller_frame_info == cur_thread_state.frame_info);
     swapContext(&g->context, g->returnContext, 0);
@@ -256,7 +257,8 @@ template <ExceptionStyle S> static Box* generatorSend(Box* s, Box* v) noexcept(S
     return rtn;
 }
 
-Box* generatorThrow(Box* s, BoxedClass* exc_cls, Box* exc_val = nullptr, Box** args = nullptr) {
+template <ExceptionStyle S>
+Box* generatorThrow(Box* s, BoxedClass* exc_cls, Box* exc_val = nullptr, Box** args = nullptr) noexcept(S == CAPI) {
     assert(s->cls == generator_cls);
     BoxedGenerator* self = static_cast<BoxedGenerator*>(s);
 
@@ -264,22 +266,32 @@ Box* generatorThrow(Box* s, BoxedClass* exc_cls, Box* exc_val = nullptr, Box** a
         Py_FatalError(".throw called on generator last advanced with __hasnext__");
 
     Box* exc_tb = args ? args[0] : nullptr;
-    if (exc_tb && exc_tb != Py_None && !PyTraceBack_Check(exc_tb))
+    if (exc_tb && exc_tb != Py_None && !PyTraceBack_Check(exc_tb)) {
+        if (S == CAPI) {
+            PyErr_SetString(TypeError, "throw() third argument must be a traceback object");
+            return NULL;
+        }
         raiseExcHelper(TypeError, "throw() third argument must be a traceback object");
+    }
     if (!exc_val)
         exc_val = Py_None;
     if (!exc_tb)
         exc_tb = Py_None;
 
     ExcInfo exc_info = excInfoForRaise(incref(exc_cls), incref(exc_val), incref(exc_tb));
-    if (self->entryExited)
+    if (self->entryExited) {
+        if (S == CAPI) {
+            setCAPIException(exc_info);
+            return NULL;
+        }
         throw exc_info;
+    }
 
     self->exception = exc_info;
-    return generatorSend<CXX>(self, Py_None);
+    return generatorSend<S>(self, Py_None);
 }
 
-Box* generatorClose(Box* s) {
+template <ExceptionStyle S> Box* generatorClose(Box* s) noexcept(S == CAPI) {
     assert(s->cls == generator_cls);
     BoxedGenerator* self = static_cast<BoxedGenerator*>(s);
 
@@ -287,15 +299,28 @@ Box* generatorClose(Box* s) {
     if (self->entryExited)
         return incref(Py_None);
 
-    try {
-        autoDecref(generatorThrow(self, GeneratorExit, nullptr, nullptr));
-        raiseExcHelper(RuntimeError, "generator ignored GeneratorExit");
-    } catch (ExcInfo e) {
-        if (e.matches(StopIteration) || e.matches(GeneratorExit)) {
-            e.clear();
+    if (S == CAPI) {
+        Box* rtn = generatorThrow<S>(self, GeneratorExit, nullptr, nullptr);
+        if (rtn) {
+            PyErr_SetString(RuntimeError, "generator ignored GeneratorExit");
+            return NULL;
+        }
+        if (PyErr_ExceptionMatches(PyExc_StopIteration) || PyErr_ExceptionMatches(PyExc_GeneratorExit)) {
+            PyErr_Clear();
             return incref(Py_None);
         }
-        throw e;
+        return NULL;
+    } else {
+        try {
+            autoDecref(generatorThrow<S>(self, GeneratorExit, nullptr, nullptr));
+            raiseExcHelper(RuntimeError, "generator ignored GeneratorExit");
+        } catch (ExcInfo e) {
+            if (e.matches(StopIteration) || e.matches(GeneratorExit)) {
+                e.clear();
+                return incref(Py_None);
+            }
+            throw e;
+        }
     }
     assert(0); // unreachable
 }
@@ -331,25 +356,9 @@ Box* generatorHasnext(Box* s) {
     return boxBool(generatorHasnextUnboxed(s));
 }
 
-extern "C" Box* yield_capi(BoxedGenerator* obj, STOLEN(Box*) value, int num_live_values, ...) noexcept {
-    try {
-        llvm::SmallVector<Box*, 8> live_values;
-        live_values.reserve(num_live_values);
-        va_list ap;
-        va_start(ap, num_live_values);
-        for (int i = 0; i < num_live_values; ++i) {
-            live_values.push_back(va_arg(ap, Box*));
-        }
-        va_end(ap);
-
-        return yield(obj, value, live_values);
-    } catch (ExcInfo e) {
-        setCAPIException(e);
-        return NULL;
-    }
-}
-
-extern "C" Box* yield(BoxedGenerator* obj, STOLEN(Box*) value, llvm::ArrayRef<Box*> live_values) {
+template <ExceptionStyle S>
+static Box* yieldInternal(BoxedGenerator* obj, STOLEN(Box*) value,
+                          llvm::ArrayRef<Box*> live_values) noexcept(S == CAPI) {
     STAT_TIMER(t0, "us_timer_generator_switching", 0);
 
     assert(obj->cls == generator_cls);
@@ -357,8 +366,6 @@ extern "C" Box* yield(BoxedGenerator* obj, STOLEN(Box*) value, llvm::ArrayRef<Bo
     assert(!self->returnValue);
     self->returnValue = value;
 
-    threading::popGenerator();
-
     FrameInfo* generator_frame_info = (FrameInfo*)cur_thread_state.frame_info;
     // a generator will only switch back (yield/unhandled exception) to its caller when it is one frame away from the
     // caller
@@ -384,13 +391,15 @@ extern "C" Box* yield(BoxedGenerator* obj, STOLEN(Box*) value, llvm::ArrayRef<Bo
     }
     cur_thread_state.frame_info = generator_frame_info;
 
-    threading::pushGenerator(obj, obj->stack_begin, obj->returnContext);
-
     // if the generator receives a exception from the caller we have to throw it
     if (self->exception.type) {
         ExcInfo e = self->exception;
         self->exception = ExcInfo(NULL, NULL, NULL);
         Py_CLEAR(self->returnValue);
+        if (S == CAPI) {
+            setCAPIException(e);
+            return NULL;
+        }
         throw e;
     }
 
@@ -399,6 +408,22 @@ extern "C" Box* yield(BoxedGenerator* obj, STOLEN(Box*) value, llvm::ArrayRef<Bo
     return r;
 }
 
+extern "C" Box* yield_capi(BoxedGenerator* obj, STOLEN(Box*) value, int num_live_values, ...) noexcept {
+    Box** live_values = (Box**)alloca(sizeof(Box*) * num_live_values);
+    va_list ap;
+    va_start(ap, num_live_values);
+    for (int i = 0; i < num_live_values; ++i) {
+        live_values[i] = va_arg(ap, Box*);
+    }
+    va_end(ap);
+
+    return yieldInternal<CAPI>(obj, value, llvm::makeArrayRef(live_values, num_live_values));
+}
+
+extern "C" Box* yield(BoxedGenerator* obj, STOLEN(Box*) value, llvm::ArrayRef<Box*> live_values) {
+    return yieldInternal<CXX>(obj, value, live_values);
+}
+
 extern "C" BoxedGenerator* createGenerator(BoxedFunctionBase* function, Box* arg1, Box* arg2, Box* arg3, Box** args) {
     assert(function);
     assert(function->cls == function_cls);
@@ -538,15 +563,6 @@ extern "C" int PyGen_NeedsFinalizing(PyGenObject* gen) noexcept {
 #endif
 }
 
-static PyObject* generator_close(PyGenObject* gen, PyObject* args) noexcept {
-    try {
-        return generatorClose((Box*)gen);
-    } catch (ExcInfo e) {
-        setCAPIException(e);
-        return NULL;
-    }
-}
-
 static void generator_del(PyObject* self) noexcept {
     PyObject* res;
     PyObject* error_type, *error_value, *error_traceback;
@@ -567,7 +583,7 @@ static void generator_del(PyObject* self) noexcept {
 
     // Pyston change:
     // res = gen_close(gen, NULL);
-    res = generator_close((PyGenObject*)gen, NULL);
+    res = generatorClose<CAPI>((Box*)gen);
 
     if (res == NULL)
         PyErr_WriteUnraisable(self);
@@ -700,7 +716,9 @@ void setupGenerator() {
     generator_cls->giveAttr(
         "__iter__", new BoxedFunction(FunctionMetadata::create((void*)generatorIter, typeFromClass(generator_cls), 1)));
 
-    generator_cls->giveAttr("close", new BoxedFunction(FunctionMetadata::create((void*)generatorClose, UNKNOWN, 1)));
+    auto generator_close = FunctionMetadata::create((void*)generatorClose<CXX>, UNKNOWN, 1);
+    generator_close->addVersion((void*)generatorClose<CAPI>, UNKNOWN, CAPI);
+    generator_cls->giveAttr("close", new BoxedFunction(generator_close));
 
     auto generator_next = FunctionMetadata::create((void*)generatorNext<CXX>, UNKNOWN, 1, ParamNames::empty(), CXX);
     generator_next->addVersion((void*)generatorNext<CAPI>, UNKNOWN, CAPI);
@@ -710,15 +728,20 @@ void setupGenerator() {
     hasnext->addVersion((void*)generatorHasnext, BOXED_BOOL);
     generator_cls->giveAttr("__hasnext__", new BoxedFunction(hasnext));
 
-    generator_cls->giveAttr("send", new BoxedFunction(FunctionMetadata::create((void*)generatorSend<CXX>, UNKNOWN, 2)));
-    auto gthrow
-        = new BoxedFunction(FunctionMetadata::create((void*)generatorThrow, UNKNOWN, 4, false, false), { NULL, NULL });
-    generator_cls->giveAttr("throw", gthrow);
+    auto generator_send = FunctionMetadata::create((void*)generatorSend<CXX>, UNKNOWN, 2);
+    generator_send->addVersion((void*)generatorSend<CAPI>, UNKNOWN, CAPI);
+    generator_cls->giveAttr("send", new BoxedFunction(generator_send));
+
+    auto generator_throw = FunctionMetadata::create((void*)generatorThrow<CXX>, UNKNOWN, 4, false, false);
+    generator_throw->addVersion((void*)generatorThrow<CAPI>, UNKNOWN, CAPI);
+    generator_cls->giveAttr("throw", new BoxedFunction(generator_throw, { NULL, NULL }));
 
     generator_cls->giveAttrDescriptor("__name__", generator_name, NULL);
 
     generator_cls->freeze();
     generator_cls->tp_iter = PyObject_SelfIter;
     generator_cls->tp_del = generator_del; // don't do giveAttr("__del__") because it should not be visible from python
+    generator_cls->tpp_hasnext = generatorHasnextUnboxed;
+    generator_cls->tp_iternext = generatorNext<CAPI>;
 }
 }

src/runtime/generator.h

@@ -26,7 +26,7 @@ struct Context;
 extern BoxedClass* generator_cls;
 
 void setupGenerator();
-void generatorEntry(BoxedGenerator* g);
+void generatorEntry(BoxedGenerator* g) noexcept;
 Context* getReturnContextForGeneratorFrame(void* frame_addr);
 
 extern "C" Box* yield(BoxedGenerator* obj, STOLEN(Box*) value, llvm::ArrayRef<Box*> live_values = {});