From 027b6a2332edd11d40b2e64b876c49519eeda287 Mon Sep 17 00:00:00 2001
From: Martijn Courteaux <courteauxmartijn@gmail.com>
Date: Fri, 1 Mar 2024 23:03:04 +0100
Subject: [PATCH] Add support for setting the default allocator and deallocator
 functions in Halide::Runtime::Buffer.

---
 src/runtime/HalideBuffer.h | 26 ++++++++++++++++++++++----
 1 file changed, 22 insertions(+), 4 deletions(-)

diff --git a/src/runtime/HalideBuffer.h b/src/runtime/HalideBuffer.h
index 4ac2317278bc..db69a0c0bbed 100644
--- a/src/runtime/HalideBuffer.h
+++ b/src/runtime/HalideBuffer.h
@@ -142,8 +142,8 @@ struct AllInts<float, Args...> : std::false_type {};
 template<typename... Args>
 struct AllInts<double, Args...> : std::false_type {};
 
-// A helper to detect if there are any zeros in a container
 namespace Internal {
+// A helper to detect if there are any zeros in a container
 template<typename Container>
 bool any_zero(const Container &c) {
     for (int i : c) {
@@ -153,6 +153,15 @@ bool any_zero(const Container &c) {
     }
     return false;
 }
+
+template<typename U>
+struct DefaultAllocatorFns {};
+
+template<>
+struct DefaultAllocatorFns<void> {
+    void *(*default_allocate_fn)(size_t) = malloc;
+    void (*default_deallocate_fn)(void *) = free;
+};
 }  // namespace Internal
 
 /** A struct acting as a header for allocations owned by the Buffer
@@ -710,7 +719,16 @@ class Buffer {
                       "Can't convert from a Buffer with static dimensionality to a Buffer with different static dimensionality");
     }
 
+    static inline Internal::DefaultAllocatorFns<T> default_allocator_fns = {};
+
 public:
+    static void set_default_allocate_fn(void *(*allocate_fn)(size_t)) {
+        Buffer<>::default_allocator_fns.default_allocate_fn = allocate_fn;
+    }
+    static void set_default_deallocate_fn(void (*deallocate_fn)(void *)) {
+        Buffer<>::default_allocator_fns.default_deallocate_fn = deallocate_fn;
+    }
+
     /** Determine if a Buffer<T, Dims, InClassDimStorage> can be constructed from some other Buffer type.
      * If this can be determined at compile time, fail with a static assert; otherwise
      * return a boolean based on runtime typing. */
@@ -893,7 +911,7 @@ class Buffer {
 
 #if HALIDE_RUNTIME_BUFFER_USE_ALIGNED_ALLOC
         // Only use aligned_alloc() if no custom allocators are specified.
-        if (!allocate_fn && !deallocate_fn) {
+        if (!allocate_fn && !deallocate_fn && Buffer<>::default_allocator_fns.default_allocate_fn == malloc && Buffer<>::default_allocator_fns.default_deallocate_fn == free) {
             // As a practical matter, sizeof(AllocationHeader) is going to be no more than 16 bytes
             // on any supported platform, so we will just overallocate by 'alignment'
             // so that the user storage also starts at an aligned point. This is a bit
@@ -908,10 +926,10 @@ class Buffer {
         // else fall thru
 #endif
         if (!allocate_fn) {
-            allocate_fn = malloc;
+            allocate_fn = Buffer<>::default_allocator_fns.default_allocate_fn;
         }
         if (!deallocate_fn) {
-            deallocate_fn = free;
+            deallocate_fn = Buffer<>::default_allocator_fns.default_deallocate_fn;
         }
 
         static_assert(sizeof(AllocationHeader) <= alignment);