Use function type in fspec instead of value (#320)

wsmoses · vchuravy · web-flow · commit 3a73ef5725a4 · 2022-04-27T15:52:03.000-04:00
Co-authored-by: Valentin Churavy &lt;vchuravy@users.noreply.github.com&gt;
diff --git a/examples/kernel.jl b/examples/kernel.jl
@@ -16,7 +16,7 @@ GPUCompiler.runtime_module(::CompilerJob{<:Any,TestCompilerParams}) = TestRuntim
 kernel() = nothing
 
 function main()
-    source = FunctionSpec(kernel)
+    source = FunctionSpec(typeof(kernel))
     target = NativeCompilerTarget()
     params = TestCompilerParams()
     job = CompilerJob(target, source, params)
diff --git a/src/driver.jl b/src/driver.jl
@@ -147,11 +147,14 @@ end
     @timeit_debug to "Julia front-end" begin
 
         # get the method instance
-        meth = which(job.source.f, job.source.tt)
-        sig = Base.signature_type(job.source.f, job.source.tt)::Type
+        sig = typed_signature(job)
+        meth = which(sig)
+
         (ti, env) = ccall(:jl_type_intersection_with_env, Any,
                           (Any, Any), sig, meth.sig)::Core.SimpleVector
+
         meth = Base.func_for_method_checked(meth, ti, env)
+
         method_instance = ccall(:jl_specializations_get_linfo, Ref{Core.MethodInstance},
                       (Any, Any, Any, UInt), meth, ti, env, job.source.world)
 
diff --git a/src/interface.jl b/src/interface.jl
@@ -62,7 +62,7 @@ export FunctionSpec
 # what we'll be compiling
 
 struct FunctionSpec{F,TT}
-    f::F
+    f::Type{F}
     tt::Type{TT}
     kernel::Bool
     name::Union{Nothing,String}
@@ -85,8 +85,11 @@ end
 #      world age intersection when querying the compilation cache. once we do, callers
 #      should probably provide the world age of the calling code (!= the current world age)
 #      so that querying the cache from, e.g. `cufuncton` is a fully static operation.
+FunctionSpec(f::Type, tt=Tuple{}, kernel=true, name=nothing, world_age=-1%UInt) =
+    FunctionSpec{f,tt}(f, tt, kernel, name, world_age)
+
 FunctionSpec(f, tt=Tuple{}, kernel=true, name=nothing, world_age=-1%UInt) =
-    FunctionSpec{typeof(f),tt}(f, tt, kernel, name, world_age)
+    FunctionSpec(Core.Typeof(f), tt, kernel, name, world_age)
 
 function Base.getproperty(@nospecialize(spec::FunctionSpec), sym::Symbol)
     if sym == :world
diff --git a/src/irgen.jl b/src/irgen.jl
@@ -302,7 +302,8 @@ end
 end
 
 function classify_arguments(@nospecialize(job::CompilerJob), codegen_ft::LLVM.FunctionType)
-    source_sig = Base.signature_type(job.source.f, job.source.tt)::Type
+    source_sig = typed_signature(job)
+
     source_types = [source_sig.parameters...]
 
     codegen_types = parameters(codegen_ft)
diff --git a/src/reflection.jl b/src/reflection.jl
@@ -47,6 +47,11 @@ end
 code_lowered(@nospecialize(job::CompilerJob); kwargs...) =
     InteractiveUtils.code_lowered(job.source.f, job.source.tt; kwargs...)
 
+@inline function typed_signature(@nospecialize(job::CompilerJob))
+    u = Base.unwrap_unionall(job.source.tt)
+    return Base.rewrap_unionall(Tuple{job.source.f, u.parameters...}, job.source.tt)
+end
+
 function code_typed(@nospecialize(job::CompilerJob); interactive::Bool=false, kwargs...)
     # TODO: use the compiler driver to get the Julia method instance (we might rewrite it)
     if interactive
@@ -58,9 +63,9 @@ function code_typed(@nospecialize(job::CompilerJob); interactive::Bool=false, kw
         descend_code_typed(job.source.f, job.source.tt; interp, kwargs...)
     elseif VERSION >= v"1.7-"
         interp = get_interpreter(job)
-        InteractiveUtils.code_typed(job.source.f, job.source.tt; interp, kwargs...)
+        Base.code_typed_by_type(typed_signature(job); interp, kwargs...)
     else
-        InteractiveUtils.code_typed(job.source.f, job.source.tt; kwargs...)
+        Base.code_typed_by_type(typed_signature(job); kwargs...)
     end
 end
 
diff --git a/src/rtlib.jl b/src/rtlib.jl
@@ -112,7 +112,7 @@ function build_runtime(@nospecialize(job::CompilerJob); ctx)
 
     # the compiler job passed into here is identifies the job that requires the runtime.
     # derive a job that represents the runtime itself (notably with kernel=false).
-    source = FunctionSpec(identity, Tuple{Nothing}, false, nothing, job.source.world_age)
+    source = FunctionSpec(typeof(identity), Tuple{Nothing}, false, nothing, job.source.world_age)
     job = CompilerJob(job.target, source, job.params)
 
     for method in values(Runtime.methods)
@@ -122,7 +122,7 @@ function build_runtime(@nospecialize(job::CompilerJob); ctx)
         else
             method.def
         end
-        emit_function!(mod, job, def, method; ctx)
+        emit_function!(mod, job, typeof(def), method; ctx)
     end
 
     # we cannot optimize the runtime library, because the code would then be optimized again
diff --git a/src/validation.jl b/src/validation.jl
@@ -2,26 +2,36 @@
 
 export InvalidIRError
 
+function get_method_matches(@nospecialize(job::CompilerJob))
+    tt = typed_signature(job)
+
+    ms = Core.MethodMatch[]
+    for m in Base._methods_by_ftype(tt, -1, job.source.world)::Vector
+        m = m::Core.MethodMatch
+        push!(ms, m)
+    end
+
+    return ms
+end
+
+
 function check_method(@nospecialize(job::CompilerJob))
     isa(job.source.f, Core.Builtin) && throw(KernelError(job, "function is not a generic function"))
 
     # get the method
-    ms = Base.methods(job.source.f, job.source.tt)
+    ms = get_method_matches(job)
     isempty(ms)   && throw(KernelError(job, "no method found"))
     length(ms)!=1 && throw(KernelError(job, "no unique matching method"))
-    m = first(ms)
 
     # kernels can't return values
     if job.source.kernel
         cache = ci_cache(job)
         mt = method_table(job)
         interp = GPUInterpreter(cache, mt, job.source.world)
-        @static if v"1.8-beta2" <= VERSION < v"1.9-" || VERSION ≥ v"1.9.0-DEV.190"
-            # https://github.com/JuliaLang/julia/pull/44515
-            rt = Base.return_types(job.source.f, job.source.tt; interp)[1]
-        else
-            rt = Base.return_types(job.source.f, job.source.tt, interp)[1]
-        end
+        m = only(ms)
+        ty = Core.Compiler.typeinf_type(interp, m.method, m.spec_types, m.sparams)
+        rt = something(ty, Any)
+
         if rt != Nothing
             throw(KernelError(job, "kernel returns a value of type `$rt`",
                 """Make sure your kernel function ends in `return`, `return nothing` or `nothing`.
@@ -61,7 +71,9 @@ end
 function check_invocation(@nospecialize(job::CompilerJob))
     # make sure any non-isbits arguments are unused
     real_arg_i = 0
-    sig = Base.signature_type(job.source.f, job.source.tt)::Type
+
+    sig = typed_signature(job)
+
     for (arg_i,dt) in enumerate(sig.parameters)
         isghosttype(dt) && continue
         Core.Compiler.isconstType(dt) && continue
diff --git a/test/definitions/gcn.jl b/test/definitions/gcn.jl
@@ -8,7 +8,7 @@ end
 # create a GCN-based test compiler, and generate reflection methods for it
 
 function gcn_job(@nospecialize(func), @nospecialize(types); kernel::Bool=false, kwargs...)
-    source = FunctionSpec(func, Base.to_tuple_type(types), kernel)
+    source = FunctionSpec(typeof(func), Base.to_tuple_type(types), kernel)
     target = GCNCompilerTarget(dev_isa="gfx900")
     params = TestCompilerParams()
     CompilerJob(target, source, params), kwargs
diff --git a/test/definitions/native.jl b/test/definitions/native.jl
@@ -18,8 +18,8 @@ end
 
 GPUCompiler.method_table(@nospecialize(job::NativeCompilerJob)) = method_table
 
-function native_job(@nospecialize(func), @nospecialize(types); kernel::Bool=false, entry_abi=:specfunc, kwargs...)
-    source = FunctionSpec(func, Base.to_tuple_type(types), kernel)
+function native_job(@nospecialize(f_type), @nospecialize(types); kernel::Bool=false, entry_abi=:specfunc, kwargs...)
+    source = FunctionSpec(f_type, Base.to_tuple_type(types), kernel)
     target = NativeCompilerTarget(always_inline=true)
     params = TestCompilerParams()
     CompilerJob(target, source, params, entry_abi), kwargs
@@ -249,8 +249,8 @@ module LazyCodegen
     end
 
     import GPUCompiler: deferred_codegen_jobs
-    @generated function deferred_codegen(::Val{f}, ::Val{tt}) where {f,tt}
-        job, _ = native_job(f, tt)
+    @generated function deferred_codegen(f::F, ::Val{tt}) where {F,tt}
+        job, _ = native_job(F, tt)
 
         addr = get_trampoline(job)
         trampoline = pointer(addr)
@@ -265,7 +265,7 @@ module LazyCodegen
         end
     end
 
-    @generated function abi_call(f::Ptr{Cvoid}, rt::Type{RT}, tt::Type{T}, args::Vararg{Any, N}) where {T, RT, N}
+    @generated function abi_call(f::Ptr{Cvoid}, rt::Type{RT}, tt::Type{T}, func::F, args::Vararg{Any, N}) where {T, RT, F, N}
         argtt    = tt.parameters[1]
         rettype  = rt.parameters[1]
         argtypes = DataType[argtt.parameters...]
@@ -276,8 +276,27 @@ module LazyCodegen
         before = :()
         after = :(ret)
 
+
         # Note this follows: emit_call_specfun_other
         JuliaContext() do ctx
+
+            if !isghosttype(F) && !Core.Compiler.isconstType(F)
+                isboxed = GPUCompiler.deserves_argbox(F)
+                argexpr = :(func)
+                if isboxed
+                    push!(ccall_types, Any)
+                else
+                    et = convert(LLVMType, func; ctx)
+                    if isa(et, LLVM.SequentialType) # et->isAggregateType
+                        push!(ccall_types, Ptr{F})
+                        argexpr = Expr(:call, GlobalRef(Base, :Ref), argexpr)
+                    else
+                        push!(ccall_types, F)
+                    end
+                end
+                push!(argexprs, argexpr)
+            end
+
             T_jlvalue = LLVM.StructType(LLVMType[],;ctx)
             T_prjlvalue = LLVM.PointerType(T_jlvalue, #= AddressSpace::Tracked =# 10)
 
@@ -330,7 +349,7 @@ module LazyCodegen
     @inline function call_delayed(f::F, args...) where F
         tt = Tuple{map(Core.Typeof, args)...}
         rt = Core.Compiler.return_type(f, tt)
-        ptr = deferred_codegen(Val(f), Val(tt))
-        abi_call(ptr, rt, tt, args...)
+        ptr = deferred_codegen(f, Val(tt))
+        abi_call(ptr, rt, tt, f, args...)
     end
 end
diff --git a/test/native.jl b/test/native.jl
@@ -11,6 +11,9 @@ include("definitions/native.jl")
 
         (CI, rt) = native_code_typed(MyCallable(), (Int, Int), kernel=false)[1]
         @test CI.slottypes[1] == Core.Compiler.Const(MyCallable())
+
+        (CI, rt) = native_code_typed(typeof(MyCallable()), (Int, Int), kernel=false)[1]
+        @test CI.slottypes[1] == Core.Compiler.Const(MyCallable())
     end
 
     @testset "Compilation database" begin
@@ -339,11 +342,7 @@ end
 
     # Test ABI removal
     ir = sprint(io->native_code_llvm(io, call_real, Tuple{ComplexF64}))
-    if VERSION < v"1.8-" || v"1.8-beta2" <= VERSION < v"1.9-" || VERSION ≥ v"1.9.0-DEV.190"
-        @test !occursin("alloca", ir)
-    else
-        @test_broken !occursin("alloca", ir)
-    end
+    @test !occursin("alloca", ir)
 
     ghostly_identity(x, y) = y
     @test call_delayed(ghostly_identity, nothing, 1) == 1
@@ -358,6 +357,20 @@ end
     throws(arr, i) = arr[i]
     @test call_delayed(throws, [1], 1) == 1
     @test_throws BoundsError call_delayed(throws, [1], 0)
+
+    struct Closure
+        x::Int64
+    end
+    (c::Closure)(b) = c.x+b
+
+    @test call_delayed(Closure(3), 5) == 8
+
+    struct Closure2
+        x::Integer
+    end
+    (c::Closure2)(b) = c.x+b
+
+    @test call_delayed(Closure2(3), 5) == 8
 end
 
 ############################################################################################
diff --git a/test/ptx.jl b/test/ptx.jl
@@ -83,11 +83,10 @@ end
     kernel() = return
 
     ir = sprint(io->ptx_code_llvm(io, kernel, Tuple{}))
-    @test occursin(r"@.*kernel.+\(\)", ir)
+    @test any(occursin(r"@.*kernel.+\(\)", a) for a in split(ir, "\n"))
 
     ir = sprint(io->ptx_code_llvm(io, kernel, Tuple{}; kernel=true))
-    @test occursin(r"@.*kernel.+\(\[1 x i64\] %state\)", ir)
-
+    @test any(occursin(r"@.*kernel.+\(\[1 x i64\] %state\)", a) for a in split(ir, "\n"))
 
     # state should only passed to device functions that use it
 
@@ -107,16 +106,16 @@ end
                                   kernel=true, dump_module=true))
 
     # kernel should take state argument before all else
-    @test occursin(r"@.*kernel.+\(\[1 x i64\] %state", ir)
+    @test any(occursin(r"@.*kernel.+\(\[1 x i64\] %state", a) for a in split(ir, "\n"))
 
     # child1 doesn't use the state
-    @test occursin(r"@.*child1.+\(i64", ir)
+    @test any(occursin(r"@.*child1.+\(i64", a) for a in split(ir, "\n"))
 
     # child2 does
-    @test occursin(r"@.*child2.+\(\[1 x i64\] %state", ir)
+    @test any(occursin(r"@.*child2.+\(\[1 x i64\] %state", a) for a in split(ir, "\n"))
 
     # can't have the unlowered intrinsic
-    @test !occursin("julia.gpu.state_getter", ir)
+    @test !any(occursin("julia.gpu.state_getter", a) for a in split(ir, "\n"))
 end
 end
 
