Fanon2 (#2136)

* Fewer anonymous funcs

* minor cleanup

* cleanup

* fix

* Aggressively noinfer

* more type annotations

Author: wsmoses

src/absint.jl

@@ -1,7 +1,7 @@
 # Abstractly interpret julia from LLVM
 
 # Return (bool if could interpret, julia object interpreted to)
-function absint(arg::LLVM.Value, partial::Bool = false)
+function absint(@nospecialize(arg::LLVM.Value), partial::Bool = false)::Tuple{Bool,Any}
     if isa(arg, LLVM.BitCastInst) || isa(arg, LLVM.AddrSpaceCastInst)
         return absint(operands(arg)[1], partial)
     end
@@ -165,7 +165,7 @@ function absint(arg::LLVM.Value, partial::Bool = false)
     return (false, nothing)
 end
 
-function actual_size(@nospecialize(typ2))
+function actual_size(@nospecialize(typ2))::Int
     @static if VERSION < v"1.11-"
         if typ2 <: Array
             return sizeof(Ptr{Cvoid}) + 2 + 2 + 4 + 2 * sizeof(Csize_t) + sizeof(Csize_t)
@@ -184,10 +184,10 @@ function actual_size(@nospecialize(typ2))
     end
 end
 
-@inline function first_non_ghost(@nospecialize(typ2))
+@inline function first_non_ghost(@nospecialize(typ2))::Tuple{Int, Int}
     @static if VERSION < v"1.11-"
         if typ2 <: Array
-            return (1, typed_fieldtype(typ2, 1))
+            return (1, 0)
         end
     end
     fc = fieldcount(typ2)
@@ -204,7 +204,7 @@ end
     return (-1, 0)
 end
 
-function should_recurse(@nospecialize(typ2), arg_t, byref, dl)
+function should_recurse(@nospecialize(typ2), @nospecialize(arg_t::LLVM.LLVMType), byref::GPUCompiler.ArgumentCC, dl::LLVM.DataLayout)::Bool
     sz = sizeof(dl, arg_t)
     if byref != GPUCompiler.BITS_VALUE
         if sz != sizeof(Int)
@@ -228,7 +228,7 @@ function should_recurse(@nospecialize(typ2), arg_t, byref, dl)
     end
 end
 
-function get_base_and_offset(larg::LLVM.Value; offsetAllowed=true, inttoptr=false)::Tuple{LLVM.Value, Int}
+function get_base_and_offset(@nospecialize(larg::LLVM.Value); offsetAllowed::Bool=true, inttoptr::Bool=false)::Tuple{LLVM.Value, Int}
     offset = 0
     while true
         if isa(larg, LLVM.ConstantExpr)
@@ -277,7 +277,7 @@ function get_base_and_offset(larg::LLVM.Value; offsetAllowed=true, inttoptr=fals
 end
 
 function abs_typeof(
-    arg::LLVM.Value,
+    @nospecialize(arg::LLVM.Value),
     partial::Bool = false, seenphis=Set{LLVM.PHIInst}()
 )::Union{Tuple{Bool,Type,GPUCompiler.ArgumentCC},Tuple{Bool,Nothing,Nothing}}
     if isa(arg, LLVM.BitCastInst) || isa(arg, LLVM.AddrSpaceCastInst)
@@ -729,7 +729,7 @@ function abs_typeof(
     return (false, nothing, nothing)
 end
 
-function abs_cstring(arg::LLVM.Value)::Tuple{Bool,String}
+function abs_cstring(@nospecialize(arg::LLVM.Value))::Tuple{Bool,String}
     if isa(arg, ConstantExpr)
         ce = arg
 	    while isa(ce, ConstantExpr)

src/compiler.jl

@@ -1125,7 +1125,11 @@ struct Return2
 end
 
 function force_recompute!(mod::LLVM.Module)
-    for f in functions(mod), bb in blocks(f), inst in collect(instructions(bb))
+    for f in functions(mod), bb in blocks(f)
+    iter = LLVM.API.LLVMGetFirstInstruction(bb)
+    while iter != C_NULL
+        inst = LLVM.Instruction(iter)
+        iter = LLVM.API.LLVMGetNextInstruction(iter)
         if isa(inst, LLVM.LoadInst)
             has_loaded = false
             for u in LLVM.uses(inst)
@@ -1170,6 +1174,7 @@ function force_recompute!(mod::LLVM.Module)
             end
         end
     end
+    end
 end
 
 function permit_inlining!(f::LLVM.Function)
@@ -3275,7 +3280,7 @@ end
 # Enzyme compiler step
 ##
 
-function annotate!(mod, mode)
+function annotate!(mod::LLVM.Module)
     inactive = LLVM.StringAttribute("enzyme_inactive", "")
     active = LLVM.StringAttribute("enzyme_active", "")
     no_escaping_alloc = LLVM.StringAttribute("enzyme_no_escaping_allocation")
@@ -3891,7 +3896,7 @@ function enzyme_extract_world(fn::LLVM.Function)::UInt
     throw(AssertionError("Enzyme: could not find world in $(string(fn))"))
 end
 
-function enzyme_custom_extract_mi(orig::LLVM.Instruction, error::Bool = true)
+function enzyme_custom_extract_mi(orig::LLVM.CallInst, error::Bool = true)
     operand = LLVM.called_operand(orig)
     if isa(operand, LLVM.Function)
         return enzyme_custom_extract_mi(operand::LLVM.Function, error)
@@ -6144,7 +6149,7 @@ end
 
 using Random
 # returns arg, return
-function no_type_setting(@nospecialize(specTypes); world = nothing)
+function no_type_setting(@nospecialize(specTypes::Type{<:Tuple}); world = nothing)
     # Even though the julia type here is ptr{int8}, the actual data can be something else
     if specTypes.parameters[1] == typeof(Random.XoshiroSimd.xoshiro_bulk_simd)
         return (true, false)
@@ -7037,7 +7042,7 @@ end
     end
 
     # annotate
-    annotate!(mod, mode)
+    annotate!(mod)
     for name in ("gpu_report_exception", "report_exception")
         if haskey(functions(mod), name)
             exc = functions(mod)[name]
@@ -8012,9 +8017,6 @@ end
     ::Type{TapeType},
     args::Vararg{Any,N},
 ) where {RawCall,PT,FA,T,RT,TapeType,N,CC,width,returnPrimal}
-
-    JuliaContext() do ctx
-        Base.@_inline_meta
         F = eltype(FA)
         is_forward =
             CC <: AugmentedForwardThunk || CC <: ForwardModeThunk || CC <: PrimalErrorThunk
@@ -8263,6 +8265,10 @@ end
             i += 1
         end
 
+    ts_ctx = JuliaContext()
+    ctx = context(ts_ctx)
+    activate(ctx)
+    (ir, fn, combinedReturn) = try
 
         if is_adjoint
             NT = Tuple{ActiveRetTypes...}
@@ -8441,31 +8447,35 @@ end
 
         ir = string(mod)
         fn = LLVM.name(llvm_f)
+        (ir, fn, combinedReturn)
+    finally
+        deactivate(ctx)
+        dispose(ts_ctx)
+    end
 
-        @assert length(types) == length(ccexprs)
+    @assert length(types) == length(ccexprs)
 
 
-        if !(GPUCompiler.isghosttype(PT) || Core.Compiler.isconstType(PT))
-            return quote
-                Base.@_inline_meta
-                Base.llvmcall(
-                    ($ir, $fn),
-                    $combinedReturn,
-                    Tuple{$PT,$(types...)},
-                    fptr,
-                    $(ccexprs...),
-                )
-            end
-        else
-            return quote
-                Base.@_inline_meta
-                Base.llvmcall(
-                    ($ir, $fn),
-                    $combinedReturn,
-                    Tuple{$(types...)},
-                    $(ccexprs...),
-                )
-            end
+    if !(GPUCompiler.isghosttype(PT) || Core.Compiler.isconstType(PT))
+        return quote
+            Base.@_inline_meta
+            Base.llvmcall(
+                ($ir, $fn),
+                $combinedReturn,
+                Tuple{$PT,$(types...)},
+                fptr,
+                $(ccexprs...),
+            )
+        end
+    else
+        return quote
+            Base.@_inline_meta
+            Base.llvmcall(
+                ($ir, $fn),
+                $combinedReturn,
+                Tuple{$(types...)},
+                $(ccexprs...),
+            )
         end
     end
 end
@@ -9071,7 +9081,10 @@ include("compiler/reflection.jl")
         )
         copysetfn = meta.entry
         blk = first(blocks(copysetfn))
-        for inst in collect(instructions(blk))
+        iter = LLVM.API.LLVMGetFirstInstruction(blk)
+        while iter != C_NULL
+            inst = LLVM.Instruction(iter)
+            iter = LLVM.API.LLVMGetNextInstruction(iter)
             if isa(inst, LLVM.FenceInst)
                 eraseInst(blk, inst)
             end

src/compiler/interpreter.jl

@@ -123,7 +123,7 @@ Core.Compiler.verbose_stmt_info(@nospecialize(::EnzymeInterpreter)) = false
 Core.Compiler.method_table(@nospecialize(interp::EnzymeInterpreter), sv::InferenceState) =
     Core.Compiler.OverlayMethodTable(interp.world, interp.method_table)
 
-function is_alwaysinline_func(@nospecialize(TT))
+function is_alwaysinline_func(@nospecialize(TT))::Bool
     isa(TT, DataType) || return false
     @static if VERSION ≥ v"1.11-"
     if TT.parameters[1] == typeof(Core.memoryref)
@@ -133,7 +133,7 @@ function is_alwaysinline_func(@nospecialize(TT))
     return false
 end
 
-function is_primitive_func(@nospecialize(TT))
+function is_primitive_func(@nospecialize(TT))::Bool
     isa(TT, DataType) || return false
     ft = TT.parameters[1]
     if ft == typeof(Enzyme.pmap)
@@ -156,11 +156,11 @@ function is_primitive_func(@nospecialize(TT))
     return false
 end
 
-function isKWCallSignature(@nospecialize(TT))
+function isKWCallSignature(@nospecialize(TT))::Bool
     return TT <: Tuple{typeof(Core.kwcall),Any,Any,Vararg}
 end
 
-function simplify_kw(@nospecialize specTypes)
+function simplify_kw(@nospecialize(specTypes))
     if isKWCallSignature(specTypes)
         return Base.tuple_type_tail(Base.tuple_type_tail(specTypes))
     else
@@ -742,15 +742,15 @@ end
     end
 end
 
-@inline function array_or_number(@nospecialize(Ty))
+@inline function array_or_number(@nospecialize(Ty))::Bool
     return Ty <: AbstractArray || Ty <: Number
 end
 
-@inline function isa_array_or_number(@nospecialize(x))
+@inline function isa_array_or_number(@nospecialize(x))::Bool
     return x isa AbstractArray || x isa Number
 end
 
-@inline function num_or_eltype(@nospecialize(Ty))
+@inline function num_or_eltype(@nospecialize(Ty))::Type
     if Ty <: AbstractArray
         eltype(Ty)
     else

src/compiler/validation.jl

@@ -20,7 +20,7 @@ function get_blas_symbols()
     return symbols
 end
 
-function lookup_blas_symbol(name)
+function lookup_blas_symbol(name::String)
     Libdl.dlsym(blas_handle::Ptr{Cvoid}, name; throw_error = false)
 end
 end
@@ -127,7 +127,7 @@ function __init__()
     end
 end
 
-function memoize!(ptr, fn)
+function memoize!(ptr::Ptr{Cvoid}, fn::String)::String
     fn = get(ptr_map, ptr, fn)
     if !haskey(ptr_map, ptr)
         ptr_map[ptr] = fn
@@ -140,7 +140,7 @@ end
 
 import GPUCompiler: IRError, InvalidIRError
 
-function restore_lookups(mod::LLVM.Module)
+function restore_lookups(mod::LLVM.Module)::Nothing
     T_size_t = convert(LLVM.LLVMType, Int)
     for (v, k) in FFI.ptr_map
         if haskey(functions(mod), k)
@@ -421,7 +421,11 @@ function check_ir!(@nospecialize(job::CompilerJob), errors::Vector{IRError}, imp
     calls = LLVM.CallInst[]
     isInline = API.EnzymeGetCLBool(cglobal((:EnzymeInline, API.libEnzyme))) != 0
     mod = LLVM.parent(f)
-    for bb in blocks(f), inst in collect(instructions(bb))
+    for bb in blocks(f)
+    iter = LLVM.API.LLVMGetFirstInstruction(bb)
+    while iter != C_NULL
+        inst = LLVM.Instruction(iter)
+        iter = LLVM.API.LLVMGetNextInstruction(iter)
         if isa(inst, LLVM.CallInst)
             push!(calls, inst)
             # remove illegal invariant.load and jtbaa_const invariants
@@ -489,7 +493,11 @@ function check_ir!(@nospecialize(job::CompilerJob), errors::Vector{IRError}, imp
                     newf, _ = get_function!(mod, fname, FT)
                 else
                     found = nothing
-                    for lbb in blocks(initfn), linst in collect(instructions(lbb))
+                    for lbb in blocks(initfn)
+                    liter = LLVM.API.LLVMGetFirstInstruction(lbb)
+                    while liter != C_NULL
+                        linst = LLVM.Instruction(liter)
+                        liter = LLVM.API.LLVMGetNextInstruction(liter)
                         if !isa(linst, LLVM.CallInst)
                             continue
                         end
@@ -502,6 +510,7 @@ function check_ir!(@nospecialize(job::CompilerJob), errors::Vector{IRError}, imp
                             break
                         end
                     end
+                    end
                     if found == nothing
                         msg = sprint() do io::IO
                             println(
@@ -630,6 +639,7 @@ function check_ir!(@nospecialize(job::CompilerJob), errors::Vector{IRError}, imp
             end
         end
     end
+    end
 
     while length(calls) > 0
         inst = pop!(calls)

src/jlrt.jl

@@ -326,7 +326,7 @@ function load_if_mixed(oval::OT, val::VT) where {OT, VT}
     end
 end
 
-function val_from_byref_if_mixed(B::LLVM.IRBuilder, gutils::GradientUtils, @nospecialize(oval::LLVM.Value), @nospecialize(val::LLVM.Value))
+function val_from_byref_if_mixed(B::LLVM.IRBuilder, gutils::GradientUtils, @nospecialize(oval::LLVM.Value), @nospecialize(val::LLVM.Value))::LLVM.Value
     world = enzyme_extract_world(LLVM.parent(position(B)))
     legal, TT, _ = abs_typeof(oval)
     if !legal
@@ -374,7 +374,7 @@ function ref_if_mixed(val::VT) where {VT}
     end
 end
 
-function byref_from_val_if_mixed(B::LLVM.IRBuilder, @nospecialize(val::LLVM.Value))
+function byref_from_val_if_mixed(B::LLVM.IRBuilder, @nospecialize(val::LLVM.Value))::LLVM.Value
     world = enzyme_extract_world(LLVM.parent(position(B)))
     legal, TT, _ = abs_typeof(val)
     if !legal

src/rules/allocrules.jl

@@ -1,7 +1,4 @@
-
-function array_inner(::Type{<:Array{T}}) where {T}
-    return T
-end
+@inline LLT_ALIGN(x::Int, sz::Int) = (((x) + (sz) - 1) & ~((sz) - 1))
 function array_shadow_handler(
     B::LLVM.API.LLVMBuilderRef,
     OrigCI::LLVM.API.LLVMValueRef,
@@ -52,8 +49,6 @@ function array_shadow_handler(
 
     isunion = typ isa Union
 
-    LLT_ALIGN(x, sz) = (((x) + (sz) - 1) & ~((sz) - 1))
-
     if !isunboxed
         elsz = sizeof(Ptr{Cvoid})
         al = elsz