Loading...
Searching...
No Matches
JITModule.h
Go to the documentation of this file.
#ifndef HALIDE_JIT_MODULE_H
#define HALIDE_JIT_MODULE_H
/** \file
* Defines the struct representing lifetime and dependencies of
* a JIT compiled halide pipeline
*/
#include <map>
#include <memory>
#include <vector>
#include "IntrusivePtr.h"
#include "Target.h"
#include "Type.h"
#include "WasmExecutor.h"
#include "runtime/HalideRuntime.h"
namespace llvm {
class Module;
}
namespace Halide {
struct ExternCFunction;
struct JITExtern;
class Module;
struct JITUserContext;
/** A set of custom overrides of runtime functions. These only apply
* when JIT-compiling code. If you are doing AOT compilation, see
* HalideRuntime.h for instructions on how to replace runtime
* functions. */
struct JITHandlers {
/** Set the function called to print messages from the runtime. */
void (*custom_print)(JITUserContext *, const char *){nullptr};
/** A custom malloc and free for halide to use. Malloc should
* return 32-byte aligned chunks of memory, and it should be safe
* for Halide to read slightly out of bounds (up to 8 bytes before
* the start or beyond the end). */
// @{
void *(*custom_malloc)(JITUserContext *, size_t){nullptr};
void (*custom_free)(JITUserContext *, void *){nullptr};
// @}
/** A custom task handler to be called by the parallel for
* loop. It is useful to set this if you want to do some
* additional bookkeeping at the granularity of parallel
* tasks. The default implementation does this:
\code
extern "C" int halide_do_task(JITUserContext *user_context,
int (*f)(void *, int, uint8_t *),
int idx, uint8_t *state) {
return f(user_context, idx, state);
}
\endcode
*
* If you're trying to use a custom parallel runtime, you probably
* don't want to call this. See instead custom_do_par_for.
*/
int (*custom_do_task)(JITUserContext *, int (*)(JITUserContext *, int, uint8_t *), int, uint8_t *){nullptr};
/** A custom parallel for loop launcher. Useful if your app
* already manages a thread pool. The default implementation is
* equivalent to this:
\code
extern "C" int halide_do_par_for(JITUserContext *user_context,
int (*f)(void *, int, uint8_t *),
int min, int extent, uint8_t *state) {
int exit_status = 0;
parallel for (int idx = min; idx < min+extent; idx++) {
int job_status = halide_do_task(user_context, f, idx, state);
if (job_status) exit_status = job_status;
}
return exit_status;
}
\endcode
*
* However, notwithstanding the above example code, if one task
* fails, we may skip over other tasks, and if two tasks return
* different error codes, we may select one arbitrarily to return.
*/
int (*custom_do_par_for)(JITUserContext *, int (*)(JITUserContext *, int, uint8_t *), int, int, uint8_t *){nullptr};
/** The error handler function that be called in the case of
* runtime errors during halide pipelines. */
void (*custom_error)(JITUserContext *, const char *){nullptr};
/** A custom routine to call when tracing is enabled. Call this
* on the output Func of your pipeline. This then sets custom
* routines for the entire pipeline, not just calls to this
* Func. */
int32_t (*custom_trace)(JITUserContext *, const halide_trace_event_t *){nullptr};
/** A method to use for Halide to resolve symbol names dynamically
* in the calling process or library from within the Halide
* runtime. Equivalent to dlsym with a null first argument. */
void *(*custom_get_symbol)(const char *name){nullptr};
/** A method to use for Halide to dynamically load libraries from
* within the runtime. Equivalent to dlopen. Returns a handle to
* the opened library. */
void *(*custom_load_library)(const char *name){nullptr};
/** A method to use for Halide to dynamically find a symbol within
* an opened library. Equivalent to dlsym. Takes a handle
* returned by custom_load_library as the first argument. */
void *(*custom_get_library_symbol)(void *lib, const char *name){nullptr};
/** A custom method for the Halide runtime acquire a cuda
* context. The cuda context is treated as a void * to avoid a
* dependence on the cuda headers. If the create argument is set
* to true, a context should be created if one does not already
* exist. */
int32_t (*custom_cuda_acquire_context)(JITUserContext *user_context, void **cuda_context_ptr, bool create){nullptr};
/** The Halide runtime calls this when it is done with a cuda
* context. The default implementation does nothing. */
int32_t (*custom_cuda_release_context)(JITUserContext *user_context){nullptr};
/** A custom method for the Halide runtime to acquire a cuda
* stream to use. The cuda context and stream are both modelled
* as a void *, to avoid a dependence on the cuda headers. */
int32_t (*custom_cuda_get_stream)(JITUserContext *user_context, void *cuda_context, void **stream_ptr){nullptr};
};
namespace Internal {
struct JITErrorBuffer;
}
/** A context to be passed to Pipeline::realize. Inherit from this to
* pass your own custom context object. Modify the handlers field to
* override runtime functions per-call to realize. */
struct JITUserContext {
Internal::JITErrorBuffer *error_buffer{nullptr};
JITHandlers handlers;
};
namespace Internal {
class JITModuleContents;
struct LoweredFunc;
struct JITModule {
IntrusivePtr<JITModuleContents> jit_module;
struct Symbol {
void *address = nullptr;
Symbol() = default;
explicit Symbol(void *address)
: address(address) {
}
};
JITModule();
JITModule(const Module &m, const LoweredFunc &fn,
const std::vector<JITModule> &dependencies = std::vector<JITModule>());
/** Take a list of JITExterns and generate trampoline functions
* which can be called dynamically via a function pointer that
* takes an array of void *'s for each argument and the return
* value.
*/
static JITModule make_trampolines_module(const Target &target,
const std::map<std::string, JITExtern> &externs,
const std::string &suffix,
const std::vector<JITModule> &deps);
/** The exports map of a JITModule contains all symbols which are
* available to other JITModules which depend on this one. For
* runtime modules, this is all of the symbols exported from the
* runtime. For a JITted Func, it generally only contains the main
* result Func of the compilation, which takes its name directly
* from the Func declaration. One can also make a module which
* contains no code itself but is just an exports maps providing
* arbitrary pointers to functions or global variables to JITted
* code. */
const std::map<std::string, Symbol> &exports() const;
/** A pointer to the raw halide function. Its true type depends
* on the Argument vector passed to CodeGen_LLVM::compile. Image
* parameters become (halide_buffer_t *), and scalar parameters become
* pointers to the appropriate values. The final argument is a
* pointer to the halide_buffer_t defining the output. This will be nullptr for
* a JITModule which has not yet been compiled or one that is not
* a Halide Func compilation at all. */
void *main_function() const;
/** Returns the Symbol structure for the routine documented in
* main_function. Returning a Symbol allows access to the LLVM
* type as well as the address. The address and type will be nullptr
* if the module has not been compiled. */
Symbol entrypoint_symbol() const;
/** Returns the Symbol structure for the argv wrapper routine
* corresponding to the entrypoint. The argv wrapper is callable
* via an array of void * pointers to the arguments for the
* call. Returning a Symbol allows access to the LLVM type as well
* as the address. The address and type will be nullptr if the module
* has not been compiled. */
Symbol argv_entrypoint_symbol() const;
/** A slightly more type-safe wrapper around the raw halide
* module. Takes it arguments as an array of pointers that
* correspond to the arguments to \ref main_function . This will
* be nullptr for a JITModule which has not yet been compiled or one
* that is not a Halide Func compilation at all. */
// @{
typedef int (*argv_wrapper)(const void *const *args);
argv_wrapper argv_function() const;
// @}
/** Add another JITModule to the dependency chain. Dependencies
* are searched to resolve symbols not found in the current
* compilation unit while JITting. */
void add_dependency(JITModule &dep);
/** Registers a single Symbol as available to modules which depend
* on this one. The Symbol structure provides both the address and
* the LLVM type for the function, which allows type safe linkage of
* extenal routines. */
void add_symbol_for_export(const std::string &name, const Symbol &extern_symbol);
/** Registers a single function as available to modules which
* depend on this one. This routine converts the ExternSignature
* info into an LLVM type, which allows type safe linkage of
* external routines. */
void add_extern_for_export(const std::string &name,
const ExternCFunction &extern_c_function);
/** Look up a symbol by name in this module or its dependencies. */
Symbol find_symbol_by_name(const std::string &) const;
/** Take an llvm module and compile it. The requested exports will
be available via the exports method. */
void compile_module(std::unique_ptr<llvm::Module> mod,
const std::string &function_name, const Target &target,
const std::vector<JITModule> &dependencies = std::vector<JITModule>(),
const std::vector<std::string> &requested_exports = std::vector<std::string>());
/** See JITSharedRuntime::memoization_cache_set_size */
void memoization_cache_set_size(int64_t size) const;
/** See JITSharedRuntime::memoization_cache_evict */
void memoization_cache_evict(uint64_t eviction_key) const;
/** See JITSharedRuntime::reuse_device_allocations */
void reuse_device_allocations(bool) const;
/** See JITSharedRuntime::get_num_threads */
int get_num_threads() const;
/** See JITSharedRuntime::set_num_threads */
int set_num_threads(int) const;
/** Return true if compile_module has been called on this module. */
bool compiled() const;
};
class JITSharedRuntime {
public:
// Note only the first llvm::Module passed in here is used. The same shared runtime is used for all JIT.
static std::vector<JITModule> get(llvm::Module *m, const Target &target, bool create = true);
static void populate_jit_handlers(JITUserContext *jit_user_context, const JITHandlers &handlers);
static JITHandlers set_default_handlers(const JITHandlers &handlers);
/** Set the maximum number of bytes used by memoization caching.
* If you are compiling statically, you should include HalideRuntime.h
* and call halide_memoization_cache_set_size() instead.
*/
static void memoization_cache_set_size(int64_t size);
/** Evict all cache entries that were tagged with the given
* eviction_key in the memoize scheduling directive. If you are
* compiling statically, you should include HalideRuntime.h and
* call halide_memoization_cache_evict() instead.
*/
static void memoization_cache_evict(uint64_t eviction_key);
/** Set whether or not Halide may hold onto and reuse device
* allocations to avoid calling expensive device API allocation
* functions. If you are compiling statically, you should include
* HalideRuntime.h and call halide_reuse_device_allocations
* instead. */
static void reuse_device_allocations(bool);
static void release_all();
/** Get the number of threads in the Halide thread pool. Includes the
* calling thread. Meaningless if a custom_do_par_for has been set. */
static int get_num_threads();
/** Set the number of threads to use in the Halide thread pool, inclusive of
* the calling thread. Pass zero to use a reasonable default (typically the
* number of CPUs online). Calling this is meaningless if custom_do_par_for
* has been set. Halide may launch more threads than this if necessary to
* avoid deadlock when using the async scheduling directive. Returns the old
* number. */
static int set_num_threads(int);
};
void *get_symbol_address(const char *s);
struct JITCache {
Target jit_target;
// Arguments for all inputs and outputs
std::vector<Argument> arguments;
std::map<std::string, JITExtern> jit_externs;
JITModule jit_module;
WasmModule wasm_module;
JITCache() = default;
JITCache(Target jit_target,
std::vector<Argument> arguments,
std::map<std::string, JITExtern> jit_externs,
JITModule jit_module,
WasmModule wasm_module);
Target get_compiled_jit_target() const;
int call_jit_code(const void *const *args);
void finish_profiling(JITUserContext *context);
};
struct JITErrorBuffer {
enum { MaxBufSize = 4096 };
char buf[MaxBufSize];
std::atomic<size_t> end{0};
void concat(const char *message);
std::string str() const;
static void handler(JITUserContext *ctx, const char *message);
};
struct JITFuncCallContext {
JITErrorBuffer error_buffer;
JITUserContext *context;
bool custom_error_handler;
JITFuncCallContext(JITUserContext *context, const JITHandlers &pipeline_handlers);
void finalize(int exit_status);
};
} // namespace Internal
} // namespace Halide
#endif
HalideRuntime.h
This file declares the routines used by Halide internally in its runtime.
IntrusivePtr.h
Support classes for reference-counting via intrusive shared pointers.
Target.h
Defines the structure that describes a Halide target.
Type.h
Defines halide types.
WasmExecutor.h
Support for running Halide-compiled Wasm code in-process.
Halide::Internal::JITSharedRuntime
Definition JITModule.h:260
Halide::Internal::JITSharedRuntime::get_num_threads
static int get_num_threads()
Get the number of threads in the Halide thread pool.
Halide::Internal::JITSharedRuntime::memoization_cache_evict
static void memoization_cache_evict(uint64_t eviction_key)
Evict all cache entries that were tagged with the given eviction_key in the memoize scheduling direct...
Halide::Internal::JITSharedRuntime::memoization_cache_set_size
static void memoization_cache_set_size(int64_t size)
Set the maximum number of bytes used by memoization caching.
Halide::Internal::JITSharedRuntime::set_num_threads
static int set_num_threads(int)
Set the number of threads to use in the Halide thread pool, inclusive of the calling thread.
Halide::Internal::JITSharedRuntime::set_default_handlers
static JITHandlers set_default_handlers(const JITHandlers &handlers)
Halide::Internal::JITSharedRuntime::reuse_device_allocations
static void reuse_device_allocations(bool)
Set whether or not Halide may hold onto and reuse device allocations to avoid calling expensive devic...
Halide::Internal::JITSharedRuntime::populate_jit_handlers
static void populate_jit_handlers(JITUserContext *jit_user_context, const JITHandlers &handlers)
Halide::Internal::JITSharedRuntime::release_all
static void release_all()
Halide::Internal::JITSharedRuntime::get
static std::vector< JITModule > get(llvm::Module *m, const Target &target, bool create=true)
Halide::Module
A halide module.
Definition Module.h:142
Halide::Internal
Definition AbstractGenerator.h:24
Halide::Internal::get_symbol_address
void * get_symbol_address(const char *s)
Halide
This file defines the class FunctionDAG, which is our representation of a Halide pipeline,...
Definition AbstractGenerator.h:19
llvm
Definition CodeGen_Internal.h:18
uint64_t
unsigned __INT64_TYPE__ uint64_t
Definition runtime_internal.h:23
int64_t
signed __INT64_TYPE__ int64_t
Definition runtime_internal.h:22
int32_t
signed __INT32_TYPE__ int32_t
Definition runtime_internal.h:24
uint8_t
unsigned __INT8_TYPE__ uint8_t
Definition runtime_internal.h:29
size_t
__SIZE_TYPE__ size_t
Definition runtime_internal.h:31
Halide::ExternCFunction
Definition Pipeline.h:557
Halide::Internal::IntrusivePtr
Intrusive shared pointers have a reference count (a RefCount object) stored in the class itself.
Definition IntrusivePtr.h:71
Halide::Internal::JITCache::call_jit_code
int call_jit_code(const void *const *args)
Halide::Internal::JITCache::jit_module
JITModule jit_module
Definition JITModule.h:309
Halide::Internal::JITCache::get_compiled_jit_target
Target get_compiled_jit_target() const
Halide::Internal::JITCache::wasm_module
WasmModule wasm_module
Definition JITModule.h:310
Halide::Internal::JITCache::jit_externs
std::map< std::string, JITExtern > jit_externs
Definition JITModule.h:308
Halide::Internal::JITCache::arguments
std::vector< Argument > arguments
Definition JITModule.h:307
Halide::Internal::JITCache::jit_target
Target jit_target
Definition JITModule.h:305
Halide::Internal::JITCache::JITCache
JITCache(Target jit_target, std::vector< Argument > arguments, std::map< std::string, JITExtern > jit_externs, JITModule jit_module, WasmModule wasm_module)
Halide::Internal::JITCache::finish_profiling
void finish_profiling(JITUserContext *context)
Halide::Internal::JITCache::JITCache
JITCache()=default
Halide::Internal::JITErrorBuffer
Definition JITModule.h:326
Halide::Internal::JITErrorBuffer::str
std::string str() const
Halide::Internal::JITErrorBuffer::concat
void concat(const char *message)
Halide::Internal::JITErrorBuffer::MaxBufSize
@ MaxBufSize
Definition JITModule.h:327
Halide::Internal::JITErrorBuffer::buf
char buf[MaxBufSize]
Definition JITModule.h:328
Halide::Internal::JITErrorBuffer::handler
static void handler(JITUserContext *ctx, const char *message)
Halide::Internal::JITErrorBuffer::end
std::atomic< size_t > end
Definition JITModule.h:329
Halide::Internal::JITFuncCallContext::finalize
void finalize(int exit_status)
Halide::Internal::JITFuncCallContext::custom_error_handler
bool custom_error_handler
Definition JITModule.h:341
Halide::Internal::JITFuncCallContext::context
JITUserContext * context
Definition JITModule.h:340
Halide::Internal::JITFuncCallContext::error_buffer
JITErrorBuffer error_buffer
Definition JITModule.h:339
Halide::Internal::JITFuncCallContext::JITFuncCallContext
JITFuncCallContext(JITUserContext *context, const JITHandlers &pipeline_handlers)
Halide::Internal::JITModule::Symbol
Definition JITModule.h:149
Halide::Internal::JITModule::Symbol::address
void * address
Definition JITModule.h:150
Halide::Internal::JITModule::Symbol::Symbol
Symbol()=default
Halide::Internal::JITModule::Symbol::Symbol
Symbol(void *address)
Definition JITModule.h:152
Halide::Internal::JITModule
Definition JITModule.h:146
Halide::Internal::JITModule::memoization_cache_evict
void memoization_cache_evict(uint64_t eviction_key) const
See JITSharedRuntime::memoization_cache_evict.
Halide::Internal::JITModule::set_num_threads
int set_num_threads(int) const
See JITSharedRuntime::set_num_threads.
Halide::Internal::JITModule::memoization_cache_set_size
void memoization_cache_set_size(int64_t size) const
See JITSharedRuntime::memoization_cache_set_size.
Halide::Internal::JITModule::add_symbol_for_export
void add_symbol_for_export(const std::string &name, const Symbol &extern_symbol)
Registers a single Symbol as available to modules which depend on this one.
Halide::Internal::JITModule::JITModule
JITModule()
Halide::Internal::JITModule::compile_module
void compile_module(std::unique_ptr< llvm::Module > mod, const std::string &function_name, const Target &target, const std::vector< JITModule > &dependencies=std::vector< JITModule >(), const std::vector< std::string > &requested_exports=std::vector< std::string >())
Take an llvm module and compile it.
Halide::Internal::JITModule::get_num_threads
int get_num_threads() const
See JITSharedRuntime::get_num_threads.
Halide::Internal::JITModule::add_extern_for_export
void add_extern_for_export(const std::string &name, const ExternCFunction &extern_c_function)
Registers a single function as available to modules which depend on this one.
Halide::Internal::JITModule::exports
const std::map< std::string, Symbol > & exports() const
The exports map of a JITModule contains all symbols which are available to other JITModules which dep...
Halide::Internal::JITModule::reuse_device_allocations
void reuse_device_allocations(bool) const
See JITSharedRuntime::reuse_device_allocations.
Halide::Internal::JITModule::add_dependency
void add_dependency(JITModule &dep)
Add another JITModule to the dependency chain.
Halide::Internal::JITModule::find_symbol_by_name
Symbol find_symbol_by_name(const std::string &) const
Look up a symbol by name in this module or its dependencies.
Halide::Internal::JITModule::make_trampolines_module
static JITModule make_trampolines_module(const Target &target, const std::map< std::string, JITExtern > &externs, const std::string &suffix, const std::vector< JITModule > &deps)
Take a list of JITExterns and generate trampoline functions which can be called dynamically via a fun...
Halide::Internal::JITModule::main_function
void * main_function() const
A pointer to the raw halide function.
Halide::Internal::JITModule::argv_entrypoint_symbol
Symbol argv_entrypoint_symbol() const
Returns the Symbol structure for the argv wrapper routine corresponding to the entrypoint.
Halide::Internal::JITModule::compiled
bool compiled() const
Return true if compile_module has been called on this module.
Halide::Internal::JITModule::argv_wrapper
int(* argv_wrapper)(const void *const *args)
A slightly more type-safe wrapper around the raw halide module.
Definition JITModule.h:211
Halide::Internal::JITModule::entrypoint_symbol
Symbol entrypoint_symbol() const
Returns the Symbol structure for the routine documented in main_function.
Halide::Internal::JITModule::argv_function
argv_wrapper argv_function() const
Halide::Internal::JITModule::jit_module
IntrusivePtr< JITModuleContents > jit_module
Definition JITModule.h:147
Halide::Internal::JITModule::JITModule
JITModule(const Module &m, const LoweredFunc &fn, const std::vector< JITModule > &dependencies=std::vector< JITModule >())
Halide::Internal::LoweredFunc
Definition of a lowered function.
Definition Module.h:101
Halide::Internal::WasmModule
Handle to compiled wasm code which can be called later.
Definition WasmExecutor.h:32
Halide::JITExtern
Definition Pipeline.h:582
Halide::JITHandlers
A set of custom overrides of runtime functions.
Definition JITModule.h:35
Halide::JITHandlers::custom_do_par_for
int(* custom_do_par_for)(JITUserContext *, int(*)(JITUserContext *, int, uint8_t *), int, int, uint8_t *)
A custom parallel for loop launcher.
Definition JITModule.h:85
Halide::JITHandlers::custom_cuda_acquire_context
int32_t(* custom_cuda_acquire_context)(JITUserContext *user_context, void **cuda_context_ptr, bool create)
A custom method for the Halide runtime acquire a cuda context.
Definition JITModule.h:117
Halide::JITHandlers::custom_cuda_get_stream
int32_t(* custom_cuda_get_stream)(JITUserContext *user_context, void *cuda_context, void **stream_ptr)
A custom method for the Halide runtime to acquire a cuda stream to use.
Definition JITModule.h:126
Halide::JITHandlers::custom_error
void(* custom_error)(JITUserContext *, const char *)
The error handler function that be called in the case of runtime errors during halide pipelines.
Definition JITModule.h:89
Halide::JITHandlers::custom_cuda_release_context
int32_t(* custom_cuda_release_context)(JITUserContext *user_context)
The Halide runtime calls this when it is done with a cuda context.
Definition JITModule.h:121
Halide::JITHandlers::custom_free
void(* custom_free)(JITUserContext *, void *)
Definition JITModule.h:45
Halide::JITHandlers::custom_trace
int32_t(* custom_trace)(JITUserContext *, const halide_trace_event_t *)
A custom routine to call when tracing is enabled.
Definition JITModule.h:95
Halide::JITHandlers::custom_do_task
int(* custom_do_task)(JITUserContext *, int(*)(JITUserContext *, int, uint8_t *), int, uint8_t *)
A custom task handler to be called by the parallel for loop.
Definition JITModule.h:63
Halide::JITHandlers::custom_print
void(* custom_print)(JITUserContext *, const char *)
Set the function called to print messages from the runtime.
Definition JITModule.h:37
Halide::JITUserContext
A context to be passed to Pipeline::realize.
Definition JITModule.h:136
Halide::JITUserContext::handlers
JITHandlers handlers
Definition JITModule.h:138
Halide::JITUserContext::error_buffer
Internal::JITErrorBuffer * error_buffer
Definition JITModule.h:137
Halide::Target
A struct representing a target machine and os to generate code for.
Definition Target.h:19
halide_trace_event_t
Definition HalideRuntime.h:589