feat(stdlib): Faster memory allocator (#2124)

grain-lang · Jun 25, 2024 · 03e10c4 · 03e10c4
1 parent 416936b
commit 03e10c4
Show file tree

Hide file tree

Showing 8 changed files with 341 additions and 240 deletions.
diff --git a/compiler/test/input/mallocTight.gr b/compiler/test/input/mallocTight.gr
diff --git a/compiler/test/runner.re b/compiler/test/runner.re
@@ -39,7 +39,7 @@ let graindoc_out_file = name =>
 let gaindoc_in_file = name =>
   Filepath.to_string(Fp.At.(test_gaindoc_dir / (name ++ ".input.gr")));
 
-let compile = (~num_pages=?, ~config_fn=?, ~hook=?, name, prog) => {
+let compile = (~num_pages=?, ~max_pages=?, ~config_fn=?, ~hook=?, name, prog) => {
   Config.preserve_all_configs(() => {
     Config.with_config(
       Config.empty,
@@ -49,11 +49,10 @@ let compile = (~num_pages=?, ~config_fn=?, ~hook=?, name, prog) => {
         | None => ()
         };
         switch (num_pages) {
-        | Some(pages) =>
-          Config.initial_memory_pages := pages;
-          Config.maximum_memory_pages := Some(pages);
+        | Some(pages) => Config.initial_memory_pages := pages
         | None => ()
         };
+        Config.maximum_memory_pages := max_pages;
         Config.include_dirs :=
           [Filepath.to_string(test_libs_dir), ...Config.include_dirs^];
         let outfile = wasmfile(name);
@@ -63,7 +62,8 @@ let compile = (~num_pages=?, ~config_fn=?, ~hook=?, name, prog) => {
   });
 };
 
-let compile_file = (~num_pages=?, ~config_fn=?, ~hook=?, filename, outfile) => {
+let compile_file =
+    (~num_pages=?, ~max_pages=?, ~config_fn=?, ~hook=?, filename, outfile) => {
   Config.preserve_all_configs(() => {
     Config.with_config(
       Config.empty,
@@ -73,11 +73,10 @@ let compile_file = (~num_pages=?, ~config_fn=?, ~hook=?, filename, outfile) => {
         | None => ()
         };
         switch (num_pages) {
-        | Some(pages) =>
-          Config.initial_memory_pages := pages;
-          Config.maximum_memory_pages := Some(pages);
+        | Some(pages) => Config.initial_memory_pages := pages
         | None => ()
         };
+        Config.maximum_memory_pages := max_pages;
         Config.include_dirs :=
           [Filepath.to_string(test_libs_dir), ...Config.include_dirs^];
         compile_file(~is_root_file=true, ~hook?, ~outfile, filename);
@@ -153,18 +152,7 @@ let open_process = args => {
   (code, out, err);
 };
 
-let run = (~num_pages=?, ~extra_args=[||], file) => {
-  let mem_flags =
-    switch (num_pages) {
-    | Some(x) => [|
-        "--initial-memory-pages",
-        string_of_int(x),
-        "--maximum-memory-pages",
-        string_of_int(x),
-      |]
-    | None => [||]
-    };
-
+let run = (~extra_args=[||], file) => {
   let stdlib = Option.get(Grain_utils.Config.stdlib_dir^);
 
   let preopen =
@@ -186,7 +174,6 @@ let run = (~num_pages=?, ~extra_args=[||], file) => {
   let cmd =
     Array.concat([
       [|"grain", "run"|],
-      mem_flags,
       [|"-S", stdlib, "-I", Filepath.to_string(test_libs_dir), preopen|],
       [|file|],
       extra_args,
@@ -306,6 +293,7 @@ let makeRunner =
     (
       test,
       ~num_pages=?,
+      ~max_pages=?,
       ~config_fn=?,
       ~extra_args=?,
       ~module_header=module_header,
@@ -315,8 +303,15 @@ let makeRunner =
     ) => {
   test(name, ({expect}) => {
     Config.preserve_all_configs(() => {
-      ignore @@ compile(~num_pages?, ~config_fn?, name, module_header ++ prog);
-      let (result, _) = run(~num_pages?, ~extra_args?, wasmfile(name));
+      ignore @@
+      compile(
+        ~num_pages?,
+        ~max_pages?,
+        ~config_fn?,
+        name,
+        module_header ++ prog,
+      );
+      let (result, _) = run(~extra_args?, wasmfile(name));
       expect.string(result).toEqual(expected);
     })
   });
@@ -327,6 +322,7 @@ let makeErrorRunner =
       test,
       ~check_exists=true,
       ~num_pages=?,
+      ~max_pages=?,
       ~config_fn=?,
       ~module_header=module_header,
       name,
@@ -335,8 +331,15 @@ let makeErrorRunner =
     ) => {
   test(name, ({expect}) => {
     Config.preserve_all_configs(() => {
-      ignore @@ compile(~num_pages?, ~config_fn?, name, module_header ++ prog);
-      let (result, _) = run(~num_pages?, wasmfile(name));
+      ignore @@
+      compile(
+        ~num_pages?,
+        ~max_pages?,
+        ~config_fn?,
+        name,
+        module_header ++ prog,
+      );
+      let (result, _) = run(wasmfile(name));
       if (check_exists) {
         expect.string(result).toMatch(expected);
       } else {
@@ -347,12 +350,13 @@ let makeErrorRunner =
 };
 
 let makeFileRunner =
-    (test, ~num_pages=?, ~config_fn=?, name, filename, expected) => {
+    (test, ~num_pages=?, ~max_pages=?, ~config_fn=?, name, filename, expected) => {
   test(name, ({expect}) => {
     Config.preserve_all_configs(() => {
       let infile = grainfile(filename);
       let outfile = wasmfile(name);
-      ignore @@ compile_file(~num_pages?, ~config_fn?, infile, outfile);
+      ignore @@
+      compile_file(~num_pages?, ~max_pages?, ~config_fn?, infile, outfile);
       let (result, _) = run(outfile);
       expect.string(result).toEqual(expected);
     })

diff --git a/compiler/test/suites/basic_functionality.re b/compiler/test/suites/basic_functionality.re
@@ -377,6 +377,6 @@ describe("basic functionality", ({test, testSkip}) => {
     ~config_fn=smallestFileConfig,
     "smallest_grain_program",
     "",
-    4769,
+    5165,
   );
 });
diff --git a/compiler/test/suites/gc.re b/compiler/test/suites/gc.re
@@ -8,24 +8,25 @@ let makeGcProgram = (program, heap_size) => {
     from "runtime/malloc" include Malloc
     from "runtime/unsafe/memory" include Memory
 
-    @disableGC
-    primitive heapStart = "@heap.start"
-
-    @disableGC
-    let leak = () => {
-      use WasmI32.{ (+), (-) }
-      // find current memory pointer, subtract space for two malloc headers + 1 GC header
-      let offset = Memory.malloc(8n) - 24n
-      // Calculate how much memory is left
-      let availableMemory = offset - (Malloc._RESERVED_RUNTIME_SPACE + heapStart())
-      // Calculate how much memory to leak
-      let toLeak = availableMemory - %dn
-      // Memory is not reclaimed due to no gc context
-      // This will actually leak 16 extra bytes because of the headers
-      Memory.malloc(toLeak - 16n);
-      void
+    @unsafe
+    let _ = {
+      use WasmI32.{(*), (-), (==)}
+      // Leak all available memory
+      // The first call to malloc ensures it has been initialized
+      Malloc.malloc(8n)
+      Malloc.leakAll()
+      // Next allocation will grow the memory by 1 page (64kib)
+      // We'll manually leak all memory except what should be reserved for the test
+      // Round reserved memory to nearest block size
+      let reserved = %dn
+      // If only one unit is requested, the allocator will include it in our next malloc,
+      // so we request 2 instead
+      let reserved = if (reserved == 1n) 2n else reserved
+      // one page - 2 malloc headers - 1 gc header - extra morecore unit - reserved space
+      let toLeak = 65536n - 16n - 8n - 64n - reserved * 64n
+      Memory.malloc(toLeak)
     }
-    leak();
+
     %s
     |},
     heap_size,
@@ -43,6 +44,7 @@ describe("garbage collection", ({test, testSkip}) => {
   let assertRunGC = (name, heapSize, prog, expected) =>
     makeRunner(
       ~num_pages=1,
+      ~max_pages=2,
       test_or_skip,
       name,
       makeGcProgram(prog, heapSize),
@@ -52,38 +54,42 @@ describe("garbage collection", ({test, testSkip}) => {
     makeErrorRunner(
       test_or_skip,
       ~num_pages=1,
+      ~max_pages=2,
       name,
       makeGcProgram(prog, heapSize),
       expected,
     );
 
   // oom tests
+  // The allocator will use 2 units for the first allocation and then oom
   assertRunGCError(
     "oomgc1",
-    48,
+    2,
     "(1, (3, 4))",
     "Maximum memory size exceeded",
   );
-  assertRunGC("oomgc2", 64, "(1, (3, 4))", "");
-  assertRunGC("oomgc3", 32, "(3, 4)", "");
+  // This requires only 2 units, but if only two are requested they would be
+  // used by the first allocation
+  assertRunGC("oomgc2", 3, "(1, (3, 4))", "");
+  assertRunGC("oomgc3", 1, "(3, 4)", "");
 
   // gc tests
   assertRunGC(
     "gc1",
-    160,
+    5,
     "let f = (() => (1, 2));\n       {\n         f();\n         f();\n         f();\n         f()\n       }",
     "",
   );
   /* https://github.com/grain-lang/grain/issues/774 */
   assertRunGC(
     "gc3",
-    1024,
+    17,
     "let foo = (s: String) => void\nlet printBool = (b: Bool) => foo(if (b) \"true\" else \"false\")\n\nlet b = true\nfor (let mut i=0; i<100000; i += 1) {\n  printBool(true)\n}",
     "",
   );
   assertRunGCError(
     "fib_gc_err",
-    256,
+    5,
     {|
     let fib = x => {
       let rec fib_help = (n, acc) => {
@@ -102,7 +108,7 @@ describe("garbage collection", ({test, testSkip}) => {
   );
   assertRunGC(
     "fib_gc",
-    512,
+    9,
     {|
     let fib = x => {
       let rec fib_help = (n, acc) => {
@@ -121,7 +127,7 @@ describe("garbage collection", ({test, testSkip}) => {
   );
   assertRunGC(
     "loop_gc",
-    256,
+    5,
     {|
     for (let mut i = 0; i < 512; i += 1) {
       let string = "string"
@@ -134,7 +140,7 @@ describe("garbage collection", ({test, testSkip}) => {
   );
   assertRunGC(
     "long_lists",
-    20000,
+    350,
     {|
     from "list" include List
     use List.*
@@ -162,7 +168,6 @@ describe("garbage collection", ({test, testSkip}) => {
     |},
     "true\n",
   );
-  assertFileRun("malloc_tight", "mallocTight", "");
   assertFileRun("memory_grow1", "memoryGrow", "1000000000000\n");
   assertMemoryLimitedFileRun(
     "loop_memory_reclaim",

diff --git a/docs/contributor/memory_management.md b/docs/contributor/memory_management.md
@@ -6,7 +6,7 @@ We ultimately aim to replace Grain's bespoke memory management with the WebAssem
 
 ## Memory Allocator
 
-Grain uses a [memory allocator](https://github.com/grain-lang/grain/blob/main/stdlib/runtime/malloc.gr) derived from the `malloc`/`free` example given in Kernighan and Ritchie's ["The C Programming Language"](https://kremlin.cc/k&r.pdf) (K&R C), pages 185-188 (PDF page 199). This module exports the following values:
+More documentation about Grain's [memory allocator](https://github.com/grain-lang/grain/blob/main/stdlib/runtime/malloc.gr) can be found in that module. It exports the following values:
 
 ```grain
 /**
@@ -31,11 +31,9 @@ export let malloc: (nbytes: WasmI32) -> WasmI32
 export let free = (ap: WasmI32) => Void
 
 /**
- * Returns the current free list pointer (used for debugging)
- *
- * @returns The free list pointer
+ * Leaks all memory in all free lists; used for testing.
  */
-export let getFreePtr = () => WasmI32
+export let leakAll = () => Void
 ```
 
 These functions should be familiar to programmers who have used `malloc` and `free` in C (and C-like languages). For further reading, refer to this Wikipedia page: [C dynamic memory allocation](https://en.wikipedia.org/wiki/C_dynamic_memory_allocation). The semantics of these functions align near-identically with those of C's corresponding functions.
-Original file line number
+Diff line change
@@ Expand Up / @@ -377,6 +377,6 @@ describe("basic functionality", ({test, testSkip}) => { @@
         ~config_fn=smallestFileConfig,
         "smallest_grain_program",
         "",
-,
+,
       );
     });