NonMoving: Implement selector optimisation

5 files changed, 353 insertions, 3 deletions

diff --git a/rts/rts.cabal.in b/rts/rts.cabal.in
index a53c166849..1968f2f693 100644
--- a/rts/rts.cabal.in
+++ b/rts/rts.cabal.in
@@ -470,6 +470,7 @@ library
                sm/NonMovingCensus.c
                sm/NonMovingMark.c
                sm/NonMovingScav.c
+               sm/NonMovingShortcut.c
                sm/NonMovingSweep.c
                sm/Sanity.c
                sm/Scav.c
diff --git a/rts/sm/NonMoving.h b/rts/sm/NonMoving.h
index 17adfd6666..4458ce3bef 100644
--- a/rts/sm/NonMoving.h
+++ b/rts/sm/NonMoving.h
@@ -294,6 +294,11 @@ INLINE_HEADER bool nonmovingClosureBeingSwept(StgClosure *p)
     }
 }
 
+INLINE_HEADER bool isNonmovingClosure(StgClosure *p)
+{
+    return !HEAP_ALLOCED_GC(p) || Bdescr((P_)p)->flags & BF_NONMOVING;
+}
+
 #if defined(DEBUG)
 
 void nonmovingPrintSegment(struct NonmovingSegment *seg);
diff --git a/rts/sm/NonMovingMark.c b/rts/sm/NonMovingMark.c
index 751406c650..acf84d44f1 100644
--- a/rts/sm/NonMovingMark.c
+++ b/rts/sm/NonMovingMark.c
@@ -11,6 +11,7 @@
 // This is sometimes declared as a register variable therefore it is necessary
 // to include the declaration so that the compiler doesn't clobber the register.
 #include "NonMovingMark.h"
+#include "NonMovingShortcut.h"
 #include "NonMoving.h"
 #include "BlockAlloc.h"  /* for countBlocks */
 #include "HeapAlloc.h"
@@ -24,7 +25,7 @@
 #include "MarkWeak.h"
 #include "sm/Storage.h"
 
-static void mark_closure (MarkQueue *queue, StgClosure *p, StgClosure **origin);
+static void mark_closure (MarkQueue *queue, const StgClosure *p, StgClosure **origin);
 static void mark_tso (MarkQueue *queue, StgTSO *tso);
 static void mark_stack (MarkQueue *queue, StgStack *stack);
 static void mark_PAP_payload (MarkQueue *queue,
@@ -1445,8 +1446,7 @@ mark_closure (MarkQueue *queue, const StgClosure *p0, StgClosure **origin)
     }
 
     case THUNK_SELECTOR:
-        PUSH_FIELD((StgSelector *) p, selectee);
-        // TODO: selector optimization
+        nonmoving_eval_thunk_selector(queue, (StgSelector*)p, origin);
         break;
 
     case AP_STACK: {
@@ -1592,6 +1592,7 @@ done:
     if (origin != NULL && (!HEAP_ALLOCED(p) || bd->flags & BF_NONMOVING)) {
         if (UNTAG_CLOSURE((StgClosure*)p0) != p && *origin == p0) {
             if (cas((StgVolatilePtr)origin, (StgWord)p0, (StgWord)TAG_CLOSURE(tag, p)) == (StgWord)p0) {
+                // debugBelch("Thunk optimization successful\n");
             }
         }
     }
diff --git a/rts/sm/NonMovingShortcut.c b/rts/sm/NonMovingShortcut.c
new file mode 100644
index 0000000000..83c4857677
--- /dev/null
+++ b/rts/sm/NonMovingShortcut.c
@@ -0,0 +1,326 @@
+/* -----------------------------------------------------------------------------
+ *
+ * (c) The GHC Team, 1998-2019
+ *
+ * Non-moving garbage collector and allocator:
+ * Indirection short-cutting and the selector optimisation
+ *
+ * ---------------------------------------------------------------------------*/
+
+#include "Rts.h"
+#include "GC.h"
+#include "SMPClosureOps.h"
+#include "NonMovingMark.h"
+#include "NonMovingShortcut.h"
+#include "Printer.h"
+
+#define MAX_THUNK_SELECTOR_DEPTH 16
+
+//#define SELECTOR_OPT_DEBUG
+
+#if defined(SELECTOR_OPT_DEBUG)
+static void
+print_selector_chain(StgClosure *p)
+{
+    debugBelch("Selector chain: %p", (void*)p);
+    StgClosure *next = p->payload[0];
+    while (next != NULL) {
+        debugBelch(", %p", next);
+        next = next->payload[0];
+    }
+    debugBelch("\n");
+}
+#endif
+
+static void
+update_selector_chain(
+        StgClosure *chain,
+        StgClosure **origin,
+        StgSelector * const p0,
+        StgClosure * const val
+) {
+    ASSERT(val != NULL);
+
+    // Make sure we don't introduce non-moving-to-moving pointers here.
+    ASSERT(isNonmovingClosure(val));
+
+    // This case we can't handle because we don't know info ptr of the closure
+    // before we locked it.
+    ASSERT(chain != val);
+
+#if defined(SELECTOR_OPT_DEBUG)
+    if (chain != NULL) {
+        print_selector_chain(chain);
+        debugBelch("Value: ");
+        printClosure(val);
+    }
+#endif
+
+    while (chain) {
+        // debugBelch("chain: %p\n", (void*)chain);
+
+        StgClosure *next = chain->payload[0];
+
+        // We only update closures in the non-moving heap
+        ASSERT(isNonmovingClosure(chain));
+
+        ((StgInd*)chain)->indirectee = val;
+        unlockClosure((StgClosure*)chain, &stg_IND_info);
+
+        chain = next;
+    }
+
+    if (origin != NULL && (StgClosure*)p0 != val) {
+        cas((StgVolatilePtr)origin, (StgWord)p0, (StgWord)val);
+    }
+}
+
+// Returns value of the selector thunk. The value is a non-moving closure. If
+// it's not possible to evaluate the selector thunk the return value will be the
+// selector itself.
+static StgClosure*
+nonmoving_eval_thunk_selector_(
+        MarkQueue *queue,
+        StgSelector * const p0,
+        StgClosure ** const origin,
+        int depth
+) {
+    // This function should only be called on non-moving objects.
+    ASSERT(HEAP_ALLOCED_GC((P_)p0) && isNonmovingClosure((StgClosure*)p0));
+
+    markQueuePushClosure(queue, (StgClosure*)p0, NULL);
+
+    // INVARIANT: A non-moving object. Locked (below).
+    StgClosure *p = (StgClosure*)p0;
+
+    // Chain of non-moving selectors to update. These will be INDs to `p` when
+    // we reach to a value. INVARIANT: All objects in the chain are locked, and
+    // in the non-moving heap.
+    StgClosure *chain = NULL;
+
+    // Variables to update: p.
+selector_changed:
+    ;
+
+    // debugBelch("Selector changed: %p\n", (void*)p);
+
+    // Lock the selector to avoid concurrent modification in mutators
+    const StgInfoTable *selector_info_ptr = lockClosure((StgClosure*)p);
+    StgInfoTable *selector_info_tbl = INFO_PTR_TO_STRUCT(selector_info_ptr);
+
+    if (selector_info_tbl->type != THUNK_SELECTOR) {
+        // Selector updated in the meantime, or we reached to a value. Update
+        // the chain.
+        unlockClosure(p, selector_info_ptr);
+        update_selector_chain(chain, origin, p0, p);
+        return p;
+    }
+
+    // The closure is locked and it's a selector thunk. If the selectee is a
+    // CONSTR we do the selection here and the In the selected value will be the
+    // value of this selector thunk.
+    //
+    // Two cases:
+    //
+    // - If the selected value is also a selector thunk, then we loop and
+    //   evaluate it. The final value will be the value of both the current
+    //   selector and the selected value (which is also a selector thunk).
+    //
+    // - If the selectee is a selector thunk, we recursively evaluate it (up to
+    //   a certain depth, specified with MAX_THUNK_SELECTOR_DEPTH), then do the
+    //   selection on the value of it.
+
+    //
+    // Do the selection
+    //
+
+    uint32_t field = selector_info_tbl->layout.selector_offset;
+    StgClosure *selectee = UNTAG_CLOSURE(((StgSelector*)p)->selectee);
+
+    // Variables to update: selectee
+selectee_changed:
+    // debugBelch("Selectee changed: %p\n", (void*)selectee);
+
+    if (!isNonmovingClosure(selectee)) {
+        // The selectee is a moving object, and it may be moved by a concurrent
+        // minor GC while we read the info table and fields, so don't try to
+        // read the fields, just update the chain.
+        unlockClosure(p, selector_info_ptr);
+        update_selector_chain(chain, origin, p0, p);
+        return p;
+    }
+
+    // Selectee is a non-moving object, mark it.
+    markQueuePushClosure(queue, selectee, NULL);
+
+    const StgInfoTable *selectee_info_tbl = get_itbl(selectee);
+    switch (selectee_info_tbl->type) {
+        case WHITEHOLE: {
+            // Probably a loop. Abort.
+            unlockClosure(p, selector_info_ptr);
+            update_selector_chain(chain, origin, p0, p);
+            return p;
+        }
+
+        case CONSTR:
+        case CONSTR_1_0:
+        case CONSTR_0_1:
+        case CONSTR_2_0:
+        case CONSTR_1_1:
+        case CONSTR_0_2:
+        case CONSTR_NOCAF: {
+            // Selectee is a constructor in the non-moving heap.
+            // Select the field.
+
+            // Check that the size is in range.
+            ASSERT(field < (StgWord32)(selectee_info_tbl->layout.payload.ptrs +
+                                       selectee_info_tbl->layout.payload.nptrs));
+
+            StgClosure *val = UNTAG_CLOSURE(selectee->payload[field]);
+
+            // `val` is the value of this selector thunk. We need to check a
+            // few cases:
+            //
+            // - If `val` is in the moving heap, we stop here and update the
+            //   chain. All updated objects should be added to the mut_list.
+            //   (TODO (osa): What happens if the value is evacuated as we do
+            //   this?)
+            //
+            // - If `val` is in the non-moving heap, we check if it's also a
+            //   selector. If it is we add it to the chain and loop.
+
+            // Follow indirections. Variables to update: `val`.
+        val_changed:
+            if (!isNonmovingClosure(val)) {
+                // The selected value is a moving object, so we won't be
+                // updating the chain to this object.
+                unlockClosure(p, selector_info_ptr);
+                update_selector_chain(chain, origin, p0, p);
+                return p;
+            }
+
+            switch (get_itbl(val)->type) {
+            case IND:
+            case IND_STATIC:
+                ;
+                // Follow the indirection
+                StgClosure *indirectee = UNTAG_CLOSURE(((StgInd*)val)->indirectee);
+                if (isNonmovingClosure(indirectee)) {
+                    val = UNTAG_CLOSURE(((StgInd*)val)->indirectee);
+                    goto val_changed;
+                } else {
+                    unlockClosure(p, selector_info_ptr);
+                    update_selector_chain(chain, origin, p0, p);
+                    return p;
+                }
+            case THUNK_SELECTOR:
+                // Value of the selector thunk is again a selector thunk in the
+                // non-moving heap. Add the current selector to the chain and
+                // loop.
+                p->payload[0] = chain;
+                chain = p;
+                p = val;
+                goto selector_changed;
+            default:
+                // Found a value, add the current selector to the chain and
+                // update it.
+                p->payload[0] = chain;
+                chain = p;
+                update_selector_chain(chain, origin, p0, val);
+                return val;
+            }
+        }
+
+        case IND:
+        case IND_STATIC: {
+            StgClosure *indirectee = UNTAG_CLOSURE(((StgInd *)selectee)->indirectee);
+            if (isNonmovingClosure(indirectee)) {
+                selectee = indirectee;
+                goto selectee_changed;
+            } else {
+                unlockClosure(p, selector_info_ptr);
+                update_selector_chain(chain, origin, p0, p);
+                return p;
+            }
+        }
+
+        case BLACKHOLE: {
+            StgClosure *indirectee = ((StgInd*)selectee)->indirectee;
+
+            if (!isNonmovingClosure(UNTAG_CLOSURE(indirectee))) {
+                unlockClosure(p, selector_info_ptr);
+                update_selector_chain(chain, origin, p0, p);
+                return p;
+            }
+
+            // Establish whether this BH has been updated, and is now an
+            // indirection, as in evacuate().
+            if (GET_CLOSURE_TAG(indirectee) == 0) {
+                const StgInfoTable *i = indirectee->header.info;
+                if (i == &stg_TSO_info
+                    || i == &stg_WHITEHOLE_info
+                    || i == &stg_BLOCKING_QUEUE_CLEAN_info
+                    || i == &stg_BLOCKING_QUEUE_DIRTY_info) {
+                    unlockClosure(p, selector_info_ptr);
+                    update_selector_chain(chain, origin, p0, p);
+                    return (StgClosure*)p;
+                }
+                ASSERT(i != &stg_IND_info); // TODO not sure about this part
+            }
+
+            // It's an indirection, follow it.
+            selectee = UNTAG_CLOSURE(indirectee);
+            goto selectee_changed;
+        }
+
+        case AP:
+        case AP_STACK:
+        case THUNK:
+        case THUNK_1_0:
+        case THUNK_0_1:
+        case THUNK_2_0:
+        case THUNK_1_1:
+        case THUNK_0_2:
+        case THUNK_STATIC: {
+            // Not evaluated yet
+            unlockClosure(p, selector_info_ptr);
+            update_selector_chain(chain, origin, p0, p);
+            return (StgClosure*)p;
+        }
+
+        case THUNK_SELECTOR: {
+            // Selectee is a selector thunk. Evaluate it if we haven't reached
+            // the recursion limit yet.
+            if (depth < MAX_THUNK_SELECTOR_DEPTH) {
+                StgClosure *new_selectee =
+                    UNTAG_CLOSURE(nonmoving_eval_thunk_selector_(
+                                queue, (StgSelector*)selectee, NULL, depth+1));
+                ASSERT(isNonmovingClosure(new_selectee));
+                if (selectee == new_selectee) {
+                    unlockClosure(p, selector_info_ptr);
+                    update_selector_chain(chain, origin, p0, p);
+                    return (StgClosure*)p;
+                } else {
+                    selectee = new_selectee;
+                    goto selectee_changed;
+                }
+            } else {
+                // Recursion limit reached
+                unlockClosure(p, selector_info_ptr);
+                update_selector_chain(chain, origin, p0, p);
+                return (StgClosure*)p;
+            }
+        }
+
+        default: {
+            barf("nonmoving_eval_thunk_selector: strange selectee %d",
+                 (int)(selectee_info_tbl->type));
+        }
+    }
+}
+
+void
+nonmoving_eval_thunk_selector(MarkQueue *queue, StgSelector *p, StgClosure **origin)
+{
+    nonmoving_eval_thunk_selector_(queue, p, origin, 0);
+}
diff --git a/rts/sm/NonMovingShortcut.h b/rts/sm/NonMovingShortcut.h
new file mode 100644
index 0000000000..72297884aa
--- /dev/null
+++ b/rts/sm/NonMovingShortcut.h
@@ -0,0 +1,17 @@
+/* -----------------------------------------------------------------------------
+ *
+ * (c) The GHC Team, 1998-2019
+ *
+ * Non-moving garbage collector and allocator:
+ * Indirection short-cutting and the selector optimisation
+ *
+ * ---------------------------------------------------------------------------*/
+
+#pragma once
+
+#include "BeginPrivate.h"
+
+void
+nonmoving_eval_thunk_selector(MarkQueue *queue, StgSelector *p, StgClosure **origin);
+
+#include "EndPrivate.h"