Add --target-sanitize=(memory|thread|address) (#59035)

Adds three hidden options that can enable the Memory / Thread /
AddressSanitizer LLVM passes even when Julia is not built with any
sanitizers enabled. In conjunction with #59033 this lets us build a
ThreadSanitizer version of Julia in only a few minutes rather than
hours.

- Adds the `--target-sanitize=(memory|thread|address)` flags to enable
those LLVM passes when compiling code ahead of time with `--output-*`.
- JITted code will enable the passes for a sanitizer only if the
runtime is built with it.
- Adds the corresponding options to `CodegenParams`.

Author: xal-0

base/options.jl

@@ -70,6 +70,9 @@ struct JLOptions
     gc_sweep_always_full::Int8
     compress_sysimage::Int8
     alert_on_critical_error::Int8
+    target_sanitize_memory::Int8
+    target_sanitize_thread::Int8
+    target_sanitize_address::Int8
 end
 
 # This runs early in the sysimage != is not defined yet

base/reflection.jl

@@ -170,17 +170,32 @@ struct CodegenParams
     """
     force_emit_all::Cint
 
+    """
+    When enabled, run the MemorySanitizer pass.
+    """
+    sanitize_memory::Cint
+    """
+    When enabled, run the ThreadSanitizer pass.
+    """
+    sanitize_thread::Cint
+    """
+    When enabled, run the AddressSanitizer pass.
+    """
+    sanitize_address::Cint
+
     function CodegenParams(; track_allocations::Bool=true, code_coverage::Bool=true,
                    prefer_specsig::Bool=false,
                    gnu_pubnames::Bool=true, debug_info_kind::Cint = default_debug_info_kind(),
                    debug_info_level::Cint = Cint(JLOptions().debug_level), safepoint_on_entry::Bool=true,
-                   gcstack_arg::Bool=true, use_jlplt::Bool=true, force_emit_all::Bool=false)
+                   gcstack_arg::Bool=true, use_jlplt::Bool=true, force_emit_all::Bool=false,
+                   sanitize_memory::Bool=false, sanitize_thread::Bool=false, sanitize_address::Bool=false)
         return new(
             Cint(track_allocations), Cint(code_coverage),
             Cint(prefer_specsig),
             Cint(gnu_pubnames), debug_info_kind,
             debug_info_level, Cint(safepoint_on_entry),
-            Cint(gcstack_arg), Cint(use_jlplt), Cint(force_emit_all))
+            Cint(gcstack_arg), Cint(use_jlplt), Cint(force_emit_all),
+            Cint(sanitize_memory), Cint(sanitize_thread), Cint(sanitize_address))
     end
 end
 

src/aotcompile.cpp

@@ -477,7 +477,7 @@ static void resolve_workqueue(jl_codegen_params_t &params, egal_set &method_root
             // emit specsig-to-(jl)invoke conversion
             proto.decl->setLinkage(GlobalVariable::InternalLinkage);
             //protodecl->setAlwaysInline();
-            jl_init_function(proto.decl, params.TargetTriple);
+            jl_init_function(proto.decl, params);
             jl_method_instance_t *mi = jl_get_ci_mi(codeinst);
             size_t nrealargs = jl_nparams(mi->specTypes); // number of actual arguments being passed
             bool is_opaque_closure = jl_is_method(mi->def.value) && mi->def.method->is_for_opaque_closure;
@@ -777,6 +777,10 @@ void *jl_emit_native_impl(jl_array_t *codeinfos, LLVMOrcThreadSafeModuleRef llvm
     CreateNativeMax.updateMax(jl_array_nrows(codeinfos));
     if (cgparams == NULL)
         cgparams = &jl_default_cgparams;
+    jl_cgparams_t target_cgparams = *cgparams;
+    target_cgparams.sanitize_memory = jl_options.target_sanitize_memory;
+    target_cgparams.sanitize_thread = jl_options.target_sanitize_thread;
+    target_cgparams.sanitize_address = jl_options.target_sanitize_address;
     jl_native_code_desc_t *data = new jl_native_code_desc_t;
     orc::ThreadSafeContext ctx;
     orc::ThreadSafeModule backing;
@@ -795,7 +799,7 @@ void *jl_emit_native_impl(jl_array_t *codeinfos, LLVMOrcThreadSafeModuleRef llvm
     jl_codegen_params_t params(ctxt, std::move(target_info.first), std::move(target_info.second));
     if (!llvmmod)
         params.getContext().setDiscardValueNames(true);
-    params.params = cgparams;
+    params.params = &target_cgparams;
     assert(params.imaging_mode); // `_imaging_mode` controls if broken features like code-coverage are disabled
     params.external_linkage = external_linkage;
     params.temporary_roots = jl_alloc_array_1d(jl_array_any_type, 0);
@@ -1565,7 +1569,11 @@ static AOTOutputs add_output_impl(Module &M, TargetMachine &SourceTM, ShardTimer
                 SourceTM.getCodeModel(),
                 SourceTM.getOptLevel()));
         fixupTM(*PMTM);
-        NewPM optimizer{std::move(PMTM), getOptLevel(jl_options.opt_level), OptimizationOptions::defaults(true, true)};
+        auto options = OptimizationOptions::defaults(true, true);
+        options.sanitize_memory = jl_options.target_sanitize_memory;
+        options.sanitize_thread = jl_options.target_sanitize_thread;
+        options.sanitize_address = jl_options.target_sanitize_address;
+        NewPM optimizer{std::move(PMTM), getOptLevel(jl_options.opt_level), options};
         optimizer.run(M);
         assert(!verifyLLVMIR(M));
         bool inject_aliases = false;
@@ -2566,7 +2574,11 @@ void jl_get_llvmf_defn_impl(jl_llvmf_dump_t *dump, jl_method_instance_t *mi, jl_
                 }
                 assert(!verifyLLVMIR(*m.getModuleUnlocked()));
                 if (optimize) {
-                    NewPM PM{jl_ExecutionEngine->cloneTargetMachine(), getOptLevel(jl_options.opt_level)};
+                    auto opts = OptimizationOptions::defaults();
+                    opts.sanitize_memory = params.sanitize_memory;
+                    opts.sanitize_thread = params.sanitize_thread;
+                    opts.sanitize_address = params.sanitize_address;
+                    NewPM PM{jl_ExecutionEngine->cloneTargetMachine(), getOptLevel(jl_options.opt_level), opts};
                     //Safe b/c context lock is held by output
                     PM.run(*m.getModuleUnlocked());
                     assert(!verifyLLVMIR(*m.getModuleUnlocked()));

src/codegen.cpp

@@ -2713,8 +2713,9 @@ static void jl_name_jlfuncparams_args(jl_codegen_params_t &params, Function *F)
     F->getArg(3)->setName("sparams::Any");
 }
 
-void jl_init_function(Function *F, const Triple &TT) JL_NOTSAFEPOINT
+void jl_init_function(Function *F, const jl_codegen_params_t &params) JL_NOTSAFEPOINT
 {
+    auto &TT = params.TargetTriple;
     // set any attributes that *must* be set on all functions
     AttrBuilder attr(F->getContext());
     if (TT.isOSWindows() && TT.getArch() == Triple::x86) {
@@ -2727,22 +2728,18 @@ void jl_init_function(Function *F, const Triple &TT) JL_NOTSAFEPOINT
         attr.addUWTableAttr(llvm::UWTableKind::Default); // force NeedsWinEH
     }
     attr.addAttribute("frame-pointer", "all");
-    if (!TT.isOSWindows()) {
-#if !defined(_COMPILER_ASAN_ENABLED_)
+    if (!TT.isOSWindows() && !JL_FEAT_TEST(params, sanitize_address)) {
         // ASAN won't like us accessing undefined memory causing spurious issues,
         // and Windows has platform-specific handling which causes it to mishandle
         // this annotation. Other platforms should just ignore this if they don't
         // implement it.
         attr.addAttribute("probe-stack", "inline-asm");
         //attr.addAttribute("stack-probe-size", "4096"); // can use this to change the default
-#endif
     }
-#if defined(_COMPILER_ASAN_ENABLED_)
-    attr.addAttribute(Attribute::SanitizeAddress);
-#endif
-#if defined(_COMPILER_MSAN_ENABLED_)
-    attr.addAttribute(Attribute::SanitizeMemory);
-#endif
+    if (JL_FEAT_TEST(params, sanitize_address))
+        attr.addAttribute(Attribute::SanitizeAddress);
+    if (JL_FEAT_TEST(params, sanitize_memory))
+        attr.addAttribute(Attribute::SanitizeMemory);
     F->addFnAttrs(attr);
 }
 
@@ -6247,7 +6244,7 @@ static std::pair<Function*, Function*> get_oc_function(jl_codectx_t &ctx, jl_met
         F = Function::Create(get_func_sig(ctx.builder.getContext()),
                              Function::ExternalLinkage,
                              proto_oc, jl_Module);
-        jl_init_function(F, ctx.emission_context.TargetTriple);
+        jl_init_function(F, ctx.emission_context);
         jl_name_jlfunc_args(ctx.emission_context, F);
         F->setAttributes(AttributeList::get(ctx.builder.getContext(), {get_func_attrs(ctx.builder.getContext()), F->getAttributes()}));
     }
@@ -6770,7 +6767,7 @@ static Function *emit_modifyhelper(jl_codectx_t &ctx2, const jl_cgval_t &op, con
         ArgTy.push_back(ctx.builder.getPtrTy());
     FunctionType *FT = FunctionType::get(elty, ArgTy, false);
     Function *w = Function::Create(FT, GlobalVariable::PrivateLinkage, "", M);
-    jl_init_function(w, ctx.emission_context.TargetTriple);
+    jl_init_function(w, ctx.emission_context);
     w->addFnAttr(Attribute::AlwaysInline);
     w->setUnnamedAddr(GlobalValue::UnnamedAddr::Global);
     Function::arg_iterator AI = w->arg_begin();
@@ -6829,7 +6826,7 @@ Function *emit_tojlinvoke(jl_code_instance_t *codeinst, Value *theFunc, Module *
     Function *f = Function::Create(ctx.types().T_jlfunc,
             GlobalVariable::InternalLinkage,
             name, M);
-    jl_init_function(f, params.TargetTriple);
+    jl_init_function(f, params);
     jl_name_jlfunc_args(params, f);
     //f->setAlwaysInline();
     ctx.f = f; // for jl_Module
@@ -7027,7 +7024,7 @@ void emit_specsig_to_fptr1(
 static void emit_fptr1_wrapper(Module *M, StringRef gf_thunk_name, Value *target, jl_value_t *rettype_const, jl_value_t *declrt, jl_value_t *jlrettype, jl_codegen_params_t &params)
 {
     Function *w = Function::Create(get_func_sig(M->getContext()), GlobalVariable::ExternalLinkage, gf_thunk_name, M);
-    jl_init_function(w, params.TargetTriple);
+    jl_init_function(w, params);
     w->setAttributes(AttributeList::get(M->getContext(), {get_func_attrs(M->getContext()), w->getAttributes()}));
     w->addFnAttr(Attribute::OptimizeNone);
     w->addFnAttr(Attribute::NoInline);
@@ -7071,7 +7068,7 @@ static void emit_specsig_to_specsig(
 {
     jl_returninfo_t returninfo = get_specsig_function(params, M, nullptr, gf_thunk_name, calltype, rettype, is_for_opaque_closure);
     Function *gf_thunk = cast<Function>(returninfo.decl.getCallee());
-    jl_init_function(gf_thunk, params.TargetTriple);
+    jl_init_function(gf_thunk, params);
     gf_thunk->setAttributes(AttributeList::get(gf_thunk->getContext(), {returninfo.attrs, gf_thunk->getAttributes()}));
     emit_specsig_to_specsig(gf_thunk, returninfo.cc, returninfo.return_roots, calltype, rettype, is_for_opaque_closure, nargs, params, target, targetsig, targetrt, targetspec, rettype_const);
 }
@@ -7317,7 +7314,7 @@ static Function *gen_cfun_wrapper(
     Function *cw = Function::Create(functype,
             GlobalVariable::ExternalLinkage,
             funcName, M);
-    jl_init_function(cw, params.TargetTriple);
+    jl_init_function(cw, params);
     cw->setAttributes(AttributeList::get(M->getContext(), {attributes, cw->getAttributes()}));
 
     jl_codectx_t ctx(M->getContext(), params, 0, 0);
@@ -7539,7 +7536,7 @@ static Function *gen_cfun_wrapper(
                 FunctionType::get(getPointerTy(ctx.builder.getContext()), { getPointerTy(ctx.builder.getContext()), ctx.types().T_ppjlvalue }, false),
                 GlobalVariable::ExternalLinkage,
                 funcName, M);
-        jl_init_function(cw_make, ctx.emission_context.TargetTriple);
+        jl_init_function(cw_make, ctx.emission_context);
         cw_make->getArg(0)->setName("wrapper");
         cw_make->getArg(1)->setName("newval");
         BasicBlock *b0 = BasicBlock::Create(ctx.builder.getContext(), "top", cw_make);
@@ -7799,7 +7796,7 @@ static void gen_invoke_wrapper(jl_method_instance_t *lam, jl_value_t *abi, jl_va
 {
     ++GeneratedInvokeWrappers;
     Function *w = Function::Create(get_func_sig(M->getContext()), GlobalVariable::ExternalLinkage, funcName, M);
-    jl_init_function(w, params.TargetTriple);
+    jl_init_function(w, params);
     jl_name_jlfunc_args(params, w);
     w->setAttributes(AttributeList::get(M->getContext(), {get_func_attrs(M->getContext()), w->getAttributes()}));
     w->addFnAttr(Attribute::OptimizeNone);
@@ -8040,7 +8037,7 @@ static jl_returninfo_t get_specsig_function(jl_codegen_params_t &params, Module
         Function *f = M ? cast_or_null<Function>(M->getNamedValue(name)) : NULL;
         if (f == NULL) {
             f = Function::Create(ftype, GlobalVariable::ExternalLinkage, name, M);
-            jl_init_function(f, params.TargetTriple);
+            jl_init_function(f, params);
             if (params.params->debug_info_level >= 2) {
                 ios_t sigbuf;
                 ios_mem(&sigbuf, 0);
@@ -8327,7 +8324,7 @@ static jl_llvm_functions_t
                                           ArgNames, nreq);
         f = cast<Function>(returninfo.decl.getCallee());
         has_sret = (returninfo.cc == jl_returninfo_t::SRet || returninfo.cc == jl_returninfo_t::Union);
-        jl_init_function(f, ctx.emission_context.TargetTriple);
+        jl_init_function(f, ctx.emission_context);
 
         // common pattern: see if all return statements are an argument in that
         // case the apply-generic call can re-use the original box for the return
@@ -8366,7 +8363,7 @@ static jl_llvm_functions_t
         f = Function::Create(needsparams ? ctx.types().T_jlfuncparams : ctx.types().T_jlfunc,
                              GlobalVariable::ExternalLinkage,
                              declarations.specFunctionObject, M);
-        jl_init_function(f, ctx.emission_context.TargetTriple);
+        jl_init_function(f, ctx.emission_context);
         if (needsparams)
             jl_name_jlfuncparams_args(ctx.emission_context, f);
         else
@@ -8402,11 +8399,9 @@ static jl_llvm_functions_t
     FnAttrs.addAttribute(Attribute::StackProtectStrong);
 #endif
 
-#ifdef _COMPILER_TSAN_ENABLED_
-    // TODO: enable this only when a argument like `-race` is passed to Julia
-    //       add a macro for no_sanitize_thread
-    FnAttrs.addAttribute(llvm::Attribute::SanitizeThread);
-#endif
+    // TODO: add a macro for no_sanitize_thread
+    if (JL_FEAT_TEST(ctx, sanitize_thread))
+        FnAttrs.addAttribute(llvm::Attribute::SanitizeThread);
 
     // add the optimization level specified for this module, if any
     int optlevel = jl_get_module_optlevel(ctx.module);
@@ -9773,7 +9768,7 @@ jl_llvm_functions_t jl_emit_codedecls(
                 Function *f = Function::Create(needsparams ? JuliaType::get_jlfuncparams_ty(M.getContext()) : JuliaType::get_jlfunc_ty(M.getContext()),
                                      GlobalVariable::ExternalLinkage,
                                      decls.specFunctionObject, M);
-                jl_init_function(f, params.TargetTriple);
+                jl_init_function(f, params);
                 f->setAttributes(AttributeList::get(M.getContext(), {get_func_attrs(M.getContext()), f->getAttributes()}));
             }
         });
@@ -9836,7 +9831,7 @@ static jl_llvm_functions_t jl_emit_oc_wrapper(orc::ThreadSafeModule &m, jl_codeg
         std::string funcName = get_function_name(true, false, ctx.name, ctx.emission_context.TargetTriple);
         jl_returninfo_t returninfo = get_specsig_function(params, M, NULL, funcName, mi->specTypes, rettype, true);
         Function *gf_thunk = cast<Function>(returninfo.decl.getCallee());
-        jl_init_function(gf_thunk, ctx.emission_context.TargetTriple);
+        jl_init_function(gf_thunk, ctx.emission_context);
         size_t nrealargs = jl_nparams(mi->specTypes);
         emit_specsig_to_fptr1(gf_thunk, returninfo.cc, returninfo.return_roots,
                 mi->specTypes, rettype, true, nrealargs, ctx.emission_context,

src/init.c

@@ -658,7 +658,23 @@ JL_DLLEXPORT jl_cgparams_t jl_default_cgparams = {
         /* safepoint_on_entry */ 1,
         /* gcstack_arg */ 1,
         /* use_jlplt*/ 1 ,
-        /*force_emit_all=*/ 0};
+        /*force_emit_all=*/ 0,
+#ifdef _COMPILER_MSAN_ENABLED_
+        /* sanitize_memory */ 1,
+#else
+        /* sanitize_memory */ 0,
+#endif
+#ifdef _COMPILER_TSAN_ENABLED_
+        /* sanitize_thread */ 1,
+#else
+        /* sanitize_thread */ 0,
+#endif
+#ifdef _COMPILER_ASAN_ENABLED_
+        /* sanitize_address */ 1,
+#else
+        /* sanitize_address */ 0,
+#endif
+};
 
 static void init_global_mutexes(void) {
     JL_MUTEX_INIT(&jl_modules_mutex, "jl_modules_mutex");

src/jitlayers.cpp

@@ -482,7 +482,7 @@ static int jl_analyze_workqueue(jl_code_instance_t *callee, jl_codegen_params_t
                         // emit specsig-to-(jl)invoke conversion
                         proto.decl->setLinkage(GlobalVariable::InternalLinkage);
                         //protodecl->setAlwaysInline();
-                        jl_init_function(proto.decl, params.TargetTriple);
+                        jl_init_function(proto.decl, params);
                         // TODO: maybe this can be cached in codeinst->specfptr?
                         int8_t gc_state = jl_gc_unsafe_enter(ct->ptls); // codegen may contain safepoints (such as jl_subtype calls)
                         jl_method_instance_t *mi = jl_get_ci_mi(codeinst);
@@ -1450,7 +1450,7 @@ namespace {
         auto operator()() JL_NOTSAFEPOINT {
             auto TM = cantFail(JTMB.createTargetMachine());
             fixupTM(*TM);
-            auto NPM = std::make_unique<NewPM>(std::move(TM), O);
+            auto NPM = std::make_unique<NewPM>(std::move(TM), O, OptimizationOptions::defaults());
             // TODO this needs to be locked, as different resource pools may add to the printer vector at the same time
             {
                 std::lock_guard<std::mutex> lock(llvm_printing_mutex);

src/jitlayers.h

@@ -90,6 +90,9 @@ struct OptimizationOptions {
     bool remove_ni;
     bool cleanup;
     bool warn_missed_transformations;
+    bool sanitize_memory;
+    bool sanitize_thread;
+    bool sanitize_address;
 
     static constexpr OptimizationOptions defaults(
         bool lower_intrinsics=true,
@@ -104,12 +107,29 @@ struct OptimizationOptions {
         bool enable_vector_pipeline=true,
         bool remove_ni=true,
         bool cleanup=true,
-        bool warn_missed_transformations=false) {
+        bool warn_missed_transformations=false,
+#ifdef _COMPILER_MSAN_ENABLED_
+        bool sanitize_memory=true,
+#else
+        bool sanitize_memory=false,
+#endif
+#ifdef _COMPILER_TSAN_ENABLED_
+        bool sanitize_thread=true,
+#else
+        bool sanitize_thread=false,
+#endif
+#ifdef _COMPILER_ASAN_ENABLED_
+        bool sanitize_address=true
+#else
+        bool sanitize_address=false
+#endif
+) JL_NOTSAFEPOINT {
         return {lower_intrinsics, dump_native, external_use, llvm_only,
                 always_inline, enable_early_simplifications,
                 enable_early_optimizations, enable_scalar_optimizations,
                 enable_loop_optimizations, enable_vector_pipeline,
-                remove_ni, cleanup, warn_missed_transformations};
+                remove_ni, cleanup, warn_missed_transformations,
+                sanitize_memory, sanitize_thread, sanitize_address};
     }
 };
 
@@ -335,7 +355,7 @@ void emit_specsig_to_fptr1(
         jl_codegen_params_t &params,
         Function *target) JL_NOTSAFEPOINT;
 Function *get_or_emit_fptr1(StringRef Name, Module *M) JL_NOTSAFEPOINT;
-void jl_init_function(Function *F, const Triple &TT) JL_NOTSAFEPOINT;
+void jl_init_function(Function *F, const jl_codegen_params_t &params) JL_NOTSAFEPOINT;
 
 void add_named_global(StringRef name, void *addr) JL_NOTSAFEPOINT;