Rewrite optimize pipeline for NewPM

[vchuravy]

No description provided.

[github-actions]

Your PR requires formatting changes to meet the project's style guidelines.
Please consider running Runic (git runic vc/split_newpm2) to apply these changes.

Click here to view the suggested changes.

diff --git a/src/compiler/optimize.jl b/src/compiler/optimize.jl
index 027f1b22..11846c06 100644
--- a/src/compiler/optimize.jl
+++ b/src/compiler/optimize.jl
@@ -73,89 +73,89 @@ function optimize!(mod::LLVM.Module, tm::LLVM.TargetMachine)
         run!(pb, mod, tm)
     end
 
-    function middle_optimize!(second_stage=false)
-    @dispose pb = NewPMPassBuilder() begin
-        registerEnzymeAndPassPipeline!(pb)
-        register!(pb, RewriteGenericMemoryPass())
-        add!(pb, NewPMAAManager()) do aam
-            add!(aam, ScopedNoAliasAA())
-            add!(aam, TypeBasedAA())
-            add!(aam, BasicAA())
+    function middle_optimize!(second_stage = false)
+        return @dispose pb = NewPMPassBuilder() begin
+            registerEnzymeAndPassPipeline!(pb)
+            register!(pb, RewriteGenericMemoryPass())
+            add!(pb, NewPMAAManager()) do aam
+                add!(aam, ScopedNoAliasAA())
+                add!(aam, TypeBasedAA())
+                add!(aam, BasicAA())
         end
-        add!(pb, NewPMModulePassManager()) do mpm
-            add!(mpm, RewriteGenericMemoryPass())
-            add!(mpm, CPUFeaturesPass()) # why is this duplicated?
-
-            add!(mpm, NewPMFunctionPassManager()) do fpm
-                add!(fpm, InstCombinePass())
-                add!(fpm, JLInstSimplifyPass())
-                add!(fpm, SimplifyCFGPass())
-                add!(fpm, SROAPass())
-                add!(fpm, InstCombinePass())
-                add!(fpm, JLInstSimplifyPass())
-                add!(fpm, JumpThreadingPass())
-                add!(fpm, CorrelatedValuePropagationPass())
-                add!(fpm, InstCombinePass())
-                add!(fpm, JLInstSimplifyPass())
-                add!(fpm, ReassociatePass())
-                add!(fpm, EarlyCSEPass())
-                add!(fpm, AllocOptPass())
-
-                add!(fpm, NewPMLoopPassManager(use_memory_ssa=true)) do lpm
-                    add!(lpm, LoopIdiomRecognizePass())
-                    add!(lpm, LoopRotatePass())
-                    add!(lpm, LowerSIMDLoopPass())
-                    add!(lpm, LICMPass())
-                    add!(lpm, JuliaLICMPass())
-                    add!(lpm, SimpleLoopUnswitchPass())
-                end
-
-                add!(fpm, InstCombinePass())
-                add!(fpm, JLInstSimplifyPass())
-                add!(fpm, NewPMLoopPassManager()) do lpm
-                    add!(lpm, IndVarSimplifyPass())
-                    add!(lpm, LoopDeletionPass())
+            add!(pb, NewPMModulePassManager()) do mpm
+                add!(mpm, RewriteGenericMemoryPass())
+                add!(mpm, CPUFeaturesPass()) # why is this duplicated?
+
+                add!(mpm, NewPMFunctionPassManager()) do fpm
+                    add!(fpm, InstCombinePass())
+                    add!(fpm, JLInstSimplifyPass())
+                    add!(fpm, SimplifyCFGPass())
+                    add!(fpm, SROAPass())
+                    add!(fpm, InstCombinePass())
+                    add!(fpm, JLInstSimplifyPass())
+                    add!(fpm, JumpThreadingPass())
+                    add!(fpm, CorrelatedValuePropagationPass())
+                    add!(fpm, InstCombinePass())
+                    add!(fpm, JLInstSimplifyPass())
+                    add!(fpm, ReassociatePass())
+                    add!(fpm, EarlyCSEPass())
+                    add!(fpm, AllocOptPass())
+
+                    add!(fpm, NewPMLoopPassManager(use_memory_ssa = true)) do lpm
+                        add!(lpm, LoopIdiomRecognizePass())
+                        add!(lpm, LoopRotatePass())
+                        add!(lpm, LowerSIMDLoopPass())
+                        add!(lpm, LICMPass())
+                        add!(lpm, JuliaLICMPass())
+                        add!(lpm, SimpleLoopUnswitchPass())
+                    end
+
+                    add!(fpm, InstCombinePass())
+                    add!(fpm, JLInstSimplifyPass())
+                    add!(fpm, NewPMLoopPassManager()) do lpm
+                        add!(lpm, IndVarSimplifyPass())
+                        add!(lpm, LoopDeletionPass())
+                    end
+                    add!(fpm, LoopUnrollPass(opt_level = 2)) # what opt level?
+                    add!(fpm, AllocOptPass())
+                    add!(fpm, SROAPass())
+                    add!(fpm, GVNPass())
+
+                    # This InstCombine needs to be after GVN
+                    # Otherwise it will generate load chains in GPU code...
+                    add!(fpm, InstCombinePass())
+                    add!(fpm, JLInstSimplifyPass())
+                    add!(fpm, MemCpyOptPass())
+                    add!(fpm, SCCPPass())
+                    add!(fpm, InstCombinePass())
+                    add!(fpm, JLInstSimplifyPass())
+                    add!(fpm, JumpThreadingPass())
+                    add!(fpm, DSEPass())
+                    add!(fpm, AllocOptPass())
+                    add!(fpm, SimplifyCFGPass())
+
+
+                    add!(fpm, NewPMLoopPassManager()) do lpm
+                        add!(lpm, LoopIdiomRecognizePass())
+                        add!(lpm, LoopDeletionPass())
                 end
-                add!(fpm, LoopUnrollPass(opt_level=2)) # what opt level?
-                add!(fpm, AllocOptPass())
-                add!(fpm, SROAPass())
-                add!(fpm, GVNPass())
-
-                # This InstCombine needs to be after GVN
-                # Otherwise it will generate load chains in GPU code...
-                add!(fpm, InstCombinePass())
-                add!(fpm, JLInstSimplifyPass())
-                add!(fpm, MemCpyOptPass())
-                add!(fpm, SCCPPass())
-                add!(fpm, InstCombinePass())
-                add!(fpm, JLInstSimplifyPass())
-                add!(fpm, JumpThreadingPass())
-                add!(fpm, DSEPass())
-                add!(fpm, AllocOptPass())
-                add!(fpm, SimplifyCFGPass())
-
-
-                add!(fpm, NewPMLoopPassManager()) do lpm
-                    add!(lpm, LoopIdiomRecognizePass())
-                    add!(lpm, LoopDeletionPass())
+                    add!(fpm, JumpThreadingPass())
+                    add!(fpm, CorrelatedValuePropagationPass())
+                    if second_stage
+
+                        add!(fpm, ADCEPass())
+                        add!(fpm, InstCombinePass())
+                        add!(fpm, JLInstSimplifyPass())
+
+                        # GC passes
+                        add!(fpm, GCInvariantVerifierPass(strong = false))
+                        add!(fpm, SimplifyCFGPass())
+                        add!(fpm, InstCombinePass())
+                        add!(fpm, JLInstSimplifyPass())
+                    end # second_stage
                 end
-                add!(fpm, JumpThreadingPass())
-                add!(fpm, CorrelatedValuePropagationPass())
-                if second_stage
-
-                add!(fpm, ADCEPass())
-                add!(fpm, InstCombinePass())
-                add!(fpm, JLInstSimplifyPass())
-
-                # GC passes
-                add!(fpm, GCInvariantVerifierPass(strong=false))
-                add!(fpm, SimplifyCFGPass())
-                add!(fpm, InstCombinePass())
-                add!(fpm, JLInstSimplifyPass())
-                end # second_stage
-            end
         end
-        run!(pb, mod, tm)
+            run!(pb, mod, tm)
     end
     end # middle_optimize!
 
diff --git a/src/llvm/transforms.jl b/src/llvm/transforms.jl
index 5462fd43..67255c3a 100644
--- a/src/llvm/transforms.jl
+++ b/src/llvm/transforms.jl
@@ -2372,7 +2372,7 @@ end
 function rewrite_generic_memory!(mod::LLVM.Module)
     @static if VERSION < v"1.11-"
         return false
-    else    
+    else
         for f in functions(mod), bb in blocks(f)
             iter = LLVM.API.LLVMGetFirstInstruction(bb)
             while iter != C_NULL
@@ -2381,7 +2381,7 @@ function rewrite_generic_memory!(mod::LLVM.Module)
                 if !isa(inst, LLVM.LoadInst)
                     continue
                 end
-        
+
                 if isa(operands(inst)[1], LLVM.ConstantExpr)
                     legal2, obj = absint(inst)
                     if legal2 && obj isa Memory && obj == typeof(obj).instance

[codecov]

Codecov Report

❌ Patch coverage is 99.25094% with 2 lines in your changes missing coverage. Please review.
⚠️ Please upload report for BASE (vc/split_newpm2@623c4c6). Learn more about missing BASE report.

Files with missing lines	Patch %	Lines
src/compiler/optimize.jl	99.54%	1 Missing ⚠️
src/llvm/transforms.jl	97.87%	1 Missing ⚠️

Additional details and impacted files

@@                Coverage Diff                 @@
##             vc/split_newpm2    #2713   +/-   ##
==================================================
  Coverage                   ?   72.38%           
==================================================
  Files                      ?       58           
  Lines                      ?    18628           
  Branches                   ?        0           
==================================================
  Hits                       ?    13483           
  Misses                     ?     5145           
  Partials                   ?        0           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[github-actions]

Benchmark Results

	main	90fe74a...	main / 90fe74a...
basics/make_zero/namedtuple	0.0531 ± 0.0032 μs	0.0575 ± 0.0052 μs	0.922 ± 0.1
basics/make_zero/struct	0.257 ± 0.0045 μs	0.266 ± 0.0069 μs	0.97 ± 0.03
basics/overhead	5.56 ± 0.01 ns	5.26 ± 0.01 ns	1.06 ± 0.0028
basics/remake_zero!/namedtuple	0.237 ± 0.012 μs	0.239 ± 0.013 μs	0.991 ± 0.075
basics/remake_zero!/struct	0.235 ± 0.0095 μs	0.237 ± 0.018 μs	0.992 ± 0.086
fold_broadcast/multidim_sum_bcast/1D	10.3 ± 0.28 μs	10.6 ± 1.9 μs	0.972 ± 0.17
fold_broadcast/multidim_sum_bcast/2D	12.2 ± 0.26 μs	12.2 ± 0.3 μs	1 ± 0.033
time_to_load	1.25 ± 0.009 s	1.37 ± 0.011 s	0.917 ± 0.0099

Benchmark Plots

A plot of the benchmark results has been uploaded as an artifact at https://github.com/EnzymeAD/Enzyme.jl/actions/runs/18921523721/artifacts/4410793269.

[vchuravy]

@wsmoses this now shows the same failure as before, but is just the rewrite of optimize!

[vchuravy]

It turns out that

Enzyme.jl/src/compiler/optimize.jl

Lines 441 to 458 in 107b327

	LLVM.run!(pm, mod)
	aggressive_dce!(pm)
	instruction_combining!(pm)
	jl_inst_simplify!(pm)
	# Loop Vectorize -- not for Enzyme
	# InstCombine
	# GC passes
	barrier_noop!(pm)
	gc_invariant_verifier_tm!(pm, tm, false)
	# FIXME: Currently crashes printing
	cfgsimplification!(pm)
	instruction_combining!(pm) # Extra for Enzyme
	jl_inst_simplify!(pm)
	LLVM.run!(pm, mod)

was running the middle pipeline twice. I figured out that you only run a few extra passes to get the test over the hill

[wsmoses]

oh interesting.

that said, can we keep exact equivalent for the current new pm change, and then we can do pass pipeline tuning as a separate follow up?

[wsmoses]

that way itll be easier to debug (e.g. git blame/figure out) if anything else goes awry/differently that wasnt caught by the CI tests