[K/N] Tweak aligned allocation on all platforms ^KT-37272

Merge-request: KT-MR-6341
Merged-by: Alexander Shabalin <Alexander.Shabalin@jetbrains.com>

kotlin-native/runtime/src/main/cpp/Natives.cpp

@@ -21,6 +21,7 @@
 #include <limits>
 #include <type_traits>
 
+#include "Alignment.hpp"
 #include "KAssert.h"
 #include "KString.h"
 #include "StackTrace.hpp"
@@ -87,20 +88,18 @@ void* Kotlin_interop_malloc(KLong size, KInt align) {
   if (size < 0 || static_cast<std::make_unsigned_t<decltype(size)>>(size) > std::numeric_limits<size_t>::max()) {
     return nullptr;
   }
-  RuntimeAssert(align > 0, "Unsupported alignment");
-  RuntimeAssert((align & (align - 1)) == 0, "Alignment must be power of two");
+  RuntimeAssert(align > 0, "Invalid alignment %d", align);
+  size_t actualAlign = static_cast<size_t>(align);
+  size_t actualSize = AlignUp(static_cast<size_t>(size), actualAlign);
 
-  void* result = std_support::aligned_calloc(align, 1, size);
-  if ((reinterpret_cast<uintptr_t>(result) & (align - 1)) != 0) {
-    // Unaligned!
-    RuntimeAssert(false, "unsupported alignment");
-  }
+  void* result = std::memset(std_support::aligned_malloc(actualAlign, actualSize), 0, actualSize);
+  RuntimeAssert(IsAligned(result, actualAlign), "aligned_malloc result %p is not aligned to %zu", result, actualAlign);
 
   return result;
 }
 
 void Kotlin_interop_free(void* ptr) {
-    std_support::free(ptr);
+    std_support::aligned_free(ptr);
 }
 
 void Kotlin_system_exitProcess(KInt status) {

kotlin-native/runtime/src/main/cpp/std_support/CStdlib.cpp

@@ -9,45 +9,50 @@
 #include <cstdlib>
 #include <unistd.h>
 
+#include "Alignment.hpp"
+#include "KAssert.h"
+
 using namespace kotlin;
 
 #if KONAN_INTERNAL_DLMALLOC
 extern "C" void* dlmalloc(size_t);
 extern "C" void* dlcalloc(size_t, size_t);
 extern "C" void* dlrealloc(void*, size_t);
+extern "C" void* dlmemalign(size_t, size_t);
 extern "C" void dlfree(void*);
+
 #define malloc_impl dlmalloc
-#define aligned_alloc_impl(alignment, size) dlmalloc(size)
+#define aligned_malloc_impl dlmemalign
 #define calloc_impl dlcalloc
-#define aligned_calloc_impl(alignment, num, size) dlcalloc(num, size)
 #define realloc_impl dlrealloc
 #define free_impl dlfree
+#define aligned_free_impl dlfree
 #else
+#include <mm_malloc.h>
+
 #define malloc_impl std::malloc
-// TODO: Consider using std::aligned_alloc
-#define aligned_alloc_impl(alignment, size) std::malloc(size)
+#define aligned_malloc_impl(alignment, size) ::_mm_malloc(size, alignment)
 #define calloc_impl std::calloc
-#define aligned_calloc_impl(alignment, num, size) std::calloc(num, size)
 #define realloc_impl std::realloc
 #define free_impl std::free
+#define aligned_free_impl ::_mm_free
 #endif
 
 void* std_support::malloc(std::size_t size) noexcept {
     return malloc_impl(size);
 }
 
-void* std_support::aligned_alloc(std::size_t alignment, std::size_t size) noexcept {
-    return aligned_alloc_impl(alignment, size);
+void* std_support::aligned_malloc(std::size_t alignment, std::size_t size) noexcept {
+    // Enforcing alignment requirements of std::aligned_alloc.
+    RuntimeAssert(IsValidAlignment(alignment), "Invalid alignment %zu", alignment);
+    RuntimeAssert(IsAligned(size, alignment), "Size %zu must be aligned to %zu", size, alignment);
+    return aligned_malloc_impl(alignment, size);
 }
 
 void* std_support::calloc(std::size_t num, std::size_t size) noexcept {
     return calloc_impl(num, size);
 }
 
-void* std_support::aligned_calloc(std::size_t alignment, std::size_t num, std::size_t size) noexcept {
-    return aligned_calloc_impl(alignment, num, size);
-}
-
 void* std_support::realloc(void* ptr, std::size_t size) noexcept {
     return realloc_impl(ptr, size);
 }
@@ -56,6 +61,10 @@ void std_support::free(void* ptr) noexcept {
     return free_impl(ptr);
 }
 
+void std_support::aligned_free(void* ptr) noexcept {
+    return aligned_free_impl(ptr);
+}
+
 namespace konan {
 
 #if KONAN_INTERNAL_DLMALLOC

kotlin-native/runtime/src/main/cpp/std_support/CStdlib.hpp

@@ -10,13 +10,16 @@
 namespace kotlin::std_support {
 
 void* malloc(std::size_t size) noexcept;
-void* aligned_alloc(std::size_t alignment, std::size_t size) noexcept;
+// TODO: Replace with `aligned_alloc` that's compatible with normal `free`.
+// Allocate aligned memory. Must be freed with `aligned_free`.
+void* aligned_malloc(std::size_t alignment, std::size_t size) noexcept;
 
 void* calloc(std::size_t num, std::size_t size) noexcept;
-void* aligned_calloc(std::size_t alignment, std::size_t num, std::size_t size) noexcept;
 
 void* realloc(void* ptr, std::size_t size) noexcept;
 
 void free(void* ptr) noexcept;
+// Free memory allocated with aligned_malloc.
+void aligned_free(void* ptr) noexcept;
 
 } // namespace kotlin::std_support

kotlin-native/runtime/src/main/cpp/std_support/CStdlibTest.cpp

@@ -0,0 +1,143 @@
+/*
+ * Copyright 2010-2022 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the LICENSE file.
+ */
+
+#include "std_support/CStdlib.hpp"
+
+#include <cstdint>
+
+#include "gmock/gmock.h"
+#include "gtest/gtest.h"
+
+#include "Alignment.hpp"
+
+using namespace kotlin;
+
+namespace {
+
+inline constexpr auto mallocAlignment = alignof(std::max_align_t);
+
+struct Struct {
+    int32_t x;
+};
+
+using Vector4f = float __attribute__((__vector_size__(16)));
+
+struct alignas(32) OverAlignedStruct {
+    int32_t x;
+};
+static_assert(
+        alignof(OverAlignedStruct) > mallocAlignment,
+        "OverAlignedStruct should require alignment of more than supported by malloc");
+
+} // namespace
+
+TEST(StdSupportCStdlibTest, Malloc) {
+    {
+        Struct* ptr = (Struct*)std_support::malloc(sizeof(Struct));
+        EXPECT_TRUE(IsAligned(ptr, mallocAlignment));
+        ptr->x = 123;
+        EXPECT_THAT(ptr->x, 123);
+        std_support::free(ptr);
+    }
+    {
+        void* ptr = std_support::malloc(0);
+        EXPECT_TRUE(IsAligned(ptr, mallocAlignment));
+        std_support::free(ptr);
+    }
+}
+
+TEST(StdSupportCStdlibTest, Free) {
+    std_support::free(nullptr);
+}
+
+TEST(StdSupportCStdlibTest, Calloc) {
+    {
+        Struct* ptr = (Struct*)std_support::calloc(sizeof(Struct), 2);
+        EXPECT_TRUE(IsAligned(ptr, mallocAlignment));
+        EXPECT_THAT(ptr[0].x, 0);
+        EXPECT_THAT(ptr[1].x, 0);
+        ptr[0].x = 123;
+        ptr[1].x = 42;
+        EXPECT_THAT(ptr[0].x, 123);
+        EXPECT_THAT(ptr[1].x, 42);
+        std_support::free(ptr);
+    }
+    {
+        void* ptr = std_support::calloc(0, 2);
+        EXPECT_TRUE(IsAligned(ptr, mallocAlignment));
+        std_support::free(ptr);
+    }
+    {
+        void* ptr = std_support::calloc(sizeof(Struct), 0);
+        EXPECT_TRUE(IsAligned(ptr, mallocAlignment));
+        std_support::free(ptr);
+    }
+}
+
+TEST(StdSupportCStdlibTest, Realloc) {
+    {
+        Struct* ptr = (Struct*)std_support::malloc(sizeof(Struct));
+        ptr->x = 123;
+        Struct* newPtr = (Struct*)std_support::realloc(ptr, sizeof(Struct) * 3);
+        EXPECT_TRUE(IsAligned(newPtr, mallocAlignment));
+        EXPECT_THAT(newPtr[0].x, 123);
+        newPtr[1].x = 42;
+        newPtr[2].x = 13;
+        EXPECT_THAT(newPtr[1].x, 42);
+        EXPECT_THAT(newPtr[2].x, 13);
+        std_support::free(newPtr);
+    }
+    {
+        Struct* ptr = (Struct*)std_support::calloc(sizeof(Struct), 2);
+        EXPECT_TRUE(IsAligned(ptr, mallocAlignment));
+        ptr[0].x = 123;
+        ptr[1].x = 42;
+        EXPECT_THAT(ptr[0].x, 123);
+        EXPECT_THAT(ptr[1].x, 42);
+        Struct* newPtr = (Struct*)std_support::realloc(ptr, sizeof(Struct) * 3);
+        EXPECT_TRUE(IsAligned(newPtr, mallocAlignment));
+        EXPECT_THAT(newPtr[0].x, 123);
+        EXPECT_THAT(newPtr[1].x, 42);
+        newPtr[2].x = 13;
+        EXPECT_THAT(newPtr[2].x, 13);
+        std_support::free(newPtr);
+    }
+}
+
+TEST(StdSupportCStdlibTest, AlignedMalloc) {
+    {
+        Struct* ptr = (Struct*)std_support::aligned_malloc(alignof(Struct), sizeof(Struct));
+        EXPECT_TRUE(IsAligned(ptr, alignof(Struct)));
+        ptr->x = 123;
+        EXPECT_THAT(ptr->x, 123);
+        std_support::aligned_free(ptr);
+    }
+    {
+        void* ptr = std_support::aligned_malloc(alignof(int), 0);
+        EXPECT_TRUE(IsAligned(ptr, alignof(int)));
+        std_support::aligned_free(ptr);
+    }
+    {
+        Vector4f* ptr = (Vector4f*)std_support::aligned_malloc(alignof(Vector4f), sizeof(Vector4f));
+        EXPECT_TRUE(IsAligned(ptr, alignof(Vector4f)));
+        *ptr = {1.1f, 2.2f, 3.3f, 4.4f};
+        EXPECT_THAT((*ptr)[0], 1.1f);
+        EXPECT_THAT((*ptr)[1], 2.2f);
+        EXPECT_THAT((*ptr)[2], 3.3f);
+        EXPECT_THAT((*ptr)[3], 4.4f);
+        std_support::aligned_free(ptr);
+    }
+    {
+        OverAlignedStruct* ptr = (OverAlignedStruct*)std_support::aligned_malloc(alignof(OverAlignedStruct), sizeof(OverAlignedStruct));
+        EXPECT_TRUE(IsAligned(ptr, alignof(OverAlignedStruct)));
+        ptr->x = 123;
+        EXPECT_THAT(ptr->x, 123);
+        std_support::aligned_free(ptr);
+    }
+}
+
+TEST(StdSupportCStdlibTest, AlignedFree) {
+    std_support::aligned_free(nullptr);
+}

kotlin-native/runtime/src/main/cpp/std_support/README.md

@@ -12,8 +12,8 @@ Proposals:
 
 Adjustments:
 * `CStdlib.hpp` -
-  `std_support::malloc`, `std_support::aligned_alloc`, `std_support::calloc`, `std_support::realloc`, `std_support::free` that use custom allocation scheme,
-  `std_support::aligned_calloc` as a version of `calloc` that allows changing alignment.
+  `std_support::malloc`, `std_support::calloc`, `std_support::realloc`, `std_support::free` that use custom allocation scheme,
+  `std_support::aligned_malloc` and `std_support::aligned_free` as a version of `malloc` and `free` that allows changing alignment.
 * `Memory.hpp` -
   `std_support::allocator` using `std_support::calloc`/`std_support::free`,
   `std_support::default_delete` that uses `std_support::free`,