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,57 @@ 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)
+    # Check argument count: _task(func, size) or _task(func, size, ci)
+    if la < 3 || la > 4
+        return Future(CallMeta(Bottom, Any, EFFECTS_THROWS, NoCallInfo()))
+    end
+    # Check that size argument is an Int
+    size_arg = argtypes[3]
+    if !(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)
+    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 = invoke_result.rt
+            fetch_error = 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
+
+    # Fallback to 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 = callinfo.rt
+        fetch_error = 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)

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)

Compiler/src/ssair/inlining.jl

@@ -1219,6 +1219,27 @@ 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 = 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
+    edge = info.edges[1]
+    edge === nothing && return nothing
+    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._task(func, size, ci)
+    push!(stmt.args, case.invoke)
+    ir[SSAValue(idx)][:stmt] = stmt
+    return nothing
 end
 
 # As a matter of convenience, this pass also computes effect-freenes.
@@ -1288,7 +1309,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,7 +1560,7 @@ 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)
+function handle_modifyop!_call!(ir::IRCode, idx::Int, stmt::Expr, info::IndirectCallInfo, state::InliningState)
     info = info.info
     info isa MethodResultPure && (info = info.info)
     info isa ConstCallInfo && (info = info.call)
@@ -1556,7 +1578,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
@@ -1651,14 +1673,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)

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}
@@ -502,4 +500,5 @@ function add_edges_impl(edges::Vector{Any}, info::GlobalAccessInfo)
     push!(edges, info.b)
 end
 
+
 @specialize

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,12 @@ end
             end
         end
         s00 = s
+    elseif isa(s00, PartialTask)
+        # Special case: accessing the 'result' field of a PartialTask returns the fetch_type or error_type or some other value set by the user
+        #if isa(name, Const) && name.val === :result
+        #    return s00.fetch_type
+        #end
+        s00 = Task
     end
     return _getfield_tfunc(widenlattice(𝕃), s00, name, setfield)
 end
@@ -1436,7 +1443,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
@@ -3050,7 +3057,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)

Compiler/src/typelattice.jl

@@ -4,9 +4,9 @@
 # structs/constants #
 #####################
 
-# N.B.: Const/PartialStruct/InterConditional are defined in Core, to allow them to be used
+# N.B.: Const/PartialStruct/InterConditional/PartialTask are defined in Core, to allow them to be used
 # inside the global code cache.
-import Core: Const, InterConditional, PartialStruct
+import Core: Const, InterConditional, PartialStruct, PartialTask
 
 function may_form_limited_typ(@nospecialize(aty), @nospecialize(bty), @nospecialize(xty))
     if aty isa LimitedAccuracy
@@ -514,6 +514,14 @@ end
     elseif isa(b, PartialOpaque)
         return false
     end
+    if isa(a, PartialTask)
+        if isa(b, PartialTask)
+            return ⊑(lattice, a.fetch_type, b.fetch_type) && ⊑(lattice, a.fetch_error, b.fetch_error)
+        end
+        return ⊑(widenlattice(lattice), Task, b)
+    elseif isa(b, PartialTask)
+        return false
+    end
     return ⊑(widenlattice(lattice), a, b)
 end
 
@@ -581,6 +589,11 @@ end
         return is_lattice_equal(lattice, a.env, b.env)
     end
     isa(b, PartialOpaque) && return false
+    if isa(a, PartialTask)
+        isa(b, PartialTask) || return false
+        return is_lattice_equal(lattice, a.fetch_type, b.fetch_type) && is_lattice_equal(lattice, a.fetch_error, b.fetch_error)
+    end
+    isa(b, PartialTask) && return false
     return is_lattice_equal(widenlattice(lattice), a, b)
 end
 
@@ -643,6 +656,18 @@ end
         ti = typeintersect(widev, t)
         valid_as_lattice(ti, true) || return Bottom
         return PartialOpaque(ti, v.env, v.parent, v.source)
+    elseif isa(v, PartialTask)
+        has_free_typevars(t) && return v
+        if Task <: t
+            return v
+        end
+        ti = typeintersect(Task, t)
+        valid_as_lattice(ti, true) || return Bottom
+        if ti === Task
+            return v
+        else
+            return Bottom  # PartialTask can only be a Task
+        end
     end
     return tmeet(widenlattice(lattice), v, t)
 end
@@ -704,6 +729,7 @@ widenconst(c::Const) = (v = c.val; isa(v, Type) ? Type{v} : typeof(v))
 widenconst(::PartialTypeVar) = TypeVar
 widenconst(t::Core.PartialStruct) = t.typ
 widenconst(t::PartialOpaque) = t.typ
+widenconst(t::PartialTask) = Task
 @nospecializeinfer widenconst(@nospecialize t::Type) = t
 widenconst(::TypeVar) = error("unhandled TypeVar")
 widenconst(::TypeofVararg) = error("unhandled Vararg")