add C test; add version that is partially working

Makefile

@@ -1,4 +1,6 @@
-all:
+CFLAGS=-Wall -Wextra -Werror -g -std=c17
+
+all: aligned_alloc_demo
 	cargo fmt
 	cargo test
 	cargo check
@@ -13,5 +15,7 @@ all:
 		-A clippy::cast-sign-loss \
 		-A clippy::cast-possible-truncation
 
+	clang-format -i '-style={BasedOnStyle: Google, ColumnLimit: 100}' *.c
+
 run_native:
 	RUSTFLAGS="-C target-cpu=native" cargo run --profile=release-nativecpu

README.md

@@ -0,0 +1,14 @@
+# Huge Page Demo
+
+This is a demonstration of using huge pages on Linux to get better performance.
+
+This will compile and run on non-Linux platforms, but won't use huge pages.
+
+For more details, see [Reliably allocating huge pages in Linux](https://mazzo.li/posts/check-huge-page.html).
+
+
+# Malloc/Mmap behaviour
+
+On Ubuntu 20.04.5 with kernel 5.15.0-1023-aws and glibc 2.31-0ubuntu9.9, malloc 4 GiB calls mmap to allocate 4 GiB + 4 KiB, then returns a pointer that is +0x10 (+16) from the pointer actually returned by mmap. Using aligned_alloc calls mmap to allocate 5 GiB + 4 KiB (size + alignment + 1 page?), then returns an aligned pointer. Calling mmap to allocate 4 GiB returns a pointer that is not aligned. E.g. On my system, I get one that is 32 kiB aligned.
+
+On Mac OS X 13.1 on an M1 ARM CPU, using mmap to request 4 GiB of memory returns a block that is aligned to a 1 GiB boundary. The same appears to be true for using malloc. I didn't fight to get dtruss to work to see what malloc is actually doing.

aligned_alloc_demo.c

@@ -0,0 +1,25 @@
+#include <assert.h>
+#include <stdint.h>
+#include <stdio.h>
+#include <stdlib.h>
+
+int main() {
+  static const size_t FOUR_GIB_IN_BYTES = UINT64_C(4) << 30;
+  static const size_t HUGE_2MIB_ALIGNMENT = UINT64_C(2) << 20;
+  static const size_t HUGE_2MIB_MASK = HUGE_2MIB_ALIGNMENT - 1;
+  static const size_t HUGE_1GIB_ALIGNMENT = UINT64_C(1) << 30;
+  static const size_t HUGE_1GIB_MASK = HUGE_1GIB_ALIGNMENT - 1;
+
+  void *plain_malloc = malloc(FOUR_GIB_IN_BYTES);
+  printf("malloc 4GiB = %p; 2MiB aligned? %d; 1GiB aligned? %d\n", plain_malloc,
+         ((size_t)plain_malloc & HUGE_2MIB_MASK) == 0,
+         ((size_t)plain_malloc & HUGE_1GIB_MASK) == 0);
+  free(plain_malloc);
+
+  void *aligned = aligned_alloc(HUGE_1GIB_ALIGNMENT, FOUR_GIB_IN_BYTES);
+  printf("aligned_alloc 4GiB = %p; 2MiB aligned? %d; 1GiB aligned? %d\n", aligned,
+         ((size_t)aligned & HUGE_2MIB_MASK) == 0, ((size_t)aligned & HUGE_1GIB_MASK) == 0);
+  free(aligned);
+
+  return 0;
+}

src/main.rs

@@ -1,16 +1,18 @@
 use core::slice;
 use memory_stats::memory_stats;
-use nix::libc::uintptr_t;
 use nix::sys::mman::{MapFlags, ProtFlags};
-#[cfg(any(test, target_os = "linux"))]
-use nix::unistd::SysconfVar;
 use rand::distributions::Distribution;
 use rand::{distributions::Uniform, RngCore, SeedableRng};
 use std::error::Error;
 use std::num::NonZeroUsize;
 use std::os::raw::c_void;
 use std::time::Instant;
 
+#[cfg(target_os = "linux")]
+use nix::sys::mman::MmapAdvise;
+#[cfg(any(test, target_os = "linux"))]
+use nix::unistd::SysconfVar;
+
 const FILLED: u64 = 0x42;
 
 fn main() -> Result<(), Box<dyn Error>> {
@@ -35,7 +37,7 @@ fn main() -> Result<(), Box<dyn Error>> {
     rnd_accesses(&mut rng, &v);
     let mem_after = memory_stats().unwrap();
     println!(
-        "RSS before: {}; RSS after: {}; diff: {}",
+        "RSS before: {}; RSS after: {}; diff: {}\n",
         humanunits::bytes_string(mem_before.physical_mem),
         humanunits::bytes_string(mem_after.physical_mem),
         humanunits::bytes_string(mem_after.physical_mem - mem_before.physical_mem)
@@ -51,7 +53,7 @@ fn main() -> Result<(), Box<dyn Error>> {
     let end = Instant::now();
     let duration = end - start;
     println!(
-        "\nMmapSlice: alloc and filled {TEST_SIZE_GIB} GiB in {duration:?}; {}",
+        "MmapSlice: alloc and filled {TEST_SIZE_GIB} GiB in {duration:?}; {}",
         humanunits::byte_rate_string(TEST_SIZE_BYTES, duration)
     );
     rnd_accesses(&mut rng, v.slice());
@@ -67,39 +69,101 @@ fn main() -> Result<(), Box<dyn Error>> {
     Ok(())
 }
 
-struct MmapU64Slice<'a> {
-    slice: &'a mut [u64],
+struct MmapAligned {
+    mmap_pointer: *mut c_void,
+    aligned_size: usize,
 }
 
-impl<'a> MmapU64Slice<'a> {
-    fn new_zero(items: usize) -> Result<Self, nix::errno::Errno> {
-        const HUGE_2MIB_MASK: uintptr_t = (2 << 20) - 1;
-        const HUGE_1GIB_MASK: uintptr_t = (1 << 30) - 1;
+impl MmapAligned {
+    // argument order is the same as aligned_alloc.
+    fn new(alignment: usize, size: usize) -> Result<Self, nix::errno::Errno> {
+        // worse case alignment: mmap returns 1 byte off the alignment, we must waste alignment-1 bytes.
+        // To ensure we can do this, we request size+alignment bytes.
+        // This shouldn't be so bad: untouched pages won't actually be allocated.
+        let aligned_size =
+            NonZeroUsize::new(size + alignment).expect("BUG: alignment and size must be > 0");
 
-        let mem_size = NonZeroUsize::new(items * 8).unwrap();
-
-        let pointer: *mut c_void;
-        let slice: &mut [u64];
+        let mmap_pointer: *mut c_void;
         unsafe {
-            pointer = nix::sys::mman::mmap(
+            mmap_pointer = nix::sys::mman::mmap(
                 None,
-                mem_size,
+                aligned_size,
                 ProtFlags::PROT_READ | ProtFlags::PROT_WRITE,
                 MapFlags::MAP_ANONYMOUS | MapFlags::MAP_PRIVATE,
                 0,
                 0,
             )?;
+        }
+
+        let allocation = Self {
+            mmap_pointer,
+            aligned_size: aligned_size.get(),
+        };
+        let aligned_pointer = allocation.get_aligned_mut(alignment);
+        let allocation_end = mmap_pointer as usize + aligned_size.get();
+        assert!(aligned_pointer as usize + size <= allocation_end);
+
+        Ok(allocation)
+    }
+
+    fn get_aligned_mut(&self, alignment: usize) -> *mut c_void {
+        align_pointer_value(alignment, self.mmap_pointer as usize) as *mut c_void
+    }
+}
+
+impl Drop for MmapAligned {
+    fn drop(&mut self) {
+        println!(
+            "dropping mmap pointer={:x?} len={}...",
+            self.mmap_pointer, self.aligned_size
+        );
+
+        unsafe {
+            nix::sys::mman::munmap(self.mmap_pointer.cast::<c_void>(), self.aligned_size)
+                .expect("BUG: munmap should not fail");
+        }
+    }
+}
+
+fn align_pointer_value(alignment: usize, pointer_value: usize) -> usize {
+    // see bit hacks to check if power of two:
+    // https://graphics.stanford.edu/~seander/bithacks.html#DetermineIfPowerOf2
+    assert_eq!(0, (alignment & (alignment - 1)));
+    // round pointer_value up to nearest alignment; assumes there is sufficient space
+    let alignment_mask = !(alignment - 1);
+    (pointer_value + (alignment - 1)) & alignment_mask
+}
+
+struct MmapU64Slice<'a> {
+    // MmapAligned unmaps the mapping using the Drop trait but is otherwise not read
+    _allocation: MmapAligned,
+    slice: &'a mut [u64],
+}
+
+impl<'a> MmapU64Slice<'a> {
+    fn new_zero(items: usize) -> Result<Self, nix::errno::Errno> {
+        const HUGE_2MIB_MASK: usize = (2 << 20) - 1;
+        const HUGE_1GIB_ALIGNMENT: usize = 1 << 30;
+        const HUGE_1GIB_MASK: usize = HUGE_1GIB_ALIGNMENT - 1;
 
-            slice = slice::from_raw_parts_mut(pointer.cast::<u64>(), items);
+        let mem_size = items * 8;
+        let allocation = MmapAligned::new(HUGE_1GIB_ALIGNMENT, mem_size)?;
+        let slice_pointer = allocation.get_aligned_mut(HUGE_1GIB_ALIGNMENT);
+        let slice: &mut [u64];
+        unsafe {
+            slice = slice::from_raw_parts_mut(slice_pointer.cast::<u64>(), items);
         }
 
-        let mut m = Self { slice };
+        let mut m = Self {
+            _allocation: allocation,
+            slice,
+        };
         m.madvise_hugepages_on_linux();
 
         let (mmap_pointer, _) = m.mmap_parts();
         let ptr_usize = mmap_pointer as usize;
         println!(
-            "mmap returned {mmap_pointer:x?}; aligned to 2MiB (0x{HUGE_2MIB_MASK:x})? {}; aligned to 1GiB (0x{HUGE_1GIB_MASK:x})? {}",
+            "mmap aligned returned {mmap_pointer:x?}; aligned to 2MiB (0x{HUGE_2MIB_MASK:x})? {}; aligned to 1GiB (0x{HUGE_1GIB_MASK:x})? {}",
             ptr_usize & HUGE_2MIB_MASK == 0,
             ptr_usize & HUGE_1GIB_MASK == 0
         );
@@ -108,12 +172,13 @@ impl<'a> MmapU64Slice<'a> {
 
     #[cfg(target_os = "linux")]
     fn madvise_hugepages_on_linux(&mut self) {
-        use nix::libc::HW_PAGESIZE;
-
         let (mmap_pointer, mmap_len) = self.mmap_parts();
         let advise_flags = MmapAdvise::MADV_HUGEPAGE;
-        nix::sys::mman::madvise(mmap_pointer, mmap_len, advise_flags)
-            .expect("BUG: madvise must succeed");
+
+        unsafe {
+            nix::sys::mman::madvise(mmap_pointer.cast::<c_void>(), mmap_len, advise_flags)
+                .expect("BUG: madvise must succeed");
+        }
 
         touch_pages(self.slice);
     }
@@ -133,25 +198,13 @@ impl<'a> MmapU64Slice<'a> {
         self.slice
     }
 
-    fn mmap_parts(&self) -> (*const u64, usize) {
-        let mmap_pointer = self.slice().as_ptr();
+    fn mmap_parts(&mut self) -> (*mut u64, usize) {
+        let mmap_pointer = self.slice_mut().as_mut_ptr();
         let mmap_len = self.slice.len() * 8;
         (mmap_pointer, mmap_len)
     }
 }
 
-impl<'a> Drop for MmapU64Slice<'a> {
-    fn drop(&mut self) {
-        let (mmap_pointer, mmap_len) = self.mmap_parts();
-        // println!("dropping mmap pointer={mmap_pointer:x?} len={mmap_len} ...");
-
-        unsafe {
-            nix::sys::mman::munmap(mmap_pointer as *mut c_void, mmap_len)
-                .expect("BUG: munmap should not fail");
-        }
-    }
-}
-
 fn rnd_accesses(rng: &mut dyn RngCore, data: &[u64]) {
     const NUM_ACCESSES: usize = 200_000_000;
 
@@ -175,7 +228,7 @@ fn touch_pages(s: &mut [u64]) {
     let page_size = nix::unistd::sysconf(SysconfVar::PAGE_SIZE)
         .expect("BUG: sysconf(_SC_PAGESIZE) must work")
         .expect("BUG: page size must not be None");
-    println!("page_size={page_size}");
+    println!("touch_pages with page_size={page_size}");
 
     // write a zero every stride elements, which should fault every page
     let stride = page_size as usize / 8;
@@ -195,4 +248,38 @@ mod test {
         let mut v: Vec<u64> = vec![0; SIZE];
         touch_pages(&mut v);
     }
+
+    #[test]
+    fn test_align_pointer_value() {
+        const ONE_GIB: usize = 1 << 30;
+        const SEVEN_GIB: usize = 7 * ONE_GIB;
+        const EIGHT_GIB: usize = 8 * ONE_GIB;
+        assert_eq!(SEVEN_GIB, align_pointer_value(ONE_GIB, SEVEN_GIB));
+        assert_eq!(EIGHT_GIB, align_pointer_value(ONE_GIB, SEVEN_GIB + 1));
+        assert_eq!(
+            EIGHT_GIB,
+            align_pointer_value(ONE_GIB, SEVEN_GIB + (ONE_GIB - 1))
+        );
+        assert_eq!(EIGHT_GIB, align_pointer_value(ONE_GIB, SEVEN_GIB + ONE_GIB));
+    }
+
+    #[test]
+    fn test_mmap_aligned() {
+        const ONE_GIB: usize = 1 << 30;
+        const ONE_MIB: usize = 1 << 20;
+        let aligned_alloc = MmapAligned::new(ONE_GIB, ONE_MIB).unwrap();
+        let aligned_pointer = aligned_alloc.get_aligned_mut(ONE_GIB);
+
+        // check that we can write to the slice
+        let slice: &mut [u64];
+        unsafe {
+            slice = slice::from_raw_parts_mut(aligned_pointer.cast::<u64>(), ONE_MIB / 8);
+        }
+        slice[0] = 0x42;
+        slice[slice.len() - 1] = 0x42;
+        assert_eq!(0x42, slice[0]);
+        assert_eq!(0, slice[1]);
+        assert_eq!(0, slice[slice.len() - 2]);
+        assert_eq!(0x42, slice[slice.len() - 1]);
+    }
 }