Refactor lattice code to expose layering and enable easy extension

[Keno]

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 )
• ConditionalsLattice ( + Conditional )
• InferenceLattice ( + LimitedAccuracy, 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).

[aviatesk]

I think I like this -- this should allow us to implement base lattice operations with more composability and provide a way to extend them, in a way we keep type stability.
I guess from external users' point of view, it may still require some understanding of the base lattice design to actually extend it (e.g. we need to understand the implicit order of InferenceLattice() = ConditionalsLattice(PartialsLattice(ConstsLattice()))) if we want to implement a customized the abstract interpretation without losing the original features), but it'd be much more useful to give an opportunity rather than staying with the current state.

I will make that change and hope to merge this
in short order (because otherwise it'll accumulate massive merge
conflicts).

Sounds good. I also want to play with this and may add some example test cases on top of it.

[Keno]

@nanosoldier runtests(ALL, vs = ":master")

[Keno]

I think this is in reasonable shape. There's a few more things to do, but the bulk of the changes is done, so I'll do the usual PkgEval, etc and then I think we can merge this and do the remaining tweaks as follow-on to minimize the rebase load on other in-progress PRs.

[nanosoldier]

Your package evaluation job has completed - possible new issues were detected. A full report can be found here.

[Keno]

@nanosoldier runtests(["AWSBatch", "AlgebraOfGraphics", "AoGExtensions", "Bagyo", "Biplots", "CairoMakie", "CineFiles", "ClimateModels", "ClusterManagers", "ClusteredLowRankSolver", "CombinatorialSpaces", "ConScape", "ConceptnetNumberbatch", "CountdownNumbers", "DiracNotation", "EarlyStopping", "FlameGraphs", "Folds", "FunSQL", "FunctionOperators", "GadgetIO", "GraphMakie", "GridLayoutBase", "IncompressibleNavierStokes", "InteractiveDynamics", "Kahuna", "LatticeDiracOperators", "LatticeQCD", "Lighthouse", "LogicToolkit", "Makie", "MakieLayout", "Mamba", "MeasureBase", "Meshes", "MetapopulationDynamics", "MolecularGraph", "Mux", "MzPlots", "NeuralQuantumState", "Pitaya", "PrettyPrinting", "QML", "RegularizedLeastSquares", "ReliabilityDiagrams", "RetroCap", "SpinGlassEngine", "SpinGlassTensors", "StrBase", "Strs", "SuiteSparseGraphBLAS", "SwagUI", "Syslogs", "SystemBenchmark", "TensorOperations", "TextParse", "TheNumberLine", "TrueRandom", "TypeDBClient", "Yunir"], vs = ":master")

[Keno]

@nanosoldier runbenchmarks("inference", vs=":master")

[nanosoldier]

Your benchmark job has completed - possible performance regressions were detected. A full report can be found here.

[nanosoldier]

Your package evaluation job has completed - possible new issues were detected. A full report can be found here.

[Keno]

@nanosoldier runbenchmarks("inference", vs=":master")

[Keno]

@nanosoldier runtests(["CairoMakie", "CineFiles", "ClusterManagers", "ConceptnetNumberbatch", "CountdownNumbers", "FlameGraphs", "Folds", "Kahuna", "LogicToolkit", "RetroCap", "StrBase", "SuiteSparseGraphBLAS", "TrueRandom"], vs = ":master")

[nanosoldier]

Your benchmark job has completed - possible performance regressions were detected. A full report can be found here.

[Keno]

@nanosoldier runbenchmarks("inference", vs=":master")

[nanosoldier]

Your benchmark job has completed - possible performance regressions were detected. A full report can be found here.

[nanosoldier]

Your package evaluation job has completed - possible new issues were detected. A full report can be found here.

[vtjnash]

These limits are not heuristic, but rather they are form the strict basis for the definition of tmerge here, and are mandatory. In particular, the tmerge operation must simultaneously satisfy the lattice join condition over the lattice argument, and the global complexity condition, both giving a loose upper bound, which together create a least upper bound under the dual lattice constraints.

[Keno]

Maybe heuristic is the wrong word, but I don't think we make any efforts to actually ensure that the result we return is actually the least upper bound under the joint lattice. Computing that is quite hard.

[vtjnash]

Why is this not a lattice operation, given that it is itself just a call to ⊑?

[Keno]

It probably should be, I'll fix that up.

[vtjnash]

It feels a little odd to merge this with the existing function of this name

help?> widen
  widen(x)

  If x is a type, return a "larger" type, defined so that arithmetic operations + and - are guaranteed not to overflow nor lose precision for any combination of
  values that type x can hold.

[Keno]

I can do widenlattice instead.

[vtjnash]

This is only usually legal at a few explicit parts of the code (in IPO and typeinf_local), since we didn't implement handling for it in most of the rest of the code. I guess we theoretically could implement that handling though, so it is permitted here?

[Keno]

It had support in the various lattice operations, so I added a layer for it. Of course, it technically violates a whole bunch of lattice assumptions, but that's a different problem. I think it's fine as a lattice layer, though of course the optimizer does not use it.

[aviatesk]

As a target example I'm currently writing TaintInterpreter that is supposed to do a simple taint analysis using this lattice extension framework. And I found that in addition to several minor fixes on this PR, we will need the following additional efforts for this to work properly:

• pass AbstractInterpreter around in the optimization phase so that optimizer can work in the existence of user-extended lattice elements
• make all tfuncs accept AbstractLattice to allow external consumers to extend it appropriately
• make caching-related functions accept AbstractLattice (like CodeInstance()) to allow external consumers to cache customized IPO-information

TaintInterpreter

using Core: SlotNumber, Argument
using Core.Compiler: slot_id, tmerge_fast_path
import .CC:
    AbstractLattice, BaseInferenceLattice, IPOResultLattice, InferenceLattice, OptimizerLattice,
    widen, is_valid_lattice, typeinf_lattice, ipo_lattice, optimizer_lattice,
    widenconst, tmeet, tmerge, ⊑, abstract_eval_special_value, widenreturn

@newinterp TaintInterpreter
struct TaintLattice{PL<:AbstractLattice} <: CC.AbstractLattice
    parent::PL
end
CC.widen(𝕃::TaintLattice) = 𝕃.parent
CC.is_valid_lattice(𝕃::TaintLattice, @nospecialize(elm)) =
    is_valid_lattice(widen(𝕃), elem) || isa(elm, Taint)

struct InterTaintLattice{PL<:AbstractLattice} <: CC.AbstractLattice
    parent::PL
end
CC.widen(𝕃::InterTaintLattice) = 𝕃.parent
CC.is_valid_lattice(𝕃::InterTaintLattice, @nospecialize(elm)) =
    is_valid_lattice(widen(𝕃), elem) || isa(elm, InterTaint)

const AnyTaintLattice{L} = Union{TaintLattice{L},InterTaintLattice{L}}

CC.typeinf_lattice(::TaintInterpreter) = InferenceLattice(TaintLattice(BaseInferenceLattice.instance))
CC.ipo_lattice(::TaintInterpreter) = InferenceLattice(InterTaintLattice(IPOResultLattice.instance))
CC.optimizer_lattice(::TaintInterpreter) = InterTaintLattice(OptimizerLattice())

struct Taint
    typ
    slots::BitSet
    function Taint(@nospecialize(typ), slots::BitSet)
        if typ isa Taint
            slots = typ.slots ∪ slots
            typ = typ.typ
        end
        return new(typ, slots)
    end
end
Taint(@nospecialize(typ), id::Int) = Taint(typ, push!(BitSet(), id))

struct InterTaint
    typ
    slots::BitSet
    function InterTaint(@nospecialize(typ), slots::BitSet)
        if typ isa InterTaint
            slots = typ.slots ∪ slots
            typ = typ.typ
        end
        return new(typ, slots)
    end
end
InterTaint(@nospecialize(typ), id::Int) = InterTaint(typ, push!(BitSet(), id))

const AnyTaint = Union{Taint, InterTaint}

function CC.tmeet(𝕃::AnyTaintLattice, @nospecialize(v), @nospecialize(t::Type))
    T = isa(𝕃, TaintLattice) ? Taint : InterTaint
    if isa(v, T)
        v = v.typ
    end
    return tmeet(widen(𝕃), v, t)
end
function CC.tmerge(𝕃::AnyTaintLattice, @nospecialize(typea), @nospecialize(typeb))
    r = tmerge_fast_path(𝕃, typea, typeb)
    r !== nothing && return r
    # type-lattice for Taint
    T = isa(𝕃, TaintLattice) ? Taint : InterTaint
    if isa(typea, T)
        if isa(typeb, T)
            return T(
                tmerge(widen(𝕃), typea.typ, typeb),
                typea.slots ∪ typeb.slots)
        else
            typea = typea.typ
        end
    elseif isa(typeb, T)
        typeb = typeb.typ
    end
    return tmerge(widen(𝕃), typea, typeb)
end
function CC.:⊑(𝕃::AnyTaintLattice, @nospecialize(typea), @nospecialize(typeb))
    T = isa(𝕃, TaintLattice) ? Taint : InterTaint
    if isa(typea, T)
        if isa(typeb, T)
            typea.slots ⊆ typeb.slots || return false
            return ⊑(widen(𝕃), typea.typ, typeb.typ)
        end
        typea = typea.typ
    elseif isa(typeb, T)
        return false
    end
    return ⊑(widen(𝕃), typea, typeb)
end
CC.widenconst(taint::AnyTaint) = widenconst(taint.typ)

function CC.abstract_eval_special_value(interp::TaintInterpreter,
    @nospecialize(e), vtypes::CC.VarTable, sv::CC.InferenceState)
    ret = @invoke CC.abstract_eval_special_value(interp::CC.AbstractInterpreter,
        e::Any, vtypes::CC.VarTable, sv::CC.InferenceState)
    if isa(e, SlotNumber) || isa(e, Argument)
        return Taint(ret, slot_id(e))
    end
    return ret
end

function CC.widenreturn(𝕃::InferenceLattice{<:InterTaintLattice}, @nospecialize(rt), @nospecialize(bestguess), nargs::Int, slottypes::Vector{Any}, changes::CC.VarTable)
    if isa(rt, Taint)
        return InterTaint(rt.typ, BitSet((id for id in rt.slots if id ≤ nargs)))
    end
    return CC.widenreturn(widen(𝕃), rt, bestguess, nargs, slottypes, changes)
end

code_typed(ifelse, (Bool, Int, Int); interp=TaintInterpreter())

[Keno]

Yes, some of those changes are potential future work, I just wanted to get the lattice infrastructure in place first.

[Keno]

@nanosoldier runbenchmarks("inference", vs=":master")

[nanosoldier]

Your benchmark job has completed - possible performance regressions were detected. A full report can be found here.