Add code_ircode (#45306)

To match `typeinf_ircode` with how typeinf lock is used ATM (i.e.,
optimizer is run inside the lock), we can manually lock it because the
lock is re-entrant.

Co-authored-by: Shuhei Kadowaki <aviatesk@gmail.com>

base/compiler/optimize.jl

@@ -534,21 +534,44 @@ function ipo_escape_cache(mi_cache::MICache) where MICache
 end
 null_escape_cache(linfo::Union{InferenceResult,MethodInstance}) = nothing
 
-function run_passes(ci::CodeInfo, sv::OptimizationState, caller::InferenceResult)
-    @timeit "convert"   ir = convert_to_ircode(ci, sv)
-    @timeit "slot2reg"  ir = slot2reg(ir, ci, sv)
+macro pass(name, expr)
+    optimize_until = esc(:optimize_until)
+    stage = esc(:__stage__)
+    macrocall = :(@timeit $(esc(name)) $(esc(expr)))
+    macrocall.args[2] = __source__  # `@timeit` may want to use it
+    quote
+        $macrocall
+        matchpass($optimize_until, ($stage += 1), $(esc(name))) && $(esc(:(@goto __done__)))
+    end
+end
+
+matchpass(optimize_until::Int, stage, _name) = optimize_until < stage
+matchpass(optimize_until::String, _stage, name) = optimize_until == name
+matchpass(::Nothing, _, _) = false
+
+function run_passes(
+    ci::CodeInfo,
+    sv::OptimizationState,
+    caller::InferenceResult,
+    optimize_until = nothing,  # run all passes by default
+)
+    __stage__ = 1  # used by @pass
+    # NOTE: The pass name MUST be unique for `optimize_until::AbstractString` to work
+    @pass "convert"   ir = convert_to_ircode(ci, sv)
+    @pass "slot2reg"  ir = slot2reg(ir, ci, sv)
     # TODO: Domsorting can produce an updated domtree - no need to recompute here
-    @timeit "compact 1" ir = compact!(ir)
-    @timeit "Inlining"  ir = ssa_inlining_pass!(ir, ir.linetable, sv.inlining, ci.propagate_inbounds)
+    @pass "compact 1" ir = compact!(ir)
+    @pass "Inlining"  ir = ssa_inlining_pass!(ir, ir.linetable, sv.inlining, ci.propagate_inbounds)
     # @timeit "verify 2" verify_ir(ir)
-    @timeit "compact 2" ir = compact!(ir)
-    @timeit "SROA"      ir = sroa_pass!(ir)
-    @timeit "ADCE"      ir = adce_pass!(ir)
-    @timeit "type lift" ir = type_lift_pass!(ir)
-    @timeit "compact 3" ir = compact!(ir)
+    @pass "compact 2" ir = compact!(ir)
+    @pass "SROA"      ir = sroa_pass!(ir)
+    @pass "ADCE"      ir = adce_pass!(ir)
+    @pass "type lift" ir = type_lift_pass!(ir)
+    @pass "compact 3" ir = compact!(ir)
     if JLOptions().debug_level == 2
         @timeit "verify 3" (verify_ir(ir); verify_linetable(ir.linetable))
     end
+    @label __done__  # used by @pass
     return ir
 end
 

base/compiler/typeinfer.jl

@@ -918,6 +918,39 @@ function typeinf_code(interp::AbstractInterpreter, method::Method, @nospecialize
     return code, rt
 end
 
+"""
+    typeinf_ircode(
+        interp::AbstractInterpreter,
+        method::Method,
+        atype,
+        sparams::SimpleVector,
+        optimize_until::Union{Integer,AbstractString,Nothing},
+    ) -> (ir::Union{IRCode,Nothing}, returntype::Type)
+
+Infer a `method` and return an `IRCode` with inferred `returntype` on success.
+"""
+function typeinf_ircode(
+    interp::AbstractInterpreter,
+    method::Method,
+    @nospecialize(atype),
+    sparams::SimpleVector,
+    optimize_until::Union{Integer,AbstractString,Nothing},
+)
+    ccall(:jl_typeinf_begin, Cvoid, ())
+    frame = typeinf_frame(interp, method, atype, sparams, false)
+    if frame === nothing
+        ccall(:jl_typeinf_end, Cvoid, ())
+        return nothing, Any
+    end
+    (; result) = frame
+    opt_params = OptimizationParams(interp)
+    opt = OptimizationState(frame, opt_params, interp)
+    ir = run_passes(opt.src, opt, result, optimize_until)
+    rt = widenconst(ignorelimited(result.result))
+    ccall(:jl_typeinf_end, Cvoid, ())
+    return ir, rt
+end
+
 # compute an inferred frame
 function typeinf_frame(interp::AbstractInterpreter, method::Method, @nospecialize(atype), sparams::SimpleVector, run_optimizer::Bool)
     mi = specialize_method(method, atype, sparams)::MethodInstance

base/reflection.jl

@@ -1291,6 +1291,104 @@ function code_typed_opaque_closure(@nospecialize(oc::Core.OpaqueClosure);
     end
 end
 
+"""
+    code_ircode(f, [types])
+
+Return an array of pairs of `IRCode` and inferred return type if type inference succeeds.
+The `Method` is included instead of `IRCode` otherwise.
+
+See also: [`code_typed`](@ref)
+
+# Internal Keyword Arguments
+
+This section should be considered internal, and is only for who understands Julia compiler
+internals.
+
+- `world=Base.get_world_counter()`: optional, controls the world age to use when looking up
+  methods, use current world age if not specified.
+- `interp=Core.Compiler.NativeInterpreter(world)`: optional, controls the interpreter to
+  use, use the native interpreter Julia uses if not specified.
+- `optimize_until`: optional, controls the optimization passes to run.  If it is a string,
+  it specifies the name of the pass up to which the optimizer is run.  If it is an integer,
+  it specifies the number of passes to run.  If it is `nothing` (default), all passes are
+  run.
+
+# Example
+
+One can put the argument types in a tuple to get the corresponding `code_ircode`.
+
+```jldoctest
+julia> Base.code_ircode(+, (Float64, Int64))
+1-element Vector{Any}:
+ 388 1 ─ %1 = Base.sitofp(Float64, _3)::Float64
+    │   %2 = Base.add_float(_2, %1)::Float64
+    └──      return %2
+     => Float64
+
+julia> Base.code_ircode(+, (Float64, Int64); optimize_until = "compact 1")
+1-element Vector{Any}:
+ 388 1 ─ %1 = Base.promote(_2, _3)::Tuple{Float64, Float64}
+    │   %2 = Core._apply_iterate(Base.iterate, Base.:+, %1)::Float64
+    └──      return %2
+     => Float64
+```
+"""
+function code_ircode(
+    @nospecialize(f),
+    @nospecialize(types = default_tt(f));
+    world = get_world_counter(),
+    interp = Core.Compiler.NativeInterpreter(world),
+    optimize_until::Union{Integer,AbstractString,Nothing} = nothing,
+)
+    if isa(f, Core.OpaqueClosure)
+        error("OpaqueClosure not supported")
+    end
+    ft = Core.Typeof(f)
+    if isa(types, Type)
+        u = unwrap_unionall(types)
+        tt = rewrap_unionall(Tuple{ft,u.parameters...}, types)
+    else
+        tt = Tuple{ft,types...}
+    end
+    return code_ircode_by_type(tt; world, interp, optimize_until)
+end
+
+"""
+    code_ircode_by_type(types::Type{<:Tuple}; ...)
+
+Similar to [`code_ircode`](@ref), except the argument is a tuple type describing
+a full signature to query.
+"""
+function code_ircode_by_type(
+    @nospecialize(tt::Type);
+    world = get_world_counter(),
+    interp = Core.Compiler.NativeInterpreter(world),
+    optimize_until::Union{Integer,AbstractString,Nothing} = nothing,
+)
+    ccall(:jl_is_in_pure_context, Bool, ()) &&
+        error("code reflection cannot be used from generated functions")
+    tt = to_tuple_type(tt)
+    matches = _methods_by_ftype(tt, -1, world)::Vector
+    asts = []
+    for match in matches
+        match = match::Core.MethodMatch
+        meth = func_for_method_checked(match.method, tt, match.sparams)
+        (code, ty) = Core.Compiler.typeinf_ircode(
+            interp,
+            meth,
+            match.spec_types,
+            match.sparams,
+            optimize_until,
+        )
+        if code === nothing
+            push!(asts, meth => Any)
+        else
+            push!(asts, code => ty)
+        end
+    end
+    return asts
+end
+
 function return_types(@nospecialize(f), @nospecialize(types=default_tt(f));
                       world = get_world_counter(),
                       interp = Core.Compiler.NativeInterpreter(world))

test/compiler/ssair.jl

@@ -335,6 +335,23 @@ f_if_typecheck() = (if nothing; end; unsafe_load(Ptr{Int}(0)))
     success(pipeline(Cmd(cmd); stdout=stdout, stderr=stderr)) && isempty(String(take!(stderr)))
 end
 
+@testset "code_ircode" begin
+    @test first(only(Base.code_ircode(+, (Float64, Float64)))) isa Compiler.IRCode
+    @test first(only(Base.code_ircode(+, (Float64, Float64); optimize_until = 3))) isa
+          Compiler.IRCode
+    @test first(only(Base.code_ircode(+, (Float64, Float64); optimize_until = "SROA"))) isa
+          Compiler.IRCode
+
+    function demo(f)
+        f()
+        f()
+        f()
+    end
+    @test first(only(Base.code_ircode(demo))) isa Compiler.IRCode
+    @test first(only(Base.code_ircode(demo; optimize_until = 3))) isa Compiler.IRCode
+    @test first(only(Base.code_ircode(demo; optimize_until = "SROA"))) isa Compiler.IRCode
+end
+
 let
     function test_useref(stmt, v, op)
         if isa(stmt, Expr)