1 #include "HalideRuntime.h"
2 #include "runtime_internal.h"
3 #include "scoped_mutex_lock.h"
4 
5 namespace Halide {
6 namespace Runtime {
7 namespace Internal {
8 
9 WEAK bool halide_reuse_device_allocations_flag = true;
10 
11 WEAK halide_mutex allocation_pools_lock;
12 WEAK halide_device_allocation_pool *device_allocation_pools = NULL;
13 
14 }  // namespace Internal
15 }  // namespace Runtime
16 }  // namespace Halide
17 
18 extern "C" {
19 
halide_reuse_device_allocations(void * user_context,bool flag)20 WEAK int halide_reuse_device_allocations(void *user_context, bool flag) {
21     halide_reuse_device_allocations_flag = flag;
22 
23     int err = 0;
24     if (!flag) {
25         ScopedMutexLock lock(&allocation_pools_lock);
26         for (halide_device_allocation_pool *p = device_allocation_pools; p != NULL; p = p->next) {
27             int ret = p->release_unused(user_context);
28             if (ret) {
29                 err = ret;
30             }
31         }
32     }
33     return err;
34 }
35 
36 /** Determines whether on device_free the memory is returned
37  * immediately to the device API, or placed on a free list for future
38  * use. Override and switch based on the user_context for
39  * finer-grained control. By default just returns the value most
40  * recently set by the method above. */
halide_can_reuse_device_allocations(void * user_context)41 WEAK bool halide_can_reuse_device_allocations(void *user_context) {
42     return halide_reuse_device_allocations_flag;
43 }
44 
halide_register_device_allocation_pool(struct halide_device_allocation_pool * pool)45 WEAK void halide_register_device_allocation_pool(struct halide_device_allocation_pool *pool) {
46     ScopedMutexLock lock(&allocation_pools_lock);
47     pool->next = device_allocation_pools;
48     device_allocation_pools = pool;
49 }
50 }
51