Use ABI queries from Julia

src/GPUCompiler.jl

@@ -67,6 +67,7 @@ include("runtime.jl")
 # compiler implementation
 include("deprecated.jl")
 include("jlgen.jl")
+include("abi.jl")
 include("irgen.jl")
 include("optim.jl")
 include("validation.jl")

src/abi.jl

@@ -0,0 +1,163 @@
+# Julia's specsig ABI, as reported by Julia itself.
+#
+# Everything here mirrors the `jl_get_specsig_layout` interface in `src/julia.h`.
+# Where it is available we ask Julia how it lowers a signature instead of
+# re-deriving the rules; the fallbacks in `irgen.jl` implement the same rules by
+# hand for older Julia versions.
+
+const JL_ABI_LAYOUT_VERSION = UInt32(1)
+
+# jl_abi_retcc_t
+const JL_ABI_RET_BOXED    = Int32(0)
+const JL_ABI_RET_REGISTER = Int32(1)
+const JL_ABI_RET_SRET     = Int32(2)
+const JL_ABI_RET_UNION    = Int32(3)
+const JL_ABI_RET_GHOSTS   = Int32(4)
+
+# jl_abi_argcc_t
+const JL_ABI_ARG_ELIDED   = Int32(0)
+const JL_ABI_ARG_VALUE    = Int32(1)
+const JL_ABI_ARG_INDIRECT = Int32(2)
+const JL_ABI_ARG_BOXED    = Int32(3)
+
+# jl_abi_elide_t
+const JL_ABI_ELIDE_NONE      = Int32(0)
+const JL_ABI_ELIDE_GHOST     = Int32(1)
+const JL_ABI_ELIDE_UNIQUEREP = Int32(2)
+
+struct JLAbiArgInfo
+    typ::Ptr{Cvoid}
+    cc::Int32
+    param_idx::Int32
+    roots_idx::Int32
+    elide_reason::Int32
+    _reserved::Int32
+end
+
+struct JLAbiLayout
+    version::UInt32
+    specsig::Int32
+    needsparams::Int32
+    sigt::Ptr{Cvoid}
+    rettype::Ptr{Cvoid}
+    rettype_cc::Int32
+    return_roots::UInt32
+    all_roots::Int32
+    union_bytes::Csize_t
+    union_align::Csize_t
+    union_minalign::Csize_t
+    sret_idx::Int32
+    return_roots_idx::Int32
+    pgcstack_idx::Int32
+    nprefix_params::Int32
+    nargs::Int32
+    nparams::Int32
+end
+JLAbiLayout() = JLAbiLayout(JL_ABI_LAYOUT_VERSION, 0, 0, C_NULL, C_NULL, 0, 0, 0,
+                            0, 0, 0, -1, -1, -1, 0, 0, 0)
+
+struct JLAbiQuery
+    version::UInt32
+    ci::Ptr{Cvoid}
+    sigt::Ptr{Cvoid}
+    rt::Ptr{Cvoid}
+    is_opaque_closure::Int32
+    cgparams::Ptr{Base.CodegenParams}
+    mod::Ptr{Cvoid}
+    datalayout::Ptr{UInt8}
+    triple::Ptr{UInt8}
+    name::Ptr{UInt8}
+    decl_out::Ptr{Ptr{Cvoid}}
+end
+
+function _have_symbol(lib::String, sym::Symbol)
+    handle = try
+        Libdl.dlopen(Libdl.dlpath(lib))
+    catch
+        return false
+    end
+    return Libdl.dlsym(handle, sym; throw_error=false) !== nothing
+end
+
+const _LIBJULIA_CODEGEN = Base.isdebugbuild() ? "libjulia-codegen-debug" : "libjulia-codegen"
+const _LIBJULIA_INTERNAL = Base.isdebugbuild() ? "libjulia-internal-debug" : "libjulia-internal"
+
+"""
+Whether this Julia can report its own specsig ABI. Probed rather than
+version-gated so that the feature is picked up by backports too.
+"""
+const HAS_ABI_LAYOUT = _have_symbol(_LIBJULIA_CODEGEN, :jl_get_specsig_layout)
+
+"""
+Whether Julia exports the boxing predicates that decide the specsig ABI; if not,
+`irgen.jl` falls back to its own copy of the rules.
+"""
+const HAS_DESERVES_CCALL = _have_symbol(_LIBJULIA_INTERNAL, :jl_deserves_stack)
+
+# jl_value_t* of an arbitrary object, including immutable ones like Types, which
+# `pointer_from_objref` refuses
+_value_ptr(@nospecialize x) = ccall(:jl_value_ptr, Ptr{Cvoid}, (Any,), x)
+
+"""
+    abi_layout(job; mod=nothing)
+
+Ask Julia for the specsig layout of `job.source`, using this job's
+[`codegen_params`](@ref) — `gcstack_arg=false` removes a leading parameter and
+so shifts every index. Returns `(layout::JLAbiLayout, args::Vector{JLAbiArgInfo})`.
+
+Pass `mod` to have the declaration built with the target's data layout and
+triple. That only changes the address spaces of the emitted pointers, not the
+parameter count or ordering, so callers that just want the index mapping (such
+as [`classify_arguments`](@ref)) can leave it out.
+"""
+function abi_layout(@nospecialize(job::CompilerJob); mod::Union{Nothing,LLVM.Module}=nothing)
+    sigt = abi_signature(job.source)
+    rt = typeinf_type(job.source; interp=get_interpreter(job))
+    return abi_layout(sigt, rt; params=codegen_params(job), mod)
+end
+
+function abi_layout(@nospecialize(sigt), @nospecialize(rt);
+                    params::Base.CodegenParams,
+                    mod::Union{Nothing,LLVM.Module}=nothing,
+                    is_opaque_closure::Bool=false)
+    HAS_ABI_LAYOUT ||
+        error("this Julia does not provide jl_get_specsig_layout")
+    nargs = length((sigt::DataType).parameters)
+    args = Vector{JLAbiArgInfo}(undef, max(nargs, 1))
+    layout = Ref(JLAbiLayout())
+    pparams = Ref(params)
+    local ret
+    GC.@preserve sigt rt args layout pparams begin
+        query = Ref(JLAbiQuery(JL_ABI_LAYOUT_VERSION,
+                               C_NULL, _value_ptr(sigt), _value_ptr(rt),
+                               Int32(is_opaque_closure),
+                               Base.unsafe_convert(Ptr{Base.CodegenParams}, pparams),
+                               mod === nothing ? C_NULL : convert(Ptr{Cvoid}, mod.ref),
+                               C_NULL, C_NULL, C_NULL, C_NULL))
+        ret = @ccall jl_get_specsig_layout(query::Ptr{JLAbiQuery}, layout::Ptr{JLAbiLayout},
+                                           pointer(args)::Ptr{JLAbiArgInfo},
+                                           Int32(nargs)::Int32)::Cint
+    end
+    ret == 0 || error("jl_get_specsig_layout failed for $sigt -> $rt (code $ret)")
+    l = layout[]
+    return l, args[1:l.nargs]
+end
+
+"""
+    abi_signature(source)
+
+The signature `source` is compiled against. This is its `specTypes` unless a
+`Core.ABIOverride` replaces it, which `specTypes` alone would miss.
+"""
+abi_signature(mi::Core.MethodInstance) = mi.specTypes
+@static if isdefined(Core, :ABIOverride)
+    abi_signature(ci::Core.CodeInstance) =
+        @static if isdefined(Base, :get_ci_abi)
+            Base.get_ci_abi(ci)
+        else
+            def = ci.def
+            def isa Core.ABIOverride ? def.abi : (def::Core.MethodInstance).specTypes
+        end
+else
+    abi_signature(ci::Core.CodeInstance) = (ci.def::Core.MethodInstance).specTypes
+end

src/gcn.jl

@@ -97,9 +97,11 @@ function add_kernarg_address_spaces!(
     )
     ft = function_type(f)
 
-    # find the byref parameters by checking for the byref attribute directly,
-    # rather than re-classifying arguments (which can fail on typed-pointer LLVM
-    # due to element type mismatches in classify_arguments assertions).
+    # find the byref parameters by checking for the byref attribute directly.
+    # This pass runs after optimization, so the parameter list is no longer the
+    # one Julia emitted; reading the attributes we ourselves applied in `irgen`
+    # is the only thing that stays valid. (It also sidesteps the assertions in
+    # `_classify_arguments_legacy`, which fire on typed-pointer LLVM.)
     byref_kind = LLVM.API.LLVMGetEnumAttributeKindForName("byref", 5)
     byref_mask = BitVector(undef, length(parameters(ft)))
     for i in 1:length(parameters(ft))

src/irgen.jl

@@ -62,7 +62,7 @@ function irgen(@nospecialize(job::CompilerJob))
     if job.config.name !== nothing
         LLVM.name!(entry, safe_name(job.config.name))
     elseif job.config.kernel
-        LLVM.name!(entry, mangle_sig(job.source.specTypes))
+        LLVM.name!(entry, mangle_sig(abi_signature(job.source)))
     end
     if job.config.entry_abi === :specfunc
         func = compiled[job.source].func
@@ -547,30 +547,90 @@ end
 # - `name`: the name of the argument
 # - `idx`: the index of the argument in the LLVM function type, or `nothing` if the argument
 #          is not passed at the LLVM level.
+# - `roots_idx`: the index of the extra `.roots.` shadow parameter codegen emits for an
+#          aggregate holding some-but-not-all tracked pointers, or `nothing`.
 function classify_arguments(@nospecialize(job::CompilerJob), codegen_ft::LLVM.FunctionType;
                             post_optimization::Bool=false)
-    source_sig = job.source.specTypes
-    source_types = [source_sig.parameters...]
+    # `post_optimization` asks a different question: by then our own passes have
+    # rewritten the parameter list (`lower_byval`, Metal's `pass_by_reference!`,
+    # ...), and callers want the convention of the function *as it now stands*.
+    # Only the untouched entry point is described by Julia's specsig ABI.
+    if HAS_ABI_LAYOUT && !post_optimization
+        return _classify_arguments_abi(job)
+    else
+        return _classify_arguments_legacy(job, codegen_ft; post_optimization)
+    end
+end
 
-    source_argnames = Base.method_argnames(job.source.def)
-    while length(source_argnames) < length(source_types)
+# argument names come from the method, not from Julia's codegen: the back-ends
+# (Metal in particular) want the names a user would recognize, and codegen only
+# reports the CodeInfo slot names.
+function source_argnames(@nospecialize(job::CompilerJob), nargs::Int)
+    names = Base.method_argnames(job.source.def)
+    while length(names) < nargs
         # this is probably due to a trailing vararg; repeat its name
-        push!(source_argnames, source_argnames[end])
+        push!(names, names[end])
     end
+    return names
+end
+
+function _classify_arguments_abi(@nospecialize(job::CompilerJob))
+    layout, arginfos = abi_layout(job)
+    layout.specsig != 0 ||
+        error("$(job.source) does not use the specialized signature; cannot classify arguments")
+    # GPUCompiler always requests gcstack_arg=false, so the only leading
+    # parameters that can appear are the return slots (kernels return nothing
+    # and so have none, but `:specfunc` entry points for ordinary functions do)
+    @assert layout.pgcstack_idx == -1
+
+    names = source_argnames(job, Int(layout.nargs))
+
+    # `param_idx` already counts the leading return slots; the kernel-state
+    # parameter only exists after optimization, which this path does not serve
+    args = []
+    for (i, info) in enumerate(arginfos)
+        typ = unsafe_pointer_to_objref(info.typ)
+        cc = if info.cc == JL_ABI_ARG_ELIDED
+            GHOST
+        elseif info.cc == JL_ABI_ARG_VALUE
+            BITS_VALUE
+        elseif info.cc == JL_ABI_ARG_INDIRECT
+            BITS_REF
+        else
+            MUT_REF
+        end
+        idx = info.param_idx < 0 ? nothing : Int(info.param_idx) + 1
+        roots_idx = info.roots_idx < 0 ? nothing : Int(info.roots_idx) + 1
+        push!(args, (cc=cc, typ=typ, name=names[i], idx=idx, roots_idx=roots_idx))
+    end
+    return args
+end
+
+# The pre-`jl_get_specsig_layout` implementation: reconstruct the mapping by
+# walking the signature in lockstep with the LLVM function type codegen produced.
+# Kept for older Julia versions, and as the differential-test reference.
+function _classify_arguments_legacy(@nospecialize(job::CompilerJob), codegen_ft::LLVM.FunctionType;
+                                    post_optimization::Bool=false)
+    source_sig = abi_signature(job.source)
+    source_types = [source_sig.parameters...]
+
+    argnames = source_argnames(job, length(source_types))
 
     codegen_types = parameters(codegen_ft)
 
     if post_optimization && kernel_state_type(job) !== Nothing
         args = []
-        push!(args, (cc=KERNEL_STATE, typ=kernel_state_type(job), name=:kernel_state, idx=1))
+        push!(args, (cc=KERNEL_STATE, typ=kernel_state_type(job), name=:kernel_state,
+                     idx=1, roots_idx=nothing))
         codegen_i = 2
     else
         args = []
         codegen_i = 1
     end
-    for (source_typ, source_name) in zip(source_types, source_argnames)
+    for (source_typ, source_name) in zip(source_types, argnames)
         if isghosttype(source_typ) || Core.Compiler.isconstType(source_typ)
-            push!(args, (cc=GHOST, typ=source_typ, name=source_name, idx=nothing))
+            push!(args, (cc=GHOST, typ=source_typ, name=source_name, idx=nothing,
+                         roots_idx=nothing))
             continue
         end
 
@@ -582,19 +642,23 @@ function classify_arguments(@nospecialize(job::CompilerJob), codegen_ft::LLVM.Fu
             # - literal pointer values
             if source_typ <: Ptr || source_typ <: Core.LLVMPtr
                 @assert llvm_source_typ == codegen_typ
-                push!(args, (cc=BITS_VALUE, typ=source_typ, name=source_name, idx=codegen_i))
+                push!(args, (cc=BITS_VALUE, typ=source_typ, name=source_name, idx=codegen_i,
+                             roots_idx=nothing))
             # - boxed values
             #   XXX: use `deserves_retbox` instead?
             elseif llvm_source_typ isa LLVM.PointerType
                 @assert llvm_source_typ == codegen_typ
-                push!(args, (cc=MUT_REF, typ=source_typ, name=source_name, idx=codegen_i))
+                push!(args, (cc=MUT_REF, typ=source_typ, name=source_name, idx=codegen_i,
+                             roots_idx=nothing))
             # - references to aggregates
             else
                 @assert llvm_source_typ != codegen_typ
-                push!(args, (cc=BITS_REF, typ=source_typ, name=source_name, idx=codegen_i))
+                push!(args, (cc=BITS_REF, typ=source_typ, name=source_name, idx=codegen_i,
+                             roots_idx=nothing))
             end
         else
-            push!(args, (cc=BITS_VALUE, typ=source_typ, name=source_name, idx=codegen_i))
+            push!(args, (cc=BITS_VALUE, typ=source_typ, name=source_name, idx=codegen_i,
+                         roots_idx=nothing))
         end
 
         codegen_i += 1
@@ -608,18 +672,9 @@ function is_immutable_datatype(T::Type)
 end
 
 function is_inlinealloc(T::Type)
-    mayinlinealloc = (T.name.flags >> 2) & 1 == true
-    # FIXME: To simple
-    if mayinlinealloc
-        if !Base.datatype_pointerfree(T)
-            t_name(dt::DataType)=dt.name
-            if t_name(T).n_uninitialized != 0
-                return false
-            end
-        end
-        return true
-    end
-    return false
+    # jl_datatype_isinlinealloc has been exported for a long time; no need to
+    # reimplement the mayinlinealloc/n_uninitialized/fielddesc rules here
+    ccall(:jl_datatype_isinlinealloc, Cint, (Any, Cint), T, 0) != 0
 end
 
 function is_concrete_immutable(T::Type)
@@ -634,14 +689,17 @@ function is_pointerfree(T::Type)
 end
 
 function deserves_stack(@nospecialize(T))
+    if HAS_DESERVES_CCALL
+        return ccall(:jl_deserves_stack, Cint, (Any,), T) != 0
+    end
     if !is_concrete_immutable(T)
         return false
     end
     return is_inlinealloc(T)
 end
 
-deserves_argbox(T) = !deserves_stack(T)
-deserves_retbox(T) = deserves_argbox(T)
+deserves_argbox(@nospecialize(T)) = !deserves_stack(T)
+deserves_retbox(@nospecialize(T)) = deserves_argbox(T)
 function deserves_sret(T, llvmT)
     @assert isa(T,DataType)
     sizeof(T) > sizeof(Ptr{Cvoid}) && !isa(llvmT, LLVM.FloatingPointType) && !isa(llvmT, LLVM.VectorType)