Refactor lattice code to expose layering and enable easy extension

There's been two threads of work involving the compiler's notion of
the inference lattice. One is that the lattice has gotten to complicated
and with too many internal constraints that are not manifest in the
type system. #42596 attempted to address this, but it's quite disruptive
as it changes the lattice types and all the signatures of the lattice
operations, which are used quite extensively throughout the ecosystem
(despite being internal), so that change is quite disruptive (and
something we'd ideally only make the ecosystem do once).

The other thread of work is that people would like to experiment with
a variety of extended lattices outside of base (either to prototype
potential additions to the lattice in base or to do custom abstract
interpretation over the Julia code). At the moment, the lattice is
quite closely interwoven with the rest of the abstract interpreter.
In response to this request in #40992, I had proposed a `CustomLattice`
element with callbacks, but this doesn't compose particularly well,
is cumbersome and imposes overhead on some of the hottest parts of
the compiler, so it's a bit of a tough sell to merge into `Base`.

In this PR, I'd like to propose a refactoring that is relatively
non-invasive to non-Base users, but I think would allow easier
experimentation with changes to the lattice for these two use
cases. In essence, we're splitting the lattice into a ladder of
5 different lattices, each containing the previous lattice as a
sub-lattice. These 5 lattices are:

- `JLTypeLattice` (Anything that's a `Type`)
- `ConstsLattice` ( + `Const`, `PartialTypeVar`)
- `PartialsLattice` ( + `PartialStruct`, `PartialOpaque` )
- `ConditionalsLattice` ( + `Conditional` )
- `InferenceLattice` ( + `LimitedAccuracy` )
- `OptimizerLattice` ( + `MaybeUndef` )

The idea is that where a lattice element contains another lattice
element (e.g. in `PartialStruct` or `Conditional`), the element
contained may only be from a wider lattice. In this PR, this
is not enforced by the type system. This is quite deliberate, as
I want to retain the types and object layouts of the lattice elements,
but of course a future #42596-like change could add such type
enforcement.

Of particular note is that the `PartialsLattice` and `ConditionalsLattice`
is parameterized and additional layers may be added in the stack.
For example, in #40992, I had proposed a lattice element that refines
`Int` and tracks symbolic expressions. In this setup, this could
be accomplished by adding an appropriate lattice in between the
`ConstsLattice` and the `PartialsLattice` (of course, additional
hooks would be required to make the tfuncs work, but that is
outside the scope of this PR).

I don't think this is a full solution, but I think it'll help us
play with some of these extended lattice options over the next
6-12 months in the packages that want to do this sort of thing.
Presumably once we know what all the potential lattice extensions
look like, we will want to take another look at this (likely
together with whatever solution we come up with for the
AbstractInterpreter composability problem and a rebase of #42596).

WIP because I didn't bother updating and plumbing through the lattice
in all the call sites yet, but that's mostly mechanical, so if we
like this direction, I will make that change and hope to merge this
in short order (because otherwise it'll accumulate massive merge
conflicts).

Author: Keno

base/compiler/abstractinterpretation.jl

@@ -0,0 +1,173 @@
+abstract type AbstractLattice; end
+function widenlattice end
+
+"""
+    struct JLTypeLattice
+
+A singleton type representing the lattice of Julia types, without any inference
+extensions.
+"""
+struct JLTypeLattice <: AbstractLattice; end
+widenlattice(::JLTypeLattice) = error("Type lattice is the least-precise lattice available")
+is_valid_lattice(::JLTypeLattice, @nospecialize(elem)) = isa(elem, Type)
+
+"""
+    struct ConstsLattice
+
+A lattice extending `JLTypeLattice` and adjoining `Const` and `PartialTypeVar`.
+"""
+struct ConstsLattice <: AbstractLattice; end
+widenlattice(::ConstsLattice) = JLTypeLattice()
+is_valid_lattice(lattice::ConstsLattice, @nospecialize(elem)) =
+    is_valid_lattice(widenlattice(lattice), elem) || isa(elem, Const) || isa(elem, PartialTypeVar)
+
+"""
+    struct PartialsLattice{L}
+
+A lattice extending lattice `L` and adjoining `PartialStruct` and `PartialOpaque`.
+"""
+struct PartialsLattice{L <: AbstractLattice} <: AbstractLattice
+    parent::L
+end
+widenlattice(L::PartialsLattice) = L.parent
+is_valid_lattice(lattice::PartialsLattice, @nospecialize(elem)) =
+    is_valid_lattice(widenlattice(lattice), elem) ||
+    isa(elem, PartialStruct) || isa(elem, PartialOpaque)
+
+"""
+    struct ConditionalsLattice{L}
+
+A lattice extending lattice `L` and adjoining `Conditional`.
+"""
+struct ConditionalsLattice{L <: AbstractLattice} <: AbstractLattice
+    parent::L
+end
+widenlattice(L::ConditionalsLattice) = L.parent
+is_valid_lattice(lattice::ConditionalsLattice, @nospecialize(elem)) =
+    is_valid_lattice(widenlattice(lattice), elem) || isa(elem, Conditional)
+
+struct InterConditionalsLattice{L <: AbstractLattice} <: AbstractLattice
+    parent::L
+end
+widenlattice(L::InterConditionalsLattice) = L.parent
+is_valid_lattice(lattice::InterConditionalsLattice, @nospecialize(elem)) =
+    is_valid_lattice(widenlattice(lattice), elem) || isa(elem, InterConditional)
+
+const AnyConditionalsLattice{L} = Union{ConditionalsLattice{L}, InterConditionalsLattice{L}}
+const BaseInferenceLattice = typeof(ConditionalsLattice(PartialsLattice(ConstsLattice())))
+const IPOResultLattice = typeof(InterConditionalsLattice(PartialsLattice(ConstsLattice())))
+
+"""
+    struct InferenceLattice{L}
+
+The full lattice used for abstract interpration during inference. Takes
+a base lattice and adjoins `LimitedAccuracy`.
+"""
+struct InferenceLattice{L} <: AbstractLattice
+    parent::L
+end
+widenlattice(L::InferenceLattice) = L.parent
+is_valid_lattice(lattice::InferenceLattice, @nospecialize(elem)) =
+    is_valid_lattice(widenlattice(lattice), elem) || isa(elem, LimitedAccuracy)
+
+"""
+    struct OptimizerLattice
+
+The lattice used by the optimizer. Extends
+`BaseInferenceLattice` with `MaybeUndef`.
+"""
+struct OptimizerLattice <: AbstractLattice; end
+widenlattice(L::OptimizerLattice) = BaseInferenceLattice.instance
+is_valid_lattice(lattice::OptimizerLattice, @nospecialize(elem)) =
+    is_valid_lattice(widenlattice(lattice), elem) || isa(elem, MaybeUndef)
+
+"""
+    tmeet(lattice, a, b::Type)
+
+Compute the lattice meet of lattice elements `a` and `b` over the lattice
+`lattice`. If `lattice` is `JLTypeLattice`, this is equiavalent to type
+intersection. Note that currently `b` is restricted to being a type (interpreted
+as a lattice element in the JLTypeLattice sub-lattice of `lattice`).
+"""
+function tmeet end
+
+function tmeet(::JLTypeLattice, @nospecialize(a::Type), @nospecialize(b::Type))
+    ti = typeintersect(a, b)
+    valid_as_lattice(ti) || return Bottom
+    return ti
+end
+
+"""
+    tmerge(lattice, a, b)
+
+Compute a lattice join of elements `a` and `b` over the lattice `lattice`.
+Note that the computed element need not be the least upper bound of `a` and
+`b`, but rather, we impose additional limitations on the complexity of the
+joined element, ideally without losing too much precision in common cases and
+remaining mostly associative and commutative.
+"""
+function tmerge end
+
+"""
+    ⊑(lattice, a, b)
+
+Compute the lattice ordering (i.e. less-than-or-equal) relationship between
+lattice elements `a` and `b` over the lattice `lattice`. If `lattice` is
+`JLTypeLattice`, this is equiavalent to subtyping.
+"""
+function ⊑ end
+
+⊑(::JLTypeLattice, @nospecialize(a::Type), @nospecialize(b::Type)) = a <: b
+
+"""
+    ⊏(lattice, a, b) -> Bool
+
+The strict partial order over the type inference lattice.
+This is defined as the irreflexive kernel of `⊑`.
+"""
+⊏(lattice::AbstractLattice, @nospecialize(a), @nospecialize(b)) = ⊑(lattice, a, b) && !⊑(lattice, b, a)
+
+"""
+    ⋤(lattice, a, b) -> Bool
+
+This order could be used as a slightly more efficient version of the strict order `⊏`,
+where we can safely assume `a ⊑ b` holds.
+"""
+⋤(lattice::AbstractLattice, @nospecialize(a), @nospecialize(b)) = !⊑(lattice, b, a)
+
+"""
+    is_lattice_equal(lattice, a, b) -> Bool
+
+Check if two lattice elements are partial order equivalent.
+This is basically `a ⊑ b && b ⊑ a` but (optionally) with extra performance optimizations.
+"""
+function is_lattice_equal(lattice::AbstractLattice, @nospecialize(a), @nospecialize(b))
+    a === b && return true
+    ⊑(lattice, a, b) && ⊑(lattice, b, a)
+end
+
+"""
+    has_nontrivial_const_info(lattice, t) -> Bool
+
+Determine whether the given lattice element `t` of `lattice` has non-trivial
+constant information that would not be available from the type itself.
+"""
+has_nontrivial_const_info(lattice::AbstractLattice, @nospecialize t) =
+    has_nontrivial_const_info(widenlattice(lattice), t)
+has_nontrivial_const_info(::JLTypeLattice, @nospecialize(t)) = false
+
+# Curried versions
+⊑(lattice::AbstractLattice) = (@nospecialize(a), @nospecialize(b)) -> ⊑(lattice, a, b)
+⊏(lattice::AbstractLattice) = (@nospecialize(a), @nospecialize(b)) -> ⊏(lattice, a, b)
+⋤(lattice::AbstractLattice) = (@nospecialize(a), @nospecialize(b)) -> ⋤(lattice, a, b)
+
+# Fallbacks for external packages using these methods
+const fallback_lattice = InferenceLattice(BaseInferenceLattice.instance)
+const fallback_ipo_lattice = InferenceLattice(IPOResultLattice.instance)
+
+⊑(@nospecialize(a), @nospecialize(b)) = ⊑(fallback_lattice, a, b)
+tmeet(@nospecialize(a), @nospecialize(b)) = tmeet(fallback_lattice, a, b)
+tmerge(@nospecialize(a), @nospecialize(b)) = tmerge(fallback_lattice, a, b)
+⊏(@nospecialize(a), @nospecialize(b)) = ⊏(fallback_lattice, a, b)
+⋤(@nospecialize(a), @nospecialize(b)) = ⋤(fallback_lattice, a, b)
+is_lattice_equal(@nospecialize(a), @nospecialize(b)) = is_lattice_equal(fallback_lattice, a, b)

base/compiler/abstractlattice.jl

@@ -153,6 +153,8 @@ include("compiler/ssair/basicblock.jl")
 include("compiler/ssair/domtree.jl")
 include("compiler/ssair/ir.jl")
 
+include("compiler/abstractlattice.jl")
+
 include("compiler/inferenceresult.jl")
 include("compiler/inferencestate.jl")
 

base/compiler/compiler.jl

@@ -1,10 +1,11 @@
 # This file is a part of Julia. License is MIT: https://julialang.org/license
 
-function is_argtype_match(@nospecialize(given_argtype),
+function is_argtype_match(lattice::AbstractLattice,
+                          @nospecialize(given_argtype),
                           @nospecialize(cache_argtype),
                           overridden_by_const::Bool)
     if is_forwardable_argtype(given_argtype)
-        return is_lattice_equal(given_argtype, cache_argtype)
+        return is_lattice_equal(lattice, given_argtype, cache_argtype)
     end
     return !overridden_by_const
 end
@@ -91,7 +92,7 @@ function matching_cache_argtypes(
     for i in 1:nargs
         given_argtype = given_argtypes[i]
         cache_argtype = cache_argtypes[i]
-        if !is_argtype_match(given_argtype, cache_argtype, false)
+        if !is_argtype_match(fallback_lattice, given_argtype, cache_argtype, false)
             # prefer the argtype we were given over the one computed from `linfo`
             cache_argtypes[i] = given_argtype
             overridden_by_const[i] = true
@@ -207,7 +208,7 @@ function matching_cache_argtypes(linfo::MethodInstance, ::Nothing)
     return cache_argtypes, falses(length(cache_argtypes))
 end
 
-function cache_lookup(linfo::MethodInstance, given_argtypes::Vector{Any}, cache::Vector{InferenceResult})
+function cache_lookup(lattice::AbstractLattice, linfo::MethodInstance, given_argtypes::Vector{Any}, cache::Vector{InferenceResult})
     method = linfo.def::Method
     nargs::Int = method.nargs
     method.isva && (nargs -= 1)
@@ -218,15 +219,15 @@ function cache_lookup(linfo::MethodInstance, given_argtypes::Vector{Any}, cache:
         cache_argtypes = cached_result.argtypes
         cache_overridden_by_const = cached_result.overridden_by_const
         for i in 1:nargs
-            if !is_argtype_match(given_argtypes[i],
+            if !is_argtype_match(lattice, given_argtypes[i],
                                  cache_argtypes[i],
                                  cache_overridden_by_const[i])
                 cache_match = false
                 break
             end
         end
         if method.isva && cache_match
-            cache_match = is_argtype_match(tuple_tfunc(given_argtypes[(nargs + 1):end]),
+            cache_match = is_argtype_match(lattice, tuple_tfunc(lattice, given_argtypes[(nargs + 1):end]),
                                            cache_argtypes[end],
                                            cache_overridden_by_const[end])
         end

base/compiler/inferenceresult.jl

@@ -92,7 +92,7 @@ function inlining_policy(interp::AbstractInterpreter, @nospecialize(src), stmt_f
         # inferred source in the local cache
         # we still won't find a source for recursive call because the "single-level" inlining
         # seems to be more trouble and complex than it's worth
-        inf_result = cache_lookup(mi, argtypes, get_inference_cache(interp))
+        inf_result = cache_lookup(optimizer_lattice(interp), mi, argtypes, get_inference_cache(interp))
         inf_result === nothing && return nothing
         src = inf_result.src
         if isa(src, CodeInfo)
@@ -215,7 +215,7 @@ function stmt_effect_flags(@nospecialize(stmt), @nospecialize(rt), src::Union{IR
     isa(stmt, PhiNode) && return (true, true)
     isa(stmt, ReturnNode) && return (false, true)
     isa(stmt, GotoNode) && return (false, true)
-    isa(stmt, GotoIfNot) && return (false, argextype(stmt.cond, src) ⊑ Bool)
+    isa(stmt, GotoIfNot) && return (false, argextype(stmt.cond, src) ⊑ₒ Bool)
     isa(stmt, Slot) && return (false, false) # Slots shouldn't occur in the IR at this point, but let's be defensive here
     if isa(stmt, GlobalRef)
         nothrow = isdefined(stmt.mod, stmt.name)
@@ -248,7 +248,7 @@ function stmt_effect_flags(@nospecialize(stmt), @nospecialize(rt), src::Union{IR
                 return (total, total)
             end
             rt === Bottom && return (false, false)
-            nothrow = _builtin_nothrow(f, Any[argextype(args[i], src) for i = 2:length(args)], rt)
+            nothrow = _builtin_nothrow(f, Any[argextype(args[i], src) for i = 2:length(args)], rt, OptimizerLattice())
             nothrow || return (false, false)
             return (contains_is(_EFFECT_FREE_BUILTINS, f), nothrow)
         elseif head === :new
@@ -262,7 +262,7 @@ function stmt_effect_flags(@nospecialize(stmt), @nospecialize(rt), src::Union{IR
             for fld_idx in 1:(length(args) - 1)
                 eT = argextype(args[fld_idx + 1], src)
                 fT = fieldtype(typ, fld_idx)
-                eT ⊑ fT || return (false, false)
+                eT ⊑ₒ fT || return (false, false)
             end
             return (true, true)
         elseif head === :foreigncall
@@ -277,11 +277,11 @@ function stmt_effect_flags(@nospecialize(stmt), @nospecialize(rt), src::Union{IR
             typ = argextype(args[1], src)
             typ, isexact = instanceof_tfunc(typ)
             isexact || return (false, false)
-            typ ⊑ Tuple || return (false, false)
+            typ ⊑ₒ Tuple || return (false, false)
             rt_lb = argextype(args[2], src)
             rt_ub = argextype(args[3], src)
             source = argextype(args[4], src)
-            if !(rt_lb ⊑ Type && rt_ub ⊑ Type && source ⊑ Method)
+            if !(rt_lb ⊑ₒ Type && rt_ub ⊑ₒ Type && source ⊑ₒ Method)
                 return (false, false)
             end
             return (true, true)
@@ -448,7 +448,7 @@ function finish(interp::AbstractInterpreter, opt::OptimizationState,
         else
             # compute the cost (size) of inlining this code
             cost_threshold = default = params.inline_cost_threshold
-            if result ⊑ Tuple && !isconcretetype(widenconst(result))
+            if ⊑(optimizer_lattice(interp), result, Tuple) && !isconcretetype(widenconst(result))
                 cost_threshold += params.inline_tupleret_bonus
             end
             # if the method is declared as `@inline`, increase the cost threshold 20x

base/compiler/optimize.jl

@@ -114,6 +114,8 @@ function CFGInliningState(ir::IRCode)
     )
 end
 
+⊑ₒ(@nospecialize(a), @nospecialize(b)) = ⊑(OptimizerLattice(), a, b)
+
 # Tells the inliner that we're now inlining into block `block`, meaning
 # all previous blocks have been processed and can be added to the new cfg
 function inline_into_block!(state::CFGInliningState, block::Int)
@@ -381,7 +383,7 @@ function ir_inline_item!(compact::IncrementalCompact, idx::Int, argexprs::Vector
             nonva_args = argexprs[1:end-1]
             va_arg = argexprs[end]
             tuple_call = Expr(:call, TOP_TUPLE, def, nonva_args...)
-            tuple_type = tuple_tfunc(Any[argextype(arg, compact) for arg in nonva_args])
+            tuple_type = tuple_tfunc(OptimizerLattice(), Any[argextype(arg, compact) for arg in nonva_args])
             tupl = insert_node_here!(compact, NewInstruction(tuple_call, tuple_type, topline))
             apply_iter_expr = Expr(:call, Core._apply_iterate, iterate, Core._compute_sparams, tupl, va_arg)
             sparam_vals = insert_node_here!(compact,
@@ -476,7 +478,7 @@ function fix_va_argexprs!(compact::IncrementalCompact,
         push!(tuple_call.args, arg)
         push!(tuple_typs, argextype(arg, compact))
     end
-    tuple_typ = tuple_tfunc(tuple_typs)
+    tuple_typ = tuple_tfunc(OptimizerLattice(), tuple_typs)
     tuple_inst = NewInstruction(tuple_call, tuple_typ, line_idx)
     push!(newargexprs, insert_node_here!(compact, tuple_inst))
     return newargexprs
@@ -1080,8 +1082,8 @@ function inline_apply!(
             nonempty_idx = 0
             for i = (arg_start + 1):length(argtypes)
                 ti = argtypes[i]
-                ti ⊑ Tuple{} && continue
-                if ti ⊑ Tuple && nonempty_idx == 0
+                ti ⊑ₒ Tuple{} && continue
+                if ti ⊑ₒ Tuple && nonempty_idx == 0
                     nonempty_idx = i
                     continue
                 end
@@ -1123,9 +1125,9 @@ end
 # TODO: this test is wrong if we start to handle Unions of function types later
 is_builtin(s::Signature) =
     isa(s.f, IntrinsicFunction) ||
-    s.ft ⊑ IntrinsicFunction ||
+    s.ft ⊑ₒ IntrinsicFunction ||
     isa(s.f, Builtin) ||
-    s.ft ⊑ Builtin
+    s.ft ⊑ₒ Builtin
 
 function inline_invoke!(
     ir::IRCode, idx::Int, stmt::Expr, info::InvokeCallInfo, flag::UInt8,
@@ -1165,7 +1167,7 @@ function narrow_opaque_closure!(ir::IRCode, stmt::Expr, @nospecialize(info), sta
         ub, exact = instanceof_tfunc(ubt)
         exact || return
         # Narrow opaque closure type
-        newT = widenconst(tmeet(tmerge(lb, info.unspec.rt), ub))
+        newT = widenconst(tmeet(OptimizerLattice(), tmerge(OptimizerLattice(), lb, info.unspec.rt), ub))
         if newT != ub
             # N.B.: Narrowing the ub requires a backdge on the mi whose type
             # information we're using, since a change in that function may
@@ -1222,7 +1224,7 @@ function process_simple!(ir::IRCode, idx::Int, state::InliningState, todo::Vecto
         ir.stmts[idx][:inst] = earlyres.val
         return nothing
     end
-    if (sig.f === modifyfield! || sig.ft ⊑ typeof(modifyfield!)) && 5 <= length(stmt.args) <= 6
+    if (sig.f === modifyfield! || sig.ft ⊑ₒ typeof(modifyfield!)) && 5 <= length(stmt.args) <= 6
         let info = ir.stmts[idx][:info]
             info isa MethodResultPure && (info = info.info)
             info isa ConstCallInfo && (info = info.call)
@@ -1240,7 +1242,7 @@ function process_simple!(ir::IRCode, idx::Int, state::InliningState, todo::Vecto
     end
 
     if check_effect_free!(ir, idx, stmt, rt)
-        if sig.f === typeassert || sig.ft ⊑ typeof(typeassert)
+        if sig.f === typeassert || sig.ft ⊑ₒ typeof(typeassert)
             # typeassert is a no-op if effect free
             ir.stmts[idx][:inst] = stmt.args[2]
             return nothing
@@ -1648,7 +1650,7 @@ function early_inline_special_case(
         elseif ispuretopfunction(f) || contains_is(_PURE_BUILTINS, f)
             return SomeCase(quoted(val))
         elseif contains_is(_EFFECT_FREE_BUILTINS, f)
-            if _builtin_nothrow(f, argtypes[2:end], type)
+            if _builtin_nothrow(f, argtypes[2:end], type, OptimizerLattice())
                 return SomeCase(quoted(val))
             end
         elseif f === Core.get_binding_type
@@ -1694,17 +1696,17 @@ function late_inline_special_case!(
     elseif length(argtypes) == 3 && istopfunction(f, :(>:))
         # special-case inliner for issupertype
         # that works, even though inference generally avoids inferring the `>:` Method
-        if isa(type, Const) && _builtin_nothrow(<:, Any[argtypes[3], argtypes[2]], type)
+        if isa(type, Const) && _builtin_nothrow(<:, Any[argtypes[3], argtypes[2]], type, OptimizerLattice())
             return SomeCase(quoted(type.val))
         end
         subtype_call = Expr(:call, GlobalRef(Core, :(<:)), stmt.args[3], stmt.args[2])
         return SomeCase(subtype_call)
-    elseif f === TypeVar && 2 <= length(argtypes) <= 4 && (argtypes[2] ⊑ Symbol)
+    elseif f === TypeVar && 2 <= length(argtypes) <= 4 && (argtypes[2] ⊑ₒ Symbol)
         typevar_call = Expr(:call, GlobalRef(Core, :_typevar), stmt.args[2],
             length(stmt.args) < 4 ? Bottom : stmt.args[3],
             length(stmt.args) == 2 ? Any : stmt.args[end])
         return SomeCase(typevar_call)
-    elseif f === UnionAll && length(argtypes) == 3 && (argtypes[2] ⊑ TypeVar)
+    elseif f === UnionAll && length(argtypes) == 3 && (argtypes[2] ⊑ₒ TypeVar)
         unionall_call = Expr(:foreigncall, QuoteNode(:jl_type_unionall), Any, svec(Any, Any),
             0, QuoteNode(:ccall), stmt.args[2], stmt.args[3])
         return SomeCase(unionall_call)