[codegen] custom cache lookup

codegen-lookup.jl

using InteractiveUtils, Test
using Core: MethodInstance, CodeInstance
using Base: CodegenParams
const CC = Core.Compiler
include("test/compiler/newinterp.jl")

@newinterp ConstInvokeInterp

function CC.concrete_eval_eligible(interp::ConstInvokeInterp,
    @nospecialize(f), result::CC.MethodCallResult, arginfo::CC.ArgInfo, sv::CC.AbsIntState)
    ret = @invoke CC.concrete_eval_eligible(interp::CC.AbstractInterpreter,
        f::Any, result::CC.MethodCallResult, arginfo::CC.ArgInfo, sv::CC.AbsIntState)
    if ret === :semi_concrete_eval
        return :none # disable semi-concrete interpretation
    end
    return ret
end

Base.@constprop :aggressive @noinline function target(c::Bool, x::Float64)
    if c
        y = sin(x)
        z = nothing
    else
        y = cos(x)
        z = missing
    end
    return y, z
end
function context(x::Float64)
    return target(true, x)
end

world = Base.get_world_counter()
interp = ConstInvokeInterp(; world)
code_typed(context; world, interp)
function custom_lookup(mi::MethodInstance, min_world::UInt, max_world::UInt)
    local matched_mi = nothing
    for inf_result in interp.inf_cache
        if inf_result.linfo === mi
            if CC.any(inf_result.overridden_by_const)
                return CodeInstance(interp, inf_result, inf_result.src, inf_result.valid_worlds)
            elseif matched_mi === nothing
                matched_mi = inf_result.linfo
            end
        end
    end
    matched_mi === nothing && return nothing
    return interp.code_cache.dict[matched_mi]
end

target_mi = CC.specialize_method(only(methods(target)), Tuple{typeof(target),Bool,Float64}, Core.svec())
target_ci = custom_lookup(target_mi, world, world)
@test target_ci.rettype === Tuple{Float64,Nothing} # constprop'ed source
@ccall jl_uncompress_ir(target_ci.def.def::Any, C_NULL::Ptr{Cvoid}, target_ci.inferred::Any)::Any

raw = false
lookup = @cfunction(custom_lookup, Any, (Any,Csize_t,Csize_t))
params = CodegenParams(;
    debug_info_kind=Cint(0),
    safepoint_on_entry=raw,
    gcstack_arg=raw,
    lookup)
code_llvm(context, (Float64,); params, dump_module=true)

; ModuleID = 'context'
source_filename = "context"
target datalayout = "e-m:o-i64:64-i128:128-n32:64-S128"
target triple = "aarch64-apple-darwin22.4.0"

@"+Core.Tuple#3442" = private unnamed_addr constant {}* @"+Core.Tuple#3442.jit"

@"+Core.Tuple#3442.jit" = private alias {}, inttoptr (i64 4413830608 to {}*)

; Function Signature: context(Float64)
;  @ /Users/aviatesk/julia/julia/codegen-lookup.jl:29 within `context`
define { double } @julia_context_3440(double %"x::Float64") #0 {
top:
;  @ /Users/aviatesk/julia/julia/codegen-lookup.jl:30 within `context`
  %0 = call { double } @j_target_3444(i8 zeroext 1, double %"x::Float64")
  ret { double } %0
}

; Function Attrs: noinline optnone
define nonnull {}* @jfptr_context_3441({}* %"function::Core.Function", {}** noalias nocapture noundef readonly %"args::Any[]", i32 %"nargs::UInt32") #1 {
top:
  %gcframe = alloca {}*, i32 3, align 16
  %0 = bitcast {}** %gcframe to i8*
  call void @llvm.memset.p0i8.i64(i8* align 16 %0, i8 0, i64 24, i1 true)
  %1 = alloca { double }, align 8
  %pgcstack = call {}*** inttoptr (i64 6652341776 to {}*** (i64)*)(i64 261) #12
  %frame.nroots = getelementptr inbounds {}*, {}** %gcframe, i32 0
  %frame.nroots4 = bitcast {}** %frame.nroots to i64*
  store i64 4, i64* %frame.nroots4, align 8
  %task.gcstack = load {}**, {}*** %pgcstack, align 8
  %frame.prev = getelementptr inbounds {}*, {}** %gcframe, i32 1
  %2 = bitcast {}** %frame.prev to {}***
  store {}** %task.gcstack, {}*** %2, align 8
  %3 = bitcast {}*** %pgcstack to {}***
  store {}** %gcframe, {}*** %3, align 8
  %4 = getelementptr inbounds {}*, {}** %"args::Any[]", i32 0
  %5 = load {}*, {}** %4, align 8
  %6 = bitcast {}* %5 to double*
  %7 = load double, double* %6, align 8
  %8 = call { double } @julia_context_3440(double %7)
  store { double } %8, { double }* %1, align 8
  %"+Core.Tuple#3442" = load {}*, {}** @"+Core.Tuple#3442", align 8
  %Tuple = ptrtoint {}* %"+Core.Tuple#3442" to i64
  %9 = inttoptr i64 %Tuple to {}*
  %10 = bitcast {}*** %pgcstack to {}**
  %current_task = getelementptr inbounds {}*, {}** %10, i64 -14
  %gc_slot_addr_0 = getelementptr inbounds {}*, {}** %gcframe, i32 2
  store {}* %9, {}** %gc_slot_addr_0, align 8
  call void @llvm.assume(i1 true) [ "align"({}* %9, i64 16) ]
  %ptls_field = getelementptr inbounds {}*, {}** %current_task, i64 16
  %ptls_load = load {}*, {}** %ptls_field, align 8
  %ptls = bitcast {}* %ptls_load to {}**
  %11 = bitcast {}** %ptls to i8*
  %12 = ptrtoint {}* %9 to i64
  %box = call noalias nonnull dereferenceable(16) {}* @ijl_gc_pool_alloc_instrumented(i8* %11, i32 1136, i32 16, i64 %12) #7
  %13 = bitcast {}* %box to {}**
  %box.tag_addr = getelementptr inbounds {}*, {}** %13, i64 -1
  store atomic {}* %9, {}** %box.tag_addr unordered, align 8
  %14 = bitcast {}* %box to i8*
  %15 = bitcast { double }* %1 to i8*
  call void @llvm.memcpy.p0i8.p0i8.i64(i8* align 8 %14, i8* %15, i64 8, i1 false)
  %16 = getelementptr inbounds {}*, {}** %gcframe, i32 1
  %frame.prev5 = load {}*, {}** %16, align 8
  %17 = bitcast {}*** %pgcstack to {}**
  store {}* %frame.prev5, {}** %17, align 8
  ret {}* %box
}

; Function Attrs: allockind("alloc,uninitialized,aligned") allocsize(1)
declare noalias nonnull {}* @julia.gc_alloc_obj({}**, i64, {}*) #2

; Function Signature: target(Bool, Float64)
declare { double } @j_target_3444(i8 zeroext, double) #3

; Function Attrs: nocallback nofree nosync nounwind readnone speculatable willreturn
declare void @llvm.dbg.value(metadata, metadata, metadata) #4

; Function Attrs: argmemonly nocallback nofree nosync nounwind willreturn
declare void @llvm.lifetime.start.p0i8(i64 immarg, i8* nocapture) #5

; Function Attrs: argmemonly nocallback nofree nosync nounwind willreturn
declare void @llvm.lifetime.end.p0i8(i64 immarg, i8* nocapture) #5

; Function Attrs: inaccessiblemem_or_argmemonly
declare void @ijl_gc_queue_root({}*) #6

; Function Attrs: allocsize(2)
declare noalias nonnull {}* @ijl_gc_pool_alloc_instrumented(i8*, i32, i32, i64) #7

; Function Attrs: allocsize(1)
declare noalias nonnull {}* @ijl_gc_big_alloc_instrumented(i8*, i64, i64) #8

; Function Attrs: allocsize(1)
declare noalias nonnull {}* @ijl_gc_alloc_typed(i8*, i64, i64) #8

declare noalias nonnull {}** @julia.new_gc_frame(i32)

declare void @julia.push_gc_frame({}**, i32)

declare {}** @julia.get_gc_frame_slot({}**, i32)

declare void @julia.pop_gc_frame({}**)

; Function Attrs: inaccessiblememonly nocallback nofree nosync nounwind willreturn
declare void @llvm.assume(i1 noundef) #9

; Function Attrs: allocsize(1)
declare noalias nonnull {}* @julia.gc_alloc_bytes(i8*, i64, i64) #8

; Function Attrs: argmemonly nocallback nofree nounwind willreturn writeonly
declare void @llvm.memset.p0i8.i64(i8* nocapture writeonly, i8, i64, i1 immarg) #10

; Function Attrs: argmemonly nocallback nofree nounwind willreturn
declare void @llvm.memcpy.p0i8.p0i8.i64(i8* noalias nocapture writeonly, i8* noalias nocapture readonly, i64, i1 immarg) #11

attributes #0 = { "julia.fsig"="context(Float64)" "probe-stack"="inline-asm" }
attributes #1 = { noinline optnone "probe-stack"="inline-asm" }
attributes #2 = { allockind("alloc,uninitialized,aligned") allocsize(1) }
attributes #3 = { "julia.fsig"="target(Bool, Float64)" "probe-stack"="inline-asm" }
attributes #4 = { nocallback nofree nosync nounwind readnone speculatable willreturn }
attributes #5 = { argmemonly nocallback nofree nosync nounwind willreturn }
attributes #6 = { inaccessiblemem_or_argmemonly }
attributes #7 = { allocsize(2) }
attributes #8 = { allocsize(1) }
attributes #9 = { inaccessiblememonly nocallback nofree nosync nounwind willreturn }
attributes #10 = { argmemonly nocallback nofree nounwind willreturn writeonly }
attributes #11 = { argmemonly nocallback nofree nounwind willreturn }
attributes #12 = { nounwind readnone }

!llvm.module.flags = !{!0, !1, !2}

!0 = !{i32 2, !"Dwarf Version", i32 4}
!1 = !{i32 2, !"Debug Info Version", i32 3}
!2 = !{i32 2, !"julia.debug_level", i32 2}