[Inference] limit inference timing recording to NativeInterpreter only

[vchuravy]

Allows consumers to separate native inference from inference coming
from a custom AbstractInterpreter.

x-ref: timholy/SnoopCompile.jl#338

I choose typeof(interp) since the interp may be arbitrarily heavy and we might not want to hold on to them.

[maleadt]

SGTM

[aviatesk]

Alternatively we can do something like:

diff --git a/base/compiler/typeinfer.jl b/base/compiler/typeinfer.jl
index 1eec73d043..eb88a2081e 100644
--- a/base/compiler/typeinfer.jl
+++ b/base/compiler/typeinfer.jl
@@ -205,8 +205,7 @@ __set_measure_typeinf(onoff::Bool) = __measure_typeinf__[] = onoff
 const __measure_typeinf__ = fill(false)
 
 # Wrapper around _typeinf that optionally records the exclusive time for each invocation.
-function typeinf(interp::AbstractInterpreter, frame::InferenceState)
-    interp = switch_from_irinterp(interp)
+function typeinf(interp::NativeInterpreter, frame::InferenceState)
     if __measure_typeinf__[]
         Timings.enter_new_timer(frame)
         v = _typeinf(interp, frame)
@@ -216,6 +215,8 @@ function typeinf(interp::AbstractInterpreter, frame::InferenceState)
         return _typeinf(interp, frame)
     end
 end
+# disable recording timings for external `AbstractInterpreter`s
+typeinf(interp::AbstractInterpreter, frame::InferenceState) = _typeinf(interp, frame)
 
 function finish!(interp::AbstractInterpreter, caller::InferenceResult)
     # If we didn't transform the src for caching, we may have to transform
@@ -242,6 +243,7 @@ function finish!(interp::AbstractInterpreter, caller::InferenceResult)
 end
 
 function _typeinf(interp::AbstractInterpreter, frame::InferenceState)
+    interp = switch_from_irinterp(interp)
     typeinf_nocycle(interp, frame) || return false # frame is now part of a higher cycle
     # with no active ip's, frame is done
     frames = frame.callers_in_cycle

?

[vchuravy]

Alternatively we can do something like:

I thought about that, but I actually want to be able to see how much time it took to infer CUDA code.

[aviatesk]

I thought about that, but I actually want to be able to see how much time it took to infer CUDA code.

GPUCompiler can implement its own measurement mechanism for that?

I'm not opposed to the current approach, but it seems somewhat risky to leave Core.Compiler.Timings to support an external AbstractInterpreters. The reason being that the external AbstractInterpreter's compilation pipeline may trigger native compilation, which could invalidate the current assumption of being able to find children/parent frames. So I believe it would be cleaner and more correct to have Core.Compiler.Timings only support NativeInterpreter by default, and have any external AbstractInterpreter optionally use the data structure of Core.Compiler.Timings to inspect its own compilation performance.

[vchuravy]

The major reason would be that it would be great if SnoopCompile could be extended to custom abstract interpreters.
For me that's the next frontier in TTFX. Oceananigans has 10+ minutes of compilation time.

[vchuravy]

The reason being that the external AbstractInterpreter's compilation pipeline may trigger native compilation, which could invalidate the current assumption of being able to find children/parent frames.

I don't know why that would be the case? I would say it is correct? You might have arbitrary interleavings?

[aviatesk]

I don't know why that would be the case? I would say it is correct? You might have arbitrary interleavings?

It's possible that there are type unstable parts within some external AbstractInterpreter implementation, and then invoking the external compilation pipeline may trigger dynamic dispatch to native-compile the compilation pipeline itself.

[aviatesk]

The major reason would be that it would be great if SnoopCompile could be extended to custom abstract interpreters.
For me that's the next frontier in TTFX. Oceananigans has 10+ minutes of compilation time.

Yeah, I agree. But would it be cleaner if external AbstractInterpreter uses SnoopCompile utilities like this?

const gpu_timings = CC.Timings.Timing[]

function CC.typeinf(interp::GPUInterpreter, frame::InferenceState)
    [ push timing into gpu_timings ]
    @invoke CC.typeinf(interp::AbstractInterpreter, frame::InferenceState)
end

macro snoopi_deep_gpu(ex) _snoopi_deep_gpu(ex) end

function _snoopi_deep_gpu(cmd::Expr)
    return quote
        start_gpu_deep_timing()
        try
            $(esc(cmd))
        finally
            stop_gpu_deep_timing()
        end
        CC.Timings.InferenceTimingNode(gpu_timings[1])
    end 
end

[vchuravy]

You would want to do both at the same time though...

…

On Tue, Apr 18, 2023, 11:56 Shuhei Kadowaki ***@***.***> wrote: The major reason would be that it would be great if SnoopCompile could be extended to custom abstract interpreters. For me that's the next frontier in TTFX. Oceananigans has 10+ minutes of compilation time. Yeah, I agree. But would it be cleaner if external AbstractInterpreter uses SnoopCompile utilities like this? const gpu_timings = CC.Timings.Timing[] function CC.typeinf(interp::GPUInterpreter, frame::InferenceState) [ push timing into gpu_timings ] @invoke CC.typeinf(interp::AbstractInterpreter, frame::InferenceState)end macro snoopi_deep_gpu(ex) _snoopi_deep_gpu(ex) end function _snoopi_deep_gpu(cmd::Expr) return quote start_gpu_deep_timing() try $(esc(cmd)) finally stop_gpu_deep_timing() end CC.Timings.InferenceTimingNode(gpu_timings[1]) end end — Reply to this email directly, view it on GitHub <#49391 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AABDO2WTCOU2UCQKTAQCSTTXB22R3ANCNFSM6AAAAAAXBT4PGE> . You are receiving this because you authored the thread.Message ID: ***@***.***>

[aviatesk]

Leaving a comment from slack for future reference:

Shuhei Kadowaki
I think it is problematic that NativeInterpreter and external AbstractInterpreters share the same data structure (Core.Compiler.Timings._timings), since Core.Compiler.Timings seems to assume that the inference graph is constructed by the same interpreter.

E.g.

julia/base/compiler/typeinfer.jl

Lines 165 to 166 in a34261f

	# Prepare to unwind one level of the stack and record in the parent
	parent_timer = _timings[end]

finds the parent frame just by looking at _timings[end], and currently we may get an inference graph like:

Timing(::GPUInterpreter, MethodInstance for xxx, ...)
 Timing(::GPUInterpreter, MethodInstance for yyy, ...)
  Timing(::NativeInterpreter, MethodInstance for GPUCompiler.abstract_eval_zzz, ...)
   Timing(::NativeInterpreter, MethodInstance for GPUCompiler.some_utility_for_abstract_eval_zzz, ...)
  Timing(::GPUInterpreter, MethodInstance for zzz, ...)

, which seems to be a bit tricky to handle?

If we isolate Core.Compiler.Timings._timings to NativeInterpreter and make external AbstractInterpreter maintain its own timings, we will get cleaner graphs like

Timing(::GPUInterpreter, MethodInstance for xxx, ...)
 Timing(::GPUInterpreter, MethodInstance for yyy, ...)
  Timing(::GPUInterpreter, MethodInstance for zzz, ...)

and

Timing(::NativeInterpreter, MethodInstance for GPUCompiler.abstract_eval_zzz, ...)
 Timing(::NativeInterpreter, MethodInstance for GPUCompiler.some_utility_for_abstract_eval_zzz, ...)