rts: Introduce stgMallocAlignedBytes

3 files changed, 57 insertions, 13 deletions

diff --git a/rts/Capability.c b/rts/Capability.c
index 98c80887e5..f5514cbeee 100644
--- a/rts/Capability.c
+++ b/rts/Capability.c
@@ -1274,7 +1274,7 @@ freeCapabilities (void)
         Capability *cap = getCapability(i);
         freeCapability(cap);
         if (cap != &MainCapability) {
-            stgFree(cap);
+            stgFreeAligned(cap);
         }
     }
 #else
diff --git a/rts/RtsUtils.c b/rts/RtsUtils.c
index d3dcdf3092..4cac10ba15 100644
--- a/rts/RtsUtils.c
+++ b/rts/RtsUtils.c
@@ -57,9 +57,9 @@ extern char *ctime_r(const time_t *, char *);
 void *
 stgMallocBytes (size_t n, char *msg)
 {
-    void *space;
+    void *space = malloc(n);
 
-    if ((space = malloc(n)) == NULL) {
+    if (space == NULL) {
       /* Quoting POSIX.1-2008 (which says more or less the same as ISO C99):
        *
        *   "Upon successful completion with size not equal to 0, malloc() shall
@@ -128,6 +128,53 @@ stgFree(void* p)
   free(p);
 }
 
+// N.B. Allocations resulting from this function must be freed by
+// `stgFreeAligned`, not `stgFree`. This is necessary due to the properties of Windows' `_aligned_malloc`
+void *
+stgMallocAlignedBytes (size_t n, size_t align, char *msg)
+{
+    void *space;
+
+#if defined(mingw32_HOST_OS)
+    space = _aligned_malloc(n, align);
+#else
+    if (posix_memalign(&space, align, n)) {
+        space = NULL; // Allocation failed
+    }
+#endif
+
+    if (space == NULL) {
+      /* Quoting POSIX.1-2008 (which says more or less the same as ISO C99):
+       *
+       *   "Upon successful completion with size not equal to 0, malloc() shall
+       *   return a pointer to the allocated space. If size is 0, either a null
+       *   pointer or a unique pointer that can be successfully passed to free()
+       *   shall be returned. Otherwise, it shall return a null pointer and set
+       *   errno to indicate the error."
+       *
+       * Consequently, a NULL pointer being returned by `malloc()` for a 0-size
+       * allocation is *not* to be considered an error.
+       */
+      if (n == 0) return NULL;
+
+      /* don't fflush(stdout); WORKAROUND bug in Linux glibc */
+      rtsConfig.mallocFailHook((W_) n, msg);
+      stg_exit(EXIT_INTERNAL_ERROR);
+    }
+    IF_DEBUG(zero_on_gc, memset(space, 0xbb, n));
+    return space;
+}
+
+void
+stgFreeAligned (void *p)
+{
+#if defined(mingw32_HOST_OS)
+    _aligned_free(p);
+#else
+    free(p);
+#endif
+}
+
 /* -----------------------------------------------------------------------------
    Stack/heap overflow
    -------------------------------------------------------------------------- */
diff --git a/rts/RtsUtils.h b/rts/RtsUtils.h
index 1c5aac82a1..c78f9d7ab3 100644
--- a/rts/RtsUtils.h
+++ b/rts/RtsUtils.h
@@ -29,16 +29,9 @@ void *stgMallocBytes(size_t n, char *msg)
  * See: https://gitlab.haskell.org/ghc/ghc/-/issues/22380
  */
 
-void *stgReallocBytes(void *p, size_t n, char *msg)
-    STG_MALLOC1(stgFree)
-    STG_ALLOC_SIZE1(2)
-    STG_RETURNS_NONNULL;
-/* Note: `stgRallocBytes` can *not* be tagged as `STG_MALLOC`
- * since its return value *can* alias an existing pointer (i.e.,
- * the given pointer `p`).
- * See the documentation of the `malloc` attribute in the GCC manual
- * for more information.
- */
+void *stgMallocAlignedBytes(size_t n, size_t align, char *msg);
+
+void *stgReallocBytes(void *p, size_t n, char *msg);
 
 void *stgCallocBytes(size_t count, size_t size, char *msg)
     STG_MALLOC STG_MALLOC1(stgFree)
@@ -48,6 +41,10 @@ void *stgCallocBytes(size_t count, size_t size, char *msg)
 char *stgStrndup(const char *s, size_t n)
     STG_MALLOC STG_MALLOC1(stgFree);
 
+void *stgMallocAlignedBytes(size_t n, size_t align, char *msg);
+
+void stgFreeAligned(void *p);
+
 /* -----------------------------------------------------------------------------
  * Misc other utilities
  * -------------------------------------------------------------------------- */