From 11608c9e245c39202dc921380846a2ff62f850aa Mon Sep 17 00:00:00 2001
From: "William S. Moses" <gh@wsmoses.com>
Date: Mon, 4 Mar 2024 22:53:07 -0500
Subject: [PATCH] Remove invariant loads on primal computation

---
 src/compiler.jl          | 3 +++
 src/compiler/optimize.jl | 1 +
 2 files changed, 4 insertions(+)

diff --git a/src/compiler.jl b/src/compiler.jl
index 771dc896db..52ccb6290f 100644
--- a/src/compiler.jl
+++ b/src/compiler.jl
@@ -4366,6 +4366,9 @@ function GPUCompiler.codegen(output::Symbol, job::CompilerJob{<:EnzymeTarget};
 
     mod, meta = GPUCompiler.codegen(:llvm, primal_job; optimize=false, toplevel=toplevel, cleanup=false, validate=false, parent_job=parent_job)
     prepare_llvm(mod, primal_job, meta)
+    for f in functions(mod)
+        permit_inlining!(f)
+    end
 
     LLVM.ModulePassManager() do pm
         API.AddPreserveNVVMPass!(pm, #=Begin=#true)
diff --git a/src/compiler/optimize.jl b/src/compiler/optimize.jl
index f430830dd4..70ee9b6912 100644
--- a/src/compiler/optimize.jl
+++ b/src/compiler/optimize.jl
@@ -1690,6 +1690,7 @@ function post_optimze!(mod, tm, machine=true)
         addOptimizationPasses!(pm)
         run!(pm, mod)
     end
+    println(string(mod))
     if machine
         # TODO enable validate_return_roots
         # validate_return_roots!(mod)