Improve accuracy of get/setAllocationCounter

Summary:
get/setAllocationCounter didn't take into account allocations in the
current block. This was known at the time, but it turns out to be
important to have more accuracy when using these in a fine-grained
way.

Test Plan:
New unit test to test incrementally larger allocaitons.  Before I got
results like this:

```
+0
+0
+0
+0
+0
+4096
+0
+0
+0
+0
+0
+4064
+0
+0
+4088
+4056
+0
+0
+0
+4088
+4096
+4056
+4096
```

Notice how the results aren't always monotonically increasing.  After
this patch:

```
+344
+416
+488
+560
+632
+704
+776
+848
+920
+992
+1064
+1136
+1208
+1280
+1352
+1424
+1496
+1568
+1640
+1712
+1784
+1856
+1928
+2000
+2072
+2144
```

Reviewers: hvr, erikd, simonmar, jrtc27, trommler

Reviewed By: simonmar

Subscribers: trommler, jrtc27, rwbarton, thomie, carter

Differential Revision: https://phabricator.haskell.org/D4363

3 files changed, 23 insertions, 14 deletions

diff --git a/rts/PrimOps.cmm b/rts/PrimOps.cmm
index 67a238488c..e3f6e4cd19 100644
--- a/rts/PrimOps.cmm
+++ b/rts/PrimOps.cmm
@@ -2491,3 +2491,23 @@ stg_traceMarkerzh ( W_ msg )
     return ();
 }
 
+
+stg_getThreadAllocationCounterzh ()
+{
+    // Account for the allocation in the current block
+    W_ offset;
+    offset = Hp - bdescr_start(CurrentNursery);
+    return (StgTSO_alloc_limit(CurrentTSO) - TO_I64(offset));
+}
+
+stg_setThreadAllocationCounterzh ( I64 counter )
+{
+    // Allocation in the current block will be subtracted by
+    // getThreadAllocationCounter#, so we have to offset any existing
+    // allocation here.  See also openNursery/closeNursery in
+    // compiler/codeGen/StgCmmForeign.hs.
+    W_ offset;
+    offset = Hp - bdescr_start(CurrentNursery);
+    StgTSO_alloc_limit(CurrentTSO) = counter + TO_I64(offset);
+    return ();
+}
diff --git a/rts/RtsSymbols.c b/rts/RtsSymbols.c
index e53a056a4c..d5800fd336 100644
--- a/rts/RtsSymbols.c
+++ b/rts/RtsSymbols.c
@@ -743,8 +743,6 @@
       SymI_HasProto(rts_isProfiled)                                     \
       SymI_HasProto(rts_isDynamic)                                      \
       SymI_HasProto(rts_setInCallCapability)                            \
-      SymI_HasProto(rts_getThreadAllocationCounter)                     \
-      SymI_HasProto(rts_setThreadAllocationCounter)                     \
       SymI_HasProto(rts_enableThreadAllocationLimit)                    \
       SymI_HasProto(rts_disableThreadAllocationLimit)                   \
       SymI_HasProto(rts_setMainThread)                                  \
@@ -895,6 +893,8 @@
       SymI_HasProto(stg_traceCcszh)                                     \
       SymI_HasProto(stg_traceEventzh)                                   \
       SymI_HasProto(stg_traceMarkerzh)                                  \
+      SymI_HasProto(stg_getThreadAllocationCounterzh)                   \
+      SymI_HasProto(stg_setThreadAllocationCounterzh)                   \
       SymI_HasProto(getMonotonicNSec)                                   \
       SymI_HasProto(lockFile)                                           \
       SymI_HasProto(unlockFile)                                         \
diff --git a/rts/Threads.c b/rts/Threads.c
index b76917773a..be6962246d 100644
--- a/rts/Threads.c
+++ b/rts/Threads.c
@@ -165,19 +165,8 @@ rts_getThreadId(StgPtr tso)
 }
 
 /* ---------------------------------------------------------------------------
- * Getting & setting the thread allocation limit
+ * Enabling and disabling the thread allocation limit
  * ------------------------------------------------------------------------ */
-HsInt64 rts_getThreadAllocationCounter(StgPtr tso)
-{
-    // NB. doesn't take into account allocation in the current nursery
-    // block, so it might be off by up to 4k.
-    return PK_Int64((W_*)&(((StgTSO *)tso)->alloc_limit));
-}
-
-void rts_setThreadAllocationCounter(StgPtr tso, HsInt64 i)
-{
-    ASSIGN_Int64((W_*)&(((StgTSO *)tso)->alloc_limit), i);
-}
 
 void rts_enableThreadAllocationLimit(StgPtr tso)
 {