OptAnalyzer: make function_filter work on function object directly

aviatesk · May 12, 2023 · 30078f0 · 30078f0
1 parent 2636533
commit 30078f0
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 25 deletions.
diff --git a/src/analyzers/optanalyzer.jl b/src/analyzers/optanalyzer.jl
@@ -110,14 +110,14 @@ that are specific to the optimization analysis.
            └──────────
 
 ---
-- `function_filter = @nospecialize(ft)->true`:\\
-  A predicate which takes a function type and returns `false` to skip runtime dispatch
+- `function_filter = @nospecialize(f)->true`:\\
+  A predicate which takes a function and returns `false` to skip runtime dispatch
   analysis on the function call. This configuration is particularly useful when your program
   uses a function that is intentionally written to use runtime dispatch.
 
   ```julia-repl
   # ignores `Core.Compiler.widenconst` calls (since it's designed to be runtime-dispatched):
-  julia> function_filter(@nospecialize(ft)) = ft !== typeof(Core.Compiler.widenconst)
+  julia> function_filter(@nospecialize(f)) = f !== Core.Compiler.widenconst
 
   julia> @test_opt function_filter=function_filter f(args...)
   ...
@@ -219,7 +219,7 @@ const OPT_ANALYZER_CACHE = IdDict{UInt, AnalysisCache}()
 
 struct OptAnalysisPass <: ReportPass end
 
-optanalyzer_function_filter(@nospecialize ft) = true
+optanalyzer_function_filter(@nospecialize f) = true
 
 # TODO better to work only `finish!`
 function CC.finish(frame::InferenceState, analyzer::OptAnalyzer)
@@ -337,12 +337,14 @@ function (::OptAnalysisPass)(::Type{RuntimeDispatchReport}, analyzer::OptAnalyze
             CC.is_stmt_throw_block(src.ssaflags[pc]) && continue
         end
         if isexpr(x, :call)
-            ft = widenconst(argextype(first(x.args), src, sptypes, slottypes))
-            ft <: Builtin && continue # ignore `:call`s of language intrinsics
-            if analyzer.function_filter(ft)
-                add_new_report!(analyzer, caller, RuntimeDispatchReport((opt, pc)))
-                reported |= true
+            ft = argextype(first(x.args), src, sptypes, slottypes)
+            f = singleton_type(ft)
+            if f !== nothing
+                f isa Builtin && continue # ignore `:call`s of language intrinsics
+                analyzer.function_filter(f) || continue # ignore user-specified functions
             end
+            add_new_report!(analyzer, caller, RuntimeDispatchReport((opt, pc)))
+            reported |= true
         end
     end
     return reported

diff --git a/test/analyzers/test_optanalyzer.jl b/test/analyzers/test_optanalyzer.jl
@@ -200,7 +200,7 @@ end
             @test !isempty(get_reports_with_test(result))
             @test any(r->isa(r,RuntimeDispatchReport), get_reports_with_test(result))
         end
-        let function_filter(@nospecialize ft) = ft !== typeof(with_runtime_dispatch)
+        let function_filter(@nospecialize f) = f !== with_runtime_dispatch
             test_opt((Vector{Any},); function_filter) do xs
                 with_runtime_dispatch(xs[1])
             end
@@ -213,7 +213,7 @@ end
             @test !isempty(get_reports_with_test(result))
             @test any(r->isa(r,RuntimeDispatchReport), get_reports_with_test(result))
         end
-        let function_filter(@nospecialize ft) = ft !== Type{WithRuntimeDispatch}
+        let function_filter(@nospecialize f) = f !== WithRuntimeDispatch
             test_opt((Vector{Any},); function_filter) do xs
                 WithRuntimeDispatch(xs[1])
             end

diff --git a/test/self_check.jl b/test/self_check.jl
@@ -33,20 +33,20 @@ let target_modules = (JET,)
 
     # ignore some dynamically-designed functions
     # TODO implement `signature_filter` and limit the ignorance scope
-    function function_filter(@nospecialize ft)
-        if (ft === typeof(JET.widenconst) ||
-            ft === typeof(JET.ignorelimited) ||
-            ft === typeof(JET.print) ||
-            ft === typeof(Base.CoreLogging.handle_message) ||
-            ft === typeof(get) ||
-            ft === typeof(sprint) ||
-            ft === typeof(string) ||
-            ft === typeof(zero) ||
-            ft === typeof(JET.copy_report) ||
-            ft === typeof(JET.handle_sig!) ||
-            ft === typeof(JET._which) ||
-            ft === typeof(JET.rewrap_unionall) ||
-            ft === typeof(JET.nameof) ||
+    function function_filter(@nospecialize f)
+        if (f === JET.widenconst ||
+            f === JET.ignorelimited ||
+            f === JET.print ||
+            f === Base.CoreLogging.handle_message ||
+            f === get ||
+            f === sprint ||
+            f === string ||
+            f === zero ||
+            f === JET.copy_report ||
+            f === JET.handle_sig! ||
+            f === JET._which ||
+            f === JET.rewrap_unionall ||
+            f === JET.nameof ||
             false)
             return false
         end