Merge pull request #305 from JuliaGPU/tb/asserts

maleadt · web-flow · commit 3eabb5f6d03b · 2022-03-07T12:03:37.000+01:00
Fix context usage and add assertions CI
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -7,27 +7,44 @@ on:
   workflow_dispatch:
 jobs:
   test:
-    name: Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }}
+    name: Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }} - assertions=${{ matrix.assertions }}
     runs-on: ${{ matrix.os }}
     strategy:
       fail-fast: false
       matrix:
-        version:
-          - '1.6'
-          - '^1.7.0-beta2'
-          - 'nightly'
-        os:
-          - ubuntu-latest
-          - macOS-latest
-          - windows-latest
-        arch:
-          - x64
+        version: ['1.6', '1.7', '^1.8.0-beta1', 'nightly']
+        os: [ubuntu-latest, macOS-latest, windows-latest]
+        arch: [x64]
+        assertions: [false]
+        include:
+          # special test with LLVM assertions enabled
+          # TODO: enable this across all versions
+          #       (needs julia-actions/setup-julia support)
+          - os: ubuntu-latest
+            arch: x64
+            version: '1.7'
+            assertions: true
+          - os: ubuntu-latest
+            arch: x64
+            version: '1.8'
+            assertions: true
     steps:
       - uses: actions/checkout@v2
       - uses: julia-actions/setup-julia@v1
+        if: ${{ ! matrix.assertions }}
         with:
           version: ${{ matrix.version }}
           arch: ${{ matrix.arch }}
+      - name: Download Julia with assertions
+        if: ${{ matrix.assertions }}
+        env:
+          version: ${{ matrix.version }}
+          arch: ${{ matrix.arch }}
+        run: |
+          wget https://julialangnightlies.s3.amazonaws.com/assert_bin/linux/$arch/$version/julia-latest-linux64.tar.gz
+          tar -xvzf julia-latest-linux64.tar.gz
+          rm -rf julia-latest-linux64.tar.gz
+          echo $PWD/julia-*/bin >> $GITHUB_PATH
       - uses: actions/cache@v1
         env:
           cache-name: cache-artifacts
diff --git a/src/driver.jl b/src/driver.jl
@@ -1,5 +1,7 @@
 # compiler driver and main interface
 
+export JuliaContext
+
 # NOTE: the keyword arguments to compile/codegen control those aspects of compilation that
 #       might have to be changed (e.g. set libraries=false when recursing, or set
 #       strip=true for reflection). What remains defines the compilation job itself,
@@ -31,19 +33,54 @@ Other keyword arguments can be found in the documentation of [`cufunction`](@ref
 function compile(target::Symbol, @nospecialize(job::CompilerJob);
                  libraries::Bool=true, deferred_codegen::Bool=true,
                  optimize::Bool=true, strip::Bool=false, validate::Bool=true,
-                 only_entry::Bool=false)
+                 only_entry::Bool=false, ctx::Union{Context,Nothing}=nothing)
     if compile_hook[] !== nothing
         compile_hook[](job)
     end
 
     return codegen(target, job;
-                   libraries, deferred_codegen, optimize, strip, validate, only_entry)
+                   libraries, deferred_codegen, optimize, strip, validate, only_entry, ctx)
+end
+
+# transitionary feature to deal versions of Julia that rely on a global context
+#
+# Julia 1.9 removed the global LLVM context, requiring to pass a context to codegen APIs,
+# so the GPUCompiler APIs have been adapted to require passing a Context object as well.
+# however, on older versions of Julia we cannot make codegen emit into that context. we
+# could use a hack (serialize + deserialize) to move code into the correct context, however
+# as it turns out some of our optimization passes erroneously rely on the context being
+# global and unique, resulting in segfaults when we use a local context instead.
+#
+# to work around this mess, and still present a reasonably unified API, we introduce the
+# JuliaContext helper below, which returns a local context on Julia 1.9, and the global
+# unique context on all other versions. Once we only support Julia 1.9, we'll deprecate
+# this helper to a regular `Contxet()` call.
+function JuliaContext()
+    if VERSION >= v"1.9.0-DEV.115"
+        # Julia 1.9 knows how to deal with arbitrary contexts
+        Context()
+    else
+        # earlier versions of Julia claim so, but actually use a global context
+        isboxed_ref = Ref{Bool}()
+        typ = LLVMType(ccall(:jl_type_to_llvm, LLVM.API.LLVMTypeRef,
+                       (Any, Ptr{Bool}), Any, isboxed_ref))
+        context(typ)
+    end
+end
+function JuliaContext(f)
+    if VERSION >= v"1.9.0-DEV.115"
+        Context(f)
+    else
+        f(JuliaContext())
+        # we cannot dispose of the global unique context
+    end
 end
 
 function codegen(output::Symbol, @nospecialize(job::CompilerJob);
                  libraries::Bool=true, deferred_codegen::Bool=true, optimize::Bool=true,
                  strip::Bool=false, validate::Bool=true, only_entry::Bool=false,
-                 parent_job::Union{Nothing, CompilerJob} = nothing)
+                 parent_job::Union{Nothing, CompilerJob}=nothing,
+                 ctx::Union{Context,Nothing}=nothing)
     ## Julia IR
 
     mi, mi_meta = emit_julia(job)
@@ -52,37 +89,56 @@ function codegen(output::Symbol, @nospecialize(job::CompilerJob);
         return mi, mi_meta
     end
 
+    temporary_context = ctx === nothing
+    if temporary_context && output == :llvm
+        # if we return IR structures, we cannot construct a temporary context
+        error("Request to return LLVM IR; please provide a context")
+    end
 
-    ## LLVM IR
-
-    ir, ir_meta = emit_llvm(job, mi; libraries, deferred_codegen, optimize, only_entry)
+    try
+        ## LLVM IR
 
-    if output == :llvm
-        if strip
-            @timeit_debug to "strip debug info" strip_debuginfo!(ir)
+        if temporary_context
+            ctx = JuliaContext()
+        elseif VERSION < v"1.9.0-DEV.115" && ctx != JuliaContext()
+            error("""Julia <1.9 does not suppport generating code in an arbitrary LLVM context.
+                     Use a JuliaContext instead.""")
         end
 
-        return ir, ir_meta
-    end
+        ir, ir_meta = emit_llvm(job, mi; libraries, deferred_codegen, optimize, only_entry, ctx)
 
+        if output == :llvm
+            if strip
+                @timeit_debug to "strip debug info" strip_debuginfo!(ir)
+            end
 
-    ## machine code
+            return ir, ir_meta
+        end
 
-    format = if output == :asm
-        LLVM.API.LLVMAssemblyFile
-    elseif output == :obj
-        LLVM.API.LLVMObjectFile
-    else
-        error("Unknown assembly format $output")
-    end
-    asm, asm_meta = emit_asm(job, ir; strip, validate, format)
 
-    if output == :asm || output == :obj
-        return asm, asm_meta
-    end
+        ## machine code
 
+        format = if output == :asm
+            LLVM.API.LLVMAssemblyFile
+        elseif output == :obj
+            LLVM.API.LLVMObjectFile
+        else
+            error("Unknown assembly format $output")
+        end
+        asm, asm_meta = emit_asm(job, ir; strip, validate, format)
 
-    error("Unknown compilation output $output")
+        if output == :asm || output == :obj
+            return asm, asm_meta
+        end
+
+
+        error("Unknown compilation output $output")
+    finally
+        if temporary_context && VERSION >= v"1.9.0-DEV.115"
+            @assert ctx != JuliaContext()
+            dispose(ctx)
+        end
+    end
 end
 
 @locked function emit_julia(@nospecialize(job::CompilerJob))
@@ -135,7 +191,7 @@ const __llvm_initialized = Ref(false)
 
 @locked function emit_llvm(@nospecialize(job::CompilerJob), @nospecialize(method_instance);
                            libraries::Bool=true, deferred_codegen::Bool=true, optimize::Bool=true,
-                           only_entry::Bool=false)
+                           only_entry::Bool=false, ctx::Context)
     if !__llvm_initialized[]
         InitializeAllTargets()
         InitializeAllTargetInfos()
@@ -146,8 +202,7 @@ const __llvm_initialized = Ref(false)
     end
 
     @timeit_debug to "IR generation" begin
-        ir, compiled = irgen(job, method_instance)
-        ctx = context(ir)
+        ir, compiled = irgen(job, method_instance; ctx)
         entry_fn = compiled[method_instance].specfunc
         entry = functions(ir)[entry_fn]
     end
@@ -226,7 +281,7 @@ const __llvm_initialized = Ref(false)
                 # cached compilation
                 dyn_entry_fn = get!(cache, dyn_job) do
                     dyn_ir, dyn_meta = codegen(:llvm, dyn_job; optimize=false,
-                                               deferred_codegen=false, parent_job=job)
+                                               deferred_codegen=false, parent_job=job, ctx)
                     dyn_entry_fn = LLVM.name(dyn_meta.entry)
                     merge!(compiled, dyn_meta.compiled)
                     @assert context(dyn_ir) == ctx
diff --git a/src/irgen.jl b/src/irgen.jl
@@ -1,9 +1,9 @@
 # LLVM IR generation
 
-function irgen(@nospecialize(job::CompilerJob), method_instance::Core.MethodInstance)
-    mod, compiled = @timeit_debug to "emission" compile_method_instance(job, method_instance)
+function irgen(@nospecialize(job::CompilerJob), method_instance::Core.MethodInstance;
+               ctx::Context)
+    mod, compiled = @timeit_debug to "emission" compile_method_instance(job, method_instance; ctx)
     entry_fn = compiled[method_instance].specfunc
-    ctx = context(mod)
 
     # clean up incompatibilities
     @timeit_debug to "clean-up" begin
diff --git a/src/jlgen.jl b/src/jlgen.jl
@@ -326,7 +326,7 @@ function _lookup_fun(mi, min_world, max_world)
 end
 
 function compile_method_instance(@nospecialize(job::CompilerJob),
-                                 method_instance::MethodInstance)
+                                 method_instance::MethodInstance; ctx::Context)
     # populate the cache
     cache = ci_cache(job)
     mt = method_table(job)
@@ -362,11 +362,17 @@ function compile_method_instance(@nospecialize(job::CompilerJob),
 
     # generate IR
     GC.@preserve lookup_cb begin
-        native_code = if VERSION >= v"1.8.0-DEV.661"
+        native_code = if VERSION >= v"1.9.0-DEV.115"
+            ccall(:jl_create_native, Ptr{Cvoid},
+                  (Vector{MethodInstance}, LLVM.API.LLVMContextRef, Ptr{Base.CodegenParams}, Cint),
+                  [method_instance], ctx, Ref(params), #=extern policy=# 1)
+        elseif VERSION >= v"1.8.0-DEV.661"
+            @assert ctx == JuliaContext()
             ccall(:jl_create_native, Ptr{Cvoid},
                   (Vector{MethodInstance}, Ptr{Base.CodegenParams}, Cint),
                   [method_instance], Ref(params), #=extern policy=# 1)
         else
+            @assert ctx == JuliaContext()
             ccall(:jl_create_native, Ptr{Cvoid},
                   (Vector{MethodInstance}, Base.CodegenParams, Cint),
                   [method_instance], params, #=extern policy=# 1)
diff --git a/src/precompile.jl b/src/precompile.jl
@@ -59,7 +59,8 @@ function _precompile_()
     #@assert precompile(Tuple{typeof(GPUCompiler.split_kwargs),Tuple{},Vector{Symbol},Vararg{Vector{Symbol}, N} where N})
     let fbody = try __lookup_kwbody__(which(GPUCompiler.compile, (Symbol,GPUCompiler.CompilerJob,))) catch missing end
         if !ismissing(fbody)
-            @assert precompile(fbody, (Bool,Bool,Bool,Bool,Bool,Bool,typeof(GPUCompiler.compile),Symbol,GPUCompiler.CompilerJob,))
+            @assert precompile(fbody, (Bool,Bool,Bool,Bool,Bool,Bool,Context,typeof(GPUCompiler.compile),Symbol,GPUCompiler.CompilerJob,))
+            @assert precompile(fbody, (Bool,Bool,Bool,Bool,Bool,Bool,Nothing,typeof(GPUCompiler.compile),Symbol,GPUCompiler.CompilerJob,))
         end
     end
 end
diff --git a/src/reflection.jl b/src/reflection.jl
@@ -100,10 +100,12 @@ See also: [`@device_code_llvm`](@ref), `InteractiveUtils.code_llvm`
 function code_llvm(io::IO, @nospecialize(job::CompilerJob); optimize::Bool=true, raw::Bool=false,
                    debuginfo::Symbol=:default, dump_module::Bool=false)
     # NOTE: jl_dump_function_ir supports stripping metadata, so don't do it in the driver
-    ir, meta = codegen(:llvm, job; optimize=optimize, strip=false, validate=false)
-    str = ccall(:jl_dump_function_ir, Ref{String},
-                (LLVM.API.LLVMValueRef, Bool, Bool, Ptr{UInt8}),
-                meta.entry, !raw, dump_module, debuginfo)
+    str = JuliaContext() do ctx
+        ir, meta = codegen(:llvm, job; optimize=optimize, strip=false, validate=false, ctx)
+        ccall(:jl_dump_function_ir, Ref{String},
+              (LLVM.API.LLVMValueRef, Bool, Bool, Ptr{UInt8}),
+              meta.entry, !raw, dump_module, debuginfo)
+    end
     highlight(io, str, "llvm")
 end
 code_llvm(@nospecialize(job::CompilerJob); kwargs...) = code_llvm(stdout, job; kwargs...)
diff --git a/src/rtlib.jl b/src/rtlib.jl
@@ -75,10 +75,10 @@ end
 
 ## functionality to build the runtime library
 
-function emit_function!(mod, @nospecialize(job::CompilerJob), f, method)
+function emit_function!(mod, @nospecialize(job::CompilerJob), f, method; ctx::Context)
     tt = Base.to_tuple_type(method.types)
     new_mod, meta = codegen(:llvm, similar(job, FunctionSpec(f, tt, #=kernel=# false));
-                            optimize=false, libraries=false)
+                            optimize=false, libraries=false, ctx)
     ft = eltype(llvmtype(meta.entry))
     expected_ft = convert(LLVM.FunctionType, method; ctx=context(new_mod))
     if return_type(ft) != return_type(expected_ft)
@@ -92,12 +92,6 @@ function emit_function!(mod, @nospecialize(job::CompilerJob), f, method)
     dispose(pm)
 
     temp_name = LLVM.name(meta.entry)
-    # FIXME: on 1.6, there's no single global LLVM context anymore,
-    #        but there's no API yet to pass a context to codegen.
-    # round-trip the module through serialization to get it in the proper context.
-    buf = convert(MemoryBuffer, new_mod)
-    new_mod = parse(LLVM.Module, buf; ctx=context(mod))
-    @assert context(mod) == context(new_mod)
     link!(mod, new_mod)
     entry = functions(mod)[temp_name]
 
@@ -128,7 +122,7 @@ function build_runtime(@nospecialize(job::CompilerJob); ctx)
         else
             method.def
         end
-        emit_function!(mod, job, def, method)
+        emit_function!(mod, job, def, method; ctx)
     end
 
     # we cannot optimize the runtime library, because the code would then be optimized again
diff --git a/test/definitions/native.jl b/test/definitions/native.jl
@@ -101,10 +101,10 @@ module LazyCodegen
             @assert !cc.compiled
             job = cc.job
 
-            ir, meta = GPUCompiler.codegen(:llvm, job; validate=false)
-            entry_name = name(meta.entry)
-
-            jitted_mod = compile!(orc, ir)
+            entry_name, jitted_mod = JuliaContext() do ctx
+                ir, meta = GPUCompiler.codegen(:llvm, job; validate=false, ctx)
+                name(meta.entry), compile!(orc, ir)
+            end
 
             addr = addressin(orc, jitted_mod, entry_name)
             ptr  = pointer(addr)
@@ -195,21 +195,24 @@ module LazyCodegen
             sym = LLVM.API.LLVMOrcCSymbolFlagsMapPair(mangle(lljit, target_sym), flags)
 
             function materialize(mr)
-                ir, meta = GPUCompiler.codegen(:llvm, job; validate=false)
+                JuliaContext() do ctx
+                    ir, meta = GPUCompiler.codegen(:llvm, job; validate=false, ctx)
 
-                # Rename entry to match target_sym
-                LLVM.name!(meta.entry, target_sym)
+                    # Rename entry to match target_sym
+                    LLVM.name!(meta.entry, target_sym)
 
-                # So 1. serialize the module
-                buf = convert(MemoryBuffer, ir)
+                    # So 1. serialize the module
+                    buf = convert(MemoryBuffer, ir)
 
-                # 2. deserialize and wrap by a ThreadSafeModule
-                ctx = ThreadSafeContext()
-                mod = parse(LLVM.Module, buf; ctx=context(ctx))
-                tsm = ThreadSafeModule(mod; ctx)
+                    # 2. deserialize and wrap by a ThreadSafeModule
+                    ThreadSafeContext() do ctx
+                        mod = parse(LLVM.Module, buf; ctx=context(ctx))
+                        tsm = ThreadSafeModule(mod; ctx)
 
-                il = LLVM.IRTransformLayer(lljit)
-                LLVM.emit(il, mr, tsm)
+                        il = LLVM.IRTransformLayer(lljit)
+                        LLVM.emit(il, mr, tsm)
+                    end
+                end
 
                 return nothing
             end
@@ -252,7 +255,7 @@ module LazyCodegen
         after = :(ret)
 
         # Note this follows: emit_call_specfun_other
-        Context() do ctx
+        JuliaContext() do ctx
             T_jlvalue = LLVM.StructType(LLVMType[],;ctx)
             T_prjlvalue = LLVM.PointerType(T_jlvalue, #= AddressSpace::Tracked =# 10)
 
diff --git a/test/native.jl b/test/native.jl
