Tighten alignment promises for halide_malloc()

Makefile

@@ -753,8 +753,8 @@ RUNTIME_CPP_COMPONENTS = \
   aarch64_cpu_features \
   alignment_128 \
   alignment_32 \
-  allocation_cache \
   alignment_64 \
+  allocation_cache \
   android_clock \
   android_host_cpu_count \
   android_io \
@@ -778,8 +778,8 @@ RUNTIME_CPP_COMPONENTS = \
   halide_buffer_t \
   hexagon_cache_allocator \
   hexagon_cpu_features \
-  hexagon_dma_pool \
   hexagon_dma \
+  hexagon_dma_pool \
   hexagon_host \
   ios_io \
   linux_clock \
@@ -794,14 +794,15 @@ RUNTIME_CPP_COMPONENTS = \
   msan \
   msan_stubs \
   opencl \
-  openglcompute \
   opengl_egl_context \
   opengl_glx_context \
+  openglcompute \
   osx_clock \
   osx_get_symbol \
   osx_host_cpu_count \
   osx_opengl_context \
   osx_yield \
+  posix_aligned_alloc \
   posix_allocator \
   posix_clock \
   posix_error_handler \

apps/hannk/interpreter/interpreter.cpp

@@ -3,15 +3,13 @@
 #include "interpreter/transforms.h"
 #include "util/error_util.h"
 
+#include "HalideBuffer.h"  // for HALIDE_RUNTIME_BUFFER_ALLOCATION_ALIGNMENT
 #include "HalideRuntime.h"
 
 #include <map>
 #include <set>
 #include <unordered_set>
 
-// TODO: apparently not part of the public Halide API. Should it be?
-extern "C" int halide_malloc_alignment();
-
 namespace hannk {
 
 Interpreter::Interpreter(OpPtr m, InterpreterOptions options)
@@ -108,8 +106,10 @@ std::unique_ptr<char[]> allocate_tensors(const Op *root, const InterpreterOption
     // Feed this info to the allocation planner.
     // Let's assume that whatever alignment halide_malloc() needs is necessary here, too.
     // (Note that TFLite will complain if alignment is less than 64...)
+    // Let's assume that whatever alignment Halide::Runtime::Buffer needs is necessary here, too.
     constexpr int kTfLiteDefaultTensorAlignment = 64;
-    const size_t alignment = (size_t)std::max(halide_malloc_alignment(), kTfLiteDefaultTensorAlignment);
+    constexpr int kHalideBufferAlignment = HALIDE_RUNTIME_BUFFER_ALLOCATION_ALIGNMENT;
+    constexpr size_t alignment = (size_t)std::max(kHalideBufferAlignment, kTfLiteDefaultTensorAlignment);
     AllocationPlanner planner(alignment);
     for (auto &it : find_tensors.tensor_info) {
         auto &info = it.second;

src/LLVM_Runtime_Linker.cpp

@@ -119,6 +119,7 @@ DECLARE_CPP_INITMOD(osx_get_symbol)
 DECLARE_CPP_INITMOD(osx_host_cpu_count)
 DECLARE_CPP_INITMOD(osx_opengl_context)
 DECLARE_CPP_INITMOD(osx_yield)
+DECLARE_CPP_INITMOD(posix_aligned_alloc)
 DECLARE_CPP_INITMOD(posix_allocator)
 DECLARE_CPP_INITMOD(posix_clock)
 DECLARE_CPP_INITMOD(posix_error_handler)
@@ -761,6 +762,7 @@ std::unique_ptr<llvm::Module> link_with_wasm_jit_runtime(llvm::LLVMContext *c, c
     vector<std::unique_ptr<llvm::Module>> modules;
     modules.push_back(std::move(extra_module));
     modules.push_back(get_initmod_fake_thread_pool(c, bits_64, debug));
+    modules.push_back(get_initmod_posix_aligned_alloc(c, bits_64, debug));
     modules.push_back(get_initmod_posix_allocator(c, bits_64, debug));
     modules.push_back(get_initmod_halide_buffer_t(c, bits_64, debug));
     modules.push_back(get_initmod_destructors(c, bits_64, debug));
@@ -820,10 +822,16 @@ std::unique_ptr<llvm::Module> get_initial_module_for_target(Target t, llvm::LLVM
 
     vector<std::unique_ptr<llvm::Module>> modules;
 
+    const auto add_allocator = [&]() {
+        modules.push_back(get_initmod_posix_aligned_alloc(c, bits_64, debug));
+        modules.push_back(get_initmod_posix_allocator(c, bits_64, debug));
+    };
+
     if (module_type != ModuleGPU) {
         if (module_type != ModuleJITInlined && module_type != ModuleAOTNoRuntime) {
             // OS-dependent modules
             if (t.os == Target::Linux) {
+                add_allocator();
                 modules.push_back(get_initmod_posix_allocator(c, bits_64, debug));
                 modules.push_back(get_initmod_posix_error_handler(c, bits_64, debug));
                 modules.push_back(get_initmod_posix_print(c, bits_64, debug));
@@ -842,7 +850,7 @@ std::unique_ptr<llvm::Module> get_initial_module_for_target(Target t, llvm::LLVM
                 }
                 modules.push_back(get_initmod_posix_get_symbol(c, bits_64, debug));
             } else if (t.os == Target::WebAssemblyRuntime) {
-                modules.push_back(get_initmod_posix_allocator(c, bits_64, debug));
+                add_allocator();
                 modules.push_back(get_initmod_posix_error_handler(c, bits_64, debug));
                 modules.push_back(get_initmod_posix_print(c, bits_64, debug));
                 modules.push_back(get_initmod_posix_clock(c, bits_64, debug));
@@ -857,7 +865,7 @@ std::unique_ptr<llvm::Module> get_initial_module_for_target(Target t, llvm::LLVM
                 }
                 modules.push_back(get_initmod_fake_get_symbol(c, bits_64, debug));
             } else if (t.os == Target::OSX) {
-                modules.push_back(get_initmod_posix_allocator(c, bits_64, debug));
+                add_allocator();
                 modules.push_back(get_initmod_posix_error_handler(c, bits_64, debug));
                 modules.push_back(get_initmod_posix_print(c, bits_64, debug));
                 modules.push_back(get_initmod_osx_clock(c, bits_64, debug));
@@ -872,7 +880,7 @@ std::unique_ptr<llvm::Module> get_initial_module_for_target(Target t, llvm::LLVM
                 modules.push_back(get_initmod_osx_get_symbol(c, bits_64, debug));
                 modules.push_back(get_initmod_osx_host_cpu_count(c, bits_64, debug));
             } else if (t.os == Target::Android) {
-                modules.push_back(get_initmod_posix_allocator(c, bits_64, debug));
+                add_allocator();
                 modules.push_back(get_initmod_posix_error_handler(c, bits_64, debug));
                 modules.push_back(get_initmod_posix_print(c, bits_64, debug));
                 if (t.arch == Target::ARM) {
@@ -890,6 +898,7 @@ std::unique_ptr<llvm::Module> get_initial_module_for_target(Target t, llvm::LLVM
                 }
                 modules.push_back(get_initmod_posix_get_symbol(c, bits_64, debug));
             } else if (t.os == Target::Windows) {
+                modules.push_back(get_initmod_posix_aligned_alloc(c, bits_64, debug));
                 modules.push_back(get_initmod_posix_allocator(c, bits_64, debug));
                 modules.push_back(get_initmod_posix_error_handler(c, bits_64, debug));
                 modules.push_back(get_initmod_posix_print(c, bits_64, debug));
@@ -903,7 +912,7 @@ std::unique_ptr<llvm::Module> get_initial_module_for_target(Target t, llvm::LLVM
                 }
                 modules.push_back(get_initmod_windows_get_symbol(c, bits_64, debug));
             } else if (t.os == Target::IOS) {
-                modules.push_back(get_initmod_posix_allocator(c, bits_64, debug));
+                add_allocator();
                 modules.push_back(get_initmod_posix_error_handler(c, bits_64, debug));
                 modules.push_back(get_initmod_posix_print(c, bits_64, debug));
                 modules.push_back(get_initmod_posix_clock(c, bits_64, debug));
@@ -916,6 +925,7 @@ std::unique_ptr<llvm::Module> get_initial_module_for_target(Target t, llvm::LLVM
                     modules.push_back(get_initmod_posix_threads(c, bits_64, debug));
                 }
             } else if (t.os == Target::QuRT) {
+                modules.push_back(get_initmod_posix_aligned_alloc(c, bits_64, debug));
                 modules.push_back(get_initmod_qurt_allocator(c, bits_64, debug));
                 modules.push_back(get_initmod_qurt_yield(c, bits_64, debug));
                 if (tsan) {
@@ -930,11 +940,12 @@ std::unique_ptr<llvm::Module> get_initial_module_for_target(Target t, llvm::LLVM
                 // NoRuntime, as OS-agnostic modules like tracing are
                 // still included below.
                 if (t.arch == Target::Hexagon) {
+                    modules.push_back(get_initmod_posix_aligned_alloc(c, bits_64, debug));
                     modules.push_back(get_initmod_qurt_allocator(c, bits_64, debug));
                 }
                 modules.push_back(get_initmod_fake_thread_pool(c, bits_64, debug));
             } else if (t.os == Target::Fuchsia) {
-                modules.push_back(get_initmod_posix_allocator(c, bits_64, debug));
+                add_allocator();
                 modules.push_back(get_initmod_posix_error_handler(c, bits_64, debug));
                 modules.push_back(get_initmod_posix_print(c, bits_64, debug));
                 modules.push_back(get_initmod_fuchsia_clock(c, bits_64, debug));

src/runtime/CMakeLists.txt

@@ -52,6 +52,7 @@ set(RUNTIME_CPP
     osx_host_cpu_count
     osx_opengl_context
     osx_yield
+    posix_aligned_alloc
     posix_allocator
     posix_clock
     posix_error_handler
@@ -81,8 +82,8 @@ set(RUNTIME_CPP
     wasm_cpu_features
     windows_clock
     windows_cuda
-    windows_d3d12compute_x86
     windows_d3d12compute_arm
+    windows_d3d12compute_x86
     windows_get_symbol
     windows_io
     windows_opencl

src/runtime/HalideBuffer.h

@@ -52,6 +52,9 @@
 #define HALIDE_RUNTIME_BUFFER_ALLOCATION_ALIGNMENT 128
 #endif
 
+static_assert(((HALIDE_RUNTIME_BUFFER_ALLOCATION_ALIGNMENT & (HALIDE_RUNTIME_BUFFER_ALLOCATION_ALIGNMENT - 1)) == 0),
+              "HALIDE_RUNTIME_BUFFER_ALLOCATION_ALIGNMENT must be a power of 2.");
+
 // Unfortunately, not all C++17 runtimes support aligned_alloc
 // (it may depends on OS/SDK version); this is provided as an opt-out
 // if you are compiling on a platform that doesn't provide a (good)

src/runtime/HalideRuntime.h

@@ -365,10 +365,10 @@ extern int halide_set_num_threads(int n);
  *
  * Note that halide_malloc must return a pointer aligned to the
  * maximum meaningful alignment for the platform for the purpose of
- * vector loads and stores. The default implementation uses 32-byte
- * alignment, which is safe for arm and x86. Additionally, it must be
- * safe to read at least 8 bytes before the start and beyond the
- * end.
+ * vector loads and stores, *and* with an allocated size that is (at least)
+ * an integral multiple of that same alignment. The default implementation
+ * uses 32-byte alignment on arm and 64-byte alignment on x86. Additionally,
+ * it must be safe to read at least 8 bytes before the start and beyond the end.
  */
 //@{
 extern void *halide_malloc(void *user_context, size_t x);

src/runtime/alignment_128.cpp

@@ -1,5 +1,8 @@
 #include "runtime_internal.h"
 
-extern "C" WEAK_INLINE int halide_malloc_alignment() {
+extern "C" {
+
+WEAK_INLINE int halide_internal_malloc_alignment() {
     return 128;
 }
+}  // extern "C"

src/runtime/alignment_32.cpp

@@ -1,5 +1,8 @@
 #include "runtime_internal.h"
 
-extern "C" WEAK_INLINE int halide_malloc_alignment() {
+extern "C" {
+
+WEAK_INLINE int halide_internal_malloc_alignment() {
     return 32;
 }
+}  // extern "C"

src/runtime/alignment_64.cpp

@@ -1,5 +1,8 @@
 #include "runtime_internal.h"
 
-extern "C" WEAK_INLINE int halide_malloc_alignment() {
+extern "C" {
+
+WEAK_INLINE int halide_internal_malloc_alignment() {
     return 64;
 }
+}  // extern "C"

src/runtime/cache.cpp

@@ -106,7 +106,7 @@ struct CacheBlockHeader {
 // the hash entry.
 WEAK __attribute((always_inline)) size_t header_bytes() {
     size_t s = sizeof(CacheBlockHeader);
-    size_t mask = halide_malloc_alignment() - 1;
+    size_t mask = ::halide_internal_malloc_alignment() - 1;
     return (s + mask) & ~mask;
 }
 

src/runtime/posix_aligned_alloc.cpp

@@ -0,0 +1,43 @@
+#include "HalideRuntime.h"
+#include "runtime_internal.h"
+
+extern "C" {
+
+extern void *malloc(size_t);
+extern void free(void *);
+
+// An implementation of aligned_alloc() that is layered on top of malloc()/free().
+WEAK_INLINE void *halide_internal_aligned_alloc(size_t alignment, size_t size) {
+    // Alignment must be a power of two and >= sizeof(void*)
+    halide_debug_assert(nullptr, is_power_of_two(alignment) && alignment >= sizeof(void *));
+
+    // Allocate enough space for aligning the pointer we return.
+    //
+    // Always round allocations up to alignment size,
+    // so that all allocators follow the behavior of aligned_alloc() and
+    // return aligned pointer *and* aligned length.
+    const size_t aligned_size = align_up(size + alignment, alignment);
+
+    void *orig = ::malloc(aligned_size);
+    if (orig == nullptr) {
+        // Will result in a failed assertion and a call to halide_error
+        return nullptr;
+    }
+
+    // malloc() and friends must return a pointer aligned to at least
+    // alignof(std::max_align_t); we can't reasonably check that in
+    // the runtime, but we can realistically assume it's at least
+    // 8-aligned.
+    halide_debug_assert(nullptr, (((uintptr_t)orig) % 8) == 0);
+
+    // We want to store the original pointer prior to the pointer we return.
+    void *ptr = (void *)align_up((uintptr_t)orig + sizeof(void *), alignment);
+    ((void **)ptr)[-1] = orig;
+    return ptr;
+}
+
+WEAK_INLINE void halide_internal_aligned_free(void *ptr) {
+    ::free(((void **)ptr)[-1]);
+}
+
+}  // extern "C"

src/runtime/posix_allocator.cpp

@@ -1,29 +1,18 @@
 #include "HalideRuntime.h"
 #include "runtime_internal.h"
 
-#include "printer.h"
-
 extern "C" {
 
 extern void *malloc(size_t);
 extern void free(void *);
 
 WEAK void *halide_default_malloc(void *user_context, size_t x) {
-    // Allocate enough space for aligning the pointer we return.
-    const size_t alignment = halide_malloc_alignment();
-    void *orig = malloc(x + alignment);
-    if (orig == nullptr) {
-        // Will result in a failed assertion and a call to halide_error
-        return nullptr;
-    }
-    // We want to store the original pointer prior to the pointer we return.
-    void *ptr = (void *)(((size_t)orig + alignment + sizeof(void *) - 1) & ~(alignment - 1));
-    ((void **)ptr)[-1] = orig;
-    return ptr;
+    const size_t alignment = ::halide_internal_malloc_alignment();
+    return ::halide_internal_aligned_alloc(alignment, x);
 }
 
 WEAK void halide_default_free(void *user_context, void *ptr) {
-    free(((void **)ptr)[-1]);
+    ::halide_internal_aligned_free(ptr);
 }
 }
 

src/runtime/qurt_allocator.cpp

@@ -1,7 +1,7 @@
 #include "HalideRuntime.h"
+#include "runtime_internal.h"
 
 extern "C" {
-
 extern void *malloc(size_t);
 extern void free(void *);
 }
@@ -10,26 +10,13 @@ namespace Halide {
 namespace Runtime {
 namespace Internal {
 
-WEAK void *aligned_malloc(size_t alignment, size_t size) {
-    // We also need to align the size of the buffer.
-    size = (size + alignment - 1) & ~(alignment - 1);
-
-    // Allocate enough space for aligning the pointer we return.
-    void *orig = malloc(size + alignment);
-    if (orig == nullptr) {
-        // Will result in a failed assertion and a call to halide_error
-        return nullptr;
-    }
-    // We want to store the original pointer prior to the pointer we return.
-    void *ptr = (void *)(((size_t)orig + alignment + sizeof(void *) - 1) & ~(alignment - 1));
-    ((void **)ptr)[-1] = orig;
+ALWAYS_INLINE void *aligned_malloc(size_t alignment, size_t size) {
+    void *ptr = ::halide_internal_aligned_alloc(alignment, size);
     return ptr;
 }
 
-WEAK void aligned_free(void *ptr) {
-    if (ptr) {
-        free(((void **)ptr)[-1]);
-    }
+ALWAYS_INLINE void aligned_free(void *ptr) {
+    ::halide_internal_aligned_free(ptr);
 }
 
 // We keep a small pool of small pre-allocated buffers for use by Halide
@@ -59,8 +46,7 @@ WEAK __attribute__((destructor)) void halide_allocator_cleanup() {
 }  // namespace Halide
 
 WEAK void *halide_default_malloc(void *user_context, size_t x) {
-    // Hexagon needs up to 128 byte alignment.
-    const size_t alignment = 128;
+    const size_t alignment = ::halide_internal_malloc_alignment();
 
     if (x <= buffer_size) {
         for (int i = 0; i < num_buffers; ++i) {