halide 19.0.0__cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl

halide/share/tools/halide_malloc_trace.h

@@ -0,0 +1,102 @@
+#ifndef HALIDE_MALLOC_TRACE_H
+#define HALIDE_MALLOC_TRACE_H
+
+//---------------------------------------------------------------------------
+// The custom trace allocator can be used in an application by calling:
+//
+//   halide_enable_malloc_trace();
+//
+// When the app is run, calls to halide_malloc/free will produce output like:
+//
+//   halide_malloc => [0x9e400, 0xa27ff], # size:17408, align:1K
+//   halide-header => [0x9e390, 0x9e3ff], # size:112, align:16
+//   halide_malloc => [0xa2880, 0xa6e9f], # size:17952, align:128
+//   halide-header => [0xa2820, 0xa287f], # size:96, align:32
+//   halide_free   => [0x9e390, 0x9e3ff], # size:112, align:16
+//   halide_free   => [0xa2820, 0xa287f], # size:96, align:32
+//
+//---------------------------------------------------------------------------
+
+#include <cstdlib>
+#include <iostream>
+#include <memory>
+
+namespace Halide {
+namespace Tools {
+
+static inline void print_meminfoalign(intptr_t val) {
+    intptr_t align_chk = 1024 * 1024;
+    while (align_chk > 0) {
+        if ((val & (align_chk - 1)) == 0) {
+            char aunit = ' ';
+            if (align_chk >= 1024) {
+                align_chk >>= 10;
+                aunit = 'K';
+            }
+            if (align_chk >= 1024) {
+                align_chk >>= 10;
+                aunit = 'M';
+            }
+            std::cout << "align:" << align_chk;
+            if (aunit != ' ') {
+                std::cout << aunit;
+            }
+            break;
+        }
+        align_chk >>= 1;
+    }
+}
+
+void *halide_malloc_trace(void *user_context, size_t x) {
+    // Halide requires halide_malloc to allocate memory that can be
+    // read 8 bytes before the start to store the original pointer.
+    // Additionally, we also need to align it to the natural vector
+    // width.
+    void *orig = malloc(x + 128);
+    if (orig == nullptr) {
+        // Will result in a failed assertion and a call to halide_error
+        return nullptr;
+    }
+    // Round up to next multiple of 128.
+    void *ptr = (void *)((((size_t)orig + 128) >> 7) << 7);
+    ((void **)ptr)[-1] = orig;
+
+    void *headend = (orig == ptr) ? orig : (char *)ptr - 1;
+    std::cout << "halide_malloc => [0x" << std::hex
+              << (intptr_t)ptr << ", 0x"
+              << (intptr_t)ptr + x - 1 << std::dec
+              << "], # size:"
+              << (intptr_t)x << ", ";
+    print_meminfoalign((intptr_t)ptr);
+    std::cout << std::endl;
+
+    std::cout << "halide-header => [0x" << std::hex
+              << (intptr_t)orig << ", 0x"
+              << (intptr_t)headend << std::dec
+              << "], # size:"
+              << (intptr_t)ptr - (intptr_t)orig << ", ";
+    print_meminfoalign((intptr_t)orig);
+    std::cout << std::endl;
+    return ptr;
+}
+
+void halide_free_trace(void *user_context, void *ptr) {
+    std::cout << "halide_free => [0x" << std::hex
+              << (intptr_t)((void **)ptr)[-1] << ", 0x"
+              << (intptr_t)ptr - 1 << std::dec
+              << "], # size:"
+              << (intptr_t)ptr - (intptr_t)((void **)ptr)[-1] << ", ";
+    print_meminfoalign((intptr_t)((void **)ptr)[-1]);
+    std::cout << std::endl;
+    free(((void **)ptr)[-1]);
+}
+
+void halide_enable_malloc_trace(void) {
+    halide_set_custom_malloc(halide_malloc_trace);
+    halide_set_custom_free(halide_free_trace);
+}
+
+}  // namespace Tools
+}  // namespace Halide
+
+#endif  // HALIDE_MALLOC_TRACE_H

halide/share/tools/halide_thread_pool.h

@@ -0,0 +1,161 @@
+#ifndef HALIDE_THREAD_POOL_H
+#define HALIDE_THREAD_POOL_H
+
+#include <cassert>
+#include <condition_variable>
+#include <functional>
+#include <future>
+#include <mutex>
+#include <queue>
+#include <thread>
+#include <utility>
+
+#ifdef _MSC_VER
+#else
+#include <unistd.h>
+#endif
+
+/** \file
+ * Define a simple thread pool utility that is modeled on the api of
+ * std::async(); since implementation details of std::async
+ * can vary considerably, with no control over thread spawning, this class
+ * allows us to use the same model but with precise control over thread usage.
+ *
+ * A ThreadPool is created with a specific number of threads, which will never
+ * vary over the life of the ThreadPool. (If created without a specific number
+ * of threads, it will attempt to use threads == number-of-cores.)
+ *
+ * Each async request will go into a queue, and will be serviced by the next
+ * available thread from the pool.
+ *
+ * The ThreadPool's dtor will block until all currently-executing tasks
+ * to finish (but won't schedule any more).
+ *
+ * Note that this is a fairly simpleminded ThreadPool, meant for tasks
+ * that are fairly coarse (e.g. different tasks in a test); it is specifically
+ * *not* intended to be the underlying implementation for Halide runtime threads
+ */
+namespace Halide {
+namespace Tools {
+
+template<typename T>
+class ThreadPool {
+    struct Job {
+        std::function<T()> func;
+        std::promise<T> result;
+
+        void run_unlocked(std::unique_lock<std::mutex> &unique_lock);
+    };
+
+    // all fields are protected by this mutex.
+    std::mutex mutex;
+
+    // Queue of Jobs.
+    std::queue<Job> jobs;
+
+    // Broadcast whenever items are added to the Job queue.
+    std::condition_variable wakeup_threads;
+
+    // Keep track of threads so they can be joined at shutdown
+    std::vector<std::thread> threads;
+
+    // True if the pool is shutting down.
+    bool shutting_down{false};
+
+    void worker_thread() {
+        std::unique_lock<std::mutex> unique_lock(mutex);
+        while (!shutting_down) {
+            if (jobs.empty()) {
+                // There are no jobs pending. Wait until more jobs are enqueued.
+                wakeup_threads.wait(unique_lock);
+            } else {
+                // Grab the next job.
+                Job cur_job = std::move(jobs.front());
+                jobs.pop();
+                cur_job.run_unlocked(unique_lock);
+            }
+        }
+    }
+
+public:
+    static size_t num_processors_online() {
+#ifdef _WIN32
+        char *num_cores = getenv("NUMBER_OF_PROCESSORS");
+        return num_cores ? atoi(num_cores) : 8;
+#else
+        return sysconf(_SC_NPROCESSORS_ONLN);
+#endif
+    }
+
+    // Default to number of available cores if not specified otherwise
+    ThreadPool(size_t desired_num_threads = num_processors_online()) {
+        // This file doesn't depend on anything else in libHalide, so
+        // we'll use assert, not internal_assert.
+        assert(desired_num_threads > 0);
+
+        std::lock_guard<std::mutex> lock(mutex);
+
+        // Create all the threads.
+        for (size_t i = 0; i < desired_num_threads; ++i) {
+            threads.emplace_back([this] { worker_thread(); });
+        }
+    }
+
+    ~ThreadPool() {
+        // Wake everyone up and tell them the party's over and it's time to go home
+        {
+            std::lock_guard<std::mutex> lock(mutex);
+            shutting_down = true;
+            wakeup_threads.notify_all();
+        }
+
+        // Wait until they leave
+        for (auto &t : threads) {
+            t.join();
+        }
+    }
+
+    template<typename Func, typename... Args>
+    std::future<T> async(Func func, Args... args) {
+        std::lock_guard<std::mutex> lock(mutex);
+
+        Job job;
+        // Don't use std::forward here: we never want args passed by reference,
+        // since they will be accessed from an arbitrary thread.
+        //
+        // Some versions of GCC won't allow capturing variadic arguments in a lambda;
+        //
+        //     job.func = [func, args...]() -> T { return func(args...); };  // Nope, sorry
+        //
+        // fortunately, we can use std::bind() to accomplish the same thing.
+        job.func = std::bind(func, args...);
+        jobs.emplace(std::move(job));
+        std::future<T> result = jobs.back().result.get_future();
+
+        // Wake up our threads.
+        wakeup_threads.notify_all();
+
+        return result;
+    }
+};
+
+template<typename T>
+inline void ThreadPool<T>::Job::run_unlocked(std::unique_lock<std::mutex> &unique_lock) {
+    unique_lock.unlock();
+    T r = func();
+    unique_lock.lock();
+    result.set_value(std::move(r));
+}
+
+template<>
+inline void ThreadPool<void>::Job::run_unlocked(std::unique_lock<std::mutex> &unique_lock) {
+    unique_lock.unlock();
+    func();
+    unique_lock.lock();
+    result.set_value();
+}
+
+}  // namespace Tools
+}  // namespace Halide
+
+#endif  // HALIDE_THREAD_POOL_H