From fa4635fbc59697056476e05a87486c82f2ea0c13 Mon Sep 17 00:00:00 2001 From: pw0908 Date: Sun, 21 Jun 2026 15:30:25 -0700 Subject: [PATCH 1/5] Use more flexible typing to allow for compatibility with Metal architecture. --- src/compiler/optimize.jl | 6 +++--- src/llvm/transforms.jl | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/compiler/optimize.jl b/src/compiler/optimize.jl index f119f3ef32..ee145ed35b 100644 --- a/src/compiler/optimize.jl +++ b/src/compiler/optimize.jl @@ -30,7 +30,7 @@ SafeAtomicToRegularStorePass() = NewPMFunctionPass("safe_atomic_to_regular_store Addr13NoAliasPass() = NewPMModulePass("addr13_noalias", addr13NoAlias) RemoveAlwaysInlineRootsPass() = NewPMModulePass("remove_alwaysinline_roots", remove_alwaysinline_roots!) -function optimize!(mod::LLVM.Module, tm::LLVM.TargetMachine) +function optimize!(mod::LLVM.Module, tm::Union{LLVM.TargetMachine, Nothing}) @dispose pb = NewPMPassBuilder() begin registerEnzymeAndPassPipeline!(pb) register!(pb, Addr13NoAliasPass()) @@ -395,7 +395,7 @@ end const DumpPreCallConv = Ref(false) const DumpPostCallConv = Ref(false) -function fixup_callconv!(mod::LLVM.Module, tm::LLVM.TargetMachine) +function fixup_callconv!(mod::LLVM.Module, tm::Union{LLVM.TargetMachine, Nothing}) addr13NoAlias(mod) removeDeadArgs!(mod, tm, #=post_gc_fixup=#false) @@ -456,7 +456,7 @@ function fixup_callconv!(mod::LLVM.Module, tm::LLVM.TargetMachine) return end -function post_optimize!(mod::LLVM.Module, tm::LLVM.TargetMachine, machine::Bool = true; callconv::Bool = true) +function post_optimize!(mod::LLVM.Module, tm::Union{LLVM.TargetMachine, Nothing}, machine::Bool = true; callconv::Bool = true) if callconv fixup_callconv!(mod, tm) end diff --git a/src/llvm/transforms.jl b/src/llvm/transforms.jl index 89ce6a55b6..5b945a09e3 100644 --- a/src/llvm/transforms.jl +++ b/src/llvm/transforms.jl @@ -2566,7 +2566,7 @@ function checkNoAssumeFalse(mod::LLVM.Module, shouldshow::Bool = false) end end -function removeDeadArgs!(mod::LLVM.Module, tm::LLVM.TargetMachine, post_gc_fixup::Bool) +function removeDeadArgs!(mod::LLVM.Module, tm::Union{LLVM.TargetMachine, Nothing}, post_gc_fixup::Bool) # We need to run globalopt first. This is because remove dead args will otherwise # take internal functions and replace their args with undef. Then on LLVM up to # and including 12 (but fixed 13+), Attributor will incorrectly change functions that From c7858f4b52240f95b161b99efdc0899b1d79bc70 Mon Sep 17 00:00:00 2001 From: pw0908 Date: Mon, 22 Jun 2026 16:48:52 -0700 Subject: [PATCH 2/5] Add compatibility with transcendental functions in Metal --- src/compiler.jl | 186 ++++++++++++++++++++++++++++++++++++- src/compiler/optimize.jl | 21 +++-- src/internal_rules/math.jl | 1 - 3 files changed, 198 insertions(+), 10 deletions(-) diff --git a/src/compiler.jl b/src/compiler.jl index 286f8d0a61..b9b6e543f5 100644 --- a/src/compiler.jl +++ b/src/compiler.jl @@ -91,6 +91,7 @@ end GPUCompiler.llvm_triple(target::EnzymeTarget) = GPUCompiler.llvm_triple(target.target) GPUCompiler.llvm_datalayout(target::EnzymeTarget) = GPUCompiler.llvm_datalayout(target.target) GPUCompiler.llvm_machine(target::EnzymeTarget) = GPUCompiler.llvm_machine(target.target) +GPUCompiler.llvm_targetinfo(target::EnzymeTarget) = GPUCompiler.llvm_targetinfo(target.target) GPUCompiler.nest_target(::EnzymeTarget, other::AbstractCompilerTarget) = EnzymeTarget(other) GPUCompiler.have_fma(target::EnzymeTarget, T::Type) = GPUCompiler.have_fma(target.target, T) GPUCompiler.dwarf_version(target::EnzymeTarget) = GPUCompiler.dwarf_version(target.target) @@ -294,6 +295,173 @@ const known_ops = Dict{DataType,Tuple{Symbol,Int,Union{Nothing,Tuple{Symbol,Data typeof(Base.FastMath.tanh_fast) => (:tanh, 1, nothing), typeof(Base.fma_emulated) => (:fma, 3, nothing), ) + +# Maps GPU math intrinsic names (e.g. Metal AIR) to their libm equivalents. +# Used by annotate_gpu_math_intrinsics! so Enzyme's C++ core can apply its +# built-in derivative rules to opaque GPU intrinsics without Julia dispatch. +const GPU_MATH_INTRINSIC_MAP = Dict{String,String}( + "air.log.f32" => "logf", "air.log.f64" => "log", + "air.log2.f32" => "log2f", "air.log2.f64" => "log2", + "air.log10.f32" => "log10f", "air.log10.f64" => "log10", + "air.log1p.f32" => "log1pf", "air.log1p.f64" => "log1p", + "air.exp.f32" => "expf", "air.exp.f64" => "exp", + "air.exp2.f32" => "exp2f", "air.exp2.f64" => "exp2", + "air.expm1.f32" => "expm1f", "air.expm1.f64" => "expm1", + "air.sqrt.f32" => "sqrtf", "air.sqrt.f64" => "sqrt", + "air.cbrt.f32" => "cbrtf", "air.cbrt.f64" => "cbrt", + "air.sin.f32" => "sinf", "air.sin.f64" => "sin", + "air.cos.f32" => "cosf", "air.cos.f64" => "cos", + "air.tan.f32" => "tanf", "air.tan.f64" => "tan", + "air.asin.f32" => "asinf", "air.asin.f64" => "asin", + "air.acos.f32" => "acosf", "air.acos.f64" => "acos", + "air.atan.f32" => "atanf", "air.atan.f64" => "atan", + "air.atan2.f32" => "atan2f", "air.atan2.f64" => "atan2", + "air.sinh.f32" => "sinhf", "air.sinh.f64" => "sinh", + "air.cosh.f32" => "coshf", "air.cosh.f64" => "cosh", + "air.tanh.f32" => "tanhf", "air.tanh.f64" => "tanh", + "air.fma.f32" => "fmaf", "air.fma.f64" => "fma", + "air.pow.f32" => "powf", "air.pow.f64" => "pow", + "air.fast_log.f32" => "logf", + "air.fast_exp.f32" => "expf", + "air.fast_sin.f32" => "sinf", + "air.fast_cos.f32" => "cosf", + "air.fast_tan.f32" => "tanf", + "air.fast_tanh.f32" => "tanhf", + "air.fast_sqrt.f32" => "sqrtf", + "air.fast_asin.f32" => "asinf", + "air.fast_acos.f32" => "acosf", + "air.fast_atan.f32" => "atanf", + "air.fast_atan2.f32" => "atan2f", + "air.fast_sinh.f32" => "sinhf", + "air.fast_cosh.f32" => "coshf", + "air.fast_acosh.f32" => "acoshf", + "air.fast_asinh.f32" => "asinhf", +) + +# Annotate external GPU math intrinsic declarations (e.g. air.log.f32) with +# enzyme_math so the Enzyme C++ core uses its built-in derivative formulas. +# This is needed because the Metal log/exp/etc. implementations are compiled +# via @device_override and may be inlined away before set_module_types! sees +# them — leaving only the bare AIR intrinsic declaration without annotation. +function annotate_gpu_math_intrinsics!(mod::LLVM.Module) + for fn in functions(mod) + fname = LLVM.name(fn) + isempty(blocks(fn)) || continue # external declarations only + libm_name = get(GPU_MATH_INTRINSIC_MAP, fname, nothing) + libm_name === nothing && continue + attrs = function_attributes(fn) + any(isa(a, LLVM.StringAttribute) && kind(a) == "enzyme_math" + for a in collect(attrs)) && continue # already annotated + push!(attrs, StringAttribute("enzyme_math", libm_name)) + push!(attrs, StringAttribute("enzyme_shouldrecompute")) + push!(attrs, StringAttribute("enzyme_preserve_primal", "*")) + push!(attrs, EnumAttribute("willreturn")) + push!(attrs, EnumAttribute("nosync")) + push!(attrs, EnumAttribute("nofree")) + if LLVM.version().major <= 15 + push!(attrs, EnumAttribute("readnone")) + else + push!(attrs, EnumAttribute("memory", NoEffects.data)) + end + end +end + +# Build a reverse lookup: libm_name => air_name (e.g. "cosf" => "air.cos.f32"). +# Only single-precision entries so we don't duplicate (log => log has f32 + f64). +const GPU_LIBM_TO_AIR_MAP = Dict{String,String}( + v => k for (k, v) in GPU_MATH_INTRINSIC_MAP if endswith(k, ".f32") +) + +# After enzyme! runs, derivative IR may contain calls to LLVM math intrinsics +# (e.g. llvm.cos.f32 for d/dx sin) or libm-style names (e.g. coshf for +# d/dx tanh) that Metal's AIR compiler cannot lower. This combined map +# drives a post-enzyme pass that replaces both categories with AIR equivalents. +const GPU_MATH_POSTENZYME_MAP = let + llvm_to_air = Dict{String,String}( + "llvm.sin.f32" => "air.sin.f32", + "llvm.cos.f32" => "air.cos.f32", + "llvm.tan.f32" => "air.tan.f32", + "llvm.exp.f32" => "air.exp.f32", + "llvm.exp2.f32" => "air.exp2.f32", + "llvm.log.f32" => "air.log.f32", + "llvm.log2.f32" => "air.log2.f32", + "llvm.log10.f32" => "air.log10.f32", + "llvm.pow.f32" => "air.pow.f32", + "llvm.sin.f64" => "air.sin.f64", + "llvm.cos.f64" => "air.cos.f64", + "llvm.exp.f64" => "air.exp.f64", + "llvm.exp2.f64" => "air.exp2.f64", + "llvm.log.f64" => "air.log.f64", + "llvm.log2.f64" => "air.log2.f64", + "llvm.pow.f64" => "air.pow.f64", + ) + # Also include libm → AIR (Enzyme emits e.g. coshf for d/dx tanh) + merge(llvm_to_air, GPU_LIBM_TO_AIR_MAP) +end + +# Replace unresolvable math function references in Enzyme-generated IR with +# their AIR equivalents. Handles both LLVM intrinsics (llvm.cos.f32) and +# libm-style names (coshf) that Enzyme C++ emits in derivative computations. +function replace_gpu_math_references!(mod::LLVM.Module) + for (old_name, air_name) in GPU_MATH_POSTENZYME_MAP + haskey(functions(mod), old_name) || continue + old_fn = functions(mod)[old_name] + isempty(blocks(old_fn)) || continue # external declarations only + ftype = LLVM.function_type(old_fn) + air_fn = if haskey(functions(mod), air_name) + functions(mod)[air_name] + else + air_new = LLVM.Function(mod, air_name, ftype) + attrs = function_attributes(air_new) + push!(attrs, EnumAttribute("willreturn")) + push!(attrs, EnumAttribute("nosync")) + push!(attrs, EnumAttribute("nofree")) + if LLVM.version().major <= 15 + push!(attrs, EnumAttribute("readnone")) + else + push!(attrs, EnumAttribute("memory", NoEffects.data)) + end + air_new + end + for use in collect(LLVM.uses(old_fn)) + inst = LLVM.user(use) + isa(inst, LLVM.CallInst) || continue + LLVM.called_operand(inst) == old_fn || continue + args = collect(operands(inst))[1:end-1] + LLVM.@dispose builder=IRBuilder() begin + position!(builder, inst) + new_call = call!(builder, ftype, air_fn, args) + LLVM.replace_uses!(inst, new_call) + end + LLVM.API.LLVMInstructionEraseFromParent(inst) + end + end +end + +# Create stub definitions for libm functions that Enzyme C++ emits when computing +# derivatives of annotated GPU intrinsics (e.g. d/dx sin(x) = cos(x) → Enzyme +# generates a call to "cosf" which must be resolved in the Metal module). +function create_gpu_math_libm_stubs!(mod::LLVM.Module) + for (libm_name, air_name) in GPU_LIBM_TO_AIR_MAP + haskey(functions(mod), air_name) || continue # AIR fn not in module + haskey(functions(mod), libm_name) && continue # stub already exists + air_fn = functions(mod)[air_name] + isempty(blocks(air_fn)) || continue # must be a declaration + ftype = LLVM.function_type(air_fn) + stub = LLVM.Function(mod, libm_name, ftype) + args = collect(parameters(stub)) + block = BasicBlock(stub, "entry") + LLVM.@dispose builder=IRBuilder() begin + position!(builder, block) + result = call!(builder, ftype, air_fn, args) + ret!(builder, result) + end + # Mark stub as private + alwaysinline so it disappears after inlining + linkage!(stub, LLVM.API.LLVMPrivateLinkage) + push!(function_attributes(stub), EnumAttribute("alwaysinline", 0)) + end +end + @inline function find_math_method(@nospecialize(func::Type), sparam_vals::Core.SimpleVector) if func ∈ keys(known_ops) name, arity, toinject = known_ops[func] @@ -5595,6 +5763,7 @@ function GPUCompiler.compile_unhooked(output::Symbol, job::CompilerJob{<:EnzymeT # target_machine = JIT.get_tm() # else target_machine = GPUCompiler.llvm_machine(job.config.target) + target_info = GPUCompiler.llvm_targetinfo(job.config.target) parallel = false process_module = false @@ -5629,7 +5798,7 @@ function GPUCompiler.compile_unhooked(output::Symbol, job::CompilerJob{<:EnzymeT end # Run early pipeline - optimize!(mod, target_machine) + optimize!(mod, target_machine, target_info) if process_module GPUCompiler.optimize_module!(primal_job, mod) @@ -6008,6 +6177,15 @@ end # Generate the adjoint memcpy_alloca_to_loadstore(mod) force_recompute!(mod) + # Annotate GPU math intrinsics (e.g. air.log.f32 for Metal) so the + # Enzyme C++ core can apply its built-in derivative rules. Must run + # after all inlining passes so bare AIR declarations are visible. + annotate_gpu_math_intrinsics!(mod) + # Provide stub definitions for libm names (e.g. cosf, sinf) that + # Enzyme C++ emits when computing derivatives of the annotated + # intrinsics. Without stubs, symbols like "cosf" are unresolved in + # the Metal module and the AIR compiler rejects the kernel. + create_gpu_math_libm_stubs!(mod) API.EnzymeDetectReadonlyOrThrow(mod) adjointf, augmented_primalf, TapeType = enzyme!( @@ -6035,6 +6213,10 @@ end LLVM.link!(mod, otherMod) end empty!(enzyme_context.modules_to_link) + # Replace LLVM math intrinsics (e.g. llvm.cos.f32) and libm-style + # names (e.g. coshf for d/dx tanh) that Enzyme C++ emits in derivative + # code with their AIR equivalents for Metal targets. + replace_gpu_math_references!(mod) toremove = String[] # Inline the wrapper for f in functions(mod) @@ -6184,7 +6366,7 @@ end if !(primal_target isa GPUCompiler.NativeCompilerTarget) reinsert_gcmarker!(adjointf) augmented_primalf !== nothing && reinsert_gcmarker!(augmented_primalf) - post_optimize!(mod, target_machine, false) #=machine=# + post_optimize!(mod, target_machine, false; tti=target_info) #=machine=# end adjointf = functions(mod)[adjointf_name] diff --git a/src/compiler/optimize.jl b/src/compiler/optimize.jl index ee145ed35b..3f4c7a805e 100644 --- a/src/compiler/optimize.jl +++ b/src/compiler/optimize.jl @@ -30,8 +30,9 @@ SafeAtomicToRegularStorePass() = NewPMFunctionPass("safe_atomic_to_regular_store Addr13NoAliasPass() = NewPMModulePass("addr13_noalias", addr13NoAlias) RemoveAlwaysInlineRootsPass() = NewPMModulePass("remove_alwaysinline_roots", remove_alwaysinline_roots!) -function optimize!(mod::LLVM.Module, tm::Union{LLVM.TargetMachine, Nothing}) +function optimize!(mod::LLVM.Module, tm::Union{LLVM.TargetMachine, Nothing}, tti=nothing) @dispose pb = NewPMPassBuilder() begin + tti === nothing || LLVM.target_transform_info!(pb, tti) registerEnzymeAndPassPipeline!(pb) register!(pb, Addr13NoAliasPass()) register!(pb, RestoreAllocaType()) @@ -67,6 +68,7 @@ function optimize!(mod::LLVM.Module, tm::Union{LLVM.TargetMachine, Nothing}) # Globalopt is separated as it can delete functions, which invalidates the Julia hardcoded pointers to # known functions @dispose pb = NewPMPassBuilder() begin + tti === nothing || LLVM.target_transform_info!(pb, tti) add!(pb, NewPMAAManager()) do aam add!(aam, ScopedNoAliasAA()) add!(aam, TypeBasedAA()) @@ -84,6 +86,7 @@ function optimize!(mod::LLVM.Module, tm::Union{LLVM.TargetMachine, Nothing}) function middle_optimize!(second_stage=false) @dispose pb = NewPMPassBuilder() begin + tti === nothing || LLVM.target_transform_info!(pb, tti) registerEnzymeAndPassPipeline!(pb) register!(pb, RestoreAllocaType()) add!(pb, NewPMAAManager()) do aam @@ -184,6 +187,7 @@ function optimize!(mod::LLVM.Module, tm::Union{LLVM.TargetMachine, Nothing}) # Globalopt is separated as it can delete functions, which invalidates the Julia hardcoded pointers to # known functions @dispose pb = NewPMPassBuilder() begin + tti === nothing || LLVM.target_transform_info!(pb, tti) add!(pb, NewPMAAManager()) do aam add!(aam, ScopedNoAliasAA()) add!(aam, TypeBasedAA()) @@ -197,9 +201,9 @@ function optimize!(mod::LLVM.Module, tm::Union{LLVM.TargetMachine, Nothing}) end run!(pb, mod, tm) end - + run!(GCInvariantVerifierPass(strong=false), mod) - + removeDeadArgs!(mod, tm, #=post_gc_fixup=#false) run!(GCInvariantVerifierPass(strong=false), mod) @@ -395,9 +399,9 @@ end const DumpPreCallConv = Ref(false) const DumpPostCallConv = Ref(false) -function fixup_callconv!(mod::LLVM.Module, tm::Union{LLVM.TargetMachine, Nothing}) +function fixup_callconv!(mod::LLVM.Module, tm::Union{LLVM.TargetMachine, Nothing}, tti=nothing) addr13NoAlias(mod) - + removeDeadArgs!(mod, tm, #=post_gc_fixup=#false) memcpy_sret_split!(mod) @@ -408,6 +412,7 @@ function fixup_callconv!(mod::LLVM.Module, tm::Union{LLVM.TargetMachine, Nothing run!(InstCombinePass(), mod) # GVN actually forwards @dispose pb = NewPMPassBuilder() begin + tti === nothing || LLVM.target_transform_info!(pb, tti) registerEnzymeAndPassPipeline!(pb) add!(pb, SimpleGVNPass()) run!(pb, mod, tm) @@ -418,6 +423,7 @@ function fixup_callconv!(mod::LLVM.Module, tm::Union{LLVM.TargetMachine, Nothing end @dispose pb = NewPMPassBuilder() begin + tti === nothing || LLVM.target_transform_info!(pb, tti) registerEnzymeAndPassPipeline!(pb) add!(pb, "enzyme-fixup-batched-julia") if VERSION < v"1.12" @@ -456,9 +462,9 @@ function fixup_callconv!(mod::LLVM.Module, tm::Union{LLVM.TargetMachine, Nothing return end -function post_optimize!(mod::LLVM.Module, tm::Union{LLVM.TargetMachine, Nothing}, machine::Bool = true; callconv::Bool = true) +function post_optimize!(mod::LLVM.Module, tm::Union{LLVM.TargetMachine, Nothing}, machine::Bool = true; callconv::Bool = true, tti=nothing) if callconv - fixup_callconv!(mod, tm) + fixup_callconv!(mod, tm, tti) end for f in functions(mod) @@ -479,6 +485,7 @@ function post_optimize!(mod::LLVM.Module, tm::Union{LLVM.TargetMachine, Nothing} removeDeadArgs!(mod, tm, #=post_gc_fixup=#true) @dispose pb = NewPMPassBuilder() begin + tti === nothing || LLVM.target_transform_info!(pb, tti) registerEnzymeAndPassPipeline!(pb) register!(pb, ReinsertGCMarkerPass()) register!(pb, SafeAtomicToRegularStorePass()) diff --git a/src/internal_rules/math.jl b/src/internal_rules/math.jl index cb05effa18..83d4f362ff 100644 --- a/src/internal_rules/math.jl +++ b/src/internal_rules/math.jl @@ -100,4 +100,3 @@ function EnzymeRules.reverse( dxs = map(x -> _hypotreverse(x, w, dret, n), xs) return (dx, dy, dz, dxs...) end - From f13c5da90869cd61d9e49adfccc26c1298f3548b Mon Sep 17 00:00:00 2001 From: pw0908 Date: Wed, 24 Jun 2026 00:46:35 -0700 Subject: [PATCH 3/5] Update pipeline to test Metal --- .buildkite/pipeline.yml | 54 +++++++++++++++++++++-------------------- 1 file changed, 28 insertions(+), 26 deletions(-) diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 05af54d9a3..43079ce16c 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -51,29 +51,31 @@ steps: # env: # JULIA_PKG_SERVER_REGISTRY_PREFERENCE: eager # -# - label: "Metal Julia v{{matrix.version}}" -# matrix: -# setup: -# version: -# - "1.10" -# plugins: -# - JuliaCI/julia#v1: -# version: "{{matrix.version}}" -# agents: -# queue: "juliaecosystem" -# os: "macos" -# arch: "aarch64" -# if: build.message !~ /\[skip tests\]/ -# timeout_in_minutes: 60 -# commands: | -# echo "--- Setup Julia packages" -# julia --color=yes -e ' -# using Pkg -# pkgs = [PackageSpec(; path) for path in (".", "lib/EnzymeCore", "lib/EnzymeTestUtils")] -# push!(pkgs, PackageSpec(; name="Metal")) -# Pkg.develop(pkgs)' || exit 3 -# -# echo "+++ Run tests" -# julia --color=yes test/metal.jl -# env: -# JULIA_PKG_SERVER_REGISTRY_PREFERENCE: eager + - label: "Metal Julia v{{matrix.version}}" + matrix: + setup: + version: + - "1.10" + - "1.11" + - "1.12" + plugins: + - JuliaCI/julia#v1: + version: "{{matrix.version}}" + agents: + queue: "juliaecosystem" + os: "macos" + arch: "aarch64" + if: build.message !~ /\[skip tests\]/ + timeout_in_minutes: 60 + commands: | + echo "--- Setup Julia packages" + julia --color=yes -e ' + using Pkg + pkgs = [PackageSpec(; path) for path in (".", "lib/EnzymeCore", "lib/EnzymeTestUtils")] + push!(pkgs, PackageSpec(; name="Metal")) + Pkg.develop(pkgs)' || exit 3 + + echo "+++ Run tests" + julia --color=yes test/metal.jl + env: + JULIA_PKG_SERVER_REGISTRY_PREFERENCE: eager From 20f0c8971e8be1d59367ad5377185e4e74b8d3da Mon Sep 17 00:00:00 2001 From: Valentin Churavy Date: Wed, 24 Jun 2026 10:03:23 +0200 Subject: [PATCH 4/5] Apply suggestion from @vchuravy --- .buildkite/pipeline.yml | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 43079ce16c..da495adc60 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -62,9 +62,7 @@ steps: - JuliaCI/julia#v1: version: "{{matrix.version}}" agents: - queue: "juliaecosystem" - os: "macos" - arch: "aarch64" + queue: "metal" if: build.message !~ /\[skip tests\]/ timeout_in_minutes: 60 commands: | From 2bf19127e8569926e3fa3ac72c9e771a0b8d8fa0 Mon Sep 17 00:00:00 2001 From: pw0908 Date: Wed, 24 Jun 2026 16:22:31 -0700 Subject: [PATCH 5/5] Changing logic for tti --- src/compiler/optimize.jl | 28 +++++++++++++++++++++------- 1 file changed, 21 insertions(+), 7 deletions(-) diff --git a/src/compiler/optimize.jl b/src/compiler/optimize.jl index 3f4c7a805e..9f9f73892f 100644 --- a/src/compiler/optimize.jl +++ b/src/compiler/optimize.jl @@ -32,7 +32,9 @@ RemoveAlwaysInlineRootsPass() = NewPMModulePass("remove_alwaysinline_roots", rem function optimize!(mod::LLVM.Module, tm::Union{LLVM.TargetMachine, Nothing}, tti=nothing) @dispose pb = NewPMPassBuilder() begin - tti === nothing || LLVM.target_transform_info!(pb, tti) + if tti !== nothing + LLVM.target_transform_info!(pb, tti) + end registerEnzymeAndPassPipeline!(pb) register!(pb, Addr13NoAliasPass()) register!(pb, RestoreAllocaType()) @@ -68,7 +70,9 @@ function optimize!(mod::LLVM.Module, tm::Union{LLVM.TargetMachine, Nothing}, tti # Globalopt is separated as it can delete functions, which invalidates the Julia hardcoded pointers to # known functions @dispose pb = NewPMPassBuilder() begin - tti === nothing || LLVM.target_transform_info!(pb, tti) + if tti !== nothing + LLVM.target_transform_info!(pb, tti) + end add!(pb, NewPMAAManager()) do aam add!(aam, ScopedNoAliasAA()) add!(aam, TypeBasedAA()) @@ -86,7 +90,9 @@ function optimize!(mod::LLVM.Module, tm::Union{LLVM.TargetMachine, Nothing}, tti function middle_optimize!(second_stage=false) @dispose pb = NewPMPassBuilder() begin - tti === nothing || LLVM.target_transform_info!(pb, tti) + if tti !== nothing + LLVM.target_transform_info!(pb, tti) + end registerEnzymeAndPassPipeline!(pb) register!(pb, RestoreAllocaType()) add!(pb, NewPMAAManager()) do aam @@ -187,7 +193,9 @@ function optimize!(mod::LLVM.Module, tm::Union{LLVM.TargetMachine, Nothing}, tti # Globalopt is separated as it can delete functions, which invalidates the Julia hardcoded pointers to # known functions @dispose pb = NewPMPassBuilder() begin - tti === nothing || LLVM.target_transform_info!(pb, tti) + if tti !== nothing + LLVM.target_transform_info!(pb, tti) + end add!(pb, NewPMAAManager()) do aam add!(aam, ScopedNoAliasAA()) add!(aam, TypeBasedAA()) @@ -412,7 +420,9 @@ function fixup_callconv!(mod::LLVM.Module, tm::Union{LLVM.TargetMachine, Nothing run!(InstCombinePass(), mod) # GVN actually forwards @dispose pb = NewPMPassBuilder() begin - tti === nothing || LLVM.target_transform_info!(pb, tti) + if tti !== nothing + LLVM.target_transform_info!(pb, tti) + end registerEnzymeAndPassPipeline!(pb) add!(pb, SimpleGVNPass()) run!(pb, mod, tm) @@ -423,7 +433,9 @@ function fixup_callconv!(mod::LLVM.Module, tm::Union{LLVM.TargetMachine, Nothing end @dispose pb = NewPMPassBuilder() begin - tti === nothing || LLVM.target_transform_info!(pb, tti) + if tti !== nothing + LLVM.target_transform_info!(pb, tti) + end registerEnzymeAndPassPipeline!(pb) add!(pb, "enzyme-fixup-batched-julia") if VERSION < v"1.12" @@ -485,7 +497,9 @@ function post_optimize!(mod::LLVM.Module, tm::Union{LLVM.TargetMachine, Nothing} removeDeadArgs!(mod, tm, #=post_gc_fixup=#true) @dispose pb = NewPMPassBuilder() begin - tti === nothing || LLVM.target_transform_info!(pb, tti) + if tti !== nothing + LLVM.target_transform_info!(pb, tti) + end registerEnzymeAndPassPipeline!(pb) register!(pb, ReinsertGCMarkerPass()) register!(pb, SafeAtomicToRegularStorePass())