From dcf1b4406bd5a9c29296df4b00d9b756cdb3df65 Mon Sep 17 00:00:00 2001
From: Steven Johnson <srj@google.com>
Date: Mon, 5 Dec 2022 17:07:59 -0800
Subject: [PATCH] Promote halide_malloc_alignment() to be user-overridable

Currently, halide_malloc_alignment() is WEAK_INLINE, so even if you attempt to override it, it won't have any effect (the call sites in halide_malloc() will have inlined the implementation selected by LLVM_Runtime_Linker).

This promotes halide_malloc_alignment() to a documented piece of the runtime you can override, with the caveat that it must return the same value every time it is called. To reduce overhead in halide_malloc(), we copy the value into a global at startup time and use that instead.

By itself, this isn't that useful, but it's essential to have in order for #7189 to be robust; I'm pulling this change into a separate PR so that it can land separately, enabling downstream users to be able to update their halide_malloc() implementations ahead of the "real" change.
---
 apps/hannk/interpreter/interpreter.cpp |  3 ---
 src/runtime/HalideRuntime.h            | 12 +++++++-----
 src/runtime/alignment_128.cpp          |  2 +-
 src/runtime/alignment_32.cpp           |  2 +-
 src/runtime/alignment_64.cpp           |  2 +-
 src/runtime/posix_allocator.cpp        | 15 ++++++++++++++-
 src/runtime/qurt_allocator.cpp         |  8 ++++++--
 src/runtime/runtime_api.cpp            |  1 +
 src/runtime/runtime_internal.h         |  2 +-
 9 files changed, 32 insertions(+), 15 deletions(-)

diff --git a/apps/hannk/interpreter/interpreter.cpp b/apps/hannk/interpreter/interpreter.cpp
index ac1bd3329396..49badd849c5a 100644
--- a/apps/hannk/interpreter/interpreter.cpp
+++ b/apps/hannk/interpreter/interpreter.cpp
@@ -9,9 +9,6 @@
 #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)
diff --git a/src/runtime/HalideRuntime.h b/src/runtime/HalideRuntime.h
index 532874b0056a..03b15a065634 100644
--- a/src/runtime/HalideRuntime.h
+++ b/src/runtime/HalideRuntime.h
@@ -364,13 +364,15 @@ extern int halide_set_num_threads(int n);
  * halide_default_malloc/free.
  *
  * 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.
+ * the value returned by halide_malloc_alignment().
+ *
+ * The default implementation of halide_malloc_alignment() will be chosen to be
+ * the maximum meaningful alignment for the platform for the purpose of
+ * vector loads and stores for a given architecture. You may override this
+ * function, but bear in mind that it must return the same value for every call.
  */
 //@{
+extern int halide_malloc_alignment();
 extern void *halide_malloc(void *user_context, size_t x);
 extern void halide_free(void *user_context, void *ptr);
 extern void *halide_default_malloc(void *user_context, size_t x);
diff --git a/src/runtime/alignment_128.cpp b/src/runtime/alignment_128.cpp
index 96b13c1341a6..310f6312a646 100644
--- a/src/runtime/alignment_128.cpp
+++ b/src/runtime/alignment_128.cpp
@@ -1,5 +1,5 @@
 #include "runtime_internal.h"
 
-extern "C" WEAK_INLINE int halide_malloc_alignment() {
+extern "C" WEAK int halide_malloc_alignment() {
     return 128;
 }
diff --git a/src/runtime/alignment_32.cpp b/src/runtime/alignment_32.cpp
index e7a62f48ad3d..a30069cded6e 100644
--- a/src/runtime/alignment_32.cpp
+++ b/src/runtime/alignment_32.cpp
@@ -1,5 +1,5 @@
 #include "runtime_internal.h"
 
-extern "C" WEAK_INLINE int halide_malloc_alignment() {
+extern "C" WEAK int halide_malloc_alignment() {
     return 32;
 }
diff --git a/src/runtime/alignment_64.cpp b/src/runtime/alignment_64.cpp
index 1833ff489ba7..e2512b2fd151 100644
--- a/src/runtime/alignment_64.cpp
+++ b/src/runtime/alignment_64.cpp
@@ -1,5 +1,5 @@
 #include "runtime_internal.h"
 
-extern "C" WEAK_INLINE int halide_malloc_alignment() {
+extern "C" WEAK int halide_malloc_alignment() {
     return 64;
 }
diff --git a/src/runtime/posix_allocator.cpp b/src/runtime/posix_allocator.cpp
index 5b39ab1c6531..d9cb739b973d 100644
--- a/src/runtime/posix_allocator.cpp
+++ b/src/runtime/posix_allocator.cpp
@@ -3,6 +3,19 @@
 
 #include "printer.h"
 
+namespace Halide {
+namespace Runtime {
+namespace Internal {
+
+// Read into a global to avoid making a call to halide_malloc_alignment()
+// in every halide_malloc() call (halide_malloc_alignment() is required to
+// return the same value every time).
+WEAK size_t _alignment = (size_t)halide_malloc_alignment();
+
+}  // namespace Internal
+}  // namespace Runtime
+}  // namespace Halide
+
 extern "C" {
 
 extern void *malloc(size_t);
@@ -10,7 +23,7 @@ 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();
+    const size_t alignment = _alignment;
     void *orig = malloc(x + alignment);
     if (orig == nullptr) {
         // Will result in a failed assertion and a call to halide_error
diff --git a/src/runtime/qurt_allocator.cpp b/src/runtime/qurt_allocator.cpp
index 961646d8c91e..67cfabb7d233 100644
--- a/src/runtime/qurt_allocator.cpp
+++ b/src/runtime/qurt_allocator.cpp
@@ -1,7 +1,6 @@
 #include "HalideRuntime.h"
 
 extern "C" {
-
 extern void *malloc(size_t);
 extern void free(void *);
 }
@@ -54,13 +53,18 @@ WEAK __attribute__((destructor)) void halide_allocator_cleanup() {
     }
 }
 
+// Read into a global to avoid making a call to halide_malloc_alignment()
+// in every halide_malloc() call (halide_malloc_alignment() is required to
+// return the same value every time).
+WEAK size_t _alignment = (size_t)halide_malloc_alignment();
+
 }  // namespace Internal
 }  // namespace Runtime
 }  // 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::Runtime::Internal::_alignment;
 
     if (x <= buffer_size) {
         for (int i = 0; i < num_buffers; ++i) {
diff --git a/src/runtime/runtime_api.cpp b/src/runtime/runtime_api.cpp
index e3a360f013e6..8e701c8a0621 100644
--- a/src/runtime/runtime_api.cpp
+++ b/src/runtime/runtime_api.cpp
@@ -114,6 +114,7 @@ extern "C" __attribute__((used)) void *halide_runtime_api_functions[] = {
     (void *)&halide_join_thread,
     (void *)&halide_load_library,
     (void *)&halide_malloc,
+    (void *)&halide_malloc_alignment,
     (void *)&halide_memoization_cache_cleanup,
     (void *)&halide_memoization_cache_evict,
     (void *)&halide_memoization_cache_lookup,
diff --git a/src/runtime/runtime_internal.h b/src/runtime/runtime_internal.h
index a2de721079c8..0be69ec3f747 100644
--- a/src/runtime/runtime_internal.h
+++ b/src/runtime/runtime_internal.h
@@ -195,7 +195,7 @@ struct halide_pseudostack_slot_t {
 WEAK void halide_use_jit_module();
 WEAK void halide_release_jit_module();
 
-WEAK_INLINE int halide_malloc_alignment();
+WEAK int halide_malloc_alignment();
 
 void halide_thread_yield();
 
