@@ -168,32 +168,33 @@ function generate_obj_for_compile(f, tt, external = true, path::String = tempnam
168168 config = GPUCompiler. CompilerConfig (NativeCompilerTarget (target... ), params, name = name, kernel = false )
169169 job = GPUCompiler. CompilerJob (GPUCompiler. methodinstance (typeof (f), tt), config)
170170
171- mod, meta = GPUCompiler. JuliaContext () do context
172- GPUCompiler. codegen (:llvm , job; strip= strip_llvm, only_entry= false , validate= false , optimize= false , ctx= context)
171+ table = GPUCompiler. JuliaContext () do context
172+ mod, meta = GPUCompiler. codegen (:llvm , job; strip= strip_llvm, only_entry= false , validate= false , optimize= false )
173+ # Use Enzyme's annotation and optimization pipeline
174+ annotate! (mod)
175+ tm = GPUCompiler. llvm_machine (external ? ExternalNativeCompilerTarget (target... ) : NativeCompilerTarget (target... ))
176+ optimize! (mod, tm)
177+
178+ # Scoop up all the pointers in the optimized module, and replace them with unitialized global variables.
179+ # `table` is a dictionary where the keys are julia objects that are needed by the function, and the values
180+ # of the dictionary are the names of their associated LLVM GlobalVariable names.
181+ table = relocation_table! (mod)
182+
183+ # Now that we've removed all the pointers from the code, we can (hopefully) safely lower all the instrinsics
184+ # (again, using Enzyme's pipeline)
185+ post_optimize! (mod, tm; remove_julia_addrspaces)
186+
187+ # Make sure we didn't make any glaring errors
188+ LLVM. verify (mod)
189+ obj, _ = GPUCompiler. emit_asm (job, mod; strip= strip_asm, validate= false , format= LLVM. API. LLVMObjectFile)
190+ # Compile the LLVM module to native code and save it to disk
191+ open (obj_path, " w" ) do io
192+ write (io, obj)
193+ end
194+ table
173195 end
174196
175- # Use Enzyme's annotation and optimization pipeline
176- annotate! (mod)
177- tm = GPUCompiler. llvm_machine (external ? ExternalNativeCompilerTarget (target... ) : NativeCompilerTarget (target... ))
178- optimize! (mod, tm)
179-
180- # Scoop up all the pointers in the optimized module, and replace them with unitialized global variables.
181- # `table` is a dictionary where the keys are julia objects that are needed by the function, and the values
182- # of the dictionary are the names of their associated LLVM GlobalVariable names.
183- table = relocation_table! (mod)
184-
185- # Now that we've removed all the pointers from the code, we can (hopefully) safely lower all the instrinsics
186- # (again, using Enzyme's pipeline)
187- post_optimize! (mod, tm; remove_julia_addrspaces)
188-
189- # Make sure we didn't make any glaring errors
190- LLVM. verify (mod)
191197
192- # Compile the LLVM module to native code and save it to disk
193- obj, _ = GPUCompiler. emit_asm (job, mod; strip= strip_asm, validate= false , format= LLVM. API. LLVMObjectFile)
194- open (obj_path, " w" ) do io
195- write (io, obj)
196- end
197198 path, name, table
198199end
199200
@@ -586,7 +587,7 @@ function native_llvm_module(f, tt, name=fix_name(f); demangle, kwargs...)
586587 end
587588 job, kwargs = native_job (f, tt, true ; name, kwargs... )
588589 m, _ = GPUCompiler. JuliaContext () do context
589- GPUCompiler. codegen (:llvm , job; strip= true , only_entry= false , validate= false , ctx = context )
590+ GPUCompiler. codegen (:llvm , job; strip= true , only_entry= false , validate= false )
590591 end
591592 return m
592593end
0 commit comments