customize the opt pass

Author: vchuravy

src/jlgen.jl

@@ -487,6 +487,50 @@ function CC.abstract_call_known(interp::GPUInterpreter, @nospecialize(f),
         max_methods::Int)
 end
 
+# Customize the optimization pipeline
+Base.iterate(compact::CC.IncrementalCompact, state=nothing) = CC.iterate(compact, state)
+Base.getindex(compact::CC.IncrementalCompact, idx) = CC.getindex(compact, idx)
+
+function deferred_pass!(ir)
+    compact = CC.IncrementalCompact(ir)
+    for ((old_idx, idx), stmt) in compact
+        if CC.is_known_call(stmt, var"gpuc.deferred", compact)
+            @safe_show stmt
+        end
+    end
+    CC.non_dce_finish!(compact)
+    CC.simple_dce!(compact)
+    ir = CC.complete(compact)
+    return ir
+end
+
+function CC.optimize(interp::GPUInterpreter, opt::CC.OptimizationState, caller::CC.InferenceResult)
+    CC.@timeit "optimizer" ir = CC.run_passes(opt.src, opt, caller)
+    # Customizing the ipo_safe pipeline is a pain
+    ir = deferred_pass!(ir)
+    @static if VERSION >= v"1.11.0-"
+        CC.ipo_dataflow_analysis!(interp, ir, caller)
+    end
+    return CC.finish(interp, opt, ir, caller)
+end
+
+function CC.typeinf_ircode(interp::GPUInterpreter, mi::CC.MethodInstance,
+                        optimize_until::Union{Integer,AbstractString,Nothing})
+    start_time = ccall(:jl_typeinf_timing_begin, UInt64, ())
+    frame = CC.typeinf_frame(interp, mi, false)
+    if frame === nothing
+        ccall(:jl_typeinf_timing_end, Cvoid, (UInt64,), start_time)
+        return nothing, Any
+    end
+    (; result) = frame
+    opt = CC.OptimizationState(frame, interp)
+    ir = CC.run_passes(opt.src, opt, result, optimize_until)
+    ir = deferred_pass!(ir)
+    rt = CC.widenconst(CC.ignorelimited(result.result))
+    ccall(:jl_typeinf_timing_end, Cvoid, (UInt64,), start_time)
+    return ir, rt
+end
+
 ## world view of the cache
 using Core.Compiler: WorldView