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: niteria, bgamari, hvr, erikd

Subscribers: rwbarton, thomie, carter

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

1 files changed, 20 insertions, 0 deletions

diff --git a/rts/PrimOps.cmm b/rts/PrimOps.cmm
index 2b3a304d06..1caa0c3343 100644
--- a/rts/PrimOps.cmm
+++ b/rts/PrimOps.cmm
@@ -2495,3 +2495,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) - 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 + offset;
+    return ();
+}