Add codegen hooks to override code generation behavior.

Author: maleadt

base/reflection.jl

@@ -493,6 +493,18 @@ function uncompressed_ast(m::Method, s::CodeInfo)
     return s
 end
 
+# this type mirrors jl_cghooks_t (documented in julia.h)
+immutable CodegenHooks
+    module_setup::Ptr{Void}
+    module_activation::Ptr{Void}
+    raise_exception::Ptr{Void}
+
+    CodegenHooks(;module_setup=nothing, module_activation=nothing, raise_exception=nothing) =
+        new(pointer_from_objref(module_setup),
+            pointer_from_objref(module_activation),
+            pointer_from_objref(raise_exception))
+end
+
 # this type mirrors jl_cgparams_t (documented in julia.h)
 immutable CodegenParams
     cached::Cint
@@ -504,14 +516,18 @@ immutable CodegenParams
     static_alloc::Cint
     dynamic_alloc::Cint
 
+    hooks::CodegenHooks
+
     CodegenParams(;cached::Bool=true,
                    runtime::Bool=true, exceptions::Bool=true,
                    track_allocations::Bool=true, code_coverage::Bool=true,
-                   static_alloc::Bool=true, dynamic_alloc::Bool=true) =
+                   static_alloc::Bool=true, dynamic_alloc::Bool=true,
+                   hooks::CodegenHooks=CodegenHooks()) =
         new(Cint(cached),
             Cint(runtime), Cint(exceptions),
             Cint(track_allocations), Cint(code_coverage),
-            Cint(static_alloc), Cint(dynamic_alloc))
+            Cint(static_alloc), Cint(dynamic_alloc),
+            hooks)
 end
 
 # Printing code representations in IR and assembly

src/cgutils.cpp

@@ -38,6 +38,24 @@ static Value *prepare_call(Value *Callee)
                   __FUNCTION__, (ctx)->file.str().c_str(), *(ctx)->line);
 
 
+// --- hook checks ---
+
+#define JL_HOOK_TEST(params,hook) ((params)->hooks.hook != jl_nothing)
+
+#define JL_HOOK_CALL(params,hook,argc,...) \
+    _hook_call<argc>((params)->hooks.hook, {{__VA_ARGS__}});
+template<int N>
+static inline void _hook_call(jl_value_t *hook, std::array<jl_value_t*,N> args) {
+    jl_value_t **argv;
+    JL_GC_PUSHARGS(argv, N+1);
+    argv[0] = hook;
+    for (int i = 0; i < N; i++)
+        argv[i+1] = args[i];
+    jl_apply(argv, N+1);
+    JL_GC_POP();
+}
+
+
 // --- string constants ---
 static StringMap<GlobalVariable*> stringConstants;
 static Value *stringConstPtr(IRBuilder<> &builder, const std::string &txt)
@@ -672,12 +690,18 @@ static void error_unless(Value *cond, const std::string &msg, jl_codectx_t *ctx)
 static void raise_exception(Value *exc, jl_codectx_t *ctx,
                             BasicBlock *contBB=nullptr)
 {
-    JL_FEAT_REQUIRE(ctx, runtime);
+    if (JL_HOOK_TEST(ctx->params, raise_exception)) {
+        JL_HOOK_CALL(ctx->params, raise_exception, 2,
+                     jl_box_voidpointer(wrap(builder.GetInsertBlock())),
+                     jl_box_voidpointer(wrap(exc)));
+    } else {
+        JL_FEAT_REQUIRE(ctx, runtime);
 #if JL_LLVM_VERSION >= 30700
-    builder.CreateCall(prepare_call(jlthrow_func), { exc });
+        builder.CreateCall(prepare_call(jlthrow_func), { exc });
 #else
-    builder.CreateCall(prepare_call(jlthrow_func), exc);
+        builder.CreateCall(prepare_call(jlthrow_func), exc);
 #endif
+    }
     builder.CreateUnreachable();
     if (!contBB) {
         contBB = BasicBlock::Create(jl_LLVMContext, "after_throw", ctx->f);
@@ -1840,3 +1864,19 @@ static Value *emit_defer_signal(jl_codectx_t *ctx)
     Constant *offset = ConstantInt::getSigned(T_int32, offsetof(jl_tls_states_t, defer_signal) / sizeof(sig_atomic_t));
     return builder.CreateGEP(ptls, ArrayRef<Value*>(offset), "jl_defer_signal");
 }
+
+static int compare_cgparams(const jl_cgparams_t *a, const jl_cgparams_t *b)
+{
+    return (a->cached == b->cached) &&
+           // language features
+           (a->runtime == b->runtime) &&
+           (a->exceptions == b->exceptions) &&
+           (a->track_allocations == b->track_allocations) &&
+           (a->code_coverage == b->code_coverage) &&
+           (a->static_alloc == b->static_alloc) &&
+           (a->dynamic_alloc == b->dynamic_alloc) &&
+           // hooks
+           (a->hooks.module_setup == b->hooks.module_setup) &&
+           (a->hooks.module_activation == b->hooks.module_activation) &&
+           (a->hooks.raise_exception == b->hooks.raise_exception);
+}

src/codegen.cpp

@@ -27,6 +27,7 @@
 #include <sstream>
 #include <fstream>
 #include <map>
+#include <array>
 #include <vector>
 #include <set>
 #include <cstdio>
@@ -834,6 +835,11 @@ jl_llvm_functions_t jl_compile_linfo(jl_method_instance_t *li, jl_code_info_t *s
     JL_TIMING(CODEGEN);
     assert(jl_is_method_instance(li));
     jl_llvm_functions_t decls = {};
+
+    if (params != &jl_default_cgparams /* fast path */ &&
+        !compare_cgparams(params, &jl_default_cgparams) && params->cached)
+        jl_error("functions compiled with custom codegen params mustn't be cached");
+
     // Step 1. See if it is already compiled,
     //         Get the codegen lock,
     //         And get the source
@@ -920,25 +926,30 @@ jl_llvm_functions_t jl_compile_linfo(jl_method_instance_t *li, jl_code_info_t *s
     f = (Function*)decls.functionObject;
     specf = (Function*)decls.specFunctionObject;
 
-    // Step 4. Prepare debug info to receive this function
-    // record that this function name came from this linfo,
-    // so we can build a reverse mapping for debug-info.
-    bool toplevel = li->def == NULL;
-    if (!toplevel) {
-        const DataLayout &DL =
-#if JL_LLVM_VERSION >= 30500
-            m->getDataLayout();
-#else
-            *jl_data_layout;
-#endif
-        // but don't remember toplevel thunks because
-        // they may not be rooted in the gc for the life of the program,
-        // and the runtime doesn't notify us when the code becomes unreachable :(
-        jl_add_linfo_in_flight((specf ? specf : f)->getName(), li, DL);
-    }
 
-    // Step 5. Add the result to the execution engine now
-    jl_finalize_module(m.release(), !toplevel);
+    if (JL_HOOK_TEST(params, module_activation)) {
+        JL_HOOK_CALL(params, module_activation, 1, jl_box_voidpointer(wrap(m.release())));
+    } else {
+        // Step 4. Prepare debug info to receive this function
+        // record that this function name came from this linfo,
+        // so we can build a reverse mapping for debug-info.
+        bool toplevel = li->def == NULL;
+        if (!toplevel) {
+            const DataLayout &DL =
+    #if JL_LLVM_VERSION >= 30500
+                m->getDataLayout();
+    #else
+                *jl_data_layout;
+    #endif
+            // but don't remember toplevel thunks because
+            // they may not be rooted in the gc for the life of the program,
+            // and the runtime doesn't notify us when the code becomes unreachable :(
+            jl_add_linfo_in_flight((specf ? specf : f)->getName(), li, DL);
+        }
+
+        // Step 5. Add the result to the execution engine now
+        jl_finalize_module(m.release(), !toplevel);
+    }
 
     if (li->jlcall_api != 2) {
         // if not inlineable, code won't be needed again
@@ -983,8 +994,13 @@ static Value *getModuleFlag(Module *m, StringRef Key)
 #define getModuleFlag(m,str) m->getModuleFlag(str)
 #endif
 
-static void jl_setup_module(Module *m)
+static void jl_setup_module(Module *m, const jl_cgparams_t *params = &jl_default_cgparams)
 {
+    if (JL_HOOK_TEST(params, module_setup)) {
+        JL_HOOK_CALL(params, module_setup, 1, jl_box_voidpointer(wrap(m)));
+        return;
+    }
+
     // Some linkers (*cough* OS X) don't understand DWARF v4, so we use v2 in
     // imaging mode. The structure of v4 is slightly nicer for debugging JIT
     // code.
@@ -4256,7 +4272,7 @@ static std::unique_ptr<Module> emit_function(jl_method_instance_t *lam, jl_code_
 
     ctx.sret = false;
     Module *M = new Module(ctx.name, jl_LLVMContext);
-    jl_setup_module(M);
+    jl_setup_module(M, params);
     if (specsig) { // assumes !va and !needsparams
         std::vector<Type*> fsig(0);
         Type *rt;

src/jltypes.c

@@ -68,7 +68,7 @@ JL_DLLEXPORT jl_value_t *jl_emptytuple=NULL;
 jl_svec_t *jl_emptysvec;
 jl_value_t *jl_nothing;
 
-const jl_cgparams_t jl_default_cgparams = {1, 1, 1, 1, 1, 1, 1};
+jl_cgparams_t jl_default_cgparams = {1, 1, 1, 1, 1, 1, 1, {NULL, NULL, NULL}};
 
 // --- type properties and predicates ---
 
@@ -3523,6 +3523,10 @@ void jl_init_types(void)
     jl_methtable_type = jl_new_uninitialized_datatype();
     jl_nothing = jl_gc_alloc(ptls, 0, NULL);
 
+    jl_default_cgparams.hooks.module_setup = jl_nothing;
+    jl_default_cgparams.hooks.module_activation = jl_nothing;
+    jl_default_cgparams.hooks.raise_exception = jl_nothing;
+
     jl_emptysvec = (jl_svec_t*)jl_gc_alloc(ptls, sizeof(void*),
                                            jl_simplevector_type);
     jl_svec_set_len_unsafe(jl_emptysvec, 0);

src/julia.h

@@ -1755,6 +1755,25 @@ typedef struct {
 
 // codegen interface ----------------------------------------------------------
 
+typedef struct {
+    // to disable a hook: set to NULL or nothing
+
+    // module setup: prepare a module for code emission (data layout, DWARF version, ...)
+    // parameters: LLVMModuleRef as Ptr{Void}
+    // return value: none
+    jl_value_t *module_setup;
+
+    // module activation: registers debug info, adds module to JIT
+    // parameters: LLVMModuleRef as Ptr{Void}
+    // return value: none
+    jl_value_t *module_activation;
+
+    // exception raising: emit LLVM instructions to raise an exception
+    // parameters: LLVMBasicBlockRef as Ptr{Void}, LLVMValueRef as Ptr{Void}
+    // return value: none
+    jl_value_t *raise_exception;
+} jl_cghooks_t;
+
 typedef struct {
     int cached;             // can the compiler use/populate the compilation cache?
 
@@ -1765,8 +1784,10 @@ typedef struct {
     int code_coverage;      // can we measure coverage (don't if disallowed)?
     int static_alloc;       // is the compiler allowed to allocate statically?
     int dynamic_alloc;      // is the compiler allowed to allocate dynamically (requires runtime)?
+
+    jl_cghooks_t hooks;
 } jl_cgparams_t;
-extern JL_DLLEXPORT const jl_cgparams_t jl_default_cgparams;
+extern JL_DLLEXPORT jl_cgparams_t jl_default_cgparams;
 
 #ifdef __cplusplus
 }