Remove buggy linearization pass (#604)

There is only one bug in the base linearization pass that needs to be
handled explicitly, otherwise linearization is fully guaranteed for IR.

---------

Co-authored-by: Shuhei Kadowaki <[email protected]>

Author: vtjnash

src/interpret.jl

@@ -9,22 +9,6 @@ function lookup_var(frame, slot::SlotNumber)
     throw(UndefVarError(frame.framecode.src.slotnames[slot.id]))
 end
 
-function lookup_expr(frame, e::Expr)
-    head = e.head
-    head === :the_exception && return frame.framedata.last_exception[]
-    if head === :static_parameter
-        arg = e.args[1]::Int
-        if isassigned(frame.framedata.sparams, arg)
-            return frame.framedata.sparams[arg]
-        else
-            syms = sparam_syms(frame.framecode.scope::Method)
-            throw(UndefVarError(syms[arg]))
-        end
-    end
-    head === :boundscheck && length(e.args) == 0 && return true
-    error("invalid lookup expr ", e)
-end
-
 """
     rhs = @lookup(frame, node)
     rhs = @lookup(mod, frame, node)
@@ -67,6 +51,32 @@ macro lookup(args...)
     end
 end
 
+function lookup_expr(frame, e::Expr)
+    head = e.head
+    head === :the_exception && return frame.framedata.last_exception[]
+    if head === :static_parameter
+        arg = e.args[1]::Int
+        if isassigned(frame.framedata.sparams, arg)
+            return frame.framedata.sparams[arg]
+        else
+            syms = sparam_syms(frame.framecode.scope::Method)
+            throw(UndefVarError(syms[arg]))
+        end
+    end
+    head === :boundscheck && length(e.args) == 0 && return true
+    if head === :call
+        f = @lookup frame e.args[1]
+        if (@static VERSION < v"1.11.0-DEV.1180" && true) && f === Core.svec
+            # work around for a linearization bug in Julia (https://github.com/JuliaLang/julia/pull/52497)
+            return f(Any[@lookup(frame, e.args[i]) for i in 2:length(e.args)]...)
+        elseif f === Core.tuple
+            # handling for ccall literal syntax
+            return f(Any[@lookup(frame, e.args[i]) for i in 2:length(e.args)]...)
+        end
+    end
+    error("invalid lookup expr ", e)
+end
+
 # This is used only for new struct/abstract/primitive nodes.
 # The most important issue is that in these expressions, :call Exprs can be nested,
 # and hence our re-use of the `callargs` field of Frame would introduce
@@ -91,18 +101,26 @@ function lookup_or_eval(@nospecialize(recurse), frame, @nospecialize(node))
         if ex.head === :call
             f = ex.args[1]
             if f === Core.svec
-                return Core.svec(ex.args[2:end]...)
+                popfirst!(ex.args)
+                return Core.svec(ex.args...)
             elseif f === Core.apply_type
-                return Core.apply_type(ex.args[2:end]...)
-            elseif f === Core.typeof
-                return Core.typeof(ex.args[2])
-            elseif f === Base.getproperty
+                popfirst!(ex.args)
+                return Core.apply_type(ex.args...)
+            elseif f === typeof && length(ex.args) == 2
+                return typeof(ex.args[2])
+            elseif f === typeassert && length(ex.args) == 3
+                return typeassert(ex.args[2], ex.args[3])
+            elseif f === Base.getproperty && length(ex.args) == 3
                 return Base.getproperty(ex.args[2], ex.args[3])
+            elseif f === Core.Compiler.Val && length(ex.args) == 2
+                return Core.Compiler.Val(ex.args[2])
+            elseif f === Val && length(ex.args) == 2
+                return Val(ex.args[2])
             else
-                Base.invokelatest(error, "unknown call f ", f)
+                Base.invokelatest(error, "unknown call f introduced by ccall lowering ", f)
             end
         else
-            error("unknown expr ", ex)
+            return lookup_expr(frame, ex)
         end
     elseif isa(node, Int) || isa(node, Number)   # Number is slow, requires subtyping
         return node

src/optimize.jl

@@ -1,94 +1,5 @@
-const calllike = (:call, :foreigncall)
-
 const compiled_calls = Dict{Any,Any}()
 
-function extract_inner_call!(stmt::Expr, idx, once::Bool=false)
-    (stmt.head === :toplevel || stmt.head === :thunk) && return nothing
-    once |= stmt.head ∈ calllike
-    for (i, a) in enumerate(stmt.args)
-        isa(a, Expr) || continue
-        # Make sure we don't "damage" special syntax that requires literals
-        if i == 1 && stmt.head === :foreigncall
-            continue
-        end
-        if i == 2 && stmt.head === :call && stmt.args[1] === :cglobal
-            continue
-        end
-        ret = extract_inner_call!(a, idx, once) # doing this first extracts innermost calls
-        ret !== nothing && return ret
-        iscalllike = a.head ∈ calllike
-        if once && iscalllike
-            stmt.args[i] = NewSSAValue(idx)
-            return a
-        end
-    end
-    return nothing
-end
-
-function replace_ssa(stmt::Expr, ssalookup)
-    return Expr(stmt.head, Any[
-        if isa(a, SSAValue)
-            SSAValue(ssalookup[a.id])
-        elseif isa(a, NewSSAValue)
-            SSAValue(a.id)
-        elseif isa(a, Expr)
-            replace_ssa(a, ssalookup)
-        else
-            a
-        end
-        for a in stmt.args
-    ]...)
-end
-
-function renumber_ssa!(stmts::Vector{Any}, ssalookup)
-    # When updating jumps, when lines get split into multiple lines
-    # (see "Un-nest :call expressions" below), we need to jump to the first of them.
-    # Consequently we use the previous "old-code" offset and add one.
-    # Fixes #455.
-    jumplookup(l, idx) = idx > 1 ? l[idx-1] + 1 : idx
-
-    for (i, stmt) in enumerate(stmts)
-        if isa(stmt, GotoNode)
-            stmts[i] = GotoNode(jumplookup(ssalookup, stmt.label))
-        elseif isa(stmt, SSAValue)
-            stmts[i] = SSAValue(ssalookup[stmt.id])
-        elseif isa(stmt, NewSSAValue)
-            stmts[i] = SSAValue(stmt.id)
-        elseif isexpr(stmt, :enter)
-            stmt.args[end] = jumplookup(ssalookup, stmt.args[1]::Int)
-        elseif isa(stmt, Expr)
-            stmts[i] = replace_ssa(stmt, ssalookup)
-        elseif isa(stmt, GotoIfNot)
-            cond = stmt.cond
-            if isa(cond, SSAValue)
-                cond = SSAValue(ssalookup[cond.id])
-            end
-            stmts[i] = GotoIfNot(cond, jumplookup(ssalookup, stmt.dest))
-        elseif isa(stmt, ReturnNode)
-            val = stmt.val
-            if isa(val, SSAValue)
-                stmts[i] = ReturnNode(SSAValue(ssalookup[val.id]))
-            end
-        elseif @static (isdefined(Core.IR, :EnterNode) && true) && isa(stmt, Core.IR.EnterNode)
-            stmts[i] = Core.IR.EnterNode(jumplookup(ssalookup, stmt.catch_dest))
-        end
-    end
-    return stmts
-end
-
-function compute_ssa_mapping_delete_statements!(code::CodeInfo, stmts::Vector{Int})
-    stmts = unique!(sort!(stmts))
-    ssalookup = collect(1:length(codelocs(code)))
-    cnt = 1
-    for i in 1:length(stmts)
-        start = stmts[i] + 1
-        stop = i == length(stmts) ? length(codelocs(code)) : stmts[i+1]
-        ssalookup[start:stop] .-= cnt
-        cnt += 1
-    end
-    return ssalookup
-end
-
 # Pre-frame-construction lookup
 function lookup_stmt(stmts, arg)
     if isa(arg, SSAValue)
@@ -179,7 +90,8 @@ function optimize!(code::CodeInfo, scope)
 
     # Replace :llvmcall and :foreigncall with compiled variants. See
     # https://github.com/JuliaDebug/JuliaInterpreter.jl/issues/13#issuecomment-464880123
-    foreigncalls_idx = Int[]
+    # Insert the foreigncall wrappers at the updated idxs
+    methodtables = Vector{Union{Compiled,DispatchableMethod}}(undef, length(code.code))
     for (idx, stmt) in enumerate(code.code)
         # Foregincalls can be rhs of assignments
         if isexpr(stmt, :(=))
@@ -192,47 +104,16 @@ function optimize!(code::CodeInfo, scope)
                 if (arg1 === :llvmcall || lookup_stmt(code.code, arg1) === Base.llvmcall) && isempty(sparams) && scope isa Method
                     # Call via `invokelatest` to avoid compiling it until we need it
                     Base.invokelatest(build_compiled_llvmcall!, stmt, code, idx, evalmod)
-                    push!(foreigncalls_idx, idx)
+                    methodtables[idx] = Compiled()
                 end
             elseif stmt.head === :foreigncall && scope isa Method
                 # Call via `invokelatest` to avoid compiling it until we need it
                 Base.invokelatest(build_compiled_foreigncall!, stmt, code, sparams, evalmod)
-                push!(foreigncalls_idx, idx)
+                methodtables[idx] = Compiled()
             end
         end
     end
 
-    ## Un-nest :call expressions (so that there will be only one :call per line)
-    # This will allow us to re-use args-buffers rather than having to allocate new ones each time.
-    old_code, old_codelocs = code.code, codelocs(code)
-    code.code = new_code = eltype(old_code)[]
-    code.codelocs = new_codelocs = Int32[]
-    ssainc = fill(1, length(old_code))
-    for (i, stmt) in enumerate(old_code)
-        loc = old_codelocs[i]
-        if isa(stmt, Expr)
-            inner = extract_inner_call!(stmt, length(new_code)+1)
-            while inner !== nothing
-                push!(new_code, inner)
-                push!(new_codelocs, loc)
-                ssainc[i] += 1
-                inner = extract_inner_call!(stmt, length(new_code)+1)
-            end
-        end
-        push!(new_code, stmt)
-        push!(new_codelocs, loc)
-    end
-    # Fix all the SSAValues and GotoNodes
-    ssalookup = cumsum(ssainc)
-    renumber_ssa!(new_code, ssalookup)
-    code.ssavaluetypes = length(new_code)
-
-    # Insert the foreigncall wrappers at the updated idxs
-    methodtables = Vector{Union{Compiled,DispatchableMethod}}(undef, length(code.code))
-    for idx in foreigncalls_idx
-        methodtables[ssalookup[idx]] = Compiled()
-    end
-
     return code, methodtables
 end
 
@@ -255,7 +136,7 @@ function parametric_type_to_expr(@nospecialize(t::Type))
     return t
 end
 
-function build_compiled_llvmcall!(stmt::Expr, code, idx, evalmod)
+function build_compiled_llvmcall!(stmt::Expr, code::CodeInfo, idx::Int, evalmod::Module)
     # Run a mini-interpreter to extract the types
     framecode = FrameCode(CompiledCalls, code; optimize=false)
     frame = Frame(framecode, prepare_framedata(framecode, []))
@@ -292,9 +173,8 @@ function build_compiled_llvmcall!(stmt::Expr, code, idx, evalmod)
     append!(stmt.args, args)
 end
 
-
 # Handle :llvmcall & :foreigncall (issue #28)
-function build_compiled_foreigncall!(stmt::Expr, code, sparams::Vector{Symbol}, evalmod)
+function build_compiled_foreigncall!(stmt::Expr, code::CodeInfo, sparams::Vector{Symbol}, evalmod::Module)
     TVal = evalmod == Core.Compiler ? Core.Compiler.Val : Val
     cfunc, RetType, ArgType = lookup_stmt(code.code, stmt.args[1]), stmt.args[2], stmt.args[3]::SimpleVector
 

src/types.jl

@@ -6,11 +6,6 @@ which will cause all calls to be evaluated via the interpreter.
 struct Compiled end
 Base.similar(::Compiled, sz) = Compiled()  # to support similar(stack, 0)
 
-# A type used transiently in renumbering CodeInfo SSAValues (to distinguish a new SSAValue from an old one)
-struct NewSSAValue
-    id::Int
-end
-
 # Our own replacements for Core types. We need to do this to ensure we can tell the difference
 # between "data" (Core types) and "code" (our types) if we step into Core.Compiler
 struct SSAValue

src/utils.jl

@@ -85,6 +85,10 @@ function scan_ssa_use!(used::BitSet, @nospecialize(stmt))
     while iterval !== nothing
         useref, state = iterval
         val = Core.Compiler.getindex(useref)
+        if (@static VERSION < v"1.11.0-DEV.1180" && true) && isexpr(val, :call)
+            # work around for a linearization bug in Julia (https://github.com/JuliaLang/julia/pull/52497)
+            scan_ssa_use!(used, val)
+        end
         if isa(val, SSAValue)
             push!(used, val.id)
         end