more per-module compiler options (#37041)

NEWS.md

@@ -8,6 +8,8 @@ New language features
   `(Foo{T} where T)(x) = ...`.
 * `<--` and `<-->` are now available as infix operators, with the same precedence
   and associativity as other arrow-like operators ([#36666]).
+* Compilation and type inference can now be enabled or disabled at the module level
+  using the experimental macro `Base.Experimental.@compiler_options` ([#37041]).
 
 Language changes
 ----------------

base/compiler/typeinfer.jl

@@ -513,6 +513,9 @@ function typeinf_edge(interp::AbstractInterpreter, method::Method, @nospecialize
             return code.rettype, mi
         end
     end
+    if ccall(:jl_get_module_infer, Cint, (Any,), method.module) == 0
+        return Any, nothing
+    end
     if !caller.cached && caller.parent === nothing
         # this caller exists to return to the user
         # (if we asked resolve_call_cyle, it might instead detect that there is a cycle that it can't merge)
@@ -617,6 +620,9 @@ function typeinf_ext(interp::AbstractInterpreter, mi::MethodInstance)
             end
         end
     end
+    if ccall(:jl_get_module_infer, Cint, (Any,), method.module) == 0
+        return retrieve_code_info(mi)
+    end
     lock_mi_inference(interp, mi)
     frame = InferenceState(InferenceResult(mi), #=cached=#true, interp)
     frame === nothing && return nothing

base/experimental.jl

@@ -120,6 +120,46 @@ macro optlevel(n::Int)
     return Expr(:meta, :optlevel, n)
 end
 
+"""
+    Experimental.@compiler_options optimize={0,1,2,3} compile={yes,no,all,min} infer={yes,no}
+
+Set compiler options for code in the enclosing module. Options correspond directly to
+command-line options with the same name, where applicable. The following options
+are currently supported:
+
+  * `optimize`: Set optimization level.
+  * `compile`: Toggle native code compilation. Currently only `min` is supported, which
+    requests the minimum possible amount of compilation.
+  * `infer`: Enable or disable type inference. If disabled, implies [`@nospecialize`](@ref).
+"""
+macro compiler_options(args...)
+    opts = Expr(:block)
+    for ex in args
+        if isa(ex, Expr) && ex.head === :(=) && length(ex.args) == 2
+            if ex.args[1] === :optimize
+                push!(opts.args, Expr(:meta, :optlevel, ex.args[2]::Int))
+            elseif ex.args[1] === :compile
+                a = ex.args[2]
+                a = #a === :no  ? 0 :
+                    #a === :yes ? 1 :
+                    #a === :all ? 2 :
+                    a === :min ? 3 : error("invalid argument to \"compile\" option")
+                push!(opts.args, Expr(:meta, :compile, a))
+            elseif ex.args[1] === :infer
+                a = ex.args[2]
+                a = a === false || a === :no  ? 0 :
+                    a === true  || a === :yes ? 1 : error("invalid argument to \"infer\" option")
+                push!(opts.args, Expr(:meta, :infer, a))
+            else
+                error("unknown option \"$(ex.args[1])\"")
+            end
+        else
+            error("invalid option syntax")
+        end
+    end
+    return opts
+end
+
 # UI features for errors
 
 """

src/ast.c

@@ -64,6 +64,7 @@ jl_sym_t *coverageeffect_sym; jl_sym_t *escape_sym;
 jl_sym_t *aliasscope_sym; jl_sym_t *popaliasscope_sym;
 jl_sym_t *optlevel_sym; jl_sym_t *thismodule_sym;
 jl_sym_t *atom_sym; jl_sym_t *statement_sym; jl_sym_t *all_sym;
+jl_sym_t *compile_sym; jl_sym_t *infer_sym;
 
 static uint8_t flisp_system_image[] = {
 #include <julia_flisp.boot.inc>
@@ -382,6 +383,8 @@ void jl_init_common_symbols(void)
     nospecialize_sym = jl_symbol("nospecialize");
     specialize_sym = jl_symbol("specialize");
     optlevel_sym = jl_symbol("optlevel");
+    compile_sym = jl_symbol("compile");
+    infer_sym = jl_symbol("infer");
     macrocall_sym = jl_symbol("macrocall");
     escape_sym = jl_symbol("escape");
     hygienicscope_sym = jl_symbol("hygienic-scope");

src/dump.c

@@ -375,7 +375,9 @@ static void jl_serialize_module(jl_serializer_state *s, jl_module_t *m)
     write_uint64(s->s, m->build_id);
     write_int32(s->s, m->counter);
     write_int32(s->s, m->nospecialize);
-    write_int32(s->s, m->optlevel);
+    write_uint8(s->s, m->optlevel);
+    write_uint8(s->s, m->compile);
+    write_uint8(s->s, m->infer);
 }
 
 static void jl_serialize_value_(jl_serializer_state *s, jl_value_t *v, int as_literal) JL_GC_DISABLED
@@ -1528,7 +1530,9 @@ static jl_value_t *jl_deserialize_value_module(jl_serializer_state *s) JL_GC_DIS
     m->build_id = read_uint64(s->s);
     m->counter = read_int32(s->s);
     m->nospecialize = read_int32(s->s);
-    m->optlevel = read_int32(s->s);
+    m->optlevel = read_int8(s->s);
+    m->compile = read_int8(s->s);
+    m->infer = read_int8(s->s);
     m->primary_world = jl_world_counter;
     return (jl_value_t*)m;
 }

src/gf.c

@@ -1848,9 +1848,18 @@ jl_code_instance_t *jl_compile_method_internal(jl_method_instance_t *mi, size_t
     jl_code_instance_t *codeinst = jl_method_compiled(mi, world);
     if (codeinst)
         return codeinst;
-
-    if (jl_options.compile_enabled == JL_OPTIONS_COMPILE_OFF ||
-        jl_options.compile_enabled == JL_OPTIONS_COMPILE_MIN) {
+    int compile_option = jl_options.compile_enabled;
+    jl_method_t *def = mi->def.method;
+    // disabling compilation per-module can override global setting
+    if (jl_is_method(def)) {
+        int mod_setting = jl_get_module_compile(((jl_method_t*)def)->module);
+        if (mod_setting == JL_OPTIONS_COMPILE_OFF ||
+            mod_setting == JL_OPTIONS_COMPILE_MIN)
+            compile_option = ((jl_method_t*)def)->module->compile;
+    }
+
+    if (compile_option == JL_OPTIONS_COMPILE_OFF ||
+        compile_option == JL_OPTIONS_COMPILE_MIN) {
         // copy fptr from the template method definition
         jl_method_t *def = mi->def.method;
         if (jl_is_method(def) && def->unspecialized) {
@@ -1876,7 +1885,7 @@ jl_code_instance_t *jl_compile_method_internal(jl_method_instance_t *mi, size_t
             jl_mi_cache_insert(mi, codeinst);
             return codeinst;
         }
-        if (jl_options.compile_enabled == JL_OPTIONS_COMPILE_OFF) {
+        if (compile_option == JL_OPTIONS_COMPILE_OFF) {
             jl_printf(JL_STDERR, "code missing for ");
             jl_static_show(JL_STDERR, (jl_value_t*)mi);
             jl_printf(JL_STDERR, " : sysimg may not have been built with --compile=all\n");

src/interpreter.c

@@ -536,10 +536,22 @@ static jl_value_t *eval_body(jl_array_t *stmts, interpreter_state *s, size_t ip,
                     if (jl_expr_nargs(stmt) == 1 && jl_exprarg(stmt, 0) == (jl_value_t*)specialize_sym) {
                         jl_set_module_nospecialize(s->module, 0);
                     }
-                    if (jl_expr_nargs(stmt) == 2 && jl_exprarg(stmt, 0) == (jl_value_t*)optlevel_sym) {
-                        if (jl_is_long(jl_exprarg(stmt, 1))) {
-                            int n = jl_unbox_long(jl_exprarg(stmt, 1));
-                            jl_set_module_optlevel(s->module, n);
+                    if (jl_expr_nargs(stmt) == 2) {
+                        if (jl_exprarg(stmt, 0) == (jl_value_t*)optlevel_sym) {
+                            if (jl_is_long(jl_exprarg(stmt, 1))) {
+                                int n = jl_unbox_long(jl_exprarg(stmt, 1));
+                                jl_set_module_optlevel(s->module, n);
+                            }
+                        }
+                        else if (jl_exprarg(stmt, 0) == (jl_value_t*)compile_sym) {
+                            if (jl_is_long(jl_exprarg(stmt, 1))) {
+                                jl_set_module_compile(s->module, jl_unbox_long(jl_exprarg(stmt, 1)));
+                            }
+                        }
+                        else if (jl_exprarg(stmt, 0) == (jl_value_t*)infer_sym) {
+                            if (jl_is_long(jl_exprarg(stmt, 1))) {
+                                jl_set_module_infer(s->module, jl_unbox_long(jl_exprarg(stmt, 1)));
+                            }
                         }
                     }
                 }

src/julia.h

@@ -513,7 +513,9 @@ typedef struct _jl_module_t {
     size_t primary_world;
     uint32_t counter;
     int32_t nospecialize;  // global bit flags: initialization for new methods
-    int32_t optlevel;
+    int8_t optlevel;
+    int8_t compile;
+    int8_t infer;
     uint8_t istopmod;
     jl_mutex_t lock;
 } jl_module_t;
@@ -1518,6 +1520,10 @@ JL_DLLEXPORT jl_module_t *jl_new_module(jl_sym_t *name);
 JL_DLLEXPORT void jl_set_module_nospecialize(jl_module_t *self, int on);
 JL_DLLEXPORT void jl_set_module_optlevel(jl_module_t *self, int lvl);
 JL_DLLEXPORT int jl_get_module_optlevel(jl_module_t *m);
+JL_DLLEXPORT void jl_set_module_compile(jl_module_t *self, int value);
+JL_DLLEXPORT int jl_get_module_compile(jl_module_t *m);
+JL_DLLEXPORT void jl_set_module_infer(jl_module_t *self, int value);
+JL_DLLEXPORT int jl_get_module_infer(jl_module_t *m);
 // get binding for reading
 JL_DLLEXPORT jl_binding_t *jl_get_binding(jl_module_t *m JL_PROPAGATES_ROOT, jl_sym_t *var);
 JL_DLLEXPORT jl_binding_t *jl_get_binding_or_error(jl_module_t *m, jl_sym_t *var);

src/julia_internal.h

@@ -1163,7 +1163,8 @@ extern jl_sym_t *colon_sym; extern jl_sym_t *hygienicscope_sym;
 extern jl_sym_t *throw_undef_if_not_sym; extern jl_sym_t *getfield_undefref_sym;
 extern jl_sym_t *gc_preserve_begin_sym; extern jl_sym_t *gc_preserve_end_sym;
 extern jl_sym_t *coverageeffect_sym; extern jl_sym_t *escape_sym;
-extern jl_sym_t *optlevel_sym;
+extern jl_sym_t *optlevel_sym; extern jl_sym_t *compile_sym;
+extern jl_sym_t *infer_sym;
 extern jl_sym_t *atom_sym; extern jl_sym_t *statement_sym; extern jl_sym_t *all_sym;
 
 struct _jl_sysimg_fptrs_t;

src/module.c

@@ -36,6 +36,8 @@ JL_DLLEXPORT jl_module_t *jl_new_module(jl_sym_t *name)
     m->counter = 1;
     m->nospecialize = 0;
     m->optlevel = -1;
+    m->compile = -1;
+    m->infer = -1;
     JL_MUTEX_INIT(&m->lock);
     htable_new(&m->bindings, 0);
     arraylist_new(&m->usings, 0);
@@ -88,6 +90,39 @@ JL_DLLEXPORT int jl_get_module_optlevel(jl_module_t *m)
     return lvl;
 }
 
+JL_DLLEXPORT void jl_set_module_compile(jl_module_t *self, int value)
+{
+    self->compile = value;
+}
+
+JL_DLLEXPORT int jl_get_module_compile(jl_module_t *m)
+{
+    int value = m->compile;
+    while (value == -1 && m->parent != m && m != jl_base_module) {
+        m = m->parent;
+        value = m->compile;
+    }
+    return value;
+}
+
+JL_DLLEXPORT void jl_set_module_infer(jl_module_t *self, int value)
+{
+    self->infer = value;
+    // no reason to specialize if inference is off
+    if (!value)
+        jl_set_module_nospecialize(self, 1);
+}
+
+JL_DLLEXPORT int jl_get_module_infer(jl_module_t *m)
+{
+    int value = m->infer;
+    while (value == -1 && m->parent != m && m != jl_base_module) {
+        m = m->parent;
+        value = m->infer;
+    }
+    return value;
+}
+
 JL_DLLEXPORT void jl_set_istopmod(jl_module_t *self, uint8_t isprimary)
 {
     self->istopmod = 1;

src/precompile.c

@@ -46,13 +46,17 @@ void jl_write_compiler_output(void)
         jl_value_t *f = jl_get_global((jl_module_t*)m, jl_symbol("__init__"));
         if (f) {
             jl_array_ptr_1d_push(jl_module_init_order, m);
-            // TODO: this would be better handled if moved entirely to jl_precompile
-            // since it's a slightly duplication of effort
-            jl_value_t *tt = jl_is_type(f) ? (jl_value_t*)jl_wrap_Type(f) : jl_typeof(f);
-            JL_GC_PUSH1(&tt);
-            tt = (jl_value_t*)jl_apply_tuple_type_v(&tt, 1);
-            jl_compile_hint((jl_tupletype_t*)tt);
-            JL_GC_POP();
+            int setting = jl_get_module_compile((jl_module_t*)m);
+            if (setting != JL_OPTIONS_COMPILE_OFF &&
+                setting != JL_OPTIONS_COMPILE_MIN) {
+                // TODO: this would be better handled if moved entirely to jl_precompile
+                // since it's a slightly duplication of effort
+                jl_value_t *tt = jl_is_type(f) ? (jl_value_t*)jl_wrap_Type(f) : jl_typeof(f);
+                JL_GC_PUSH1(&tt);
+                tt = (jl_value_t*)jl_apply_tuple_type_v(&tt, 1);
+                jl_compile_hint((jl_tupletype_t*)tt);
+                JL_GC_POP();
+            }
         }
     }
 

src/toplevel.c

@@ -816,7 +816,9 @@ jl_value_t *jl_toplevel_eval_flex(jl_module_t *JL_NONNULL m, jl_value_t *e, int
     jl_value_t *result;
     if (has_intrinsics || (!has_defs && fast && has_loops &&
                            jl_options.compile_enabled != JL_OPTIONS_COMPILE_OFF &&
-                           jl_options.compile_enabled != JL_OPTIONS_COMPILE_MIN)) {
+                           jl_options.compile_enabled != JL_OPTIONS_COMPILE_MIN &&
+                           jl_get_module_compile(m) != JL_OPTIONS_COMPILE_OFF &&
+                           jl_get_module_compile(m) != JL_OPTIONS_COMPILE_MIN)) {
         // use codegen
         mfunc = method_instance_for_thunk(thk, m);
         jl_resolve_globals_in_ir((jl_array_t*)thk->code, m, NULL, 0);
@@ -826,7 +828,7 @@ jl_value_t *jl_toplevel_eval_flex(jl_module_t *JL_NONNULL m, jl_value_t *e, int
         // helps in common cases.
         size_t world = jl_world_counter;
         ptls->world_age = world;
-        if (!has_defs) {
+        if (!has_defs && jl_get_module_infer(m) != 0) {
             (void)jl_type_infer(mfunc, world, 0);
         }
         result = jl_invoke(/*func*/NULL, /*args*/NULL, /*nargs*/0, mfunc);