Compiler: add Task call, invoke, and fetch optimizations and operations

Compiler/src/abstractinterpretation.jl

@@ -2345,7 +2345,7 @@ function abstract_finalizer(interp::AbstractInterpreter, argtypes::Vector{Any},
         finalizer_argvec = Any[argtypes[2], argtypes[3]]
         call = abstract_call(interp, ArgInfo(nothing, finalizer_argvec), StmtInfo(false, false), sv, #=max_methods=#1)::Future
         return Future{CallMeta}(call, interp, sv) do call, interp, sv
-            return CallMeta(Nothing, Any, Effects(), FinalizerInfo(call.info, call.effects))
+            return CallMeta(Nothing, Any, Effects(), IndirectCallInfo(call.info, call.effects, false))
         end
     end
     return Future(CallMeta(Nothing, Any, Effects(), NoCallInfo()))
@@ -2679,6 +2679,8 @@ function abstract_call_known(interp::AbstractInterpreter, @nospecialize(f),
             return Future(abstract_eval_isdefinedglobal(interp, sv, si.saw_latestworld, argtypes))
         elseif f === Core.get_binding_type
             return Future(abstract_eval_get_binding_type(interp, sv, argtypes))
+        elseif f === Core._task
+            return abstract_eval_task_builtin(interp, arginfo, si, sv)
         end
         rt = abstract_call_builtin(interp, f, arginfo, sv)
         ft = popfirst!(argtypes)
@@ -3208,6 +3210,58 @@ function abstract_eval_splatnew(interp::AbstractInterpreter, e::Expr, sstate::St
     return RTEffects(rt, Any, effects)
 end
 
+function abstract_eval_task_builtin(interp::AbstractInterpreter, arginfo::ArgInfo, si::StmtInfo, sv::AbsIntState)
+    (; fargs, argtypes) = arginfo
+    la = length(argtypes)
+    𝕃ᵢ = typeinf_lattice(interp)
+    if !isempty(argtypes) && !isvarargtype(argtypes[end])
+        if !(3 <= la <= 4)
+            return Future(CallMeta(Bottom, Any, EFFECTS_THROWS, NoCallInfo()))
+        end
+    elseif isempty(argtypes) || la > 5
+        return Future(CallMeta(Bottom, Any, EFFECTS_THROWS, NoCallInfo()))
+    end
+    size_arg = argtypes[3]
+    if !hasintersect(widenconst(size_arg), Int)
+        return Future(CallMeta(Bottom, Any, EFFECTS_THROWS, NoCallInfo()))
+    end
+    func_arg = argtypes[2]
+
+    # Handle optional Method/CodeInstance/Type argument (4th parameter) as invoke
+    if la == 4
+        invoke_args = Any[Const(Core.invoke), func_arg, argtypes[4]]
+        invoke_arginfo = ArgInfo(nothing, invoke_args)
+        invoke_future = abstract_invoke(interp, invoke_arginfo, si, sv)
+        return Future{CallMeta}(invoke_future, interp, sv) do invoke_result, interp, sv
+            fetch_type = widenconst(invoke_result.rt)
+            fetch_error = widenconst(invoke_result.exct)
+            task_effects = invoke_result.effects
+            if fetch_type === Any && fetch_error === Any
+                rt_result = Task
+            else
+                rt_result = PartialTask(fetch_type, fetch_error)
+            end
+            info_result = IndirectCallInfo(invoke_result.info, task_effects, true)
+            return CallMeta(rt_result, Any, Effects(), info_result)
+        end
+    end
+
+    # Otherwise use abstract_call for function analysis
+    callinfo_future = abstract_call(interp, ArgInfo(nothing, Any[func_arg]), StmtInfo(true, si.saw_latestworld), sv, #=max_methods=#1)
+    return Future{CallMeta}(callinfo_future, interp, sv) do callinfo, interp, sv
+        fetch_type = widenconst(callinfo.rt)
+        fetch_error = widenconst(callinfo.exct)
+        task_effects = callinfo.effects
+        if fetch_type === Any && fetch_error === Any
+            rt_result = Task
+        else
+            rt_result = PartialTask(fetch_type, fetch_error)
+        end
+        info_result = IndirectCallInfo(callinfo.info, task_effects, true)
+        return CallMeta(rt_result, Any, Effects(), info_result)
+    end
+end
+
 function abstract_eval_new_opaque_closure(interp::AbstractInterpreter, e::Expr, sstate::StatementState,
                                           sv::AbsIntState)
     𝕃ᵢ = typeinf_lattice(interp)
@@ -4030,6 +4084,8 @@ end
             fields[i] = a
         end
         anyrefine && return PartialStruct(𝕃ᵢ, rt.typ, _getundefs(rt), fields)
+    elseif isa(rt, PartialTask)
+        return rt # already widened, by construction
     end
     if isa(rt, PartialOpaque)
         return rt # XXX: this case was missed in #39512

Compiler/src/abstractlattice.jl

@@ -24,13 +24,13 @@ is_valid_lattice_norec(::ConstsLattice, @nospecialize(elem)) = isa(elem, Const)
 """
     struct PartialsLattice{𝕃<:AbstractLattice} <: AbstractLattice
 
-A lattice extending a base lattice `𝕃` and adjoining `PartialStruct` and `PartialOpaque`.
+A lattice extending a base lattice `𝕃` and adjoining `PartialStruct`, `PartialOpaque`, and `PartialTask`.
 """
 struct PartialsLattice{𝕃<:AbstractLattice} <: AbstractLattice
     parent::𝕃
 end
 widenlattice(𝕃::PartialsLattice) = 𝕃.parent
-is_valid_lattice_norec(::PartialsLattice, @nospecialize(elem)) = isa(elem, PartialStruct) || isa(elem, PartialOpaque)
+is_valid_lattice_norec(::PartialsLattice, @nospecialize(elem)) = isa(elem, PartialStruct) || isa(elem, PartialOpaque) || isa(elem, PartialTask)
 
 """
     struct ConditionalsLattice{𝕃<:AbstractLattice} <: AbstractLattice
@@ -191,6 +191,7 @@ information that would not be available from the type itself.
 @nospecializeinfer function has_nontrivial_extended_info(𝕃::PartialsLattice, @nospecialize t)
     isa(t, PartialStruct) && return true
     isa(t, PartialOpaque) && return true
+    isa(t, PartialTask) && return true
     return has_nontrivial_extended_info(widenlattice(𝕃), t)
 end
 @nospecializeinfer function has_nontrivial_extended_info(𝕃::ConstsLattice, @nospecialize t)
@@ -223,6 +224,7 @@ that should be forwarded along with constant propagation.
         # return false
     end
     isa(t, PartialOpaque) && return true
+    isa(t, PartialTask) && return true
     return is_const_prop_profitable_arg(widenlattice(𝕃), t)
 end
 @nospecializeinfer function is_const_prop_profitable_arg(𝕃::ConstsLattice, @nospecialize t)
@@ -246,6 +248,7 @@ end
 @nospecializeinfer function is_forwardable_argtype(𝕃::PartialsLattice, @nospecialize x)
     isa(x, PartialStruct) && return true
     isa(x, PartialOpaque) && return true
+    isa(x, PartialTask) && return true
     return is_forwardable_argtype(widenlattice(𝕃), x)
 end
 @nospecializeinfer function is_forwardable_argtype(𝕃::ConstsLattice, @nospecialize x)

Compiler/src/ssair/inlining.jl

@@ -1219,6 +1219,37 @@ function narrow_opaque_closure!(ir::IRCode, stmt::Expr, @nospecialize(info::Call
             stmt.args[3] = newT
         end
     end
+    return nothing
+end
+
+function handle_task_call!(ir::IRCode, idx::Int, stmt::Expr, info::IndirectCallInfo, state::InliningState)
+    length(stmt.args) == 3 || return
+    # Extract the CodeInstance from the inference result if available
+    info_edge = extract_indirect_invoke(info; check_fully_covers=false)
+    info_edge === nothing && return nothing
+    info, edge = info_edge
+    case = compileable_specialization(edge, Effects(), InliningEdgeTracker(state), info, state)
+    case === nothing && return nothing
+    # Append the CodeInstance as a third argument to the _task call
+    # Core._task(func, size) becomes Core._tak(func, size, ci)
+    push!(stmt.args, case.invoke)
+    ir[SSAValue(idx)][:stmt] = stmt
+    return nothing
+end
+
+function extract_indirect_invoke(info::IndirectCallInfo; check_fully_covers::Bool)
+    info = info.info
+    info isa MethodResultPure && (info = info.info)
+    info isa ConstCallInfo && (info = info.call)
+    info isa MethodMatchInfo || return nothing
+    length(info.edges) == length(info.results) == 1 || return nothing
+    match = info.results[1]::MethodMatch
+    if check_fully_covers
+        match.fully_covers || return nothing
+    end
+    edge = info.edges[1]
+    edge === nothing && return nothing
+    return info, edge
 end
 
 # As a matter of convenience, this pass also computes effect-freenes.
@@ -1288,7 +1319,8 @@ function process_simple!(todo::Vector{Pair{Int,Any}}, ir::IRCode, idx::Int, flag
                 f !== modifyfield! &&
                 f !== Core.modifyglobal! &&
                 f !== Core.memoryrefmodify! &&
-                f !== atomic_pointermodify)
+                f !== atomic_pointermodify &&
+                f !== Core._task)
                 # No inlining defined for most builtins (just invoke/apply/typeassert/finalizer), so attempt an early exit for them
                 return nothing
             end
@@ -1538,16 +1570,10 @@ function handle_opaque_closure_call!(todo::Vector{Pair{Int,Any}},
     return nothing
 end
 
-function handle_modifyop!_call!(ir::IRCode, idx::Int, stmt::Expr, info::ModifyOpInfo, state::InliningState)
-    info = info.info
-    info isa MethodResultPure && (info = info.info)
-    info isa ConstCallInfo && (info = info.call)
-    info isa MethodMatchInfo || return nothing
-    length(info.edges) == length(info.results) == 1 || return nothing
-    match = info.results[1]::MethodMatch
-    match.fully_covers || return nothing
-    edge = info.edges[1]
-    edge === nothing && return nothing
+function handle_modifyop!_call!(ir::IRCode, idx::Int, stmt::Expr, info::IndirectCallInfo, state::InliningState)
+    info_edge = extract_indirect_invoke(info; check_fully_covers=true)
+    info_edge === nothing && return nothing
+    info, edge = info_edge
     case = compileable_specialization(edge, Effects(), InliningEdgeTracker(state), info, state)
     case === nothing && return nothing
     stmt.head = :invoke_modify
@@ -1556,7 +1582,7 @@ function handle_modifyop!_call!(ir::IRCode, idx::Int, stmt::Expr, info::ModifyOp
     return nothing
 end
 
-function handle_finalizer_call!(ir::IRCode, idx::Int, stmt::Expr, info::FinalizerInfo,
+function handle_finalizer_call!(ir::IRCode, idx::Int, stmt::Expr, info::IndirectCallInfo,
                                 state::InliningState)
     # Finalizers don't return values, so if their execution is not observable,
     # we can just not register them
@@ -1649,14 +1675,22 @@ function assemble_inline_todo!(ir::IRCode, state::InliningState)
         end
 
         # handle special cased builtins
+        f = sig.f
         if isa(info, OpaqueClosureCallInfo)
             handle_opaque_closure_call!(todo, ir, idx, stmt, info, flag, sig, state)
-        elseif isa(info, ModifyOpInfo)
-            handle_modifyop!_call!(ir, idx, stmt, info, state)
-        elseif sig.f === Core.invoke
+        elseif isa(info, IndirectCallInfo)
+            if f === Core.finalizer
+                handle_finalizer_call!(ir, idx, stmt, info, state)
+            elseif f === modifyfield! ||
+                   f === Core.modifyglobal! ||
+                   f === Core.memoryrefmodify! ||
+                   f === atomic_pointermodify
+                handle_modifyop!_call!(ir, idx, stmt, info, state)
+            elseif f === Core._task
+                handle_task_call!(ir, idx, stmt, info, state)
+            end
+        elseif f === Core.invoke
             handle_invoke_call!(todo, ir, idx, stmt, info, flag, sig, state)
-        elseif isa(info, FinalizerInfo)
-            handle_finalizer_call!(ir, idx, stmt, info, state)
         else
             # cascade to the generic (and extendable) handler
             handle_call!(todo, ir, idx, stmt, info, flag, sig, state)
@@ -1718,6 +1752,8 @@ function early_inline_special_case(ir::IRCode, stmt::Expr, flag::UInt32,
         elseif ⊑(optimizer_lattice(state.interp), cond, Bool) && stmt.args[3] === stmt.args[4]
             return SomeCase(stmt.args[3])
         end
+    elseif f === Core.task_result_type
+        return SomeCase(quoted(instanceof_tfunc(type)[1]))
     end
     return nothing
 end

Compiler/src/ssair/passes.jl

@@ -1298,7 +1298,7 @@ function sroa_pass!(ir::IRCode, inlining::Union{Nothing,InliningState}=nothing)
         elseif is_known_call(stmt, Core.finalizer, compact)
             3 <= length(stmt.args) <= 5 || continue
             info = compact[SSAValue(idx)][:info]
-            if isa(info, FinalizerInfo)
+            if isa(info, IndirectCallInfo)
                 is_finalizer_inlineable(info.effects) || continue
             else
                 # Inlining performs legality checks on the finalizer to determine
@@ -1673,7 +1673,7 @@ function try_resolve_finalizer!(ir::IRCode, alloc_idx::Int, finalizer_idx::Int,
 
     finalizer_stmt = ir[SSAValue(finalizer_idx)][:stmt]
     argexprs = Any[finalizer_stmt.args[2], finalizer_stmt.args[3]]
-    flag = info isa FinalizerInfo ? flags_for_effects(info.effects) : IR_FLAG_NULL
+    flag = isa(info, IndirectCallInfo) ? flags_for_effects(info.effects) : IR_FLAG_NULL
     if length(finalizer_stmt.args) >= 4
         inline = finalizer_stmt.args[4]
         if inline === nothing

Compiler/src/stmtinfo.jl

@@ -454,33 +454,31 @@ end
 add_edges_impl(edges::Vector{Any}, info::ReturnTypeCallInfo) = add_edges!(edges, info.info)
 
 """
-    info::FinalizerInfo <: CallInfo
-
-Represents the information of a potential (later) call to the finalizer on the given
-object type.
-"""
-struct FinalizerInfo <: CallInfo
-    info::CallInfo   # the callinfo for the finalizer call
-    effects::Effects # the effects for the finalizer call
-end
-# merely allocating a finalizer does not imply edges (unless it gets inlined later)
-add_edges_impl(::Vector{Any}, ::FinalizerInfo) = nothing
-
-"""
-    info::ModifyOpInfo <: CallInfo
-
-Represents a resolved call of one of:
- - `modifyfield!(obj, name, op, x, [order])`
- - `modifyglobal!(mod, var, op, x, order)`
- - `memoryrefmodify!(memref, op, x, order, boundscheck)`
- - `Intrinsics.atomic_pointermodify(ptr, op, x, order)`
-
-`info.info` wraps the call information of `op(getval(), x)`.
-"""
-struct ModifyOpInfo <: CallInfo
-    info::CallInfo # the callinfo for the `op(getval(), x)` call
+    info::IndirectCallInfo <: CallInfo
+
+Represents information about a call that involves an indirect/nested function call.
+Used for:
+ - `modifyfield!(obj, name, op, x, [order])` where `op(getval(), x)` is called
+ - `modifyglobal!(mod, var, op, x, order)` where `op(getval(), x)` is called
+ - `memoryrefmodify!(memref, op, x, order, boundscheck)` where `op(getval(), x)` is called
+ - `Intrinsics.atomic_pointermodify(ptr, op, x, order)` where `op(getval(), x)` is called
+ - `Core._task(f, size)` where `f()` will be called when the task runs
+ - `Core.finalizer(f, obj)` where `f(obj)` will be called during garbage collection
+
+Contains the `CallInfo` for the indirect function call, its effects, and whether
+the indirect call should contribute edges for invalidation tracking.
+"""
+struct IndirectCallInfo <: CallInfo
+    info::CallInfo   # the callinfo for the indirect function call
+    effects::Effects # the effects for the indirect function call
+    add_edges::Bool  # whether to add edges for invalidation tracking
+end
+function add_edges_impl(edges::Vector{Any}, info::IndirectCallInfo)
+    if info.add_edges
+        add_edges!(edges, info.info)
+    end
+    # otherwise add no edges (e.g., for finalizers that don't imply edges unless inlined)
 end
-add_edges_impl(edges::Vector{Any}, info::ModifyOpInfo) = add_edges!(edges, info.info)
 
 struct VirtualMethodMatchInfo <: CallInfo
     info::Union{MethodMatchInfo,UnionSplitInfo,InvokeCallInfo}

Compiler/src/tfuncs.jl

@@ -579,6 +579,7 @@ add_tfunc(Core.sizeof, 1, 1, sizeof_tfunc, 1)
 end
 add_tfunc(nfields, 1, 1, nfields_tfunc, 1)
 add_tfunc(Core._expr, 1, INT_INF, @nospecs((𝕃::AbstractLattice, args...)->Expr), 100)
+
 add_tfunc(svec, 0, INT_INF, @nospecs((𝕃::AbstractLattice, args...)->SimpleVector), 20)
 @nospecs function _svec_ref_tfunc(𝕃::AbstractLattice, s, i)
     if isa(s, Const) && isa(i, Const)
@@ -1143,6 +1144,8 @@ end
             end
         end
         s00 = s
+    elseif isa(s00, PartialTask)
+        s00 = Task
     end
     return _getfield_tfunc(widenlattice(𝕃), s00, name, setfield)
 end
@@ -1436,7 +1439,7 @@ end
             elseif isconcretetype(RT) && has_nontrivial_extended_info(𝕃ᵢ, TF2) # isconcrete condition required to form a PartialStruct
                 RT = PartialStruct(fallback_lattice, RT, Union{Nothing,Bool}[false,false], Any[TF, TF2])
             end
-            info = ModifyOpInfo(callinfo.info)
+            info = IndirectCallInfo(callinfo.info, callinfo.effects, true)
             return CallMeta(RT, Any, Effects(), info)
         end
     end
@@ -2346,6 +2349,7 @@ const _PURE_BUILTINS = Any[
     tuple,
     svec,
     ===,
+    Core.task_result_type,
     typeof,
     nfields,
 ]
@@ -2407,6 +2411,7 @@ const _INACCESSIBLEMEM_BUILTINS = Any[
     fieldtype,
     isa,
     nfields,
+    Core.task_result_type,
     throw,
     Core.throw_methoderror,
     tuple,
@@ -2615,6 +2620,7 @@ const _EFFECTS_KNOWN_BUILTINS = Any[
     # setglobalonce!,
     swapfield!,
     # swapglobal!,
+    Core.task_result_type,
     throw,
     tuple,
     typeassert,
@@ -3050,7 +3056,7 @@ function intrinsic_effects(f::IntrinsicFunction, argtypes::Vector{Any})
         # llvmcall can do arbitrary things
         return Effects()
     elseif f === atomic_pointermodify
-        # atomic_pointermodify has memory effects, plus any effects from the ModifyOpInfo
+        # atomic_pointermodify has memory effects, plus any effects from the IndirectCallInfo
         return Effects()
     end
     is_effect_free = _is_effect_free_infer(f)
@@ -3262,6 +3268,15 @@ add_tfunc(replaceglobal!, 4, 6, @nospecs((𝕃::AbstractLattice, args...)->Any),
 add_tfunc(setglobalonce!, 3, 5, @nospecs((𝕃::AbstractLattice, args...)->Bool), 3)
 add_tfunc(Core.get_binding_type, 2, 2, @nospecs((𝕃::AbstractLattice, args...)->Type), 0)
 
+@nospecs function task_result_type_tfunc(𝕃::AbstractLattice, T)
+    hasintersect(widenconst(T), Task) || return Union{}
+    if T isa PartialTask
+        return Type{widenconst(T.fetch_type)}
+    end
+    return Type
+end
+add_tfunc(Core.task_result_type, 1, 1, task_result_type_tfunc, 0)
+
 # foreigncall
 # ===========